diff --git a/tests/_strategies.py b/tests/_strategies.py index 13f1dfc..0695d27 100644 --- a/tests/_strategies.py +++ b/tests/_strategies.py @@ -26,7 +26,7 @@ ) import zodipy -from zodipy._line_of_sight import COMPONENT_CUTOFFS +from zodipy.line_of_sight import COMPONENT_CUTOFFS from zodipy.model_registry import model_registry MIN_FREQ = u.Quantity(10, u.GHz) diff --git a/zodipy/_contour.py b/zodipy/_contour.py index 6b40514..16d3dae 100644 --- a/zodipy/_contour.py +++ b/zodipy/_contour.py @@ -5,8 +5,8 @@ import astropy.units as u import numpy as np -from zodipy._ipd_dens_funcs import construct_density_partials from zodipy.model_registry import model_registry +from zodipy.number_density import construct_density_partials if TYPE_CHECKING: import numpy.typing as npt diff --git a/zodipy/blackbody.py b/zodipy/blackbody.py index 33db6ce..13fdd2f 100644 --- a/zodipy/blackbody.py +++ b/zodipy/blackbody.py @@ -11,6 +11,23 @@ blackbody = BlackBody(temperatures) +def get_dust_grain_temperature( + R: npt.NDArray[np.float64], T_0: float, delta: float +) -> npt.NDArray[np.float64]: + """Return the dust grain temperature given a radial distance from the Sun. + + Args: + R: Radial distance from the sun in ecliptic heliocentric coordinates [AU / 1AU]. + T_0: Temperature of dust grains located 1 AU from the Sun [K]. + delta: Powerlaw index. + + Returns: + Dust grain temperature [K]. + + """ + 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( diff --git a/zodipy/_emission.py b/zodipy/brightness.py similarity index 81% rename from zodipy/_emission.py rename to zodipy/brightness.py index f53cc5d..fb55098 100644 --- a/zodipy/_emission.py +++ b/zodipy/brightness.py @@ -5,31 +5,27 @@ import numpy as np import numpy.typing as npt -from zodipy._ipd_model import RRM, InterplanetaryDustModel, Kelsall -from zodipy._source_funcs import ( - get_dust_grain_temperature, - get_phase_function, - get_scattering_angle, -) +from zodipy.blackbody import get_dust_grain_temperature +from zodipy.scattering import get_phase_function, get_scattering_angle if TYPE_CHECKING: - from zodipy._ipd_dens_funcs import ComponentDensityFn + from zodipy.number_density import ComponentNumberDensityCallable """ Function that return the zodiacal emission at a step along all lines of sight given a zodiacal model. """ -GetCompEmissionAtStepFn = Callable[..., npt.NDArray[np.float64]] +BrightnessAtStepCallable = Callable[..., npt.NDArray[np.float64]] -def kelsall( +def kelsall_brightness_at_step( r: npt.NDArray[np.float64], start: np.float64, stop: npt.NDArray[np.float64], X_obs: npt.NDArray[np.float64], u_los: npt.NDArray[np.float64], bp_interpolation_table: npt.NDArray[np.float64], - get_density_function: ComponentDensityFn, + get_density_function: ComponentNumberDensityCallable, T_0: float, delta: float, emissivity: np.float64, @@ -59,14 +55,14 @@ def kelsall( return emission * get_density_function(X_helio) -def rrm( +def rrm_brightness_at_step( r: npt.NDArray[np.float64], start: npt.NDArray[np.float64], stop: npt.NDArray[np.float64], X_obs: npt.NDArray[np.float64], u_los: npt.NDArray[np.float64], bp_interpolation_table: npt.NDArray[np.float64], - get_density_function: ComponentDensityFn, + get_density_function: ComponentNumberDensityCallable, T_0: float, delta: float, calibration: np.float64, @@ -84,7 +80,7 @@ def rrm( return blackbody_emission * get_density_function(X_helio) * calibration -EMISSION_MAPPING: dict[type[InterplanetaryDustModel], GetCompEmissionAtStepFn] = { - Kelsall: kelsall, - RRM: rrm, -} +# EMISSION_MAPPING: dict[type[InterplanetaryDustModel], GetCompEmissionAtStepFn] = { +# Kelsall: kelsall, +# RRM: rrm, +# } diff --git a/zodipy/comps.py b/zodipy/comps.py index b44226b..1025866 100644 --- a/zodipy/comps.py +++ b/zodipy/comps.py @@ -1,12 +1,11 @@ from __future__ import annotations from zodipy._constants import R_ASTEROID_BELT, R_KUIPER_BELT, R_MARS -from zodipy._ipd_comps import ( +from zodipy.zodiacal_component import ( Band, BroadBand, Cloud, Comet, - Component, ComponentLabel, Fan, Feature, @@ -15,9 +14,10 @@ NarrowBand, Ring, RingRRM, + ZodiacalComponent, ) -DIRBE: dict[ComponentLabel, Component] = { +DIRBE: dict[ComponentLabel, ZodiacalComponent] = { ComponentLabel.CLOUD: Cloud( x_0=0.011887800744346281, y_0=0.0054765064662263777, @@ -97,7 +97,7 @@ PLANCK.pop(ComponentLabel.FEATURE) -RRM: dict[ComponentLabel, Component] = { +RRM: dict[ComponentLabel, ZodiacalComponent] = { ComponentLabel.FAN: Fan( x_0=0.0, y_0=0.0, diff --git a/zodipy/interpolate.py b/zodipy/interpolate.py index 0fd292a..e2fe68d 100644 --- a/zodipy/interpolate.py +++ b/zodipy/interpolate.py @@ -7,8 +7,8 @@ from astropy import units from scipy import integrate -from zodipy._ipd_model import RRM, InterplanetaryDustModel, Kelsall from zodipy.comps import ComponentLabel +from zodipy.zodiacal_light_model import RRM, Kelsall, ZodiacalLightModel CompParamDict = dict[ComponentLabel, dict[str, Any]] CommonParamDict = dict[str, Any] @@ -121,20 +121,20 @@ def interpolate_spectral_parameter( return interpolated_parameter -T = TypeVar("T", contravariant=True, bound=InterplanetaryDustModel) +T = TypeVar("T", contravariant=True, bound=ZodiacalLightModel) CallableModelToDicts = Callable[ [units.Quantity, Union[units.Quantity, None], T], tuple[CompParamDict, CommonParamDict] ] -MODEL_INTERPOLATION_MAPPING: dict[type[InterplanetaryDustModel], CallableModelToDicts] = { +MODEL_INTERPOLATION_MAPPING: dict[type[ZodiacalLightModel], CallableModelToDicts] = { Kelsall: kelsall_params_to_dicts, RRM: rrm_params_to_dicts, } def get_model_to_dicts_callable( - model: InterplanetaryDustModel, + model: ZodiacalLightModel, ) -> CallableModelToDicts: """Get the appropriate parameter unpacker for the model.""" return MODEL_INTERPOLATION_MAPPING[type(model)] diff --git a/zodipy/_line_of_sight.py b/zodipy/line_of_sight.py similarity index 87% rename from zodipy/_line_of_sight.py rename to zodipy/line_of_sight.py index 92a46a0..e72eb80 100644 --- a/zodipy/_line_of_sight.py +++ b/zodipy/line_of_sight.py @@ -1,12 +1,12 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Iterable +from typing import TYPE_CHECKING, Callable, Iterable import numpy as np from zodipy._constants import R_0, R_EARTH, R_JUPITER, R_KUIPER_BELT -from zodipy._ipd_comps import ComponentLabel from zodipy.comps import RRM +from zodipy.zodiacal_component import ComponentLabel if TYPE_CHECKING: import numpy.typing as npt @@ -47,6 +47,15 @@ COMPONENT_CUTOFFS = {**DIRBE_CUTOFFS, **RRM_CUTOFFS} +def integrate_gauss_legendre( + fn: Callable[[float], npt.NDArray[np.float64]], + points: npt.NDArray[np.float64], + weights: npt.NDArray[np.float64], +) -> npt.NDArray[np.float64]: + """Integrate a function using Gauss-Legendre quadrature.""" + return np.squeeze(sum(fn(x) * w for x, w in zip(points, weights))) + + def get_sphere_intersection( obs_pos: npt.NDArray[np.float64], unit_vectors: npt.NDArray[np.float64], diff --git a/zodipy/model_registry.py b/zodipy/model_registry.py index ddd918f..502bbe9 100644 --- a/zodipy/model_registry.py +++ b/zodipy/model_registry.py @@ -1,5 +1,5 @@ from zodipy import comps, source_params -from zodipy._ipd_model import RRM, Kelsall, model_registry +from zodipy.zodiacal_light_model import RRM, Kelsall, model_registry model_registry.register_model( name="dirbe", diff --git a/zodipy/_ipd_dens_funcs.py b/zodipy/number_density.py similarity index 89% rename from zodipy/_ipd_dens_funcs.py rename to zodipy/number_density.py index 7109a99..4e9a21b 100644 --- a/zodipy/_ipd_dens_funcs.py +++ b/zodipy/number_density.py @@ -8,12 +8,11 @@ import numpy as np import numpy.typing as npt # type: ignore -from zodipy._ipd_comps import ( +from zodipy.zodiacal_component import ( Band, BroadBand, Cloud, Comet, - Component, ComponentLabel, Fan, Feature, @@ -22,6 +21,7 @@ NarrowBand, Ring, RingRRM, + ZodiacalComponent, ) """The density functions for the different types of components. @@ -36,7 +36,7 @@ class must have all the parameters as instance variables. ComputeDensityFunc = Callable[..., npt.NDArray[np.float64]] -def compute_cloud_density( +def cloud_number_density( X_helio: npt.NDArray[np.float64], X_0: npt.NDArray[np.float64], sin_Omega_rad: float, @@ -66,7 +66,7 @@ def compute_cloud_density( return n_0 * R_cloud**-alpha * np.exp(-beta * g**gamma) -def compute_band_density( +def band_number_density( X_helio: npt.NDArray[np.float64], X_0: npt.NDArray[np.float64], sin_Omega_rad: float, @@ -103,7 +103,7 @@ def compute_band_density( return term1 * term2 * term3 * term4 -def compute_ring_density( +def ring_number_density( X_helio: npt.NDArray[np.float64], X_0: npt.NDArray[np.float64], sin_Omega_rad: float, @@ -132,7 +132,7 @@ def compute_ring_density( return n_0 * np.exp(term1 - term2) -def compute_feature_density( +def feature_number_density( X_helio: npt.NDArray[np.float64], X_0: npt.NDArray[np.float64], X_earth: npt.NDArray[np.float64], @@ -174,7 +174,7 @@ def compute_feature_density( return n_0 * np.exp(-exp_term) -def compute_fan_density( # * +def fan_number_density( X_helio: npt.NDArray[np.float64], X_0: npt.NDArray[np.float64], sin_Omega_rad: float, @@ -211,7 +211,7 @@ def compute_fan_density( # * return density -def compute_comet_density( +def comet_number_density( X_helio: npt.NDArray[np.float64], X_0: npt.NDArray[np.float64], sin_Omega_rad: float, @@ -249,14 +249,15 @@ def compute_comet_density( return density -def compute_interstellar_density( +def interstellar_number_density( X_helio: npt.NDArray[np.float64], amp: float, ) -> npt.NDArray[np.float64]: + """Interstellar constant number density.""" return np.array([amp]) -def compute_narrow_band_density( +def narrow_band_number_density( X_helio: npt.NDArray[np.float64], X_0: npt.NDArray[np.float64], sin_Omega_rad: float, @@ -296,7 +297,7 @@ def compute_narrow_band_density( return density -def compute_broad_band_density( +def broad_band_number_density( X_helio: npt.NDArray[np.float64], X_0: npt.NDArray[np.float64], sin_Omega_rad: float, @@ -334,7 +335,7 @@ def compute_broad_band_density( return density -def compute_ring_density_rmm( +def rrm_ring_number_density( X_helio: npt.NDArray[np.float64], X_0: npt.NDArray[np.float64], sin_Omega_rad: float, @@ -347,7 +348,8 @@ def compute_ring_density_rmm( sigma_z: float, A: float, ) -> npt.NDArray[np.float64]: - return A * compute_ring_density( + """RRM ring is just K98 ring with a scaling factor.""" + return A * ring_number_density( X_helio=X_helio, X_0=X_0, sin_Omega_rad=sin_Omega_rad, @@ -361,7 +363,7 @@ def compute_ring_density_rmm( ) -def compute_feature_density_rmm( +def rrm_feature_number_density( X_helio: npt.NDArray[np.float64], X_0: npt.NDArray[np.float64], X_earth: npt.NDArray[np.float64], @@ -377,7 +379,8 @@ def compute_feature_density_rmm( sigma_theta_rad: float, A: float, ) -> npt.NDArray[np.float64]: - return A * compute_feature_density( + """RRM feature is just K98 feature with a scaling factor.""" + return A * feature_number_density( X_helio=X_helio, X_0=X_0, sin_Omega_rad=sin_Omega_rad, @@ -395,36 +398,39 @@ def compute_feature_density_rmm( # Mapping of implemented zodiacal component data classes and their density functions. -DENSITY_FUNCS: dict[type[Component], ComputeDensityFunc] = { - Cloud: compute_cloud_density, - Band: compute_band_density, - Ring: compute_ring_density, - Feature: compute_feature_density, - Fan: compute_fan_density, - Comet: compute_comet_density, - Interstellar: compute_interstellar_density, - NarrowBand: compute_narrow_band_density, - BroadBand: compute_broad_band_density, - RingRRM: compute_ring_density_rmm, - FeatureRRM: compute_feature_density_rmm, +DENSITY_FUNCS: dict[type[ZodiacalComponent], ComputeDensityFunc] = { + Cloud: cloud_number_density, + Band: band_number_density, + Ring: ring_number_density, + Feature: feature_number_density, + Fan: fan_number_density, + Comet: comet_number_density, + Interstellar: interstellar_number_density, + NarrowBand: narrow_band_number_density, + BroadBand: broad_band_number_density, + RingRRM: rrm_ring_number_density, + FeatureRRM: rrm_feature_number_density, } -class ComponentDensityFn(Protocol): - def __call__(self, X_helio: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: ... +class ComponentNumberDensityCallable(Protocol): + """Protocol for a zodiacal components number density function.""" + + def __call__(self, X_helio: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + """Return the number density of the component at the heliocentric position.""" def construct_density_partials( - comps: Sequence[Component], + comps: Sequence[ZodiacalComponent], dynamic_params: dict[str, Any], -) -> tuple[ComponentDensityFn, ...]: +) -> tuple[ComponentNumberDensityCallable, ...]: """Return density partials for the components. 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: list[ComponentDensityFn] = [] + partial_density_funcs: list[ComponentNumberDensityCallable] = [] for comp in comps: comp_dict = asdict(comp) func_params = inspect.signature(DENSITY_FUNCS[type(comp)]).parameters.keys() @@ -454,9 +460,9 @@ def construct_density_partials( def construct_density_partials_comps( - comps: Mapping[ComponentLabel, Component], + comps: Mapping[ComponentLabel, ZodiacalComponent], dynamic_params: dict[str, Any], -) -> dict[ComponentLabel, ComponentDensityFn]: +) -> dict[ComponentLabel, ComponentNumberDensityCallable]: """Construct density partials for components. Return a tuple of the density expressions above which has been prepopulated with @@ -464,7 +470,7 @@ def construct_density_partials_comps( Raises exception for incorrectly defined components or component density functions. """ - partial_density_funcs: dict[ComponentLabel, ComponentDensityFn] = {} + partial_density_funcs: dict[ComponentLabel, ComponentNumberDensityCallable] = {} for comp_label, comp in comps.items(): comp_dict = asdict(comp) func_params = inspect.signature(DENSITY_FUNCS[type(comp)]).parameters.keys() diff --git a/zodipy/_source_funcs.py b/zodipy/scattering.py similarity index 79% rename from zodipy/_source_funcs.py rename to zodipy/scattering.py index 9753ba6..92ce320 100644 --- a/zodipy/_source_funcs.py +++ b/zodipy/scattering.py @@ -9,23 +9,6 @@ import numpy.typing as npt -def get_dust_grain_temperature( - R: npt.NDArray[np.float64], T_0: float, delta: float -) -> npt.NDArray[np.float64]: - """Return the dust grain temperature given a radial distance from the Sun. - - Args: - R: Radial distance from the sun in ecliptic heliocentric coordinates [AU / 1AU]. - T_0: Temperature of dust grains located 1 AU from the Sun [K]. - delta: Powerlaw index. - - Returns: - Dust grain temperature [K]. - - """ - return T_0 * R**-delta - - def get_scattering_angle( R_los: float | npt.NDArray[np.float64], R_helio: npt.NDArray[np.float64], diff --git a/zodipy/_coords.py b/zodipy/skycoords.py similarity index 100% rename from zodipy/_coords.py rename to zodipy/skycoords.py diff --git a/zodipy/source_params.py b/zodipy/source_params.py index 6381ff2..677987c 100644 --- a/zodipy/source_params.py +++ b/zodipy/source_params.py @@ -1,6 +1,6 @@ import astropy.units as u -from zodipy._ipd_comps import ComponentLabel +from zodipy.zodiacal_component import ComponentLabel T_0_DIRBE = 286 DELTA_DIRBE = 0.46686259861486573 diff --git a/zodipy/_ipd_comps.py b/zodipy/zodiacal_component.py similarity index 88% rename from zodipy/_ipd_comps.py rename to zodipy/zodiacal_component.py index 2efc914..69b16f6 100644 --- a/zodipy/_ipd_comps.py +++ b/zodipy/zodiacal_component.py @@ -12,7 +12,7 @@ @dataclass -class Component(ABC): +class ZodiacalComponent(ABC): """Base class for storing common model parameters for zodiacal components. Args: @@ -45,7 +45,7 @@ def __post_init__(self) -> None: @dataclass -class Cloud(Component): +class Cloud(ZodiacalComponent): """DIRBE diffuse cloud. Args: @@ -65,7 +65,7 @@ class Cloud(Component): @dataclass -class Band(Component): +class Band(ZodiacalComponent): """DIRBE asteroidal dust band. Args: @@ -90,7 +90,7 @@ def __post_init__(self) -> None: @dataclass -class Ring(Component): +class Ring(ZodiacalComponent): """DIRBE circum-solar ring (excluding the Earth-trailing Feature). Args: @@ -108,7 +108,7 @@ class Ring(Component): @dataclass -class Feature(Component): +class Feature(ZodiacalComponent): """DIRBE Earth-trailing Feature. Args: @@ -137,7 +137,9 @@ def __post_init__(self) -> None: @dataclass -class Fan(Component): +class Fan(ZodiacalComponent): + """RRM fan.""" + gamma: float Z_0: float Q: float @@ -146,7 +148,9 @@ class Fan(Component): @dataclass -class Comet(Component): +class Comet(ZodiacalComponent): + """RRM comet.""" + gamma: float Z_0: float P: float @@ -156,12 +160,16 @@ class Comet(Component): @dataclass -class Interstellar(Component): +class Interstellar(ZodiacalComponent): + """RRM interstellar dust.""" + amp: float @dataclass -class NarrowBand(Component): +class NarrowBand(ZodiacalComponent): + """RRM narrow band.""" + gamma: float A: float G: float @@ -171,7 +179,9 @@ class NarrowBand(Component): @dataclass -class BroadBand(Component): +class BroadBand(ZodiacalComponent): + """RRM broad band.""" + gamma: float A: float R_inner: float @@ -182,11 +192,15 @@ class BroadBand(Component): @dataclass class RingRRM(Ring): + """RRM circum-solar ring.""" + A: float @dataclass class FeatureRRM(Feature): + """RRM Earth-trailing Feature.""" + A: float diff --git a/zodipy/_ipd_model.py b/zodipy/zodiacal_light_model.py similarity index 68% rename from zodipy/_ipd_model.py rename to zodipy/zodiacal_light_model.py index 01d22d2..fa78a25 100644 --- a/zodipy/_ipd_model.py +++ b/zodipy/zodiacal_light_model.py @@ -1,18 +1,24 @@ from __future__ import annotations -from abc import ABC +import abc from dataclasses import dataclass, field from typing import TYPE_CHECKING, Mapping, Sequence import numpy as np from astropy import units +from zodipy.brightness import ( + BrightnessAtStepCallable, + kelsall_brightness_at_step, + rrm_brightness_at_step, +) + if TYPE_CHECKING: - from zodipy._ipd_comps import Component, ComponentLabel + from zodipy.zodiacal_component import ComponentLabel, ZodiacalComponent -@dataclass -class InterplanetaryDustModel(ABC): +@dataclass(repr=False) +class ZodiacalLightModel(abc.ABC): """Base class for interplanetary dust models. Args: @@ -21,9 +27,15 @@ class InterplanetaryDustModel(ABC): """ - comps: Mapping[ComponentLabel, Component] + comps: Mapping[ComponentLabel, ZodiacalComponent] spectrum: units.Quantity + @property + @abc.abstractmethod + def brightness_at_step_callable(cls) -> BrightnessAtStepCallable: + """Return the callable that computes the brightness at a step.""" + ... + def to_dict(self) -> dict: """Return a dictionary representation of the model.""" _dict: dict = {} @@ -45,6 +57,7 @@ def to_dict(self) -> dict: @property def ncomps(self) -> int: + """Return the number of components in the model.""" return len(self.comps) def is_valid_at(self, wavelengths: units.Quantity) -> np.bool_: @@ -53,8 +66,8 @@ def is_valid_at(self, wavelengths: units.Quantity) -> np.bool_: return np.all((self.spectrum.min() <= wavelengths) & (wavelengths <= self.spectrum.max())) -@dataclass -class Kelsall(InterplanetaryDustModel): +@dataclass(repr=False) +class Kelsall(ZodiacalLightModel): """Kelsall et al. (1998) model.""" T_0: float @@ -66,38 +79,50 @@ class Kelsall(InterplanetaryDustModel): C2: Sequence[float] | None = None C3: Sequence[float] | None = None + @property + def brightness_at_step_callable(cls) -> BrightnessAtStepCallable: + """Kellsall brightness at a step fuction.""" + return kelsall_brightness_at_step -@dataclass -class RRM(InterplanetaryDustModel): + +@dataclass(repr=False) +class RRM(ZodiacalLightModel): """Rowan-Robinson and May (2013) model.""" T_0: Mapping[ComponentLabel, float] delta: Mapping[ComponentLabel, float] calibration: Sequence[float] + @property + def brightness_at_step_callable(cls) -> BrightnessAtStepCallable: + """RRM brightness at a step fuction.""" + return rrm_brightness_at_step + @dataclass class InterplanetaryDustModelRegistry: """Container for registered models.""" - _registry: dict[str, InterplanetaryDustModel] = field(init=False, default_factory=dict) + _registry: dict[str, ZodiacalLightModel] = field(init=False, default_factory=dict) @property def models(self) -> list[str]: + """Return a list of registered models.""" return list(self._registry.keys()) def register_model( self, name: str, - model: InterplanetaryDustModel, + model: ZodiacalLightModel, ) -> None: + """Register a model with the registry.""" if (name := name.lower()) in self._registry: msg = f"a model by the name {name!s} is already registered." raise ValueError(msg) - self._registry[name] = model - def get_model(self, name: str) -> InterplanetaryDustModel: + def get_model(self, name: str) -> ZodiacalLightModel: + """Return a model from the registry.""" if (name := name.lower()) not in self._registry: msg = ( f"{name!r} is not a registered Interplanetary Dust model. " diff --git a/zodipy/zodipy.py b/zodipy/zodipy.py index 8ede134..c7a65c8 100644 --- a/zodipy/zodipy.py +++ b/zodipy/zodipy.py @@ -3,21 +3,20 @@ import functools import multiprocessing import platform -from typing import TYPE_CHECKING, Callable +from typing import TYPE_CHECKING import numpy as np from astropy import coordinates as coords from astropy import units from scipy import integrate -from zodipy._coords import get_earth_skycoord, get_obs_skycoord -from zodipy._emission import EMISSION_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_range from zodipy.blackbody import tabulate_bandpass_integrated_bnu, tabulate_center_wavelength_bnu from zodipy.interpolate import get_model_to_dicts_callable +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.zodiacal_component import ComponentLabel if TYPE_CHECKING: import numpy.typing as npt @@ -27,13 +26,7 @@ class Model: - """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. - """ + """Main interface to ZodiPy.""" def __init__( self, @@ -53,7 +46,8 @@ def __init__( `weights` must be provided. weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are assumed to be given in spectral radiance units (Jy/sr). - name: Interplanetary dust model to use. Defaults to 'dirbe' which is the DIRBE model. + name: Interplanetary dust model to use. For a list of available models, see + https://cosmoglobe.github.io/zodipy/introduction/. Defaults to 'dirbe'. gauss_quad_degree: Order of the Gaussian-Legendre quadrature used to evaluate the line-of-sight integral in the simulations. Default is 50 points. extrapolate: If `True` all spectral quantities in the selected model are extrapolated to @@ -75,6 +69,7 @@ def __init__( raise ValueError(msg) self._interplanetary_dust_model = model_registry.get_model(name) + if not extrapolate and not self._interplanetary_dust_model.is_valid_at(x): msg = ( "The requested frequencies are outside the valid range of the model." @@ -152,11 +147,13 @@ def evaluate( ) skycoord = skycoord.transform_to(coords.BarycentricMeanEcliptic) + unit_vector = skycoord.cartesian.xyz.value + obspos = obs_skycoord.cartesian.xyz.to_value(units.AU) start, stop = get_line_of_sight_range( components=self._interplanetary_dust_model.comps.keys(), - unit_vectors=skycoord.cartesian.xyz.value, - obs_pos=obs_skycoord.cartesian.xyz.to_value(units.AU), + unit_vectors=unit_vector, + obs_pos=obspos, ) density_partials = construct_density_partials_comps( @@ -167,17 +164,15 @@ def evaluate( ] }, ) - common_integrand = functools.partial( - EMISSION_MAPPING[type(self._interplanetary_dust_model)], - X_obs=obs_skycoord.cartesian.xyz.to_value(units.AU)[:, np.newaxis, np.newaxis], + self._interplanetary_dust_model.brightness_at_step_callable, + X_obs=obspos[:, 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(skycoord.cartesian.xyz.value, self._n_proc, axis=-1) + unit_vector_chunks = np.array_split(unit_vector, self._n_proc, axis=-1) integrated_comp_emission = np.zeros( (self._interplanetary_dust_model.ncomps, skycoord.size) ) @@ -194,20 +189,20 @@ def evaluate( comp_integrands = [ functools.partial( common_integrand, - u_los=np.expand_dims(unit_vectors, axis=-1), + u_los=np.expand_dims(unit_vector, axis=-1), start=np.expand_dims(start, axis=-1), stop=np.expand_dims(stop, axis=-1), get_density_function=density_partials[comp_label], **self._comp_parameters[comp_label], ) - for unit_vectors, start, stop in zip( + for unit_vector, start, stop in zip( unit_vector_chunks, start_chunks, stop_chunks ) ] proc_chunks = [ pool.apply_async( - _integrate_gauss_quad, + integrate_gauss_legendre, args=(comp_integrand, *self._gauss_points_and_weights), ) for comp_integrand in comp_integrands @@ -223,12 +218,11 @@ def evaluate( integrated_comp_emission = np.zeros( (self._interplanetary_dust_model.ncomps, skycoord.size) ) - unit_vectors_expanded = np.expand_dims(skycoord.cartesian.xyz.value, axis=-1) for idx, comp_label in enumerate(self._interplanetary_dust_model.comps.keys()): comp_integrand = functools.partial( common_integrand, - u_los=unit_vectors_expanded, + u_los=np.expand_dims(unit_vector, 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], @@ -236,7 +230,7 @@ def evaluate( ) integrated_comp_emission[idx] = ( - _integrate_gauss_quad(comp_integrand, *self._gauss_points_and_weights) + integrate_gauss_legendre(comp_integrand, *self._gauss_points_and_weights) * 0.5 * (stop[comp_label] - start[comp_label]) ) @@ -246,12 +240,6 @@ def evaluate( 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.""" - 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.""" return self._interplanetary_dust_model.to_dict() @@ -286,12 +274,3 @@ def __repr__(self) -> str: repr_str += f"{attribute_name}={attribute!r}, " return repr_str[:-2] + ")" - - -def _integrate_gauss_quad( - fn: Callable[[float], npt.NDArray[np.float64]], - points: npt.NDArray[np.float64], - weights: npt.NDArray[np.float64], -) -> npt.NDArray[np.float64]: - """Integrate a function using Gauss-Legendre quadrature.""" - return np.squeeze(sum(fn(x) * w for x, w in zip(points, weights)))