Merge pull request #2 from lincc-frameworks/more_complex_model

Allow PhysicalModels to inherit key properties from their host

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()

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)

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