Add a helper function for assembling a snia graph

src/tdastro/astro_utils/snia_utils.py

@@ -5,6 +5,7 @@
 from scipy.stats.sampling import NumericalInversePolynomial
 
 from tdastro.base_models import FunctionNode
+from tdastro.math_nodes.np_random import NumpyRandomFunc
 from tdastro.math_nodes.scipy_random import NumericalInversePolynomialFunc
 
 
@@ -320,3 +321,58 @@ def _distmod_from_redshift(self, redshift):
             The distance modulus (in mag)
         """
         return self.cosmo.distmod(redshift).value
+
+
+def snia_x0_x1_from_host(
+    host,
+    H0=73.0,
+    Omega_m=0.3,
+    alpha=0.14,
+    beta=3.1,
+    c_func=None,
+    m_abs_func=None,
+):
+    """Constructs multiple interdependent FunctionNodes to generate
+    the parameters for an SNIA supernova.
+
+    Parameters
+    ----------
+    host : PhysicalModel
+        The PhysicalModel for the host, including information on redshift and hostmass.
+    H0 : constant
+        The Hubble constant. Default: 73.0
+    Omega_m : constant
+        The matter density Omega_m. Default: 0.3
+    alpha : function or constant
+        The alpha parameter in the Tripp relation. Default: 0.14
+    beta : function or constant
+        The beta parameter in the Tripp relation. Default: 3.1
+    c_func : function, constant, or None, optional
+        The function or constant providing the c value. If None then samples
+        from a normal distribution with loc=0, scale=0.02. Default: None
+    m_abs_func : function, constant, or None, optional
+        The function or constant providing the m_abs value. If None then samples
+        from a normal distribution with loc=-19.3, scale=0.1. Default: None
+
+    Returns
+    -------
+    x0_func : X0FromDistMod
+        The FunctionNode for computing the x0 parameter from the host's distmod.
+    x1_func : HostmassX1Func
+        The FunctionNode for computing the x1 parameter from the host's mass.
+    """
+    distmod_func = DistModFromRedshift(host.redshift, H0=H0, Omega_m=Omega_m, node_label="distmod_func")
+    x1_func = HostmassX1Func(host.hostmass, node_label="x1_func")
+    c_func = c_func if c_func is not None else NumpyRandomFunc("normal", loc=0, scale=0.02)
+    m_abs_func = m_abs_func if m_abs_func is not None else NumpyRandomFunc("normal", loc=-19.3, scale=0.1)
+
+    x0_func = X0FromDistMod(
+        distmod=distmod_func,
+        x1=x1_func,
+        c=c_func,
+        alpha=alpha,
+        beta=beta,
+        m_abs=m_abs_func,
+        node_label="x0_func",
+    )
+    return x0_func, x1_func

src/tdastro/example_runs/simulate_snia.py

@@ -8,10 +8,8 @@
 from tdastro.astro_utils.noise_model import apply_noise
 from tdastro.astro_utils.passbands import PassbandGroup
 from tdastro.astro_utils.snia_utils import (
-    DistModFromRedshift,
-    HostmassX1Func,
-    X0FromDistMod,
     num_snia_per_redshift_bin,
+    snia_x0_x1_from_host,
 )
 from tdastro.astro_utils.unit_utils import flam_to_fnu, fnu_to_flam
 from tdastro.math_nodes.np_random import NumpyRandomFunc
@@ -40,8 +38,7 @@ def construct_snia_source(oversampled_observations, zpdf):
     logger.info("Creating the source model.")
 
     # Get the range in which t0 can occur.
-    t_min = oversampled_observations["observationStartMJD"].min()
-    t_max = oversampled_observations["observationStartMJD"].max()
+    t_min, t_max = oversampled_observations.time_bounds()
 
     # TODO: Extract the ra and dec center from the opsim in case that changes.
     # Currently we are relying on the fact the test opsim has all pointings
@@ -56,18 +53,15 @@ def construct_snia_source(oversampled_observations, zpdf):
         node_label="host",
     )
 
-    distmod_func = DistModFromRedshift(host.redshift, H0=73.0, Omega_m=0.3)
-    x1_func = HostmassX1Func(host.hostmass)
     c_func = NumpyRandomFunc("normal", loc=0, scale=0.02)
-    m_abs_func = NumpyRandomFunc("normal", loc=-19.3, scale=0.1)
-    x0_func = X0FromDistMod(
-        distmod=distmod_func,
-        x1=x1_func,
-        c=c_func,
+    x0_func, x1_func = snia_x0_x1_from_host(
+        host,
+        H0=73.0,
+        Omega_m=0.3,
         alpha=0.14,
         beta=3.1,
-        m_abs=m_abs_func,
-        node_label="x0_func",
+        c_func=c_func,
+        m_abs_func=NumpyRandomFunc("normal", loc=-19.3, scale=0.1),
     )
 
     sncosmo_modelname = "salt3"

tests/tdastro/astro_utils/test_snia_utils.py

@@ -1,8 +1,15 @@
 import numpy as np
 import pytest
 from scipy.stats import norm
-from tdastro.astro_utils.snia_utils import DistModFromRedshift, HostmassX1Distr, HostmassX1Func
+from tdastro.astro_utils.snia_utils import (
+    DistModFromRedshift,
+    HostmassX1Distr,
+    HostmassX1Func,
+    X0FromDistMod,
+    snia_x0_x1_from_host,
+)
 from tdastro.math_nodes.np_random import NumpyRandomFunc
+from tdastro.sources.snia_host import SNIaHost
 
 
 def test_hostmass_x1_distr():
@@ -107,3 +114,50 @@ def test_dist_mod_from_redshift():
         node = DistModFromRedshift(redshift=z, H0=73.0, Omega_m=0.3)
         state = node.sample_parameters(num_samples=1)
         assert node.get_param(state, "function_node_result") == pytest.approx(expected[idx])
+
+
+def test_snia_x0_x1_from_host():
+    """Test that we can assemble the x0 and x1 functions from the host data."""
+    static_host = SNIaHost(
+        ra=0.0,
+        dec=0.0,
+        hostmass=8.0,
+        redshift=0.3,
+        node_label="host",
+    )
+
+    # Manually create x0_func and x1_func.
+    distmod_func = DistModFromRedshift(static_host.redshift, H0=73.0, Omega_m=0.3)
+    x1_func_a = HostmassX1Func(static_host.hostmass, node_label="x1_func")
+    c_func = NumpyRandomFunc("normal", loc=0, scale=0.02)
+    m_abs_func = NumpyRandomFunc("normal", loc=-19.3, scale=0.1)
+    x0_func_a = X0FromDistMod(
+        distmod=distmod_func,
+        x1=x1_func_a,
+        c=c_func,
+        alpha=0.14,
+        beta=3.1,
+        m_abs=m_abs_func,
+        node_label="x0_func",
+    )
+
+    # Sample using a random number generator with a specified seed (for consistency).
+    rng_a = np.random.default_rng(seed=100)
+    state_a = x0_func_a.sample_parameters(rng_info=rng_a)
+
+    # Create the functions with snia_x0_x1_from_host().
+    x0_func_b, _ = snia_x0_x1_from_host(
+        static_host,
+        H0=73.0,
+        Omega_m=0.3,
+        alpha=0.14,
+        beta=3.1,
+    )
+
+    # Sample using a random number generator with the SAME specified seed.
+    rng_b = np.random.default_rng(seed=100)
+    state_b = x0_func_b.sample_parameters(rng_info=rng_b)
+
+    # Test that we get the same results.
+    assert state_a["x1_func"]["function_node_result"] == state_b["x1_func"]["function_node_result"]
+    assert state_a["x0_func"]["function_node_result"] == state_b["x0_func"]["function_node_result"]