diff --git a/src/tdastro/base_models.py b/src/tdastro/base_models.py index a43319d2..0bd53b07 100644 --- a/src/tdastro/base_models.py +++ b/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) diff --git a/src/tdastro/sources/galaxy_models.py b/src/tdastro/sources/galaxy_models.py new file mode 100644 index 00000000..9aea8bfc --- /dev/null +++ b/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) diff --git a/tests/tdastro/sources/test_galaxy_models.py b/tests/tdastro/sources/test_galaxy_models.py new file mode 100644 index 00000000..e127494d --- /dev/null +++ b/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