Merge pull request #1 from lincc-frameworks/jeremys_proposal

An initial framework for sources

pyproject.toml

@@ -16,6 +16,7 @@ classifiers = [
 dynamic = ["version"]
 requires-python = ">=3.9"
 dependencies = [
+    "numpy",
 ]
 
 [project.urls]
@@ -86,7 +87,6 @@ select = [
     "D300",
     "D417",
     "D419",
-
 ]
 
 ignore = [
@@ -97,6 +97,7 @@ ignore = [
     "SIM117", # Allow nested with
     "UP015", # Allow redundant open parameters
     "UP028", # Allow yield in for loop
+    "E721", # Allow direct type comparison
 ]
 
 [tool.coverage.run]

src/tdastro/__init__.py

@@ -1,3 +0,0 @@
-from .example_module import greetings, meaning
-
-__all__ = ["greetings", "meaning"]

src/tdastro/base_models.py

@@ -0,0 +1,126 @@
+class PhysicalModel:
+    """A physical model of a source of flux.
+
+    Attributes
+    ----------
+    host : `PhysicalModel`
+        A physical model of the current source's host.
+    effects : `list`
+        A list of effects to apply to an observations.
+    """
+
+    def __init__(self, host=None, **kwargs):
+        """Create a PhysicalModel object.
+
+        Parameters
+        ----------
+        host : `PhysicalModel`, optional
+            A physical model of the current source's host.
+        **kwargs : `dict`, optional
+           Any additional keyword arguments.
+        """
+        self.host = host
+        self.effects = []
+
+    def add_effect(self, effect):
+        """Add a transformational effect to the PhysicalModel.
+        Effects are applied in the order in which they are added.
+
+        Parameters
+        ----------
+        effect : `EffectModel`
+            The effect to apply.
+
+        Raises
+        ------
+        Raises a ``AttributeError`` if the PhysicalModel does not have all of the
+        required attributes.
+        """
+        required: list = effect.required_parameters()
+        for parameter in required:
+            # Raise an AttributeError if the parameter is missing or set to None.
+            if getattr(self, parameter) is None:
+                raise AttributeError(f"Parameter {parameter} unset for model {type(self).__name__}")
+
+        self.effects.append(effect)
+
+    def _evaluate(self, times, bands=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.
+        **kwargs : `dict`, optional
+           Any additional keyword arguments.
+
+        Returns
+        -------
+        flux_density : `numpy.ndarray`
+            The results.
+        """
+        raise NotImplementedError()
+
+    def evaluate(self, times, bands=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.
+        **kwargs : `dict`, optional
+           Any additional keyword arguments.
+
+        Returns
+        -------
+        flux_density : `numpy.ndarray`
+            The results.
+        """
+        flux_density = self._evaluate(times, bands, **kwargs)
+        for effect in self.effects:
+            flux_density = effect.apply(flux_density, bands, self, **kwargs)
+        return flux_density
+
+
+class EffectModel:
+    """A physical or systematic effect to apply to an observation."""
+
+    def __init__(self, **kwargs):
+        pass
+
+    def required_parameters(self):
+        """Returns a list of the parameters of a PhysicalModel
+        that this effect needs to access.
+
+        Returns
+        -------
+        parameters : `list` of `str`
+            A list of every required parameter the effect needs.
+        """
+        return []
+
+    def apply(self, flux_density, bands=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.
+        physical_model : `PhysicalModel`
+            A PhysicalModel from which the effect may query parameters
+            such as redshift, position, or distance.
+        **kwargs : `dict`, optional
+           Any additional keyword arguments.
+
+        Returns
+        -------
+        flux_density : `numpy.ndarray`
+            The results.
+        """
+        raise NotImplementedError()

src/tdastro/effects/white_noise.py

@@ -0,0 +1,48 @@
+import numpy as np
+
+from tdastro.base_models import EffectModel
+
+
+class WhiteNoise(EffectModel):
+    """A white noise model.
+
+    Attributes
+    ----------
+    scale : `float`
+        The scale of the noise.
+    """
+
+    def __init__(self, scale, **kwargs):
+        """Create a WhiteNoise effect model.
+
+        Parameters
+        ----------
+        scale : `float`
+            The scale of the noise.
+        **kwargs : `dict`, optional
+           Any additional keyword arguments.
+        """
+        super().__init__(**kwargs)
+        self.scale = scale
+
+    def apply(self, flux_density, bands=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.
+        physical_model : `PhysicalModel`
+            A PhysicalModel from which the effect may query parameters
+            such as redshift, position, or distance.
+        **kwargs : `dict`, optional
+           Any additional keyword arguments.
+
+        Returns
+        -------
+        flux_density : `numpy.ndarray`
+            The results.
+        """
+        return np.random.normal(loc=flux_density, scale=self.scale)

src/tdastro/sources/static_source.py

@@ -0,0 +1,56 @@
+import types
+
+import numpy as np
+
+from tdastro.base_models import PhysicalModel
+
+
+class StaticSource(PhysicalModel):
+    """A static source.
+
+    Attributes
+    ----------
+    brightness : `float`
+        The inherent brightness
+    """
+
+    def __init__(self, brightness, **kwargs):
+        """Create a StaticSource object.
+
+        Parameters
+        ----------
+        brightness : `float`, `function`, or `None`
+            The inherent brightness
+        **kwargs : `dict`, optional
+           Any additional keyword arguments.
+        """
+        super().__init__(**kwargs)
+
+        if brightness is None:
+            # If we were not given the parameter, use a default sampling function.
+            self.brightness = np.random.rand(10.0, 20.0)
+        elif isinstance(brightness, types.FunctionType):
+            # If we were given a sampling function, use it.
+            self.brightness = brightness(**kwargs)
+        else:
+            # Otherwise assume we were given the parameter itself.
+            self.brightness = brightness
+
+    def _evaluate(self, times, bands=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.
+        **kwargs : `dict`, optional
+           Any additional keyword arguments.
+
+        Returns
+        -------
+        flux : `numpy.ndarray`
+           The results.
+        """
+        return np.full_like(times, self.brightness)

tests/tdastro/test_gaussian_noise.py

@@ -0,0 +1,19 @@
+import numpy as np
+from tdastro.effects.white_noise import WhiteNoise
+from tdastro.sources.static_source import StaticSource
+
+
+def brightness_generator():
+    """A test generator function."""
+    return 10.0 + 0.5 * np.random.rand(1)
+
+
+def test_white_noise() -> None:
+    """Test that we can sample and create a WhiteNoise object."""
+    model = StaticSource(brightness=brightness_generator)
+    model.add_effect(WhiteNoise(scale=0.01))
+
+    values = model.evaluate(np.array([1, 2, 3, 4, 5]))
+    assert len(values) == 5
+    assert not np.all(values == 10.0)
+    assert np.all(np.abs(values - 10.0) < 1.0)