From ce4ac118f8bcb70d19274015e1eb4ae460b17c5c Mon Sep 17 00:00:00 2001 From: Metin San Date: Thu, 25 Apr 2024 10:19:47 +0200 Subject: [PATCH 01/17] Move solar_cut argument from the ZodiPy class to the binned function signatures, as these really only make sense in that context (I think). Also, start the migration from healpy to astropy-healpix --- zodipy/zodipy.py | 74 +++++++++++++++++++++++++----------------------- 1 file changed, 38 insertions(+), 36 deletions(-) diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index f7ed6b5..219fa56 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -6,9 +6,10 @@ from typing import TYPE_CHECKING, Callable, Literal, Sequence import astropy.units as u +import astropy_healpix as ahpx import healpy as hp import numpy as np -from astropy.coordinates import solar_system_ephemeris +from astropy.coordinates import SkyCoord, solar_system_ephemeris from scipy.interpolate import interp1d from zodipy._bandpass import get_bandpass_interpolation_table, validate_and_get_bandpass @@ -57,11 +58,6 @@ class Zodipy: Earth. Defaults to 'de432s', which requires downloading (and caching) a ~10MB file. For more information on available ephemeridis, please visit the [Astropy documentation](https://docs.astropy.org/en/stable/coordinates/solarsystem.html) - solar_cut (u.Quantity[u.deg]): Cutoff angle from the sun in degrees. The emission - for all the pointing with angular distance between the sun smaller than - `solar_cutoff` are masked. Defaults to `None`. - solar_cut_fill_value (float): Fill value for pixels masked with `solar_cut`. - Defaults to `np.nan`. 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. @@ -74,8 +70,6 @@ def __init__( extrapolate: bool = False, interp_kind: str = "linear", ephemeris: str = "de432s", - solar_cut: u.Quantity[u.deg] | None = None, - solar_cut_fill_value: float = np.nan, n_proc: int = 1, ) -> None: self.model = model @@ -83,8 +77,6 @@ def __init__( self.extrapolate = extrapolate self.interp_kind = interp_kind self.ephemeris = ephemeris - self.solar_cut = solar_cut.to(u.rad) if solar_cut is not None else solar_cut - self.solar_cut_fill_value = solar_cut_fill_value self.n_proc = n_proc self._interpolator = partial( @@ -232,8 +224,8 @@ def get_emission_pix( emission: Simulated zodiacal emission in units of 'MJy/sr'. """ + # healpix = ahpx.HEALPix(nside=nside, order="ring") pixels = get_validated_pix(pixels=pixels, nside=nside) - unique_pixels, indicies = np.unique(pixels, return_inverse=True) unit_vectors = get_unit_vectors_from_pixels( coord_in=coord_in, @@ -250,7 +242,6 @@ def get_emission_pix( unit_vectors=unit_vectors, indicies=indicies, pixels=unique_pixels, - nside=nside, return_comps=return_comps, ) @@ -267,6 +258,7 @@ def get_binned_emission_ang( lonlat: bool = False, return_comps: bool = False, coord_in: Literal["E", "G", "C"] = "E", + solar_cut: u.Quantity[u.deg] | None = None, ) -> u.Quantity[u.MJy / u.sr]: """Return the simulated binned zodiacal emission given angles on the sky. @@ -298,13 +290,16 @@ def get_binned_emission_ang( return_comps: If True, the emission is returned component-wise. Defaults to False. coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic coordinates) by default. + solar_cut (u.Quantity[u.deg]): Cutoff angle from the sun in degrees. The emission + for all the pointing with angular distance between the sun smaller than + `solar_cutoff` are masked. Defaults to `None`. Returns: emission: Simulated zodiacal emission in units of 'MJy/sr'. """ + healpix = ahpx.HEALPix(nside=nside, order="ring") theta, phi = get_validated_ang(theta=theta, phi=phi, lonlat=lonlat) - unique_angles, counts = np.unique(np.asarray([theta, phi]), return_counts=True, axis=1) unique_pixels = hp.ang2pix(nside, *unique_angles, lonlat=lonlat) unit_vectors = get_unit_vectors_from_ang( @@ -313,7 +308,7 @@ def get_binned_emission_ang( phi=unique_angles[1], lonlat=lonlat, ) - + solcar_cut = solar_cut.to(u.rad) if solar_cut is not None else solar_cut return self._compute_emission( freq=freq, weights=weights, @@ -322,10 +317,10 @@ def get_binned_emission_ang( obs_pos=obs_pos, unit_vectors=unit_vectors, indicies=counts, - binned=True, pixels=unique_pixels, - nside=nside, + healpix=healpix, return_comps=return_comps, + solar_cut=solcar_cut, ) def get_binned_emission_pix( @@ -339,6 +334,7 @@ def get_binned_emission_pix( weights: Sequence[float] | npt.NDArray[np.floating] | None = None, return_comps: bool = False, coord_in: Literal["E", "G", "C"] = "E", + solar_cut: u.Quantity[u.deg] | None = None, ) -> u.Quantity[u.MJy / u.sr]: """Return the simulated binned zodiacal Emission given pixel numbers. @@ -363,11 +359,15 @@ def get_binned_emission_pix( return_comps: If True, the emission is returned component-wise. Defaults to False. coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic coordinates) by default. + solar_cut (u.Quantity[u.deg]): Cutoff angle from the sun in degrees. The emission + for all the pointing with angular distance between the sun smaller than + `solar_cutoff` are masked. Defaults to `None`. Returns: emission: Simulated zodiacal emission in units of 'MJy/sr'. """ + healpix = ahpx.HEALPix(nside=nside, order="ring") pixels = get_validated_pix(pixels=pixels, nside=nside) unique_pixels, counts = np.unique(pixels, return_counts=True) @@ -376,7 +376,7 @@ def get_binned_emission_pix( pixels=unique_pixels, nside=nside, ) - + solcar_cut = solar_cut.to(u.rad) if solar_cut is not None else solar_cut return self._compute_emission( freq=freq, weights=weights, @@ -385,10 +385,10 @@ def get_binned_emission_pix( obs_pos=obs_pos, unit_vectors=unit_vectors, indicies=counts, - binned=True, pixels=unique_pixels, - nside=nside, + healpix=healpix, return_comps=return_comps, + solar_cut=solcar_cut, ) def _compute_emission( @@ -399,11 +399,11 @@ def _compute_emission( obs_time: Time, unit_vectors: npt.NDArray[np.float64], indicies: npt.NDArray[np.int64], - binned: bool = False, obs_pos: u.Quantity[u.AU] | None = None, pixels: npt.NDArray[np.int64] | None = None, - nside: int | None = None, + healpix: ahpx.HEALPix | None = None, return_comps: bool = False, + solar_cut: u.Quantity[u.rad] | None = None, ) -> u.Quantity[u.MJy / u.sr]: """Compute the component-wise zodiacal emission.""" bandpass = validate_and_get_bandpass( @@ -506,25 +506,27 @@ def _compute_emission( * (stop[comp_label] - start[comp_label]) ) - emission = np.zeros( - ( - len(self._ipd_model.comps), - hp.nside2npix(nside) if binned else indicies.size, - ) - ) - if binned: + # Output is requested to be binned + if healpix: + emission = np.zeros((len(self._ipd_model.comps), healpix.npix)) emission[:, pixels] = integrated_comp_emission + if solar_cut is not None: + x, y, z = observer_position.flatten() * u.AU + coord = SkyCoord( + x=x, y=y, z=z, representation_type="cartesian", frame="heliocentricmeanecliptic" + ) + lon, lat = coord.spherical.lon, coord.spherical.lat + lon += 180 * u.deg + solar_mask = healpix.cone_search_lonlat(lon, lat, solar_cut) + if pixels is not None: + emission[:, pixels[solar_mask]] = np.nan + else: + emission[:, solar_mask[indicies]] = np.nan + else: + emission = np.zeros((len(self._ipd_model.comps), indicies.size)) emission = integrated_comp_emission[:, indicies] - if self.solar_cut is not None: - ang_dist = hp.rotator.angdist(-observer_position.flatten(), unit_vectors) - solar_mask = ang_dist < self.solar_cut.value - if binned and pixels is not None: - emission[:, pixels[solar_mask]] = self.solar_cut_fill_value - else: - emission[:, solar_mask[indicies]] = self.solar_cut_fill_value - emission = (emission << SPECIFIC_INTENSITY_UNITS).to(u.MJy / u.sr) return emission if return_comps else emission.sum(axis=0) From 0bf63f88e2365b35a9993292235d9c2df7e5aeef Mon Sep 17 00:00:00 2001 From: Metin San Date: Fri, 26 Apr 2024 09:24:05 +0200 Subject: [PATCH 02/17] Use angular separation instead of cone search. Turns out healpix-cone search is much slower than to just compute the angular separation between the pointing --- .../get_binned_emission_solar_cutoff.py | 3 +- docs/usage.md | 7 ++--- tests/_strategies.py | 1 - tests/test_get_emission.py | 6 ++-- zodipy/zodipy.py | 31 ++++++++++++------- 5 files changed, 28 insertions(+), 20 deletions(-) diff --git a/docs/examples/get_binned_emission_solar_cutoff.py b/docs/examples/get_binned_emission_solar_cutoff.py index c1fa9ca..b4acbc9 100644 --- a/docs/examples/get_binned_emission_solar_cutoff.py +++ b/docs/examples/get_binned_emission_solar_cutoff.py @@ -6,7 +6,7 @@ from zodipy import Zodipy -model = Zodipy("dirbe", solar_cut=60 * u.deg) +model = Zodipy("dirbe") nside = 256 binned_emission = model.get_binned_emission_pix( @@ -15,6 +15,7 @@ nside=nside, obs_time=Time("2020-01-01"), obs="earth", + solar_cut=60 * u.deg, ) hp.mollview( diff --git a/docs/usage.md b/docs/usage.md index 76ccc9e..130852a 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -49,10 +49,9 @@ Note that the bandpass weights must be in spectral radiance units (Jy/sr), even ![Center frequency emission](img/center_freq.png) ![Bandpass integrated emission](img/bandpass_integrated.png) -### Solar cutoff angle -Few experiments look directly in towards the Sun. We can initialize `Zodipy` with the `solar_cut` -argument to mask all input pointing that looks in towards the sun with an angular distance smaller -than the `solar_cut` value. +### Solar cutoff +In the case where the user wishes to bin a set of observations into a HEALPix map, using the `*_binned_*` API functions, +a solar cutoff angle can be specified to mask out the all pixels closer to the Sun by some angle given by `solar_cut` ```python hl_lines="9" {!examples/get_binned_emission_solar_cutoff.py!} diff --git a/tests/_strategies.py b/tests/_strategies.py index 52a9c14..05d4b98 100644 --- a/tests/_strategies.py +++ b/tests/_strategies.py @@ -198,7 +198,6 @@ def any_obs(draw: DrawFn, model: zodipy.Zodipy) -> str: "model": sampled_from(AVAILABLE_MODELS), "gauss_quad_degree": integers(min_value=1, max_value=200), "extrapolate": booleans(), - "solar_cut": quantities(min_value=0, max_value=360, unit=u.deg), } diff --git a/tests/test_get_emission.py b/tests/test_get_emission.py index a6a8758..ded988d 100644 --- a/tests/test_get_emission.py +++ b/tests/test_get_emission.py @@ -6,7 +6,7 @@ import pytest from astropy.time import Time, TimeDelta from hypothesis import given, settings -from hypothesis.strategies import DataObject, data, integers +from hypothesis.strategies import DataObject, booleans, data, integers from zodipy.zodipy import Zodipy @@ -17,6 +17,7 @@ nside, obs, pixels, + quantities, random_freq, time, weights, @@ -60,13 +61,14 @@ def test_get_binned_emission_pix( observer = data.draw(obs(model, time)) frequency = data.draw(freq(model)) pix = data.draw(pixels(nside)) - + cut_solar = data.draw(booleans()) emission_binned = model.get_binned_emission_pix( frequency, pixels=pix, nside=nside, obs_time=time, obs=observer, + solar_cut=data.draw(quantities(20, 50, u.deg)) if cut_solar else None, ) assert emission_binned.shape == (hp.nside2npix(nside),) diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index 219fa56..be0ffca 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -5,11 +5,11 @@ from functools import partial from typing import TYPE_CHECKING, Callable, Literal, Sequence +import astropy.coordinates as coords import astropy.units as u import astropy_healpix as ahpx import healpy as hp import numpy as np -from astropy.coordinates import SkyCoord, solar_system_ephemeris from scipy.interpolate import interp1d from zodipy._bandpass import get_bandpass_interpolation_table, validate_and_get_bandpass @@ -90,7 +90,7 @@ def __init__( @property def supported_observers(self) -> list[str]: """Return a list of available observers given an ephemeris.""" - return [*list(solar_system_ephemeris.bodies), "semb-l2"] + return [*list(coords.solar_system_ephemeris.bodies), "semb-l2"] def get_parameters(self) -> ParameterDict: """Return a dictionary containing the interplanetary dust model parameters.""" @@ -308,7 +308,6 @@ def get_binned_emission_ang( phi=unique_angles[1], lonlat=lonlat, ) - solcar_cut = solar_cut.to(u.rad) if solar_cut is not None else solar_cut return self._compute_emission( freq=freq, weights=weights, @@ -320,7 +319,7 @@ def get_binned_emission_ang( pixels=unique_pixels, healpix=healpix, return_comps=return_comps, - solar_cut=solcar_cut, + solar_cut=solar_cut, ) def get_binned_emission_pix( @@ -376,7 +375,6 @@ def get_binned_emission_pix( pixels=unique_pixels, nside=nside, ) - solcar_cut = solar_cut.to(u.rad) if solar_cut is not None else solar_cut return self._compute_emission( freq=freq, weights=weights, @@ -388,7 +386,7 @@ def get_binned_emission_pix( pixels=unique_pixels, healpix=healpix, return_comps=return_comps, - solar_cut=solcar_cut, + solar_cut=solar_cut, ) def _compute_emission( @@ -511,13 +509,22 @@ def _compute_emission( emission = np.zeros((len(self._ipd_model.comps), healpix.npix)) emission[:, pixels] = integrated_comp_emission if solar_cut is not None: - x, y, z = observer_position.flatten() * u.AU - coord = SkyCoord( - x=x, y=y, z=z, representation_type="cartesian", frame="heliocentricmeanecliptic" + observer_coords = coords.SkyCoord( + *(observer_position.flatten() * u.AU), + representation_type="cartesian", + frame=coords.HeliocentricTrueEcliptic, ) - lon, lat = coord.spherical.lon, coord.spherical.lat - lon += 180 * u.deg - solar_mask = healpix.cone_search_lonlat(lon, lat, solar_cut) + sun_coords = observer_coords.copy() + sun_lon = sun_coords.spherical.lon + sun_lon += 180 * u.deg + + pointing_angles = coords.SkyCoord( + *unit_vectors, + representation_type="cartesian", + frame=coords.HeliocentricTrueEcliptic, + ) + angular_distance = pointing_angles.separation(sun_coords) + solar_mask = angular_distance < solar_cut if pixels is not None: emission[:, pixels[solar_mask]] = np.nan else: From b484d7439591171c944799686334cc561bf4a8e1 Mon Sep 17 00:00:00 2001 From: Metin San Date: Fri, 26 Apr 2024 22:57:04 +0200 Subject: [PATCH 03/17] Add tests for skycoord api --- tests/_strategies.py | 53 ++++-- tests/test_get_emission.py | 279 +++++++++++++++---------------- tests/test_params.py | 6 +- tests/test_reprs.py | 4 +- tests/test_tabulated_density.py | 4 +- tests/test_validators.py | 16 +- zodipy/_sky_coords.py | 31 ++-- zodipy/_types.py | 2 + zodipy/_unit_vectors.py | 45 ++--- zodipy/_validators.py | 21 +-- zodipy/zodipy.py | 288 ++++++++++++++++++-------------- 11 files changed, 403 insertions(+), 346 deletions(-) diff --git a/tests/_strategies.py b/tests/_strategies.py index 05d4b98..319b5af 100644 --- a/tests/_strategies.py +++ b/tests/_strategies.py @@ -5,12 +5,11 @@ from math import log2 from typing import Any, Callable, Sequence +import astropy.coordinates as coords import astropy.units as u import healpy as hp import numpy as np import numpy.typing as npt -from astropy.coordinates import HeliocentricMeanEcliptic, get_body -from astropy.time import Time from hypothesis.extra.numpy import arrays from hypothesis.strategies import ( DrawFn, @@ -66,15 +65,39 @@ def quantities( @composite -def time(draw: DrawFn) -> Time: - return draw(datetimes(min_value=MIN_DATE, max_value=MAX_DATE).map(Time)) +def times(draw: DrawFn) -> times.Time: + return draw(datetimes(min_value=MIN_DATE, max_value=MAX_DATE).map(times.Time)) @composite -def nside(draw: Callable[[SearchStrategy[int]], int]) -> int: +def nsides(draw: Callable[[SearchStrategy[int]], int]) -> int: return draw(integers(min_value=MIN_NSIDE_EXP, max_value=MAX_NSIDE_EXP).map(partial(pow, 2))) +@composite +def frames(draw: DrawFn) -> type[coords.BaseCoordinateFrame]: + return draw(sampled_from([coords.ICRS, coords.Galactic, coords.HeliocentricMeanEcliptic])) + + +@composite +def sky_coords(draw: DrawFn) -> coords.SkyCoord: + theta_strategy = floats(min_value=0, max_value=360) + phi_strategy = floats(min_value=-90, max_value=90) + + shape = draw(integers(min_value=1, max_value=MAX_ANGELS_LEN)) + + theta_array_strategy = arrays(dtype=float, shape=shape, elements=theta_strategy).map( + partial(u.Quantity, unit=u.deg) + ) + phi_array_strategy = arrays(dtype=float, shape=shape, elements=phi_strategy).map( + partial(u.Quantity, unit=u.deg) + ) + frame = draw(frames()) + lon = draw(theta_array_strategy) + lat = draw(phi_array_strategy) + return coords.SkyCoord(lon, lat, frame=frame) + + @composite def pixels(draw: DrawFn, nside: int) -> int | list[int] | npt.NDArray[np.integer]: npix = hp.nside2npix(nside) @@ -110,7 +133,7 @@ def angles(draw: DrawFn, lonlat: bool = False) -> tuple[u.Quantity[u.deg], u.Qua @composite -def freq(draw: DrawFn, model: zodipy.Zodipy) -> u.Quantity[u.GHz] | u.Quantity[u.micron]: +def freqs(draw: DrawFn, model: zodipy.Zodipy) -> u.Quantity[u.GHz] | u.Quantity[u.micron]: if model.extrapolate: return draw(sampled_from(FREQ_LOG_RANGE).map(np.exp).map(partial(u.Quantity, unit=u.GHz))) @@ -127,7 +150,7 @@ def freq(draw: DrawFn, model: zodipy.Zodipy) -> u.Quantity[u.GHz] | u.Quantity[u @composite -def random_freq(draw: DrawFn, unit: u.Unit | None = None, bandpass: bool = False) -> u.Quantity: +def random_freqs(draw: DrawFn, unit: u.Unit | None = None, bandpass: bool = False) -> u.Quantity: random_freq = draw( sampled_from(FREQ_LOG_RANGE).map(np.exp).map(partial(u.Quantity, unit=u.GHz)) ) @@ -166,15 +189,21 @@ def normalize_array( @composite -def obs(draw: DrawFn, model: zodipy.Zodipy, obs_time: Time) -> str: - def get_obs_dist(obs: str, obs_time: Time) -> u.Quantity[u.AU]: +def obs_positions(draw: DrawFn, model: zodipy.Zodipy, obs_time: times.Time) -> str: + def get_obs_dist(obs: str, obs_time: times.Time) -> u.Quantity[u.AU]: if obs == "semb-l2": obs_pos = ( - get_body("earth", obs_time).transform_to(HeliocentricMeanEcliptic).cartesian.xyz + coords.get_body("earth", obs_time) + .transform_to(coords.HeliocentricMeanEcliptic) + .cartesian.xyz ) obs_pos += 0.01 * u.AU else: - obs_pos = get_body(obs, obs_time).transform_to(HeliocentricMeanEcliptic).cartesian.xyz + obs_pos = ( + coords.get_body(obs, obs_time) + .transform_to(coords.HeliocentricMeanEcliptic) + .cartesian.xyz + ) return u.Quantity(np.linalg.norm(obs_pos.value), u.AU) los_dist_cut = min( @@ -202,7 +231,7 @@ def any_obs(draw: DrawFn, model: zodipy.Zodipy) -> str: @composite -def model(draw: DrawFn, **static_params: dict[str, Any]) -> zodipy.Zodipy: +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: diff --git a/tests/test_get_emission.py b/tests/test_get_emission.py index ded988d..a2834df 100644 --- a/tests/test_get_emission.py +++ b/tests/test_get_emission.py @@ -1,33 +1,55 @@ from __future__ import annotations +import astropy.coordinates as coords import astropy.units as u import healpy as hp import numpy as np import pytest from astropy.time import Time, TimeDelta from hypothesis import given, settings -from hypothesis.strategies import DataObject, booleans, data, integers +from hypothesis.strategies import DataObject, booleans, data from zodipy.zodipy import Zodipy from ._strategies import ( angles, - freq, - model, - nside, - obs, + freqs, + nsides, + obs_positions, pixels, quantities, - random_freq, - time, + random_freqs, + sky_coords, + times, weights, + zodipy_models, ) from ._tabulated_dirbe import DAYS, LAT, LON, TABULATED_DIRBE_EMISSION DIRBE_START_DAY = Time("1990-01-01") -@given(model(), time(), nside(), data()) +@given(zodipy_models(), times(), sky_coords(), data()) +@settings(deadline=None) +def test_get_emission_skycoord( + model: Zodipy, + time: Time, + coordinates: coords.SkyCoord, + data: DataObject, +) -> None: + """Property test for get_emission_skycoord.""" + observer = data.draw(obs_positions(model, time)) + frequency = data.draw(freqs(model)) + emission = model.get_emission_skycoord( + coordinates, + freq=frequency, + obs_time=time, + obs_pos=observer, + ) + assert emission.size == 1 + + +@given(zodipy_models(), times(), nsides(), data()) @settings(deadline=None) def test_get_emission_pix( model: Zodipy, @@ -36,20 +58,20 @@ def test_get_emission_pix( data: DataObject, ) -> None: """Property test for get_emission_pix.""" - observer = data.draw(obs(model, time)) - frequency = data.draw(freq(model)) + observer = data.draw(obs_positions(model, time)) + frequency = data.draw(freqs(model)) pix = data.draw(pixels(nside)) emission = model.get_emission_pix( - frequency, - pixels=pix, + pix, nside=nside, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) assert np.size(emission) == np.size(pix) -@given(model(), time(), nside(), data()) +@given(zodipy_models(), times(), nsides(), data()) @settings(deadline=None) def test_get_binned_emission_pix( model: Zodipy, @@ -58,22 +80,22 @@ def test_get_binned_emission_pix( data: DataObject, ) -> None: """Property test for get_binned_emission_pix.""" - observer = data.draw(obs(model, time)) - frequency = data.draw(freq(model)) + observer = data.draw(obs_positions(model, time)) + frequency = data.draw(freqs(model)) pix = data.draw(pixels(nside)) cut_solar = data.draw(booleans()) emission_binned = model.get_binned_emission_pix( - frequency, - pixels=pix, + pix, nside=nside, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, solar_cut=data.draw(quantities(20, 50, u.deg)) if cut_solar else None, ) assert emission_binned.shape == (hp.nside2npix(nside),) -@given(model(), time(), angles(), data()) +@given(zodipy_models(), times(), angles(), data()) @settings(deadline=None) def test_get_emission_ang( model: Zodipy, @@ -84,20 +106,20 @@ def test_get_emission_ang( """Property test for get_emission_ang.""" theta, phi = angles - observer = data.draw(obs(model, time)) - frequency = data.draw(freq(model)) + observer = data.draw(obs_positions(model, time)) + frequency = data.draw(freqs(model)) emission = model.get_emission_ang( - frequency, - theta=theta, - phi=phi, + theta, + phi, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) assert emission.size == theta.size == phi.size -@given(model(), time(), nside(), angles(), data()) +@given(zodipy_models(), times(), nsides(), angles(), data()) @settings(deadline=None) def test_get_binned_emission_ang( model: Zodipy, @@ -109,21 +131,21 @@ def test_get_binned_emission_ang( """Property test for get_binned_emission_pix.""" theta, phi = angles - observer = data.draw(obs(model, time)) - frequency = data.draw(freq(model)) + observer = data.draw(obs_positions(model, time)) + frequency = data.draw(freqs(model)) emission_binned = model.get_binned_emission_ang( - frequency, - theta=theta, - phi=phi, + theta, + phi, nside=nside, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) assert emission_binned.shape == (hp.nside2npix(nside),) -@given(model(extrapolate=False), time(), angles(lonlat=True), nside(), data()) +@given(zodipy_models(extrapolate=False), times(), angles(lonlat=True), nsides(), data()) @settings(deadline=None) def test_invalid_freq( model: Zodipy, @@ -134,47 +156,47 @@ def test_invalid_freq( ) -> None: """Property test checking for unsupported spectral range.""" theta, phi = angles - observer = data.draw(obs(model, time)) + observer = data.draw(obs_positions(model, time)) pix = data.draw(pixels(nside)) - freq = data.draw(random_freq(unit=model._ipd_model.spectrum.unit)) + 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( - freq, - pixels=pix, + pix, + freq=freq, nside=nside, obs_time=time, - obs=observer, + obs_pos=observer, ) with pytest.raises(ValueError): model.get_emission_ang( - freq, - theta=theta, - phi=phi, + theta, + phi, + freq=freq, obs_time=time, - obs=observer, + obs_pos=observer, lonlat=True, ) with pytest.raises(ValueError): model.get_binned_emission_pix( - freq, - pixels=pix, + pix, nside=nside, + freq=freq, obs_time=time, - obs=observer, + obs_pos=observer, ) with pytest.raises(ValueError): model.get_binned_emission_ang( - freq, - theta=theta, - phi=phi, + theta, + phi, nside=nside, + freq=freq, obs_time=time, - obs=observer, + obs_pos=observer, lonlat=True, ) @@ -191,49 +213,16 @@ def test_compare_to_dirbe_idl() -> None: time = DIRBE_START_DAY + TimeDelta(day - 1, format="jd") emission = model.get_emission_ang( - frequency * u.micron, - theta=lon * u.deg, - phi=lat * u.deg, + lon * u.deg, + lat * u.deg, + freq=frequency * u.micron, lonlat=True, - obs="earth", + obs_pos="earth", obs_time=time, ) assert emission.value == pytest.approx(tabulated_emission[idx], rel=0.01) -@given(model(), time(), nside(), integers(min_value=1, max_value=100), data()) -@settings(max_examples=10, deadline=None) -def test_invalid_pixel( - model: Zodipy, - time: Time, - nside: int, - random_integer: int, - data: DataObject, -) -> None: - """Tests that an error is raised when an invalid pixel is passed to get_emission_pix.""" - frequency = data.draw(freq(model)) - observer = data.draw(obs(model, time)) - npix = hp.nside2npix(nside) - - with pytest.raises(ValueError): - model.get_emission_pix( - frequency, - pixels=npix + random_integer, - nside=nside, - obs_time=time, - obs=observer, - ) - - with pytest.raises(ValueError): - model.get_emission_pix( - frequency, - pixels=-random_integer, - nside=nside, - obs_time=time, - obs=observer, - ) - - def test_multiprocessing() -> None: """Tests multiprocessing with for zodipy. @@ -252,69 +241,69 @@ def test_multiprocessing() -> None: phi = np.random.uniform(0, 360, size=1000) * u.deg emission_pix = model.get_emission_pix( - frequency, - pixels=pix, + pix, nside=nside, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) emission_pix_parallel = model_parallel.get_emission_pix( - frequency, - pixels=pix, + pix, nside=nside, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) assert np.allclose(emission_pix, emission_pix_parallel) emission_binned_pix = model.get_binned_emission_pix( - frequency, - pixels=pix, + pix, nside=nside, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) emission_binned_pix_parallel = model_parallel.get_binned_emission_pix( - frequency, - pixels=pix, + pix, nside=nside, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) assert np.allclose(emission_binned_pix, emission_binned_pix_parallel) emission_ang = model.get_emission_ang( - frequency, - theta=theta, - phi=phi, + theta, + phi, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) emission_ang_parallel = model_parallel.get_emission_ang( - frequency, - theta=theta, - phi=phi, + theta, + phi, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) assert np.allclose(emission_ang, emission_ang_parallel) emission_binned_ang = model.get_binned_emission_ang( - frequency, - theta=theta, - phi=phi, + theta, + phi, + freq=frequency, nside=nside, obs_time=time, - obs=observer, + obs_pos=observer, ) emission_binned_ang_parallel = model_parallel.get_binned_emission_ang( - frequency, - theta=theta, - phi=phi, + theta, + phi, + freq=frequency, nside=nside, obs_time=time, - obs=observer, + obs_pos=observer, ) assert np.allclose(emission_binned_ang.value, emission_binned_ang_parallel.value) @@ -332,28 +321,28 @@ def test_inner_radial_cutoff_multiprocessing() -> None: pix = np.random.randint(0, hp.nside2npix(nside), size=1000) emission_pix = model.get_emission_pix( - frequency, - pixels=pix, + pix, nside=nside, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) emission_pix_parallel = model_parallel.get_emission_pix( - frequency, - pixels=pix, + pix, nside=nside, + freq=frequency, obs_time=time, - obs=observer, + obs_pos=observer, ) assert np.allclose(emission_pix, emission_pix_parallel) @given( - model(extrapolate=True), - time(), - nside(), + zodipy_models(extrapolate=True), + times(), + nsides(), angles(), - random_freq(bandpass=True), + random_freqs(bandpass=True), data(), ) @settings(deadline=None) @@ -367,26 +356,26 @@ def test_bandpass_integration( ) -> None: """Property test for bandpass integrations.""" theta, phi = angles - observer = data.draw(obs(model, time)) + observer = data.draw(obs_positions(model, time)) bp_weights = data.draw(weights(freqs)) emission_binned = model.get_binned_emission_ang( - freqs, + theta, + phi, + freq=freqs, weights=bp_weights, - theta=theta, - phi=phi, nside=nside, obs_time=time, - obs=observer, + obs_pos=observer, ) assert emission_binned.shape == (hp.nside2npix(nside),) @given( - model(extrapolate=True), - time(), - nside(), + zodipy_models(extrapolate=True), + times(), + nsides(), angles(), - random_freq(bandpass=True), + random_freqs(bandpass=True), data(), ) @settings(deadline=None) @@ -400,17 +389,17 @@ def test_weights( ) -> None: """Property test for bandpass weights.""" theta, phi = angles - observer = data.draw(obs(model, time)) + observer = data.draw(obs_positions(model, time)) bp_weights = data.draw(weights(freqs)) model.get_binned_emission_ang( - freqs, + theta, + phi, weights=bp_weights, - theta=theta, - phi=phi, + freq=freqs, nside=nside, obs_time=time, - obs=observer, + obs_pos=observer, ) @@ -427,12 +416,12 @@ def test_custom_weights() -> None: weights /= np.trapz(weights, freqs.value) model.get_emission_pix( + pix, freq=freqs, weights=weights, - pixels=pix, nside=nside, obs_time=time, - obs="earth", + obs_pos="earth", ) @@ -444,16 +433,16 @@ def test_custom_obs_pos() -> None: pix = np.arange(hp.nside2npix(nside)) model.get_emission_pix( + pix, freq=234 * u.micron, - pixels=pix, nside=nside, obs_time=time, obs_pos=[0.1, 0.2, 1] * u.AU, ) model.get_emission_pix( + pix, freq=234 * u.micron, - pixels=pix, nside=nside, obs_time=time, obs_pos=[2, 0.1, 4] * u.AU, @@ -461,8 +450,8 @@ def test_custom_obs_pos() -> None: with pytest.raises(TypeError): model.get_emission_pix( + pix, freq=234 * u.micron, - pixels=pix, nside=nside, obs_time=time, obs_pos=[2, 0.1, 4], @@ -470,8 +459,8 @@ def test_custom_obs_pos() -> None: with pytest.raises(u.UnitsError): model.get_emission_pix( + pix, freq=234 * u.micron, - pixels=pix, nside=nside, obs_time=time, obs_pos=[2, 0.1, 4] * u.s, diff --git a/tests/test_params.py b/tests/test_params.py index c60ed46..6a6adca 100644 --- a/tests/test_params.py +++ b/tests/test_params.py @@ -4,16 +4,16 @@ from zodipy.zodipy import Zodipy -from ._strategies import model +from ._strategies import zodipy_models -@given(model()) +@given(zodipy_models()) def test_get_parameters(model: Zodipy) -> None: """Tests that the parameters are returned as a dictionary.""" assert isinstance(model.get_parameters(), dict) -@given(model()) +@given(zodipy_models()) def test_update_model(model: Zodipy) -> None: """Tests that the model can be updated.""" parameters = model.get_parameters() diff --git a/tests/test_reprs.py b/tests/test_reprs.py index ef05271..19c7602 100644 --- a/tests/test_reprs.py +++ b/tests/test_reprs.py @@ -2,10 +2,10 @@ from zodipy.zodipy import Zodipy -from ._strategies import model +from ._strategies import zodipy_models -@given(model()) +@given(zodipy_models()) def test_ipd_model_repr(model: Zodipy) -> None: """Tests that the IPD model has a userfriendly repr.""" repr(model) diff --git a/tests/test_tabulated_density.py b/tests/test_tabulated_density.py index 5572cb3..fbb5248 100644 --- a/tests/test_tabulated_density.py +++ b/tests/test_tabulated_density.py @@ -7,11 +7,11 @@ from zodipy import model_registry, tabulate_density from zodipy.zodipy import Zodipy -from ._strategies import model +from ._strategies import zodipy_models @given( - model(), + zodipy_models(), integers(min_value=10, max_value=200), floats(min_value=2, max_value=10), floats(min_value=2, max_value=10), diff --git a/tests/test_validators.py b/tests/test_validators.py index 53d070e..79330cf 100644 --- a/tests/test_validators.py +++ b/tests/test_validators.py @@ -9,7 +9,7 @@ from zodipy._validators import get_validated_and_normalized_weights, get_validated_freq from zodipy.zodipy import Zodipy -from ._strategies import model, random_freq, weights +from ._strategies import random_freqs, weights, zodipy_models BANDPASS_FREQUENCIES = np.linspace(95, 105, 11) * u.GHz BANDPASS_WAVELENGTHS = np.linspace(20, 25, 11) * u.micron @@ -17,7 +17,7 @@ OBS_TIME = Time("2021-01-01T00:00:00") -@given(model(extrapolate=False)) +@given(zodipy_models(extrapolate=False)) def test_validate_frequencies(model: Zodipy) -> None: """Tests that the frequencies are validated.""" with pytest.raises(TypeError): @@ -46,7 +46,7 @@ def test_validate_frequencies(model: Zodipy) -> None: ) -@given(random_freq(bandpass=True), data()) +@given(random_freqs(bandpass=True), data()) def test_validate_weights(freq: FrequencyOrWavelength, data: DataObject) -> None: """Tests that the bandpass weights are validated.""" bp_weights = data.draw(weights(freq)) @@ -115,22 +115,22 @@ 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(400 * u.micron, pixels=[1, 4, 5], nside=32, obs_time=OBS_TIME) + model.get_emission_pix([1, 4, 5], freq=400 * u.micron, nside=32, obs_time=OBS_TIME) model = Zodipy("dirbe", extrapolate=True) - model.get_emission_pix(400 * u.micron, pixels=[1, 4, 5], nside=32, obs_time=OBS_TIME) + model.get_emission_pix([1, 4, 5], freq=400 * u.micron, 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(27 * u.micron, pixels=[1, 4, 5], nside=32, obs_time=OBS_TIME) + linear = model.get_emission_pix([1, 4, 5], freq=27 * u.micron, nside=32, obs_time=OBS_TIME) model = Zodipy("dirbe", interp_kind="quadratic") - quadratic = model.get_emission_pix(27 * u.micron, pixels=[1, 4, 5], nside=32, obs_time=OBS_TIME) + quadratic = model.get_emission_pix([1, 4, 5], freq=27 * u.micron, 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(27 * u.micron, pixels=[1, 4, 5], nside=32, obs_time=OBS_TIME) + model.get_emission_pix(pixels=[1, 4, 5], freq=27 * u.micron, nside=32, obs_time=OBS_TIME) diff --git a/zodipy/_sky_coords.py b/zodipy/_sky_coords.py index 44d7dde..dc4dd6f 100644 --- a/zodipy/_sky_coords.py +++ b/zodipy/_sky_coords.py @@ -1,18 +1,15 @@ -from typing import Tuple, Union - +import astropy.coordinates as coords import astropy.units as u import numpy as np -import numpy.typing as npt -from astropy.coordinates import HeliocentricMeanEcliptic, get_body from astropy.time import Time from zodipy._constants import DISTANCE_FROM_EARTH_TO_L2 +from zodipy._types import NumpyArray -@u.quantity_input def get_obs_and_earth_positions( - obs: str, obs_time: Time, obs_pos: Union[u.Quantity[u.AU], None] -) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]: + obs_pos: u.Quantity[u.AU] | str, obs_time: Time +) -> tuple[NumpyArray, NumpyArray]: """Return the position of the observer and the Earth (3, `n_pointing`, `n_gauss_quad_degree`). The lagrange point SEMB-L2 is not included in any of the current available @@ -21,18 +18,24 @@ def get_obs_and_earth_positions( pointing to Earth from the Sun. """ earth_position = _get_earth_position(obs_time) - if obs_pos is None: - obs_position = _get_observer_position(obs, obs_time, earth_position) + if isinstance(obs_pos, str): + obs_position = _get_observer_position(obs_pos, obs_time, earth_position) else: - obs_position = obs_pos.to(u.AU) + try: + obs_position = obs_pos.to(u.AU) + except AttributeError: + msg = ( + "Observer position must be a string or an astropy Quantity with units of distance." + ) + raise TypeError(msg) from AttributeError return obs_position.reshape(3, 1, 1).value, earth_position.reshape(3, 1, 1).value def _get_earth_position(obs_time: Time) -> u.Quantity[u.AU]: """Return the position of the Earth given an ephemeris and observation time.""" - earth_skycoordinate = get_body("earth", obs_time) - earth_skycoordinate = earth_skycoordinate.transform_to(HeliocentricMeanEcliptic) + earth_skycoordinate = coords.get_body("earth", obs_time) + earth_skycoordinate = earth_skycoordinate.transform_to(coords.HeliocentricMeanEcliptic) return earth_skycoordinate.cartesian.xyz.to(u.AU) @@ -43,8 +46,8 @@ def _get_observer_position( if obs.lower() == "semb-l2": return _get_sun_earth_moon_barycenter(earth_pos) - observer_skycoordinate = get_body(obs, obs_time) - observer_skycoordinate = observer_skycoordinate.transform_to(HeliocentricMeanEcliptic) + observer_skycoordinate = coords.get_body(obs, obs_time) + observer_skycoordinate = observer_skycoordinate.transform_to(coords.HeliocentricMeanEcliptic) return observer_skycoordinate.cartesian.xyz.to(u.AU) diff --git a/zodipy/_types.py b/zodipy/_types.py index 61014ff..cd2b8ed 100644 --- a/zodipy/_types.py +++ b/zodipy/_types.py @@ -8,3 +8,5 @@ SkyAngles = Union[u.Quantity[u.deg], u.Quantity[u.rad]] FrequencyOrWavelength = Union[u.Quantity[u.Hz], u.Quantity[u.m]] ParameterDict = dict +NumpyArray = Union[npt.NDArray[np.float64], npt.NDArray[np.int64]] +MJySr = u.Quantity[u.MJy / u.sr] diff --git a/zodipy/_unit_vectors.py b/zodipy/_unit_vectors.py index 8c0cca3..797de75 100644 --- a/zodipy/_unit_vectors.py +++ b/zodipy/_unit_vectors.py @@ -1,29 +1,30 @@ -from __future__ import annotations +# from __future__ import annotations -from typing import TYPE_CHECKING, Sequence +# from typing import TYPE_CHECKING, Sequence -import healpy as hp -import numpy as np +# import healpy as hp +# import numpy as np -if TYPE_CHECKING: - import astropy.units as u - import numpy.typing as npt +# if TYPE_CHECKING: +# import astropy.units as u +# import numpy.typing as npt -def get_unit_vectors_from_pixels( - coord_in: str, pixels: Sequence[int] | npt.NDArray[np.int64], nside: int -) -> npt.NDArray[np.float64]: - """Return ecliptic unit vectors from HEALPix pixels representing some pointing.""" - unit_vectors = np.asarray(hp.pix2vec(nside, pixels)) - return np.asarray(hp.Rotator(coord=[coord_in, "E"])(unit_vectors)) +# def get_unit_vectors_from_pixels( +# coord_in: str, pixels: Sequence[int] | npt.NDArray[np.int64], nside: int +# ) -> npt.NDArray[np.float64]: +# """Return ecliptic unit vectors from HEALPix pixels representing some pointing.""" +# unit_vectors = np.asarray(hp.pix2vec(nside, pixels)) +# return np.asarray(hp.Rotator(coord=[coord_in, "E"])(unit_vectors)) -def get_unit_vectors_from_ang( - coord_in: str, - phi: u.Quantity[u.rad] | u.Quantity[u.deg], - theta: u.Quantity[u.rad] | u.Quantity[u.deg], - lonlat: bool = False, -) -> npt.NDArray[np.float64]: - """Return ecliptic unit vectors from sky angles representing some pointing.""" - unit_vectors = np.asarray(hp.ang2vec(theta, phi, lonlat=lonlat)).transpose() - return np.asarray(hp.Rotator(coord=[coord_in, "E"])(unit_vectors)) +# def get_unit_vectors_from_ang( +# coord_in: str, +# phi: u.Quantity[u.rad] | u.Quantity[u.deg], +# theta: u.Quantity[u.rad] | u.Quantity[u.deg], +# lonlat: bool = False, +# ) -> npt.NDArray[np.float64]: +# """Return ecliptic unit vectors from sky angles representing some pointing.""" +# unit_vectors = np.asarray(hp.ang2vec(theta, phi, lonlat=lonlat)).transpose() +# return np.asarray(hp.Rotator(coord=[coord_in, "E"])(unit_vectors)) + diff --git a/zodipy/_validators.py b/zodipy/_validators.py index 51697a1..3f9f7e7 100644 --- a/zodipy/_validators.py +++ b/zodipy/_validators.py @@ -1,12 +1,13 @@ from typing import Sequence, Tuple, Union import astropy.units as u -import healpy as hp + +# import healpy as hp import numpy as np import numpy.typing as npt from ._ipd_model import InterplanetaryDustModel -from ._types import FrequencyOrWavelength, Pixels, SkyAngles +from ._types import FrequencyOrWavelength, SkyAngles def get_validated_freq( @@ -90,13 +91,13 @@ def get_validated_ang( return theta, phi -def get_validated_pix(pixels: Pixels, nside: int) -> Pixels: - """Validate user inputted pixels.""" - if (np.max(pixels) > hp.nside2npix(nside)) or (np.min(pixels) < 0): - msg = "invalid pixel number given nside" - raise ValueError(msg) +# def get_validated_pix(pixels: Pixels, nside: int) -> Pixels: +# """Validate user inputted pixels.""" +# if (np.max(pixels) > hp.nside2npix(nside)) or (np.min(pixels) < 0): +# msg = "invalid pixel number given nside" +# raise ValueError(msg) - if np.ndim(pixels) == 0: - return np.expand_dims(pixels, axis=0) +# if np.ndim(pixels) == 0: +# return np.expand_dims(pixels, axis=0) - return pixels +# return pixels diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index be0ffca..4e8b9d7 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -1,16 +1,15 @@ from __future__ import annotations +import functools import multiprocessing import platform -from functools import partial -from typing import TYPE_CHECKING, Callable, Literal, Sequence +from typing import TYPE_CHECKING, Callable, Sequence import astropy.coordinates as coords import astropy.units as u -import astropy_healpix as ahpx -import healpy as hp +import astropy_healpix as hp import numpy as np -from scipy.interpolate import interp1d +from scipy import interpolate from zodipy._bandpass import get_bandpass_interpolation_table, validate_and_get_bandpass from zodipy._constants import SPECIFIC_INTENSITY_UNITS @@ -20,8 +19,7 @@ from zodipy._ipd_dens_funcs import construct_density_partials_comps from zodipy._line_of_sight import get_line_of_sight_start_and_stop_distances from zodipy._sky_coords import get_obs_and_earth_positions -from zodipy._unit_vectors import get_unit_vectors_from_ang, get_unit_vectors_from_pixels -from zodipy._validators import get_validated_ang, get_validated_pix +from zodipy._validators import get_validated_ang from zodipy.model_registry import model_registry if TYPE_CHECKING: @@ -79,8 +77,8 @@ def __init__( self.ephemeris = ephemeris self.n_proc = n_proc - self._interpolator = partial( - interp1d, + self._interpolator = functools.partial( + interpolate.interp1d, kind=self.interp_kind, fill_value="extrapolate" if self.extrapolate else np.nan, ) @@ -118,84 +116,135 @@ def update_parameters(self, parameters: ParameterDict) -> None: self._ipd_model = self._ipd_model.__class__(**_dict) - def get_emission_ang( + def get_emission_skycoord( self, + coord: coords.SkyCoord, + *, + obs_time: Time, freq: FrequencyOrWavelength, + obs_pos: u.Quantity[u.AU] | str = "earth", + weights: Sequence[float] | npt.NDArray[np.floating] | None = None, + return_comps: bool = False, + ) -> u.Quantity[u.MJy / u.sr]: + """Return the simulated zodiacal light for observations in an Astropy `SkyCoord` object. + + The pointing, for which to compute the emission, is specified in form of angles on + the sky given by `theta` and `phi`. + + Args: + coord: Astropy `SkyCoord` object representing the pointing on the sky. This + object must have a frame attribute representing the coordinate frame of the + input pointing, for example `astropy.coordinates.Galactic`. The frame must + be convertible to `BarycentricMeanEcliptic`. + obs_time: Time of observation. + 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. + obs_pos: The heliocentric ecliptic position of the observer in AU, or a string + representing an observer in the `astropy.coordinates.solar_system_ephemeris`. + Defaults to 'earth'. + weights: Bandpass weights corresponding the the frequencies in `freq`. The weights + are assumed to be given in spectral radiance units (Jy/sr). + return_comps: If True, the emission is returned component-wise. Defaults to False. + + Returns: + emission: Simulated zodiacal light in units of 'MJy/sr'. + + """ + (unique_lon, unique_lat), indicies = np.unique( + np.vstack([coord.spherical.lon.value, coord.spherical.lat.value]), + return_inverse=True, + axis=1, + ) + coord = coords.SkyCoord(unique_lon * u.deg, unique_lat * u.deg, frame=coord.frame) + + return self._compute_emission( + freq=freq, + weights=weights, + obs_time=obs_time, + obs_pos=obs_pos, + coordinates=coord, + indicies=indicies, + return_comps=return_comps, + ) + + def get_emission_ang( + self, theta: SkyAngles, phi: SkyAngles, + *, + freq: FrequencyOrWavelength, obs_time: Time, - obs: str = "earth", - obs_pos: u.Quantity[u.AU] | None = None, + obs_pos: u.Quantity[u.AU] | str = "earth", + frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, weights: Sequence[float] | npt.NDArray[np.floating] | None = None, lonlat: bool = False, return_comps: bool = False, - coord_in: Literal["E", "G", "C"] = "E", ) -> u.Quantity[u.MJy / u.sr]: """Return the simulated zodiacal emission given angles on the sky. The pointing, for which to compute the emission, is specified in form of angles on - the sky given by `theta` and `phi`. + the sky given by `theta` and `phi`, matching the healpy convention. 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. theta: Angular co-latitude coordinate of a point, or a sequence of points, on the celestial sphere. Must be in the range [0, π] rad. Units must be either radians or degrees. phi: Angular longitude coordinate of a point, or a sequence of points, on the celestial sphere. Must be in the range [0, 2π] rad. Units must be either radians or degrees. + 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. obs_time: Time of observation. - obs: Name of the Solar System observer. A list of all support observers (for a - given ephemeridis) is specified in `supported_observers` attribute of the - `Zodipy` instance. Defaults to 'earth'. - obs_pos: The heliocentric ecliptic cartesian position of the observer in AU. - Overrides the `obs` argument. Default is None. + obs_pos: The heliocentric ecliptic position of the observer in AU, or a string + representing an observer in the `astropy.coordinates.solar_system_ephemeris`. + Defaults to 'earth'. + frame: Astropy coordinate frame representing the coordinate frame of the input pointing. + Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other + alternatives are `Galactic` and `ICRS`. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). lonlat: If True, input angles (`theta`, `phi`) are assumed to be longitude and latitude, otherwise, they are co-latitude and longitude. return_comps: If True, the emission is returned component-wise. Defaults to False. - coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic - coordinates) by default. Returns: emission: Simulated zodiacal emission in units of 'MJy/sr'. """ theta, phi = get_validated_ang(theta=theta, phi=phi, lonlat=lonlat) - - unique_angles, indicies = np.unique(np.asarray([theta, phi]), return_inverse=True, axis=1) - unit_vectors = get_unit_vectors_from_ang( - coord_in=coord_in, - theta=unique_angles[0], - phi=unique_angles[1], - lonlat=lonlat, + if not lonlat: + theta, phi = phi, (np.pi / 2) * u.rad - theta + + (theta, phi), indicies = np.unique(np.stack([theta, phi]), return_inverse=True, axis=1) + coordinates = coords.SkyCoord( + theta, + phi, + frame=frame, ) return self._compute_emission( freq=freq, weights=weights, - obs=obs, obs_time=obs_time, obs_pos=obs_pos, - unit_vectors=unit_vectors, + coordinates=coordinates, indicies=indicies, return_comps=return_comps, ) def get_emission_pix( self, - freq: FrequencyOrWavelength, pixels: Pixels, + *, nside: int, + freq: FrequencyOrWavelength, obs_time: Time, - obs: str = "earth", - obs_pos: u.Quantity[u.AU] | None = None, + obs_pos: u.Quantity[u.AU] | str = "earth", + frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, weights: Sequence[float] | npt.NDArray[np.floating] | None = None, return_comps: bool = False, - coord_in: Literal["E", "G", "C"] = "E", ) -> u.Quantity[u.MJy / u.sr]: """Return the simulated zodiacal emission given pixel numbers. @@ -203,93 +252,85 @@ def get_emission_pix( given by `nside`. Args: + pixels: HEALPix pixel indicies representing points on the celestial sphere. + nside: HEALPix resolution parameter of the pixels and the binned map. 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. - pixels: HEALPix pixel indicies representing points on the celestial sphere. - nside: HEALPix resolution parameter of the pixels and the binned map. obs_time: Time of observation. - obs: Name of the Solar System observer. A list of all support observers (for a - given ephemeridis) is specified in `supported_observers` attribute of the - `Zodipy` instance. Defaults to 'earth'. - obs_pos: The heliocentric ecliptic cartesian position of the observer in AU. - Overrides the `obs` argument. Default is None. + obs_pos: The heliocentric ecliptic position of the observer in AU, or a string + representing an observer in the `astropy.coordinates.solar_system_ephemeris`. + Defaults to 'earth'. + frame: Astropy coordinate frame representing the coordinate frame of the input pointing. + Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other + alternatives are `Galactic` and `ICRS`. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. - coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic - coordinates) by default. Returns: emission: Simulated zodiacal emission in units of 'MJy/sr'. """ - # healpix = ahpx.HEALPix(nside=nside, order="ring") - pixels = get_validated_pix(pixels=pixels, nside=nside) + healpix = hp.HEALPix(nside=nside, order="ring", frame=frame) + unique_pixels, indicies = np.unique(pixels, return_inverse=True) - unit_vectors = get_unit_vectors_from_pixels( - coord_in=coord_in, - pixels=unique_pixels, - nside=nside, - ) + coordinates = healpix.healpix_to_skycoord(unique_pixels) return self._compute_emission( freq=freq, weights=weights, - obs=obs, obs_time=obs_time, obs_pos=obs_pos, - unit_vectors=unit_vectors, + coordinates=coordinates, indicies=indicies, - pixels=unique_pixels, return_comps=return_comps, ) def get_binned_emission_ang( self, - freq: FrequencyOrWavelength, theta: SkyAngles, phi: SkyAngles, + *, nside: int, + freq: FrequencyOrWavelength, obs_time: Time, - obs: str = "earth", - obs_pos: u.Quantity[u.AU] | None = None, + obs_pos: u.Quantity[u.AU] | str = "earth", + frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, weights: Sequence[float] | npt.NDArray[np.floating] | None = None, lonlat: bool = False, return_comps: bool = False, - coord_in: Literal["E", "G", "C"] = "E", solar_cut: u.Quantity[u.deg] | None = None, ) -> u.Quantity[u.MJy / u.sr]: """Return the simulated binned zodiacal emission given angles on the sky. The pointing, for which to compute the emission, is specified in form of angles on - the sky given by `theta` and `phi`. The emission is binned to a HEALPix map with - resolution given by `nside`. + the sky given by `theta` and `phi`, matching the healpy convention. The emission is + binned to a HEALPix map with resolution given by `nside`. 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. theta: Angular co-latitude coordinate of a point, or a sequence of points, on the celestial sphere. Must be in the range [0, π] rad. Units must be either radians or degrees. phi: Angular longitude coordinate of a point, or a sequence of points, on the celestial sphere. Must be in the range [0, 2π] rad. Units must be either radians or degrees. - nside: HEALPix resolution parameter of the binned map. + nside: HEALPix resolution parameter of the pixels and the binned map. + 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. obs_time: Time of observation. - obs: Name of the Solar System observer. A list of all support observers (for a - given ephemeridis) is specified in `supported_observers` attribute of the - `Zodipy` instance. Defaults to 'earth'. - obs_pos: The heliocentric ecliptic cartesian position of the observer in AU. - Overrides the `obs` argument. Default is None. + obs_pos: The heliocentric ecliptic position of the observer in AU, or a string + representing an observer in the `astropy.coordinates.solar_system_ephemeris`. + Defaults to 'earth'. + frame: Astropy coordinate frame representing the coordinate frame of the input pointing. + Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other + alternatives are `Galactic` and `ICRS`. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). - lonlat: If True, input angles `theta`, `phi` are assumed to be longitude and + lonlat: If True, input angles (`theta`, `phi`) are assumed to be longitude and latitude, otherwise, they are co-latitude and longitude. return_comps: If True, the emission is returned component-wise. Defaults to False. - coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic - coordinates) by default. solar_cut (u.Quantity[u.deg]): Cutoff angle from the sun in degrees. The emission for all the pointing with angular distance between the sun smaller than `solar_cutoff` are masked. Defaults to `None`. @@ -298,25 +339,26 @@ def get_binned_emission_ang( emission: Simulated zodiacal emission in units of 'MJy/sr'. """ - healpix = ahpx.HEALPix(nside=nside, order="ring") theta, phi = get_validated_ang(theta=theta, phi=phi, lonlat=lonlat) - unique_angles, counts = np.unique(np.asarray([theta, phi]), return_counts=True, axis=1) - unique_pixels = hp.ang2pix(nside, *unique_angles, lonlat=lonlat) - unit_vectors = get_unit_vectors_from_ang( - coord_in=coord_in, - theta=unique_angles[0], - phi=unique_angles[1], - lonlat=lonlat, + + if not lonlat: + theta, phi = phi, (np.pi / 2) * u.rad - theta + + healpix = hp.HEALPix(nside, order="ring", frame=frame) + (theta, phi), counts = np.unique(np.vstack([theta, phi]), return_counts=True, axis=1) + coordinates = coords.SkyCoord( + theta, + phi, + frame=frame, ) + return self._compute_emission( freq=freq, weights=weights, - obs=obs, obs_time=obs_time, obs_pos=obs_pos, - unit_vectors=unit_vectors, + coordinates=coordinates, indicies=counts, - pixels=unique_pixels, healpix=healpix, return_comps=return_comps, solar_cut=solar_cut, @@ -324,15 +366,15 @@ def get_binned_emission_ang( def get_binned_emission_pix( self, - freq: FrequencyOrWavelength, pixels: Pixels, + *, nside: int, + freq: FrequencyOrWavelength, obs_time: Time, - obs: str = "earth", - obs_pos: u.Quantity[u.AU] | None = None, + obs_pos: u.Quantity[u.AU] | str = "earth", + frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, weights: Sequence[float] | npt.NDArray[np.floating] | None = None, return_comps: bool = False, - coord_in: Literal["E", "G", "C"] = "E", solar_cut: u.Quantity[u.deg] | None = None, ) -> u.Quantity[u.MJy / u.sr]: """Return the simulated binned zodiacal Emission given pixel numbers. @@ -342,48 +384,40 @@ def get_binned_emission_pix( `nside`. Args: + pixels: HEALPix pixel indicies representing points on the celestial sphere. + nside: HEALPix resolution parameter of the pixels and the binned map. 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. - pixels: HEALPix pixel indicies representing points on the celestial sphere. - nside: HEALPix resolution parameter of the pixels and the binned map. obs_time: Time of observation. - obs: Name of the Solar System observer. A list of all support observers (for a - given ephemeridis) is specified in `supported_observers` attribute of the - `Zodipy` instance. Defaults to 'earth'. - obs_pos: The heliocentric ecliptic cartesian position of the observer in AU. - Overrides the `obs` argument. Default is None. + obs_pos: The heliocentric ecliptic position of the observer in AU, or a string + representing an observer in the `astropy.coordinates.solar_system_ephemeris`. + Defaults to 'earth'. + frame: Astropy coordinate frame representing the coordinate frame of the input pointing. + Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other + alternatives are `Galactic` and `ICRS`. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. - coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic - coordinates) by default. solar_cut (u.Quantity[u.deg]): Cutoff angle from the sun in degrees. The emission - for all the pointing with angular distance between the sun smaller than - `solar_cutoff` are masked. Defaults to `None`. + for all the pointing with angular distance between the sun smaller than + `solar_cutoff` are masked. Defaults to `None`. Returns: emission: Simulated zodiacal emission in units of 'MJy/sr'. """ - healpix = ahpx.HEALPix(nside=nside, order="ring") - pixels = get_validated_pix(pixels=pixels, nside=nside) - + healpix = hp.HEALPix(nside=nside, order="ring", frame=frame) unique_pixels, counts = np.unique(pixels, return_counts=True) - unit_vectors = get_unit_vectors_from_pixels( - coord_in=coord_in, - pixels=unique_pixels, - nside=nside, - ) + coordinates = healpix.healpix_to_skycoord(unique_pixels) + return self._compute_emission( freq=freq, weights=weights, - obs=obs, obs_time=obs_time, obs_pos=obs_pos, - unit_vectors=unit_vectors, + coordinates=coordinates, indicies=counts, - pixels=unique_pixels, healpix=healpix, return_comps=return_comps, solar_cut=solar_cut, @@ -393,13 +427,11 @@ def _compute_emission( self, freq: FrequencyOrWavelength, weights: Sequence[float] | npt.NDArray[np.floating] | None, - obs: str, obs_time: Time, - unit_vectors: npt.NDArray[np.float64], + obs_pos: u.Quantity[u.AU] | str, + coordinates: coords.SkyCoord, indicies: npt.NDArray[np.int64], - obs_pos: u.Quantity[u.AU] | None = None, - pixels: npt.NDArray[np.int64] | None = None, - healpix: ahpx.HEALPix | None = None, + healpix: hp.HEALPix | None = None, return_comps: bool = False, solar_cut: u.Quantity[u.rad] | None = None, ) -> u.Quantity[u.MJy / u.sr]: @@ -417,9 +449,11 @@ def _compute_emission( bandpass, self._ipd_model, self._interpolator ) - observer_position, earth_position = get_obs_and_earth_positions( - obs=obs, obs_time=obs_time, obs_pos=obs_pos - ) + observer_position, earth_position = get_obs_and_earth_positions(obs_pos, obs_time) + + # Rotate to ecliptic coordinates to evaluate zodiacal light model + coordinates = coordinates.transform_to(coords.BarycentricMeanEcliptic) + 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. @@ -437,7 +471,7 @@ def _compute_emission( # Make table of pre-computed bandpass integrated blackbody emission. bandpass_interpolatation_table = get_bandpass_interpolation_table(bandpass) - common_integrand = partial( + common_integrand = functools.partial( EMISSION_MAPPING[type(self._ipd_model)], X_obs=observer_position, bp_interpolation_table=bandpass_interpolatation_table, @@ -457,7 +491,7 @@ def _compute_emission( else: start_chunks = np.array_split(start[comp_label], self.n_proc, axis=-1) comp_integrands = [ - partial( + functools.partial( common_integrand, u_los=np.expand_dims(unit_vectors, axis=-1), start=np.expand_dims(start, axis=-1), @@ -489,7 +523,7 @@ def _compute_emission( unit_vectors_expanded = np.expand_dims(unit_vectors, axis=-1) for idx, comp_label in enumerate(self._ipd_model.comps.keys()): - comp_integrand = partial( + comp_integrand = functools.partial( common_integrand, u_los=unit_vectors_expanded, start=np.expand_dims(start[comp_label], axis=-1), @@ -506,13 +540,14 @@ def _compute_emission( # Output is requested to be binned if healpix: + pixels = healpix.skycoord_to_healpix(coordinates) emission = np.zeros((len(self._ipd_model.comps), healpix.npix)) emission[:, pixels] = integrated_comp_emission if solar_cut is not None: observer_coords = coords.SkyCoord( *(observer_position.flatten() * u.AU), representation_type="cartesian", - frame=coords.HeliocentricTrueEcliptic, + frame=coords.BarycentricMeanEcliptic, ) sun_coords = observer_coords.copy() sun_lon = sun_coords.spherical.lon @@ -521,14 +556,11 @@ def _compute_emission( pointing_angles = coords.SkyCoord( *unit_vectors, representation_type="cartesian", - frame=coords.HeliocentricTrueEcliptic, + frame=coords.BarycentricMeanEcliptic, ) angular_distance = pointing_angles.separation(sun_coords) solar_mask = angular_distance < solar_cut - if pixels is not None: - emission[:, pixels[solar_mask]] = np.nan - else: - emission[:, solar_mask[indicies]] = np.nan + emission[:, pixels[solar_mask]] = np.nan else: emission = np.zeros((len(self._ipd_model.comps), indicies.size)) From 771a576098a111002870f15f551f88dae4aae1e1 Mon Sep 17 00:00:00 2001 From: Metin San Date: Sat, 27 Apr 2024 11:20:46 +0200 Subject: [PATCH 04/17] Update tests and example code to use new API --- .../get_bandpass_integrated_emission.py | 11 +- docs/examples/get_binned_emission.py | 11 +- .../get_binned_emission_solar_cutoff.py | 11 +- docs/examples/get_binned_gal_emission.py | 13 +- docs/examples/get_comp_binned_emission.py | 11 +- docs/examples/get_density_contour.py | 1 - docs/examples/get_emission_ang.py | 9 +- docs/examples/get_parallel_emission.py | 7 +- tests/_strategies.py | 30 ++- tests/test_get_emission.py | 2 +- tests/test_validators.py | 34 ++-- zodipy/_bandpass.py | 4 +- zodipy/_constants.py | 2 +- zodipy/_coords.py | 52 +++++ zodipy/_sky_coords.py | 63 ------ zodipy/_types.py | 2 - zodipy/_unit_vectors.py | 30 --- zodipy/_validators.py | 27 +-- zodipy/zodipy.py | 185 ++++++++---------- 19 files changed, 226 insertions(+), 279 deletions(-) create mode 100644 zodipy/_coords.py delete mode 100644 zodipy/_sky_coords.py delete mode 100644 zodipy/_unit_vectors.py diff --git a/docs/examples/get_bandpass_integrated_emission.py b/docs/examples/get_bandpass_integrated_emission.py index 132c3c6..25bd5c4 100644 --- a/docs/examples/get_bandpass_integrated_emission.py +++ b/docs/examples/get_bandpass_integrated_emission.py @@ -1,3 +1,5 @@ +from multiprocessing import cpu_count + import astropy.units as u import healpy as hp import matplotlib.pyplot as plt @@ -15,16 +17,15 @@ plt.plot(freqs, weights) plt.xlabel("Frequency [GHz]") plt.ylabel("Weights") -plt.savefig("../img/bandpass.png", dpi=300) -model = Zodipy(model="planck18") +model = Zodipy(model="planck18", n_proc=cpu_count()) emission_central_freq = model.get_binned_emission_pix( freq=center_freq, pixels=np.arange(hp.nside2npix(nside)), nside=nside, obs_time=Time("2022-03-10"), - obs="SEMB-L2", + obs_pos="SEMB-L2", ) emission_bandpass_integrated = model.get_binned_emission_pix( @@ -33,7 +34,7 @@ pixels=np.arange(hp.nside2npix(nside)), nside=nside, obs_time=Time("2022-03-10"), - obs="SEMB-L2", + obs_pos="SEMB-L2", ) hp.mollview( @@ -43,7 +44,6 @@ cmap="afmhot", norm="log", ) -plt.savefig("../img/center_freq.png", dpi=300) hp.mollview( emission_bandpass_integrated, @@ -52,5 +52,4 @@ cmap="afmhot", norm="log", ) -plt.savefig("../img/bandpass_integrated.png", dpi=300) plt.show() diff --git a/docs/examples/get_binned_emission.py b/docs/examples/get_binned_emission.py index ace5eea..fcd7858 100644 --- a/docs/examples/get_binned_emission.py +++ b/docs/examples/get_binned_emission.py @@ -1,3 +1,5 @@ +from multiprocessing import cpu_count + import astropy.units as u import healpy as hp import matplotlib.pyplot as plt @@ -6,15 +8,15 @@ from zodipy import Zodipy -model = Zodipy("planck18") +model = Zodipy("planck18", n_proc=cpu_count()) nside = 256 binned_emission = model.get_binned_emission_pix( - 857 * u.GHz, - pixels=np.arange(hp.nside2npix(nside)), + np.arange(hp.nside2npix(nside)), + freq=857 * u.GHz, nside=nside, obs_time=Time("2022-06-14"), - obs="earth", + obs_pos="earth", ) hp.mollview( @@ -25,5 +27,4 @@ max=1, cmap="afmhot", ) -plt.savefig("../img/binned.png", dpi=300) plt.show() diff --git a/docs/examples/get_binned_emission_solar_cutoff.py b/docs/examples/get_binned_emission_solar_cutoff.py index b4acbc9..8250d9f 100644 --- a/docs/examples/get_binned_emission_solar_cutoff.py +++ b/docs/examples/get_binned_emission_solar_cutoff.py @@ -1,3 +1,5 @@ +from multiprocessing import cpu_count + import astropy.units as u import healpy as hp import matplotlib.pyplot as plt @@ -6,15 +8,15 @@ from zodipy import Zodipy -model = Zodipy("dirbe") +model = Zodipy("dirbe", n_proc=cpu_count()) nside = 256 binned_emission = model.get_binned_emission_pix( - 25 * u.micron, - pixels=np.arange(hp.nside2npix(nside)), + np.arange(hp.nside2npix(nside)), nside=nside, + freq=25 * u.micron, obs_time=Time("2020-01-01"), - obs="earth", + obs_pos="earth", solar_cut=60 * u.deg, ) @@ -26,5 +28,4 @@ coord="E", cmap="afmhot", ) -plt.savefig("../img/binned_solar_cutoff.png", dpi=300) plt.show() diff --git a/docs/examples/get_binned_gal_emission.py b/docs/examples/get_binned_gal_emission.py index f7f8f6b..46fa3bd 100644 --- a/docs/examples/get_binned_gal_emission.py +++ b/docs/examples/get_binned_gal_emission.py @@ -1,3 +1,5 @@ +from multiprocessing import cpu_count + import astropy.units as u import healpy as hp import matplotlib.pyplot as plt @@ -6,16 +8,16 @@ from zodipy import Zodipy -model = Zodipy("planck18") +model = Zodipy("planck18", n_proc=cpu_count()) nside = 256 binned_emission = model.get_binned_emission_pix( - 857 * u.GHz, - pixels=np.arange(hp.nside2npix(nside)), + np.arange(hp.nside2npix(nside)), + freq=857 * u.GHz, nside=nside, obs_time=Time("2022-02-20"), - obs="earth", - coord_in="G", # Coordinates of the input pointing + obs_pos="earth", + frame="galactic", # Coordinates of the input pointing ) hp.mollview( @@ -26,5 +28,4 @@ min=0, max=1, ) -plt.savefig("../img/binned_gal.png", dpi=300) plt.show() diff --git a/docs/examples/get_comp_binned_emission.py b/docs/examples/get_comp_binned_emission.py index 1e89be3..ffdde11 100644 --- a/docs/examples/get_comp_binned_emission.py +++ b/docs/examples/get_comp_binned_emission.py @@ -1,3 +1,5 @@ +from multiprocessing import cpu_count + import astropy.units as u import healpy as hp import matplotlib.pyplot as plt @@ -6,15 +8,15 @@ from zodipy import Zodipy -model = Zodipy("dirbe") +model = Zodipy("dirbe", n_proc=cpu_count()) nside = 256 binned_emission = model.get_binned_emission_pix( - 25 * u.micron, - pixels=np.arange(hp.nside2npix(nside)), + np.arange(hp.nside2npix(nside)), + freq=25 * u.micron, nside=nside, obs_time=Time("2022-01-01"), - obs="earth", + obs_pos="earth", return_comps=True, ) fig = plt.figure(figsize=(8, 6.5), constrained_layout=True) @@ -29,5 +31,4 @@ sub=(3, 2, idx + 1), fig=fig, ) -# plt.savefig("../img/binned_comp.png", dpi=300) plt.show() diff --git a/docs/examples/get_density_contour.py b/docs/examples/get_density_contour.py index 74fdc48..733a65a 100644 --- a/docs/examples/get_density_contour.py +++ b/docs/examples/get_density_contour.py @@ -26,5 +26,4 @@ plt.title("Cross section of the interplanetary dust density (yz-plane)") plt.xlabel("x [AU]") plt.ylabel("z [AU]") -# plt.savefig("../img/density_grid.png", dpi=300) plt.show() diff --git a/docs/examples/get_emission_ang.py b/docs/examples/get_emission_ang.py index 146271a..608947b 100644 --- a/docs/examples/get_emission_ang.py +++ b/docs/examples/get_emission_ang.py @@ -1,3 +1,5 @@ +from multiprocessing import cpu_count + import astropy.units as u import matplotlib.pyplot as plt import numpy as np @@ -5,23 +7,22 @@ from zodipy import Zodipy -model = Zodipy("dirbe") +model = Zodipy("dirbe", n_proc=cpu_count()) latitudes = np.linspace(-90, 90, 10000) * u.deg longitudes = np.zeros_like(latitudes) emission = model.get_emission_ang( - 30 * u.micron, theta=longitudes, phi=latitudes, + freq=30 * u.micron, lonlat=True, obs_time=Time("2022-06-14"), - obs="earth", + obs_pos="earth", ) plt.plot(latitudes, emission) plt.xlabel("Latitude [deg]") plt.ylabel("Emission [MJy/sr]") -plt.savefig("../img/timestream.png", dpi=300) plt.show() diff --git a/docs/examples/get_parallel_emission.py b/docs/examples/get_parallel_emission.py index e5be445..a8f1afd 100644 --- a/docs/examples/get_parallel_emission.py +++ b/docs/examples/get_parallel_emission.py @@ -1,4 +1,5 @@ import time +from multiprocessing import cpu_count import astropy.units as u import healpy as hp @@ -10,15 +11,15 @@ nside = 256 pixels = np.arange(hp.nside2npix(nside)) obs_time = Time("2020-01-01") -n_proc = 8 +n_proc = cpu_count() model = Zodipy() model_parallel = Zodipy(n_proc=n_proc) start = time.perf_counter() emission = model.get_binned_emission_pix( - 40 * u.micron, pixels=pixels, + freq=40 * u.micron, nside=nside, obs_time=obs_time, ) @@ -27,9 +28,9 @@ start = time.perf_counter() emission_parallel = model_parallel.get_binned_emission_pix( - 40 * u.micron, pixels=pixels, nside=nside, + freq=40 * u.micron, obs_time=obs_time, ) print( diff --git a/tests/_strategies.py b/tests/_strategies.py index 319b5af..b09a742 100644 --- a/tests/_strategies.py +++ b/tests/_strategies.py @@ -7,9 +7,10 @@ import astropy.coordinates as coords import astropy.units as u -import healpy as hp +import astropy_healpix as hp import numpy as np import numpy.typing as npt +from astropy import time from hypothesis.extra.numpy import arrays from hypothesis.strategies import ( DrawFn, @@ -65,8 +66,8 @@ def quantities( @composite -def times(draw: DrawFn) -> times.Time: - return draw(datetimes(min_value=MIN_DATE, max_value=MAX_DATE).map(times.Time)) +def times(draw: DrawFn) -> time.Time: + return draw(datetimes(min_value=MIN_DATE, max_value=MAX_DATE).map(time.Time)) @composite @@ -76,7 +77,18 @@ def nsides(draw: Callable[[SearchStrategy[int]], int]) -> int: @composite def frames(draw: DrawFn) -> type[coords.BaseCoordinateFrame]: - return draw(sampled_from([coords.ICRS, coords.Galactic, coords.HeliocentricMeanEcliptic])) + return draw( + sampled_from( + [ + coords.BarycentricTrueEcliptic, + coords.ICRS, + coords.Galactic, + "galactic", + "barycentrictrueecliptic", + "icrs", + ] + ) + ) @composite @@ -100,10 +112,10 @@ def sky_coords(draw: DrawFn) -> coords.SkyCoord: @composite def pixels(draw: DrawFn, nside: int) -> int | list[int] | npt.NDArray[np.integer]: - npix = hp.nside2npix(nside) - pixel_strategy = integers(min_value=0, max_value=npix - 1) + healpix = hp.HEALPix(nside=nside) + pixel_strategy = integers(min_value=0, max_value=healpix.npix - 1) - shape = draw(integers(min_value=1, max_value=npix - 1)) + shape = draw(integers(min_value=1, max_value=healpix.npix - 1)) list_stategy = lists(pixel_strategy, min_size=1) array_strategy = arrays(dtype=int, shape=shape, elements=pixel_strategy) @@ -189,8 +201,8 @@ def normalize_array( @composite -def obs_positions(draw: DrawFn, model: zodipy.Zodipy, obs_time: times.Time) -> str: - def get_obs_dist(obs: str, obs_time: times.Time) -> u.Quantity[u.AU]: +def obs_positions(draw: DrawFn, model: zodipy.Zodipy, obs_time: time.Time) -> str: + def get_obs_dist(obs: str, obs_time: time.Time) -> u.Quantity[u.AU]: if obs == "semb-l2": obs_pos = ( coords.get_body("earth", obs_time) diff --git a/tests/test_get_emission.py b/tests/test_get_emission.py index a2834df..8f8f22a 100644 --- a/tests/test_get_emission.py +++ b/tests/test_get_emission.py @@ -46,7 +46,7 @@ def test_get_emission_skycoord( obs_time=time, obs_pos=observer, ) - assert emission.size == 1 + assert emission.size == coordinates.size @given(zodipy_models(), times(), nsides(), data()) diff --git a/tests/test_validators.py b/tests/test_validators.py index 79330cf..f4e5832 100644 --- a/tests/test_validators.py +++ b/tests/test_validators.py @@ -1,20 +1,18 @@ -import astropy.units as u import numpy as np import pytest -from astropy.time import Time +from astropy import time, units from hypothesis import given from hypothesis.strategies import DataObject, data -from zodipy._types import FrequencyOrWavelength from zodipy._validators import get_validated_and_normalized_weights, get_validated_freq from zodipy.zodipy import Zodipy from ._strategies import random_freqs, weights, zodipy_models -BANDPASS_FREQUENCIES = np.linspace(95, 105, 11) * u.GHz -BANDPASS_WAVELENGTHS = np.linspace(20, 25, 11) * u.micron +BANDPASS_FREQUENCIES = np.linspace(95, 105, 11) * units.GHz +BANDPASS_WAVELENGTHS = np.linspace(20, 25, 11) * units.micron BANDPASS_WEIGHTS = np.array([2, 3, 5, 9, 11, 12, 11, 9, 5, 3, 2]) -OBS_TIME = Time("2021-01-01T00:00:00") +OBS_TIME = time.Time("2021-01-01T00:00:00") @given(zodipy_models(extrapolate=False)) @@ -32,22 +30,22 @@ def test_validate_frequencies(model: Zodipy) -> None: model=model._ipd_model, extrapolate=model.extrapolate, ) - with pytest.raises(u.UnitsError): + with pytest.raises(units.UnitsError): get_validated_freq( - freq=BANDPASS_FREQUENCIES.value * u.g, + freq=BANDPASS_FREQUENCIES.value * units.g, model=model._ipd_model, extrapolate=model.extrapolate, ) - with pytest.raises(u.UnitsError): + with pytest.raises(units.UnitsError): get_validated_freq( - freq=25 * u.g, + freq=25 * units.g, model=model._ipd_model, extrapolate=model.extrapolate, ) @given(random_freqs(bandpass=True), data()) -def test_validate_weights(freq: FrequencyOrWavelength, data: DataObject) -> None: +def test_validate_weights(freq: units.Quantity, data: DataObject) -> None: """Tests that the bandpass weights are validated.""" bp_weights = data.draw(weights(freq)) bp_weights = get_validated_and_normalized_weights( @@ -115,22 +113,26 @@ 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 * u.micron, nside=32, obs_time=OBS_TIME) + model.get_emission_pix([1, 4, 5], freq=400 * units.micron, nside=32, obs_time=OBS_TIME) model = Zodipy("dirbe", extrapolate=True) - model.get_emission_pix([1, 4, 5], freq=400 * u.micron, nside=32, obs_time=OBS_TIME) + model.get_emission_pix([1, 4, 5], freq=400 * units.micron, 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 * u.micron, nside=32, obs_time=OBS_TIME) + linear = model.get_emission_pix([1, 4, 5], freq=27 * units.micron, nside=32, obs_time=OBS_TIME) model = Zodipy("dirbe", interp_kind="quadratic") - quadratic = model.get_emission_pix([1, 4, 5], freq=27 * u.micron, nside=32, obs_time=OBS_TIME) + quadratic = model.get_emission_pix( + [1, 4, 5], freq=27 * units.micron, 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 * u.micron, nside=32, obs_time=OBS_TIME) + model.get_emission_pix( + pixels=[1, 4, 5], freq=27 * units.micron, nside=32, obs_time=OBS_TIME + ) diff --git a/zodipy/_bandpass.py b/zodipy/_bandpass.py index 610201a..e5ecc5e 100644 --- a/zodipy/_bandpass.py +++ b/zodipy/_bandpass.py @@ -1,7 +1,7 @@ from __future__ import annotations from dataclasses import dataclass -from typing import TYPE_CHECKING, Sequence +from typing import TYPE_CHECKING import astropy.units as u import numpy as np @@ -44,7 +44,7 @@ def switch_convention(self) -> None: def validate_and_get_bandpass( freq: FrequencyOrWavelength, - weights: Sequence[float] | npt.NDArray[np.floating] | None, + weights: npt.ArrayLike | None, model: InterplanetaryDustModel, extrapolate: bool, ) -> Bandpass: diff --git a/zodipy/_constants.py b/zodipy/_constants.py index d630969..d513e49 100644 --- a/zodipy/_constants.py +++ b/zodipy/_constants.py @@ -18,7 +18,7 @@ R_NARROW_BANDS = 3.015 R_JUPITER = 5.2 R_KUIPER_BELT = 30 -DISTANCE_FROM_EARTH_TO_L2 = u.Quantity(0.009896235034000056, u.AU) +DISTANCE_FROM_EARTH_TO_SEMB_L2 = u.Quantity(0.009896235034000056, u.AU) N_INTERPOLATION_POINTS = 100 MIN_INTERPOLATION_GRID_TEMPERATURE = 40 diff --git a/zodipy/_coords.py b/zodipy/_coords.py new file mode 100644 index 0000000..0ef977f --- /dev/null +++ b/zodipy/_coords.py @@ -0,0 +1,52 @@ +import astropy.coordinates as coords +from astropy import time, units + +from zodipy._constants import DISTANCE_FROM_EARTH_TO_SEMB_L2 + +__all__ = ["get_earth_skycoord", "get_obs_skycoord"] + + +def get_sun_earth_moon_barycenter_skycoord( + earth_skycoord: coords.SkyCoord, +) -> coords.SkyCoord: + """Return a SkyCoord of the heliocentric position of the SEMB-L2 point. + + Note that this is an approximate position, as the SEMB-L2 point is not included in + any of the current available ephemerides. We assume that SEMB-L2 is at all times + located at a fixed distance from Earth along the vector pointing to Earth from the Sun. + """ + earth_distance = earth_skycoord.cartesian.norm() + SEMB_L2_distance = earth_distance + DISTANCE_FROM_EARTH_TO_SEMB_L2 + earth_unit_vector = earth_skycoord.cartesian.xyz / earth_distance + + return coords.SkyCoord( + *earth_unit_vector * SEMB_L2_distance, + frame=coords.HeliocentricMeanEcliptic, + representation_type="cartesian", + ) + + +def get_earth_skycoord(obs_time: time.Time) -> coords.SkyCoord: + """Return the sky coordinates of the Earth in the heliocentric frame.""" + return coords.get_body("earth", obs_time).transform_to(coords.HeliocentricMeanEcliptic) + + +def get_obs_skycoord( + obs_pos: units.Quantity | str, + obs_time: time.Time, + earth_skycoord: coords.SkyCoord, +) -> coords.SkyCoord: + """Return the sky coordinates of the observer in the heliocentric frame.""" + if isinstance(obs_pos, str): + if obs_pos.lower() == "semb-l2": + return get_sun_earth_moon_barycenter_skycoord(earth_skycoord) + return coords.get_body(obs_pos, obs_time).transform_to(coords.HeliocentricMeanEcliptic) + try: + return coords.SkyCoord( + *obs_pos.to(units.AU), + frame=coords.HeliocentricMeanEcliptic, + representation_type="cartesian", + ) + except AttributeError: + msg = "Observer position (`obs_pos`) must be a string or an astropy Quantity." + raise TypeError(msg) from AttributeError diff --git a/zodipy/_sky_coords.py b/zodipy/_sky_coords.py deleted file mode 100644 index dc4dd6f..0000000 --- a/zodipy/_sky_coords.py +++ /dev/null @@ -1,63 +0,0 @@ -import astropy.coordinates as coords -import astropy.units as u -import numpy as np -from astropy.time import Time - -from zodipy._constants import DISTANCE_FROM_EARTH_TO_L2 -from zodipy._types import NumpyArray - - -def get_obs_and_earth_positions( - obs_pos: u.Quantity[u.AU] | str, obs_time: Time -) -> tuple[NumpyArray, NumpyArray]: - """Return the position of the observer and the Earth (3, `n_pointing`, `n_gauss_quad_degree`). - - The lagrange point SEMB-L2 is not included in any of the current available - ephemerides. We implement an approximation to its position, assuming that - SEMB-L2 is at all times located at a fixed distance from Earth along the vector - pointing to Earth from the Sun. - """ - earth_position = _get_earth_position(obs_time) - if isinstance(obs_pos, str): - obs_position = _get_observer_position(obs_pos, obs_time, earth_position) - else: - try: - obs_position = obs_pos.to(u.AU) - except AttributeError: - msg = ( - "Observer position must be a string or an astropy Quantity with units of distance." - ) - raise TypeError(msg) from AttributeError - - return obs_position.reshape(3, 1, 1).value, earth_position.reshape(3, 1, 1).value - - -def _get_earth_position(obs_time: Time) -> u.Quantity[u.AU]: - """Return the position of the Earth given an ephemeris and observation time.""" - earth_skycoordinate = coords.get_body("earth", obs_time) - earth_skycoordinate = earth_skycoordinate.transform_to(coords.HeliocentricMeanEcliptic) - return earth_skycoordinate.cartesian.xyz.to(u.AU) - - -def _get_observer_position( - obs: str, obs_time: Time, earth_pos: u.Quantity[u.AU] -) -> u.Quantity[u.AU]: - """Return the position of the Earth and the observer.""" - if obs.lower() == "semb-l2": - return _get_sun_earth_moon_barycenter(earth_pos) - - observer_skycoordinate = coords.get_body(obs, obs_time) - observer_skycoordinate = observer_skycoordinate.transform_to(coords.HeliocentricMeanEcliptic) - - return observer_skycoordinate.cartesian.xyz.to(u.AU) - - -def _get_sun_earth_moon_barycenter( - earth_position: u.Quantity[u.AU], -) -> u.Quantity[u.AU]: - """Return the approximated position of SEMB-L2 given Earth's position.""" - r_earth = np.linalg.norm(earth_position) - earth_unit_vector = earth_position / r_earth - semb_l2_length = r_earth + DISTANCE_FROM_EARTH_TO_L2 - - return earth_unit_vector * semb_l2_length diff --git a/zodipy/_types.py b/zodipy/_types.py index cd2b8ed..61014ff 100644 --- a/zodipy/_types.py +++ b/zodipy/_types.py @@ -8,5 +8,3 @@ SkyAngles = Union[u.Quantity[u.deg], u.Quantity[u.rad]] FrequencyOrWavelength = Union[u.Quantity[u.Hz], u.Quantity[u.m]] ParameterDict = dict -NumpyArray = Union[npt.NDArray[np.float64], npt.NDArray[np.int64]] -MJySr = u.Quantity[u.MJy / u.sr] diff --git a/zodipy/_unit_vectors.py b/zodipy/_unit_vectors.py deleted file mode 100644 index 797de75..0000000 --- a/zodipy/_unit_vectors.py +++ /dev/null @@ -1,30 +0,0 @@ -# from __future__ import annotations - -# from typing import TYPE_CHECKING, Sequence - -# import healpy as hp -# import numpy as np - -# if TYPE_CHECKING: -# import astropy.units as u -# import numpy.typing as npt - - -# def get_unit_vectors_from_pixels( -# coord_in: str, pixels: Sequence[int] | npt.NDArray[np.int64], nside: int -# ) -> npt.NDArray[np.float64]: -# """Return ecliptic unit vectors from HEALPix pixels representing some pointing.""" -# unit_vectors = np.asarray(hp.pix2vec(nside, pixels)) -# return np.asarray(hp.Rotator(coord=[coord_in, "E"])(unit_vectors)) - - -# def get_unit_vectors_from_ang( -# coord_in: str, -# phi: u.Quantity[u.rad] | u.Quantity[u.deg], -# theta: u.Quantity[u.rad] | u.Quantity[u.deg], -# lonlat: bool = False, -# ) -> npt.NDArray[np.float64]: -# """Return ecliptic unit vectors from sky angles representing some pointing.""" -# unit_vectors = np.asarray(hp.ang2vec(theta, phi, lonlat=lonlat)).transpose() -# return np.asarray(hp.Rotator(coord=[coord_in, "E"])(unit_vectors)) - diff --git a/zodipy/_validators.py b/zodipy/_validators.py index 3f9f7e7..edee47e 100644 --- a/zodipy/_validators.py +++ b/zodipy/_validators.py @@ -1,8 +1,6 @@ -from typing import Sequence, Tuple, Union +from typing import Tuple import astropy.units as u - -# import healpy as hp import numpy as np import numpy.typing as npt @@ -49,7 +47,7 @@ def get_validated_freq( def get_validated_and_normalized_weights( - weights: Union[Sequence[float], npt.NDArray[np.floating], None], + weights: npt.ArrayLike | None, freq: FrequencyOrWavelength, ) -> npt.NDArray[np.float64]: """Validate user inputted weights.""" @@ -58,14 +56,14 @@ def get_validated_and_normalized_weights( raise ValueError(msg) if weights is not None: - if freq.size != len(weights): + normalized_weights = np.asarray(weights, dtype=np.float64) + if freq.size != len(normalized_weights): msg = "Number of frequencies and weights must be the same." raise ValueError(msg) if np.any(np.diff(freq) < 0): msg = "Bandpass frequencies must be strictly increasing." raise ValueError(msg) - normalized_weights = np.asarray(weights, dtype=np.float64) else: normalized_weights = np.array([1], dtype=np.float64) @@ -79,7 +77,7 @@ def get_validated_and_normalized_weights( def get_validated_ang( theta: SkyAngles, phi: SkyAngles, lonlat: bool ) -> Tuple[SkyAngles, SkyAngles]: - """Validate user inputted sky angles.""" + """Validate user inputted sky angles and make sure it adheres to the healpy convention.""" theta = theta.to(u.deg) if lonlat else theta.to(u.rad) phi = phi.to(u.deg) if lonlat else phi.to(u.rad) @@ -88,16 +86,7 @@ def get_validated_ang( if phi.isscalar: phi = u.Quantity([phi]) - return theta, phi - - -# def get_validated_pix(pixels: Pixels, nside: int) -> Pixels: -# """Validate user inputted pixels.""" -# if (np.max(pixels) > hp.nside2npix(nside)) or (np.min(pixels) < 0): -# msg = "invalid pixel number given nside" -# raise ValueError(msg) + if not lonlat: + theta, phi = phi, (np.pi / 2) * u.rad - theta -# if np.ndim(pixels) == 0: -# return np.expand_dims(pixels, axis=0) - -# return pixels + return theta, phi diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index 4e8b9d7..74bd7ad 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -3,31 +3,28 @@ import functools import multiprocessing import platform -from typing import TYPE_CHECKING, Callable, Sequence +from typing import TYPE_CHECKING, Callable -import astropy.coordinates as coords -import astropy.units as u import astropy_healpix as hp import numpy as np +from astropy import coordinates as coords +from astropy import units from scipy import interpolate 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_obs_skycoord from zodipy._emission import EMISSION_MAPPING from zodipy._interpolate_source import SOURCE_PARAMS_MAPPING from zodipy._ipd_comps import ComponentLabel from zodipy._ipd_dens_funcs import construct_density_partials_comps from zodipy._line_of_sight import get_line_of_sight_start_and_stop_distances -from zodipy._sky_coords import get_obs_and_earth_positions from zodipy._validators import get_validated_ang from zodipy.model_registry import model_registry if TYPE_CHECKING: import numpy.typing as npt - from astropy.time import Time - - from zodipy._types import FrequencyOrWavelength, ParameterDict, Pixels, SkyAngles - + from astropy import time PLATFORM = platform.system().lower() SYS_PROC_START_METHOD = "fork" if "windows" not in PLATFORM else None @@ -90,11 +87,11 @@ def supported_observers(self) -> list[str]: """Return a list of available observers given an ephemeris.""" return [*list(coords.solar_system_ephemeris.bodies), "semb-l2"] - def get_parameters(self) -> ParameterDict: + def get_parameters(self) -> dict: """Return a dictionary containing the interplanetary dust model parameters.""" return self._ipd_model.to_dict() - def update_parameters(self, parameters: ParameterDict) -> None: + def update_parameters(self, parameters: dict) -> None: """Update the interplanetary dust model parameters. Args: @@ -120,12 +117,12 @@ def get_emission_skycoord( self, coord: coords.SkyCoord, *, - obs_time: Time, - freq: FrequencyOrWavelength, - obs_pos: u.Quantity[u.AU] | str = "earth", - weights: Sequence[float] | npt.NDArray[np.floating] | None = None, + obs_time: time.Time, + freq: units.Quantity, + obs_pos: units.Quantity | str = "earth", + weights: npt.ArrayLike | None = None, return_comps: bool = False, - ) -> u.Quantity[u.MJy / u.sr]: + ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated zodiacal light for observations in an Astropy `SkyCoord` object. The pointing, for which to compute the emission, is specified in form of angles on @@ -136,13 +133,13 @@ def get_emission_skycoord( object must have a frame attribute representing the coordinate frame of the input pointing, for example `astropy.coordinates.Galactic`. The frame must be convertible to `BarycentricMeanEcliptic`. - obs_time: Time of observation. + obs_time: Time of observation. This should be a single observational time. 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. obs_pos: The heliocentric ecliptic position of the observer in AU, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - Defaults to 'earth'. + This should correspond to a single position. Defaults to 'earth'. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. @@ -156,7 +153,7 @@ def get_emission_skycoord( return_inverse=True, axis=1, ) - coord = coords.SkyCoord(unique_lon * u.deg, unique_lat * u.deg, frame=coord.frame) + coord = coords.SkyCoord(unique_lon * units.deg, unique_lat * units.deg, frame=coord.frame) return self._compute_emission( freq=freq, @@ -170,17 +167,17 @@ def get_emission_skycoord( def get_emission_ang( self, - theta: SkyAngles, - phi: SkyAngles, + theta: units.Quantity, + phi: units.Quantity, *, - freq: FrequencyOrWavelength, - obs_time: Time, - obs_pos: u.Quantity[u.AU] | str = "earth", + freq: units.Quantity, + obs_time: time.Time, + obs_pos: units.Quantity | str = "earth", frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, - weights: Sequence[float] | npt.NDArray[np.floating] | None = None, + weights: npt.ArrayLike | None = None, lonlat: bool = False, return_comps: bool = False, - ) -> u.Quantity[u.MJy / u.sr]: + ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated zodiacal emission given angles on the sky. The pointing, for which to compute the emission, is specified in form of angles on @@ -188,18 +185,18 @@ def get_emission_ang( Args: theta: Angular co-latitude coordinate of a point, or a sequence of points, on - the celestial sphere. Must be in the range [0, π] rad. Units must be either - radians or degrees. + the celestial sphere. Must be in the range [0, π] rad. Units must be compatible + with degrees. phi: Angular longitude coordinate of a point, or a sequence of points, on the - celestial sphere. Must be in the range [0, 2π] rad. Units must be either - radians or degrees. + celestial sphere. Must be in the range [0, 2π] rad. Units must be compatible + with degrees. 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. - obs_time: Time of observation. + obs_time: Time of observation. This should be a single observational time. obs_pos: The heliocentric ecliptic position of the observer in AU, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - Defaults to 'earth'. + This should correspond to a single position. Defaults to 'earth'. frame: Astropy coordinate frame representing the coordinate frame of the input pointing. Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other alternatives are `Galactic` and `ICRS`. @@ -214,8 +211,6 @@ def get_emission_ang( """ theta, phi = get_validated_ang(theta=theta, phi=phi, lonlat=lonlat) - if not lonlat: - theta, phi = phi, (np.pi / 2) * u.rad - theta (theta, phi), indicies = np.unique(np.stack([theta, phi]), return_inverse=True, axis=1) coordinates = coords.SkyCoord( @@ -236,16 +231,16 @@ def get_emission_ang( def get_emission_pix( self, - pixels: Pixels, + pixels: npt.ArrayLike, *, nside: int, - freq: FrequencyOrWavelength, - obs_time: Time, - obs_pos: u.Quantity[u.AU] | str = "earth", + freq: units.Quantity, + obs_time: time.Time, + obs_pos: units.Quantity | str = "earth", frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, - weights: Sequence[float] | npt.NDArray[np.floating] | None = None, + weights: npt.ArrayLike | None = None, return_comps: bool = False, - ) -> u.Quantity[u.MJy / u.sr]: + ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated zodiacal emission given pixel numbers. The pixel numbers represent the pixel indicies on a HEALPix grid with resolution @@ -257,10 +252,10 @@ def get_emission_pix( 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. - obs_time: Time of observation. + obs_time: Time of observation. This should be a single observational time. obs_pos: The heliocentric ecliptic position of the observer in AU, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - Defaults to 'earth'. + This should correspond to a single position. Defaults to 'earth'. frame: Astropy coordinate frame representing the coordinate frame of the input pointing. Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other alternatives are `Galactic` and `ICRS`. @@ -273,7 +268,6 @@ def get_emission_pix( """ healpix = hp.HEALPix(nside=nside, order="ring", frame=frame) - unique_pixels, indicies = np.unique(pixels, return_inverse=True) coordinates = healpix.healpix_to_skycoord(unique_pixels) @@ -289,19 +283,19 @@ def get_emission_pix( def get_binned_emission_ang( self, - theta: SkyAngles, - phi: SkyAngles, + theta: units.Quantity, + phi: units.Quantity, *, nside: int, - freq: FrequencyOrWavelength, - obs_time: Time, - obs_pos: u.Quantity[u.AU] | str = "earth", + freq: units.Quantity, + obs_time: time.Time, + obs_pos: units.Quantity[units.AU] | str = "earth", frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, - weights: Sequence[float] | npt.NDArray[np.floating] | None = None, + weights: npt.ArrayLike | None = None, lonlat: bool = False, return_comps: bool = False, - solar_cut: u.Quantity[u.deg] | None = None, - ) -> u.Quantity[u.MJy / u.sr]: + solar_cut: units.Quantity | None = None, + ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated binned zodiacal emission given angles on the sky. The pointing, for which to compute the emission, is specified in form of angles on @@ -319,10 +313,10 @@ def get_binned_emission_ang( 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. - obs_time: Time of observation. + obs_time: Time of observation. This should be a single observational time. obs_pos: The heliocentric ecliptic position of the observer in AU, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - Defaults to 'earth'. + This should correspond to a single position. Defaults to 'earth'. frame: Astropy coordinate frame representing the coordinate frame of the input pointing. Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other alternatives are `Galactic` and `ICRS`. @@ -331,19 +325,14 @@ def get_binned_emission_ang( lonlat: If True, input angles (`theta`, `phi`) are assumed to be longitude and latitude, otherwise, they are co-latitude and longitude. return_comps: If True, the emission is returned component-wise. Defaults to False. - solar_cut (u.Quantity[u.deg]): Cutoff angle from the sun in degrees. The emission - for all the pointing with angular distance between the sun smaller than - `solar_cutoff` are masked. Defaults to `None`. + solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular + distance between the sun smaller than `solar_cut` are masked. Defaults to `None`. Returns: emission: Simulated zodiacal emission in units of 'MJy/sr'. """ theta, phi = get_validated_ang(theta=theta, phi=phi, lonlat=lonlat) - - if not lonlat: - theta, phi = phi, (np.pi / 2) * u.rad - theta - healpix = hp.HEALPix(nside, order="ring", frame=frame) (theta, phi), counts = np.unique(np.vstack([theta, phi]), return_counts=True, axis=1) coordinates = coords.SkyCoord( @@ -366,17 +355,17 @@ def get_binned_emission_ang( def get_binned_emission_pix( self, - pixels: Pixels, + pixels: npt.ArrayLike, *, nside: int, - freq: FrequencyOrWavelength, - obs_time: Time, - obs_pos: u.Quantity[u.AU] | str = "earth", - frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, - weights: Sequence[float] | npt.NDArray[np.floating] | None = None, + freq: units.Quantity, + obs_time: time.Time, + obs_pos: units.Quantity | str = "earth", + frame: type[coords.BaseCoordinateFrame] | str = coords.BarycentricMeanEcliptic, + weights: npt.ArrayLike | None = None, return_comps: bool = False, - solar_cut: u.Quantity[u.deg] | None = None, - ) -> u.Quantity[u.MJy / u.sr]: + solar_cut: units.Quantity | None = None, + ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated binned zodiacal Emission given pixel numbers. The pixel numbers represent the pixel indicies on a HEALPix grid with resolution @@ -389,19 +378,18 @@ def get_binned_emission_pix( 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. - obs_time: Time of observation. + obs_time: Time of observation. This should be a single observational time. obs_pos: The heliocentric ecliptic position of the observer in AU, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - Defaults to 'earth'. + This should correspond to a single position. Defaults to 'earth'. frame: Astropy coordinate frame representing the coordinate frame of the input pointing. Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other alternatives are `Galactic` and `ICRS`. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. - solar_cut (u.Quantity[u.deg]): Cutoff angle from the sun in degrees. The emission - for all the pointing with angular distance between the sun smaller than - `solar_cutoff` are masked. Defaults to `None`. + solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular + distance between the sun smaller than `solar_cut` are masked. Defaults to `None`. Returns: emission: Simulated zodiacal emission in units of 'MJy/sr'. @@ -425,17 +413,17 @@ def get_binned_emission_pix( def _compute_emission( self, - freq: FrequencyOrWavelength, - weights: Sequence[float] | npt.NDArray[np.floating] | None, - obs_time: Time, - obs_pos: u.Quantity[u.AU] | str, + freq: units.Quantity, + weights: npt.ArrayLike | None, + obs_time: time.Time, + obs_pos: units.Quantity | str, coordinates: coords.SkyCoord, - indicies: npt.NDArray[np.int64], + indicies: npt.NDArray, healpix: hp.HEALPix | None = None, return_comps: bool = False, - solar_cut: u.Quantity[u.rad] | None = None, - ) -> u.Quantity[u.MJy / u.sr]: - """Compute the component-wise zodiacal emission.""" + solar_cut: units.Quantity | None = None, + ) -> units.Quantity[units.MJy / units.sr]: + """Compute the zodiacal light for a given configuration.""" bandpass = validate_and_get_bandpass( freq=freq, weights=weights, @@ -449,7 +437,8 @@ def _compute_emission( bandpass, self._ipd_model, self._interpolator ) - observer_position, earth_position = get_obs_and_earth_positions(obs_pos, obs_time) + earth_skycoord = get_earth_skycoord(obs_time) + obs_skycoord = get_obs_skycoord(obs_pos, obs_time, earth_skycoord) # Rotate to ecliptic coordinates to evaluate zodiacal light model coordinates = coordinates.transform_to(coords.BarycentricMeanEcliptic) @@ -460,25 +449,28 @@ def _compute_emission( start, stop = get_line_of_sight_start_and_stop_distances( components=self._ipd_model.comps.keys(), unit_vectors=unit_vectors, - obs_pos=observer_position, + obs_pos=obs_skycoord.cartesian.xyz.value, ) density_partials = construct_density_partials_comps( comps=self._ipd_model.comps, - dynamic_params={"X_earth": earth_position}, + dynamic_params={ + "X_earth": earth_skycoord.cartesian.xyz.value[:, np.newaxis, np.newaxis] + }, ) - # Make table of pre-computed bandpass integrated blackbody emission. bandpass_interpolatation_table = get_bandpass_interpolation_table(bandpass) common_integrand = functools.partial( EMISSION_MAPPING[type(self._ipd_model)], - X_obs=observer_position, + X_obs=obs_skycoord.cartesian.xyz.value[:, np.newaxis, np.newaxis], bp_interpolation_table=bandpass_interpolatation_table, **source_parameters["common"], ) - if self.n_proc > 1: + # 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: 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])) with multiprocessing.get_context(SYS_PROC_START_METHOD).Pool( @@ -540,25 +532,16 @@ def _compute_emission( # Output is requested to be binned if healpix: - pixels = healpix.skycoord_to_healpix(coordinates) emission = np.zeros((len(self._ipd_model.comps), healpix.npix)) + pixels = healpix.skycoord_to_healpix(coordinates) emission[:, pixels] = integrated_comp_emission if solar_cut is not None: - observer_coords = coords.SkyCoord( - *(observer_position.flatten() * u.AU), - representation_type="cartesian", - frame=coords.BarycentricMeanEcliptic, - ) - sun_coords = observer_coords.copy() - sun_lon = sun_coords.spherical.lon - sun_lon += 180 * u.deg - - pointing_angles = coords.SkyCoord( - *unit_vectors, - representation_type="cartesian", + sun_skycoord = coords.SkyCoord( + obs_skycoord.spherical.lon + 180 * units.deg, + obs_skycoord.spherical.lat, frame=coords.BarycentricMeanEcliptic, ) - angular_distance = pointing_angles.separation(sun_coords) + angular_distance = coordinates.separation(sun_skycoord) solar_mask = angular_distance < solar_cut emission[:, pixels[solar_mask]] = np.nan @@ -566,7 +549,7 @@ def _compute_emission( emission = np.zeros((len(self._ipd_model.comps), indicies.size)) emission = integrated_comp_emission[:, indicies] - emission = (emission << SPECIFIC_INTENSITY_UNITS).to(u.MJy / u.sr) + emission = (emission << SPECIFIC_INTENSITY_UNITS).to(units.MJy / units.sr) return emission if return_comps else emission.sum(axis=0) From f6ef2bc1faf17f8245267cbeb4f1c14257f84099 Mon Sep 17 00:00:00 2001 From: Metin San Date: Sat, 27 Apr 2024 11:22:41 +0200 Subject: [PATCH 05/17] Update pyproject stubs --- poetry.lock | 29 ++++++++++++++++++++++++++++- pyproject.toml | 3 ++- 2 files changed, 30 insertions(+), 2 deletions(-) diff --git a/poetry.lock b/poetry.lock index cc32659..e7e5320 100644 --- a/poetry.lock +++ b/poetry.lock @@ -52,6 +52,33 @@ recommended = ["matplotlib (>=3.3,!=3.4.0,!=3.5.2)", "scipy (>=1.5)"] test = ["pytest (>=7.0)", "pytest-astropy (>=0.10)", "pytest-astropy-header (>=0.2.1)", "pytest-doctestplus (>=0.12)", "pytest-xdist", "threadpoolctl"] test-all = ["astropy[test]", "coverage[toml]", "ipython (>=4.2)", "objgraph", "sgp4 (>=2.3)", "skyfield (>=1.20)"] +[[package]] +name = "astropy-healpix" +version = "1.0.3" +description = "BSD-licensed HEALPix for Astropy" +optional = false +python-versions = ">=3.9" +files = [ + {file = "astropy_healpix-1.0.3-cp39-abi3-macosx_10_9_x86_64.whl", hash = "sha256:1a523b4bda2618a94efad12b60431c28c0f5ef4b5afb890d21a858aef2ed779a"}, + {file = "astropy_healpix-1.0.3-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:fa34adecf7f5168d65cda05d40db95d4a1aad3e6b66f43489ed4aac471f76496"}, + {file = "astropy_healpix-1.0.3-cp39-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6a7d2163421ba66e06ed40b7ce7ce89fea60a342d035974c98c2c06c6cf187e9"}, + {file = "astropy_healpix-1.0.3-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a5b03a5b18c28ab7ac030a22157cfcd8ff951db7fc751e29e854c17d3da52537"}, + {file = "astropy_healpix-1.0.3-cp39-abi3-win32.whl", hash = "sha256:ec083f48e091baa5ce97c2e96e94b9fc8921a8c370768a92ac14e86fa3465248"}, + {file = "astropy_healpix-1.0.3-cp39-abi3-win_amd64.whl", hash = "sha256:e8c32aabfa1fa7d6fd629e3918680cd0585656cfca6ec0c5dafe9e1c7a304508"}, + {file = "astropy_healpix-1.0.3-pp39-pypy39_pp73-macosx_10_9_x86_64.whl", hash = "sha256:cce4687deb102d1e9dc40f381fcd33027cd9bd0ca7955aa0d70e464128c6914e"}, + {file = "astropy_healpix-1.0.3-pp39-pypy39_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f4d4e3e86d8d010eb50bd31ca0ba438c1b50cdc9ed251da97665a7821880206e"}, + {file = "astropy_healpix-1.0.3-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:59b1dd358c5bcd6196213845c007cee4920b2e1936234374afc7d8cfd6506cfc"}, + {file = "astropy_healpix-1.0.3.tar.gz", hash = "sha256:de5d2a7ec97b167045066971f25a150d1e4061f07159be23649634db36e79746"}, +] + +[package.dependencies] +astropy = ">=3" +numpy = ">=1.19" + +[package.extras] +docs = ["matplotlib", "sphinx-astropy"] +test = ["hypothesis", "pytest-astropy"] + [[package]] name = "astropy-iers-data" version = "0.2024.4.8.0.32.4" @@ -1894,4 +1921,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = ">=3.9" -content-hash = "e8cbfe076ced3accc312f14cb6b0b4a45bc46cd36b87c1027074b45dce5f6ea7" \ No newline at end of file +content-hash = "1ac6ed62ac4175c6cac380f70fb5d8338489432afd21906921799ecc9eeff4fc" diff --git a/pyproject.toml b/pyproject.toml index af76c4c..d317f77 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,6 +33,7 @@ healpy = "^1.16.6" astropy = ">=5.0.1" jplephem = "^2.17" scipy = "^1.13.0" +astropy-healpix = "^1.0.3" [tool.poetry.group.dev.dependencies] pytest = "^7.2.1" @@ -58,8 +59,8 @@ build-backend = "poetry.core.masonry.api" disable_error_code = ["misc"] plugins = "numpy.typing.mypy_plugin" overrides = [ - { module = "healpy.*", ignore_missing_imports = true }, { module = "astropy.*", ignore_missing_imports = true }, + { module = "astropy_healpix.*", ignore_missing_imports = true }, { module = "scipy.*", ignore_missing_imports = true }, { module = "pkg_resources.*", ignore_missing_imports = true }, ] From 13cf38d7b2cbacf5ff811498494447fa03613c17 Mon Sep 17 00:00:00 2001 From: Metin San Date: Sat, 27 Apr 2024 16:29:14 +0200 Subject: [PATCH 06/17] "test mypy precomit" --- .pre-commit-config.yaml | 6 +++++- zodipy/_coords.py | 29 ++++++++++++++++++++++------- zodipy/_validators.py | 5 ++--- 3 files changed, 29 insertions(+), 11 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 14e02c9..dfed5f2 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -7,12 +7,16 @@ repos: langauge_version: python3 - repo: https://github.com/astral-sh/ruff-pre-commit - # Ruff version. rev: v0.3.7 hooks: - id: ruff - id: ruff-format + - repo: https://github.com/pre-commit/mirrors-mypy + rev: v1.10.0 + hooks: + - id: mypy + additional_dependencies: [numpy>=1.21.0] hooks: - id: poetry-export args: ["--dev", "-f", "requirements.txt", "-o", "requirements-dev.txt"] \ No newline at end of file diff --git a/zodipy/_coords.py b/zodipy/_coords.py index 0ef977f..7c2e6c0 100644 --- a/zodipy/_coords.py +++ b/zodipy/_coords.py @@ -1,14 +1,14 @@ +from __future__ import annotations + import astropy.coordinates as coords from astropy import time, units from zodipy._constants import DISTANCE_FROM_EARTH_TO_SEMB_L2 -__all__ = ["get_earth_skycoord", "get_obs_skycoord"] +__all__ = ["get_earth_skycoord", "get_obs_skycoord", "get_frame_from_string"] -def get_sun_earth_moon_barycenter_skycoord( - earth_skycoord: coords.SkyCoord, -) -> coords.SkyCoord: +def get_sun_earth_moon_barycenter_skycoord(earth_skycoord: coords.SkyCoord) -> coords.SkyCoord: """Return a SkyCoord of the heliocentric position of the SEMB-L2 point. Note that this is an approximate position, as the SEMB-L2 point is not included in @@ -32,15 +32,14 @@ def get_earth_skycoord(obs_time: time.Time) -> coords.SkyCoord: def get_obs_skycoord( - obs_pos: units.Quantity | str, - obs_time: time.Time, - earth_skycoord: coords.SkyCoord, + obs_pos: units.Quantity | str, obs_time: time.Time, earth_skycoord: coords.SkyCoord ) -> coords.SkyCoord: """Return the sky coordinates of the observer in the heliocentric frame.""" if isinstance(obs_pos, str): if obs_pos.lower() == "semb-l2": return get_sun_earth_moon_barycenter_skycoord(earth_skycoord) return coords.get_body(obs_pos, obs_time).transform_to(coords.HeliocentricMeanEcliptic) + try: return coords.SkyCoord( *obs_pos.to(units.AU), @@ -50,3 +49,19 @@ def get_obs_skycoord( except AttributeError: msg = "Observer position (`obs_pos`) must be a string or an astropy Quantity." raise TypeError(msg) from AttributeError + + +def get_frame_from_string(frame_literal: int) -> type[coords.BaseCoordinateFrame]: + """Return the appropriate astropy coordinate frame class from a string literal.""" + if frame_literal == "E": + return coords.BarycentricMeanEcliptic + if frame_literal == "G": + return coords.Galactic + if frame_literal == "C": + return coords.ICRS + + msg = ( + f"Invalid frame literal: {frame_literal}. Must be one of 'E' (Ecliptic)," + "'G' (Galactic), or 'C' (Celestial)." + ) + raise ValueError(msg) diff --git a/zodipy/_validators.py b/zodipy/_validators.py index edee47e..9e708b9 100644 --- a/zodipy/_validators.py +++ b/zodipy/_validators.py @@ -1,4 +1,4 @@ -from typing import Tuple +from __future__ import annotations import astropy.units as u import numpy as np @@ -73,10 +73,9 @@ def get_validated_and_normalized_weights( return normalized_weights -@u.quantity_input(theta=[u.deg, u.rad], phi=[u.deg, u.rad]) def get_validated_ang( theta: SkyAngles, phi: SkyAngles, lonlat: bool -) -> Tuple[SkyAngles, SkyAngles]: +) -> tuple[SkyAngles, SkyAngles]: """Validate user inputted sky angles and make sure it adheres to the healpy convention.""" theta = theta.to(u.deg) if lonlat else theta.to(u.rad) phi = phi.to(u.deg) if lonlat else phi.to(u.rad) From 824edbdc9cbcf3648561b5eeda70452e3572da61 Mon Sep 17 00:00:00 2001 From: Metin San Date: Sat, 27 Apr 2024 16:31:33 +0200 Subject: [PATCH 07/17] Revert some breaking API changes for the healpy compatible functions --- poetry.lock | 56 +++++++-------- tests/_strategies.py | 9 ++- tests/test_get_emission.py | 50 +++++++++++--- zodipy/_coords.py | 2 +- zodipy/zodipy.py | 137 +++++++++++++++++++++++++------------ 5 files changed, 173 insertions(+), 81 deletions(-) diff --git a/poetry.lock b/poetry.lock index e7e5320..c36d214 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1105,38 +1105,38 @@ mkdocstrings = ">=0.20" [[package]] name = "mypy" -version = "1.9.0" +version = "1.10.0" description = "Optional static typing for Python" optional = false python-versions = ">=3.8" files = [ - {file = "mypy-1.9.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:f8a67616990062232ee4c3952f41c779afac41405806042a8126fe96e098419f"}, - {file = "mypy-1.9.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:d357423fa57a489e8c47b7c85dfb96698caba13d66e086b412298a1a0ea3b0ed"}, - {file = "mypy-1.9.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:49c87c15aed320de9b438ae7b00c1ac91cd393c1b854c2ce538e2a72d55df150"}, - {file = "mypy-1.9.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:48533cdd345c3c2e5ef48ba3b0d3880b257b423e7995dada04248725c6f77374"}, - {file = "mypy-1.9.0-cp310-cp310-win_amd64.whl", hash = "sha256:4d3dbd346cfec7cb98e6cbb6e0f3c23618af826316188d587d1c1bc34f0ede03"}, - {file = "mypy-1.9.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:653265f9a2784db65bfca694d1edd23093ce49740b2244cde583aeb134c008f3"}, - {file = "mypy-1.9.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3a3c007ff3ee90f69cf0a15cbcdf0995749569b86b6d2f327af01fd1b8aee9dc"}, - {file = "mypy-1.9.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2418488264eb41f69cc64a69a745fad4a8f86649af4b1041a4c64ee61fc61129"}, - {file = "mypy-1.9.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:68edad3dc7d70f2f17ae4c6c1b9471a56138ca22722487eebacfd1eb5321d612"}, - {file = "mypy-1.9.0-cp311-cp311-win_amd64.whl", hash = "sha256:85ca5fcc24f0b4aeedc1d02f93707bccc04733f21d41c88334c5482219b1ccb3"}, - {file = "mypy-1.9.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:aceb1db093b04db5cd390821464504111b8ec3e351eb85afd1433490163d60cd"}, - {file = "mypy-1.9.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:0235391f1c6f6ce487b23b9dbd1327b4ec33bb93934aa986efe8a9563d9349e6"}, - {file = "mypy-1.9.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d4d5ddc13421ba3e2e082a6c2d74c2ddb3979c39b582dacd53dd5d9431237185"}, - {file = "mypy-1.9.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:190da1ee69b427d7efa8aa0d5e5ccd67a4fb04038c380237a0d96829cb157913"}, - {file = "mypy-1.9.0-cp312-cp312-win_amd64.whl", hash = "sha256:fe28657de3bfec596bbeef01cb219833ad9d38dd5393fc649f4b366840baefe6"}, - {file = "mypy-1.9.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:e54396d70be04b34f31d2edf3362c1edd023246c82f1730bbf8768c28db5361b"}, - {file = "mypy-1.9.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:5e6061f44f2313b94f920e91b204ec600982961e07a17e0f6cd83371cb23f5c2"}, - {file = "mypy-1.9.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:81a10926e5473c5fc3da8abb04119a1f5811a236dc3a38d92015cb1e6ba4cb9e"}, - {file = "mypy-1.9.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:b685154e22e4e9199fc95f298661deea28aaede5ae16ccc8cbb1045e716b3e04"}, - {file = "mypy-1.9.0-cp38-cp38-win_amd64.whl", hash = "sha256:5d741d3fc7c4da608764073089e5f58ef6352bedc223ff58f2f038c2c4698a89"}, - {file = "mypy-1.9.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:587ce887f75dd9700252a3abbc9c97bbe165a4a630597845c61279cf32dfbf02"}, - {file = "mypy-1.9.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:f88566144752999351725ac623471661c9d1cd8caa0134ff98cceeea181789f4"}, - {file = "mypy-1.9.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:61758fabd58ce4b0720ae1e2fea5cfd4431591d6d590b197775329264f86311d"}, - {file = "mypy-1.9.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:e49499be624dead83927e70c756970a0bc8240e9f769389cdf5714b0784ca6bf"}, - {file = "mypy-1.9.0-cp39-cp39-win_amd64.whl", hash = "sha256:571741dc4194b4f82d344b15e8837e8c5fcc462d66d076748142327626a1b6e9"}, - {file = "mypy-1.9.0-py3-none-any.whl", hash = "sha256:a260627a570559181a9ea5de61ac6297aa5af202f06fd7ab093ce74e7181e43e"}, - {file = "mypy-1.9.0.tar.gz", hash = "sha256:3cc5da0127e6a478cddd906068496a97a7618a21ce9b54bde5bf7e539c7af974"}, + {file = "mypy-1.10.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:da1cbf08fb3b851ab3b9523a884c232774008267b1f83371ace57f412fe308c2"}, + {file = "mypy-1.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:12b6bfc1b1a66095ab413160a6e520e1dc076a28f3e22f7fb25ba3b000b4ef99"}, + {file = "mypy-1.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9e36fb078cce9904c7989b9693e41cb9711e0600139ce3970c6ef814b6ebc2b2"}, + {file = "mypy-1.10.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:2b0695d605ddcd3eb2f736cd8b4e388288c21e7de85001e9f85df9187f2b50f9"}, + {file = "mypy-1.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:cd777b780312ddb135bceb9bc8722a73ec95e042f911cc279e2ec3c667076051"}, + {file = "mypy-1.10.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3be66771aa5c97602f382230165b856c231d1277c511c9a8dd058be4784472e1"}, + {file = "mypy-1.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:8b2cbaca148d0754a54d44121b5825ae71868c7592a53b7292eeb0f3fdae95ee"}, + {file = "mypy-1.10.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1ec404a7cbe9fc0e92cb0e67f55ce0c025014e26d33e54d9e506a0f2d07fe5de"}, + {file = "mypy-1.10.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e22e1527dc3d4aa94311d246b59e47f6455b8729f4968765ac1eacf9a4760bc7"}, + {file = "mypy-1.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:a87dbfa85971e8d59c9cc1fcf534efe664d8949e4c0b6b44e8ca548e746a8d53"}, + {file = "mypy-1.10.0-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:a781f6ad4bab20eef8b65174a57e5203f4be627b46291f4589879bf4e257b97b"}, + {file = "mypy-1.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b808e12113505b97d9023b0b5e0c0705a90571c6feefc6f215c1df9381256e30"}, + {file = "mypy-1.10.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:8f55583b12156c399dce2df7d16f8a5095291354f1e839c252ec6c0611e86e2e"}, + {file = "mypy-1.10.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:4cf18f9d0efa1b16478c4c129eabec36148032575391095f73cae2e722fcf9d5"}, + {file = "mypy-1.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:bc6ac273b23c6b82da3bb25f4136c4fd42665f17f2cd850771cb600bdd2ebeda"}, + {file = "mypy-1.10.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:9fd50226364cd2737351c79807775136b0abe084433b55b2e29181a4c3c878c0"}, + {file = "mypy-1.10.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:f90cff89eea89273727d8783fef5d4a934be2fdca11b47def50cf5d311aff727"}, + {file = "mypy-1.10.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fcfc70599efde5c67862a07a1aaf50e55bce629ace26bb19dc17cece5dd31ca4"}, + {file = "mypy-1.10.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:075cbf81f3e134eadaf247de187bd604748171d6b79736fa9b6c9685b4083061"}, + {file = "mypy-1.10.0-cp38-cp38-win_amd64.whl", hash = "sha256:3f298531bca95ff615b6e9f2fc0333aae27fa48052903a0ac90215021cdcfa4f"}, + {file = "mypy-1.10.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:fa7ef5244615a2523b56c034becde4e9e3f9b034854c93639adb667ec9ec2976"}, + {file = "mypy-1.10.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:3236a4c8f535a0631f85f5fcdffba71c7feeef76a6002fcba7c1a8e57c8be1ec"}, + {file = "mypy-1.10.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4a2b5cdbb5dd35aa08ea9114436e0d79aceb2f38e32c21684dcf8e24e1e92821"}, + {file = "mypy-1.10.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:92f93b21c0fe73dc00abf91022234c79d793318b8a96faac147cd579c1671746"}, + {file = "mypy-1.10.0-cp39-cp39-win_amd64.whl", hash = "sha256:28d0e038361b45f099cc086d9dd99c15ff14d0188f44ac883010e172ce86c38a"}, + {file = "mypy-1.10.0-py3-none-any.whl", hash = "sha256:f8c083976eb530019175aabadb60921e73b4f45736760826aa1689dda8208aee"}, + {file = "mypy-1.10.0.tar.gz", hash = "sha256:3d087fcbec056c4ee34974da493a826ce316947485cef3901f511848e687c131"}, ] [package.dependencies] diff --git a/tests/_strategies.py b/tests/_strategies.py index b09a742..8bdf8e0 100644 --- a/tests/_strategies.py +++ b/tests/_strategies.py @@ -91,11 +91,16 @@ def frames(draw: DrawFn) -> type[coords.BaseCoordinateFrame]: ) +@composite +def coords_in(draw: DrawFn) -> str: + return draw(sampled_from(["E", "G", "C"])) + + @composite def sky_coords(draw: DrawFn) -> coords.SkyCoord: theta_strategy = floats(min_value=0, max_value=360) phi_strategy = floats(min_value=-90, max_value=90) - + obs_time = draw(times()) shape = draw(integers(min_value=1, max_value=MAX_ANGELS_LEN)) theta_array_strategy = arrays(dtype=float, shape=shape, elements=theta_strategy).map( @@ -107,7 +112,7 @@ def sky_coords(draw: DrawFn) -> coords.SkyCoord: frame = draw(frames()) lon = draw(theta_array_strategy) lat = draw(phi_array_strategy) - return coords.SkyCoord(lon, lat, frame=frame) + return coords.SkyCoord(lon, lat, frame=frame, obstime=obs_time) @composite diff --git a/tests/test_get_emission.py b/tests/test_get_emission.py index 8f8f22a..a19b9a4 100644 --- a/tests/test_get_emission.py +++ b/tests/test_get_emission.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import Literal + import astropy.coordinates as coords import astropy.units as u import healpy as hp @@ -13,6 +15,7 @@ from ._strategies import ( angles, + coords_in, freqs, nsides, obs_positions, @@ -29,33 +32,55 @@ DIRBE_START_DAY = Time("1990-01-01") -@given(zodipy_models(), times(), sky_coords(), data()) +@given(zodipy_models(), sky_coords(), data()) @settings(deadline=None) def test_get_emission_skycoord( model: Zodipy, - time: Time, coordinates: coords.SkyCoord, data: DataObject, ) -> None: """Property test for get_emission_skycoord.""" - observer = data.draw(obs_positions(model, time)) + observer = data.draw(obs_positions(model, coordinates.obstime)) frequency = data.draw(freqs(model)) emission = model.get_emission_skycoord( coordinates, freq=frequency, - obs_time=time, obs_pos=observer, ) assert emission.size == coordinates.size -@given(zodipy_models(), times(), nsides(), data()) +@given(zodipy_models(), sky_coords(), nsides(), data()) +@settings(deadline=None) +def test_get_binned_skycoord( + model: Zodipy, + coordinates: coords.SkyCoord, + nside: int, + data: DataObject, +) -> 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=nside, + solar_cut=data.draw(quantities(20, 50, u.deg)) if cut_solar else None, + ) + assert emission_binned.shape == (hp.nside2npix(nside),) + + +@given(zodipy_models(), times(), nsides(), data(), coords_in()) @settings(deadline=None) def test_get_emission_pix( model: Zodipy, time: Time, nside: int, data: DataObject, + coord_in: Literal["E", "G", "C"], ) -> None: """Property test for get_emission_pix.""" observer = data.draw(obs_positions(model, time)) @@ -67,17 +92,19 @@ def test_get_emission_pix( freq=frequency, obs_time=time, obs_pos=observer, + coord_in=coord_in, ) assert np.size(emission) == np.size(pix) -@given(zodipy_models(), times(), nsides(), data()) +@given(zodipy_models(), times(), nsides(), data(), coords_in()) @settings(deadline=None) def test_get_binned_emission_pix( model: Zodipy, time: Time, nside: int, data: DataObject, + coord_in: Literal["E", "G", "C"], ) -> None: """Property test for get_binned_emission_pix.""" observer = data.draw(obs_positions(model, time)) @@ -91,17 +118,19 @@ def test_get_binned_emission_pix( obs_time=time, obs_pos=observer, solar_cut=data.draw(quantities(20, 50, u.deg)) if cut_solar else None, + coord_in=coord_in, ) assert emission_binned.shape == (hp.nside2npix(nside),) -@given(zodipy_models(), times(), angles(), data()) +@given(zodipy_models(), times(), angles(), data(), coords_in()) @settings(deadline=None) def test_get_emission_ang( model: Zodipy, time: Time, angles: tuple[u.Quantity[u.deg], u.Quantity[u.deg]], data: DataObject, + coord_in: Literal["E", "G", "C"], ) -> None: """Property test for get_emission_ang.""" theta, phi = angles @@ -115,11 +144,12 @@ def test_get_emission_ang( freq=frequency, obs_time=time, obs_pos=observer, + coord_in=coord_in, ) assert emission.size == theta.size == phi.size -@given(zodipy_models(), times(), nsides(), angles(), data()) +@given(zodipy_models(), times(), nsides(), angles(), data(), coords_in()) @settings(deadline=None) def test_get_binned_emission_ang( model: Zodipy, @@ -127,12 +157,14 @@ def test_get_binned_emission_ang( nside: int, angles: tuple[u.Quantity[u.deg], u.Quantity[u.deg]], data: DataObject, + coord_in: Literal["E", "G", "C"], ) -> None: """Property test for get_binned_emission_pix.""" 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, @@ -141,6 +173,8 @@ def test_get_binned_emission_ang( freq=frequency, obs_time=time, obs_pos=observer, + coord_in=coord_in, + solar_cut=data.draw(quantities(20, 50, u.deg)) if cut_solar else None, ) assert emission_binned.shape == (hp.nside2npix(nside),) diff --git a/zodipy/_coords.py b/zodipy/_coords.py index 7c2e6c0..8a6a5b4 100644 --- a/zodipy/_coords.py +++ b/zodipy/_coords.py @@ -51,7 +51,7 @@ def get_obs_skycoord( raise TypeError(msg) from AttributeError -def get_frame_from_string(frame_literal: int) -> type[coords.BaseCoordinateFrame]: +def get_frame_from_string(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/zodipy.py b/zodipy/zodipy.py index 74bd7ad..52d2efb 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -3,7 +3,7 @@ import functools import multiprocessing import platform -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING, Callable, Literal import astropy_healpix as hp import numpy as np @@ -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_obs_skycoord +from zodipy._coords import get_earth_skycoord, get_frame_from_string, get_obs_skycoord from zodipy._emission import EMISSION_MAPPING from zodipy._interpolate_source import SOURCE_PARAMS_MAPPING from zodipy._ipd_comps import ComponentLabel @@ -117,31 +117,28 @@ def get_emission_skycoord( self, coord: coords.SkyCoord, *, - obs_time: time.Time, freq: units.Quantity, - obs_pos: units.Quantity | str = "earth", weights: npt.ArrayLike | None = None, + obs_pos: units.Quantity | str = "earth", return_comps: bool = False, ) -> units.Quantity[units.MJy / units.sr]: - """Return the simulated zodiacal light for observations in an Astropy `SkyCoord` object. - - The pointing, for which to compute the emission, is specified in form of angles on - the sky given by `theta` and `phi`. + """Return the simulated zodiacal light for all observations in a `SkyCoord` object. Args: - coord: Astropy `SkyCoord` object representing the pointing on the sky. This - object must have a frame attribute representing the coordinate frame of the - input pointing, for example `astropy.coordinates.Galactic`. The frame must - be convertible to `BarycentricMeanEcliptic`. - obs_time: Time of observation. This should be a single observational time. + coord: `astropy.coordinates.SkyCoord` object representing the observations for which to + simulate the zodiacal light. The `frame` and `obstime` attributes of the `SkyCoord` + object must be set. The `obstime` attribute should correspond to a single + observational time for which the zodiacal light is assumed to be stationary. + Additionally, the frame must be convertible to + `astropy.coordinates.BarycentricMeanEcliptic`. 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). obs_pos: The heliocentric ecliptic position of the observer in AU, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. - weights: Bandpass weights corresponding the the frequencies in `freq`. The weights - are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. Returns: @@ -153,13 +150,15 @@ def get_emission_skycoord( return_inverse=True, axis=1, ) - coord = coords.SkyCoord(unique_lon * units.deg, unique_lat * units.deg, frame=coord.frame) + coord = coords.SkyCoord( + unique_lon * units.deg, unique_lat * units.deg, frame=coord.frame, obstime=coord.obstime + ) return self._compute_emission( freq=freq, weights=weights, - obs_time=obs_time, obs_pos=obs_pos, + obs_time=coord.obstime, coordinates=coord, indicies=indicies, return_comps=return_comps, @@ -170,12 +169,12 @@ def get_emission_ang( theta: units.Quantity, phi: units.Quantity, *, + lonlat: bool = False, freq: units.Quantity, obs_time: time.Time, obs_pos: units.Quantity | str = "earth", - frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, + coord_in: Literal["E", "G", "C"] = "E", weights: npt.ArrayLike | None = None, - lonlat: bool = False, return_comps: bool = False, ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated zodiacal emission given angles on the sky. @@ -190,6 +189,8 @@ def get_emission_ang( phi: Angular longitude coordinate of a point, or a sequence of points, on the celestial sphere. Must be in the range [0, 2π] rad. Units must be compatible with degrees. + lonlat: If True, input angles (`theta`, `phi`) are assumed to be longitude and + latitude, otherwise, they are co-latitude and longitude. 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. @@ -197,13 +198,10 @@ def get_emission_ang( obs_pos: The heliocentric ecliptic position of the observer in AU, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. - frame: Astropy coordinate frame representing the coordinate frame of the input pointing. - Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other - alternatives are `Galactic` and `ICRS`. + coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic + coordinates) by default. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). - lonlat: If True, input angles (`theta`, `phi`) are assumed to be longitude and - latitude, otherwise, they are co-latitude and longitude. return_comps: If True, the emission is returned component-wise. Defaults to False. Returns: @@ -213,11 +211,8 @@ 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) - coordinates = coords.SkyCoord( - theta, - phi, - frame=frame, - ) + frame = get_frame_from_string(coord_in) + coordinates = coords.SkyCoord(theta, phi, frame=frame) return self._compute_emission( freq=freq, @@ -237,7 +232,7 @@ def get_emission_pix( freq: units.Quantity, obs_time: time.Time, obs_pos: units.Quantity | str = "earth", - frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, + coord_in: Literal["E", "G", "C"] = "E", weights: npt.ArrayLike | None = None, return_comps: bool = False, ) -> units.Quantity[units.MJy / units.sr]: @@ -256,9 +251,8 @@ def get_emission_pix( obs_pos: The heliocentric ecliptic position of the observer in AU, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. - frame: Astropy coordinate frame representing the coordinate frame of the input pointing. - Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other - alternatives are `Galactic` and `ICRS`. + coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic + coordinates) by default. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. @@ -267,6 +261,7 @@ def get_emission_pix( emission: Simulated zodiacal emission in units of 'MJy/sr'. """ + frame = get_frame_from_string(coord_in) healpix = hp.HEALPix(nside=nside, order="ring", frame=frame) unique_pixels, indicies = np.unique(pixels, return_inverse=True) coordinates = healpix.healpix_to_skycoord(unique_pixels) @@ -281,18 +276,76 @@ def get_emission_pix( return_comps=return_comps, ) + def get_binned_emission_skycoord( + self, + coord: coords.SkyCoord, + *, + nside: int, + freq: units.Quantity, + weights: npt.ArrayLike | None = None, + obs_pos: units.Quantity | str = "earth", + return_comps: bool = False, + solar_cut: units.Quantity | None = None, + ) -> units.Quantity[units.MJy / units.sr]: + """Return the simulated binned zodiacal light for all observations in a `SkyCoord` object. + + Args: + coord: `astropy.coordinates.SkyCoord` object representing the observations for which to + simulate the zodiacal light. The `frame` and `obstime` attributes of the `SkyCoord` + object must be set. The `obstime` attribute should correspond to a single + observational time for which the zodiacal light is assumed to be stationary. + Additionally, the frame must be convertible to + `astropy.coordinates.BarycentricMeanEcliptic`. + nside: HEALPix resolution parameter of the pixels and the binned map. + 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). + obs_pos: The heliocentric ecliptic position of the observer in AU, or a string + representing an observer in the `astropy.coordinates.solar_system_ephemeris`. + This should correspond to a single position. Defaults to 'earth'. + return_comps: If True, the emission is returned component-wise. Defaults to False. + solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular + distance between the sun smaller than `solar_cut` are masked. Defaults to `None`. + + Returns: + emission: Simulated zodiacal light in units of 'MJy/sr'. + + """ + (unique_lon, unique_lat), indicies = np.unique( + np.vstack([coord.spherical.lon.value, coord.spherical.lat.value]), + return_inverse=True, + axis=1, + ) + coord = coords.SkyCoord( + unique_lon * units.deg, unique_lat * units.deg, frame=coord.frame, obstime=coord.obstime + ) + healpix = hp.HEALPix(nside, order="ring", frame=coord.frame) + return self._compute_emission( + freq=freq, + weights=weights, + obs_pos=obs_pos, + obs_time=coord.obstime, + coordinates=coord, + indicies=indicies, + healpix=healpix, + return_comps=return_comps, + solar_cut=solar_cut, + ) + def get_binned_emission_ang( self, theta: units.Quantity, phi: units.Quantity, *, + lonlat: bool = False, nside: int, freq: units.Quantity, obs_time: time.Time, obs_pos: units.Quantity[units.AU] | str = "earth", - frame: type[coords.BaseCoordinateFrame] = coords.BarycentricMeanEcliptic, + coord_in: Literal["E", "G", "C"] = "E", weights: npt.ArrayLike | None = None, - lonlat: bool = False, return_comps: bool = False, solar_cut: units.Quantity | None = None, ) -> units.Quantity[units.MJy / units.sr]: @@ -317,9 +370,8 @@ def get_binned_emission_ang( obs_pos: The heliocentric ecliptic position of the observer in AU, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. - frame: Astropy coordinate frame representing the coordinate frame of the input pointing. - Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other - alternatives are `Galactic` and `ICRS`. + coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic + coordinates) by default. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). lonlat: If True, input angles (`theta`, `phi`) are assumed to be longitude and @@ -333,6 +385,7 @@ def get_binned_emission_ang( """ theta, phi = get_validated_ang(theta=theta, phi=phi, lonlat=lonlat) + frame = get_frame_from_string(coord_in) healpix = hp.HEALPix(nside, order="ring", frame=frame) (theta, phi), counts = np.unique(np.vstack([theta, phi]), return_counts=True, axis=1) coordinates = coords.SkyCoord( @@ -361,7 +414,7 @@ def get_binned_emission_pix( freq: units.Quantity, obs_time: time.Time, obs_pos: units.Quantity | str = "earth", - frame: type[coords.BaseCoordinateFrame] | str = coords.BarycentricMeanEcliptic, + coord_in: Literal["E", "G", "C"] = "E", weights: npt.ArrayLike | None = None, return_comps: bool = False, solar_cut: units.Quantity | None = None, @@ -382,9 +435,8 @@ def get_binned_emission_pix( obs_pos: The heliocentric ecliptic position of the observer in AU, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. - frame: Astropy coordinate frame representing the coordinate frame of the input pointing. - Default is `BarycentricMeanEcliptic`, corresponding to ecliptic coordinates. Other - alternatives are `Galactic` and `ICRS`. + coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic + coordinates) by default. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. @@ -395,6 +447,7 @@ def get_binned_emission_pix( emission: Simulated zodiacal emission in units of 'MJy/sr'. """ + frame = get_frame_from_string(coord_in) healpix = hp.HEALPix(nside=nside, order="ring", frame=frame) unique_pixels, counts = np.unique(pixels, return_counts=True) coordinates = healpix.healpix_to_skycoord(unique_pixels) From 2062d4803241c9ea59937925473b6fdb45fe14de Mon Sep 17 00:00:00 2001 From: Metin San Date: Sat, 27 Apr 2024 16:39:07 +0200 Subject: [PATCH 08/17] Remove healpy as a dependency. Add healpy as a dev dependency for plotting --- poetry.lock | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/poetry.lock b/poetry.lock index c36d214..c7e5082 100644 --- a/poetry.lock +++ b/poetry.lock @@ -1921,4 +1921,4 @@ testing = ["big-O", "jaraco.functools", "jaraco.itertools", "more-itertools", "p [metadata] lock-version = "2.0" python-versions = ">=3.9" -content-hash = "1ac6ed62ac4175c6cac380f70fb5d8338489432afd21906921799ecc9eeff4fc" +content-hash = "017bc2cb411ec2db40e207b1a78d7be254cae5e1339851a7d391fa237d722824" diff --git a/pyproject.toml b/pyproject.toml index d317f77..91eb80a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,7 +29,6 @@ repository = "https://github.com/cosmoglobe/zodipy" [tool.poetry.dependencies] python = ">=3.9" numpy = "^1.26.4" -healpy = "^1.16.6" astropy = ">=5.0.1" jplephem = "^2.17" scipy = "^1.13.0" @@ -50,6 +49,7 @@ mkdocstrings-python = "^1.7.3" ruff = "^0.3.7" markdown = "<3.4.0" hypothesis = "^6.99.11" +healpy = "^1.16.6" [build-system] requires = ["poetry-core>=1.0.0"] From 53a1a1fa18406b3f1227d10a4a67f2da6bc0b610 Mon Sep 17 00:00:00 2001 From: Metin San Date: Sat, 27 Apr 2024 19:58:50 +0200 Subject: [PATCH 09/17] Updates docs --- README.md | 3 ++- docs/install.md | 2 +- zodipy/_line_of_sight.py | 7 +++--- zodipy/_types.py | 3 +++ zodipy/zodipy.py | 52 ++++++++++++++++++++++------------------ 5 files changed, 39 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index 4e5e002..215824c 100644 --- a/README.md +++ b/README.md @@ -57,8 +57,8 @@ ZodiPy is installed using `pip install zodipy`. ## Dependencies ZodiPy supports all Python versions >= 3.9, and has the following dependencies: - [Astropy](https://www.astropy.org/) (>=5.0.1) +- [Astropy-healpix](https://astropy-healpix.readthedocs.io/en/latest/) - [NumPy](https://numpy.org/) -- [healpy](https://healpy.readthedocs.io/en/latest/) - [jplephem](https://pypi.org/project/jplephem/) - [SciPy](https://scipy.org/) @@ -67,6 +67,7 @@ Contributing developers will need to download the following additional dependenc - pytest - pytest-cov - hypothesis +- healpy - coverage - ruff - mypy diff --git a/docs/install.md b/docs/install.md index 179e612..3a405f4 100644 --- a/docs/install.md +++ b/docs/install.md @@ -13,7 +13,7 @@ pip install zodipy ZodiPy has the following dependencies (these are automatically downloaded alongside ZodiPy): - [Astropy](https://www.astropy.org) (>= 5.0.1) +- [Astropy-healpix](https://astropy-healpix.readthedocs.io/en/latest/) - [NumPy](https://numpy.org) -- [healpy](https://healpy.readthedocs.io/en/latest/) - [jplehem](https://pypi.org/project/jplephem/) - [SciPy](https://scipy.org/) diff --git a/zodipy/_line_of_sight.py b/zodipy/_line_of_sight.py index 7f55a4a..19efb27 100644 --- a/zodipy/_line_of_sight.py +++ b/zodipy/_line_of_sight.py @@ -11,8 +11,9 @@ if TYPE_CHECKING: import numpy.typing as npt + from zodipy._types import Float -DIRBE_CUTOFFS: dict[ComponentLabel, tuple[float | np.float64, float]] = { +DIRBE_CUTOFFS: dict[ComponentLabel, tuple[Float, Float]] = { ComponentLabel.CLOUD: (R_0, R_JUPITER), ComponentLabel.BAND1: (R_0, R_JUPITER), ComponentLabel.BAND2: (R_0, R_JUPITER), @@ -21,7 +22,7 @@ ComponentLabel.FEATURE: (R_EARTH - 0.2, R_EARTH + 0.2), } -RRM_CUTOFFS: dict[ComponentLabel, tuple[float | np.float64, float]] = { +RRM_CUTOFFS: dict[ComponentLabel, tuple[Float, Float]] = { ComponentLabel.FAN: (R_0, RRM[ComponentLabel.FAN].R_outer), # type: ignore ComponentLabel.INNER_NARROW_BAND: ( RRM[ComponentLabel.INNER_NARROW_BAND].R_inner, # type: ignore @@ -78,7 +79,7 @@ def get_sphere_intersection( return np.maximum(q, c / q) -def get_line_of_sight_start_and_stop_distances( +def get_line_of_sight_range( components: Iterable[ComponentLabel], unit_vectors: npt.NDArray[np.float64], obs_pos: npt.NDArray[np.float64], diff --git a/zodipy/_types.py b/zodipy/_types.py index 61014ff..c887a5a 100644 --- a/zodipy/_types.py +++ b/zodipy/_types.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from typing import Sequence, Union import astropy.units as u @@ -8,3 +10,4 @@ SkyAngles = Union[u.Quantity[u.deg], u.Quantity[u.rad]] FrequencyOrWavelength = Union[u.Quantity[u.Hz], u.Quantity[u.m]] ParameterDict = dict +Float = Union[float, np.floating] diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index 52d2efb..3474fa4 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -18,7 +18,7 @@ from zodipy._interpolate_source import SOURCE_PARAMS_MAPPING from zodipy._ipd_comps import ComponentLabel from zodipy._ipd_dens_funcs import construct_density_partials_comps -from zodipy._line_of_sight import get_line_of_sight_start_and_stop_distances +from zodipy._line_of_sight import get_line_of_sight_range from zodipy._validators import get_validated_ang from zodipy.model_registry import model_registry @@ -136,9 +136,9 @@ def get_emission_skycoord( 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). - obs_pos: The heliocentric ecliptic position of the observer in AU, or a string - representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - This should correspond to a single position. Defaults to 'earth'. + obs_pos: The heliocentric ecliptic position of the observer, or a string representing + an observer in the `astropy.coordinates.solar_system_ephemeris`. This should + correspond to a single position. Defaults to 'earth'. return_comps: If True, the emission is returned component-wise. Defaults to False. Returns: @@ -151,7 +151,10 @@ def get_emission_skycoord( axis=1, ) coord = coords.SkyCoord( - unique_lon * units.deg, unique_lat * units.deg, frame=coord.frame, obstime=coord.obstime + unique_lon * units.deg, + unique_lat * units.deg, + frame=coord.frame, + obstime=coord.obstime, ) return self._compute_emission( @@ -195,9 +198,9 @@ def get_emission_ang( a bandpass over which to evaluate the zodiacal emission. The frequencies must be strictly increasing. obs_time: Time of observation. This should be a single observational time. - obs_pos: The heliocentric ecliptic position of the observer in AU, or a string - representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - This should correspond to a single position. Defaults to 'earth'. + obs_pos: The heliocentric ecliptic position of the observer, or a string representing + an observer in the `astropy.coordinates.solar_system_ephemeris`. This should + correspond to a single position. Defaults to 'earth'. coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic coordinates) by default. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights @@ -248,9 +251,9 @@ def get_emission_pix( a bandpass over which to evaluate the zodiacal emission. The frequencies must be strictly increasing. obs_time: Time of observation. This should be a single observational time. - obs_pos: The heliocentric ecliptic position of the observer in AU, or a string - representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - This should correspond to a single position. Defaults to 'earth'. + obs_pos: The heliocentric ecliptic position of the observer, or a string representing + an observer in the `astropy.coordinates.solar_system_ephemeris`. This should + correspond to a single position. Defaults to 'earth'. coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic coordinates) by default. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights @@ -302,9 +305,9 @@ def get_binned_emission_skycoord( 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). - obs_pos: The heliocentric ecliptic position of the observer in AU, or a string - representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - This should correspond to a single position. Defaults to 'earth'. + obs_pos: The heliocentric ecliptic position of the observer, or a string representing + an observer in the `astropy.coordinates.solar_system_ephemeris`. This should + correspond to a single position. Defaults to 'earth'. return_comps: If True, the emission is returned component-wise. Defaults to False. solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular distance between the sun smaller than `solar_cut` are masked. Defaults to `None`. @@ -319,7 +322,10 @@ def get_binned_emission_skycoord( axis=1, ) coord = coords.SkyCoord( - unique_lon * units.deg, unique_lat * units.deg, frame=coord.frame, obstime=coord.obstime + unique_lon * units.deg, + unique_lat * units.deg, + frame=coord.frame, + obstime=coord.obstime, ) healpix = hp.HEALPix(nside, order="ring", frame=coord.frame) return self._compute_emission( @@ -343,7 +349,7 @@ def get_binned_emission_ang( nside: int, freq: units.Quantity, obs_time: time.Time, - obs_pos: units.Quantity[units.AU] | str = "earth", + obs_pos: units.Quantity | str = "earth", coord_in: Literal["E", "G", "C"] = "E", weights: npt.ArrayLike | None = None, return_comps: bool = False, @@ -367,9 +373,9 @@ def get_binned_emission_ang( a bandpass over which to evaluate the zodiacal emission. The frequencies must be strictly increasing. obs_time: Time of observation. This should be a single observational time. - obs_pos: The heliocentric ecliptic position of the observer in AU, or a string - representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - This should correspond to a single position. Defaults to 'earth'. + obs_pos: The heliocentric ecliptic position of the observer, or a string representing + an observer in the `astropy.coordinates.solar_system_ephemeris`. This should + correspond to a single position. Defaults to 'earth'. coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic coordinates) by default. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights @@ -432,9 +438,9 @@ def get_binned_emission_pix( a bandpass over which to evaluate the zodiacal emission. The frequencies must be strictly increasing. obs_time: Time of observation. This should be a single observational time. - obs_pos: The heliocentric ecliptic position of the observer in AU, or a string - representing an observer in the `astropy.coordinates.solar_system_ephemeris`. - This should correspond to a single position. Defaults to 'earth'. + obs_pos: The heliocentric ecliptic position of the observer, or a string representing + an observer in the `astropy.coordinates.solar_system_ephemeris`. This should + correspond to a single position. Defaults to 'earth'. coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic coordinates) by default. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights @@ -499,7 +505,7 @@ def _compute_emission( # 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_start_and_stop_distances( + start, stop = get_line_of_sight_range( components=self._ipd_model.comps.keys(), unit_vectors=unit_vectors, obs_pos=obs_skycoord.cartesian.xyz.value, From a46cb20b7432b6d243979d80c07637afee2f2a1b Mon Sep 17 00:00:00 2001 From: Metin San Date: Sun, 28 Apr 2024 09:50:45 +0200 Subject: [PATCH 10/17] Add support for specifying healpix order --- tests/_strategies.py | 7 +- tests/test_get_emission.py | 190 ++++++++++++++++++++++++++----------- zodipy/_line_of_sight.py | 15 +-- zodipy/_validators.py | 41 ++++---- zodipy/zodipy.py | 26 +++-- 5 files changed, 184 insertions(+), 95 deletions(-) diff --git a/tests/_strategies.py b/tests/_strategies.py index 8bdf8e0..318b342 100644 --- a/tests/_strategies.py +++ b/tests/_strategies.py @@ -3,7 +3,7 @@ import datetime from functools import partial from math import log2 -from typing import Any, Callable, Sequence +from typing import Any, Callable, Literal, Sequence import astropy.coordinates as coords import astropy.units as u @@ -71,8 +71,9 @@ def times(draw: DrawFn) -> time.Time: @composite -def nsides(draw: Callable[[SearchStrategy[int]], int]) -> int: - return draw(integers(min_value=MIN_NSIDE_EXP, max_value=MAX_NSIDE_EXP).map(partial(pow, 2))) +def healpixes(draw: DrawFn, order: Literal["ring", "nested"] = "ring") -> hp.HEALPix: + nside = draw(integers(min_value=MIN_NSIDE_EXP, max_value=MAX_NSIDE_EXP).map(partial(pow, 2))) + return hp.HEALPix(nside=nside, order=order) @composite diff --git a/tests/test_get_emission.py b/tests/test_get_emission.py index a19b9a4..e772b29 100644 --- a/tests/test_get_emission.py +++ b/tests/test_get_emission.py @@ -4,7 +4,7 @@ import astropy.coordinates as coords import astropy.units as u -import healpy as hp +import astropy_healpix as hp import numpy as np import pytest from astropy.time import Time, TimeDelta @@ -17,7 +17,7 @@ angles, coords_in, freqs, - nsides, + healpixes, obs_positions, pixels, quantities, @@ -50,12 +50,12 @@ def test_get_emission_skycoord( assert emission.size == coordinates.size -@given(zodipy_models(), sky_coords(), nsides(), data()) +@given(zodipy_models(), sky_coords(), healpixes(), data()) @settings(deadline=None) def test_get_binned_skycoord( model: Zodipy, coordinates: coords.SkyCoord, - nside: int, + healpix: hp.HEALPix, data: DataObject, ) -> None: """Property test for get_binned_emission_pix.""" @@ -67,28 +67,28 @@ def test_get_binned_skycoord( coordinates, freq=frequency, obs_pos=observer, - nside=nside, + nside=healpix.nside, solar_cut=data.draw(quantities(20, 50, u.deg)) if cut_solar else None, ) - assert emission_binned.shape == (hp.nside2npix(nside),) + assert emission_binned.shape == (healpix.npix,) -@given(zodipy_models(), times(), nsides(), data(), coords_in()) +@given(zodipy_models(), times(), healpixes(), data(), coords_in()) @settings(deadline=None) def test_get_emission_pix( model: Zodipy, time: Time, - nside: int, + healpix: hp.HEALPix, data: DataObject, coord_in: Literal["E", "G", "C"], ) -> None: """Property test for get_emission_pix.""" observer = data.draw(obs_positions(model, time)) frequency = data.draw(freqs(model)) - pix = data.draw(pixels(nside)) + pix = data.draw(pixels(healpix.nside)) emission = model.get_emission_pix( pix, - nside=nside, + nside=healpix.nside, freq=frequency, obs_time=time, obs_pos=observer, @@ -97,30 +97,30 @@ def test_get_emission_pix( assert np.size(emission) == np.size(pix) -@given(zodipy_models(), times(), nsides(), data(), coords_in()) +@given(zodipy_models(), times(), healpixes(), data(), coords_in()) @settings(deadline=None) def test_get_binned_emission_pix( model: Zodipy, time: Time, - nside: int, + healpix: hp.HEALPix, data: DataObject, coord_in: Literal["E", "G", "C"], ) -> None: """Property test for get_binned_emission_pix.""" observer = data.draw(obs_positions(model, time)) frequency = data.draw(freqs(model)) - pix = data.draw(pixels(nside)) + pix = data.draw(pixels(healpix.nside)) cut_solar = data.draw(booleans()) emission_binned = model.get_binned_emission_pix( pix, - nside=nside, + 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, coord_in=coord_in, ) - assert emission_binned.shape == (hp.nside2npix(nside),) + assert emission_binned.shape == (healpix.npix,) @given(zodipy_models(), times(), angles(), data(), coords_in()) @@ -149,12 +149,12 @@ def test_get_emission_ang( assert emission.size == theta.size == phi.size -@given(zodipy_models(), times(), nsides(), angles(), data(), coords_in()) +@given(zodipy_models(), times(), healpixes(), angles(), data(), coords_in()) @settings(deadline=None) def test_get_binned_emission_ang( model: Zodipy, time: Time, - nside: int, + healpix: hp.HEALPix, angles: tuple[u.Quantity[u.deg], u.Quantity[u.deg]], data: DataObject, coord_in: Literal["E", "G", "C"], @@ -169,29 +169,29 @@ def test_get_binned_emission_ang( emission_binned = model.get_binned_emission_ang( theta, phi, - nside=nside, + nside=healpix.nside, freq=frequency, obs_time=time, obs_pos=observer, coord_in=coord_in, solar_cut=data.draw(quantities(20, 50, u.deg)) if cut_solar else None, ) - assert emission_binned.shape == (hp.nside2npix(nside),) + assert emission_binned.shape == (healpix.npix,) -@given(zodipy_models(extrapolate=False), times(), angles(lonlat=True), nsides(), data()) +@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]], - nside: int, + healpix: hp.HEALPix, data: DataObject, ) -> None: """Property test checking for unsupported spectral range.""" theta, phi = angles observer = data.draw(obs_positions(model, time)) - pix = data.draw(pixels(nside)) + 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]): @@ -199,7 +199,7 @@ def test_invalid_freq( model.get_emission_pix( pix, freq=freq, - nside=nside, + nside=healpix.nside, obs_time=time, obs_pos=observer, ) @@ -217,7 +217,7 @@ def test_invalid_freq( with pytest.raises(ValueError): model.get_binned_emission_pix( pix, - nside=nside, + nside=healpix.nside, freq=freq, obs_time=time, obs_pos=observer, @@ -227,7 +227,7 @@ def test_invalid_freq( model.get_binned_emission_ang( theta, phi, - nside=nside, + nside=healpix.nside, freq=freq, obs_time=time, obs_pos=observer, @@ -269,21 +269,21 @@ def test_multiprocessing() -> None: observer = "earth" time = Time("2020-01-01") frequency = 78 * u.micron - nside = 32 - pix = np.random.randint(0, hp.nside2npix(nside), size=1000) + healpix = hp.HEALPix(32) + pix = np.random.randint(0, healpix.npix, size=1000) theta = np.random.uniform(0, 180, size=1000) * u.deg phi = np.random.uniform(0, 360, size=1000) * u.deg emission_pix = model.get_emission_pix( pix, - nside=nside, + nside=healpix.nside, freq=frequency, obs_time=time, obs_pos=observer, ) emission_pix_parallel = model_parallel.get_emission_pix( pix, - nside=nside, + nside=healpix.nside, freq=frequency, obs_time=time, obs_pos=observer, @@ -292,14 +292,14 @@ def test_multiprocessing() -> None: emission_binned_pix = model.get_binned_emission_pix( pix, - nside=nside, + nside=healpix.nside, freq=frequency, obs_time=time, obs_pos=observer, ) emission_binned_pix_parallel = model_parallel.get_binned_emission_pix( pix, - nside=nside, + nside=healpix.nside, freq=frequency, obs_time=time, obs_pos=observer, @@ -326,7 +326,7 @@ def test_multiprocessing() -> None: theta, phi, freq=frequency, - nside=nside, + nside=healpix.nside, obs_time=time, obs_pos=observer, ) @@ -335,7 +335,7 @@ def test_multiprocessing() -> None: theta, phi, freq=frequency, - nside=nside, + nside=healpix.nside, obs_time=time, obs_pos=observer, ) @@ -351,19 +351,19 @@ def test_inner_radial_cutoff_multiprocessing() -> None: observer = "earth" time = Time("2020-01-01") frequency = 78 * u.micron - nside = 32 - pix = np.random.randint(0, hp.nside2npix(nside), size=1000) + healpix = hp.HEALPix(32) + pix = np.random.randint(0, healpix.npix, size=1000) emission_pix = model.get_emission_pix( pix, - nside=nside, + nside=healpix.nside, freq=frequency, obs_time=time, obs_pos=observer, ) emission_pix_parallel = model_parallel.get_emission_pix( pix, - nside=nside, + nside=healpix.nside, freq=frequency, obs_time=time, obs_pos=observer, @@ -374,7 +374,7 @@ def test_inner_radial_cutoff_multiprocessing() -> None: @given( zodipy_models(extrapolate=True), times(), - nsides(), + healpixes(), angles(), random_freqs(bandpass=True), data(), @@ -383,7 +383,7 @@ def test_inner_radial_cutoff_multiprocessing() -> None: def test_bandpass_integration( model: Zodipy, time: Time, - nside: int, + healpix: hp.HEALPix, angles: tuple[u.Quantity[u.deg], u.Quantity[u.deg]], freqs: u.Quantity[u.Hz] | u.Quantity[u.micron], data: DataObject, @@ -397,17 +397,17 @@ def test_bandpass_integration( phi, freq=freqs, weights=bp_weights, - nside=nside, + nside=healpix.nside, obs_time=time, obs_pos=observer, ) - assert emission_binned.shape == (hp.nside2npix(nside),) + assert emission_binned.shape == (healpix.npix,) @given( zodipy_models(extrapolate=True), times(), - nsides(), + healpixes(), angles(), random_freqs(bandpass=True), data(), @@ -416,7 +416,7 @@ def test_bandpass_integration( def test_weights( model: Zodipy, time: Time, - nside: int, + healpix: hp.HEALPix, angles: tuple[u.Quantity[u.deg], u.Quantity[u.deg]], freqs: u.Quantity[u.Hz] | u.Quantity[u.micron], data: DataObject, @@ -431,7 +431,7 @@ def test_weights( phi, weights=bp_weights, freq=freqs, - nside=nside, + nside=healpix.nside, obs_time=time, obs_pos=observer, ) @@ -441,8 +441,8 @@ def test_custom_weights() -> None: """Tests bandpass integration with custom weights.""" model = Zodipy() time = Time("2020-01-01") - nside = 32 - pix = np.arange(hp.nside2npix(nside)) + healpix = hp.HEALPix(64) + pix = np.arange(healpix.npix) central_freq = 25 sigma_freq = 3 freqs = np.linspace(central_freq - sigma_freq, central_freq + sigma_freq, 100) * u.micron @@ -453,7 +453,7 @@ def test_custom_weights() -> None: pix, freq=freqs, weights=weights, - nside=nside, + nside=healpix.nside, obs_time=time, obs_pos="earth", ) @@ -463,13 +463,13 @@ def test_custom_obs_pos() -> None: """Tests a user specified observer position.""" model = Zodipy() time = Time("2020-01-01") - nside = 64 - pix = np.arange(hp.nside2npix(nside)) + healpix = hp.HEALPix(64) + pix = np.arange(healpix.npix) model.get_emission_pix( pix, freq=234 * u.micron, - nside=nside, + nside=healpix.nside, obs_time=time, obs_pos=[0.1, 0.2, 1] * u.AU, ) @@ -477,7 +477,7 @@ def test_custom_obs_pos() -> None: model.get_emission_pix( pix, freq=234 * u.micron, - nside=nside, + nside=healpix.nside, obs_time=time, obs_pos=[2, 0.1, 4] * u.AU, ) @@ -486,7 +486,7 @@ def test_custom_obs_pos() -> None: model.get_emission_pix( pix, freq=234 * u.micron, - nside=nside, + nside=healpix.nside, obs_time=time, obs_pos=[2, 0.1, 4], ) @@ -495,7 +495,91 @@ def test_custom_obs_pos() -> None: model.get_emission_pix( pix, freq=234 * u.micron, - nside=nside, + nside=healpix.nside, obs_time=time, obs_pos=[2, 0.1, 4] * u.s, ) + + +def test_pixel_ordering() -> None: + """Tests that the pixel related functions work for both healpix orderings.""" + model = Zodipy() + time = Time("2020-01-01") + + healpix_ring = hp.HEALPix(32) + healpix_nested = hp.HEALPix(32) + + pix_ring = np.arange(healpix_ring.npix) + lon, lat = healpix_ring.healpix_to_lonlat(pix_ring) + lon = u.Quantity(lon) + lat = u.Quantity(lat) + pix_nested = healpix_nested.lonlat_to_healpix(lon, lat) + + 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", + ) + + 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", + ) + + model.get_binned_emission_ang( + lon, + lat, + lonlat=True, + freq=234 * u.micron, + nside=healpix_ring.nside, + obs_time=time, + order="ring", + ) + model.get_binned_emission_ang( + lon, + lat, + lonlat=True, + freq=234 * u.micron, + nside=healpix_nested.nside, + obs_time=time, + order="nested", + ) + + sky_coord = coords.SkyCoord( + lon, + lat, + frame=coords.BarycentricMeanEcliptic, + obstime=time, + ) + + 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/zodipy/_line_of_sight.py b/zodipy/_line_of_sight.py index 19efb27..92a46a0 100644 --- a/zodipy/_line_of_sight.py +++ b/zodipy/_line_of_sight.py @@ -11,9 +11,7 @@ if TYPE_CHECKING: import numpy.typing as npt - from zodipy._types import Float - -DIRBE_CUTOFFS: dict[ComponentLabel, tuple[Float, Float]] = { +DIRBE_CUTOFFS: dict[ComponentLabel, tuple[float | np.floating, float | np.floating]] = { ComponentLabel.CLOUD: (R_0, R_JUPITER), ComponentLabel.BAND1: (R_0, R_JUPITER), ComponentLabel.BAND2: (R_0, R_JUPITER), @@ -22,7 +20,7 @@ ComponentLabel.FEATURE: (R_EARTH - 0.2, R_EARTH + 0.2), } -RRM_CUTOFFS: dict[ComponentLabel, tuple[Float, Float]] = { +RRM_CUTOFFS: dict[ComponentLabel, tuple[float | np.floating, float | np.floating]] = { ComponentLabel.FAN: (R_0, RRM[ComponentLabel.FAN].R_outer), # type: ignore ComponentLabel.INNER_NARROW_BAND: ( RRM[ComponentLabel.INNER_NARROW_BAND].R_inner, # type: ignore @@ -54,12 +52,7 @@ def get_sphere_intersection( unit_vectors: npt.NDArray[np.float64], cutoff: float | np.float64, ) -> npt.NDArray[np.float64]: - """Get RMAX per pixel. - - Given the observer position, return distance from the observer to the - intersection between the line-of-sights and a heliocentric sphere with radius `cutoff`. - - """ + """Returns the distance from the observer to a heliocentric sphere with radius `cutoff`.""" x, y, z = obs_pos.flatten() r_obs = np.sqrt(x**2 + y**2 + z**2) @@ -87,7 +80,7 @@ def get_line_of_sight_range( dict[ComponentLabel, npt.NDArray[np.float64]], dict[ComponentLabel, npt.NDArray[np.float64]], ]: - """Get the start and stop distances for each component.""" + """Return the component-wise start- and end-point of each line of sight.""" start = { comp: get_sphere_intersection(obs_pos, unit_vectors, COMPONENT_CUTOFFS[comp][0]) for comp in components diff --git a/zodipy/_validators.py b/zodipy/_validators.py index 9e708b9..0461226 100644 --- a/zodipy/_validators.py +++ b/zodipy/_validators.py @@ -1,33 +1,32 @@ from __future__ import annotations -import astropy.units as u import numpy as np import numpy.typing as npt +from astropy import units from ._ipd_model import InterplanetaryDustModel -from ._types import FrequencyOrWavelength, SkyAngles def get_validated_freq( - freq: FrequencyOrWavelength, model: InterplanetaryDustModel, extrapolate: bool -) -> FrequencyOrWavelength: + freq: units.Quantity, model: InterplanetaryDustModel, extrapolate: bool +) -> units.Quantity: """Validate user inputted frequency.""" - if not isinstance(freq, u.Quantity): + if not isinstance(freq, units.Quantity): msg = "Frequency must be an astropy Quantity." raise TypeError(msg) - if freq.unit.is_equivalent(u.Hz): - freq = freq.to(u.Hz) - elif freq.unit.is_equivalent(u.micron): - freq = freq.to(u.micron) + if freq.unit.is_equivalent(units.Hz): + freq = freq.to(units.Hz) + elif freq.unit.is_equivalent(units.micron): + freq = freq.to(units.micron) else: msg = "Frequency must be in units compatible with Hz or micron." - raise u.UnitsError(msg) + raise units.UnitsError(msg) if extrapolate: return freq - freq_in_spectrum_units = freq.to(model.spectrum.unit, equivalencies=u.spectral()) + freq_in_spectrum_units = freq.to(model.spectrum.unit, equivalencies=units.spectral()) lower_freq_range = model.spectrum.min() upper_freq_range = model.spectrum.max() @@ -48,7 +47,7 @@ def get_validated_freq( def get_validated_and_normalized_weights( weights: npt.ArrayLike | None, - freq: FrequencyOrWavelength, + freq: units.Quantity, ) -> npt.NDArray[np.float64]: """Validate user inputted weights.""" if weights is None and freq.size > 1: @@ -74,18 +73,22 @@ def get_validated_and_normalized_weights( def get_validated_ang( - theta: SkyAngles, phi: SkyAngles, lonlat: bool -) -> tuple[SkyAngles, SkyAngles]: + theta: units.Quantity, phi: units.Quantity, lonlat: bool +) -> tuple[units.Quantity, units.Quantity]: """Validate user inputted sky angles and make sure it adheres to the healpy convention.""" - theta = theta.to(u.deg) if lonlat else theta.to(u.rad) - phi = phi.to(u.deg) if lonlat else phi.to(u.rad) + try: + theta = theta.to(units.deg) if lonlat else theta.to(units.rad) + phi = phi.to(units.deg) if lonlat else phi.to(units.rad) + except AttributeError: + msg = "Sky angles `theta` and `phi` must be astropy Quantities." + raise TypeError(msg) from AttributeError if theta.isscalar: - theta = u.Quantity([theta]) + theta = np.expand_dims(theta, axis=0) if phi.isscalar: - phi = u.Quantity([phi]) + phi = np.expand_dims(phi, axis=0) if not lonlat: - theta, phi = phi, (np.pi / 2) * u.rad - theta + theta, phi = phi, (np.pi / 2) * units.rad - theta return theta, phi diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index 3474fa4..b36c691 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -150,18 +150,18 @@ def get_emission_skycoord( return_inverse=True, axis=1, ) + obs_time = coord.obstime coord = coords.SkyCoord( unique_lon * units.deg, unique_lat * units.deg, frame=coord.frame, - obstime=coord.obstime, ) return self._compute_emission( freq=freq, weights=weights, obs_pos=obs_pos, - obs_time=coord.obstime, + obs_time=obs_time, coordinates=coord, indicies=indicies, return_comps=return_comps, @@ -238,6 +238,7 @@ def get_emission_pix( coord_in: Literal["E", "G", "C"] = "E", weights: npt.ArrayLike | None = None, return_comps: bool = False, + order: Literal["ring", "nested"] = "ring", ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated zodiacal emission given pixel numbers. @@ -259,13 +260,14 @@ def get_emission_pix( weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. + order: Order of the HEALPix grid. Returns: emission: Simulated zodiacal emission in units of 'MJy/sr'. """ frame = get_frame_from_string(coord_in) - healpix = hp.HEALPix(nside=nside, order="ring", frame=frame) + 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) @@ -288,6 +290,7 @@ def get_binned_emission_skycoord( weights: npt.ArrayLike | None = None, obs_pos: units.Quantity | str = "earth", return_comps: bool = False, + order: Literal["ring", "nested"] = "ring", solar_cut: units.Quantity | None = None, ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated binned zodiacal light for all observations in a `SkyCoord` object. @@ -309,6 +312,7 @@ def get_binned_emission_skycoord( an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. return_comps: If True, the emission is returned component-wise. Defaults to False. + order: Order of the HEALPix grid. solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular distance between the sun smaller than `solar_cut` are masked. Defaults to `None`. @@ -321,18 +325,18 @@ def get_binned_emission_skycoord( return_inverse=True, axis=1, ) + obs_time = coord.obstime coord = coords.SkyCoord( unique_lon * units.deg, unique_lat * units.deg, frame=coord.frame, - obstime=coord.obstime, ) - healpix = hp.HEALPix(nside, order="ring", frame=coord.frame) + healpix = hp.HEALPix(nside, order=order, frame=coord.frame) return self._compute_emission( freq=freq, weights=weights, obs_pos=obs_pos, - obs_time=coord.obstime, + obs_time=obs_time, coordinates=coord, indicies=indicies, healpix=healpix, @@ -353,6 +357,7 @@ def get_binned_emission_ang( coord_in: Literal["E", "G", "C"] = "E", weights: npt.ArrayLike | None = None, return_comps: bool = False, + order: Literal["ring", "nested"] = "ring", solar_cut: units.Quantity | None = None, ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated binned zodiacal emission given angles on the sky. @@ -383,6 +388,7 @@ def get_binned_emission_ang( lonlat: If True, input angles (`theta`, `phi`) are assumed to be longitude and latitude, otherwise, they are co-latitude and longitude. return_comps: If True, the emission is returned component-wise. Defaults to False. + order: Order of the HEALPix grid. solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular distance between the sun smaller than `solar_cut` are masked. Defaults to `None`. @@ -390,9 +396,9 @@ def get_binned_emission_ang( emission: Simulated zodiacal emission in units of 'MJy/sr'. """ - theta, phi = get_validated_ang(theta=theta, phi=phi, lonlat=lonlat) + theta, phi = get_validated_ang(theta, phi, lonlat=lonlat) frame = get_frame_from_string(coord_in) - healpix = hp.HEALPix(nside, order="ring", frame=frame) + 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( theta, @@ -423,6 +429,7 @@ def get_binned_emission_pix( coord_in: Literal["E", "G", "C"] = "E", weights: npt.ArrayLike | None = None, return_comps: bool = False, + order: Literal["ring", "nested"] = "ring", solar_cut: units.Quantity | None = None, ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated binned zodiacal Emission given pixel numbers. @@ -446,6 +453,7 @@ def get_binned_emission_pix( weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. + order: Order of the HEALPix grid. solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular distance between the sun smaller than `solar_cut` are masked. Defaults to `None`. @@ -454,7 +462,7 @@ def get_binned_emission_pix( """ frame = get_frame_from_string(coord_in) - healpix = hp.HEALPix(nside=nside, order="ring", frame=frame) + 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) From 2eea4be4428870f96159a4638b50a134dece778b Mon Sep 17 00:00:00 2001 From: Metin San Date: Sun, 28 Apr 2024 10:00:42 +0200 Subject: [PATCH 11/17] Update workflow so that joss paper only compiles if its modified --- .github/workflows/draft-pdf.yml | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/.github/workflows/draft-pdf.yml b/.github/workflows/draft-pdf.yml index e4c5c62..d3f5980 100644 --- a/.github/workflows/draft-pdf.yml +++ b/.github/workflows/draft-pdf.yml @@ -1,4 +1,9 @@ -on: [push, workflow_dispatch] +on: + push: + paths: ["paper/**"] + pull_request: + paths: ["paper/**"] + workflow_dispatch: jobs: paper: @@ -20,4 +25,4 @@ jobs: # This is the output path where Pandoc will write the compiled # PDF. Note, this should be the same directory as the input # paper.md - path: paper/paper.pdf \ No newline at end of file + path: paper/paper.pdf From b85fa59c7bbd9710649115a259f8f94f5b8b5219 Mon Sep 17 00:00:00 2001 From: Metin San Date: Sun, 28 Apr 2024 10:34:45 +0200 Subject: [PATCH 12/17] Add tests to get back to 99% coverage --- .pre-commit-config.yaml | 1 + tests/test_validators.py | 25 +++++++++++++++++++++++++ zodipy/_bandpass.py | 13 ++++++------- zodipy/_ipd_model.py | 5 +++-- zodipy/_types.py | 13 ------------- 5 files changed, 35 insertions(+), 22 deletions(-) delete mode 100644 zodipy/_types.py diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index dfed5f2..a2549ed 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -16,6 +16,7 @@ repos: rev: v1.10.0 hooks: - id: mypy + exclude: tests/ additional_dependencies: [numpy>=1.21.0] hooks: - id: poetry-export diff --git a/tests/test_validators.py b/tests/test_validators.py index f4e5832..d429233 100644 --- a/tests/test_validators.py +++ b/tests/test_validators.py @@ -136,3 +136,28 @@ def test_interp_kind() -> None: model.get_emission_pix( pixels=[1, 4, 5], freq=27 * units.micron, 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() + 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", + ) + + +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() + with pytest.raises(TypeError): + model.get_emission_ang( + 32, + 12, + freq=27 * units.micron, + obs_time=OBS_TIME, + ) diff --git a/zodipy/_bandpass.py b/zodipy/_bandpass.py index e5ecc5e..fba3428 100644 --- a/zodipy/_bandpass.py +++ b/zodipy/_bandpass.py @@ -3,8 +3,8 @@ from dataclasses import dataclass from typing import TYPE_CHECKING -import astropy.units as u import numpy as np +from astropy import units from zodipy._constants import ( MAX_INTERPOLATION_GRID_TEMPERATURE, @@ -18,12 +18,11 @@ import numpy.typing as npt from zodipy._ipd_model import InterplanetaryDustModel - from zodipy._types import FrequencyOrWavelength @dataclass class Bandpass: - frequencies: FrequencyOrWavelength + frequencies: units.Quantity weights: npt.NDArray[np.float64] def integrate(self, quantity: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: @@ -33,8 +32,8 @@ def integrate(self, quantity: npt.NDArray[np.float64]) -> npt.NDArray[np.float64 def switch_convention(self) -> None: """Switch the bandpass from frequency to wavelength or the other way around.""" self.frequencies = self.frequencies.to( - u.micron if self.frequencies.unit.is_equivalent(u.Hz) else u.Hz, - equivalencies=u.spectral(), + units.micron if self.frequencies.unit.is_equivalent(units.Hz) else units.Hz, + equivalencies=units.spectral(), ) if self.frequencies.size > 1: self.frequencies = np.flip(self.frequencies) @@ -43,7 +42,7 @@ def switch_convention(self) -> None: def validate_and_get_bandpass( - freq: FrequencyOrWavelength, + freq: units.Quantity, weights: npt.ArrayLike | None, model: InterplanetaryDustModel, extrapolate: bool, @@ -62,7 +61,7 @@ def get_bandpass_interpolation_table( max_temp: float = MAX_INTERPOLATION_GRID_TEMPERATURE, ) -> npt.NDArray[np.float64]: """Pre-compute the bandpass integrated blackbody emission for a grid of temperatures.""" - if not bandpass.frequencies.unit.is_equivalent(u.Hz): + if not bandpass.frequencies.unit.is_equivalent(units.Hz): bandpass.switch_convention() bandpass_integrals = np.zeros(n_points) diff --git a/zodipy/_ipd_model.py b/zodipy/_ipd_model.py index 4589106..f14fd47 100644 --- a/zodipy/_ipd_model.py +++ b/zodipy/_ipd_model.py @@ -5,8 +5,9 @@ from typing import TYPE_CHECKING, Mapping, Sequence if TYPE_CHECKING: + from astropy import units + from zodipy._ipd_comps import Component, ComponentLabel - from zodipy._types import FrequencyOrWavelength @dataclass @@ -20,7 +21,7 @@ class InterplanetaryDustModel(ABC): """ comps: Mapping[ComponentLabel, Component] - spectrum: FrequencyOrWavelength + spectrum: units.Quantity def to_dict(self) -> dict: """Return a dictionary representation of the model.""" diff --git a/zodipy/_types.py b/zodipy/_types.py deleted file mode 100644 index c887a5a..0000000 --- a/zodipy/_types.py +++ /dev/null @@ -1,13 +0,0 @@ -from __future__ import annotations - -from typing import Sequence, Union - -import astropy.units as u -import numpy as np -import numpy.typing as npt - -Pixels = Union[int, Sequence[int], npt.NDArray[np.integer]] -SkyAngles = Union[u.Quantity[u.deg], u.Quantity[u.rad]] -FrequencyOrWavelength = Union[u.Quantity[u.Hz], u.Quantity[u.m]] -ParameterDict = dict -Float = Union[float, np.floating] From 805f0ab905281313a4fcffa4898bf965fa63cae9 Mon Sep 17 00:00:00 2001 From: Metin San Date: Sun, 28 Apr 2024 10:51:08 +0200 Subject: [PATCH 13/17] Update example --- docs/examples/get_emission_ang.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/docs/examples/get_emission_ang.py b/docs/examples/get_emission_ang.py index 608947b..952d05d 100644 --- a/docs/examples/get_emission_ang.py +++ b/docs/examples/get_emission_ang.py @@ -1,28 +1,30 @@ -from multiprocessing import cpu_count - import astropy.units as u import matplotlib.pyplot as plt import numpy as np +from astropy.coordinates import BarycentricMeanEcliptic, SkyCoord from astropy.time import Time from zodipy import Zodipy -model = Zodipy("dirbe", n_proc=cpu_count()) +model = Zodipy() + +# Longitude and Latitude values corresponding to a scan through the eclitpic plane +lats = np.linspace(-90, 90, 100) * u.deg +lons = np.zeros_like(lats) -latitudes = np.linspace(-90, 90, 10000) * u.deg -longitudes = np.zeros_like(latitudes) +obs_time = Time("2022-06-14") -emission = model.get_emission_ang( - theta=longitudes, - phi=latitudes, - freq=30 * u.micron, - lonlat=True, - obs_time=Time("2022-06-14"), - obs_pos="earth", +# The SkyCoord object needs to include the coordinate frame and time of observation +coords = SkyCoord( + lons, + lats, + frame=BarycentricMeanEcliptic, + obstime=obs_time, ) +emission = model.get_emission_skycoord(coords, freq=30 * u.micron) -plt.plot(latitudes, emission) +plt.plot(lats, emission) plt.xlabel("Latitude [deg]") plt.ylabel("Emission [MJy/sr]") plt.show() From d0532821b9c9b6c7e77d157a835e9ffcc3be0f42 Mon Sep 17 00:00:00 2001 From: Metin San Date: Mon, 29 Apr 2024 07:31:25 +0200 Subject: [PATCH 14/17] Move bandpass calculations to Zodipy init fixing #26 and fix bug with ephem --- zodipy/_coords.py | 17 +++++-- zodipy/zodipy.py | 114 +++++++++++++--------------------------------- 2 files changed, 43 insertions(+), 88 deletions(-) diff --git a/zodipy/_coords.py b/zodipy/_coords.py index 8a6a5b4..1807548 100644 --- a/zodipy/_coords.py +++ b/zodipy/_coords.py @@ -26,23 +26,30 @@ def get_sun_earth_moon_barycenter_skycoord(earth_skycoord: coords.SkyCoord) -> c ) -def get_earth_skycoord(obs_time: time.Time) -> coords.SkyCoord: +def get_earth_skycoord(obs_time: time.Time, ephemeris: str) -> coords.SkyCoord: """Return the sky coordinates of the Earth in the heliocentric frame.""" - return coords.get_body("earth", obs_time).transform_to(coords.HeliocentricMeanEcliptic) + return coords.get_body("earth", obs_time, ephemeris=ephemeris).transform_to( + coords.HeliocentricMeanEcliptic + ) def get_obs_skycoord( - obs_pos: units.Quantity | str, obs_time: time.Time, earth_skycoord: coords.SkyCoord + obs_pos: units.Quantity | str, + obs_time: time.Time, + earth_skycoord: coords.SkyCoord, + ephemeris: str, ) -> coords.SkyCoord: """Return the sky coordinates of the observer in the heliocentric frame.""" if isinstance(obs_pos, str): if obs_pos.lower() == "semb-l2": return get_sun_earth_moon_barycenter_skycoord(earth_skycoord) - return coords.get_body(obs_pos, obs_time).transform_to(coords.HeliocentricMeanEcliptic) + return coords.get_body(obs_pos, obs_time, ephemeris=ephemeris).transform_to( + coords.HeliocentricMeanEcliptic + ) try: return coords.SkyCoord( - *obs_pos.to(units.AU), + *obs_pos, frame=coords.HeliocentricMeanEcliptic, representation_type="cartesian", ) diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index b36c691..07a7020 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -37,6 +37,11 @@ class Zodipy: 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 @@ -60,28 +65,40 @@ class Zodipy: def __init__( self, + freq: units.Quantity, + weights: npt.ArrayLike | None = None, model: str = "dirbe", gauss_quad_degree: int = 50, extrapolate: bool = False, interp_kind: str = "linear", - ephemeris: str = "de432s", + ephemeris: str = "builtin", n_proc: int = 1, ) -> None: self.model = model self.gauss_quad_degree = gauss_quad_degree self.extrapolate = extrapolate - self.interp_kind = interp_kind self.ephemeris = ephemeris self.n_proc = n_proc self._interpolator = functools.partial( interpolate.interp1d, - kind=self.interp_kind, + kind=interp_kind, fill_value="extrapolate" if self.extrapolate else np.nan, ) self._ipd_model = model_registry.get_model(model) self._gauss_points_and_weights = np.polynomial.legendre.leggauss(gauss_quad_degree) + bandpass = validate_and_get_bandpass( + freq=freq, + weights=weights, + model=self._ipd_model, + extrapolate=self.extrapolate, + ) + 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 + ) + @property def supported_observers(self) -> list[str]: """Return a list of available observers given an ephemeris.""" @@ -117,8 +134,6 @@ def get_emission_skycoord( self, coord: coords.SkyCoord, *, - freq: units.Quantity, - weights: npt.ArrayLike | None = None, obs_pos: units.Quantity | str = "earth", return_comps: bool = False, ) -> units.Quantity[units.MJy / units.sr]: @@ -131,11 +146,6 @@ def get_emission_skycoord( observational time for which the zodiacal light is assumed to be stationary. Additionally, the frame must be convertible to `astropy.coordinates.BarycentricMeanEcliptic`. - 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). obs_pos: The heliocentric ecliptic position of the observer, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. @@ -158,8 +168,6 @@ def get_emission_skycoord( ) return self._compute_emission( - freq=freq, - weights=weights, obs_pos=obs_pos, obs_time=obs_time, coordinates=coord, @@ -173,11 +181,9 @@ def get_emission_ang( phi: units.Quantity, *, lonlat: bool = False, - freq: units.Quantity, obs_time: time.Time, obs_pos: units.Quantity | str = "earth", coord_in: Literal["E", "G", "C"] = "E", - weights: npt.ArrayLike | None = None, return_comps: bool = False, ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated zodiacal emission given angles on the sky. @@ -194,17 +200,12 @@ def get_emission_ang( with degrees. lonlat: If True, input angles (`theta`, `phi`) are assumed to be longitude and latitude, otherwise, they are co-latitude and longitude. - 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. obs_time: Time of observation. This should be a single observational time. obs_pos: The heliocentric ecliptic position of the observer, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic coordinates) by default. - weights: Bandpass weights corresponding the the frequencies in `freq`. The weights - are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. Returns: @@ -218,8 +219,6 @@ def get_emission_ang( coordinates = coords.SkyCoord(theta, phi, frame=frame) return self._compute_emission( - freq=freq, - weights=weights, obs_time=obs_time, obs_pos=obs_pos, coordinates=coordinates, @@ -232,11 +231,9 @@ def get_emission_pix( pixels: npt.ArrayLike, *, nside: int, - freq: units.Quantity, obs_time: time.Time, obs_pos: units.Quantity | str = "earth", coord_in: Literal["E", "G", "C"] = "E", - weights: npt.ArrayLike | None = None, return_comps: bool = False, order: Literal["ring", "nested"] = "ring", ) -> units.Quantity[units.MJy / units.sr]: @@ -248,17 +245,12 @@ def get_emission_pix( Args: pixels: HEALPix pixel indicies representing points on the celestial sphere. nside: HEALPix resolution parameter of the pixels and the binned map. - 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. obs_time: Time of observation. This should be a single observational time. obs_pos: The heliocentric ecliptic position of the observer, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic coordinates) by default. - weights: Bandpass weights corresponding the the frequencies in `freq`. The weights - are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. order: Order of the HEALPix grid. @@ -272,8 +264,6 @@ def get_emission_pix( coordinates = healpix.healpix_to_skycoord(unique_pixels) return self._compute_emission( - freq=freq, - weights=weights, obs_time=obs_time, obs_pos=obs_pos, coordinates=coordinates, @@ -286,8 +276,6 @@ def get_binned_emission_skycoord( coord: coords.SkyCoord, *, nside: int, - freq: units.Quantity, - weights: npt.ArrayLike | None = None, obs_pos: units.Quantity | str = "earth", return_comps: bool = False, order: Literal["ring", "nested"] = "ring", @@ -303,11 +291,6 @@ def get_binned_emission_skycoord( Additionally, the frame must be convertible to `astropy.coordinates.BarycentricMeanEcliptic`. nside: HEALPix resolution parameter of the pixels and the binned map. - 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). obs_pos: The heliocentric ecliptic position of the observer, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. @@ -333,8 +316,6 @@ def get_binned_emission_skycoord( ) healpix = hp.HEALPix(nside, order=order, frame=coord.frame) return self._compute_emission( - freq=freq, - weights=weights, obs_pos=obs_pos, obs_time=obs_time, coordinates=coord, @@ -351,11 +332,9 @@ def get_binned_emission_ang( *, lonlat: bool = False, nside: int, - freq: units.Quantity, obs_time: time.Time, obs_pos: units.Quantity | str = "earth", coord_in: Literal["E", "G", "C"] = "E", - weights: npt.ArrayLike | None = None, return_comps: bool = False, order: Literal["ring", "nested"] = "ring", solar_cut: units.Quantity | None = None, @@ -374,17 +353,12 @@ def get_binned_emission_ang( celestial sphere. Must be in the range [0, 2π] rad. Units must be either radians or degrees. nside: HEALPix resolution parameter of the pixels and the binned map. - 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. obs_time: Time of observation. This should be a single observational time. obs_pos: The heliocentric ecliptic position of the observer, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic coordinates) by default. - weights: Bandpass weights corresponding the the frequencies in `freq`. The weights - are assumed to be given in spectral radiance units (Jy/sr). lonlat: If True, input angles (`theta`, `phi`) are assumed to be longitude and latitude, otherwise, they are co-latitude and longitude. return_comps: If True, the emission is returned component-wise. Defaults to False. @@ -407,8 +381,6 @@ def get_binned_emission_ang( ) return self._compute_emission( - freq=freq, - weights=weights, obs_time=obs_time, obs_pos=obs_pos, coordinates=coordinates, @@ -423,11 +395,9 @@ def get_binned_emission_pix( pixels: npt.ArrayLike, *, nside: int, - freq: units.Quantity, obs_time: time.Time, obs_pos: units.Quantity | str = "earth", coord_in: Literal["E", "G", "C"] = "E", - weights: npt.ArrayLike | None = None, return_comps: bool = False, order: Literal["ring", "nested"] = "ring", solar_cut: units.Quantity | None = None, @@ -441,17 +411,12 @@ def get_binned_emission_pix( Args: pixels: HEALPix pixel indicies representing points on the celestial sphere. nside: HEALPix resolution parameter of the pixels and the binned map. - 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. obs_time: Time of observation. This should be a single observational time. obs_pos: The heliocentric ecliptic position of the observer, or a string representing an observer in the `astropy.coordinates.solar_system_ephemeris`. This should correspond to a single position. Defaults to 'earth'. coord_in: Coordinate frame of the input pointing. Assumes 'E' (ecliptic coordinates) by default. - weights: Bandpass weights corresponding the the frequencies in `freq`. The weights - are assumed to be given in spectral radiance units (Jy/sr). return_comps: If True, the emission is returned component-wise. Defaults to False. order: Order of the HEALPix grid. solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular @@ -467,8 +432,6 @@ def get_binned_emission_pix( coordinates = healpix.healpix_to_skycoord(unique_pixels) return self._compute_emission( - freq=freq, - weights=weights, obs_time=obs_time, obs_pos=obs_pos, coordinates=coordinates, @@ -480,8 +443,6 @@ def get_binned_emission_pix( def _compute_emission( self, - freq: units.Quantity, - weights: npt.ArrayLike | None, obs_time: time.Time, obs_pos: units.Quantity | str, coordinates: coords.SkyCoord, @@ -491,21 +452,8 @@ def _compute_emission( solar_cut: units.Quantity | None = None, ) -> units.Quantity[units.MJy / units.sr]: """Compute the zodiacal light for a given configuration.""" - bandpass = validate_and_get_bandpass( - freq=freq, - weights=weights, - model=self._ipd_model, - extrapolate=self.extrapolate, - ) - - # Get model parameters, some of which have been interpolated to the given - # frequency or bandpass. - source_parameters = SOURCE_PARAMS_MAPPING[type(self._ipd_model)]( - bandpass, self._ipd_model, self._interpolator - ) - - earth_skycoord = get_earth_skycoord(obs_time) - obs_skycoord = get_obs_skycoord(obs_pos, obs_time, earth_skycoord) + 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) @@ -516,23 +464,23 @@ def _compute_emission( start, stop = get_line_of_sight_range( components=self._ipd_model.comps.keys(), unit_vectors=unit_vectors, - obs_pos=obs_skycoord.cartesian.xyz.value, + obs_pos=obs_skycoord.cartesian.xyz.to_value(units.AU), ) density_partials = construct_density_partials_comps( comps=self._ipd_model.comps, dynamic_params={ - "X_earth": earth_skycoord.cartesian.xyz.value[:, np.newaxis, np.newaxis] + "X_earth": earth_skycoord.cartesian.xyz.to_value(units.AU)[ + :, np.newaxis, np.newaxis + ] }, ) - bandpass_interpolatation_table = get_bandpass_interpolation_table(bandpass) - common_integrand = functools.partial( EMISSION_MAPPING[type(self._ipd_model)], - X_obs=obs_skycoord.cartesian.xyz.value[:, np.newaxis, np.newaxis], - bp_interpolation_table=bandpass_interpolatation_table, - **source_parameters["common"], + 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"], ) # Parallelize the line-of-sight integrals if more than one processor is used and the @@ -556,7 +504,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], - **source_parameters[comp_label], + **self.source_parameters[comp_label], ) for unit_vectors, start, stop in zip( unit_vector_chunks, start_chunks, stop_chunks @@ -588,7 +536,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], - **source_parameters[comp_label], + **self.source_parameters[comp_label], ) integrated_comp_emission[idx] = ( From 5d49729ff75782c43d9480ce71ab80682d53cbd7 Mon Sep 17 00:00:00 2001 From: Metin San Date: Mon, 29 Apr 2024 09:46:48 +0200 Subject: [PATCH 15/17] Solar cut is now used to filter out pointing instead of post-evaluation masking --- zodipy/zodipy.py | 132 ++++++++++++++++++++++++++--------------------- 1 file changed, 73 insertions(+), 59 deletions(-) diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index 07a7020..07d2fce 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -46,18 +46,18 @@ class Zodipy: 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 to interpolate relevant model parameters. - Defaults to 'linear'. For more information on available interpolation methods, - please visit the [Scipy documentation]( - https://docs.scipy.org/doc/scipy/reference/generated/scipy.interpolate.interp1d.html). + 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 to compute the positions of the observer and the - Earth. Defaults to 'de432s', which requires downloading (and caching) a ~10MB - file. For more information on available ephemeridis, please visit the [Astropy - documentation](https://docs.astropy.org/en/stable/coordinates/solarsystem.html) + 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. @@ -99,37 +99,6 @@ def __init__( bandpass, self._ipd_model, self._interpolator ) - @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"] - - def get_parameters(self) -> dict: - """Return a dictionary containing the interplanetary dust model parameters.""" - return self._ipd_model.to_dict() - - def update_parameters(self, parameters: dict) -> None: - """Update the interplanetary dust model parameters. - - Args: - parameters: Dictionary of parameters to update. The keys must be the names - of the parameters as defined in the model. To get the parameters dict - of an existing model, use `Zodipy("dirbe").get_parameters()`. - - """ - _dict = parameters.copy() - _dict["comps"] = {} - for key, value in parameters.items(): - if key == "comps": - for comp_key, comp_value in value.items(): - _dict["comps"][ComponentLabel(comp_key)] = type( - self._ipd_model.comps[ComponentLabel(comp_key)] - )(**comp_value) - elif isinstance(value, dict): - _dict[key] = {ComponentLabel(k): v for k, v in value.items()} - - self._ipd_model = self._ipd_model.__class__(**_dict) - def get_emission_skycoord( self, coord: coords.SkyCoord, @@ -160,10 +129,16 @@ def get_emission_skycoord( return_inverse=True, axis=1, ) + obs_time = coord.obstime + if obs_time is None: + msg = "The `obstime` attribute of the `SkyCoord` object must be set." + raise ValueError(msg) + coord = coords.SkyCoord( - unique_lon * units.deg, - unique_lat * units.deg, + unique_lon, + unique_lat, + unit=units.deg, frame=coord.frame, ) @@ -296,8 +271,8 @@ def get_binned_emission_skycoord( correspond to a single position. Defaults to 'earth'. return_comps: If True, the emission is returned component-wise. Defaults to False. order: Order of the HEALPix grid. - solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular - distance between the sun smaller than `solar_cut` are masked. Defaults to `None`. + solar_cut: Angular distance around the Sun for which all pointing are discarded. + Defaults to `None`. Returns: emission: Simulated zodiacal light in units of 'MJy/sr'. @@ -308,10 +283,16 @@ def get_binned_emission_skycoord( return_inverse=True, axis=1, ) + obs_time = coord.obstime + if obs_time is None: + msg = "The `obstime` attribute of the `SkyCoord` object must be set." + raise ValueError(msg) + coord = coords.SkyCoord( - unique_lon * units.deg, - unique_lat * units.deg, + unique_lon, + unique_lat, + unit=units.deg, frame=coord.frame, ) healpix = hp.HEALPix(nside, order=order, frame=coord.frame) @@ -363,8 +344,8 @@ def get_binned_emission_ang( latitude, otherwise, they are co-latitude and longitude. return_comps: If True, the emission is returned component-wise. Defaults to False. order: Order of the HEALPix grid. - solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular - distance between the sun smaller than `solar_cut` are masked. Defaults to `None`. + solar_cut: Angular distance around the Sun for which all pointing are discarded. + Defaults to `None`. Returns: emission: Simulated zodiacal emission in units of 'MJy/sr'. @@ -419,8 +400,8 @@ def get_binned_emission_pix( coordinates) by default. return_comps: If True, the emission is returned component-wise. Defaults to False. order: Order of the HEALPix grid. - solar_cut: Cutoff angle from the sun. The emission for all the pointing with angular - distance between the sun smaller than `solar_cut` are masked. Defaults to `None`. + solar_cut: Angular distance around the Sun for which all pointing are discarded. + Defaults to `None`. Returns: emission: Simulated zodiacal emission in units of 'MJy/sr'. @@ -457,6 +438,18 @@ def _compute_emission( # 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: + sun_skycoord = coords.SkyCoord( + obs_skycoord.spherical.lon + 180 * units.deg, + obs_skycoord.spherical.lat, + frame=coords.BarycentricMeanEcliptic, + ) + angular_distance = coordinates.separation(sun_skycoord) + solar_mask = angular_distance < solar_cut + coordinates = coordinates[~solar_mask] + unit_vectors = coordinates.cartesian.xyz.value # Get the integration limits for each zodiacal component (which may be @@ -550,16 +543,6 @@ def _compute_emission( emission = np.zeros((len(self._ipd_model.comps), healpix.npix)) pixels = healpix.skycoord_to_healpix(coordinates) emission[:, pixels] = integrated_comp_emission - if solar_cut is not None: - sun_skycoord = coords.SkyCoord( - obs_skycoord.spherical.lon + 180 * units.deg, - obs_skycoord.spherical.lat, - frame=coords.BarycentricMeanEcliptic, - ) - angular_distance = coordinates.separation(sun_skycoord) - solar_mask = angular_distance < solar_cut - emission[:, pixels[solar_mask]] = np.nan - else: emission = np.zeros((len(self._ipd_model.comps), indicies.size)) emission = integrated_comp_emission[:, indicies] @@ -568,6 +551,37 @@ def _compute_emission( return emission if return_comps else emission.sum(axis=0) + @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"] + + def get_parameters(self) -> dict: + """Return a dictionary containing the interplanetary dust model parameters.""" + return self._ipd_model.to_dict() + + def update_parameters(self, parameters: dict) -> None: + """Update the interplanetary dust model parameters. + + Args: + parameters: Dictionary of parameters to update. The keys must be the names + of the parameters as defined in the model. To get the parameters dict + of an existing model, use `Zodipy("dirbe").get_parameters()`. + + """ + _dict = parameters.copy() + _dict["comps"] = {} + for key, value in parameters.items(): + if key == "comps": + for comp_key, comp_value in value.items(): + _dict["comps"][ComponentLabel(comp_key)] = type( + self._ipd_model.comps[ComponentLabel(comp_key)] + )(**comp_value) + elif isinstance(value, dict): + _dict[key] = {ComponentLabel(k): v for k, v in value.items()} + + self._ipd_model = self._ipd_model.__class__(**_dict) + def __repr__(self) -> str: repr_str = f"{self.__class__.__name__}(" for attribute_name, attribute in self.__dict__.items(): From 333b3f2b4056cb5a922a306adaa2118e43eb864e Mon Sep 17 00:00:00 2001 From: Metin San Date: Mon, 29 Apr 2024 17:36:43 +0200 Subject: [PATCH 16/17] Move frequency and weights to Zodipy initializer and away from methods --- tests/_strategies.py | 55 ++++++++----- tests/test_get_emission.py | 137 +++++--------------------------- tests/test_tabulated_density.py | 4 +- tests/test_validators.py | 38 ++++----- zodipy/_coords.py | 4 +- zodipy/_ipd_model.py | 4 + zodipy/zodipy.py | 121 ++++++++++++++-------------- 7 files changed, 140 insertions(+), 223 deletions(-) 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.""" From 57056e2921a1c84736885bda17b5844e1b5ff839 Mon Sep 17 00:00:00 2001 From: Metin San Date: Mon, 29 Apr 2024 17:40:47 +0200 Subject: [PATCH 17/17] Fix bug with freqs hypothesis strategy --- tests/_strategies.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/_strategies.py b/tests/_strategies.py index a12421f..3039908 100644 --- a/tests/_strategies.py +++ b/tests/_strategies.py @@ -256,7 +256,7 @@ def zodipy_models(draw: DrawFn, **static_params: dict[str, Any]) -> zodipy.Zodip w = draw(weights(frequencies)) else: frequencies = static_params.pop( - "freq", draw(freqs(min=min_freq, max=max_freq, extrapolate=extrapolate)) + "freq", draw(freqs(min_freq=min_freq, max_freq=max_freq, extrapolate=extrapolate)) ) w = None