From 771a576098a111002870f15f551f88dae4aae1e1 Mon Sep 17 00:00:00 2001 From: Metin San Date: Sat, 27 Apr 2024 11:20:46 +0200 Subject: [PATCH] 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)