예제 #1
0
 def T_nocut(self, r):
     """
     Calculates the temperature as a function of r ignoring any interior
     cutoff.  The interior cutoff would affect T for adiabatic disks
     
     Paramters
     ---------
     
     r : array, SimArray, or float
         radius (or radii) at which to calculate the temperature
     
     Returns
     -------
     
     T : SimArray
         temperature
     """
     # Load settings
     params = self._parent.settings.physical
     
     if not hasattr(params, 'kind'):
         # Add this check for backwards compatibility.  Previous versions
         # only had one kind of temperature profile
         params.kind = 'powerlaw'
         
     T0 = params.T0
     Tmin = params.Tmin
     r0 = params.r0
     kind = params.kind
     
     # Calculate T(r)
     r = match_units(r, r0)[0]
     a = (r/r0)
     a = match_units(a, '1')[0]
     
     # Powerlaw temperature (default)
     if kind == 'powerlaw':
         
         Tpower = params.Tpower
         Tout = T0 * np.power(a, Tpower)
         Tout[Tout < Tmin] = Tmin
         
     # MQWS temperature profile
     elif (kind == 'mqws') | (kind == 'MQWS'):
         
         # NOTE: I'm not sure how exactly they generate the temperature
         # profile.  The quoted equation doesn't match their figures
         Tout = T0 * np.exp(-3*a/2) + Tmin
         
     else:
         
         raise TypeError, 'Could not find temperature kind {0}'.format(kind)
         
     # Apply max Temp cutoff
     if hasattr(params, 'Tmax'):
         Tmax = params.Tmax
         Tout[Tout > Tmax] = Tmax
     
     return Tout
예제 #2
0
    def _make_sigma(self, r_bins, sigmaBinned):
        """
        Generates the surface density as a function of r, a callable object 
        sigma(r) and assigns it to self.sigma
        
        Generates a spline interpolation of sigma vs r from the file
        defined by settings.filenames.sigmaFileName.  Returns sigma vs r as
        an cubic spline interpolation object
        (see scipy.interpolation.InterpolatedUnivariateSpline)
        
        sigma_input should be a pickled dictionary with the entries:
        'sigma': <sigma evaluated at r>
        'r':     <r for the bins>
        
        If the input sigma has units, sigma vs r will be returned in units
        of Msol/au^2
        """       
        # Convert to default units of Msol/au^2.  If no units, assign default
        sigmaBinned = match_units(sigmaBinned, 'Msol au**-2')[0]
        # Convert r_bins to default units of 'au'
        r_bins = match_units(r_bins, 'au')[0]
        # Calculate spline interpolation
        sigspline = spline(r_bins, sigmaBinned, k=3, ext='zeros')
#        print 'Calculating spline interpolation (slow for many data points)'
#        sigspline = interp1d(r_bins,sigmaBinned,kind='cubic',fill_value=0.0,\
#        bounds_error=False)
        
        def sigout(r):
            """
            Linear spline interpolation of sigma(r).  
            
            ARGUMENTS:
            
            r - can be scalar, numpy array, or sim array
            
            RETURNS: 
            
            sigma (surface density) evaluated at r
            """
            
            # Try to convert r to the units used to make sigspline ('au')
            r = match_units(r, 'au')[0]
            
            return SimArray(sigspline(r), 'Msol au**-2')
        
        self.sigma = sigout
        self.r_bins = r_bins
예제 #3
0
    def _make_sigma(self, r_bins, sigmaBinned):
        """
        Generates the surface density as a function of r, a callable object 
        sigma(r) and assigns it to self.sigma
        
        Generates a spline interpolation of sigma vs r from the file
        defined by settings.filenames.sigmaFileName.  Returns sigma vs r as
        an cubic spline interpolation object
        (see scipy.interpolation.InterpolatedUnivariateSpline)
        
        sigma_input should be a pickled dictionary with the entries:
        'sigma': <sigma evaluated at r>
        'r':     <r for the bins>
        
        If the input sigma has units, sigma vs r will be returned in units
        of Msol/au^2
        """
        # Convert to default units of Msol/au^2.  If no units, assign default
        sigmaBinned = match_units(sigmaBinned, 'Msol au**-2')[0]
        # Convert r_bins to default units of 'au'
        r_bins = match_units(r_bins, 'au')[0]
        # Calculate spline interpolation
        sigspline = spline(r_bins, sigmaBinned, k=3, ext='zeros')

        def sigout(r):
            """
            Linear spline interpolation of sigma(r).  
            
            ARGUMENTS:
            
            r - can be scalar, numpy array, or sim array
            
            RETURNS: 
            
            sigma (surface density) evaluated at r
            """

            # Try to convert r to the units used to make sigspline ('au')
            r = match_units(r, 'au')[0]

            return SimArray(sigspline(r), 'Msol au**-2')

        self.sigma = sigout
        self.r_bins = r_bins
예제 #4
0
def init_snapshot(IC):
    """
    Initialize a snapshot for the IC object.  Requires that positions have
    been created.  Also sets:
     * pos
     * metals
     * temp
     * mass
     * star eps
    
    Parameters
    ----------
    IC : ICobj
    
    Returns
    -------
    snapshot : SimSnap
    """
    # Get required settings from IC
    settings = IC.settings
    # particle positions
    r = IC.pos.r
    xyz = IC.pos.xyz
    nParticles = IC.pos.nParticles
    m_star = settings.physical.M
    m_disk = IC.sigma.m_disk
    m_disk = match_units(m_disk, m_star)[0]
    m_particles = m_disk / float(nParticles)
    metals = settings.snapshot.metals
    # re-scale the particles (allows making of lo-mass disk)
    m_particles *= settings.snapshot.mScale
    
    # Handle units
    units = setup_units(m_star, r)
    
    if xyz.units != r.units:
        
        xyz.convert_units(units['x'])
    
    # Initialize arrays
    snapshot = pynbody.new(star=1,gas=nParticles)
    snapshot['vel'].units = units['v']
    snapshot['eps'] = SimArray(0.01, units['x'])
    snapshot['rho'] = 0.
    snapshot['metals'] = metals
    # Assign array values
    snapshot.gas['pos'] = xyz
    snapshot.gas['temp'] = IC.T(r)
    snapshot.gas['mass'] = m_particles
    
    snapshot.star['pos'] = 0.
    snapshot.star['mass'] = m_star
    # Estimate the star's softening length as the closest particle distance/2
    snapshot.star['eps'] = r.min()/2.
    
    return snapshot
예제 #5
0
def init_snapshot(IC):
    """
    Initialize a snapshot for the IC object.  Requires that positions have
    been created.  Also sets:
     * pos
     * metals
     * temp
     * mass
     * star eps
    
    Parameters
    ----------
    IC : ICobj
    
    Returns
    -------
    snapshot : SimSnap
    """
    # Get required settings from IC
    settings = IC.settings
    # particle positions
    r = IC.pos.r
    xyz = IC.pos.xyz
    nParticles = IC.pos.nParticles
    m_star = settings.physical.M
    m_disk = IC.sigma.m_disk
    m_disk = match_units(m_disk, m_star)[0]
    m_particles = m_disk / float(nParticles)
    metals = settings.snapshot.metals
    # re-scale the particles (allows making of lo-mass disk)
    m_particles *= settings.snapshot.mScale
    
    # Handle units
    units = setup_units(m_star, r)
    
    if xyz.units != r.units:
        
        xyz.convert_units(units['x'])
    
    # Initialize arrays
    snapshot = pynbody.new(star=1,gas=nParticles)
    snapshot['vel'].units = units['v']
    snapshot['eps'] = SimArray(0.01, units['x'])
    snapshot['rho'] = 0.
    snapshot['metals'] = metals
    # Assign array values
    snapshot.gas['pos'] = xyz
    snapshot.gas['temp'] = IC.T(r)
    snapshot.gas['mass'] = m_particles
    
    snapshot.star['pos'] = 0.
    snapshot.star['mass'] = m_star
    # Estimate the star's softening length as the closest particle distance/2
    snapshot.star['eps'] = r.min()/2.
    
    return snapshot
예제 #6
0
 def pdf_fcn(r_in):
     """
     Normalized cubic spline interpolation of the PDF(r) from sigma(r).
     The PDF is just calculated as 2*pi*r*sigma(r).
     
     ARGUMENTS:
     
     r_in - radii at which to calculate the PDF
     
     RETURNS:
     
     probability density function from sigma(r) evaluated at r_in
     """
     # Put r_in into the units used in generating the pdf
     r_in = match_units(r_in, self.r_bins)[0]
     # Evaluate the pdf at r_in
     pdf_vals = pdfSpline(r_in)
     # Put the pdf into units of r_in.units**-1
     pdf_vals = match_units(pdf_vals, 1/r_in)[0]
     
     return pdf_vals
예제 #7
0
        def pdf_fcn(r_in):
            """
            Normalized cubic spline interpolation of the PDF(r) from sigma(r).
            The PDF is just calculated as 2*pi*r*sigma(r).
            
            ARGUMENTS:
            
            r_in - radii at which to calculate the PDF
            
            RETURNS:
            
            probability density function from sigma(r) evaluated at r_in
            """
            # Put r_in into the units used in generating the pdf
            r_in = match_units(r_in, self.r_bins)[0]
            # Evaluate the pdf at r_in
            pdf_vals = pdfSpline(r_in)
            # Put the pdf into units of r_in.units**-1
            pdf_vals = match_units(pdf_vals, 1 / r_in)[0]

            return pdf_vals
예제 #8
0
    def _disk_mass(self):
        """
        Calculate the total disk mass by integrating sigma
        """
        # Assign variables
        r = self.r_bins
        sig = self.sigma(r)
        # Now integrate
        m_disk = simps(2 * np.pi * r * sig, r)
        m_units = sig.units * (r.units)**2
        m_disk = match_units(m_disk, m_units)[0]

        self.m_disk = m_disk
예제 #9
0
 def _disk_mass(self):
     """
     Calculate the total disk mass by integrating sigma
     """
     # Assign variables
     r = self.r_bins
     sig = self.sigma(r)
     # Now integrate
     m_disk = simps(2*np.pi*r*sig, r)
     m_units = sig.units * (r.units)**2
     m_disk = match_units(m_disk, m_units)[0]
     
     self.m_disk = m_disk
예제 #10
0
    def T_adiabatic(self, r):
        """
        Estimates the adiabatic temperature profile as a function of r
        
        setup_interior must be run first
        """
#        if not self.adiabatic_ready:
#            
#            # Try to setup the adiabatic profile
#            self.setup_interior()
            
#        p = self.p
#        A = self.Tscale
#        print A
#        
#        r = match_units(r, 'au')[0]
#        sigma = self._parent.sigma(r)
#        sigma.convert_units('Msol au**-2')
#        sigma = strip_units(sigma)
#        r = strip_units(r)
#        return A * ((sigma**2)/(r**3))**p
            
        b = 1.5
        
        sigma = self._parent.sigma(r)
        r_a = self.r_a
        r = match_units(r, r_a.units)[0]
        x = (r/r_a).in_units('1')
        sigma_a = self.sigma_a
        y = (sigma/sigma_a).in_units('1')
        
        x = strip_units(x)
        y = strip_units(y)
        
        gamma = self._parent.settings.physical.gamma
        p = (gamma-1.)/(gamma+1.)
        
        T_a = self.T_a
        
        T = T_a * (y**(2*p)) * (x**-b)
        
        return T
        
#        r = strip_units(match_units(r,'au')[0])
#        
#        return A * (r**2)
예제 #11
0
 def sigout(r):
     """
     Linear spline interpolation of sigma(r).  
     
     ARGUMENTS:
     
     r - can be scalar, numpy array, or sim array
     
     RETURNS: 
     
     sigma (surface density) evaluated at r
     """
     
     # Try to convert r to the units used to make sigspline ('au')
     r = match_units(r, 'au')[0]
     
     return SimArray(sigspline(r), 'Msol au**-2')
예제 #12
0
        def sigout(r):
            """
            Linear spline interpolation of sigma(r).  
            
            ARGUMENTS:
            
            r - can be scalar, numpy array, or sim array
            
            RETURNS: 
            
            sigma (surface density) evaluated at r
            """

            # Try to convert r to the units used to make sigspline ('au')
            r = match_units(r, 'au')[0]

            return SimArray(sigspline(r), 'Msol au**-2')
예제 #13
0
def sigma(snapshot, bins=100, cmFlag=True):
    """Calculates surface density vs r (relative to the center of mass)
    
    Parameters
    ----------
    
    snapshot : tipsy snapshot
    bins : int, list, array...
        (optional) Either the number of bins to use or the binedges to use
    cmFlag : bool
        (optional) Calculate relative to the center of mass
    
    Returns
    -------
    
    sigma : SimArray
        Surface density as a function of r
    r_bins : SimArray
        Radial bin edges
    
    """
    
    if cmFlag:
        
        # Begin by subtracting off the center of mass position
        cm = (snapshot['mass'][:,None] * snapshot['pos']).sum()/(snapshot['mass'].sum())
        snapshot['pos'] -= cm
        
    r = snapshot.g['rxy']
    # particle mass
    m_gas = snapshot.gas['mass'][[0]]

    N, r_bins = np.histogram(r, bins=bins)
    r_bins = match_units(r_bins, r.units)[0]
    r_center = (r_bins[1:] + r_bins[0:-1])/2
    dr = r_bins[[1]] - r_bins[[0]]

    sig = N*m_gas/(2*np.pi*r_center*dr)
    
    if cmFlag:
        
        # Add star position back to positions
        snapshot['pos'] += cm

    return sig, r_bins
예제 #14
0
def sigma(snapshot, bins=100, cmFlag=True):
    """Calculates surface density vs r (relative to the center of mass)
    
    Parameters
    ----------
    
    snapshot : tipsy snapshot
    bins : int, list, array...
        (optional) Either the number of bins to use or the binedges to use
    cmFlag : bool
        (optional) Calculate relative to the center of mass
    
    Returns
    -------
    
    sigma : SimArray
        Surface density as a function of r
    r_bins : SimArray
        Radial bin edges
    
    """
    
    if cmFlag:
        
        # Begin by subtracting off the center of mass position
        cm = (snapshot['mass'][:,None] * snapshot['pos']).sum()/(snapshot['mass'].sum())
        snapshot['pos'] -= cm
        
    r = snapshot.g['rxy']
    # particle mass
    m_gas = snapshot.gas['mass'][[0]]

    N, r_bins = np.histogram(r, bins=bins)
    r_bins = match_units(r_bins, r.units)[0]
    r_center = (r_bins[1:] + r_bins[0:-1])/2
    dr = r_bins[[1]] - r_bins[[0]]

    sig = N*m_gas/(2*np.pi*r_center*dr)
    
    if cmFlag:
        
        # Add star position back to positions
        snapshot['pos'] += cm

    return sig, r_bins
예제 #15
0
        def finv_fcn(m_in):
            """
            The inverse CDF for sigma(r).
            
            input:
            
            0 <= m_in < 1
            
            returns:
            
            r (radius), the inverse CDF evaluated at m_in
            
            Uses a linear spline interpolation.
            """

            r_out = finv(m_in)
            r_out = match_units(r_out, r)[0]
            return r_out
예제 #16
0
 def finv_fcn(m_in):
     """
     The inverse CDF for sigma(r).
     
     input:
     
     0 <= m_in < 1
     
     returns:
     
     r (radius), the inverse CDF evaluated at m_in
     
     Uses a linear spline interpolation.
     """
     
     r_out = finv(m_in)
     r_out = match_units(r_out, r)[0]
     return r_out
예제 #17
0
def MQWS(settings, T):
    """
    Generates a surface density profile as the per method used in Mayer, Quinn,
    Wadsley, and Stadel 2004
    
    ** ARGUMENTS **
    NOTE: if units are not supplied, assumed units are AU, Msol
    
    settings : IC settings
        settings like those contained in an IC object (see ICgen_settings.py)
        
    T : callable
        A function to calculate temperature as a function of radius
        
    ** RETURNS **
    
    r : SimArray
        Radii at which sigma is calculated
    sigma : SimArray
        Surface density profile as a function of R
    """
    # Q calculation parameters:
    G = SimArray([1.0], 'G')
    kB = SimArray([1.0], 'k')

    # Load in settings
    n_points = settings.sigma.n_points
    rin = settings.sigma.rin
    rout = settings.sigma.rout
    rmax = settings.sigma.rmax
    Qmin = settings.sigma.Qmin
    m = settings.physical.m
    Mstar = settings.physical.M
    #m_disk = settings.sigma.m_disk

    rin = match_units(pynbody.units.au, rin)[1]
    rout = match_units(pynbody.units.au, rout)[1]
    #m_disk = match_units(pynbody.units.Msol, m_disk)[1]

    if rmax is None:

        rmax = 2.5 * rout

    else:

        rmax = match_units(pynbody.units.au, rmax)[1]

    r = np.linspace(0, rmax, n_points)

    a = (rin / r).in_units('1')
    b = (r / rout).in_units('1')
    sigma = (np.exp(-a**2 - b**2) / r) * Mstar.units / r.units

    # Calculate Q
    Q = np.sqrt(Mstar * kB * T(r) / (G * m * r**3)) / (np.pi * sigma)
    Q.convert_units('1')

    sigma *= np.nanmin(Q) / Qmin

    # Remove all nans
    sigma[np.isnan(sigma)] = 0.0

    return r, sigma
예제 #18
0
def powerlaw(settings, T=None):
    """
    Generates a surface density profile according to a powerlaw sigma ~ r^p
    with a smooth interior cutoff and smooth exterior exponential cutoff.
    
    **ARGUMENTS**
    
    settings : IC settings
        settings like those contained in an IC object (see ICgen_settings.py)
    T : callable function
        Function that returns temperature of the disk as a function of radius
        IF none, a powerlaw temperature is assumed
    
    **RETURNS**
    
    R : SimArray
        Radii at which sigma is calculated
    sigma : SimArray
        Surface density profile as a function of R
    """
    # Parse settings
    Rd = settings.sigma.Rd
    rin = settings.sigma.rin
    rmax = settings.sigma.rmax
    cutlength = settings.sigma.cutlength
    Mstar = settings.physical.M
    Qmin = settings.sigma.Qmin
    n_points = settings.sigma.n_points
    m = settings.physical.m
    power = settings.sigma.power
    gamma = settings.physical.gamma_cs()

    if T is None:
        # If no callable object to calculate Temperature(R) is provided,
        # default to a powerlaw T ~ R^-q

        T0 = SimArray([129.0], 'K')  # Temperature at 1 AU
        R0 = SimArray([1.0], 'au')
        q = 0.59

        def T(x):

            return T0 * np.power((x / R0).in_units('1'), -q)

    Rd = match_units(pynbody.units.au, Rd)[1]
    Mstar = match_units(pynbody.units.Msol, Mstar)[1]
    # Molecular weight
    m = match_units(m, pynbody.units.m_p)[0]
    # Maximum R to calculate sigma at (needed for the exponential cutoff region)
    Rmax = rmax * Rd

    # Q calculation parameters:
    G = SimArray([1.0], 'G')
    kB = SimArray([1.0], 'k')

    # Initialize stuff
    A = SimArray(1.0, 'Msol') / (2 * np.pi * np.power(Rd, 2))
    # dflemin3 Nov. 4, 2015
    # Made units more explicit via SimArrays
    r_units = Rd.units
    R = SimArray(np.linspace(0, Rmax, n_points), r_units)
    r = R / Rd

    # Calculate sigma
    # Powerlaw
    #dflemin3 edit 06/10/2015: Try powerlaw of the form sigma ~ r^power
    sigma = A * np.power(r, power)
    sigma[0] = 0.0
    # Exterior cutoff
    sigma[r > 1] *= np.exp(-(r[r > 1] - 1)**2 / (2 * cutlength**2))
    # Interior cutoff
    sigma[r < rin] *= smoothstep(r[r < rin], degree=21, rescale=True)

    # Calculate Q
    Q = np.sqrt(Mstar * gamma * kB * T(R) / (G * m * R**3)) / (np.pi * sigma)
    Q.convert_units('1')

    # Rescale sigma to meet the minimum Q requirement
    sigma *= Q.min() / Qmin

    # Calculate Q
    Q = np.sqrt(Mstar * gamma * kB * T(R) / (G * m * R**3)) / (np.pi * sigma)
    Q.convert_units('1')
    return R, sigma
예제 #19
0
def Q(snapshot, molecular_mass = 2.0, bins=100, use_velocity=False, \
use_omega=True):
    """Calculates the Toomre Q as a function of r, assuming radial temperature
    profile and kappa ~= omega
    
    Parameters
    ----------
    
    snapshot : tipsy snapshot
    molecular_mass : float
        Mean molecular mass (for sound speed).  Default = 2.0
    bins : int or array
        Either the number of bins or the bin edges
    use_velocity : Bool
        Determines whether to use the particles' velocities to calculate orbital
        velocity.  Useful if the circular orbital velocities are set in the
        snapshot.
    use_omega : Bool
        Default=True.  Use omega as a proxy for kappa to reduce noise

    Returns
    -------
    
    Q : array
        Toomre Q as a function of r
    r_edges : array
        Radial bin edges
    """

    # Physical constants
    kB = SimArray([1.0],'k')
    G = SimArray([1.0],'G')
    # Calculate surface density
    sig, r_edges = sigma(snapshot, bins)
    # Calculate sound speed
    m = match_units(molecular_mass,'m_p')[0]
    c_s_all = np.sqrt(kB*snapshot.g['temp']/m)
    # Bin/average sound speed
    dummy, c_s, dummy2 = binned_mean(snapshot.g['rxy'], c_s_all, binedges=r_edges)

    if use_omega:
        # Calculate keplerian angular velocity (as a proxy for the epicyclic
        # frequency, which is a noisy calculation)
        if use_velocity:
            # Calculate directly from particle's velocity
            dummy, omega, dummy2 = binned_mean(snapshot.g['rxy'], \
            snapshot.g['vt']/snapshot.g['rxy'], binedges=r_edges)

        else:
            # Estimate, from forces, using pynbody
            p = pb.analysis.profile.Profile(snapshot, bins=r_edges)
            omega = p['omega']

        kappa_calc = omega

    else:

        if use_velocity:
            # Calculate directly from particle's velocities
            kappa_calc, dummy = kappa(snapshot, r_edges)

        else:
            # Estimate, from forces, using pynbody
            p = pb.analysis.profile.Profile(snapshot, bins=r_edges)
            kappa_calc = p['kappa']

    return (kappa_calc*c_s/(np.pi*G*sig)).in_units('1'), r_edges
예제 #20
0
파일: _math.py 프로젝트: philchang/diskpy
def binned_mean(x, y, bins=10, nbins=None, binedges = None, weights=None,\
weighted_bins=False, ret_bin_edges=False, binind=None):
    """
    Bins y according to x and takes the average for each bin.

    RETURNS a tuple of (bin_centers, y_mean, y_err) if ret_bin_edges=False
    else, Returns (bin_edges, y_mean, y_err)
    
    Parameters
    ----------
    x, y : array-like
        Bin and average y according to x
    bins : int or arraylike
        Either number of bins or an array of binedges
    nbins : int
        (optional) For backward compatibility, prefer using bins
    binedges : array-like
        (optional) For backward compatibility, prefer using bins
    weights : array-like
        Weights to use for data points.  Assume weights=1 (uniform)
    weighted_bins : bool
        If true, the bin centers will be calculated as weighted centers
    ret_bin_edges : bool
        Return bin centers instead of edges.  Default = False
    binind : list of arrays
        An optional list of arrays which specify which indices of x, y belong
        to which bin.  This can be used for speed.  when using a pynbody profile,
        pynbody will generate such a list:
            >>> p = pynbody.analysis.profile.Profile(f)
            >>> p.binind # binind
        For the returned bins to be consistent, the user should supply the used
        binedges.
    
    """
    x = np.asanyarray(x)
    y = np.asanyarray(y)
    if (isinstance(bins, int)) and (nbins is None):

        nbins = bins

    elif (hasattr(bins, '__iter__')) and (binedges is None):

        binedges = bins

    if binedges is not None:

        nbins = len(binedges) - 1

    else:

        binedges = np.linspace(x.min(), (1 + np.spacing(2)) * x.max(),
                               nbins + 1)

    if weights is None:

        weights = np.ones(x.shape)

    weights = strip_units(weights)

    # Pre-factor for weighted STD:
    A = 1 / (1 - (weights**2).sum())

    # Initialize
    y_mean = np.zeros(nbins)
    y_std = np.zeros(nbins)
    if binind is None:
        # Find the index bins for each data point
        ind = np.digitize(x, binedges) - 1
        N = np.histogram(x, binedges)[0]
    else:
        N = np.array([ind.sum() for ind in binind])
    # Ignore nans
    nan_ind = np.isnan(y)
    # Initialize bin_centers (try to retain units)
    bin_centers = 0.0 * binedges[1:]

    for i in range(nbins):

        if binind is None:
            #Indices to use
            mask = (ind == i) & (~nan_ind)
            # Set up the weighting
            w = weights[mask].copy()
            yvals = y[mask]
            xvals = x[mask]
        else:
            select_ind = binind[i]
            w = weights[select_ind]
            yvals = y[select_ind]
            xvals = x[select_ind]

        w /= w.sum()
        A = 1 / (1 - (w**2).sum())
        y_mean[i] = (w * yvals).sum()
        var = A * (w * (yvals - y_mean[i])**2).sum()
        y_std[i] = np.sqrt(var)

        if weighted_bins:
            # Center of mass of x positions
            bin_centers[i] = (w * xvals).sum()

    y_mean = match_units(y_mean, y)[0]
    y_err = y_std / np.sqrt(N)
    y_err = match_units(y_err, y)[0]

    y_mean[N == 0] = np.nan
    y_err[N == 0] = np.nan

    if not weighted_bins:

        bin_centers = (binedges[0:-1] + binedges[1:]) / 2.0
        binedges = match_units(binedges, x)[0]
        bin_centers = match_units(bin_centers, x)[0]

    else:

        bin_centers[N == 0] = np.nan

    if ret_bin_edges:

        return binedges, y_mean, y_err

    else:

        return bin_centers, y_mean, y_err
예제 #21
0
파일: _math.py 프로젝트: philchang/diskpy
    def meshspline(x1, y1):
        """
        Callable interpolation function, interoplates the value of z at
        points (x1, y1)
        
        Parameters
        ----------
        
        x1, y1 : array
            x and y points to evaluate z at.  Must be the same shape.  ie,
            x1[i], y1[i] define a point (x, y).
            If @x1 or @y1 have no units, they are assumed to have the units of
            the nodes used to make the interpolator.  Otherwise they are
            converted to the proper units
            
        Returns
        -------
        
        z(x1, y1) : array
            z evaluated at @x1, @y1
        """
        # Handle units
        x1 = strip_units(match_units(x1, units[0])[0])
        y1 = strip_units(match_units(y1, units[1])[0])

        # Setup x and y points to estimate z at
        x1 = np.asarray(x1).copy()
        y1 = np.asarray(y1)

        if len(x1.shape) < 1:

            x1 = x1[None]

        if len(y1.shape) < 1:

            y1 = y1[None]

        # Flatten arrays
        shape = x1.shape
        nElements = np.prod(shape)
        x1 = np.reshape(x1, [nElements])
        y1 = np.reshape(y1, [nElements])

        # Deal with xs outside of boundaries
        x1[x1 < xmin] = xmin
        x1[x1 > xmax] = xmax
        # Find bin indices
        ind = np.digitize(x1, xedges) - 1
        ind[ind < 0] = 0
        ind[ind > nbins - 1] = nbins - 1

        # Get bin info for every point
        xlo = xedges[ind]
        xhi = xedges[ind + 1]
        dx = binsize[ind]

        # Get weights for bins (distance from bin edges)
        wlo = (xhi - x1) / dx
        whi = (x1 - xlo) / dx

        # Get function values at left and right xedges
        flo = np.zeros(x1.shape)
        fhi = np.zeros(x1.shape)

        for i in range(nbins):

            # Select everything in bin i
            mask = (ind == i)

            if np.any(mask):

                # Retrieve function values
                flo[mask] = splines[i](y1[mask])
                fhi[mask] = splines[i + 1](y1[mask])

        # Take a weighted average of the function values at left and right
        # bin edges
        fout = wlo * flo + whi * fhi

        # Unflatten fout:
        fout = np.reshape(fout, shape)

        return SimArray(fout, units[2])
예제 #22
0
def snapshot_gen(ICobj):
    """
    Generates a tipsy snapshot from the initial conditions object ICobj.
    
    Returns snapshot, param
    
        snapshot: tipsy snapshot
        param: dictionary containing info for a .param file
    Note: Code has been edited (dflemin3) such that now it returns a snapshot for a circumbinary disk
    where initial conditions generated assuming star at origin of mass M.  After gas initialized, replaced
    star at origin with binary system who's center of mass lies at the origin and who's mass m1 +m2 = M
    """
    
    print 'Generating snapshot...'
    # Constants
    G = SimArray(1.0,'G')
    # ------------------------------------
    # Load in things from ICobj
    # ------------------------------------
    print 'Accessing data from ICs'
    settings = ICobj.settings
    
    # snapshot file name
    snapshotName = settings.filenames.snapshotName
    paramName = settings.filenames.paramName   
 
    # particle positions
    r = ICobj.pos.r
    xyz = ICobj.pos.xyz
    
    # Number of particles
    nParticles = ICobj.pos.nParticles
    
    # molecular mass
    m = settings.physical.m
    
    # star mass
    m_star = settings.physical.M.copy()
    
    # disk mass
    m_disk = ICobj.sigma.m_disk.copy()
    m_disk = match_units(m_disk, m_star)[0]
    
    # mass of the gas particles
    m_particles = m_disk / float(nParticles)
    
    # re-scale the particles (allows making of low-mass disk)
    m_particles *= settings.snapshot.mScale
    
    # -------------------------------------------------
    # Assign output
    # -------------------------------------------------
    print 'Assigning data to snapshot'
    # Get units all set up
    m_unit = m_star.units
    pos_unit = r.units
    
    if xyz.units != r.units:
        
        xyz.convert_units(pos_unit)
        
    # time units are sqrt(L^3/GM)
    t_unit = np.sqrt((pos_unit**3)*np.power((G*m_unit), -1)).units
    # velocity units are L/t
    v_unit = (pos_unit/t_unit).ratio('km s**-1')
    # Make it a unit, save value for future conversion
    v_unit_vel = v_unit
    #Ensure v_unit_vel is the same as what I assume it is.
    assert(np.fabs(AddBinary.VEL_UNIT-v_unit_vel)<AddBinary.SMALL),"VEL_UNIT not equal to ChaNGa unit! Why??"			
	
    v_unit = pynbody.units.Unit('{0} km s**-1'.format(v_unit))
    
    # Other settings
    metals = settings.snapshot.metals
    star_metals = metals
    
    # Generate snapshot
    # Note that empty pos, vel, and mass arrays are created in the snapshot
    snapshot = pynbody.new(star=1,gas=nParticles)
    snapshot['vel'].units = v_unit
    snapshot['eps'] = 0.01*SimArray(np.ones(nParticles+1, dtype=np.float32), pos_unit)
    snapshot['metals'] = SimArray(np.zeros(nParticles+1, dtype=np.float32))
    snapshot['rho'] = SimArray(np.zeros(nParticles+1, dtype=np.float32))
    
    snapshot.gas['pos'] = xyz
    snapshot.gas['temp'] = ICobj.T(r)
    snapshot.gas['mass'] = m_particles
    snapshot.gas['metals'] = metals
    
    snapshot.star['pos'] = SimArray([[ 0.,  0.,  0.]],pos_unit)
    snapshot.star['vel'] = SimArray([[ 0.,  0.,  0.]], v_unit)
    snapshot.star['mass'] = m_star
    snapshot.star['metals'] = SimArray(star_metals)
    # Estimate the star's softening length as the closest particle distance
    #snapshot.star['eps'] = r.min()
    
    # Make param file
    param = make_param(snapshot, snapshotName)
    param['dMeanMolWeight'] = m
       
    gc.collect()
    
    # CALCULATE VELOCITY USING calc_velocity.py.  This also estimates the 
    # gravitational softening length eps
    print 'Calculating circular velocity'
    preset = settings.changa_run.preset
    max_particles = global_settings['misc']['max_particles']
    calc_velocity.v_xy(snapshot, param, changa_preset=preset, max_particles=max_particles)
    
    gc.collect()
  
	# -------------------------------------------------
    # Estimate time step for changa to use
    # -------------------------------------------------
    # Save param file
    configsave(param, paramName, 'param')
    # Save snapshot
    snapshot.write(filename=snapshotName, fmt=pynbody.tipsy.TipsySnap)
    # est dDelta
    dDelta = ICgen_utils.est_time_step(paramName, preset)
    param['dDelta'] = dDelta
 
	# -------------------------------------------------
    # Create director file
    # -------------------------------------------------
    # largest radius to plot
    r_director = float(0.9 * r.max())
    # Maximum surface density
    sigma_min = float(ICobj.sigma(r_director))
    # surface density at largest radius
    sigma_max = float(ICobj.sigma.input_dict['sigma'].max())
    # Create director dict
    director = make_director(sigma_min, sigma_max, r_director, filename=param['achOutName'])
    ## Save .director file
    #configsave(director, directorName, 'director')

    #Now that velocities and everything are all initialized for gas particles, create new snapshot to return in which
    #single star particle is replaced by 2, same units as above
    snapshotBinary = pynbody.new(star=2,gas=nParticles)
    snapshotBinary['eps'] = 0.01*SimArray(np.ones(nParticles+2, dtype=np.float32), pos_unit)
    snapshotBinary['metals'] = SimArray(np.zeros(nParticles+2, dtype=np.float32))
    snapshotBinary['vel'].units = v_unit
    snapshotBinary['pos'].units = pos_unit
    snapshotBinary['mass'].units = snapshot['mass'].units
    snapshotBinary['rho'] = SimArray(np.zeros(nParticles+2, dtype=np.float32))

    #Assign gas particles with calculated/given values from above
    snapshotBinary.gas['pos'] = snapshot.gas['pos']
    snapshotBinary.gas['vel'] = snapshot.gas['vel']
    snapshotBinary.gas['temp'] = snapshot.gas['temp']
    snapshotBinary.gas['rho'] = snapshot.gas['rho']
    snapshotBinary.gas['eps'] = snapshot.gas['eps']
    snapshotBinary.gas['mass'] = snapshot.gas['mass']
    snapshotBinary.gas['metals'] = snapshot.gas['metals']

    #Load Binary system obj to initialize system
    binsys = ICobj.settings.physical.binsys
    
    x1,x2,v1,v2 = binsys.generateICs()

    #Put velocity in sim units
    #!!! Note: v_unit_vel will always be 29.785598165 km/s when m_unit = Msol and r_unit = 1 AU in kpc!!!
    #conv = v_unit_vel #km/s in sim units
    #v1 /= conv
    #v2 /= conv

    #Assign position, velocity assuming CCW orbit

    snapshotBinary.star[0]['pos'] = SimArray(x1,pos_unit)
    snapshotBinary.star[0]['vel'] = SimArray(v1,v_unit)
    snapshotBinary.star[1]['pos'] = SimArray(x2,pos_unit)
    snapshotBinary.star[1]['vel'] = SimArray(v2,v_unit)

    #Set stellar masses
    #Set Mass units
    #Create simArray for mass, convert units to simulation mass units
    priMass = SimArray(binsys.m1,m_unit)
    secMass = SimArray(binsys.m2,m_unit)

    snapshotBinary.star[0]['mass'] = priMass
    snapshotBinary.star[1]['mass'] = secMass
    snapshotBinary.star['metals'] = SimArray(star_metals)

    #Estimate stars' softening length as fraction of distance to COM
    d = np.sqrt(AddBinary.dotProduct(x1-x2,x1-x2))

    snapshotBinary.star[0]['eps'] = SimArray(math.fabs(d)/4.0,pos_unit)
    snapshotBinary.star[1]['eps'] = SimArray(math.fabs(d)/4.0,pos_unit)
 
    print 'Wrapping up'
    # Now set the star particle's tform to a negative number.  This allows
    # UW ChaNGa treat it as a sink particle.
    snapshotBinary.star['tform'] = -1.0
    
    #Set Sink Radius to be mass-weighted average of Roche lobes of two stars
    r1 = AddBinary.calcRocheLobe(binsys.m1/binsys.m2,binsys.a) 
    r2 = AddBinary.calcRocheLobe(binsys.m2/binsys.m1,binsys.a)
    p = strip_units(binsys.m1/(binsys.m1 + binsys.m2))

    r_sink = (r1*p) + (r2*(1.0-p))
    param['dSinkBoundOrbitRadius'] = r_sink
    param['dSinkRadius'] = r_sink
    param['dSinkMassMin'] = 0.9 * strip_units(secMass)
    param['bDoSinks'] = 1
    
    return snapshotBinary, param, director
    
        
예제 #23
0
파일: _math.py 프로젝트: dflemin3/diskpy
    def meshspline(x1, y1):
        """
        Callable interpolation function, interoplates the value of z at
        points (x1, y1)
        
        Parameters
        ----------
        
        x1, y1 : array
            x and y points to evaluate z at.  Must be the same shape.  ie,
            x1[i], y1[i] define a point (x, y).
            If @x1 or @y1 have no units, they are assumed to have the units of
            the nodes used to make the interpolator.  Otherwise they are
            converted to the proper units
            
        Returns
        -------
        
        z(x1, y1) : array
            z evaluated at @x1, @y1
        """
        # Handle units
        x1 = strip_units(match_units(x1, units[0])[0])
        y1 = strip_units(match_units(y1, units[1])[0])

        # Setup x and y points to estimate z at
        x1 = np.asarray(x1).copy()
        y1 = np.asarray(y1)

        if len(x1.shape) < 1:

            x1 = x1[None]

        if len(y1.shape) < 1:

            y1 = y1[None]

        # Flatten arrays
        shape = x1.shape
        nElements = np.prod(shape)
        x1 = np.reshape(x1, [nElements])
        y1 = np.reshape(y1, [nElements])

        # Deal with xs outside of boundaries
        x1[x1 < xmin] = xmin
        x1[x1 > xmax] = xmax
        # Find bin indices
        ind = np.digitize(x1, xedges) - 1
        ind[ind < 0] = 0
        ind[ind > nbins - 1] = nbins - 1

        # Get bin info for every point
        xlo = xedges[ind]
        xhi = xedges[ind + 1]
        dx = binsize[ind]

        # Get weights for bins (distance from bin edges)
        wlo = (xhi - x1) / dx
        whi = (x1 - xlo) / dx

        # Get function values at left and right xedges
        flo = np.zeros(x1.shape)
        fhi = np.zeros(x1.shape)

        for i in range(nbins):

            # Select everything in bin i
            mask = ind == i

            if np.any(mask):

                # Retrieve function values
                flo[mask] = splines[i](y1[mask])
                fhi[mask] = splines[i + 1](y1[mask])

        # Take a weighted average of the function values at left and right
        # bin edges
        fout = wlo * flo + whi * fhi

        # Unflatten fout:
        fout = np.reshape(fout, shape)

        return SimArray(fout, units[2])
예제 #24
0
def MQWS(settings, T):
    """
    Generates a surface density profile as the per method used in Mayer, Quinn,
    Wadsley, and Stadel 2004
    
    ** ARGUMENTS **
    NOTE: if units are not supplied, assumed units are AU, Msol
    
    settings : IC settings
        settings like those contained in an IC object (see ICgen_settings.py)
        
    T : callable
        A function to calculate temperature as a function of radius
        
    ** RETURNS **
    
    r : SimArray
        Radii at which sigma is calculated
    sigma : SimArray
        Surface density profile as a function of R
    """
    # Q calculation parameters:
    G = SimArray([1.0],'G')
    kB = SimArray([1.0],'k')
    
    # Load in settings
    n_points = settings.sigma.n_points
    rin = settings.sigma.rin
    rout = settings.sigma.rout
    rmax = settings.sigma.rmax
    Qmin = settings.sigma.Qmin
    m = settings.physical.m
    Mstar = settings.physical.M
    #m_disk = settings.sigma.m_disk
    
    rin = match_units(pynbody.units.au, rin)[1]
    rout = match_units(pynbody.units.au, rout)[1]
    #m_disk = match_units(pynbody.units.Msol, m_disk)[1]
    
    if rmax is None:
        
        rmax = 2.5 * rout
        
    else:
        
        rmax = match_units(pynbody.units.au, rmax)[1]
        
    r = np.linspace(0, rmax, n_points)
    
    a = (rin/r).in_units('1')
    b = (r/rout).in_units('1')
    sigma = (np.exp(-a**2 - b**2)/r) * Mstar.units/r.units
    
    # Calculate Q
    Q = np.sqrt(Mstar*kB*T(r)/(G*m*r**3))/(np.pi*sigma)
    Q.convert_units('1')

    sigma *= np.nanmin(Q)/Qmin
    
    # Remove all nans
    sigma[np.isnan(sigma)] = 0.0
    
    
    return r, sigma
예제 #25
0
def snapshot_gen(ICobj):
    """
    Generates a tipsy snapshot from the initial conditions object ICobj.
    
    Returns snapshot, param
    
        snapshot: tipsy snapshot
        param: dictionary containing info for a .param file
    Note: Code has been edited (dflemin3) such that now it returns a snapshot for a circumbinary disk
    where initial conditions generated assuming star at origin of mass M.  After gas initialized, replaced
    star at origin with binary system who's center of mass lies at the origin and who's mass m1 +m2 = M
    """

    print "Generating snapshot..."
    # Constants
    G = SimArray(1.0, "G")
    # ------------------------------------
    # Load in things from ICobj
    # ------------------------------------
    print "Accessing data from ICs"
    settings = ICobj.settings

    # snapshot file name
    snapshotName = settings.filenames.snapshotName
    paramName = settings.filenames.paramName

    # Load user supplied snapshot (assumed to be in cwd)
    path = "/astro/store/scratch/tmp/dflemin3/nbodyshare/9au-Q1.05-129K/"
    snapshot = pynbody.load(path + snapshotName)

    # particle positions
    r = snapshot.gas["r"]
    xyz = snapshot.gas["pos"]

    # Number of particles
    nParticles = len(snapshot.gas)

    # molecular mass
    m = settings.physical.m

    # Pull star mass from user-supplied snapshot
    ICobj.settings.physical.M = snapshot.star["mass"]  # Total stellar mass in solar masses
    m_star = ICobj.settings.physical.M

    # disk mass
    m_disk = np.sum(snapshot.gas["mass"])
    m_disk = match_units(m_disk, m_star)[0]

    # mass of the gas particles
    m_particles = m_disk / float(nParticles)

    # re-scale the particles (allows making of low-mass disk)
    m_particles *= settings.snapshot.mScale

    # -------------------------------------------------
    # Assign output
    # -------------------------------------------------
    print "Assigning data to snapshot"
    # Get units all set up
    m_unit = m_star.units
    pos_unit = r.units

    if xyz.units != r.units:

        xyz.convert_units(pos_unit)

    # time units are sqrt(L^3/GM)
    t_unit = np.sqrt((pos_unit ** 3) * np.power((G * m_unit), -1)).units
    # velocity units are L/t
    v_unit = (pos_unit / t_unit).ratio("km s**-1")
    # Make it a unit, save value for future conversion
    v_unit_vel = v_unit
    # Ensure v_unit_vel is the same as what I assume it is.
    assert np.fabs(AddBinary.VEL_UNIT - v_unit_vel) < AddBinary.SMALL, "VEL_UNIT not equal to ChaNGa unit! Why??"

    v_unit = pynbody.units.Unit("{0} km s**-1".format(v_unit))

    # Other settings
    metals = settings.snapshot.metals
    star_metals = metals

    # Estimate the star's softening length as the closest particle distance
    eps = r.min()

    # Make param file
    param = make_param(snapshot, snapshotName)
    param["dMeanMolWeight"] = m

    gc.collect()

    # CALCULATE VELOCITY USING calc_velocity.py.  This also estimates the
    # gravitational softening length eps

    preset = settings.changa_run.preset

    # -------------------------------------------------
    # Estimate time step for changa to use
    # -------------------------------------------------
    # Save param file
    configsave(param, paramName, "param")
    # Save snapshot
    snapshot.write(filename=snapshotName, fmt=pynbody.tipsy.TipsySnap)
    # est dDelta
    dDelta = ICgen_utils.est_time_step(paramName, preset)
    param["dDelta"] = dDelta

    # -------------------------------------------------
    # Create director file
    # -------------------------------------------------
    # largest radius to plot
    r_director = float(0.9 * r.max())
    # Maximum surface density
    sigma_min = float(ICobj.sigma(r_director))
    # surface density at largest radius
    sigma_max = float(ICobj.sigma.input_dict["sigma"].max())
    # Create director dict
    director = make_director(sigma_min, sigma_max, r_director, filename=param["achOutName"])
    ## Save .director file
    # configsave(director, directorName, 'director')

    """
    Now that the gas disk is initializes around the primary (M=m1), add in the
    second star as specified by the user.
    """

    # Now that velocities and everything are all initialized for gas particles, create new snapshot to return in which
    # single star particle is replaced by 2, same units as above
    snapshotBinary = pynbody.new(star=2, gas=nParticles)
    snapshotBinary["eps"] = 0.01 * SimArray(np.ones(nParticles + 2, dtype=np.float32), pos_unit)
    snapshotBinary["metals"] = SimArray(np.zeros(nParticles + 2, dtype=np.float32))
    snapshotBinary["vel"].units = v_unit
    snapshotBinary["pos"].units = pos_unit
    snapshotBinary["mass"].units = snapshot["mass"].units
    snapshotBinary["rho"] = SimArray(np.zeros(nParticles + 2, dtype=np.float32))

    # Assign gas particles with calculated/given values from above
    snapshotBinary.gas["pos"] = snapshot.gas["pos"]
    snapshotBinary.gas["vel"] = snapshot.gas["vel"]
    snapshotBinary.gas["temp"] = snapshot.gas["temp"]
    snapshotBinary.gas["rho"] = snapshot.gas["rho"]
    snapshotBinary.gas["eps"] = snapshot.gas["eps"]
    snapshotBinary.gas["mass"] = snapshot.gas["mass"]
    snapshotBinary.gas["metals"] = snapshot.gas["metals"]

    # Load Binary system obj to initialize system
    binsys = ICobj.settings.physical.binsys
    m_disk = strip_units(np.sum(snapshotBinary.gas["mass"]))
    binsys.m1 = strip_units(m_star)
    binsys.m1 = binsys.m1 + m_disk
    # Recompute cartesian coords considering primary as m1+m_disk
    binsys.computeCartesian()

    x1, x2, v1, v2 = binsys.generateICs()

    # Assign position, velocity assuming CCW orbit
    snapshotBinary.star[0]["pos"] = SimArray(x1, pos_unit)
    snapshotBinary.star[0]["vel"] = SimArray(v1, v_unit)
    snapshotBinary.star[1]["pos"] = SimArray(x2, pos_unit)
    snapshotBinary.star[1]["vel"] = SimArray(v2, v_unit)

    """
    We have the binary positions about their center of mass, (0,0,0), so 
    shift the position, velocity of the gas disk to be around the primary.
    """
    snapshotBinary.gas["pos"] += snapshotBinary.star[0]["pos"]
    snapshotBinary.gas["vel"] += snapshotBinary.star[0]["vel"]

    # Set stellar masses: Create simArray for mass, convert units to simulation mass units
    snapshotBinary.star[0]["mass"] = SimArray(binsys.m1 - m_disk, m_unit)
    snapshotBinary.star[1]["mass"] = SimArray(binsys.m2, m_unit)
    snapshotBinary.star["metals"] = SimArray(star_metals)

    print "Wrapping up"
    # Now set the star particle's tform to a negative number.  This allows
    # UW ChaNGa treat it as a sink particle.
    snapshotBinary.star["tform"] = -1.0

    # Set sink radius, stellar smoothing length as fraction of distance
    # from primary to inner edge of the disk
    r_sink = eps
    snapshotBinary.star[0]["eps"] = SimArray(r_sink / 2.0, pos_unit)
    snapshotBinary.star[1]["eps"] = SimArray(r_sink / 2.0, pos_unit)
    param["dSinkBoundOrbitRadius"] = r_sink
    param["dSinkRadius"] = r_sink
    param["dSinkMassMin"] = 0.9 * binsys.m2
    param["bDoSinks"] = 1

    return snapshotBinary, param, director
예제 #26
0
def powerlaw(settings, T = None):
    """
    Generates a surface density profile according to a powerlaw sigma ~ r^p
    with a smooth interior cutoff and smooth exterior exponential cutoff.
    
    **ARGUMENTS**
    
    settings : IC settings
        settings like those contained in an IC object (see ICgen_settings.py)
    T : callable function
        Function that returns temperature of the disk as a function of radius
        IF none, a powerlaw temperature is assumed
    
    **RETURNS**
    
    R : SimArray
        Radii at which sigma is calculated
    sigma : SimArray
        Surface density profile as a function of R
    """
    # Parse settings
    Rd = settings.sigma.Rd
    rin = settings.sigma.rin
    rmax = settings.sigma.rmax
    cutlength = settings.sigma.cutlength
    Mstar = settings.physical.M
    Qmin = settings.sigma.Qmin
    n_points = settings.sigma.n_points
    m = settings.physical.m
    power = settings.sigma.power    

    if T is None:
        # If no callable object to calculate Temperature(R) is provided, 
        # default to a powerlaw T ~ R^-q
        
        T0 = SimArray([129.0],'K') # Temperature at 1 AU
        R0 = SimArray([1.0],'au')
        q = 0.59
        def T(x):
            
            return T0 * np.power((x/R0).in_units('1'),-q)
        
    Rd = match_units(pynbody.units.au, Rd)[1]
    Mstar = match_units(pynbody.units.Msol, Mstar)[1]
    # Molecular weight
    m = match_units(m, pynbody.units.m_p)[0]
    # Maximum R to calculate sigma at (needed for the exponential cutoff region)
    Rmax = rmax*Rd
    
    # Q calculation parameters:
    G = SimArray([1.0],'G')
    kB = SimArray([1.0],'k')
    
    # Initialize stuff
    A = SimArray(1.0,'Msol')/(2*np.pi*np.power(Rd,2))
    #R = np.linspace(0,Rmax,n_points)
    #r = np.array((R/Rd).in_units('1'))
    # dflemin3 Nov. 4, 2015
    # Made units more explicit via SimArrays
    r_units = Rd.units
    R = SimArray(np.linspace(0,Rmax,n_points),r_units)
    r = R/Rd
    
    # Calculate sigma
    # Powerlaw
    #sigma = A/r
    #dflemin3 edit 06/10/2015: Try powerlaw of the form sigma ~ r^power
    sigma = A*np.power(r,power)
    sigma[0] = 0.0
    # Exterior cutoff
    sigma[r>1] *= np.exp(-(r[r>1] - 1)**2 / (2*cutlength**2))
    # Interior cutoff
    sigma[r<rin] *= smoothstep(r[r<rin],degree=21,rescale=True)
    
    # Calculate Q
    Q = np.sqrt(Mstar*kB*T(R)/(G*m*R**3))/(np.pi*sigma)
    Q.convert_units('1')
    
    # Rescale sigma to meet the minimum Q requirement
    sigma *= Q.min()/Qmin
    
    # Calculate Q
    Q = np.sqrt(Mstar*kB*T(R)/(G*m*R**3))/(np.pi*sigma)
    Q.convert_units('1')
    return R, sigma
예제 #27
0
def snapshot_gen(ICobj):
    """
    Generates a tipsy snapshot from the initial conditions object ICobj.
    
    Returns snapshot, param
    
        snapshot: tipsy snapshot
        param: dictionary containing info for a .param file
    Note: Code has been edited (dflemin3) such that now it returns a snapshot for a circumbinary disk
    where initial conditions generated assuming star at origin of mass M.  After gas initialized, replaced
    star at origin with binary system who's center of mass lies at the origin and who's mass m1 +m2 = M
    """

    print 'Generating snapshot...'
    # Constants
    G = SimArray(1.0, 'G')
    # ------------------------------------
    # Load in things from ICobj
    # ------------------------------------
    print 'Accessing data from ICs'
    settings = ICobj.settings

    # snapshot file name
    snapshotName = settings.filenames.snapshotName
    paramName = settings.filenames.paramName

    #Load user supplied snapshot (assumed to be in cwd)
    path = "/astro/store/scratch/tmp/dflemin3/nbodyshare/9au-Q1.05-129K/"
    snapshot = pynbody.load(path + snapshotName)

    # particle positions
    r = snapshot.gas['r']
    xyz = snapshot.gas['pos']

    # Number of particles
    nParticles = len(snapshot.gas)

    # molecular mass
    m = settings.physical.m

    #Pull star mass from user-supplied snapshot
    ICobj.settings.physical.M = snapshot.star[
        'mass']  #Total stellar mass in solar masses
    m_star = ICobj.settings.physical.M

    # disk mass
    m_disk = np.sum(snapshot.gas['mass'])
    m_disk = match_units(m_disk, m_star)[0]

    # mass of the gas particles
    m_particles = m_disk / float(nParticles)

    # re-scale the particles (allows making of low-mass disk)
    m_particles *= settings.snapshot.mScale

    # -------------------------------------------------
    # Assign output
    # -------------------------------------------------
    print 'Assigning data to snapshot'
    # Get units all set up
    m_unit = m_star.units
    pos_unit = r.units

    if xyz.units != r.units:

        xyz.convert_units(pos_unit)

    # time units are sqrt(L^3/GM)
    t_unit = np.sqrt((pos_unit**3) * np.power((G * m_unit), -1)).units
    # velocity units are L/t
    v_unit = (pos_unit / t_unit).ratio('km s**-1')
    # Make it a unit, save value for future conversion
    v_unit_vel = v_unit
    #Ensure v_unit_vel is the same as what I assume it is.
    assert (np.fabs(AddBinary.VEL_UNIT - v_unit_vel) <
            AddBinary.SMALL), "VEL_UNIT not equal to ChaNGa unit! Why??"

    v_unit = pynbody.units.Unit('{0} km s**-1'.format(v_unit))

    # Other settings
    metals = settings.snapshot.metals
    star_metals = metals

    # Estimate the star's softening length as the closest particle distance
    eps = r.min()

    # Make param file
    param = make_param(snapshot, snapshotName)
    param['dMeanMolWeight'] = m

    gc.collect()

    # CALCULATE VELOCITY USING calc_velocity.py.  This also estimates the
    # gravitational softening length eps

    preset = settings.changa_run.preset

    # -------------------------------------------------
    # Estimate time step for changa to use
    # -------------------------------------------------
    # Save param file
    configsave(param, paramName, 'param')
    # Save snapshot
    snapshot.write(filename=snapshotName, fmt=pynbody.tipsy.TipsySnap)
    # est dDelta
    dDelta = ICgen_utils.est_time_step(paramName, preset)
    param['dDelta'] = dDelta

    # -------------------------------------------------
    # Create director file
    # -------------------------------------------------
    # largest radius to plot
    r_director = float(0.9 * r.max())
    # Maximum surface density
    sigma_min = float(ICobj.sigma(r_director))
    # surface density at largest radius
    sigma_max = float(ICobj.sigma.input_dict['sigma'].max())
    # Create director dict
    director = make_director(sigma_min,
                             sigma_max,
                             r_director,
                             filename=param['achOutName'])
    ## Save .director file
    #configsave(director, directorName, 'director')
    """
    Now that the gas disk is initializes around the primary (M=m1), add in the
    second star as specified by the user.
    """

    #Now that velocities and everything are all initialized for gas particles, create new snapshot to return in which
    #single star particle is replaced by 2, same units as above
    snapshotBinary = pynbody.new(star=2, gas=nParticles)
    snapshotBinary['eps'] = 0.01 * SimArray(
        np.ones(nParticles + 2, dtype=np.float32), pos_unit)
    snapshotBinary['metals'] = SimArray(
        np.zeros(nParticles + 2, dtype=np.float32))
    snapshotBinary['vel'].units = v_unit
    snapshotBinary['pos'].units = pos_unit
    snapshotBinary['mass'].units = snapshot['mass'].units
    snapshotBinary['rho'] = SimArray(np.zeros(nParticles + 2,
                                              dtype=np.float32))

    #Assign gas particles with calculated/given values from above
    snapshotBinary.gas['pos'] = snapshot.gas['pos']
    snapshotBinary.gas['vel'] = snapshot.gas['vel']
    snapshotBinary.gas['temp'] = snapshot.gas['temp']
    snapshotBinary.gas['rho'] = snapshot.gas['rho']
    snapshotBinary.gas['eps'] = snapshot.gas['eps']
    snapshotBinary.gas['mass'] = snapshot.gas['mass']
    snapshotBinary.gas['metals'] = snapshot.gas['metals']

    #Load Binary system obj to initialize system
    binsys = ICobj.settings.physical.binsys
    m_disk = strip_units(np.sum(snapshotBinary.gas['mass']))
    binsys.m1 = strip_units(m_star)
    binsys.m1 = binsys.m1 + m_disk
    #Recompute cartesian coords considering primary as m1+m_disk
    binsys.computeCartesian()

    x1, x2, v1, v2 = binsys.generateICs()

    #Assign position, velocity assuming CCW orbit
    snapshotBinary.star[0]['pos'] = SimArray(x1, pos_unit)
    snapshotBinary.star[0]['vel'] = SimArray(v1, v_unit)
    snapshotBinary.star[1]['pos'] = SimArray(x2, pos_unit)
    snapshotBinary.star[1]['vel'] = SimArray(v2, v_unit)
    """
    We have the binary positions about their center of mass, (0,0,0), so 
    shift the position, velocity of the gas disk to be around the primary.
    """
    snapshotBinary.gas['pos'] += snapshotBinary.star[0]['pos']
    snapshotBinary.gas['vel'] += snapshotBinary.star[0]['vel']

    #Set stellar masses: Create simArray for mass, convert units to simulation mass units
    snapshotBinary.star[0]['mass'] = SimArray(binsys.m1 - m_disk, m_unit)
    snapshotBinary.star[1]['mass'] = SimArray(binsys.m2, m_unit)
    snapshotBinary.star['metals'] = SimArray(star_metals)

    print 'Wrapping up'
    # Now set the star particle's tform to a negative number.  This allows
    # UW ChaNGa treat it as a sink particle.
    snapshotBinary.star['tform'] = -1.0

    #Set sink radius, stellar smoothing length as fraction of distance
    #from primary to inner edge of the disk
    r_sink = eps
    snapshotBinary.star[0]['eps'] = SimArray(r_sink / 2.0, pos_unit)
    snapshotBinary.star[1]['eps'] = SimArray(r_sink / 2.0, pos_unit)
    param['dSinkBoundOrbitRadius'] = r_sink
    param['dSinkRadius'] = r_sink
    param['dSinkMassMin'] = 0.9 * binsys.m2
    param['bDoSinks'] = 1

    return snapshotBinary, param, director
예제 #28
0
def snapshot_gen(ICobj):
    """
    Generates a tipsy snapshot from the initial conditions object ICobj.
    
    Returns snapshot, param
    
        snapshot: tipsy snapshot
        param: dictionary containing info for a .param file
    """
    
    print 'Generating snapshot...'
    # Constants
    G = SimArray(1.0,'G')
    # ------------------------------------
    # Load in things from ICobj
    # ------------------------------------
    print 'Accessing data from ICs'
    settings = ICobj.settings
    # filenames
    snapshotName = settings.filenames.snapshotName
    paramName = settings.filenames.paramName
        
    # particle positions
    r = ICobj.pos.r
    xyz = ICobj.pos.xyz
    # Number of particles
    nParticles = ICobj.pos.nParticles
    # molecular mass
    m = settings.physical.m
    # star mass
    m_star = settings.physical.M.copy()
    # disk mass
    m_disk = ICobj.sigma.m_disk.copy()
    m_disk = match_units(m_disk, m_star)[0]
    # mass of the gas particles
    m_particles = m_disk / float(nParticles)
    # re-scale the particles (allows making of lo-mass disk)
    m_particles *= settings.snapshot.mScale
    
    # -------------------------------------------------
    # Assign output
    # -------------------------------------------------
    print 'Assigning data to snapshot'
    # Get units all set up
    m_unit = m_star.units
    pos_unit = r.units
    
    if xyz.units != r.units:
        
        xyz.convert_units(pos_unit)
        
    # time units are sqrt(L^3/GM)
    t_unit = np.sqrt((pos_unit**3)*np.power((G*m_unit), -1)).units
    # velocity units are L/t
    v_unit = (pos_unit/t_unit).ratio('km s**-1')
    # Make it a unit
    v_unit = pynbody.units.Unit('{0} km s**-1'.format(v_unit))
    
    # Other settings
    metals = settings.snapshot.metals
    star_metals = metals
    
    # -------------------------------------------------
    # Initialize snapshot
    # -------------------------------------------------
    # Note that empty pos, vel, and mass arrays are created in the snapshot
    snapshot = pynbody.new(star=1,gas=nParticles)
    snapshot['vel'].units = v_unit
    snapshot['eps'] = 0.01*SimArray(np.ones(nParticles+1, dtype=np.float32), pos_unit)
    snapshot['metals'] = SimArray(np.zeros(nParticles+1, dtype=np.float32))
    snapshot['rho'] = SimArray(np.zeros(nParticles+1, dtype=np.float32))
    
    snapshot.gas['pos'] = xyz
    snapshot.gas['temp'] = ICobj.T(r)
    snapshot.gas['mass'] = m_particles
    snapshot.gas['metals'] = metals
    
    snapshot.star['pos'] = SimArray([[ 0.,  0.,  0.]],pos_unit)
    snapshot.star['vel'] = SimArray([[ 0.,  0.,  0.]], v_unit)
    snapshot.star['mass'] = m_star
    snapshot.star['metals'] = SimArray(star_metals)
    # Estimate the star's softening length as the closest particle distance
    snapshot.star['eps'] = r.min()
    
    # Make param file
    param = make_param(snapshot, snapshotName)
    param['dMeanMolWeight'] = m
    eos = (settings.physical.eos).lower()
    
    if eos == 'adiabatic':
        
        param['bGasAdiabatic'] = 1
        param['bGasIsothermal'] = 0
        
    param['dConstGamma']
       
    gc.collect()
    
    # -------------------------------------------------
    # CALCULATE VELOCITY USING calc_velocity.py.  This also estimates the 
    # gravitational softening length eps
    # -------------------------------------------------
    print 'Calculating circular velocity'
    preset = settings.changa_run.preset
    max_particles = global_settings['misc']['max_particles']
    calc_velocity.v_xy(snapshot, param, changa_preset=preset, max_particles=max_particles)
    
    gc.collect()
    
    # -------------------------------------------------
    # Estimate time step for changa to use
    # -------------------------------------------------
    # Save param file
    configsave(param, paramName, 'param')
    # Save snapshot
    snapshot.write(filename=snapshotName, fmt=pynbody.tipsy.TipsySnap)
    # est dDelta
    dDelta = ICgen_utils.est_time_step(paramName, preset)
    param['dDelta'] = dDelta
    
    # -------------------------------------------------
    # Create director file
    # -------------------------------------------------
    # largest radius to plot
    r_director = float(0.9 * r.max())
    # Maximum surface density
    sigma_min = float(ICobj.sigma(r_director))
    # surface density at largest radius
    sigma_max = float(ICobj.sigma.input_dict['sigma'].max())
    # Create director dict
    director = make_director(sigma_min, sigma_max, r_director, filename=param['achOutName'])
    ## Save .director file
    #configsave(director, directorName, 'director')
    
    # -------------------------------------------------
    # Wrap up
    # -------------------------------------------------
    print 'Wrapping up'
    # Now set the star particle's tform to a negative number.  This allows
    # UW ChaNGa treat it as a sink particle.
    snapshot.star['tform'] = -1.0
    
    # Update params
    r_sink = strip_units(r.min())
    param['dSinkBoundOrbitRadius'] = r_sink
    param['dSinkRadius'] = r_sink
    param['dSinkMassMin'] = 0.9 * strip_units(m_star)
    param['bDoSinks'] = 1
    
    return snapshot, param, director
예제 #29
0
def snapshot_gen(ICobj):
    """
    Generates a tipsy snapshot from the initial conditions object ICobj.
    
    Returns snapshot, param
    
        snapshot: tipsy snapshot
        param: dictionary containing info for a .param file
    Note: Code has been edited (dflemin3) such that now it returns a snapshot for a circumbinary disk
    where initial conditions generated assuming star at origin of mass M.  After gas initialized, replaced
    star at origin with binary system who's center of mass lies at the origin and who's mass m1 +m2 = M
    """
    
    print 'Generating snapshot...'
    # Constants
    G = SimArray(1.0,'G')
    # ------------------------------------
    # Load in things from ICobj
    # ------------------------------------
    print 'Accessing data from ICs'
    settings = ICobj.settings
    
    # snapshot file name
    snapshotName = settings.filenames.snapshotName
    paramName = settings.filenames.paramName   
 
    # particle positions
    r = ICobj.pos.r
    xyz = ICobj.pos.xyz
    
    # Number of particles
    nParticles = ICobj.pos.nParticles
    
    # molecular mass
    m = settings.physical.m
    
    # star mass
    m_star = settings.physical.M.copy()
    
    # disk mass
    m_disk = ICobj.sigma.m_disk.copy()
    m_disk = match_units(m_disk, m_star)[0]
    
    # mass of the gas particles
    m_particles = m_disk / float(nParticles)
    
    # re-scale the particles (allows making of low-mass disk)
    m_particles *= settings.snapshot.mScale
    
    # -------------------------------------------------
    # Assign output
    # -------------------------------------------------
    print 'Assigning data to snapshot'
    # Get units all set up
    m_unit = m_star.units
    pos_unit = r.units
    
    if xyz.units != r.units:
        
        xyz.convert_units(pos_unit)
        
    # time units are sqrt(L^3/GM)
    t_unit = np.sqrt((pos_unit**3)*np.power((G*m_unit), -1)).units
    # velocity units are L/t
    v_unit = (pos_unit/t_unit).ratio('km s**-1')
    # Make it a unit, save value for future conversion
    v_unit_vel = v_unit
    #Ensure v_unit_vel is the same as what I assume it is.
    assert(np.fabs(AddBinary.VEL_UNIT-v_unit_vel)<AddBinary.SMALL),"VEL_UNIT not equal to ChaNGa unit! Why??"			
	
    v_unit = pynbody.units.Unit('{0} km s**-1'.format(v_unit))
    
    # Other settings
    metals = settings.snapshot.metals
    star_metals = metals
    
    # Generate snapshot
    # Note that empty pos, vel, and mass arrays are created in the snapshot
    snapshot = pynbody.new(star=1,gas=nParticles)
    snapshot['vel'].units = v_unit
    snapshot['eps'] = 0.01*SimArray(np.ones(nParticles+1, dtype=np.float32), pos_unit)
    snapshot['metals'] = SimArray(np.zeros(nParticles+1, dtype=np.float32))
    snapshot['rho'] = SimArray(np.zeros(nParticles+1, dtype=np.float32))
    
    snapshot.gas['pos'] = xyz
    snapshot.gas['temp'] = ICobj.T(r)
    snapshot.gas['mass'] = m_particles
    snapshot.gas['metals'] = metals
    
    snapshot.star['pos'] = SimArray([[ 0.,  0.,  0.]],pos_unit)
    snapshot.star['vel'] = SimArray([[ 0.,  0.,  0.]], v_unit)
    snapshot.star['mass'] = m_star
    snapshot.star['metals'] = SimArray(star_metals)
    # Estimate the star's softening length as the closest particle distance
    #snapshot.star['eps'] = r.min()
    
    # Make param file
    param = make_param(snapshot, snapshotName)
    param['dMeanMolWeight'] = m
       
    gc.collect()
    
    # CALCULATE VELOCITY USING calc_velocity.py.  This also estimates the 
    # gravitational softening length eps
    print 'Calculating circular velocity'
    preset = settings.changa_run.preset
    max_particles = global_settings['misc']['max_particles']
    calc_velocity.v_xy(snapshot, param, changa_preset=preset, max_particles=max_particles)
    
    gc.collect()
  
	# -------------------------------------------------
    # Estimate time step for changa to use
    # -------------------------------------------------
    # Save param file
    configsave(param, paramName, 'param')
    # Save snapshot
    snapshot.write(filename=snapshotName, fmt=pynbody.tipsy.TipsySnap)
    # est dDelta
    dDelta = ICgen_utils.est_time_step(paramName, preset)
    param['dDelta'] = dDelta
 
	# -------------------------------------------------
    # Create director file
    # -------------------------------------------------
    # largest radius to plot
    r_director = float(0.9 * r.max())
    # Maximum surface density
    sigma_min = float(ICobj.sigma(r_director))
    # surface density at largest radius
    sigma_max = float(ICobj.sigma.input_dict['sigma'].max())
    # Create director dict
    director = make_director(sigma_min, sigma_max, r_director, filename=param['achOutName'])
    ## Save .director file
    #configsave(director, directorName, 'director')

    #Now that velocities and everything are all initialized for gas particles, create new snapshot to return in which
    #single star particle is replaced by 2, same units as above
    snapshotBinary = pynbody.new(star=2,gas=nParticles)
    snapshotBinary['eps'] = 0.01*SimArray(np.ones(nParticles+2, dtype=np.float32), pos_unit)
    snapshotBinary['metals'] = SimArray(np.zeros(nParticles+2, dtype=np.float32))
    snapshotBinary['vel'].units = v_unit
    snapshotBinary['pos'].units = pos_unit
    snapshotBinary['mass'].units = snapshot['mass'].units
    snapshotBinary['rho'] = SimArray(np.zeros(nParticles+2, dtype=np.float32))

    #Assign gas particles with calculated/given values from above
    snapshotBinary.gas['pos'] = snapshot.gas['pos']
    snapshotBinary.gas['vel'] = snapshot.gas['vel']
    snapshotBinary.gas['temp'] = snapshot.gas['temp']
    snapshotBinary.gas['rho'] = snapshot.gas['rho']
    snapshotBinary.gas['eps'] = snapshot.gas['eps']
    snapshotBinary.gas['mass'] = snapshot.gas['mass']
    snapshotBinary.gas['metals'] = snapshot.gas['metals']

    #Load Binary system obj to initialize system
    binsys = ICobj.settings.physical.binsys
    
    x1,x2,v1,v2 = binsys.generateICs()

    #Put velocity in sim units
    #!!! Note: v_unit_vel will always be 29.785598165 km/s when m_unit = Msol and r_unit = 1 AU in kpc!!!
    #conv = v_unit_vel #km/s in sim units
    #v1 /= conv
    #v2 /= conv

    #Assign position, velocity assuming CCW orbit

    snapshotBinary.star[0]['pos'] = SimArray(x1,pos_unit)
    snapshotBinary.star[0]['vel'] = SimArray(v1,v_unit)
    snapshotBinary.star[1]['pos'] = SimArray(x2,pos_unit)
    snapshotBinary.star[1]['vel'] = SimArray(v2,v_unit)

    #Set stellar masses
    #Set Mass units
    #Create simArray for mass, convert units to simulation mass units
    priMass = SimArray(binsys.m1,m_unit)
    secMass = SimArray(binsys.m2,m_unit)

    snapshotBinary.star[0]['mass'] = priMass
    snapshotBinary.star[1]['mass'] = secMass
    snapshotBinary.star['metals'] = SimArray(star_metals)

    #Estimate stars' softening length as fraction of distance to COM
    d = np.sqrt(AddBinary.dotProduct(x1-x2,x1-x2))

    snapshotBinary.star[0]['eps'] = SimArray(math.fabs(d)/4.0,pos_unit)
    snapshotBinary.star[1]['eps'] = SimArray(math.fabs(d)/4.0,pos_unit)
 
    print 'Wrapping up'
    # Now set the star particle's tform to a negative number.  This allows
    # UW ChaNGa treat it as a sink particle.
    snapshotBinary.star['tform'] = -1.0
    
    #Set Sink Radius to be mass-weighted average of Roche lobes of two stars
    r1 = AddBinary.calcRocheLobe(binsys.m1/binsys.m2,binsys.a) 
    r2 = AddBinary.calcRocheLobe(binsys.m2/binsys.m1,binsys.a)
    p = strip_units(binsys.m1/(binsys.m1 + binsys.m2))

    r_sink = (r1*p) + (r2*(1.0-p))
    param['dSinkBoundOrbitRadius'] = r_sink
    param['dSinkRadius'] = r_sink
    param['dSinkMassMin'] = 0.9 * strip_units(secMass)
    param['bDoSinks'] = 1
    
    return snapshotBinary, param, director
예제 #30
0
파일: _math.py 프로젝트: dflemin3/diskpy
def binned_mean(x, y, bins=10, nbins=None, binedges=None, weights=None, weighted_bins=False, ret_bin_edges=False):
    """
    Bins y according to x and takes the average for each bin.

    bins can either be an integer (the number of bins to use) or an array of
    binedges.  bins will be overridden by nbins or binedges

    Optionally (for compatibility reasons) if binedges is specified, the
    x-bins are defined by binedges.  Otherwise the x-bins are determined by
    nbins

    If weights = None, equal weights are assumed for the average, otherwise
    weights for each data point should be specified

    y_err (error in y) is calculated as the standard deviation in y for each
    bin, divided by sqrt(N), where N is the number of counts in each bin

    IF weighted_bins is True, the bin centers are calculated as a center of
    mass

    NaNs are ignored for the input.  Empty bins are returned with nans

    RETURNS a tuple of (bin_centers, y_mean, y_err) if ret_bin_edges=False
    else, Returns (bin_edges, y_mean, y_err)
    """
    if (isinstance(bins, int)) and (nbins is None):

        nbins = bins

    elif (hasattr(bins, "__iter__")) and (binedges is None):

        binedges = bins

    if binedges is not None:

        nbins = len(binedges) - 1

    else:

        binedges = np.linspace(x.min(), (1 + np.spacing(2)) * x.max(), nbins + 1)

    if weights is None:

        weights = np.ones(x.shape)

    weights = strip_units(weights)

    # Pre-factor for weighted STD:
    A = 1 / (1 - (weights ** 2).sum())

    # Initialize
    y_mean = np.zeros(nbins)
    y_std = np.zeros(nbins)
    # Find the index bins for each data point
    ind = np.digitize(x, binedges) - 1
    # Ignore nans
    nan_ind = np.isnan(y)
    N = np.histogram(x, binedges)[0]

    # Initialize bin_centers (try to retain units)
    bin_centers = 0.0 * binedges[1:]

    for i in range(nbins):

        # Indices to use
        mask = (ind == i) & (~nan_ind)
        # Set up the weighting
        w = weights[mask].copy()
        w /= w.sum()
        A = 1 / (1 - (w ** 2).sum())
        # y_mean[i] = np.nanmean(y[mask])
        y_mean[i] = (w * y[mask]).sum()
        var = A * (w * (y[mask] - y_mean[i]) ** 2).sum()
        y_std[i] = np.sqrt(var)
        # y_std[i] = np.std(y[use_ind])

        if weighted_bins:
            # Center of mass of x positions
            bin_centers[i] = (w * x[mask]).sum()

    y_mean = match_units(y_mean, y)[0]
    y_err = y_std / np.sqrt(N)
    y_err = match_units(y_err, y)[0]

    y_mean[N == 0] = np.nan
    y_err[N == 0] = np.nan

    if not weighted_bins:

        bin_centers = (binedges[0:-1] + binedges[1:]) / 2.0
        binedges = match_units(binedges, x)[0]
        bin_centers = match_units(bin_centers, x)[0]

    else:

        bin_centers[N == 0] = np.nan

    if ret_bin_edges:

        return binedges, y_mean, y_err

    else:

        return bin_centers, y_mean, y_err
예제 #31
0
def Q(snapshot, molecular_mass = 2.0, bins=100, use_velocity=False, \
use_omega=True, gamma=1.):
    """Calculates the Toomre Q as a function of r, assuming radial temperature
    profile and kappa ~= omega
    
    Parameters
    ----------
    
    snapshot : tipsy snapshot
    molecular_mass : float
        Mean molecular mass (for sound speed).  Default = 2.0
    bins : int or array
        Either the number of bins or the bin edges
    use_velocity : Bool
        Determines whether to use the particles' velocities to calculate orbital
        velocity.  Useful if the circular orbital velocities are set in the
        snapshot.
    use_omega : Bool
        Default=True.  Use omega as a proxy for kappa to reduce noise
    gamma : float
        'Effective' adiabatic index, used to calculate sound speed.  Use 
        gamma = 1 for an isothermal EOS

    Returns
    -------
    
    Q : array
        Toomre Q as a function of r
    r_edges : array
        Radial bin edges
    """

    # Physical constants
    kB = SimArray([1.0],'k')
    G = SimArray([1.0],'G')
    # Calculate surface density
    sig, r_edges = sigma(snapshot, bins)
    # Calculate sound speed
    m = match_units(molecular_mass,'m_p')[0]
    c_s_all = np.sqrt(kB*snapshot.g['temp']*gamma/m)
    # Bin/average sound speed
    dummy, c_s, dummy2 = binned_mean(snapshot.g['rxy'], c_s_all, binedges=r_edges)

    if use_omega:
        # Calculate keplerian angular velocity (as a proxy for the epicyclic
        # frequency, which is a noisy calculation)
        if use_velocity:
            # Calculate directly from particle's velocity
            dummy, omega, dummy2 = binned_mean(snapshot.g['rxy'], \
            snapshot.g['vt']/snapshot.g['rxy'], binedges=r_edges)

        else:
            # Estimate, from forces, using pynbody
            p = pb.analysis.profile.Profile(snapshot, bins=r_edges)
            omega = p['omega']

        kappa_calc = omega

    else:

        if use_velocity:
            # Calculate directly from particle's velocities
            kappa_calc, dummy = kappa(snapshot, r_edges)

        else:
            # Estimate, from forces, using pynbody
            p = pb.analysis.profile.Profile(snapshot, bins=r_edges)
            kappa_calc = p['kappa']

    return (kappa_calc*c_s/(np.pi*G*sig)).in_units('1'), r_edges