Merge pull request #31 from lincc-frameworks/galaxies

Create the ability to sample from a background model

src/tdastro/base_models.py

@@ -191,9 +191,13 @@ def sample_parameters(self, max_depth=50, **kwargs):
 
 
 class PhysicalModel(ParameterizedModel):
-    """A physical model of a source of flux. Physical models can have fixed attributes
-    (where you need to create a new model to change them) and settable attributes that
-    can be passed functions or constants.
+    """A physical model of a source of flux.
+
+    Physical models can have fixed attributes (where you need to create a new model
+    to change them) and settable attributes that can be passed functions or constants.
+    They can also have special background pointers that link to another PhysicalModel
+    producing flux. We can chain these to have a supernova in front of a star in front
+    of a static background.
 
     Attributes
     ----------
@@ -203,11 +207,13 @@ class PhysicalModel(ParameterizedModel):
         The object's declination (in degrees)
     distance : `float`
         The object's distance (in pc)
+    background : `PhysicalModel`
+        A source of background flux such as a host galaxy.
     effects : `list`
         A list of effects to apply to an observations.
     """
 
-    def __init__(self, ra=None, dec=None, distance=None, **kwargs):
+    def __init__(self, ra=None, dec=None, distance=None, background=None, **kwargs):
         super().__init__(**kwargs)
         self.effects = []
 
@@ -216,6 +222,9 @@ def __init__(self, ra=None, dec=None, distance=None, **kwargs):
         self.add_parameter("dec", dec)
         self.add_parameter("distance", distance)
 
+        # Background is an object not a sampled parameter
+        self.background = background
+
     def __str__(self):
         """Return the string representation of the model."""
         return "PhysicalModel"
@@ -296,7 +305,12 @@ def evaluate(self, times, wavelengths, resample_parameters=False, **kwargs):
         if resample_parameters:
             self.sample_parameters(kwargs)
 
+        # Compute the flux density for both the current object and add in anything
+        # behind it, such as a host galaxy.
         flux_density = self._evaluate(times, wavelengths, **kwargs)
+        if self.background is not None:
+            flux_density += self.background._evaluate(times, wavelengths, ra=self.ra, dec=self.dec, **kwargs)
+
         for effect in self.effects:
             flux_density = effect.apply(flux_density, wavelengths, self, **kwargs)
         return flux_density
@@ -313,7 +327,10 @@ def sample_parameters(self, include_effects=True, **kwargs):
             All the keyword arguments, including the values needed to sample
             parameters.
         """
+        if self.background is not None:
+            self.background.sample_parameters(include_effects, **kwargs)
         super().sample_parameters(**kwargs)
+
         if include_effects:
             for effect in self.effects:
                 effect.sample_parameters(**kwargs)

src/tdastro/sources/galaxy_models.py

@@ -0,0 +1,83 @@
+import numpy as np
+from astropy.coordinates import angular_separation
+
+from tdastro.base_models import PhysicalModel
+
+
+class GaussianGalaxy(PhysicalModel):
+    """A static source.
+
+    Attributes
+    ----------
+    radius_std : `float`
+        The standard deviation of the brightness as we move away
+        from the galaxy's center (in degrees).
+    brightness : `float`
+        The inherent brightness at the center of the galaxy.
+    """
+
+    def __init__(self, brightness, radius, **kwargs):
+        super().__init__(**kwargs)
+        self.add_parameter("galaxy_radius_std", radius, required=True, **kwargs)
+        self.add_parameter("brightness", brightness, required=True, **kwargs)
+
+    def __str__(self):
+        """Return the string representation of the model."""
+        return f"GuassianGalaxy({self.brightness}, {self.galaxy_radius_std})"
+
+    def sample_ra(self):
+        """Sample an right ascension coordinate based on the center and radius of the galaxy.
+
+        Returns
+        -------
+        ra : `float`
+            The sampled right ascension in degrees.
+        """
+        return np.random.normal(loc=self.ra, scale=self.galaxy_radius_std)
+
+    def sample_dec(self):
+        """Sample a declination coordinate based on the center and radius of the galaxy.
+
+        Returns
+        -------
+        dec : `float`
+            The sampled declination in degrees.
+        """
+        return np.random.normal(loc=self.dec, scale=self.galaxy_radius_std)
+
+    def _evaluate(self, times, wavelengths, ra=None, dec=None, **kwargs):
+        """Draw effect-free observations for this object.
+
+        Parameters
+        ----------
+        times : `numpy.ndarray`
+            A length T array of timestamps.
+        wavelengths : `numpy.ndarray`, optional
+            A length N array of wavelengths.
+        ra : `float`, optional
+            The right ascension of the observations in degrees.
+        dec : `float`, optional
+            The declination of the observations in degrees.
+        **kwargs : `dict`, optional
+           Any additional keyword arguments.
+
+        Returns
+        -------
+        flux_density : `numpy.ndarray`
+            A length T x N matrix of SED values.
+        """
+        if ra is None:
+            ra = self.ra
+        if dec is None:
+            dec = self.dec
+
+        # Scale the brightness as a Guassian function centered on the object's RA and Dec.
+        dist = angular_separation(
+            self.ra * np.pi / 180.0,
+            self.dec * np.pi / 180.0,
+            ra * np.pi / 180.0,
+            dec * np.pi / 180.0,
+        )
+        scale = np.exp(-(dist * dist) / (2.0 * self.galaxy_radius_std * self.galaxy_radius_std))
+
+        return np.full((len(times), len(wavelengths)), self.brightness * scale)

tests/tdastro/sources/test_galaxy_models.py

@@ -0,0 +1,68 @@
+import random
+
+import numpy as np
+from tdastro.sources.galaxy_models import GaussianGalaxy
+from tdastro.sources.static_source import StaticSource
+
+
+def _sample_ra(**kwargs):
+    """Return a random value between 0 and 360.
+
+    Parameters
+    ----------
+    **kwargs : `dict`, optional
+        Absorbs additional parameters
+    """
+    return 360.0 * random.random()
+
+
+def _sample_dec(**kwargs):
+    """Return a random value between -90 and 90.
+
+    Parameters
+    ----------
+    **kwargs : `dict`, optional
+        Absorbs additional parameters
+    """
+    return 180.0 * random.random() - 90.0
+
+
+def test_gaussian_galaxy() -> None:
+    """Test that we can sample and create a StaticSource object."""
+    random.seed(1001)
+
+    host = GaussianGalaxy(ra=_sample_ra, dec=_sample_dec, brightness=10.0, radius=1.0 / 3600.0)
+    host_ra = host.ra
+    host_dec = host.dec
+
+    source = StaticSource(ra=host.sample_ra, dec=host.sample_dec, background=host, brightness=100.0)
+
+    # Both RA and dec should be "close" to (but not exactly at) the center of the galaxy.
+    source_ra_offset = source.ra - host_ra
+    assert 0.0 < np.abs(source_ra_offset) < 100.0 / 3600.0
+
+    source_dec_offset = source.dec - host_dec
+    assert 0.0 < np.abs(source_dec_offset) < 100.0 / 3600.0
+
+    times = np.array([1, 2, 3, 4, 5, 10])
+    wavelengths = np.array([100.0, 200.0, 300.0])
+
+    # All the measured fluxes should have some contribution from the background object.
+    values = source.evaluate(times, wavelengths)
+    assert values.shape == (6, 3)
+    assert np.all(values > 100.0)
+    assert np.all(values <= 110.0)
+
+    # Check that if we resample the source it will propagate and correctly resample the host.
+    # the host's (RA, dec) should change and the source's should still be close.
+    source.sample_parameters()
+    assert host_ra != host.ra
+    assert host_dec != host.dec
+
+    source_ra_offset2 = source.ra - host.ra
+    assert source_ra_offset != source_ra_offset2
+    assert 0.0 < np.abs(source_ra_offset2) < 100.0 / 3600.0
+
+    source_dec_offset2 = source.dec - host.dec
+    assert source_dec_offset != source_dec_offset2
+    assert 0.0 < np.abs(source_ra_offset2) < 100.0 / 3600.0