diff --git a/src/tdastro/base_models.py b/src/tdastro/base_models.py index 0afd7b9e..8e77ecf2 100644 --- a/src/tdastro/base_models.py +++ b/src/tdastro/base_models.py @@ -5,21 +5,46 @@ class PhysicalModel: ---------- host : `PhysicalModel` A physical model of the current source's host. + ra : `float` + The object's right ascension (in degrees) + dec : `float` + The object's declination (in degrees) + distance : `float` + The object's distance (in au) effects : `list` A list of effects to apply to an observations. """ - def __init__(self, host=None, **kwargs): + def __init__(self, host=None, ra=None, dec=None, distance=None, **kwargs): """Create a PhysicalModel object. Parameters ---------- host : `PhysicalModel`, optional A physical model of the current source's host. + ra : `float`, optional + The object's right ascension (in degrees) + dec : `float`, optional + The object's declination (in degrees) + distance : `float`, optional + The object's distance (in au) **kwargs : `dict`, optional Any additional keyword arguments. """ self.host = host + + # Set RA, dec, and distance from the given value or, if it is None and + # we have a host, from the host's value. + self.ra = ra + self.dec = dec + self.distance = distance + if ra is None and host is not None: + self.ra = host.ra + if dec is None and host is not None: + self.dec = host.dec + if distance is None and host is not None: + self.distance = host.distance + self.effects = [] def add_effect(self, effect): @@ -44,45 +69,45 @@ def add_effect(self, effect): self.effects.append(effect) - def _evaluate(self, times, bands=None, **kwargs): + def _evaluate(self, times, wavelengths=None, **kwargs): """Draw effect-free observations for this object. Parameters ---------- times : `numpy.ndarray` - An array of timestamps. - bands : `numpy.ndarray`, optional - An array of bands. + A length N array of timestamps. + wavelengths : `numpy.ndarray`, optional + A length N array of wavelengths. **kwargs : `dict`, optional Any additional keyword arguments. Returns ------- flux_density : `numpy.ndarray` - The results. + A length N-array of flux densities. """ raise NotImplementedError() - def evaluate(self, times, bands=None, **kwargs): + def evaluate(self, times, wavelengths=None, **kwargs): """Draw observations for this object and apply the noise. Parameters ---------- times : `numpy.ndarray` - An array of timestamps. - bands : `numpy.ndarray`, optional - An array of bands. + A length N array of timestamps. + wavelengths : `numpy.ndarray`, optional + A length N array of wavelengths. **kwargs : `dict`, optional Any additional keyword arguments. Returns ------- flux_density : `numpy.ndarray` - The results. + A length N-array of flux densities. """ - flux_density = self._evaluate(times, bands, **kwargs) + flux_density = self._evaluate(times, wavelengths, **kwargs) for effect in self.effects: - flux_density = effect.apply(flux_density, bands, self, **kwargs) + flux_density = effect.apply(flux_density, wavelengths, self, **kwargs) return flux_density @@ -103,15 +128,15 @@ def required_parameters(self): """ return [] - def apply(self, flux_density, bands=None, physical_model=None, **kwargs): + def apply(self, flux_density, wavelengths=None, physical_model=None, **kwargs): """Apply the effect to observations (flux_density values) Parameters ---------- flux_density : `numpy.ndarray` - An array of flux density values. - bands : `numpy.ndarray`, optional - An array of bands. + A length N array of flux density values. + wavelengths : `numpy.ndarray`, optional + A length N array of wavelengths. physical_model : `PhysicalModel` A PhysicalModel from which the effect may query parameters such as redshift, position, or distance. @@ -121,6 +146,6 @@ def apply(self, flux_density, bands=None, physical_model=None, **kwargs): Returns ------- flux_density : `numpy.ndarray` - The results. + A length N-array of flux densities after the effect is applied. """ raise NotImplementedError() diff --git a/src/tdastro/sources/static_source.py b/src/tdastro/sources/static_source.py index 99a2f081..1560e514 100644 --- a/src/tdastro/sources/static_source.py +++ b/src/tdastro/sources/static_source.py @@ -36,21 +36,21 @@ def __init__(self, brightness, **kwargs): # Otherwise assume we were given the parameter itself. self.brightness = brightness - def _evaluate(self, times, bands=None, **kwargs): + def _evaluate(self, times, wavelengths=None, **kwargs): """Draw effect-free observations for this object. Parameters ---------- times : `numpy.ndarray` - An array of timestamps. - bands : `numpy.ndarray`, optional - An array of bands. If ``None`` then does something. + A length N array of timestamps. + wavelengths : `numpy.ndarray`, optional + A length N array of wavelengths. **kwargs : `dict`, optional Any additional keyword arguments. Returns ------- - flux : `numpy.ndarray` - The results. + flux_density : `numpy.ndarray` + A length N-array of flux densities. """ return np.full_like(times, self.brightness) diff --git a/tests/tdastro/test_gaussian_noise.py b/tests/tdastro/effects/test_white_noise.py similarity index 100% rename from tests/tdastro/test_gaussian_noise.py rename to tests/tdastro/effects/test_white_noise.py diff --git a/tests/tdastro/sources/test_static_source.py b/tests/tdastro/sources/test_static_source.py new file mode 100644 index 00000000..334b4d09 --- /dev/null +++ b/tests/tdastro/sources/test_static_source.py @@ -0,0 +1,26 @@ +import numpy as np +from tdastro.sources.static_source import StaticSource + + +def test_static_source() -> None: + """Test that we can sample and create a StaticSource object.""" + model = StaticSource(brightness=10.0) + assert model.brightness == 10.0 + assert model.ra is None + assert model.dec is None + assert model.distance is None + + values = model.evaluate(np.array([1, 2, 3, 4, 5, 10])) + assert len(values) == 6 + assert np.all(values == 10.0) + + +def test_static_source_host() -> None: + """Test that we can sample and create a StaticSource object with properties + derived from the host object.""" + host = StaticSource(brightness=15.0, ra=1.0, dec=2.0, distance=3.0) + model = StaticSource(brightness=10.0, host=host) + assert model.brightness == 10.0 + assert model.ra == 1.0 + assert model.dec == 2.0 + assert model.distance == 3.0