def test_nfw_mc_positions():
    """
  Compares samples with halotools and analytic density
  """
    model = NFWProfile()
    scaled_radius = np.logspace(-2, 0, 100)

    for c in [5, 10, 20]:
        distr = NFW(concentration=c, Rvir=1)

        samples = model.mc_generate_nfw_radial_positions(num_pts=int(1e6),
                                                         conc=c,
                                                         halo_radius=1)
        samples_tf = distr.sample(1e6)

        h = np.histogram(samples, 32, density=True, range=[0.01, 1])
        h_tf = np.histogram(samples_tf, 32, density=True, range=[0.01, 1])
        x = 0.5 * (h[1][:-1] + h[1][1:])

        p = distr.prob(x)

        # Comparing histograms
        assert_allclose(h[0], h_tf[0], rtol=5e-2)
        # Comparing to prob
        assert_allclose(h_tf[0], p, rtol=5e-2)
def test_nfw_mass_cdf():
    """
  Compares CDF values to halotools
  """
    model = NFWProfile()
    scaled_radius = np.logspace(-2, 0, 100)

    for c in [5, 10, 20]:
        distr = NFW(concentration=c, Rvir=1)
        y = model.cumulative_mass_PDF(scaled_radius, conc=c)
        y_tf = distr.cdf(scaled_radius)
        assert_allclose(y, y_tf.numpy(), rtol=1e-4)
Beispiel #3
0
    def _dc_fog_corr(
        self,
        abs_mag,
        halos,
        galinhalo,
        halo_centers,
        mag_threshold=-20.5,
        seedvalue=None,
    ):
        """Corrected comoving distance."""
        # array to store return results

        dcfogcorr = np.zeros(len(self.z))

        # run for each massive halo
        for i in halos.labels_h_massive[0]:
            # select only galaxies with magnitudes over than -20.5
            sat_gal_mask = (galinhalo.groups == i) * (abs_mag > mag_threshold)

            # number of galaxies for corrections
            numgal = np.sum(sat_gal_mask)

            # Monte Carlo simulation for distance
            nfw = NFWProfile(self.cosmo, halo_centers[i], mdef=self.delta_c)
            radial_positions_pos = nfw.mc_generate_nfw_radial_positions(
                num_pts=N_MONTE_CARLO,
                halo_radius=halos.radius[i],
                seed=seedvalue,
            )
            radial_positions_neg = nfw.mc_generate_nfw_radial_positions(
                num_pts=N_MONTE_CARLO,
                halo_radius=halos.radius[i],
                seed=seedvalue,
            )
            radial_positions_neg = -1 * radial_positions_neg
            radial_positions = np.r_[radial_positions_pos,
                                     radial_positions_neg]

            # random choice of distance for each galaxy
            al = np.random.RandomState(seedvalue)
            dc = al.choice(radial_positions, size=numgal)

            # combine Monte Carlo distance and distance to halo center
            dcfogcorr[sat_gal_mask] = halo_centers[i] + dc

        return dcfogcorr, halo_centers, halos.radius, galinhalo.groups
Beispiel #4
0
def HaloConcentration(mass, cosmo, redshift, mdef='vir'):
    """
    Return halo concentration from halo mass, based on the analytic fitting
    formulas presented in
    `Dutton and Maccio 2014 <https://arxiv.org/abs/1402.7073>`_.

    .. note::
        The units of the input mass are assumed to be :math:`M_{\odot}/h`

    Parameters
    ----------
    mass : array_like
        either a numpy or dask array specifying the halo mass; units
        assumed to be :math:`M_{\odot}/h`
    cosmo : :class:`~nbodykit.cosmology.cosmology.Cosmology`
        the cosmology instance used in the analytic formula
    redshift : float
        compute the c(M) relation at this redshift
    mdef : str, optional
        string specifying the halo mass definition to use; should be
        'vir' or 'XXXc' or 'XXXm' where 'XXX' is an int specifying the
        overdensity

    Returns
    -------
    concen : :class:`dask.array.Array`
        a dask array holding the analytic concentration values

    References
    ----------
    Dutton and Maccio, "Cold dark matter haloes in the Planck era: evolution
    of structural parameters for Einasto and NFW profiles", 2014, arxiv:1402.7073
    """
    from halotools.empirical_models import NFWProfile

    if not isinstance(mass, da.Array):
        mass = da.from_array(mass, chunks=100000)

    # initialize the model
    kws = {'cosmology':cosmo.to_astropy(), 'conc_mass_model':'dutton_maccio14', 'mdef':mdef, 'redshift':redshift}
    model = NFWProfile(**kws)

    return mass.map_blocks(lambda mass: model.conc_NFWmodel(prim_haloprop=mass), dtype=mass.dtype)
Beispiel #5
0
    def _group_prop(self, id_groups, groups, xyz):
        """Determine halos properties.

        Calculate Cartesian coordinates of halo centers, the halos comoving
        distances, the z-values of halo centers, halos radii and halos masses.

        """
        # select only galaxies in groups
        galincluster = id_groups[id_groups > -1]

        # arrays to store results
        xyzcenters = np.empty([len(galincluster), 3])
        dc_center = np.empty([len(galincluster)])
        hmass = np.empty([len(galincluster)])
        z_center = np.empty([len(galincluster)])
        radius = np.empty([len(galincluster)])

        # run for each group of galaxies
        for i in galincluster:
            mask = groups == i

            # halo radius
            radius[i] = self._radius(self.ra[mask], self.dec[mask],
                                     self.z[mask])

            # halo center
            x, y, z, dc, z_cen = self._centers(xyz[mask], self.z[mask])
            xyzcenters[i, 0] = x
            xyzcenters[i, 1] = y
            xyzcenters[i, 2] = z
            dc_center[i] = dc
            z_center[i] = z_cen

            # halo mass
            # use a Navarro profile (Navarro et al. 1997) [navarro97]_
            model = NFWProfile(self.cosmo, z_cen, mdef=self.delta_c)
            hmass[i] = model.halo_radius_to_halo_mass(radius[i])
        return xyzcenters, dc_center, z_center, radius, hmass
Beispiel #6
0
import halotools
from halotools.empirical_models import NFWProfile
cosmo
_critical_density_func = None


def rho0(z):
    global _rho0_func
    if _critical_density_func is None:
        zs = np.linspace(0, 10, 1000)
        rho_0 = 2.77536627e11
        density = cosmo.critical_density(zs) / cosmo.critical_density0 * rho_0
        _critical_density_func = np.interp(z, density)


nfwprofile = NFWProfile()


def F_grav(R, M200, c):
    return dtk.NFW_enclosed_mass(R, c) * M200 / R**2


def F_tidal(r, R, M200, c):
    return -F_grav(R, M200, c) + F_grav(R - r, M200, c)


plt.figure()
r = np.linspace(0, 0.2, 1000)
R = 1
M1 = 1e14
M2 = 1e11
    def mc_generate_nfw_phase_space_points(self, Ngals=int(1e4),
            conc=5, mass=1e12, b_to_a=0.7, c_to_a=0.5,
            halo_axisA_x=1.0, halo_axisA_y=0.0, halo_axisA_z=0.0,
            halo_axisC_x=1.0, halo_axisC_y=0.0, halo_axisC_z=0.0,
             verbose=True, seed=None):
        r""" Return a Monte Carlo realization of points
        in the phase space of an NFW halo in isotropic Jeans equilibrium.
        Parameters
        -----------
        Ngals : int, optional
            Number of galaxies in the Monte Carlo realization of the
            phase space distribution. Default is 1e4.
        conc : float, optional
            Concentration of the NFW profile being realized.
            Default is 5.
        mass : float, optional
            Mass of the halo whose phase space distribution is being realized
            in units of Msun/h. Default is 1e12.
        verbose : bool, optional
            If True, a message prints with an estimate of the build time.
            Default is True.
        seed : int, optional
            Random number seed used in the Monte Carlo realization.
            Default is None, which will produce stochastic results.
        Returns
        --------
        t : table
            `~astropy.table.Table` containing the Monte Carlo realization of the
            phase space distribution.
            Keys are 'x', 'y', 'z', 'vx', 'vy', 'vz', 'radial_position', 'radial_velocity'.
            Length units in Mpc/h, velocity units in km/s.
        Examples
        ---------
        >>> nfw = AnisotropicNFWPhaseSpace()
        >>> mass, conc, b_to_a, c_to_a = 1e13, 8., 0.9, 0.6
        >>> data = nfw.mc_generate_nfw_phase_space_points(Ngals=100, mass=mass, conc=conc, b_to_a=b_to_a, c_to_a=c_to_a, verbose=False)
        Now suppose you wish to compute the radial velocity dispersion of all the returned points:
        >>> vrad_disp = np.std(data['radial_velocity'])
        If you wish to do the same calculation but for points in a specific range of radius:
        >>> mask = data['radial_position'] < 0.1
        >>> vrad_disp_inner_points = np.std(data['radial_velocity'][mask])
        You may also wish to select points according to their distance to the halo center
        in units of the virial radius. In such as case, you can use the
        `~halotools.empirical_models.NFWPhaseSpace.halo_mass_to_halo_radius`
        method to scale the halo-centric distances. Here is an example
        of how to compute the velocity dispersion in the z-dimension of all points
        residing within :math:`R_{\rm vir}/2`:
        >>> halo_radius = nfw.halo_mass_to_halo_radius(mass)
        >>> scaled_radial_positions = data['radial_position']/halo_radius
        >>> mask = scaled_radial_positions < 0.5
        >>> vz_disp_inner_half = np.std(data['vz'][mask])
        """

        m = np.zeros(Ngals) + mass
        c = np.zeros(Ngals) + conc
        halo_axisA_x = np.zeros(Ngals) + halo_axisA_x
        halo_axisA_y = np.zeros(Ngals) + halo_axisA_y
        halo_axisA_z = np.zeros(Ngals) + halo_axisA_z
        halo_axisC_x = np.zeros(Ngals) + halo_axisC_x
        halo_axisC_y = np.zeros(Ngals) + halo_axisC_y
        halo_axisC_z = np.zeros(Ngals) + halo_axisC_z
        rvir = NFWProfile.halo_mass_to_halo_radius(self, total_mass=m)


        new_b_to_a, new_c_to_a = self.anisotropy_bias_response(b_to_a, c_to_a)

        print('here 1:')
        x, y, z = self.mc_halo_centric_pos(c,
            halo_radius=rvir,
            b_to_a=new_b_to_a,
            c_to_a=new_c_to_a,
            halo_axisA_x=halo_axisA_x,
            halo_axisA_y=halo_axisA_y,
            halo_axisA_z=halo_axisA_z,
            halo_axisC_x=halo_axisC_x,
            halo_axisC_y=halo_axisC_y,
            halo_axisC_z=halo_axisC_z,
            seed=seed)
        r = np.sqrt(x**2 + y**2 + z**2)
        scaled_radius = r/rvir

        if seed is not None:
            seed += 1
        vx = self.mc_radial_velocity(scaled_radius, m, c, seed=seed)
        if seed is not None:
            seed += 1
        vy = self.mc_radial_velocity(scaled_radius, m, c, seed=seed)
        if seed is not None:
            seed += 1
        vz = self.mc_radial_velocity(scaled_radius, m, c, seed=seed)

        xrel, vxrel = relative_positions_and_velocities(x, 0, v1=vx, v2=0)
        yrel, vyrel = relative_positions_and_velocities(y, 0, v1=vy, v2=0)
        zrel, vzrel = relative_positions_and_velocities(z, 0, v1=vz, v2=0)

        vrad = (xrel*vxrel + yrel*vyrel + zrel*vzrel)/r

        t = Table({'x': x, 'y': y, 'z': z,
            'vx': vx, 'vy': vy, 'vz': vz,
            'radial_position': r, 'radial_velocity': vrad})

        return t
Beispiel #8
0
def plot_richness_mass_relation():
    background.set_cosmology("wmap7")
    richness = np.logspace(np.log10(20), np.log10(300), 100)
    m200m_rykoff, r200m_rykoff = background.lambda_to_m200_r200(
        richness, 0.25, richness_mass_author="Rykoff_mean")
    m200m_simet, r200m_simet = background.lambda_to_m200_r200(
        richness, 0.25, richness_mass_author="Simet_mean")
    m200m_baxter, r200m_baxter = background.lambda_to_m200_r200(
        richness, 0.25, richness_mass_author="Baxter_mean")
    m200c_rykoff, r200c_rykoff = background.lambda_to_m200_r200(
        richness, 0.25, richness_mass_author="Rykoff_crit")
    m200c_simet, r200c_simet = background.lambda_to_m200_r200(
        richness, 0.25, richness_mass_author="Simet_crit")
    m200c_baxter, r200c_baxter = background.lambda_to_m200_r200(
        richness, 0.25, richness_mass_author="Baxter_crit")
    plt.figure()
    plt.plot(richness, m200m_rykoff, '-r')
    plt.plot(richness, m200m_simet, '-b')
    plt.plot(richness, m200m_baxter, '-g')
    plt.plot(richness, m200c_rykoff, '--r')
    plt.plot(richness, m200c_simet, '--b')
    plt.plot(richness, m200c_baxter, '--g')
    plt.plot([], [], 'r', label='rykoff')
    plt.plot([], [], 'b', label='simet')
    plt.plot([], [], 'g', label='baxter')
    plt.plot([], [], '-k', label='M200m')
    plt.plot([], [], '--k', label='M200c')
    plt.legend(loc='best', framealpha=0.3)
    plt.yscale('log')
    plt.xscale('log')
    plt.xlabel('Richness')
    plt.ylabel(r'M$_{200}$')
    plt.grid()

    plt.figure()
    plt.plot(richness, r200m_rykoff, '-r')
    plt.plot(richness, r200m_simet, '-b')
    plt.plot(richness, r200m_baxter, '-g')
    plt.plot(richness, r200c_rykoff, '--r')
    plt.plot(richness, r200c_simet, '--b')
    plt.plot(richness, r200c_baxter, '--g')
    plt.plot([], [], 'r', label='rykoff')
    plt.plot([], [], 'b', label='simet')
    plt.plot([], [], 'g', label='baxter')
    plt.plot([], [], '-k', label='R200m')
    plt.plot([], [], '--k', label='R200c')
    plt.legend(loc='best', framealpha=0.3)
    plt.yscale('log')
    plt.xscale('log')
    plt.xlabel('Richness')
    plt.ylabel(r'R$_{200}$')
    plt.grid()

    plt.figure()
    plt.plot(richness, m200c_rykoff / m200m_rykoff, '-r', label='rykoff')
    plt.plot(richness, m200c_simet / m200m_simet, '-b', label='simet')
    plt.plot(richness, m200c_baxter / m200m_baxter, '-g', label='baxter')
    plt.legend(loc='best', framealpha=0.3)
    plt.xscale('log')
    plt.xlabel('Richness')
    plt.ylabel(r'M200c/M200m')
    plt.grid()

    plt.figure()
    plt.plot(richness, r200c_rykoff / r200m_rykoff, '-r', label='rykoff')
    plt.plot(richness, r200c_simet / r200m_simet, '-b', label='simet')
    plt.plot(richness, r200c_baxter / r200m_baxter, '-g', label='baxter')
    plt.legend(loc='best', framealpha=0.3)
    plt.xscale('log')
    plt.xlabel('Richness')
    plt.ylabel(r'R200c/R200m')
    plt.grid()

    plt.figure()
    a = 1.0 / (1.0 + 0.25)
    # plt.plot(m200m_simet, r200m_simet, '-b',  label='my simet mean')
    # plt.plot(m200c_simet, r200c_simet, '--b', label='my simet crit')
    plt.plot(m200m_simet, r200m_simet, '-b', label='my simet mean')
    plt.plot(m200c_simet, r200c_simet, '--b', label='my simet crit')

    nfw_200m = NFWProfile(mdef='200m', redshift=0.25)
    nfw_200c = NFWProfile(mdef='200c', redshift=0.25)
    plt.plot(m200m_simet,
             nfw_200m.halo_mass_to_halo_radius(m200m_simet) * 1000,
             '-r',
             label='ht simet mean')
    plt.plot(m200c_simet,
             nfw_200c.halo_mass_to_halo_radius(m200c_simet) * 1000,
             '--r',
             label='ht simet crit')

    plt.plot(m200m_simet,
             M_to_R(m200m_simet, 0.25, '200m'),
             '-g',
             label='cls simet mean')
    plt.plot(m200c_simet,
             M_to_R(m200c_simet, 0.25, '200c'),
             '--g',
             label='cls simet crit')

    plt.legend(loc='best', framealpha=0.3)
    plt.xscale('log')
    plt.yscale('log')
    plt.xlabel('M200 [Msun]')
    plt.ylabel('R200 [proper kpc]')
    plt.grid()

    plt.show()

    #header list can be found in http://arxiv.org/pdf/1303.3562v2.pdf
    hdulist = pyfits.open("redmapper_dr8_public_v6.3_catalog.fits")
    tdata = hdulist[1].data

    # red_ra = tdata.field('ra')
    # red_dec= tdata.field('dec')
    red_z = tdata.field('z_lambda')
    red_lambda = tdata.field('lambda')
    # red_pcen=tdata.field('p_cen')

    plt.figure()
    h, xbins = np.histogram(red_lambda, bins=256)
    plt.plot(dtk.bins_avg(xbins), h)
    plt.yscale('log')
    plt.xscale('log')
    plt.xlabel('Richness')
    plt.ylabel('count')

    m200m_rykoff = background.lambda_to_m200m_Rykoff(red_lambda, 0.25)
    m200m_simet = background.lambda_to_m200m_Simet(red_lambda, 0.25)
    m200m_baxter = background.lambda_to_m200m_Baxter(red_lambda, 0.25)
    m200c_rykoff = background.lambda_to_m200c_Rykoff(red_lambda, 0.25)
    m200c_simet = background.lambda_to_m200c_Simet(red_lambda, 0.25)
    m200c_baxter = background.lambda_to_m200c_Baxter(red_lambda, 0.25)

    xbins = np.logspace(13.5, 15.5, 100)
    xbins_avg = dtk.bins_avg(xbins)

    plt.figure()
    h, _ = np.histogram(m200m_rykoff, bins=xbins)
    plt.plot(xbins_avg, h, label='rykoff')
    h, _ = np.histogram(m200m_simet, bins=xbins)
    plt.plot(xbins_avg, h, label='simet')
    h, _ = np.histogram(m200m_baxter, bins=xbins)
    plt.plot(xbins_avg, h, label='baxter')
    plt.yscale('log')
    plt.xscale('log')
    plt.ylabel('Counts')
    plt.xlabel('M$_{200m}$ [h$^{-1}$ M$_\odot$]')
    plt.legend(loc='best', framealpha=1.0)
    plt.grid()

    plt.figure()
    h, _ = np.histogram(m200c_rykoff, bins=xbins)
    plt.plot(xbins_avg, h, label='rykoff')
    h, _ = np.histogram(m200c_simet, bins=xbins)
    plt.plot(xbins_avg, h, label='simet')
    h, _ = np.histogram(m200c_baxter, bins=xbins)
    plt.plot(xbins_avg, h, label='baxter')
    plt.yscale('log')
    plt.xscale('log')
    plt.ylabel('Counts')
    plt.xlabel('M$_{200c}$ [h$^{-1}$ M$_\odot$]')
    plt.legend(loc='best', framealpha=1.0)
    plt.grid()

    plt.show()
Beispiel #9
0
 def get_nfw_conc(mass, redshift):
     kw1 = {}
     kw1.update(kws)
     kw1['redshift'] = redshift
     model = NFWProfile(**kw1)
     return model.conc_NFWmodel(prim_haloprop=mass)
Beispiel #10
0
# halotools.test_installation()

# determine dark mass and virial radius
model = PrebuiltSubhaloModelFactory('behroozi10', redshift = redshift)
stars_mass = 60850451172.24926
# stars_mass in units Msun from simulation
log_stars_mass = np.log10(stars_mass)
log_dark_mass = model.mean_log_halo_mass(log_stellar_mass=log_stars_mass)
dm_stellar = 10**log_dark_mass
virial = halo_mass_to_halo_radius(dm_stellar, cosmo, redshift, mdef)
virial_kpc = (virial * u.Mpc).to('kpc') / h
virial_rad = np.arange(radmin, virial_kpc.value, inc)
scaled_rad = virial_rad / np.max(virial_rad)

# calculate the density threshold for dimensionless mass density calculation
nfw = NFWProfile()
rho_thresh = density_threshold(cosmo, redshift, mdef)
# rho_thresh in units Msun*h^2/Mpc^3
rho_units = (rho_thresh * u.Msun / u.Mpc**3).to('Msun/kpc3') * h**2
dimless_massdens = nfw.dimensionless_mass_density(scaled_rad, conc)
mass_dens = dimless_massdens * rho_units
# mass_dens is in units Msun/kpc^3

# create an array of heights to take vertical components into account
height_arr = np.arange(-z, z+1)
height_scale = height_arr / z

# prepare arrays
nfw_heights = np.zeros(len(height_scale))
nfw_avg = np.zeros(len(scaled_rad))
Beispiel #11
0
        else:
            plt.plot(profile[:,0], profile[:,1], 'C0')

    # Auriga
    files = glob.glob('data/sims/Vcirc_auriga/*.txt')
    legend = False
    for file in files:
        profile = np.loadtxt(file)
        plt.plot(profile[:,0], profile[:,1], 'C1')
        if not legend:
            plt.plot(profile[:,0], profile[:,1], 'C1', label='Auriga')
            legend = True
        else:
            plt.plot(profile[:,0], profile[:,1], 'C1')

    nfw = NFWProfile()
    nfw_Vcirc = nfw.circular_velocity(profile[:,0]*10**-3, 10**12, conc=10)
    plt.plot(profile[:,0], nfw_Vcirc, 'k--',\
                label=r'NFW(10$^{12}$ M$_\odot$, c=10)')

    plt.xlabel(r'$r$ [kpc]')
    plt.ylabel(r'$V_{circ}$ [km s$^{-1}$]')
    plt.xlim(0.5, 150)
    plt.ylim(50, 300)
    plt.legend(loc='best')
    plt.xscale('log')
    plt.yscale('log')
    plt.savefig(pltpth+'vcirc.pdf', bbox_inches='tight')
    plt.close()

# # ## ### ##### ######## ############# #####################
Beispiel #12
0
 def get_nfw_conc(mass, redshift):
     kw1 = {}
     kw1.update(kws)
     kw1['redshift'] = redshift
     model = NFWProfile(**kw1)
     return model.conc_NFWmodel(prim_haloprop=mass)
    def __init__(self,
                 z,
                 m200c=None,
                 r200c=None,
                 c=None,
                 cM_err=False,
                 cosmo=cm.OuterRim_params,
                 seed=None):
        """
        Class for generating NFW test-case input files for the ray tracing modules supplied in
        the directory above. This class is constructed with a halo mass, redshift, and 
        cosmological model, and builds a HaloTools NFWProfile object. The methods provided here 
        can then populate the profile with a particle distribution realization in 3 dimensions, 
        and output the result in the form expected by the raytracing modules.
        
        Parameters
        ----------
        z : float 
            The redshift of the halo.
        m200c : float
            The mass of the halo within a radius containing 200*rho_crit, in M_sun. If not passed, 
            then r200c must be supplied
        r200c : float
            The radius of the halo enclosing a mean density of 200*rho_crit, in Mpc. If not passed,
            then m200c must be supplied.
        c : float, optional 
            The concentration of the halo. If not given, samples from a Gaussian 
            with location and scale suggested by the M-c relation of Child+2018
        cM_err : bool, optional 
            Whether or not to impose scatter on the cM relation used to draw a concentration, 
            in the case that the argument c is not passed. If False, the concentration drawn 
            will always lie exactly on the cM relation used (currently Child+ 2018). Defaults
            to True. In either case, the 'sod_halo_cdelta_error' quntity in the output halo propery
            csv file will be zero.
        cosmo : object, optional
            An AstroPy cosmology object. Defaults to OuterRim parameters.
        seed : float, optional
            Random seed to pass to HaloTools for generation of radial particle positions, and
            use for drawing concentrations and angular positions of particles. Defaults to None
            (giving stochastic output)
        
        Methods
        -------
        populate_halo(r)
            Uses HaloTools to generate a MonteCarlo realization of discrete tracers of the density
            profile (particles)
        output_particles():
            Writes out the particle positions generated by populate_halo() to a form that is prepped 
            for input to the ray tracing modules of this package.
        """

        assert (m200c is not None or r200c is not None
                ), "Either m200c (in M_sun) or r200c (in Mpc) must be supplied"
        self.m200c = m200c
        self.r200c = r200c
        self.redshift = z
        self.cosmo = cosmo
        self.seed = seed

        self.profile = NFWProfile(cosmology=self.cosmo,
                                  redshift=self.redshift,
                                  mdef='200c')

        # HaloTools and Colossus expect masses,radii with h dependence, so scale accordingly on input and output
        if (r200c is None):
            self.m200ch = m200c * cosmo.h  # M_sun/h
            self.r200ch = self.profile.halo_mass_to_halo_radius(
                self.m200ch)  #proper Mpc/h
            self.r200c = self.r200ch / cosmo.h  # proper Mpc
        if (m200c is None):
            self.r200ch = r200c * cosmo.h  # proper Mpc/h
            self.m200ch = self.profile.halo_radius_to_halo_mass(
                self.r200ch)  # M_sun/h
            self.m200c = self.m200ch / cosmo.h  # M_sun

        # these to be filled by populate_halo()
        self.r = None
        self.theta = None
        self.phi = None
        self.mpp = None
        self.max_rfrac = None
        self.populated = False

        if c is not None:
            self.c = c
            self.c_err = 0
        else:
            # if cM_err=True, draw a concentration from gaussian, otherwise use Child+2018 cM relation scatter-free
            rand = np.random.RandomState(self.seed)

            cosmo_colossus = colcos.setCosmology(
                'OuterRim', {
                    'Om0': cosmo.Om0,
                    'Ob0': cosmo.Ob0,
                    'H0': cosmo.H0.value,
                    'sigma8': 0.8,
                    'ns': 0.963,
                    'relspecies': False
                })
            c_u = mass_conc(self.m200ch, '200c', z, model='child18')
            if (cM_err):
                c_sig = c_u / 3
                self.c = rand.normal(loc=c_u, scale=c_sig)
            else:
                self.c = c_u
            self.c_err = 0
class NFW:
    def __init__(self,
                 z,
                 m200c=None,
                 r200c=None,
                 c=None,
                 cM_err=False,
                 cosmo=cm.OuterRim_params,
                 seed=None):
        """
        Class for generating NFW test-case input files for the ray tracing modules supplied in
        the directory above. This class is constructed with a halo mass, redshift, and 
        cosmological model, and builds a HaloTools NFWProfile object. The methods provided here 
        can then populate the profile with a particle distribution realization in 3 dimensions, 
        and output the result in the form expected by the raytracing modules.
        
        Parameters
        ----------
        z : float 
            The redshift of the halo.
        m200c : float
            The mass of the halo within a radius containing 200*rho_crit, in M_sun. If not passed, 
            then r200c must be supplied
        r200c : float
            The radius of the halo enclosing a mean density of 200*rho_crit, in Mpc. If not passed,
            then m200c must be supplied.
        c : float, optional 
            The concentration of the halo. If not given, samples from a Gaussian 
            with location and scale suggested by the M-c relation of Child+2018
        cM_err : bool, optional 
            Whether or not to impose scatter on the cM relation used to draw a concentration, 
            in the case that the argument c is not passed. If False, the concentration drawn 
            will always lie exactly on the cM relation used (currently Child+ 2018). Defaults
            to True. In either case, the 'sod_halo_cdelta_error' quntity in the output halo propery
            csv file will be zero.
        cosmo : object, optional
            An AstroPy cosmology object. Defaults to OuterRim parameters.
        seed : float, optional
            Random seed to pass to HaloTools for generation of radial particle positions, and
            use for drawing concentrations and angular positions of particles. Defaults to None
            (giving stochastic output)
        
        Methods
        -------
        populate_halo(r)
            Uses HaloTools to generate a MonteCarlo realization of discrete tracers of the density
            profile (particles)
        output_particles():
            Writes out the particle positions generated by populate_halo() to a form that is prepped 
            for input to the ray tracing modules of this package.
        """

        assert (m200c is not None or r200c is not None
                ), "Either m200c (in M_sun) or r200c (in Mpc) must be supplied"
        self.m200c = m200c
        self.r200c = r200c
        self.redshift = z
        self.cosmo = cosmo
        self.seed = seed

        self.profile = NFWProfile(cosmology=self.cosmo,
                                  redshift=self.redshift,
                                  mdef='200c')

        # HaloTools and Colossus expect masses,radii with h dependence, so scale accordingly on input and output
        if (r200c is None):
            self.m200ch = m200c * cosmo.h  # M_sun/h
            self.r200ch = self.profile.halo_mass_to_halo_radius(
                self.m200ch)  #proper Mpc/h
            self.r200c = self.r200ch / cosmo.h  # proper Mpc
        if (m200c is None):
            self.r200ch = r200c * cosmo.h  # proper Mpc/h
            self.m200ch = self.profile.halo_radius_to_halo_mass(
                self.r200ch)  # M_sun/h
            self.m200c = self.m200ch / cosmo.h  # M_sun

        # these to be filled by populate_halo()
        self.r = None
        self.theta = None
        self.phi = None
        self.mpp = None
        self.max_rfrac = None
        self.populated = False

        if c is not None:
            self.c = c
            self.c_err = 0
        else:
            # if cM_err=True, draw a concentration from gaussian, otherwise use Child+2018 cM relation scatter-free
            rand = np.random.RandomState(self.seed)

            cosmo_colossus = colcos.setCosmology(
                'OuterRim', {
                    'Om0': cosmo.Om0,
                    'Ob0': cosmo.Ob0,
                    'H0': cosmo.H0.value,
                    'sigma8': 0.8,
                    'ns': 0.963,
                    'relspecies': False
                })
            c_u = mass_conc(self.m200ch, '200c', z, model='child18')
            if (cM_err):
                c_sig = c_u / 3
                self.c = rand.normal(loc=c_u, scale=c_sig)
            else:
                self.c = c_u
            self.c_err = 0

    # -----------------------------------------------------------------------------------------------

    def populate_halo(self, N=10000, rfrac=1, rfrac_los=None):
        """
        Generates a 3-dimensional relization of the discreteley-sampled NFW mass distribution for
        this halo. The radial positions are obtained with the HaloTools 
        mc_generate_nfw_radial_positions module. The angular positions are drawn from a uniform 
        random distribution, the azimuthal coordiante ranging from 0 to 2pi, and the coaltitude 
        from 0 to pi.

        Parameters
        ----------
        N : int
            The number of particles to drawn
        rfrac: float, optional
            Multiplier of r200c which sets the maximum radial extent of the population
            (concentration will be scaled as well, as c=r200c/r_s). Defaults to 1
        rfrac_los: float, optional
            Multiplier of r200c which sets the maaximum extent of the population in the line-of-sight (LOS)
            dimension, i.e. it clips the halo along the LOS. If rfrac_los = 0.5 the LOS dimension of the halo 
            will be clipped 0.5*r200c toward the observer, and again away from the observer, with respect to 
            the halo center. Note that this does *not* rescale mpp, and is therefore almost totally useless...
            the argument is kept for rather specific debugging purposes, but probably should not be used. 
            Default is None, in which case no clipping is performed. Also if rfrac_los > rfrac, obviously 
            nothing will happen.
        """

        self.populated = True

        # the radial positions in proper Mpc
        r = self.profile.mc_generate_nfw_radial_positions(num_pts=N,
                                                          conc=rfrac * self.c,
                                                          halo_radius=rfrac *
                                                          self.r200ch,
                                                          seed=self.seed + 1)
        self.r = r / self.cosmo.h
        self.max_rfrac = rfrac

        # compute mass enclosed to find mass per particle
        # (this is the analytic integration of the NFW profile in terms of m_200c, assuming c=c_200c)
        rs = self.r200c / self.c
        rmax = rfrac * self.r200c
        n = np.log((rs + rmax) / rs) - rmax / (rmax + rs)
        d = np.log(1 + self.c) - self.c / (1 + self.c)
        M_enc = self.m200c * n / d
        self.mpp = M_enc / N

        # radial positions need to be in comoving comoving coordiantes, as the kappa maps in the raytracing
        # modules expect the density estimation to be done on a comoving set of particles
        self.r = self.r * (1 + self.redshift)

        # now let's add in uniform random positions in the angular coordinates as well
        # Note that this is not the same as a uniform distribution in theta and phi
        # over [0, pi] and [0, 2pi], since the area element on a sphere is a function of
        # the coaltitude! See http://mathworld.wolfram.com/SpherePointPicking.html
        rand = np.random.RandomState(self.seed)
        v = rand.uniform(low=0, high=1, size=len(r))
        self.phi = rand.uniform(low=0, high=2 * np.pi, size=len(r))
        self.theta = np.arccos(2 * v - 1)

        # finally, do los clipping if user requested (in self.output_particles below, the los dimension
        # is assumed to be the cartesian x)
        x = self.r * np.sin(self.theta) * np.cos(self.phi)
        if (rfrac_los is not None):
            los_mask = (np.abs(x) / self.r200c) <= rfrac_los
            self.r, self.theta, self.phi = self.r[los_mask], self.theta[
                los_mask], self.phi[los_mask]

    # -----------------------------------------------------------------------------------------------

    def populate_halo_fov(self, N=10000, rfrac=1, depth=None):
        """
        eventually merge with function above...
        """

        self.populated = True

        # modify requested radius to fill the encapsulating fov volume
        if (depth < rfrac or depth is None): depth = rfrac
        radius = max([rfrac, depth])
        halo_radius = np.sqrt(2) * radius * self.r200ch
        conc = np.sqrt(2) * radius * self.c

        # the radial positions in proper Mpc
        r = self.profile.mc_generate_nfw_radial_positions(
            num_pts=N, conc=conc, halo_radius=halo_radius, seed=self.seed + 1)
        self.r = r / self.cosmo.h
        self.max_rfrac = rfrac

        # compute mass enclosed to find mass per particle
        # (this is the analytic integration of the NFW profile in terms of m_200c, assuming c=c_200c)
        rs = self.r200c / self.c
        rmax = rfrac * self.r200c
        n = np.log((rs + rmax) / rs) - rmax / (rmax + rs)
        d = np.log(1 + self.c) - self.c / (1 + self.c)
        M_enc = self.m200c * n / d
        self.mpp = M_enc / N

        # radial positions need to be in comoving comoving coordiantes, as the kappa maps in the raytracing
        # modules expect the density estimation to be done on a comoving set of particles
        self.r = self.r * (1 + self.redshift)

        # now let's add in uniform random positions in the angular coordinates as well
        # Note that this is not the same as a uniform distribution in theta and phi
        # over [0, pi] and [0, 2pi], since the area element on a sphere is a function of
        # the coaltitude! See http://mathworld.wolfram.com/SpherePointPicking.html
        rand = np.random.RandomState(self.seed)
        v = rand.uniform(low=0, high=1, size=len(r))
        self.phi = rand.uniform(low=0, high=2 * np.pi, size=len(r))
        self.theta = np.arccos(2 * v - 1)

        # trim the particle population to the fov
        # move to polar coordinates, with x the los dimension
        x = self.r * np.sin(self.theta) * np.cos(self.phi)
        y = self.r * np.sin(self.theta) * np.sin(self.phi)
        z = self.r * np.cos(self.theta)
        fov_mask = np.logical_and.reduce(
            (np.abs(x) <= depth * self.r200c, np.abs(y) <= rmax * self.r200c,
             np.abs(z) <= rmax * self.r200c))
        self.r, self.theta, self.phi = self.r[fov_mask], self.theta[
            fov_mask], self.phi[fov_mask]

    # -----------------------------------------------------------------------------------------------

    def output_particles(self,
                         output_dir='./nfw_particle_realization',
                         vis_debug=False,
                         vis_output_dir=None,
                         fov_multiplier=None):
        """
        Computes three dimensional quantities for particles sampled along radial dimension. Each 
        quantity is output as little-endian binary files (expected input for ray-tracing modules
        in this package). The output quantities are x, y, z, theta, phi, redshift. In 
        cartesian space, the distribution is placed at a distance along the x-axis computed as 
        the comoving distance to the halo redshift by the input cosmology.

        Parameters
        ----------
        output_dir : string
            The desired output location for the binary files
        vis_debug : bool
            If True, display a 3d plot of the particles to be output for visual inspection
        vis_output_dir : string
            The desired output location for matplotlib figures images, if vis_debug is True
        fov_multiplier : float, optional
            The size of the field of view, which sets the scale of the eventual density 
            estimation, in projected comving Mpc, as a fraction of the largest radial particle 
            dispalcement. Defaults to None, in which case the field of view is set to be the 
            largest square that can fit entirely inside the projected halo.
        """

        if (self.populated == False):
            raise RuntimeError(
                'populate_halo must be called before output_particles')
        if (vis_output_dir is None): vis_output_dir = output_dir
        if not os.path.exists(output_dir):
            os.makedirs(output_dir, exist_ok=True)

        # now find projected positions wrt origin after pushing halo down x-axis (Mpc and arcsec)
        self.halo_r = self.cosmo.comoving_distance(self.redshift).value
        x = self.r * np.sin(self.theta) * np.cos(self.phi) + self.halo_r
        y = self.r * np.sin(self.theta) * np.sin(self.phi)
        z = self.r * np.cos(self.theta)
        r_fov = np.linalg.norm([y, z], axis=0)
        r_sky = np.linalg.norm([x, y, z], axis=0)
        theta_sky = np.arccos(z / r_sky) * 180 / np.pi * 3600
        phi_sky = np.arctan(y / x) * 180 / np.pi * 3600

        # get particle redshifts
        zmin = z_at_value(self.cosmo.comoving_distance,
                          ((r_sky.min() - 0.1) * u.Mpc))
        zmax = z_at_value(self.cosmo.comoving_distance,
                          ((r_sky.max() + 0.1) * u.Mpc))
        z_samp = np.linspace(zmin, zmax, 10)
        x_samp = self.cosmo.comoving_distance(z_samp).value
        invfunc = scipy.interpolate.interp1d(x_samp, z_samp)
        redshift = invfunc(r_sky)

        if (vis_debug):
            print(vis_output_dir)
            f = plt.figure(figsize=(12, 6))
            ax = f.add_subplot(121, projection='3d')
            ax2 = f.add_subplot(122)

            ax.scatter(x, y, z, c='k', alpha=0.25)
            max_range = np.array(
                [x.max() - x.min(),
                 y.max() - y.min(),
                 z.max() - z.min()]).max() / 2.0
            mid_x = (x.max() + x.min()) * 0.5
            mid_y = (y.max() + y.min()) * 0.5
            mid_z = (z.max() + z.min()) * 0.5
            ax.set_xlim(mid_x - max_range, mid_x + max_range)
            ax.set_ylim(mid_y - max_range, mid_y + max_range)
            ax.set_zlim(mid_z - max_range, mid_z + max_range)

            ax.set_xlabel(r'$x\>[Mpc/h]$', fontsize=16)
            ax.set_ylabel(r'$y\>[Mpc/h]$', fontsize=16)
            ax.set_zlabel(r'$z\>[Mpc/h]$', fontsize=16)

            ax2.scatter(theta_sky, phi_sky, c='k', alpha=0.2)
            ax2.set_xlabel(r'$\theta\>[\mathrm{arsec}]$', fontsize=16)
            ax2.set_ylabel(r'$\phi\>[\mathrm{arcsec}]$', fontsize=16)
            plt.savefig('{}/nfw_particles.png'.format(vis_output_dir), dpi=300)

        # write out all to binary
        x.astype('f').tofile('{}/x.bin'.format(output_dir))
        y.astype('f').tofile('{}/y.bin'.format(output_dir))
        z.astype('f').tofile('{}/z.bin'.format(output_dir))
        theta_sky.astype('f').tofile('{}/theta.bin'.format(output_dir))
        phi_sky.astype('f').tofile('{}/phi.bin'.format(output_dir))
        redshift.astype('f').tofile('{}/redshift.bin'.format(output_dir))

        if (fov_multiplier is not None):
            fov_size = fov_multiplier * np.max(self.r)
        else:
            # the halo prop file records half the radius of the square FOV, which sets the scale for the density
            # estimation... we don't want the FOV to include any space outside of the region we have populated with
            # halos, else the density estiamtion will plummet at the boundary. Above, we populated the halo with
            # particles out to rfrac * r200c. The largest square that can fit inside the projection of this NFW sphere
            # then has a side length of 2*(rfrac*r200c)/sqrt(2) --> radius = (rfrac*r200c)/sqrt(2).
            # Replace rfrac*r200c by the radial distance to the furthest particle and trim by 5%, to be safe.
            # Also note that self.r is a comoving distance, which is correct
            fov_size = 0.95 * (np.max(r_fov) / np.sqrt(2))
        self._write_prop_file(fov_size, output_dir)

    # -----------------------------------------------------------------------------------------------

    def _write_prop_file(self, fov_radius, output_dir):
        """
        Writes a csv file contining the halo properties needed by this package's ray tracing modules
        The boxRadius can really be anything, since the space around the NFW ball is empty-- here, we
        set it to correspond to a transverse comoving distance equal to R*r200 at the redshift of the
        halo.

        Parameters
        ----------
        fov_radius : float
            Half of the square FOV side length (this scale will be used in later calls to the density estaimtor)

        output_dir : string, optional
            The desired output location for the property file. Defaults to a subdir created at the 
            location of this module.
        """

        # find the angular scale corresponding to fov_r200c * r200c in proper Mpc at the redshift of the halo
        boxRadius_Mpc = fov_radius
        trans_Mpc_per_arcsec = (self.cosmo.kpc_proper_per_arcmin(
            self.redshift).value / 1e3) / 60 * (self.redshift + 1)
        boxRadius_arcsec = boxRadius_Mpc / trans_Mpc_per_arcsec

        cols = '#halo_redshift, sod_halo_mass, sod_halo_radius, sod_halo_cdelta, sod_halo_cdelta_error, '\
               'halo_lc_x, halo_lc_y, halo_lc_z, boxRadius_Mpc, boxRadius_arcsec, mpp'
        props = np.array([
            self.redshift, self.m200c, self.r200c, self.c, self.c_err, 0, 0, 0,
            boxRadius_Mpc, boxRadius_arcsec, self.mpp
        ])

        np.savetxt(
            '{}/properties.csv'.format(output_dir), [props],
            fmt='%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f,%.6f',
            delimiter=',',
            header=cols)