sncosmo simulation function example. This is just a preliminary working function that simulates SN data.

simulation.py

"""Tools for simulation of transients."""

import sys
import math
from copy import deepcopy

import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline as Spline1d
from astropy.table import Table
from astropy import cosmology
from astropy.utils import OrderedDict as odict

import sncosmo

__all__ = ['simulate_vol']

wholesky_sqdeg = 4. * np.pi * (180. / np.pi) ** 2

def simulate_vol(obs_sets, model, gen_params, vrate, obs_area=1.,
                 cosmo=None, z_range=(0., 1.), default_area=1., nsim=None,
                 nret=10, thresh=5.):
    """Simulate transient photometric data according to observations.

    Parameters
    ----------
    obs_sets : dict of `astropy.table.Table`
        A dictionary of "observation sets". Each observation set is a table
        of observations. See the notes section below for information on what
        the table must contain.
    model : `sncosmo.Model`
        The model to use in the simulation.
    gen_params : callable
        A callable that accepts a single float (redshift) and returns
        a dictionary on each call. Typically the callable would randomly
        select parameters from some underlying distribution on each call.
    vrate : callable, optional
        A callable that returns the rate as a function of redshift.
        (The default is 1.e-4.)
    cosmo : astropy.cosmology.Cosmology, optional
        Cosmology used to determine volumes and luminosity distances.
        If `None`, the default, use the current "best" cosmology.
    z_range : (float, float), optional
        Redshift range in which to generate transients.
    default_area : float, optional
        Area in deg^2 for observation sets that do not have an 'AREA'
        keyword in their metadata.
    nsim : int, optional
        Number of transients to simulate. Cannot set both `nsim` and `nret`.
        Default is `None`.
    nret : int, optional
        Number of transients to return (number simulated that pass
        flux significance threshold). Cannot set both `nsim` and `nret`.
        Default is 10. Set both `nsim` and `nret` to `None` to let 
    thresh : float, optional
        Minimum flux significant threshold for a transient to be returned.
    
    Returns
    -------
    transients : list
        List of `~astropy.table.Table`s where each table is the photometric
        data for a single simulated SN.

    Notes
    -----
    
    Each `obs_set` (values in `obs_sets`) must have the following columns:
    
    * ``mjd``
    * ``flt``
    * ``ccd_gain``
    * ``skysig``
    * ``psf1``
    * ``zptavg``

    These are currently just what the SIMLIB files from SNANA have. In the
    future these can be more flexible.
    """

    if nsim is not None and nret is not None:
        raise ValueError('cannot specify both nsim and nret')

    if cosmo is None:
        cosmo = cosmology.get_current()
    model.cosmo = cosmo

    # Get comoving volume in each redshift shell.
    z_bins = 100  # Good enough for now.
    z_min, z_max = z_range
    z_binedges = np.linspace(z_min, z_max, z_bins + 1) 
    z_binctrs = 0.5 * (z_binedges[1:] + z_binedges[:-1]) 
    sphere_vols = cosmo.comoving_volume(z_binedges) 
    shell_vols = sphere_vols[1:] - sphere_vols[:-1]

    # SN / (observer year) in shell
    shell_snrate = shell_vols * vrate(z_binctrs) / (1. + z_binctrs)

    # SN / (observer year) within z_binedges
    vol_snrate = np.zeros_like(z_binedges)
    vol_snrate[1:] = np.add.accumulate(shell_snrate)

    # Create a ppf (inverse cdf). We'll use this later to get
    # a random SN redshift from the distribution.
    snrate_cdf = vol_snrate / vol_snrate[-1]
    snrate_ppf = Spline1d(snrate_cdf, z_binedges, k=1)

    # Get obs sets' data, time ranges, areas and weights.
    # We do this now so we can weight the location of transients 
    # according to the area and time ranges of the observation sets.
    obs_sets = obs_sets.values()
    obs_sets_data = [np.asarray(obs_set) for obs_set in obs_sets]
    time_ranges = [(obs['MJD'].min() - 10. * (1. + z_max),
                    obs['MJD'].max() + 10. * (1. + z_max))
                   for obs in obs_sets_data]
    areas = [obs_set.meta['AREA'] if 'AREA' in obs_set.meta else default_area
             for obs_set in obs_sets]
    area_time_products = [a * (t[1] - t[0])
                          for a, t in zip(areas, time_ranges)]
    total_area_time = sum(area_time_products)
    weights = [a_t / total_area_time for a_t in area_time_products]
    cumweights = np.add.accumulate(np.array(weights))
    
    # How many to simulate?
    if nsim is not None:
        nret = 0
    elif nret is not None:
        nsim = 0
    else:
        nsim = total_area_time / wholesky_sqdeg * vol_snrate[-1]

    i = 0
    transients = []
    while i < nsim or len(transients) < nret:
        i += 1

        # which obs_set did this occur in?
        j = 0
        x = np.random.rand()
        while cumweights[j] < x:
            j += 1

        obsdata = obs_sets_data[j]
        time_range = time_ranges[j]

        # Get a redshift from the distribution
        z = snrate_ppf(np.random.rand())
        t0 = np.random.uniform(time_range[0], time_range[1])

        # Get rest of parameters from user-defined gen_params():
        params = gen_params(z)
        params.update(z=z, t0=t0)
        model.set(**params)

        # Get model fluxes 
        flux = model.bandflux(obsdata['FLT'], obsdata['MJD'],
                              zp=obsdata['ZPTAVG'], zpsys='ab')

        # Get flux errors
        noise_area = 4. * math.pi * obsdata['PSF1']
        bkgpixnoise = obsdata['SKYSIG']
        fluxerr = np.sqrt((noise_area * bkgpixnoise) ** 2 +
                          np.abs(flux) / obsdata['CCD_GAIN'])

        # Scatter fluxes by the fluxerr
        flux = np.random.normal(flux, fluxerr)

        # Check if any of the fluxes are significant
        if not np.any((flux / fluxerr) > thresh):
            continue

        simulated_data = odict([('date', obsdata['MJD']),
                                ('band', obsdata['FLT']),
                                ('flux', flux),
                                ('fluxerr', fluxerr),
                                ('zp', obsdata['ZPTAVG']),
                                ('zpsys', ['ab'] * len(flux))])

        params.update(modelname=model.name)
        transients.append(Table(simulated_data, meta=params))

    return transients


if __name__ == '__main__':
    
    # Read in observation set data
    meta, obs_sets = sncosmo.read_snana_simlib(
        '../sncosmo-examples/data/DES_hybrid_10field_griz.SIMLIB')

    # The observation sets have bands called 'g', 'r', 'i', 'z'. 
    # These correspond to sncosmo built-in bands 'desg', 'desr', etc.
    # Register the built-in bands under the names 'g', 'r', etc.
    for regname, newname in [('desg', 'g'), ('desr', 'r'), ('desi', 'i'),
                             ('desz', 'z')]:
        band = sncosmo.get_bandpass(regname)
        band.name = newname
        sncosmo.registry.register(band)

    # Get the model
    model = sncosmo.get_model('salt2-extended')

    # Define parmaeter generating function. 
    # This specific one is a normally distributed x1 and c with
    # a normally distributed hubble residual. It is not redshift independent.
    alpha = 0.13
    beta = 2.5
    def gen_params(z):
        x1 = np.random.normal(0., 1.)
        c = np.random.normal(0., 0.1)
        mabs = -19.3 - alpha * x1 + beta * c + np.random.normal(0., 0.15)
        return {'c': c, 'x1': x1, 'mabs': mabs}

    # Define a volumetric SN rate in SN / yr / Mpc^3
    def snrate(z):
        return 0.25e-4 * (1. + 2.5 * z)

    # Generate simulated SNe.
    sne = simulate_vol(obs_sets, model, gen_params, snrate, nret=10)

    # save first sn to a file
    sncosmo.write_lc(np.asarray(sne[0]), 'sn0.csv', meta=sne[0].meta, fmt='csv')