Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow PhysicalModels to inherit key properties from their host #2

Merged
merged 2 commits into from
Jun 21, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 43 additions & 18 deletions src/tdastro/base_models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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


Expand All @@ -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.
Expand All @@ -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()
12 changes: 6 additions & 6 deletions src/tdastro/sources/static_source.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
26 changes: 26 additions & 0 deletions tests/tdastro/sources/test_static_source.py
Original file line number Diff line number Diff line change
@@ -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
Loading