gammapy / gammapy

gammapy/astro/population/simulate.py

# Licensed under a 3-clause BSD style license - see LICENSE.rst
"""Simulate source catalogs."""
from __future__ import absolute_import, division, print_function, unicode_literals
import numpy as np
from ...extern import six
from astropy.table import Table, Column

external import "from astropy.table import Table, Column" should be placed before "from ...extern import six"

from astropy.units import Quantity

external import "from astropy.units import Quantity" should be placed before "from ...extern import six"

from astropy.coordinates import SkyCoord, spherical_to_cartesian

external import "from astropy.coordinates import SkyCoord, spherical_to_cartesian" should be placed before "from ...extern import six"

from ...utils import coordinates as astrometry
from ...utils.coordinates import D_SUN_TO_GALACTIC_CENTER
from ...utils.distributions import draw, pdf
from ...utils.random import sample_sphere, sample_sphere_distance, get_random_state
from ..source import SNR, SNRTrueloveMcKee, PWN, Pulsar
from ..population.spatial import (
    Exponential,
    FaucherSpiral,
    RMIN,
    RMAX,
    ZMIN,
    ZMAX,
    radial_distributions,
)
from ..population.velocity import VMIN, VMAX, velocity_distributions

__all__ = [
    "make_catalog_random_positions_cube",
    "make_catalog_random_positions_sphere",
    "make_base_catalog_galactic",
    "add_snr_parameters",
    "add_pulsar_parameters",
    "add_pwn_parameters",
    "add_observed_source_parameters",
    "add_observed_parameters",
]


def make_catalog_random_positions_cube(
    size=100, dimension=3, dmax=10, random_state="random-seed"

Wrong hanging indentation before block (add 4 spaces).

):
    """Make a catalog of sources randomly distributed on a line, square or cube.

    TODO: is this useful enough for general use or should we hide it as an
      internal method to generate test datasets?

    Parameters
    ----------
    size : int, optional
        Number of sources
    dimension : int, optional
        Number of dimensions
    dmax : int, optional
        Maximum distance in pc.
    random_state : {int, 'random-seed', 'global-rng', `~numpy.random.RandomState`}
        Defines random number generator initialisation.
        Passed to `~gammapy.utils.random.get_random_state`.

    Returns
    -------
    catalog : `~astropy.table.Table`
        Source catalog with columns:
    """
    random_state = get_random_state(random_state)

    # Generate positions 1D, 2D, or 3D
    if dimension == 3:
        x = random_state.uniform(-dmax, dmax, size)
        y = random_state.uniform(-dmax, dmax, size)
        z = random_state.uniform(-dmax, dmax, size)
    elif dimension == 2:
        x = random_state.uniform(-dmax, dmax, size)
        y = random_state.uniform(-dmax, dmax, size)
        z = np.zeros_like(x)
    else:
        x = random_state.uniform(-dmax, dmax, size)
        y = np.zeros_like(x)
        z = np.zeros_like(x)

    table = Table()
    table["x"] = Column(x, unit="pc", description="Galactic cartesian coordinate")
    table["y"] = Column(y, unit="pc", description="Galactic cartesian coordinate")
    table["z"] = Column(z, unit="pc", description="Galactic cartesian coordinate")

    return table


def make_catalog_random_positions_sphere(
    size, center="Earth", distance=Quantity([0, 1], "Mpc"), random_state="random-seed"

Wrong hanging indentation before block (add 4 spaces).

):
    """Sample random source locations in a sphere.

    This can be used to generate an isotropic source population
    to represent extra-galactic sources.

    Parameters
    ----------
    size : int
        Number of sources
    center : {'Earth', 'Milky Way'}
        Sphere center
    distance : `~astropy.units.Quantity` tuple
        Distance min / max range.
    random_state : {int, 'random-seed', 'global-rng', `~numpy.random.RandomState`}
        Defines random number generator initialisation.
        Passed to `~gammapy.utils.random.get_random_state`.

    Returns
    -------
    catalog : `~astropy.table.Table`
        Source catalog with columns:

        - RAJ2000, DEJ2000 (deg)
        - GLON, GLAT (deg)
        - Distance (Mpc)
    """
    random_state = get_random_state(random_state)

    lon, lat = sample_sphere(size, random_state=random_state)
    radius = sample_sphere_distance(
        distance[0], distance[1], size, random_state=random_state
    )

    # TODO: it shouldn't be necessary here to convert to cartesian ourselves ...

TODO and FIXME comments should generally be avoided.

    x, y, z = spherical_to_cartesian(radius, lat, lon)
    pos = SkyCoord(x, y, z, frame="galactocentric", representation="cartesian")

    if center == "Milky Way":
        pass
    elif center == "Earth":
        # TODO: add shift Galactic center -> Earth

TODO and FIXME comments should generally be avoided.

        raise NotImplementedError
    else:
        msg = "Invalid center: {}\n".format(center)
        msg += "Choose one of: Earth, Milky Way"
        raise ValueError(msg)

    table = Table()
    table.meta["center"] = center

    icrs = pos.transform_to("icrs")
    table["RAJ2000"] = icrs.ra.to("deg")
    table["DEJ2000"] = icrs.dec.to("deg")

    galactic = icrs.transform_to("galactic")
    table["GLON"] = galactic.l.to("deg")
    table["GLAT"] = galactic.b.to("deg")

    table["Distance"] = icrs.distance.to("Mpc")

    return table


def make_base_catalog_galactic(

Too many arguments (7/5)

This function exceeds the maximum number of variables (26/15).

    n_sources,

Wrong hanging indentation before block (add 4 spaces).

    rad_dis="YK04",

Wrong hanging indentation before block (add 4 spaces).

    vel_dis="H05",

Wrong hanging indentation before block (add 4 spaces).

    max_age=Quantity(1e6, "yr"),

Wrong hanging indentation before block (add 4 spaces).

    spiralarms=True,

Wrong hanging indentation before block (add 4 spaces).

    n_ISM=Quantity(1, "cm-3"),

Wrong hanging indentation before block (add 4 spaces).

    random_state="random-seed",

Wrong hanging indentation before block (add 4 spaces).

):
    """Make a catalog of Galactic sources, with basic source parameters.

    Choose a radial distribution, a velocity distribution, the number
    of pulsars n_pulsars, the maximal age max_age[years] and the fraction
    of the individual morphtypes. There's an option spiralarms. If set on
    True a spiralarm modelling after Faucher&Kaspi is included.

    max_age and n_sources effectively correspond to s SN rate:
    SN_rate = n_sources / max_age

    Parameters
    ----------
    n_sources : int
        Number of sources to simulate.
    rad_dis : callable
        Radial surface density distribution of sources.
    vel_dis : callable
        Proper motion velocity distribution of sources.
    max_age : `~astropy.units.Quantity`
        Maximal age of the source
    spiralarms : bool
        Include a spiralarm model in the catalog.
    n_ISM : `~astropy.units.Quantity`
        Density of the interstellar medium.
    random_state : {int, 'random-seed', 'global-rng', `~numpy.random.RandomState`}
        Defines random number generator initialisation.
        Passed to `~gammapy.utils.random.get_random_state`.

    Returns
    -------
    table : `~astropy.table.Table`
        Catalog of simulated source positions and proper velocities.
    """
    random_state = get_random_state(random_state)

    if isinstance(rad_dis, six.string_types):
        rad_dis = radial_distributions[rad_dis]

    if isinstance(vel_dis, six.string_types):
        vel_dis = velocity_distributions[vel_dis]

    # Draw random values for the age
    age = random_state.uniform(0, max_age.to_value("yr"), n_sources)
    age = Quantity(age, "yr")

    # Draw r and z values from the given distribution
    r = draw(
        RMIN.to_value("kpc"),
        RMAX.to_value("kpc"),
        n_sources,
        pdf(rad_dis()),
        random_state=random_state,
    )
    r = Quantity(r, "kpc")

    z = draw(
        ZMIN.to_value("kpc"),
        ZMAX.to_value("kpc"),
        n_sources,
        Exponential(),
        random_state=random_state,
    )
    z = Quantity(z, "kpc")

    # Apply spiralarm modelling or not
    if spiralarms:
        r, theta, spiralarm = FaucherSpiral()(r, random_state=random_state)
    else:
        theta = Quantity(random_state.uniform(0, 2 * np.pi, n_sources), "rad")
        spiralarm = None

    # Compute cartesian coordinates
    x, y = astrometry.cartesian(r, theta)

    # Draw values from velocity distribution
    v = draw(
        VMIN.to_value("km/s"),
        VMAX.to_value("km/s"),
        n_sources,
        vel_dis(),
        random_state=random_state,
    )
    v = Quantity(v, "km/s")

    # Draw random direction of initial velocity
    theta = Quantity(random_state.uniform(0, np.pi, x.size), "rad")
    phi = Quantity(random_state.uniform(0, 2 * np.pi, x.size), "rad")

    # Compute new position
    dx, dy, dz, vx, vy, vz = astrometry.motion_since_birth(v, age, theta, phi)

    # Add displacement to birth position
    x_moved = x + dx
    y_moved = y + dy
    z_moved = z + dz

    # Set environment interstellar density
    n_ISM = n_ISM * np.ones(n_sources)

    table = Table()
    table["age"] = Column(age, unit="yr", description="Age of the source")
    table["n_ISM"] = Column(
        n_ISM, unit="cm-3", description="Interstellar medium density"
    )
    if spiralarms:
        table["spiralarm"] = Column(spiralarm, description="Which spiralarm?")

    table["x_birth"] = Column(
        x, unit="kpc", description="Galactocentric x coordinate at birth"
    )
    table["y_birth"] = Column(
        y, unit="kpc", description="Galactocentric y coordinate at birth"
    )
    table["z_birth"] = Column(
        z, unit="kpc", description="Galactocentric z coordinate at birth"
    )

    table["x"] = Column(
        x_moved.to("kpc"), unit="kpc", description="Galactocentric x coordinate"
    )
    table["y"] = Column(
        y_moved.to("kpc"), unit="kpc", description="Galactocentric y coordinate"
    )
    table["z"] = Column(
        z_moved.to("kpc"), unit="kpc", description="Galactocentric z coordinate"
    )

    table["vx"] = Column(
        vx.to("km/s"), unit="km/s", description="Galactocentric velocity in x direction"
    )
    table["vy"] = Column(
        vy.to("km/s"), unit="km/s", description="Galactocentric velocity in y direction"
    )
    table["vz"] = Column(
        vz.to("km/s"), unit="km/s", description="Galactocentric velocity in z direction"
    )
    table["v_abs"] = Column(
        v, unit="km/s", description="Galactocentric velocity (absolute)"
    )

    return table


def add_snr_parameters(table):
    """Add SNR parameters to the table."""
    # Read relevant columns
    age = table["age"].quantity
    n_ISM = table["n_ISM"].quantity

    # Compute properties
    snr = SNR(n_ISM=n_ISM)
    E_SN = snr.e_sn * np.ones(len(table))
    r_out = snr.radius(age)
    r_in = snr.radius_inner(age)
    L_SNR = snr.luminosity_tev(age)

    # Add columns to table
    table["E_SN"] = Column(E_SN, unit="erg", description="SNR kinetic energy")
    table["r_out"] = Column(r_out, unit="pc", description="SNR outer radius")
    table["r_in"] = Column(r_in, unit="pc", description="SNR inner radius")
    table["L_SNR"] = Column(L_SNR, unit="s-1", description="SNR luminosity")
    return table


def add_pulsar_parameters(

Too many arguments (6/5)

This function exceeds the maximum number of variables (18/15).

    table,

Wrong hanging indentation before block (add 4 spaces).

    B_mean=12.05,

Wrong hanging indentation before block (add 4 spaces).

    B_stdv=0.55,

Wrong hanging indentation before block (add 4 spaces).

    P_mean=0.3,

Wrong hanging indentation before block (add 4 spaces).

    P_stdv=0.15,

Wrong hanging indentation before block (add 4 spaces).

    random_state="random-seed",

Wrong hanging indentation before block (add 4 spaces).

):
    """Add pulsar parameters to the table.

    For the initial normal distribution of period and logB can exist the following
    Parameters: B_mean=12.05[log Gauss], B_stdv=0.55, P_mean=0.3[s], P_stdv=0.15

    Parameters
    ----------
    random_state : {int, 'random-seed', 'global-rng', `~numpy.random.RandomState`}
        Defines random number generator initialisation.
        Passed to `~gammapy.utils.random.get_random_state`.

    """
    random_state = get_random_state(random_state)
    # Read relevant columns
    age = table["age"].quantity

    # Draw the initial values for the period and magnetic field
    def p_dist(x):
        return np.exp(-0.5 * ((x - P_mean) / P_stdv) ** 2)

    p0_birth = draw(0, 2, len(table), p_dist, random_state=random_state)
    p0_birth = Quantity(p0_birth, "s")

    logB = random_state.normal(B_mean, B_stdv, len(table))

    # Compute pulsar parameters
    psr = Pulsar(p0_birth, logB)
    p0 = psr.period(age)
    p1 = psr.period_dot(age)
    p1_birth = psr.P_dot_0
    tau = psr.tau(age)
    tau_0 = psr.tau_0
    l_psr = psr.luminosity_spindown(age)
    l0_psr = psr.L_0

    # Add columns to table
    table["P0"] = Column(p0, unit="s", description="Pulsar period")
    table["P1"] = Column(p1, unit="", description="Pulsar period derivative")
    table["P0_birth"] = Column(p0_birth, unit="s", description="Pulsar birth period")
    table["P1_birth"] = Column(
        p1_birth, unit="", description="Pulsar birth period derivative"
    )
    table["CharAge"] = Column(tau, unit="yr", description="Pulsar characteristic age")
    table["Tau0"] = Column(tau_0, unit="yr")
    table["L_PSR"] = Column(l_psr, unit="erg s-1")
    table["L0_PSR"] = Column(l0_psr, unit="erg s-1")
    table["logB"] = Column(logB, unit="Gauss")
    return table


def add_pwn_parameters(table):
    """Add PWN parameters to the table."""
    # Some of the computations (specifically `pwn.radius`) aren't vectorised
    # across all parameters; so here we loop over source parameters explicitly

    results = []

    for idx in range(len(table)):
        age = table["age"].quantity[idx]
        E_SN = table["E_SN"].quantity[idx]
        n_ISM = table["n_ISM"].quantity[idx]
        P0_birth = table["P0_birth"].quantity[idx]
        logB = table["logB"][idx]

        # Compute properties
        pulsar = Pulsar(P0_birth, logB)
        snr = SNRTrueloveMcKee(e_sn=E_SN, n_ISM=n_ISM)
        pwn = PWN(pulsar, snr)
        r_out_pwn = pwn.radius(age).to_value("pc")
        L_PWN = pwn.luminosity_tev(age).to_value("erg")
        results.append(dict(r_out_pwn=r_out_pwn, L_PWN=L_PWN))

    # Add columns to table
    table["r_out_PWN"] = Column(
        [_["r_out_pwn"] for _ in results], unit="pc", description="PWN outer radius"
    )
    table["L_PWN"] = Column(
        [_["L_PWN"] for _ in results],
        unit="erg",
        description="PWN luminosity above 1 TeV",
    )
    return table


def add_observed_source_parameters(table):

This function exceeds the maximum number of variables (16/15).

    """Add observed source parameters to the table."""
    # Read relevant columns
    distance = table["distance"]
    r_in = table["r_in"]
    r_out = table["r_out"]
    r_out_PWN = table["r_out_PWN"]
    L_SNR = table["L_SNR"]
    L_PSR = table["L_PSR"]
    L_PWN = table["L_PWN"]

    # Compute properties
    ext_in_SNR = astrometry.radius_to_angle(r_in, distance)
    ext_out_SNR = astrometry.radius_to_angle(r_out, distance)
    ext_out_PWN = astrometry.radius_to_angle(r_out_PWN, distance)

    # Ellipse parameters not used for now
    theta = np.pi / 2 * np.ones(len(table))  # Position angle?
    epsilon = np.zeros(len(table))  # Ellipticity?

    S_SNR = astrometry.luminosity_to_flux(L_SNR, distance)
    # Ld2_PSR = astrometry.luminosity_to_flux(L_PSR, distance)
    Ld2_PSR = L_PSR / distance ** 2
    S_PWN = astrometry.luminosity_to_flux(L_PWN, distance)

    # Add columns
    table["ext_in_SNR"] = Column(ext_in_SNR, unit="deg")
    table["ext_out_SNR"] = Column(ext_out_SNR, unit="deg")
    table["ext_out_PWN"] = Column(ext_out_PWN, unit="deg")
    table["theta"] = Column(theta, unit="rad")
    table["epsilon"] = Column(epsilon, unit="")
    table["S_SNR"] = Column(S_SNR, unit="cm-2 s-1")
    table["Ld2_PSR"] = Column(Ld2_PSR, unit="erg s-1 kpc-2")
    table["S_PWN"] = Column(S_PWN, unit="cm-2 s-1")
    return table


def add_observed_parameters(table, obs_pos=None):

This function exceeds the maximum number of variables (20/15).

    """Add observable parameters (such as sky position or distance).

    Input table columns: x, y, z, extension, luminosity

    Output table columns: distance, glon, glat, flux, angular_extension

    Position of observer in cartesian coordinates.
    Center of galaxy as origin, x-axis goes trough sun.

    Parameters
    ----------
    table : `~astropy.table.Table`
        Input table
    obs_pos : tuple or None
        Observation position (X, Y, Z) in Galactocentric coordinates (default: Earth)

    Returns
    -------
    table : `~astropy.table.Table`
        Modified input table with columns added
    """
    obs_pos = obs_pos or [D_SUN_TO_GALACTIC_CENTER, 0, 0]

    # Get data
    x, y, z = table["x"].quantity, table["y"].quantity, table["z"].quantity
    vx, vy, vz = table["vx"].quantity, table["vy"].quantity, table["vz"].quantity

    distance, glon, glat = astrometry.galactic(x, y, z, obs_pos=obs_pos)

    # Compute projected velocity
    v_glon, v_glat = astrometry.velocity_glon_glat(x, y, z, vx, vy, vz)

    coordinate = SkyCoord(glon, glat, unit="deg", frame="galactic").transform_to("icrs")
    ra, dec = coordinate.ra.deg, coordinate.dec.deg

    # Add columns to table
    table["distance"] = Column(
        distance, unit="pc", description="Distance observer to source center"
    )
    table["GLON"] = Column(glon, unit="deg", description="Galactic longitude")
    table["GLAT"] = Column(glat, unit="deg", description="Galactic latitude")
    table["VGLON"] = Column(
        v_glon.to("deg/Myr"),
        unit="deg/Myr",
        description="Velocity in Galactic longitude",
    )
    table["VGLAT"] = Column(
        v_glat.to("deg/Myr"),
        unit="deg/Myr",
        description="Velocity in Galactic latitude",
    )
    table["RA"] = Column(ra, unit="deg", description="Right ascension")
    table["DEC"] = Column(dec, unit="deg", description="Declination")

    try:
        luminosity = table["luminosity"]
        flux = astrometry.luminosity_to_flux(luminosity, distance)
        table["flux"] = Column(flux.value, unit=flux.unit, description="Source flux")
    except KeyError:
        pass

    try:
        extension = table["extension"]
        angular_extension = np.degrees(np.arctan(extension / distance))
        table["angular_extension"] = Column(
            angular_extension,
            unit="deg",
            description="Source angular radius (i.e. half-diameter)",
        )
    except KeyError:
        pass

    return table