diff --git a/tests/_strategies.py b/tests/_strategies.py index 318b342..a12421f 100644 --- a/tests/_strategies.py +++ b/tests/_strategies.py @@ -16,7 +16,6 @@ DrawFn, SearchStrategy, booleans, - builds, composite, datetimes, floats, @@ -28,6 +27,7 @@ import zodipy from zodipy._line_of_sight import COMPONENT_CUTOFFS +from zodipy.model_registry import model_registry MIN_FREQ = u.Quantity(10, u.GHz) MAX_FREQ = u.Quantity(0.1, u.micron).to(u.GHz, equivalencies=u.spectral()) @@ -151,17 +151,18 @@ def angles(draw: DrawFn, lonlat: bool = False) -> tuple[u.Quantity[u.deg], u.Qua @composite -def freqs(draw: DrawFn, model: zodipy.Zodipy) -> u.Quantity[u.GHz] | u.Quantity[u.micron]: - if model.extrapolate: +def freqs( + draw: DrawFn, + min_freq: u.Quantity, + max_freq: u.Quantity, + extrapolate: bool, +) -> u.Quantity[u.GHz] | u.Quantity[u.micron]: + if extrapolate: return draw(sampled_from(FREQ_LOG_RANGE).map(np.exp).map(partial(u.Quantity, unit=u.GHz))) - min_freq = model._ipd_model.spectrum[0] - max_freq = model._ipd_model.spectrum[-1] freq_range = np.geomspace(np.log(min_freq.value), np.log(max_freq.value), N_FREQS) freq_strategy = ( - sampled_from(freq_range.tolist()) - .map(np.exp) - .map(partial(u.Quantity, unit=model._ipd_model.spectrum.unit)) + sampled_from(freq_range.tolist()).map(np.exp).map(partial(u.Quantity, unit=min_freq.unit)) ) return np.clip(draw(freq_strategy), min_freq, max_freq) @@ -241,18 +242,32 @@ def any_obs(draw: DrawFn, model: zodipy.Zodipy) -> str: return draw(sampled_from(model.supported_observers)) -MODEL_STRATEGY_MAPPINGS: dict[str, SearchStrategy[Any]] = { - "model": sampled_from(AVAILABLE_MODELS), - "gauss_quad_degree": integers(min_value=1, max_value=200), - "extrapolate": booleans(), -} - - @composite def zodipy_models(draw: DrawFn, **static_params: dict[str, Any]) -> zodipy.Zodipy: - strategies = MODEL_STRATEGY_MAPPINGS.copy() - for key in static_params: - if key in strategies: - strategies.pop(key) + extrapolate = static_params.pop("extrapolate", draw(booleans())) + model = static_params.pop("model", draw(sampled_from(AVAILABLE_MODELS))) + ipd_model = model_registry.get_model(model) + min_freq = ipd_model.spectrum.min() + max_freq = ipd_model.spectrum.max() + + do_bp = static_params.pop("bandpass_integrate", None) + if do_bp is not None: + frequencies = draw(random_freqs(bandpass=True)) + w = draw(weights(frequencies)) + else: + frequencies = static_params.pop( + "freq", draw(freqs(min=min_freq, max=max_freq, extrapolate=extrapolate)) + ) + w = None - return draw(builds(partial(zodipy.Zodipy, **static_params), **strategies)) + gauss_quad_degree = static_params.pop( + "gauss_quad_degree", draw(integers(min_value=1, max_value=200)) + ) + + return zodipy.Zodipy( + freq=frequencies, + model=model, + weights=w, + gauss_quad_degree=gauss_quad_degree, + extrapolate=extrapolate, + ) diff --git a/tests/test_get_emission.py b/tests/test_get_emission.py index e772b29..bad1cfa 100644 --- a/tests/test_get_emission.py +++ b/tests/test_get_emission.py @@ -16,15 +16,12 @@ from ._strategies import ( angles, coords_in, - freqs, healpixes, obs_positions, pixels, quantities, - random_freqs, sky_coords, times, - weights, zodipy_models, ) from ._tabulated_dirbe import DAYS, LAT, LON, TABULATED_DIRBE_EMISSION @@ -41,10 +38,8 @@ def test_get_emission_skycoord( ) -> None: """Property test for get_emission_skycoord.""" observer = data.draw(obs_positions(model, coordinates.obstime)) - frequency = data.draw(freqs(model)) emission = model.get_emission_skycoord( coordinates, - freq=frequency, obs_pos=observer, ) assert emission.size == coordinates.size @@ -60,12 +55,10 @@ def test_get_binned_skycoord( ) -> None: """Property test for get_binned_emission_pix.""" observer = data.draw(obs_positions(model, coordinates.obstime)) - frequency = data.draw(freqs(model)) cut_solar = data.draw(booleans()) emission_binned = model.get_binned_emission_skycoord( coordinates, - freq=frequency, obs_pos=observer, nside=healpix.nside, solar_cut=data.draw(quantities(20, 50, u.deg)) if cut_solar else None, @@ -84,12 +77,10 @@ def test_get_emission_pix( ) -> None: """Property test for get_emission_pix.""" observer = data.draw(obs_positions(model, time)) - frequency = data.draw(freqs(model)) pix = data.draw(pixels(healpix.nside)) emission = model.get_emission_pix( pix, nside=healpix.nside, - freq=frequency, obs_time=time, obs_pos=observer, coord_in=coord_in, @@ -108,13 +99,11 @@ def test_get_binned_emission_pix( ) -> None: """Property test for get_binned_emission_pix.""" observer = data.draw(obs_positions(model, time)) - frequency = data.draw(freqs(model)) pix = data.draw(pixels(healpix.nside)) cut_solar = data.draw(booleans()) emission_binned = model.get_binned_emission_pix( pix, nside=healpix.nside, - freq=frequency, obs_time=time, obs_pos=observer, solar_cut=data.draw(quantities(20, 50, u.deg)) if cut_solar else None, @@ -136,12 +125,10 @@ def test_get_emission_ang( theta, phi = angles observer = data.draw(obs_positions(model, time)) - frequency = data.draw(freqs(model)) emission = model.get_emission_ang( theta, phi, - freq=frequency, obs_time=time, obs_pos=observer, coord_in=coord_in, @@ -163,14 +150,12 @@ def test_get_binned_emission_ang( theta, phi = angles observer = data.draw(obs_positions(model, time)) - frequency = data.draw(freqs(model)) cut_solar = data.draw(booleans()) emission_binned = model.get_binned_emission_ang( theta, phi, nside=healpix.nside, - freq=frequency, obs_time=time, obs_pos=observer, coord_in=coord_in, @@ -179,60 +164,16 @@ def test_get_binned_emission_ang( assert emission_binned.shape == (healpix.npix,) -@given(zodipy_models(extrapolate=False), times(), angles(lonlat=True), healpixes(), data()) -@settings(deadline=None) -def test_invalid_freq( - model: Zodipy, - time: Time, - angles: tuple[u.Quantity[u.deg], u.Quantity[u.deg]], - healpix: hp.HEALPix, - data: DataObject, -) -> None: +def test_invalid_freq() -> None: """Property test checking for unsupported spectral range.""" - theta, phi = angles - observer = data.draw(obs_positions(model, time)) - pix = data.draw(pixels(healpix.nside)) - - freq = data.draw(random_freqs(unit=model._ipd_model.spectrum.unit)) - if not (model._ipd_model.spectrum[0] <= freq <= model._ipd_model.spectrum[-1]): - with pytest.raises(ValueError): - model.get_emission_pix( - pix, - freq=freq, - nside=healpix.nside, - obs_time=time, - obs_pos=observer, - ) - - with pytest.raises(ValueError): - model.get_emission_ang( - theta, - phi, - freq=freq, - obs_time=time, - obs_pos=observer, - lonlat=True, - ) - - with pytest.raises(ValueError): - model.get_binned_emission_pix( - pix, - nside=healpix.nside, - freq=freq, - obs_time=time, - obs_pos=observer, - ) - - with pytest.raises(ValueError): - model.get_binned_emission_ang( - theta, - phi, - nside=healpix.nside, - freq=freq, - obs_time=time, - obs_pos=observer, - lonlat=True, - ) + with pytest.raises(ValueError): + Zodipy(freq=0.4 * u.micron, model="DIRBE", extrapolate=False) + with pytest.raises(ValueError): + Zodipy(freq=500 * u.micron, model="DIRBE", extrapolate=False) + with pytest.raises(ValueError): + Zodipy(freq=10 * u.GHz, model="Planck2018", extrapolate=False) + with pytest.raises(ValueError): + Zodipy(freq=1000 * u.micron, model="Planck2018", extrapolate=False) def test_compare_to_dirbe_idl() -> None: @@ -241,15 +182,14 @@ def test_compare_to_dirbe_idl() -> None: Zodipy should be able to reproduce the tabulated emission from the DIRBE Zoidacal Light Prediction Software with a maximum difference of 0.1%. """ - model = Zodipy("dirbe") for frequency, tabulated_emission in TABULATED_DIRBE_EMISSION.items(): + model = Zodipy(freq=frequency * u.micron, model="dirbe") for idx, (day, lon, lat) in enumerate(zip(DAYS, LON, LAT)): time = DIRBE_START_DAY + TimeDelta(day - 1, format="jd") emission = model.get_emission_ang( lon * u.deg, lat * u.deg, - freq=frequency * u.micron, lonlat=True, obs_pos="earth", obs_time=time, @@ -263,12 +203,11 @@ def test_multiprocessing() -> None: Tests that model with multiprocessing enabled returns the same value as without multiprocessing. """ - model = Zodipy() - model_parallel = Zodipy(n_proc=4) + model = Zodipy(freq=78 * u.micron) + model_parallel = Zodipy(freq=78 * u.micron, n_proc=4) observer = "earth" time = Time("2020-01-01") - frequency = 78 * u.micron healpix = hp.HEALPix(32) pix = np.random.randint(0, healpix.npix, size=1000) theta = np.random.uniform(0, 180, size=1000) * u.deg @@ -277,14 +216,12 @@ def test_multiprocessing() -> None: emission_pix = model.get_emission_pix( pix, nside=healpix.nside, - freq=frequency, obs_time=time, obs_pos=observer, ) emission_pix_parallel = model_parallel.get_emission_pix( pix, nside=healpix.nside, - freq=frequency, obs_time=time, obs_pos=observer, ) @@ -293,14 +230,12 @@ def test_multiprocessing() -> None: emission_binned_pix = model.get_binned_emission_pix( pix, nside=healpix.nside, - freq=frequency, obs_time=time, obs_pos=observer, ) emission_binned_pix_parallel = model_parallel.get_binned_emission_pix( pix, nside=healpix.nside, - freq=frequency, obs_time=time, obs_pos=observer, ) @@ -309,14 +244,12 @@ def test_multiprocessing() -> None: emission_ang = model.get_emission_ang( theta, phi, - freq=frequency, obs_time=time, obs_pos=observer, ) emission_ang_parallel = model_parallel.get_emission_ang( theta, phi, - freq=frequency, obs_time=time, obs_pos=observer, ) @@ -325,7 +258,6 @@ def test_multiprocessing() -> None: emission_binned_ang = model.get_binned_emission_ang( theta, phi, - freq=frequency, nside=healpix.nside, obs_time=time, obs_pos=observer, @@ -334,7 +266,6 @@ def test_multiprocessing() -> None: emission_binned_ang_parallel = model_parallel.get_binned_emission_ang( theta, phi, - freq=frequency, nside=healpix.nside, obs_time=time, obs_pos=observer, @@ -345,26 +276,24 @@ def test_multiprocessing() -> None: def test_inner_radial_cutoff_multiprocessing() -> None: """Testing that model with inner radial cutoffs can be parallelized.""" - model = Zodipy("RRM-experimental") - model_parallel = Zodipy("RRM-experimental", n_proc=4) + frequency = 78 * u.micron + model = Zodipy(freq=frequency, model="RRM-experimental") + model_parallel = Zodipy(freq=frequency, model="RRM-experimental", n_proc=4) observer = "earth" time = Time("2020-01-01") - frequency = 78 * u.micron healpix = hp.HEALPix(32) pix = np.random.randint(0, healpix.npix, size=1000) emission_pix = model.get_emission_pix( pix, nside=healpix.nside, - freq=frequency, obs_time=time, obs_pos=observer, ) emission_pix_parallel = model_parallel.get_emission_pix( pix, nside=healpix.nside, - freq=frequency, obs_time=time, obs_pos=observer, ) @@ -372,11 +301,10 @@ def test_inner_radial_cutoff_multiprocessing() -> None: @given( - zodipy_models(extrapolate=True), + zodipy_models(extrapolate=True, bandpass_integrate=True), times(), healpixes(), angles(), - random_freqs(bandpass=True), data(), ) @settings(deadline=None) @@ -385,18 +313,14 @@ def test_bandpass_integration( time: Time, healpix: hp.HEALPix, angles: tuple[u.Quantity[u.deg], u.Quantity[u.deg]], - freqs: u.Quantity[u.Hz] | u.Quantity[u.micron], data: DataObject, ) -> None: """Property test for bandpass integrations.""" theta, phi = angles observer = data.draw(obs_positions(model, time)) - bp_weights = data.draw(weights(freqs)) emission_binned = model.get_binned_emission_ang( theta, phi, - freq=freqs, - weights=bp_weights, nside=healpix.nside, obs_time=time, obs_pos=observer, @@ -405,11 +329,10 @@ def test_bandpass_integration( @given( - zodipy_models(extrapolate=True), + zodipy_models(extrapolate=True, bandpass_integrate=True), times(), healpixes(), angles(), - random_freqs(bandpass=True), data(), ) @settings(deadline=None) @@ -418,19 +341,15 @@ def test_weights( time: Time, healpix: hp.HEALPix, angles: tuple[u.Quantity[u.deg], u.Quantity[u.deg]], - freqs: u.Quantity[u.Hz] | u.Quantity[u.micron], data: DataObject, ) -> None: """Property test for bandpass weights.""" theta, phi = angles observer = data.draw(obs_positions(model, time)) - bp_weights = data.draw(weights(freqs)) model.get_binned_emission_ang( theta, phi, - weights=bp_weights, - freq=freqs, nside=healpix.nside, obs_time=time, obs_pos=observer, @@ -439,7 +358,6 @@ def test_weights( def test_custom_weights() -> None: """Tests bandpass integration with custom weights.""" - model = Zodipy() time = Time("2020-01-01") healpix = hp.HEALPix(64) pix = np.arange(healpix.npix) @@ -448,11 +366,10 @@ def test_custom_weights() -> None: freqs = np.linspace(central_freq - sigma_freq, central_freq + sigma_freq, 100) * u.micron weights = np.random.randn(len(freqs)) weights /= np.trapz(weights, freqs.value) + model = Zodipy(freq=freqs, weights=weights) model.get_emission_pix( pix, - freq=freqs, - weights=weights, nside=healpix.nside, obs_time=time, obs_pos="earth", @@ -461,14 +378,13 @@ def test_custom_weights() -> None: def test_custom_obs_pos() -> None: """Tests a user specified observer position.""" - model = Zodipy() + model = Zodipy(freq=234 * u.micron) time = Time("2020-01-01") healpix = hp.HEALPix(64) pix = np.arange(healpix.npix) model.get_emission_pix( pix, - freq=234 * u.micron, nside=healpix.nside, obs_time=time, obs_pos=[0.1, 0.2, 1] * u.AU, @@ -476,16 +392,14 @@ def test_custom_obs_pos() -> None: model.get_emission_pix( pix, - freq=234 * u.micron, nside=healpix.nside, obs_time=time, obs_pos=[2, 0.1, 4] * u.AU, ) - with pytest.raises(TypeError): + with pytest.raises(u.UnitConversionError): model.get_emission_pix( pix, - freq=234 * u.micron, nside=healpix.nside, obs_time=time, obs_pos=[2, 0.1, 4], @@ -494,7 +408,6 @@ def test_custom_obs_pos() -> None: with pytest.raises(u.UnitsError): model.get_emission_pix( pix, - freq=234 * u.micron, nside=healpix.nside, obs_time=time, obs_pos=[2, 0.1, 4] * u.s, @@ -503,7 +416,7 @@ def test_custom_obs_pos() -> None: def test_pixel_ordering() -> None: """Tests that the pixel related functions work for both healpix orderings.""" - model = Zodipy() + model = Zodipy(freq=234 * u.micron) time = Time("2020-01-01") healpix_ring = hp.HEALPix(32) @@ -517,14 +430,12 @@ def test_pixel_ordering() -> None: model.get_emission_pix( pix_ring, - freq=234 * u.micron, nside=healpix_ring.nside, obs_time=time, order="ring", ) model.get_emission_pix( pix_nested, - freq=234 * u.micron, nside=healpix_nested.nside, obs_time=time, order="nested", @@ -532,14 +443,12 @@ def test_pixel_ordering() -> None: model.get_binned_emission_pix( pix_ring, - freq=234 * u.micron, nside=healpix_ring.nside, obs_time=time, order="ring", ) model.get_binned_emission_pix( pix_nested, - freq=234 * u.micron, nside=healpix_nested.nside, obs_time=time, order="nested", @@ -549,7 +458,6 @@ def test_pixel_ordering() -> None: lon, lat, lonlat=True, - freq=234 * u.micron, nside=healpix_ring.nside, obs_time=time, order="ring", @@ -558,7 +466,6 @@ def test_pixel_ordering() -> None: lon, lat, lonlat=True, - freq=234 * u.micron, nside=healpix_nested.nside, obs_time=time, order="nested", @@ -573,13 +480,11 @@ def test_pixel_ordering() -> None: model.get_binned_emission_skycoord( sky_coord, - freq=234 * u.micron, nside=healpix_ring.nside, order="ring", ) model.get_binned_emission_skycoord( sky_coord, - freq=234 * u.micron, nside=healpix_nested.nside, order="nested", ) diff --git a/tests/test_tabulated_density.py b/tests/test_tabulated_density.py index fbb5248..60559d2 100644 --- a/tests/test_tabulated_density.py +++ b/tests/test_tabulated_density.py @@ -11,7 +11,7 @@ @given( - zodipy_models(), + zodipy_models(model="DIRBE"), integers(min_value=10, max_value=200), floats(min_value=2, max_value=10), floats(min_value=2, max_value=10), @@ -34,7 +34,7 @@ def test_tabulated_density( grid_array = np.asarray(grid_regular) grid = random.choice([grid_array, grid_regular]) - assert tabulate_density(grid, model=model.model).shape == ( + assert tabulate_density(grid, model="DIRBE").shape == ( len(model._ipd_model.comps), n_grid_points, n_grid_points, diff --git a/tests/test_validators.py b/tests/test_validators.py index d429233..76faaf3 100644 --- a/tests/test_validators.py +++ b/tests/test_validators.py @@ -22,25 +22,25 @@ def test_validate_frequencies(model: Zodipy) -> None: get_validated_freq( freq=BANDPASS_FREQUENCIES.value, model=model._ipd_model, - extrapolate=model.extrapolate, + extrapolate=False, ) with pytest.raises(TypeError): get_validated_freq( freq=25, model=model._ipd_model, - extrapolate=model.extrapolate, + extrapolate=False, ) with pytest.raises(units.UnitsError): get_validated_freq( freq=BANDPASS_FREQUENCIES.value * units.g, model=model._ipd_model, - extrapolate=model.extrapolate, + extrapolate=False, ) with pytest.raises(units.UnitsError): get_validated_freq( freq=25 * units.g, model=model._ipd_model, - extrapolate=model.extrapolate, + extrapolate=False, ) @@ -112,40 +112,35 @@ def test_validate_weights_shape() -> None: def test_extrapolate_raises_error() -> None: """Tests that an error is correctly raised when extrapolation is not allowed.""" with pytest.raises(ValueError): - model = Zodipy("dirbe") - model.get_emission_pix([1, 4, 5], freq=400 * units.micron, nside=32, obs_time=OBS_TIME) + model = Zodipy(freq=400 * units.micron, model="dirbe") + model.get_emission_pix([1, 4, 5], nside=32, obs_time=OBS_TIME) - model = Zodipy("dirbe", extrapolate=True) - model.get_emission_pix([1, 4, 5], freq=400 * units.micron, nside=32, obs_time=OBS_TIME) + model = Zodipy(freq=400 * units.micron, model="dirbe", extrapolate=True) + model.get_emission_pix([1, 4, 5], nside=32, obs_time=OBS_TIME) def test_interp_kind() -> None: """Tests that the interpolation kind can be passed in.""" - model = Zodipy("dirbe", interp_kind="linear") - linear = model.get_emission_pix([1, 4, 5], freq=27 * units.micron, nside=32, obs_time=OBS_TIME) + model = Zodipy(freq=27 * units.micron, model="dirbe", interp_kind="linear") + linear = model.get_emission_pix([1, 4, 5], nside=32, obs_time=OBS_TIME) - model = Zodipy("dirbe", interp_kind="quadratic") - quadratic = model.get_emission_pix( - [1, 4, 5], freq=27 * units.micron, nside=32, obs_time=OBS_TIME - ) + model = Zodipy(freq=27 * units.micron, model="dirbe", interp_kind="quadratic") + quadratic = model.get_emission_pix([1, 4, 5], nside=32, obs_time=OBS_TIME) assert not np.allclose(linear, quadratic) with pytest.raises(NotImplementedError): - model = Zodipy("dirbe", interp_kind="sdfs") - model.get_emission_pix( - pixels=[1, 4, 5], freq=27 * units.micron, nside=32, obs_time=OBS_TIME - ) + model = Zodipy(freq=27 * units.micron, model="dirbe", interp_kind="sdfs") + model.get_emission_pix(pixels=[1, 4, 5], nside=32, obs_time=OBS_TIME) def test_wrong_frame() -> None: """Tests that an error is correctly raised when an incorrect frame is passed in.""" - model = Zodipy() + model = Zodipy(freq=27 * units.micron) with pytest.raises(ValueError): model.get_emission_pix( [1, 4, 5], nside=32, - freq=27 * units.micron, obs_time=OBS_TIME, coord_in="not a valid frame", ) @@ -153,11 +148,10 @@ def test_wrong_frame() -> None: def test_non_quantity_ang_raises_error() -> None: """Tests that an error is correctly raised if the user inputed angles are not Quantities.""" - model = Zodipy() + model = Zodipy(freq=27 * units.micron) with pytest.raises(TypeError): model.get_emission_ang( 32, 12, - freq=27 * units.micron, obs_time=OBS_TIME, ) diff --git a/zodipy/_coords.py b/zodipy/_coords.py index 1807548..c59c026 100644 --- a/zodipy/_coords.py +++ b/zodipy/_coords.py @@ -5,7 +5,7 @@ from zodipy._constants import DISTANCE_FROM_EARTH_TO_SEMB_L2 -__all__ = ["get_earth_skycoord", "get_obs_skycoord", "get_frame_from_string"] +__all__ = ["get_earth_skycoord", "get_obs_skycoord", "string_to_coordinate_frame"] def get_sun_earth_moon_barycenter_skycoord(earth_skycoord: coords.SkyCoord) -> coords.SkyCoord: @@ -58,7 +58,7 @@ def get_obs_skycoord( raise TypeError(msg) from AttributeError -def get_frame_from_string(frame_literal: str) -> type[coords.BaseCoordinateFrame]: +def string_to_coordinate_frame(frame_literal: str) -> type[coords.BaseCoordinateFrame]: """Return the appropriate astropy coordinate frame class from a string literal.""" if frame_literal == "E": return coords.BarycentricMeanEcliptic diff --git a/zodipy/_ipd_model.py b/zodipy/_ipd_model.py index f14fd47..165b608 100644 --- a/zodipy/_ipd_model.py +++ b/zodipy/_ipd_model.py @@ -42,6 +42,10 @@ def to_dict(self) -> dict: return _dict + @property + def ncomps(self) -> int: + return len(self.comps) + @dataclass class Kelsall(InterplanetaryDustModel): diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index 07d2fce..44f3af4 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -13,7 +13,7 @@ from zodipy._bandpass import get_bandpass_interpolation_table, validate_and_get_bandpass from zodipy._constants import SPECIFIC_INTENSITY_UNITS -from zodipy._coords import get_earth_skycoord, get_frame_from_string, get_obs_skycoord +from zodipy._coords import get_earth_skycoord, get_obs_skycoord, string_to_coordinate_frame from zodipy._emission import EMISSION_MAPPING from zodipy._interpolate_source import SOURCE_PARAMS_MAPPING from zodipy._ipd_comps import ComponentLabel @@ -31,36 +31,12 @@ class Zodipy: - """Interface for simulating zodiacal emission. - - This class provides methods for simulating zodiacal emission given observer pointing - either in sky angles or through HEALPix pixels. - - Args: - freq: Delta frequency/wavelength or a sequence of frequencies corresponding to - a bandpass over which to evaluate the zodiacal emission. The frequencies - must be strictly increasing. - weights: Bandpass weights corresponding the the frequencies in `freq`. The weights - are assumed to be given in spectral radiance units (Jy/sr). - model (str): Name of the interplanetary dust model to use in the simulations. - Defaults to DIRBE. - gauss_quad_degree (int): Order of the Gaussian-Legendre quadrature used to evaluate - the line-of-sight integral in the simulations. Default is 50 points. - interp_kind (str): Interpolation kind used in `scipy.interpolate.interp1d` to interpolate - spectral paramters (see [Scipy documentation]( - https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.interp1d.html)). - Defaults to 'linear'. - extrapolate (bool): If `True` all spectral quantities in the selected model are - extrapolated to the requested frequencies or wavelengths. If `False`, an - exception is raised on requested frequencies/wavelengths outside of the - valid model range. Default is `False`. - ephemeris (str): Ephemeris used in `astropy.coordinates.solar_system_ephemeris` to compute - the positions of the observer and the Earth. Defaults to 'builtin'. See the - [Astropy documentation](https://docs.astropy.org/en/stable/coordinates/solarsystem.html) - for available ephemerides. - n_proc (int): Number of cores to use. If `n_proc` is greater than 1, the line-of-sight - integrals are parallelized using the `multiprocessing` module. Defaults to 1. + """Main interface to ZodiPy. + The zodiacal light simulations are configured by specifying a bandpass (`freq`, `weights`) + or a delta/center frequency (`freq`), and a string representation of a built in interplanetary + dust model (`model`). See https://cosmoglobe.github.io/zodipy/introduction/ for a list of + available models. """ def __init__( @@ -74,16 +50,40 @@ def __init__( ephemeris: str = "builtin", n_proc: int = 1, ) -> None: - self.model = model - self.gauss_quad_degree = gauss_quad_degree - self.extrapolate = extrapolate + """Initialize the Zodipy interface. + + Args: + freq: Delta frequency/wavelength or a sequence of frequencies corresponding to + a bandpass over which to evaluate the zodiacal emission. The frequencies + must be strictly increasing. + weights: Bandpass weights corresponding the the frequencies in `freq`. The weights + are assumed to be given in spectral radiance units (Jy/sr). + model (str): Name of the interplanetary dust model to use in the simulations. + Defaults to DIRBE. + gauss_quad_degree (int): Order of the Gaussian-Legendre quadrature used to evaluate + the line-of-sight integral in the simulations. Default is 50 points. + interp_kind (str): Interpolation kind used in `scipy.interpolate.interp1d` to + interpolate spectral paramters (see [Scipy documentation]( + https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.interp1d.html)). + Defaults to 'linear'. + extrapolate (bool): If `True` all spectral quantities in the selected model are + extrapolated to the requested frequencies or wavelengths. If `False`, an + exception is raised on requested frequencies/wavelengths outside of the + valid model range. Default is `False`. + ephemeris (str): Ephemeris used in `astropy.coordinates.solar_system_ephemeris` to + compute the positions of the observer and the Earth. Defaults to 'builtin'. See the + [Astropy documentation](https://docs.astropy.org/en/stable/coordinates/solarsystem.html) + for available ephemerides. + n_proc (int): Number of cores to use. If `n_proc` is greater than 1, the line-of-sight + integrals are parallelized using the `multiprocessing` module. Defaults to 1. + """ self.ephemeris = ephemeris self.n_proc = n_proc self._interpolator = functools.partial( interpolate.interp1d, kind=interp_kind, - fill_value="extrapolate" if self.extrapolate else np.nan, + fill_value="extrapolate" if extrapolate else np.nan, ) self._ipd_model = model_registry.get_model(model) self._gauss_points_and_weights = np.polynomial.legendre.leggauss(gauss_quad_degree) @@ -92,10 +92,10 @@ def __init__( freq=freq, weights=weights, model=self._ipd_model, - extrapolate=self.extrapolate, + extrapolate=extrapolate, ) - self.bandpass_interpolatation_table = get_bandpass_interpolation_table(bandpass) - self.source_parameters = SOURCE_PARAMS_MAPPING[type(self._ipd_model)]( + self._bandpass_interpolatation_table = get_bandpass_interpolation_table(bandpass) + self._source_parameters = SOURCE_PARAMS_MAPPING[type(self._ipd_model)]( bandpass, self._ipd_model, self._interpolator ) @@ -190,7 +190,7 @@ def get_emission_ang( theta, phi = get_validated_ang(theta=theta, phi=phi, lonlat=lonlat) (theta, phi), indicies = np.unique(np.stack([theta, phi]), return_inverse=True, axis=1) - frame = get_frame_from_string(coord_in) + frame = string_to_coordinate_frame(coord_in) coordinates = coords.SkyCoord(theta, phi, frame=frame) return self._compute_emission( @@ -233,7 +233,7 @@ def get_emission_pix( emission: Simulated zodiacal emission in units of 'MJy/sr'. """ - frame = get_frame_from_string(coord_in) + frame = string_to_coordinate_frame(coord_in) healpix = hp.HEALPix(nside=nside, order=order, frame=frame) unique_pixels, indicies = np.unique(pixels, return_inverse=True) coordinates = healpix.healpix_to_skycoord(unique_pixels) @@ -352,7 +352,7 @@ def get_binned_emission_ang( """ theta, phi = get_validated_ang(theta, phi, lonlat=lonlat) - frame = get_frame_from_string(coord_in) + frame = string_to_coordinate_frame(coord_in) healpix = hp.HEALPix(nside, order=order, frame=frame) (theta, phi), counts = np.unique(np.vstack([theta, phi]), return_counts=True, axis=1) coordinates = coords.SkyCoord( @@ -407,7 +407,7 @@ def get_binned_emission_pix( emission: Simulated zodiacal emission in units of 'MJy/sr'. """ - frame = get_frame_from_string(coord_in) + frame = string_to_coordinate_frame(coord_in) healpix = hp.HEALPix(nside=nside, order=order, frame=frame) unique_pixels, counts = np.unique(pixels, return_counts=True) coordinates = healpix.healpix_to_skycoord(unique_pixels) @@ -436,11 +436,13 @@ def _compute_emission( earth_skycoord = get_earth_skycoord(obs_time, ephemeris=self.ephemeris) obs_skycoord = get_obs_skycoord(obs_pos, obs_time, earth_skycoord, ephemeris=self.ephemeris) - # Rotate to ecliptic coordinates to evaluate zodiacal light model coordinates = coordinates.transform_to(coords.BarycentricMeanEcliptic) - # For binned emission, remove all coordinates within the solar cut - if healpix is not None and solar_cut is not None: + bin_output_to_healpix_map = healpix is not None + filter_coords_by_solar_cut = solar_cut is not None + distribute_to_cores = self.n_proc > 1 and coordinates.size > self.n_proc + + if bin_output_to_healpix_map and filter_coords_by_solar_cut: sun_skycoord = coords.SkyCoord( obs_skycoord.spherical.lon + 180 * units.deg, obs_skycoord.spherical.lat, @@ -452,8 +454,6 @@ def _compute_emission( unit_vectors = coordinates.cartesian.xyz.value - # Get the integration limits for each zodiacal component (which may be - # different or the same depending on the model) along all line of sights. start, stop = get_line_of_sight_range( components=self._ipd_model.comps.keys(), unit_vectors=unit_vectors, @@ -472,15 +472,14 @@ def _compute_emission( common_integrand = functools.partial( EMISSION_MAPPING[type(self._ipd_model)], X_obs=obs_skycoord.cartesian.xyz.to_value(units.AU)[:, np.newaxis, np.newaxis], - bp_interpolation_table=self.bandpass_interpolatation_table, - **self.source_parameters["common"], + bp_interpolation_table=self._bandpass_interpolatation_table, + **self._source_parameters["common"], ) - # Parallelize the line-of-sight integrals if more than one processor is used and the - # number of unique observations is greater than the number of processors. - if self.n_proc > 1 and unit_vectors.shape[-1] > self.n_proc: + if distribute_to_cores: unit_vector_chunks = np.array_split(unit_vectors, self.n_proc, axis=-1) - integrated_comp_emission = np.zeros((len(self._ipd_model.comps), unit_vectors.shape[1])) + integrated_comp_emission = np.zeros((self._ipd_model.ncomps, coordinates.size)) + with multiprocessing.get_context(SYS_PROC_START_METHOD).Pool( processes=self.n_proc ) as pool: @@ -497,7 +496,7 @@ def _compute_emission( start=np.expand_dims(start, axis=-1), stop=np.expand_dims(stop, axis=-1), get_density_function=density_partials[comp_label], - **self.source_parameters[comp_label], + **self._source_parameters[comp_label], ) for unit_vectors, start, stop in zip( unit_vector_chunks, start_chunks, stop_chunks @@ -519,7 +518,7 @@ def _compute_emission( ) else: - integrated_comp_emission = np.zeros((len(self._ipd_model.comps), unit_vectors.shape[1])) + integrated_comp_emission = np.zeros((self._ipd_model.ncomps, coordinates.size)) unit_vectors_expanded = np.expand_dims(unit_vectors, axis=-1) for idx, comp_label in enumerate(self._ipd_model.comps.keys()): @@ -529,7 +528,7 @@ def _compute_emission( start=np.expand_dims(start[comp_label], axis=-1), stop=np.expand_dims(stop[comp_label], axis=-1), get_density_function=density_partials[comp_label], - **self.source_parameters[comp_label], + **self._source_parameters[comp_label], ) integrated_comp_emission[idx] = ( @@ -538,13 +537,12 @@ def _compute_emission( * (stop[comp_label] - start[comp_label]) ) - # Output is requested to be binned - if healpix: - emission = np.zeros((len(self._ipd_model.comps), healpix.npix)) - pixels = healpix.skycoord_to_healpix(coordinates) + if bin_output_to_healpix_map: + emission = np.zeros((self._ipd_model.ncomps, healpix.npix)) # type: ignore + pixels = healpix.skycoord_to_healpix(coordinates) # type: ignore emission[:, pixels] = integrated_comp_emission else: - emission = np.zeros((len(self._ipd_model.comps), indicies.size)) + emission = np.zeros((self._ipd_model.ncomps, indicies.size)) emission = integrated_comp_emission[:, indicies] emission = (emission << SPECIFIC_INTENSITY_UNITS).to(units.MJy / units.sr) @@ -554,7 +552,8 @@ def _compute_emission( @property def supported_observers(self) -> list[str]: """Return a list of available observers given an ephemeris.""" - return [*list(coords.solar_system_ephemeris.bodies), "semb-l2"] + with coords.solar_system_ephemeris.set(self.ephemeris): + return [*list(coords.solar_system_ephemeris.bodies), "semb-l2"] def get_parameters(self) -> dict: """Return a dictionary containing the interplanetary dust model parameters."""