<commit_before>#!/usr/bin/env python

from __future__ import division

import math

import numpy as np
from numpy.lib import scimath as sm
import scipy.constants
import scipy.optimize as opt

import astropy.cosmology
from astropy import units as u


def arcsec(z):
    """Compute the inverse sec of the complex number z."""
    val1 = 1j / z
    val2 = sm.sqrt(1 - 1./z**2)
    val = 1.j * np.log(val2 + val1)
    return 0.5 * np.pi + val


def unit_checker(x, unit):
    """Check that x has units u. Convert to appropriate units if not.

    Arguments:
    - `x`: array_like
    - `u`: astro.units unit
    """
    if not isinstance(x, u.Quantity):
        return x * unit
    return x.to(unit)


class InvalidNFWValue(Exception):
    pass


class NFW(object):
    """Compute properties of an NFW halo.

    Required inputs are

    size - radius or mass of the halo in Mpc or M_sun
    c - concentration (value|"duffy|dolag")
    z - halo redshift

    optional input

    size_type - "(radius|mass)" specifies whether the halo size is given as
                radius or mass
    overdensity - the factor above the critical/mean density of the Universe
                  at which mass/radius are computed. Default 200
    overdensity_type = "(critical|mean)"
    cosmology - object, use the current astropy.cosmology if None, otherwise
                an astropy.cosmology object
    """

    def __init__(self, size, c, z, size_type="mass",
                 overdensity=200.,
                 overdensity_type="critical", cosmology=None):
        if overdensity_type != "critical":
            print "You must be kidding."
            print "Of course overdensities are defined wrt the " \
                "critical density."

        self._c = float(c)
        self._z = float(z)
        self._overdensity = overdensity
        if cosmology is not None:
            self._cosmology = cosmology
            self._var_cosmology = False
        else:
            self._cosmology = astropy.cosmology.get_current()
            self._var_cosmology = True

        self._rho_c = None
        self._r_s = None
        self._r_Delta = None
        self._rho_c = self.rho_c

        if size_type not in ['mass', 'radius']:
            raise InvalidNFWValue(size_type)
        if size_type == "mass":
            self._size = unit_checker(size, u.solMass)
        else:
            self._size = unit_checker(size, u.megaparsec)
        self._size_type = size_type

        self._r_Delta = self.r_Delta
        self._r_s = self.r_s

        return

    def _update_required(self, attr):
        """
        Check whether an attr needs updating due to a new cosmology
        Arguments:
        - `attr`: an instance attribute
        """
        if attr is None:
            return True
        if not self._var_cosmology:
            return False
        if self._cosmology == astropy.cosmology.get_current():
            return False
        return True

    @property
    def c(self):
        """Halo concentration"""
        return self._c

    @property
    def z(self):
        """Halo redshift"""
        return self._z

    @property
    def delta_c(self):
        """Characterstic overdensity
        """
        return 200/3 * self.c**3 / (np.log(1 + self.c) - self.c/(1. + self.c))

    @property
    def rho_c(self):
        """Critical density at halo redshift
        """
        if self._update_required(self._rho_c):
            if self._var_cosmology:
                self._cosmology = astropy.cosmology.get_current()
            self._rho_c = self._cosmology.critical_density(self.z)
            self._rho_c = self._rho_c.to(u.solMass / u.megaparsec**3)
            if self._r_Delta is not None:
                self._r_Delta = self.r_Delta
            if self._r_s is not None:
                self._r_s = self.r_s
        return self._rho_c

    @property
    def r_Delta(self):
        """Halo radius at initialization overdensity
        """
        if self._size_type == "mass":
            if self._update_required(self._rho_c):
                self._rho_c = self.rho_c
            self._r_Delta = (3*self._size
                             / (4 * np.pi * self._overdensity
                                * self._rho_c))**(1/3)
        else:
            self._r_Delta = self._size
        return self._r_Delta

    @property
    def r_s(self):
        """Scale radius
        """
        if self._update_required(self._r_s):
            self._r_Delta = self.r_Delta
            self._r_s = self._rDelta2rs(self._r_Delta, self._overdensity)
            if self._var_cosmology:
                self._cosmology = astropy.cosmology.get_current()
        return self._r_s

    def _rDelta2r200_zero(self, rs, r_Delta, overdensity):
        rs *= u.megaparsec
        z = (overdensity/3 * r_Delta**3) - (self.delta_c * rs**3) \
            * (math.log((rs+r_Delta) / rs) - (r_Delta / (rs+r_Delta)))
        return z.value

    def _rDelta2rs(self, r_Delta, overdensity):
        r200 = opt.brentq(self._rDelta2r200_zero, 1e-6, 100,
                          args=(r_Delta, overdensity))
        return r200 * u.megaparsec

    def __str__(self):
        prop_str = "NFW halo with concentration %.2g at redshift %.2f:\n\n" \
                   % (self.c, self.z,)
        for delta in (200, 500, 2500):
            prop_str += "M_%d = %.2e M_sun\tr_%d = %.2g Mpc\n" % \
                        (delta, self.mass_Delta(delta).value, delta,
                         self.radius_Delta(delta).value)
        return prop_str

    def _mean_density_zero(self, r, Delta):
        return (self.mean_density(r) - Delta*self.rho_c).value

    def radius_Delta(self, Delta):
        """Find the radius at which the mean density is Delta times the
        critical density. Returns radius in Mpc."""
        x0 = opt.brentq(self._mean_density_zero, 1e-6, 10, args=(Delta,))
        return x0 * u.megaparsec

    def mass_Delta(self, Delta):
        """Find the mass inside a radius inside which the mean density
        is Delta times the critical density. Returns mass in M_sun."""
        r = self.radius_Delta(Delta)
        return self.mass(r)<commit_msg>add __repr__ to NFW<commit_after>#!/usr/bin/env python

from __future__ import division

import math

import numpy as np
from numpy.lib import scimath as sm
import scipy.constants
import scipy.optimize as opt

import astropy.cosmology
from astropy import units as u


def arcsec(z):
    """Compute the inverse sec of the complex number z."""
    val1 = 1j / z
    val2 = sm.sqrt(1 - 1./z**2)
    val = 1.j * np.log(val2 + val1)
    return 0.5 * np.pi + val


def unit_checker(x, unit):
    """Check that x has units u. Convert to appropriate units if not.

    Arguments:
    - `x`: array_like
    - `u`: astro.units unit
    """
    if not isinstance(x, u.Quantity):
        return x * unit
    return x.to(unit)


class InvalidNFWValue(Exception):
    pass


class NFW(object):
    """Compute properties of an NFW halo.

    Required inputs are

    size - radius or mass of the halo in Mpc or M_sun
    c - concentration (value|"duffy|dolag")
    z - halo redshift

    optional input

    size_type - "(radius|mass)" specifies whether the halo size is given as
                radius or mass
    overdensity - the factor above the critical/mean density of the Universe
                  at which mass/radius are computed. Default 200
    overdensity_type = "(critical|mean)"
    cosmology - object, use the current astropy.cosmology if None, otherwise
                an astropy.cosmology object
    """

    def __init__(self, size, c, z, size_type="mass",
                 overdensity=200.,
                 overdensity_type="critical", cosmology=None):
        if overdensity_type != "critical":
            print "You must be kidding."
            print "Of course overdensities are defined wrt the " \
                "critical density."

        self._c = float(c)
        self._z = float(z)
        self._overdensity = overdensity
        if cosmology is not None:
            self._cosmology = cosmology
            self._var_cosmology = False
        else:
            self._cosmology = astropy.cosmology.get_current()
            self._var_cosmology = True

        self._rho_c = None
        self._r_s = None
        self._r_Delta = None
        self._rho_c = self.rho_c

        if size_type not in ['mass', 'radius']:
            raise InvalidNFWValue(size_type)
        if size_type == "mass":
            self._size = unit_checker(size, u.solMass)
        else:
            self._size = unit_checker(size, u.megaparsec)
        self._size_type = size_type

        self._r_Delta = self.r_Delta
        self._r_s = self.r_s

        return

    def _update_required(self, attr):
        """
        Check whether an attr needs updating due to a new cosmology
        Arguments:
        - `attr`: an instance attribute
        """
        if attr is None:
            return True
        if not self._var_cosmology:
            return False
        if self._cosmology == astropy.cosmology.get_current():
            return False
        return True

    @property
    def c(self):
        """Halo concentration"""
        return self._c

    @property
    def z(self):
        """Halo redshift"""
        return self._z

    @property
    def delta_c(self):
        """Characterstic overdensity
        """
        return 200/3 * self.c**3 / (np.log(1 + self.c) - self.c/(1. + self.c))

    @property
    def rho_c(self):
        """Critical density at halo redshift
        """
        if self._update_required(self._rho_c):
            if self._var_cosmology:
                self._cosmology = astropy.cosmology.get_current()
            self._rho_c = self._cosmology.critical_density(self.z)
            self._rho_c = self._rho_c.to(u.solMass / u.megaparsec**3)
            if self._r_Delta is not None:
                self._r_Delta = self.r_Delta
            if self._r_s is not None:
                self._r_s = self.r_s
        return self._rho_c

    @property
    def r_Delta(self):
        """Halo radius at initialization overdensity
        """
        if self._size_type == "mass":
            if self._update_required(self._rho_c):
                self._rho_c = self.rho_c
            self._r_Delta = (3*self._size
                             / (4 * np.pi * self._overdensity
                                * self._rho_c))**(1/3)
        else:
            self._r_Delta = self._size
        return self._r_Delta

    @property
    def r_s(self):
        """Scale radius
        """
        if self._update_required(self._r_s):
            self._r_Delta = self.r_Delta
            self._r_s = self._rDelta2rs(self._r_Delta, self._overdensity)
            if self._var_cosmology:
                self._cosmology = astropy.cosmology.get_current()
        return self._r_s

    def _rDelta2r200_zero(self, rs, r_Delta, overdensity):
        rs *= u.megaparsec
        z = (overdensity/3 * r_Delta**3) - (self.delta_c * rs**3) \
            * (math.log((rs+r_Delta) / rs) - (r_Delta / (rs+r_Delta)))
        return z.value

    def _rDelta2rs(self, r_Delta, overdensity):
        r200 = opt.brentq(self._rDelta2r200_zero, 1e-6, 100,
                          args=(r_Delta, overdensity))
        return r200 * u.megaparsec

    def __str__(self):
        prop_str = "NFW halo with concentration %.2g at redshift %.2f:\n\n" \
                   % (self.c, self.z,)
        for delta in (200, 500, 2500):
            prop_str += "M_%d = %.2e M_sun\tr_%d = %.2g Mpc\n" % \
                        (delta, self.mass_Delta(delta).value, delta,
                         self.radius_Delta(delta).value)
        return prop_str

    def __repr__(self):
        return self.__str__()

    def _mean_density_zero(self, r, Delta):
        return (self.mean_density(r) - Delta*self.rho_c).value

    def radius_Delta(self, Delta):
        """Find the radius at which the mean density is Delta times the
        critical density. Returns radius in Mpc."""
        x0 = opt.brentq(self._mean_density_zero, 1e-6, 10, args=(Delta,))
        return x0 * u.megaparsec

    def mass_Delta(self, Delta):
        """Find the mass inside a radius inside which the mean density
        is Delta times the critical density. Returns mass in M_sun."""
        r = self.radius_Delta(Delta)
        return self.mass(r)
{'fraction_non_alphanumeric': 0.0993939393939394, 'fraction_numerical': 0.011797979797979799, 'mean_word_length': 1.8768014876801489, 'pattern_counts': {'":': 6, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '837851', 'n_tokens_mistral': 3970, 'n_tokens_neox': 3701, 'n_words': 1030}