Merge pull request #43 from cgannonucm/import_galacticus_subhalos

Realization from Galacticus output

pyHalo/Halos/HaloModels/TNFWFromParams.py

@@ -0,0 +1,97 @@
+from pyHalo.Halos.halo_base import Halo
+from lenstronomy.LensModel.Profiles.tnfw import TNFW as TNFWLenstronomy
+import numpy as np
+from pyHalo.Halos.tnfw_halo_util import tnfw_mass_fraction
+from pyHalo.Halos.HaloModels.TNFW import TNFWSubhalo
+from lenstronomy.LensModel.Profiles.tnfw import TNFW
+
+class TNFWFromParams(TNFWSubhalo):
+    """
+    Creates a TNFW halo based on physical params. 
+    """
+
+    KEY_RT = "r_trunc_kpc"
+    KEY_RS = "rs"
+    KEY_RHO_S = "rho_s"
+    KEY_RV = "rv"
+
+    def __init__(self, mass, x_kpc, y_kpc, r3d, z,
+                 sub_flag, lens_cosmo_instance, args, unique_tag=None):
+        """
+        Defines a TNFW subhalo with physical params r_trunc_kpc, rs, rhos passed in the args argument
+        """
+
+        self._lens_cosmo = lens_cosmo_instance
+
+        self._kpc_per_arcsec_at_z = self._lens_cosmo.cosmo.kpc_proper_per_asec(z)
+
+        x = x_kpc / self._kpc_per_arcsec_at_z
+
+        y = y_kpc / self._kpc_per_arcsec_at_z
+
+        keys_physical = (self.KEY_RV,self.KEY_RS,self.KEY_RHO_S,self.KEY_RV,self.KEY_RT)
+        self._params_physical = {key:args[key] for key in keys_physical}
+
+        self._c = self._params_physical[self.KEY_RV] / self._params_physical[self.KEY_RS]
+
+        super(TNFWFromParams, self).__init__(mass,x,y,r3d,z,sub_flag,lens_cosmo_instance,args,None,None,unique_tag)
+
+    def density_profile_3d(self, r, params_physical=None):
+        """
+        Computes the 3-D density profile of the halo
+        :param r: distance from center of halo [kpc]
+        :return: the density profile in units M_sun / kpc^3
+        """
+
+        _params = self._params_physical if params_physical is None else params_physical
+
+        r_t = _params[self.KEY_RT]
+        r_s = _params[self.KEY_RS]
+        rho_s = _params[self.KEY_RHO_S]
+
+        n = 1
+
+        x = r / r_s
+        tau = r_t / r_s
+
+        density_nfw = (rho_s / ((x)*(1+x)**2))
+
+        return density_nfw * (tau**2 / (tau**2 + x**2))**n 
+
+
+    @property
+    def profile_args(self):
+        """
+        See documentation in base class (Halos/halo_base.py)
+        """
+        if not hasattr(self, '_profile_args'):
+            truncation_radius_kpc = self.params_physical[self.KEY_RT]
+            self._profile_args = (self.c, truncation_radius_kpc)
+        return self._profile_args
+
+
+    @property
+    def lenstronomy_params(self):
+        """
+        See documentation in base class (Halos/halo_base.py)
+        """
+        if not hasattr(self, '_kwargs_lenstronomy'):
+
+            r_t = self.params_physical[self.KEY_RT]
+            r_s = self.params_physical[self.KEY_RS]
+            rho_s = self.params_physical[self.KEY_RHO_S]
+
+            Rs_angle, theta_Rs = self.nfw_physical2angle_fromNFWparams(rho_s,r_s,self.z)
+
+            x, y = np.round(self.x, 4), np.round(self.y, 4)
+
+            Rs_angle = np.round(Rs_angle, 10)
+            theta_Rs = np.round(theta_Rs, 10)
+            r_trunc_arcsec = r_t / self._lens_cosmo.cosmo.kpc_proper_per_asec(self.z)
+
+            kwargs = [{'alpha_Rs': self._rescale_norm * theta_Rs, 'Rs': Rs_angle,
+                      'center_x': x, 'center_y': y, 'r_trunc': r_trunc_arcsec}]
+
+            self._kwargs_lenstronomy = kwargs
+
+        return self._kwargs_lenstronomy, None

pyHalo/Halos/galacticus_util/galacticus_filter.py

@@ -0,0 +1,97 @@
+from pyHalo.Halos.galacticus_util.galacticus_util import GalacticusUtil
+import numpy as np
+
+def nodedata_apply_filter(nodedata,filter,galacticus_util = None):
+    """Takes a dictionary with numpy arrays as values and apply as boolean filter to all.
+     Returns a dictionary with same keys, but a filter applied to all arrays."""
+
+    return {key:val[filter] for key,val in nodedata.items()}
+
+def nodedata_filter_tree(nodedata, treenum,galacticus_util = None):
+    """Returns a filter that excludes all nodes but nodes in the specified tree"""
+    gutil = GalacticusUtil() if galacticus_util is None else galacticus_util
+    return nodedata[gutil.PARAM_TREE_INDEX] == treenum
+
+def nodedata_filter_subhalos(nodedata,galacticus_util= None):
+    """Returns a filter that excludes all but subhalos (excludes host halos)"""
+    gutil = GalacticusUtil() if galacticus_util is None else galacticus_util
+    return nodedata[gutil.PARAM_ISOLATED] == 0
+
+def nodedata_filter_halos(nodedata,galacticus_util = None):
+    """Returns a filter that excludes all but halo (excludes sub-halos)"""
+    return np.logical_not(nodedata_filter_subhalos(nodedata))
+
+def nodedata_filter_range(nodedata,range,key,galacticus_util=None):
+    """Returns a filter that excludes nodes not within the given range for a specified parameter"""
+    return (nodedata[key] > range[0]) & (nodedata[key] < range[1]) 
+
+def nodedata_filter_virialized(nodedata,galacticus_util= None):
+    """
+    Returns a filter that excludes everything outside of the host halos virial radius
+    WARNING: Current implementation only works if there is only one host halo per tree,
+    IE we are looking at the last output from galacticus. 
+    """
+    gutil = GalacticusUtil() if galacticus_util is None else galacticus_util
+
+    #Get radial position of halos
+    rvec = np.asarray((nodedata[gutil.PARAM_X],nodedata[gutil.PARAM_Y],nodedata[gutil.PARAM_Z]))
+    r = np.linalg.norm(rvec,axis=0)
+
+    #Filter halos and get there virial radii
+    filtered_halos = nodedata_filter_halos(nodedata)
+    rv_halos = nodedata[gutil.PARAM_RADIUS_VIRIAL][filtered_halos]
+    halo_output_n = nodedata[gutil.PARAM_TREE_ORDER]
+
+    filter_virialized = r < rv_halos[halo_output_n]
+
+    return filter_virialized
+
+def nodedata_filter_r2d(nodedata,r2d_max,plane_normal,
+                        galacticus_util= None):
+
+    """
+    Filters based on projected radii.
+    """
+    gutil = GalacticusUtil() if galacticus_util is None else galacticus_util
+
+    r = np.asarray((nodedata[gutil.PARAM_X],nodedata[gutil.PARAM_Y],nodedata[gutil.PARAM_Z]))
+
+    r2d = project_r2d(*r,plane_normal)
+
+    return r2d < r2d_max
+
+def project_r2d(x,y,z,plane_normal):
+    """
+    Takes in arrays of coordinates, calculates the projected radius on the plane. 
+     
+
+    :param x: An array of x coordinates
+    :param y: An array of y coordinates
+    :param z: An array of z coordinates
+    :plane_normal: Normal vector of the plane to project onto
+    """
+    coords = np.asarray((x,y,z))
+
+    #Reshape so if scalers are passed for x,y,z we do not encounter issue
+    if coords.ndim == 1:
+        coords.reshape((3,1))
+
+    #convert to unit normal
+    un = plane_normal / np.linalg.norm(plane_normal)
+
+    #Project distance to plane
+    #Projected distance to plane is sqrt(r.r - (r.un)^2)
+    rdotr = np.linalg.norm(coords,axis=0)**2
+    rdotun = np.dot(un,coords)
+    return np.sqrt(rdotr - rdotun**2)
+
+
+
+
+
+
+
+
+
+
+