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

An initial framework for sources #1

Merged
merged 7 commits into from
Jun 21, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
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
Prev Previous commit
Next Next commit
ruff formatting
  • Loading branch information
jeremykubica committed Jun 21, 2024
commit a20273a0469508572ca674a20a279ea08440a4a2
195 changes: 98 additions & 97 deletions src/tdastro/base_models.py
Original file line number Diff line number Diff line change
@@ -1,103 +1,104 @@
import abc


class PhysicalModel(abc.ABC):
def __init__(self, ra=None, dec=None, distance=None, **kwargs):
self.ra = ra
self.dec = dec
self.distance = distance
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.

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.

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
def __init__(self, ra=None, dec=None, distance=None, **kwargs):
self.ra = ra
self.dec = dec
self.distance = distance
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.

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.

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(abc.ABC):
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.

Returns
-------
flux_density : `numpy.ndarray`
The results.
"""
raise NotImplementedError()
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.

Returns
-------
flux_density : `numpy.ndarray`
The results.
"""
raise NotImplementedError()
103 changes: 52 additions & 51 deletions src/tdastro/effects/white_noise.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,66 +3,67 @@

from tdastro.base_models import EffectModel, PhysicalModel


class WhiteNoise(EffectModel):
def __init__(self, scale, **kwargs):
super().__init__(**kwargs)
self.scale = scale
def __init__(self, scale, **kwargs):
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)
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.
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.

Returns
-------
flux_density : `numpy.ndarray`
The results.
"""
return np.random.normal(loc=flux_density, scale=self.scale)
Returns
-------
flux_density : `numpy.ndarray`
The results.
"""
return np.random.normal(loc=flux_density, scale=self.scale)


class DistanceBasedWhiteNoise(EffectModel):
def __init__(self, scale, dist_multiplier, **kwargs):
super().__init__(**kwargs)
self.scale = scale
self.dist_multiplier = dist_multiplier
def __init__(self, scale, dist_multiplier, **kwargs):
super().__init__(**kwargs)
self.scale = scale
self.dist_multiplier = dist_multiplier

def required_parameters(self):
"""Returns a list of the parameters of a PhysicalModel
that this effect needs to access.
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 ['distance']
Returns
-------
parameters : `list` of `str`
A list of every required parameter the effect needs.
"""
return ["distance"]

def apply(self, flux_density, bands=None, physical_model=None, **kwargs):
"""Apply the effect to observations (flux_density values)
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.
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.

Returns
-------
flux_density : `numpy.ndarray`
The results.
"""
scale_value = self.scale + self.dist_multiplier * physical_model.distance
return np.random.normal(loc=flux_density, scale=scale_value)
Returns
-------
flux_density : `numpy.ndarray`
The results.
"""
scale_value = self.scale + self.dist_multiplier * physical_model.distance
return np.random.normal(loc=flux_density, scale=scale_value)
52 changes: 26 additions & 26 deletions src/tdastro/sources/static_source.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,34 +3,34 @@

from tdastro.base_models import PhysicalModel


class StaticSource(PhysicalModel):
def __init__(self, brightness, **kwargs):
super().__init__(**kwargs)
def __init__(self, brightness, **kwargs):
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 type(brightness) is 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
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 type(brightness) is 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.
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.
Parameters
----------
times : `numpy.ndarray`
An array of timestamps.
bands : `numpy.ndarray`, optional
An array of bands. If ``None`` then does something.

Returns
-------
flux : `numpy.ndarray`
The results.
"""
return np.full_like(times, self.brightness)

Returns
-------
flux : `numpy.ndarray`
The results.
"""
return np.full_like(times, self.brightness)
Loading
Loading