diff --git a/zodipy/_contour.py b/zodipy/_contour.py index 16d3dae..651271a 100644 --- a/zodipy/_contour.py +++ b/zodipy/_contour.py @@ -15,7 +15,7 @@ def tabulate_density( - grid: npt.NDArray[np.floating] | Sequence[npt.NDArray[np.floating]], + grid: npt.NDArray[np.float64] | Sequence[npt.NDArray[np.float64]], model: str = "DIRBE", earth_pos: u.Quantity[u.AU] = DEFAULT_EARTH_POS, ) -> npt.NDArray[np.float64]: diff --git a/zodipy/blackbody.py b/zodipy/blackbody.py index 13fdd2f..b7c2346 100644 --- a/zodipy/blackbody.py +++ b/zodipy/blackbody.py @@ -1,3 +1,5 @@ +from __future__ import annotations + import numpy as np import numpy.typing as npt from astropy import units @@ -28,27 +30,20 @@ def get_dust_grain_temperature( return T_0 * R**-delta -def tabulate_center_wavelength_bnu(wavelength: units.Quantity) -> npt.NDArray[np.float64]: - """Tabulate blackbody specific intensity for a center wavelength.""" - return np.asarray( - [ - temperatures.to_value(units.K), - blackbody(wavelength).to_value(units.MJy / units.sr), - ] - ) - - -def tabulate_bandpass_integrated_bnu( - wavelengths: units.Quantity, normalized_weights: units.Quantity +def tabulate_blackbody_emission( + wavelengths: units.Quantity, weights: units.Quantity | None ) -> npt.NDArray[np.float64]: """Tabulate bandpass integrated blackbody specific intensity for a range of temperatures.""" - blackbody_emission = blackbody(wavelengths[:, np.newaxis]) - integrated_blackbody_emission = integrate.trapezoid( - normalized_weights * blackbody_emission.transpose(), wavelengths - ) + if weights is None: + tabulated_blackbody_emission = blackbody(wavelengths) + else: + tabulated_blackbody_emission = integrate.trapezoid( + weights * blackbody(wavelengths[:, np.newaxis]).transpose(), wavelengths + ) + return np.asarray( [ temperatures.to_value(units.K), - integrated_blackbody_emission.to_value(units.MJy / units.sr), + tabulated_blackbody_emission.to_value(units.MJy / units.sr), ] ) diff --git a/zodipy/bodies.py b/zodipy/bodies.py new file mode 100644 index 0000000..8973633 --- /dev/null +++ b/zodipy/bodies.py @@ -0,0 +1,53 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import astropy.coordinates as coords +import numpy as np +from astropy import time, units + +from zodipy._constants import DISTANCE_FROM_EARTH_TO_SEMB_L2 + +if TYPE_CHECKING: + import numpy.typing as npt + + +def get_sun_earth_moon_barycenter( + earthpos: npt.NDArray[np.float64], +) -> npt.NDArray[np.float64]: + """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 = np.linalg.norm(earthpos) + SEMB_L2_distance = earth_distance + DISTANCE_FROM_EARTH_TO_SEMB_L2 + earth_unit_vector = earthpos / earth_distance + + return earth_unit_vector * SEMB_L2_distance + + +def get_earthpos(obs_time: time.Time, ephemeris: str) -> npt.NDArray[np.float64]: + """Return the sky coordinates of the Earth in the heliocentric frame.""" + return ( + coords.get_body("earth", obs_time, ephemeris=ephemeris) + .transform_to(coords.HeliocentricMeanEcliptic) + .cartesian.xyz.to_value(units.AU) + ) + + +def get_obspos( + obs: str, + obstime: time.Time, + earthpos: npt.NDArray[np.float64], + ephemeris: str, +) -> npt.NDArray[np.float64]: + """Return the sky coordinates of the observer in the heliocentric frame.""" + if obs.lower() == "semb-l2": + return get_sun_earth_moon_barycenter(earthpos) + return ( + coords.get_body(obs, obstime, ephemeris=ephemeris) + .transform_to(coords.HeliocentricMeanEcliptic) + .cartesian.xyz.to_value(units.AU) + ) diff --git a/zodipy/brightness.py b/zodipy/brightness.py index fb55098..0a2e593 100644 --- a/zodipy/brightness.py +++ b/zodipy/brightness.py @@ -30,9 +30,9 @@ def kelsall_brightness_at_step( delta: float, emissivity: np.float64, albedo: np.float64, - C1: float, - C2: float, - C3: float, + C1: np.float64, + C2: np.float64, + C3: np.float64, solar_irradiance: np.float64, ) -> npt.NDArray[np.float64]: """Kelsall uses common line of sight grid from obs to 5.2 AU.""" @@ -78,9 +78,3 @@ def rrm_brightness_at_step( blackbody_emission = np.interp(temperature, *bp_interpolation_table) return blackbody_emission * get_density_function(X_helio) * calibration - - -# EMISSION_MAPPING: dict[type[InterplanetaryDustModel], GetCompEmissionAtStepFn] = { -# Kelsall: kelsall, -# RRM: rrm, -# } diff --git a/zodipy/line_of_sight.py b/zodipy/line_of_sight.py index e72eb80..92f378b 100644 --- a/zodipy/line_of_sight.py +++ b/zodipy/line_of_sight.py @@ -11,7 +11,7 @@ if TYPE_CHECKING: import numpy.typing as npt -DIRBE_CUTOFFS: dict[ComponentLabel, tuple[float | np.floating, float | np.floating]] = { +DIRBE_CUTOFFS: dict[ComponentLabel, tuple[float | np.float64, float | np.float64]] = { ComponentLabel.CLOUD: (R_0, R_JUPITER), ComponentLabel.BAND1: (R_0, R_JUPITER), ComponentLabel.BAND2: (R_0, R_JUPITER), @@ -20,7 +20,7 @@ ComponentLabel.FEATURE: (R_EARTH - 0.2, R_EARTH + 0.2), } -RRM_CUTOFFS: dict[ComponentLabel, tuple[float | np.floating, float | np.floating]] = { +RRM_CUTOFFS: dict[ComponentLabel, tuple[float | np.float64, float | np.float64]] = { ComponentLabel.FAN: (R_0, RRM[ComponentLabel.FAN].R_outer), # type: ignore ComponentLabel.INNER_NARROW_BAND: ( RRM[ComponentLabel.INNER_NARROW_BAND].R_inner, # type: ignore diff --git a/zodipy/number_density.py b/zodipy/number_density.py index 4e9a21b..bbc271f 100644 --- a/zodipy/number_density.py +++ b/zodipy/number_density.py @@ -468,7 +468,6 @@ def construct_density_partials_comps( Return a tuple of the density expressions above which has been prepopulated with model and configuration parameters, leaving only the `X_helio` argument to be supplied. Raises exception for incorrectly defined components or component density functions. - """ partial_density_funcs: dict[ComponentLabel, ComponentNumberDensityCallable] = {} for comp_label, comp in comps.items(): diff --git a/zodipy/scattering.py b/zodipy/scattering.py index 92ce320..8ab9ac1 100644 --- a/zodipy/scattering.py +++ b/zodipy/scattering.py @@ -33,7 +33,7 @@ def get_scattering_angle( def get_phase_function( - Theta: npt.NDArray[np.float64], C1: float, C2: float, C3: float + Theta: npt.NDArray[np.float64], C1: np.float64, C2: np.float64, C3: np.float64 ) -> npt.NDArray[np.float64]: """Return the phase function. diff --git a/zodipy/skycoords.py b/zodipy/skycoords.py deleted file mode 100644 index c59c026..0000000 --- a/zodipy/skycoords.py +++ /dev/null @@ -1,74 +0,0 @@ -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", "string_to_coordinate_frame"] - - -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, ephemeris: str) -> coords.SkyCoord: - """Return the sky coordinates of the Earth in the heliocentric frame.""" - 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, - 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, ephemeris=ephemeris).transform_to( - coords.HeliocentricMeanEcliptic - ) - - try: - return coords.SkyCoord( - *obs_pos, - 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 - - -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 - 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/interpolate.py b/zodipy/unpack_model.py similarity index 87% rename from zodipy/interpolate.py rename to zodipy/unpack_model.py index e2fe68d..5d589d6 100644 --- a/zodipy/interpolate.py +++ b/zodipy/unpack_model.py @@ -12,13 +12,16 @@ CompParamDict = dict[ComponentLabel, dict[str, Any]] CommonParamDict = dict[str, Any] +UnpackedModelDicts = tuple[CompParamDict, CommonParamDict] +T = TypeVar("T", bound=ZodiacalLightModel) +UnpackModelCallable = Callable[[units.Quantity, Union[units.Quantity, None], T], UnpackedModelDicts] -def kelsall_params_to_dicts( +def unpack_kelsall( wavelengths: units.Quantity, weights: units.Quantity | None, model: Kelsall, -) -> tuple[CompParamDict, CommonParamDict]: +) -> UnpackedModelDicts: """InterplantaryDustModelToDicts implementation for Kelsall model.""" model_spectrum = model.spectrum.to(wavelengths.unit, equivalencies=units.spectral()) @@ -82,11 +85,11 @@ def kelsall_params_to_dicts( return comp_params, common_params -def rrm_params_to_dicts( +def unpack_rrm( wavelengths: units.Quantity, weights: units.Quantity | None, model: RRM, -) -> tuple[CompParamDict, CommonParamDict]: +) -> UnpackedModelDicts: """InterplantaryDustModelToDicts implementation for Kelsall model.""" model_spectrum = model.spectrum.to(wavelengths.unit, equivalencies=units.spectral()) @@ -121,20 +124,14 @@ def interpolate_spectral_parameter( return interpolated_parameter -T = TypeVar("T", contravariant=True, bound=ZodiacalLightModel) -CallableModelToDicts = Callable[ - [units.Quantity, Union[units.Quantity, None], T], tuple[CompParamDict, CommonParamDict] -] - - -MODEL_INTERPOLATION_MAPPING: dict[type[ZodiacalLightModel], CallableModelToDicts] = { - Kelsall: kelsall_params_to_dicts, - RRM: rrm_params_to_dicts, +model_unpack_mapping: dict[type[ZodiacalLightModel], UnpackModelCallable] = { + Kelsall: unpack_kelsall, + RRM: unpack_rrm, } def get_model_to_dicts_callable( model: ZodiacalLightModel, -) -> CallableModelToDicts: +) -> UnpackModelCallable: """Get the appropriate parameter unpacker for the model.""" - return MODEL_INTERPOLATION_MAPPING[type(model)] + return model_unpack_mapping[type(model)] diff --git a/zodipy/zodiacal_component.py b/zodipy/zodiacal_component.py index 69b16f6..8212c01 100644 --- a/zodipy/zodiacal_component.py +++ b/zodipy/zodiacal_component.py @@ -1,6 +1,6 @@ from __future__ import annotations -from abc import ABC +import abc from dataclasses import dataclass, field from enum import Enum from typing import TYPE_CHECKING @@ -12,7 +12,7 @@ @dataclass -class ZodiacalComponent(ABC): +class ZodiacalComponent(abc.ABC): """Base class for storing common model parameters for zodiacal components. Args: diff --git a/zodipy/zodiacal_light_model.py b/zodipy/zodiacal_light_model.py index fa78a25..8ae684a 100644 --- a/zodipy/zodiacal_light_model.py +++ b/zodipy/zodiacal_light_model.py @@ -100,7 +100,7 @@ def brightness_at_step_callable(cls) -> BrightnessAtStepCallable: @dataclass -class InterplanetaryDustModelRegistry: +class ModelRegistry: """Container for registered models.""" _registry: dict[str, ZodiacalLightModel] = field(init=False, default_factory=dict) @@ -133,4 +133,4 @@ def get_model(self, name: str) -> ZodiacalLightModel: return self._registry[name] -model_registry = InterplanetaryDustModelRegistry() +model_registry = ModelRegistry() diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index 309fa11..ec2b6e9 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -10,12 +10,12 @@ from astropy import units from scipy import integrate -from zodipy.blackbody import tabulate_bandpass_integrated_bnu, tabulate_center_wavelength_bnu -from zodipy.interpolate import get_model_to_dicts_callable +from zodipy.blackbody import tabulate_blackbody_emission +from zodipy.bodies import get_earthpos, get_obspos from zodipy.line_of_sight import get_line_of_sight_range, integrate_gauss_legendre from zodipy.model_registry import model_registry from zodipy.number_density import construct_density_partials_comps -from zodipy.skycoords import get_earth_skycoord, get_obs_skycoord +from zodipy.unpack_model import get_model_to_dicts_callable from zodipy.zodiacal_component import ComponentLabel if TYPE_CHECKING: @@ -84,16 +84,16 @@ def __init__( msg = "Number of wavelengths and weights must be the same in the bandpass." raise ValueError(msg) normalized_weights = weights / integrate.trapezoid(weights, x) - self._b_nu_table = tabulate_bandpass_integrated_bnu(x, normalized_weights) else: - self._b_nu_table = tabulate_center_wavelength_bnu(x) normalized_weights = None + self._b_nu_table = tabulate_blackbody_emission(x, normalized_weights) # Interpolate and convert the model parameters to dictionaries which can be used to evaluate # the zodiacal light model. - self._comp_parameters, self._common_parameters = get_model_to_dicts_callable( - self._interplanetary_dust_model - )(x, normalized_weights, self._interplanetary_dust_model) + brightness_callable_dicts = get_model_to_dicts_callable(self._interplanetary_dust_model)( + x, normalized_weights, self._interplanetary_dust_model + ) + self._comp_parameters, self._common_parameters = brightness_callable_dicts self._ephemeris = ephemeris self._n_proc = n_proc @@ -105,6 +105,7 @@ def evaluate( *, obspos: units.Quantity | str = "earth", return_comps: bool = False, + contains_duplicates: bool = False, ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated zodiacal light for all observations in a `SkyCoord` object. @@ -118,6 +119,9 @@ def evaluate( 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. + contains_duplicates: If True, the input coordinates are filtered and only unique + pointing is used to calculate the emission. The output is then mapped back to the + original coordinates resulting in the same output shape. Defaults to False. Returns: emission: Simulated zodiacal light in units of 'MJy/sr'. @@ -127,48 +131,54 @@ def evaluate( msg = "The `obstime` attribute of the `SkyCoord` object must be set." raise ValueError(msg) - # Pick out unique coordinates, and only calculate the emission for these. and return the - # inverse indices to map the output back to the original coordinates. - _, index, inverse = np.unique( - [skycoord.spherical.lon, skycoord.spherical.lat], - return_index=True, - return_inverse=True, - axis=1, - ) - skycoord = skycoord[index] + ncoords = skycoord.size + if contains_duplicates: + _, index, inverse = np.unique( + [skycoord.spherical.lon, skycoord.spherical.lat], + return_index=True, + return_inverse=True, + axis=1, + ) + skycoord = skycoord[index] # filter out identical coordinates - earth_skycoord = get_earth_skycoord(skycoord.obstime, ephemeris=self._ephemeris) - obs_skycoord = get_obs_skycoord( - obspos, skycoord.obstime, earth_skycoord, ephemeris=self._ephemeris - ) + earth_xyz = get_earthpos(skycoord.obstime, ephemeris=self._ephemeris) + if isinstance(obspos, str): + obs_xyz = get_obspos(obspos, skycoord.obstime, earth_xyz, ephemeris=self._ephemeris) + else: + obs_xyz = obspos.to_value(units.AU) skycoord = skycoord.transform_to(coords.BarycentricMeanEcliptic) - unit_vector = skycoord.cartesian.xyz.value - obspos = obs_skycoord.cartesian.xyz.to_value(units.AU) + + if skycoord.isscalar: + skycoord_xyz = skycoord.cartesian.xyz.value[:, np.newaxis] + else: + skycoord_xyz = skycoord.cartesian.xyz.value start, stop = get_line_of_sight_range( components=self._interplanetary_dust_model.comps.keys(), - unit_vectors=unit_vector, - obs_pos=obspos, + unit_vectors=skycoord_xyz, + obs_pos=obs_xyz, ) + # Return a dict of partial functions corresponding to the number density each zodiacal + # component in the interplanetary dust model. density_partials = construct_density_partials_comps( comps=self._interplanetary_dust_model.comps, - dynamic_params={ - "X_earth": earth_skycoord.cartesian.xyz.to_value(units.AU)[ - :, np.newaxis, np.newaxis - ] - }, + dynamic_params={"X_earth": earth_xyz[:, np.newaxis, np.newaxis]}, ) + + # Partial function of the brightness integral at a step along the line-of-sight prepopulated + # with shared arguments between zodiacal components. common_integrand = functools.partial( self._interplanetary_dust_model.brightness_at_step_callable, - X_obs=obspos[:, np.newaxis, np.newaxis], + X_obs=obs_xyz[:, np.newaxis, np.newaxis], bp_interpolation_table=self._b_nu_table, **self._common_parameters, ) + distribute_to_cores = self._n_proc > 1 and skycoord.size > self._n_proc if distribute_to_cores: - unit_vector_chunks = np.array_split(unit_vector, self._n_proc, axis=-1) + unit_vector_chunks = np.array_split(skycoord_xyz, self._n_proc, axis=-1) integrated_comp_emission = np.zeros( (self._interplanetary_dust_model.ncomps, skycoord.size) ) @@ -211,14 +221,11 @@ def evaluate( ) else: - integrated_comp_emission = np.zeros( - (self._interplanetary_dust_model.ncomps, inverse.size) - ) - + integrated_comp_emission = np.zeros((self._interplanetary_dust_model.ncomps, ncoords)) for idx, comp_label in enumerate(self._interplanetary_dust_model.comps.keys()): comp_integrand = functools.partial( common_integrand, - u_los=np.expand_dims(unit_vector, axis=-1), + u_los=np.expand_dims(skycoord_xyz, axis=-1), 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], @@ -231,8 +238,11 @@ def evaluate( * (stop[comp_label] - start[comp_label]) ) - emission = np.zeros((self._interplanetary_dust_model.ncomps, inverse.size)) - emission = integrated_comp_emission[:, inverse] << (units.MJy / units.sr) + if contains_duplicates: + emission = np.zeros((self._interplanetary_dust_model.ncomps, ncoords)) + emission = integrated_comp_emission[:, inverse] << (units.MJy / units.sr) + else: + emission = integrated_comp_emission << (units.MJy / units.sr) return emission if return_comps else emission.sum(axis=0)