stellar-mass_halo-mass_Moster2012

# Initialize Auto-Reloading Magic
%reload_ext autoreload
%autoreload 2

import os
import sys
from importlib import reload

import astropy as ap
import numpy as np
import scipy as sp
import scipy.stats
import scipy.interpolate
import matplotlib as mpl
import matplotlib.pyplot as plt

# Silence annoying numpy errors
np.seterr(divide='ignore', invalid='ignore');

# Plotting settings
mpl.rc('font', **{'family': 'serif', 'sans-serif': ['Times'], 'size': 12})
mpl.rc('lines', solid_capstyle='round')
mpl.rc('mathtext', fontset='cm')
plt.rcParams.update({'grid.alpha': 0.5})

MSOL = ap.constants.M_sun.cgs.value

def colormap(args, cmap, log=False):
    cmap = plt.get_cmap(cmap)
    min = np.min(args)
    max = np.max(args)    

    # Create normalization
    if log:
        norm = mpl.colors.LogNorm(vmin=min, vmax=max)
    else:
        norm = mpl.colors.Normalize(vmin=min, vmax=max)

    # Create scalar-mappable
    smap = mpl.cm.ScalarMappable(norm=norm, cmap=cmap)
    # Bug-Fix something something
    smap._A = []
    
    return smap


class Moster_1205_5807():

    DATA = {
        "M1": {
            "VALS": [11.590, 1.195],
            "SIGMA": [0.236, 0.353],
            "LOG": True
        },

        "N1": {
            "VALS": [0.0351, -0.0247],
            "SIGMA": [0.0058, 0.0069],
            "LOG": False
        },

        "B1": {
            "VALS": [1.376, -0.826],
            "SIGMA": [0.153, 0.225],
            "LOG": False
        },

        "G1": {
            "VALS": [0.608, 0.329],
            "SIGMA": [0.059, 0.173],
            "LOG": False
        }
    }

    def __init__(self, redz=0.0, store=False):
        MHALO_GRID_RANGE = [10.0, 20.0]   # log10(M_h/Msol)
        MHALO_GRID_SIZE_PER_DEX = 4      # per decade
        NUM_MC = 1000

        SIGMA_GRID_RANGE = [-3.0, 3.0]   # standard-deviations
        SIGMA_GRID_SIZE_PER_STD = 4      # per std

        # Construct a grid of halo-masses
        mhalo_grid_size = np.diff(MHALO_GRID_RANGE) * MHALO_GRID_SIZE_PER_DEX
        mhalo_grid = np.logspace(*MHALO_GRID_RANGE, mhalo_grid_size) * MSOL

        # MC Calculate stellar-masses from grid of halo-masses
        # shape: (N, M) for `N` halo-masses, and `M` MC samples
        mstar_from_mhalo_grid = self.stellar_mass(mhalo_grid[:, np.newaxis], size=NUM_MC)

        # Construct a grid of standard-deviation values
        sigma_grid_size = np.diff(SIGMA_GRID_RANGE) * SIGMA_GRID_SIZE_PER_STD
        sigma_grid = np.linspace(*SIGMA_GRID_RANGE, sigma_grid_size)
        # Convert standard-deviations to percentiles
        percentiles = sp.stats.norm.cdf(sigma_grid)
        # Calculate distribution of stellar masses
        # shape: (L, N) for `L` standard-deviation values and `N` halo masses
        mstar_dist = np.percentile(mstar_from_mhalo_grid, 100*percentiles, axis=-1)

        # Find the range of valid stellar-masses
        # Minimum value is the one reached by the *highest* percentile, at the lowest halo-mass
        # Maximum value is the one reached by the *lowest* percentile, at the highest halo-mass
        mstar_range = [np.max(mstar_dist[:, 0]), np.min(mstar_dist[:, -1])]
        # Construct a grid of stellar-masses spanning this range
        mstar_grid = np.logspace(*np.log10(mstar_range), mhalo_grid.size//2)

        # Interpolate to find halo-masses corresponding to stellar-masses at each percentile
        mhalo_dist = [sp.interpolate.interp1d(np.log10(pp), np.log10(mhalo_grid))(np.log10(mstar_grid))
                      for pp in mstar_dist]
        mhalo_dist = np.power(10, mhalo_dist)

        if store:
            self._mhalo_grid = mhalo_grid
            self._sigma_grid = sigma_grid
            self._mstar_grid = mstar_grid
            self._mhalo_dist = mhalo_dist
            self._mstar_dist = mstar_dist

        # Construct 2D interpolant between stellar-mass and standard deviation, and halo-mass
        # Note: input log(stellar-mass), output log(halo-mass)
        self._log_mstar_from_log_mhalo_sigma = sp.interpolate.interp2d(
            np.log10(mhalo_grid), sigma_grid, np.log10(mstar_dist))

        # Construct 2D interpolant between stellar-mass and standard deviation, and halo-mass
        # Note: input log(stellar-mass), output log(halo-mass)
        self._log_mhalo_from_log_mstar_sigma = sp.interpolate.interp2d(
            np.log10(mstar_grid), sigma_grid[::-1], np.log10(mhalo_dist))

        return

    def mhalo_from_mstar(self, mstar, sigma=0.0):
        log_mstar = np.log10(mstar)
        log_mhalo = self._log_mhalo_from_log_mstar_sigma(log_mstar, sigma)
        mhalo = np.power(10.0, log_mhalo)
        return mhalo

    def mstar_from_mhalo(self, mhalo, sigma=0.0):
        log_mhalo = np.log10(mhalo)
        log_mstar = self._log_mstar_from_log_mhalo_sigma(log_mhalo, sigma)
        mstar = np.power(10.0, log_mstar)
        return mstar

    @classmethod
    def param(cls, name, redz=0.0, size=None):
        data = cls.DATA[name]
        vals = data['VALS']
        sigma = data['SIGMA']
        log_flag = data['LOG']

        if size is not None:
            vals = [np.random.normal(pp, ss, size=size) for pp, ss in zip(vals, sigma)]

        par = vals[0] + vals[1] * redz / (1 + redz)
        if log_flag:
            par = np.power(10.0, par) * MSOL

        return par

    @classmethod
    def stellar_mass(cls, mhalo, redz=0.0, size=None):
        m1 = cls.param("M1", redz=redz, size=size)
        norm = cls.param("N1", redz=redz, size=size)
        beta = cls.param("B1", redz=redz, size=size)
        gamma = cls.param("G1", redz=redz, size=size)

        bterm = np.power(mhalo/m1, -beta)
        gterm = np.power(mhalo/m1, gamma)
        mstar = mhalo * 2 * norm / (bterm + gterm)
        return mstar


mos = Moster_1205_5807(store=True)

labels = ['Halo Mass $[M_\odot]$', 'Stellar Mass $[M_\odot]$']

fig, axes = plt.subplots(figsize=[15, 8], ncols=2)
for ax, xlab, ylab in zip(axes, labels, labels[::-1]):
    ax.grid(True)
    ax.set(xscale='log', xlabel=xlab, yscale='log', ylabel=ylab)

sigmas = [-2.0, -1.5, -1.0, -0.5, 0.0, 0.5, 1.0, 1.5, 2.0]

percs = sp.stats.norm.cdf(sigmas)
col_percs = np.square(np.fabs(0.5 - percs))
cmap = colormap(col_percs, 'Blues_r')
colors = cmap.to_rgba(col_percs)

ax = axes[0]
mhalo_grid = np.logspace(10, 15, 20)
mstar_grid = np.logspace(8, 12, 20)
funcs = [mos.mstar_from_mhalo, mos.mhalo_from_mstar]
mass_grids = [mhalo_grid, mstar_grid]
for ii, (ax, fun, grid) in enumerate(zip(axes, funcs, mass_grids)):

    med = fun(grid * MSOL) / MSOL

    ax.plot(grid, med, 'k--', lw=4.0, alpha=0.4)
    ax.plot(grid, med, 'r--', lw=2.0)

    dist = fun(grid*MSOL, sigma=sigmas) / MSOL
    lines = []
    names = []
    for ss, ms, cc in zip(sigmas, dist, colors):
        la, = ax.plot(grid, ms, color='0.5', lw=2.0, alpha=0.5)
        lb, = ax.plot(grid, ms, color=cc)
        if ss >= 0.0:
            lines.append((la, lb))
            names.append("${:.1f} \,\, \sigma$".format(ss))
    
    if ii == 0:
        ax.legend(lines, names, loc='lower right')

test_mstar = 1e11

# NOTE: tihs is negative sigma value for mhalo_from_mstar specifically (i.e. right plot)
test_sigma = -1.0
test_mhalo = mos.mhalo_from_mstar(test_mstar*MSOL, sigma=test_sigma) / MSOL

data = [test_mhalo, test_mstar]
for ax, xx, yy in zip(axes, data, data[::-1]):
    ax.scatter(xx, yy, marker='*', color='red', alpha=0.75, s=100, zorder=10)
        
plt.show()