diff --git a/tests/test_evaluate.py b/tests/test_evaluate.py index 4ece603..5d83203 100644 --- a/tests/test_evaluate.py +++ b/tests/test_evaluate.py @@ -109,6 +109,48 @@ def test_output_shape() -> None: assert test_model.evaluate(skycoord, return_comps=False).shape == (6,) +def test_input_shape() -> None: + """Test that shape of the input sky_coord is correct and obspos and obstime are correct.""" + obstimes = time.Time([TEST_TIME.mjd + i for i in range(4)], format="mjd") + obsposes = ( + coords.get_body("earth", obstimes) + .transform_to(coords.HeliocentricMeanEcliptic) + .cartesian.xyz.to(units.AU) + ) + test_model.evaluate(SkyCoord([20, 30], [30, 40], unit=units.deg, obstime=TEST_TIME)) + test_model.evaluate( + SkyCoord([20, 30, 40, 50], [30, 40, 30, 20], unit=units.deg, obstime=obstimes) + ) + test_model.evaluate( + SkyCoord([20, 30, 40, 50], [30, 40, 30, 20], unit=units.deg, obstime=obstimes), + obspos=obsposes, + ) + + # obstime > sky_coord + with pytest.raises(ValueError): + test_model.evaluate(SkyCoord([20, 30], [30, 40], unit=units.deg, obstime=obstimes)) + + # obspos > sky_coord + with pytest.raises(ValueError): + test_model.evaluate(SkyCoord(20, 30, unit=units.deg, obstime=TEST_TIME), obspos=obsposes) + + # obspos > obstime + with pytest.raises(ValueError): + test_model.evaluate(SkyCoord(20, 30, unit=units.deg, obstime=TEST_TIME), obspos=obsposes) + + # obstime > obspos + with pytest.raises(ValueError): + test_model.evaluate( + SkyCoord([20, 30, 40, 50], [30, 40, 30, 20], unit=units.deg, obstime=obstimes), + obspos=[[1, -0.4, 0.2]] * units.AU, + ) + with pytest.raises(ValueError): + test_model.evaluate( + SkyCoord([20, 30, 40, 50], [30, 40, 30, 20], unit=units.deg, obstime=obstimes), + obspos=obsposes[:2], + ) + + def test_return_comps() -> None: """Test that the return_comps function works for valid user input.""" emission_comps = test_model.evaluate(TEST_SKY_COORD, return_comps=True) diff --git a/zodipy/__init__.py b/zodipy/__init__.py index 7f54dca..e8d30b9 100644 --- a/zodipy/__init__.py +++ b/zodipy/__init__.py @@ -1,6 +1,6 @@ +from zodipy.model import Model from zodipy.model_registry import model_registry from zodipy.number_density import grid_number_density -from zodipy.zodipy import Model __all__ = ( "Model", diff --git a/zodipy/bodies.py b/zodipy/bodies.py index 35becf3..d4e4568 100644 --- a/zodipy/bodies.py +++ b/zodipy/bodies.py @@ -5,6 +5,7 @@ import astropy.coordinates as coords import numpy as np from astropy import time, units +from scipy import interpolate if TYPE_CHECKING: import numpy.typing as npt @@ -29,11 +30,32 @@ def get_sun_earth_moon_barycenter( def get_earthpos_xyz(obstime: time.Time, ephemeris: str) -> npt.NDArray[np.float64]: """Return the sky coordinates of the Earth in the heliocentric frame.""" - return ( - coords.get_body("earth", obstime, ephemeris=ephemeris) + if obstime.size == 1: + return ( + coords.get_body("earth", obstime, ephemeris=ephemeris) + .transform_to(coords.HeliocentricMeanEcliptic) + .cartesian.xyz.to_value(units.AU) + ) + return get_interpolated_body_xyz("earth", obstime, ephemeris) + + +def get_interpolated_body_xyz( + body: str, + obstimes: time.Time, + ephemeris: str, +) -> npt.NDArray[np.float64]: + """Return interpolated Earth positions in the heliocentric frame.""" + dt = (1 * units.hour).to_value(units.day) + t0, t1 = obstimes[0].mjd, obstimes[-1].mjd + times = time.Time(np.arange(t0, max(t0 + 365, t1), dt), format="mjd") + + bodypos = ( + coords.get_body(body, times, ephemeris=ephemeris) .transform_to(coords.HeliocentricMeanEcliptic) .cartesian.xyz.to_value(units.AU) ) + f = interpolate.interp1d(times.mjd, bodypos, axis=-1) + return f(obstimes.mjd) def get_obspos_xyz( @@ -46,12 +68,17 @@ def get_obspos_xyz( if isinstance(obspos, str): if obspos.lower() == "semb-l2": return get_sun_earth_moon_barycenter(earthpos) + if obspos.lower() == "earth": + return earthpos + try: - return ( - coords.get_body(obspos, obstime, ephemeris=ephemeris) - .transform_to(coords.HeliocentricMeanEcliptic) - .cartesian.xyz.to_value(units.AU) - ) + if obstime.size == 1: + return ( + coords.get_body(obspos, obstime, ephemeris=ephemeris) + .transform_to(coords.HeliocentricMeanEcliptic) + .cartesian.xyz.to_value(units.AU) + ) + return get_interpolated_body_xyz(obspos, obstime, ephemeris) except KeyError as error: valid_obs = [*coords.solar_system_ephemeris.bodies, "semb-l2"] msg = f"Invalid observer string: '{obspos}'. Valid observers are: {valid_obs}" diff --git a/zodipy/brightness.py b/zodipy/brightness.py index b824d18..df438b4 100644 --- a/zodipy/brightness.py +++ b/zodipy/brightness.py @@ -36,7 +36,8 @@ def kelsall_brightness_at_step( solar_irradiance: np.float64, ) -> npt.NDArray[np.float64]: """Kelsall uses common line of sight grid from obs to 5.2 AU.""" - # Convert the quadrature range from [0, inf] to the true ecliptic positions + # Convert the quadrature range from [-1, 1] to the true ecliptic positions + # and back again at the end R_los = 0.5 * (stop - start) * r + 0.5 * (stop + start) X_los = R_los * u_los @@ -46,14 +47,13 @@ def kelsall_brightness_at_step( temperature = get_dust_grain_temperature(R_helio, T_0, delta) blackbody_emission = np.interp(temperature, *bp_interpolation_table) emission = (1 - albedo) * (emissivity * blackbody_emission) - if albedo != 0: solar_flux = solar_irradiance / R_helio**2 scattering_angle = get_scattering_angle(R_los, R_helio, X_los, X_helio) phase_function = get_phase_function(scattering_angle, C1, C2, C3) emission += albedo * solar_flux * phase_function - return emission * get_density_function(X_helio) + return emission * get_density_function(X_helio) * 0.5 * (stop - start) def rrm_brightness_at_step( @@ -69,7 +69,8 @@ def rrm_brightness_at_step( calibration: np.float64, ) -> npt.NDArray[np.float64]: """RRM is implented with component specific line-of-sight grids.""" - # Convert the quadrature range from [0, inf] to the true ecliptic positions + # Convert the quadrature range from [-1, 1] to the true ecliptic positions + # and back again at the end R_los = 0.5 * (stop - start) * r + 0.5 * (stop + start) X_los = R_los * u_los @@ -79,4 +80,4 @@ def rrm_brightness_at_step( temperature = get_dust_grain_temperature(R_helio, T_0, delta) blackbody_emission = np.interp(temperature, *bp_interpolation_table) - return blackbody_emission * get_density_function(X_helio) * calibration + return blackbody_emission * get_density_function(X_helio) * calibration * 0.5 * (stop - start) diff --git a/zodipy/line_of_sight.py b/zodipy/line_of_sight.py index 4114407..6992b89 100644 --- a/zodipy/line_of_sight.py +++ b/zodipy/line_of_sight.py @@ -21,8 +21,8 @@ ComponentLabel.BAND1: (R_0, R_JUPITER), ComponentLabel.BAND2: (R_0, R_JUPITER), ComponentLabel.BAND3: (R_0, R_JUPITER), - ComponentLabel.RING: (R_EARTH - 0.2, R_EARTH + 0.2), - ComponentLabel.FEATURE: (R_EARTH - 0.2, R_EARTH + 0.2), + ComponentLabel.RING: (R_0, R_EARTH + 0.2), + ComponentLabel.FEATURE: (R_0, R_EARTH + 0.2), } RRM_CUTOFFS: dict[ComponentLabel, tuple[float | np.float64, float | np.float64]] = { @@ -67,10 +67,12 @@ def get_sphere_intersection( cutoff: float | np.float64, ) -> npt.NDArray[np.float64]: """Returns the distance from the observer to a heliocentric sphere with radius `cutoff`.""" - x, y, z = obs_pos.flatten() + x, y, z = obs_pos r_obs = np.sqrt(x**2 + y**2 + z**2) - if r_obs > cutoff: - return np.asarray([np.finfo(float).eps]) + if (r_obs > cutoff).any(): + if obs_pos.ndim == 1: + return np.array([np.finfo(float).eps]) + return np.full(obs_pos.shape[-1], np.finfo(float).eps) u_x, u_y, u_z = unit_vectors lon = np.arctan2(u_y, u_x) diff --git a/zodipy/zodipy.py b/zodipy/model.py similarity index 70% rename from zodipy/zodipy.py rename to zodipy/model.py index f05b5f8..bbe74fc 100644 --- a/zodipy/zodipy.py +++ b/zodipy/model.py @@ -27,7 +27,7 @@ class Model: - """The `Model` class is the main interface to ZodiPy.""" + """Main interface to ZodiPy.""" def __init__( self, @@ -39,23 +39,23 @@ def __init__( extrapolate: bool = False, ephemeris: str = "builtin", ) -> None: - """Initialize the Zodipy interface. + """Initialize the ZodiPy model interface. Args: x: Wavelength or frequency. If `x` is a sequence, it is assumed to be a the points corresponding to a bandpass and the corresponding `weights` must be provided. - weights: Bandpass weights corresponding the the frequencies in `freq`. The weights are - assumed to represent a normalized instrument response in units of spectral radiance - (Jy/sr). - 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-laguerre quadrature used to evaluate the - line-of-sight integral in the simulations. Default is 50 points. + weights: Bandpass weights corresponding the the frequencies or wavelengths in `x`. The + weights are assumed to represent a normalized instrument response in units of + spectral radiance (Jy/sr). + name: Zodiacal light model to use for the simulations. 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 integrals in the simulations. Default is 50 points. extrapolate: 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 `x` outside of the valid model range. Default is `False`. + requested values of `x` outside of the valid model range. Default is `False`. ephemeris: Ephemeris used in `astropy.coordinates.solar_system_ephemeris` to compute the - positions of the observer and the Earth. Defaults to 'builtin'. See the + positions of the observer and Earth. Defaults to 'builtin'. See the [Astropy documentation](https://docs.astropy.org/en/stable/coordinates/solarsystem.html) for available ephemerides. @@ -107,29 +107,28 @@ def evaluate( *, obspos: units.Quantity | str = "earth", return_comps: bool = False, - contains_duplicates: bool = False, nprocesses: int = 1, ) -> units.Quantity[units.MJy / units.sr]: """Return the simulated zodiacal light. - The zodiacal light is simulated for a single, or a sequence of observations from a - position in the Solar system specified by the `obspos` argument, and at instant in - time specified in `skycoord` The `obstime` and `frame` keywords must be specified in the - `SkyCoord` object. + The zodiacal light is simulated for a single, or a sequence of observations. If a single + `obspos` and `obstime` is provided for multiple coordinates, all coordinates are assumed to + be observed from that position at that time. Otherwise, when `obspos` and `obstime` contains + multiple values, corresponding to coordinates in `skycoord`, the zodiacal light is simulated + in a time-ordered manner. Args: - skycoord: `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 the `BarycentricMeanEcliptic` frame. + skycoord: `astropy.coordinates.SkyCoord` object representing the coordinates or + observations for which to simulate the zodiacal light. The `frame` and `obstime` + attributes of the `SkyCoord` object must be set. The `obstime` attribute must be + specified, and correspond to a single, or a sequence of observational times with + length matching the number of coordinates. The frame must be convertible to the + `astropy.coordinates.BarycentricMeanEcliptic` frame. obspos: 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'. + an observer in the `astropy.coordinates.solar_system_ephemeris`. If an explicit + position is given, it must either be a single, or a sequence of positions with + shape matching the number of coordinates 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. nprocesses: Number of cores to use. If `nprocesses >= 1`, the line-of-sight integrals are parallelized using the `multiprocessing` module. Defaults to 1. @@ -137,29 +136,7 @@ def evaluate( emission: Simulated zodiacal light in units of 'MJy/sr'. """ - try: - obstime = typing.cast(time.Time, skycoord.obstime) - except AttributeError as error: - msg = "The input coordinates must be an astropy SkyCoord object." - raise TypeError(msg) from error - if obstime is None: - msg = "The `obstime` attribute of the `SkyCoord` object must be set." - raise ValueError(msg) - - if contains_duplicates: - _, index, inverse = np.unique( - typing.cast( - list[npt.NDArray[np.float64]], - [skycoord.spherical.lon.value, skycoord.spherical.lat.value], - ), - return_index=True, - return_inverse=True, - axis=1, - ) - skycoord = typing.cast( - coords.SkyCoord, skycoord[index] - ) # filter out identical coordinates - ncoords = skycoord.size + obstime = validate_user_input(skycoord, obspos) earth_xyz = get_earthpos_xyz(obstime, self._ephemeris) obs_xyz = get_obspos_xyz(obstime, obspos, earth_xyz, self._ephemeris) @@ -177,24 +154,29 @@ def evaluate( obs_pos=obs_xyz, ) - # Return a dict of partial functions corresponding to the number density each zodiacal + # If a time and obspos is provided per coordinate, we need to make arrays broadcastable. + if obs_xyz.ndim == 1: + obs_xyz = obs_xyz[:, np.newaxis] + earth_xyz = earth_xyz[:, np.newaxis] + + # Return a dict of partial functions corresponding to the number density of each zodiacal # component in the interplanetary dust model. density_callables = populate_number_density_with_model( comps=self._interplanetary_dust_model.comps, - dynamic_params={"X_earth": earth_xyz[:, np.newaxis, np.newaxis]}, + dynamic_params={"X_earth": earth_xyz}, ) - # Partial function of the brightness integral at a step along the line-of-sight prepopulated - # with shared arguments between zodiacal components. + # Create partial function of the brightness integral at a step along the line-of-sight with + # shared arguments between zodiacal components. common_integrand = functools.partial( self._interplanetary_dust_model.brightness_at_step_callable, - X_obs=obs_xyz[:, np.newaxis, np.newaxis], + X_obs=obs_xyz, bp_interpolation_table=self._b_nu_table, **self._common_parameters, ) - emission = np.zeros((self._interplanetary_dust_model.ncomps, ncoords)) - dist_to_cores = ncoords > nprocesses and nprocesses > 1 + emission = np.zeros((self._interplanetary_dust_model.ncomps, skycoord.size)) + dist_to_cores = skycoord.size > nprocesses and nprocesses > 1 if dist_to_cores: skycoord_xyz_splits = np.array_split(skycoord_xyz, nprocesses, axis=-1) with multiprocessing.get_context(_PLATFORM_METHOD).Pool(nprocesses) as pool: @@ -207,9 +189,9 @@ def evaluate( comp_integrands = [ functools.partial( common_integrand, - u_los=np.expand_dims(vec, axis=-1), - start=np.expand_dims(start, axis=-1), - stop=np.expand_dims(stop, axis=-1), + u_los=vec, + start=start, + stop=stop, get_density_function=density_callables[comp_label], **self._comp_parameters[comp_label], ) @@ -225,27 +207,19 @@ def evaluate( ] emission[idx] = np.concatenate([result.get() for result in proc_chunks]) - # Correct for change of integral limits - emission[idx] *= 0.5 * (stop[comp_label] - start[comp_label]) - else: for idx, comp_label in enumerate(self._interplanetary_dust_model.comps.keys()): comp_integrand = functools.partial( common_integrand, - 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), + u_los=skycoord_xyz, + start=start[comp_label], + stop=stop[comp_label], get_density_function=density_callables[comp_label], **self._comp_parameters[comp_label], ) emission[idx] = integrate_leggauss( comp_integrand, *self._leggauss_points_and_weights ) - # Correct for change of integral limits - emission[idx] *= 0.5 * (stop[comp_label] - start[comp_label]) - - if contains_duplicates: - emission = emission[:, inverse] emission <<= units.MJy / units.sr return emission if return_comps else emission.sum(axis=0) @@ -282,3 +256,29 @@ def update_parameters(self, parameters: dict) -> None: _dict[key] = {ComponentLabel(k): v for k, v in value.items()} self._interplanetary_dust_model = self._interplanetary_dust_model.__class__(**_dict) + + +def validate_user_input(skycoord: coords.SkyCoord, obspos: units.Quantity | str) -> time.Time: + """Validate the shapes and types of the input coordinate information.""" + try: + obstime = typing.cast(time.Time, skycoord.obstime) + except AttributeError as error: + msg = "The input coordinates must be an astropy SkyCoord object." + raise TypeError(msg) from error + if obstime is None: + msg = "The `obstime` attribute of the `SkyCoord` object must be set." + raise ValueError(msg) + + try: + if not isinstance(obspos, str) and obspos.ndim > 1 and obstime.size != obspos.shape[-1]: + msg = "The number of obstime and obspos must match." + raise ValueError(msg) + except AttributeError as error: + msg = "The observer position must be a string or an astropy Quantity." + raise TypeError(msg) from error + + if obstime.size > skycoord.size: + msg = "The number of obstime must be either 1 or the same size as ." + raise ValueError(msg) + + return obstime diff --git a/zodipy/number_density.py b/zodipy/number_density.py index 38fbac9..0b45399 100644 --- a/zodipy/number_density.py +++ b/zodipy/number_density.py @@ -58,7 +58,6 @@ def cloud_number_density( """Density of the diffuse cloud (see Eq (6). in K98).""" X_cloud = X_helio - X_0 R_cloud = np.sqrt(X_cloud[0] ** 2 + X_cloud[1] ** 2 + X_cloud[2] ** 2) - Z_cloud = ( X_cloud[0] * sin_Omega_rad * sin_i_rad - X_cloud[1] * cos_Omega_rad * sin_i_rad diff --git a/zodipy/zodiacal_component.py b/zodipy/zodiacal_component.py index 8212c01..6476dff 100644 --- a/zodipy/zodiacal_component.py +++ b/zodipy/zodiacal_component.py @@ -37,7 +37,7 @@ class ZodiacalComponent(abc.ABC): cos_Omega_rad: float = field(init=False) def __post_init__(self) -> None: - self.X_0 = np.array([self.x_0, self.y_0, self.z_0]).reshape(3, 1, 1) + self.X_0 = np.array([self.x_0, self.y_0, self.z_0]).reshape(3, 1) self.sin_i_rad = np.sin(np.radians(self.i)) self.cos_i_rad = np.cos(np.radians(self.i)) self.sin_Omega_rad = np.sin(np.radians(self.Omega))