Python pynbody.new函数代码示例

def test_monotonic_incomplete_bridge():
    # Test the monotonic bridge where not all the particles are shared
    f1 = pynbody.new(dm=10)
    f2 = pynbody.new(dm=9)
    f1['iord'] = np.arange(0, 10)
    f2['iord'] = np.array([0, 1, 2, 4, 5 ,6, 7, 8, 9])
    b = pynbody.bridge.OrderBridge(f1, f2, monotonic=False)

    assert (b(f1[:5])['iord'] == np.array([0, 1, 2, 4])).all()

def test_order_bridge():
    f1 = pynbody.new(dm=1000)
    f2 = pynbody.new(dm=3000)
    f1['iord'] = np.arange(5, 2005, 2, dtype=int)
    f2['iord'] = np.arange(3000, dtype=int)

    b = pynbody.bridge.OrderBridge(f1, f2)

    h1 = f1[:50:2]
    assert b(h1).ancestor is f2
    assert (b(h1)['iord'] == h1['iord']).all

def test_family_morphing():
    f1 = pynbody.new(dm=5, gas=5)
    f2 = pynbody.new(dm=10)

    # set dm and gas iords separately as it's possible the new command initialises them out of order
    f1.dm['iord'] = np.arange(0,5)
    f1.gas['iord'] = np.arange(5,10)
    f2['iord'] = np.array([0,2,4,6,8,1,3,5,7,9])

    b = pynbody.bridge.OrderBridge(f1,f2,monotonic=False,allow_family_change=True)

    assert (b(f2).dm['iord']==np.array([0,2,4,1,3])).all()
    assert (b(f2).gas['iord'] == np.array([6, 8, 5, 7, 9])).all()

def test_interpsnapshotKeplerPotential_eval():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 1.
    s['eps']= 0.
    sp= potential.InterpSnapshotRZPotential(s,
                                            rgrid=(0.01,2.,201),
                                            zgrid=(0.,0.2,201),
                                            logR=False,
                                            interpPot=True,
                                            zsym=True,
                                            numcores=1)
    kp= potential.KeplerPotential(amp=1.) #should be the same
    #This just tests on the grid
    rs= numpy.linspace(0.01,2.,21)
    zs= numpy.linspace(-0.2,0.2,41)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp(r,z)-kp(r,z))/kp(r,z)) < 10.**-10., 'RZPot interpolation w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g)' % (r,z)
    #This tests within the grid
    rs= numpy.linspace(0.01,2.,10)
    zs= numpy.linspace(-0.2,0.2,20)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp(r,z)-kp(r,z))/kp(r,z)) < 10.**-5., 'RZPot interpolation w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g) by %g' % (r,z,numpy.fabs((sp(r,z)-kp(r,z))/kp(r,z)))           
    #Test all at the same time to use vector evaluation
    mr,mz= numpy.meshgrid(rs,zs)
    mr= mr.flatten()
    mz= mz.flatten()
    assert numpy.all(numpy.fabs((sp(mr,mz)-kp(mr,mz))/kp(mr,mz)) < 10.**-5.), 'RZPot interpolation w/ interpRZPotential fails for vector input'
    return None

def test_interpsnapshotKeplerPotential_eval_naz():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 1.
    s['eps']= 0.
    sp= potential.InterpSnapshotRZPotential(s,
                                            rgrid=(0.01,2.,51),
                                            zgrid=(0.,0.2,51),
                                            logR=False,
                                            interpPot=True,
                                            zsym=True,
                                            numcores=1)
    spaz= potential.InterpSnapshotRZPotential(s,
                                              rgrid=(0.01,2.,51),
                                              zgrid=(0.,0.2,51),
                                              logR=False,
                                              interpPot=True,
                                              zsym=True,
                                              numcores=1,nazimuths=12)
    #This just tests on the grid
    rs= numpy.linspace(0.01,2.,21)
    zs= numpy.linspace(-0.2,0.2,41)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp(r,z)-spaz(r,z))/sp(r,z)) < 10.**-10., 'RZPot interpolation w/ InterpSnapShotPotential of KeplerPotential with different nazimuths fails at (R,z) = (%g,%g)' % (r,z)
    #This tests within the grid, with vector evaluation
    rs= numpy.linspace(0.01,2.,10)
    zs= numpy.linspace(-0.2,0.2,20)
    mr,mz= numpy.meshgrid(rs,zs)
    mr= mr.flatten()
    mz= mz.flatten()
    assert numpy.all(numpy.fabs((sp(mr,mz)-spaz(mr,mz))/sp(mr,mz)) < 10.**-5.), 'RZPot interpolation w/ interpRZPotential with different nazimimuths fails for vector input'
    return None

def test_interpsnapshotKeplerPotential_z2deriv():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 2.
    s['eps']= 0.
    sp= potential.InterpSnapshotRZPotential(s,
                                            rgrid=(0.01,2.,101),
                                            zgrid=(0.,0.2,101),
                                            logR=False,
                                            interpPot=True,
                                            interpverticalfreq=True,
                                            zsym=True)
    kp= potential.KeplerPotential(amp=2.) #should be the same
    #This just tests on the grid
    rs= numpy.linspace(0.01,2.,21)[1:]
    zs= numpy.linspace(-0.2,0.2,41)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp.z2deriv(r,z)-kp.z2deriv(r,z))/kp.z2deriv(r,z)) < 10.**-4., 'RZPot interpolation of z2deriv w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g) by %g' % (r,z,numpy.fabs((sp.z2deriv(r,z)-kp.z2deriv(r,z))/kp.z2deriv(r,z)))
    #This tests within the grid
    rs= numpy.linspace(0.01,2.,10)[1:]
    zs= numpy.linspace(-0.2,0.2,20)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp.z2deriv(r,z)-kp.z2deriv(r,z))/kp.z2deriv(r,z)) < 2.*10.**-4., 'RZPot interpolation of z2deriv w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g) by %g' % (r,z,numpy.fabs((sp.z2deriv(r,z)-kp.z2deriv(r,z))/kp.z2deriv(r,z)))
    return None

def test_interpsnapshotKeplerpotential_Rzderiv():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 2.
    s['eps']= 0.
    sp= potential.InterpSnapshotRZPotential(s,
                                            rgrid=(0.01,2.,101),
                                            zgrid=(0.,0.2,101),
                                            logR=False,
                                            interpPot=True,
                                            interpepifreq=True,
                                            interpverticalfreq=True,
                                            zsym=True)
    kp= potential.KeplerPotential(amp=2.) #should be the same
    #This just tests on the grid
    rs= numpy.linspace(0.01,2.,21)[1:]
    zs= numpy.linspace(-0.2,0.2,41)
    zs= zs[zs != 0.]# avoid zero
    # Test, but give small |z| a less constraining 
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp.Rzderiv(r,z)-kp.Rzderiv(r,z))/kp.Rzderiv(r,z)) < 10.**-4.*(1.+19.*(numpy.fabs(z) < 0.05)), 'RZPot interpolation of Rzderiv w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g) by %g; value is %g' % (r,z,numpy.fabs((sp.Rzderiv(r,z)-kp.Rzderiv(r,z))/kp.Rzderiv(r,z)),kp.Rzderiv(r,z))
    #This tests within the grid
    rs= numpy.linspace(0.01,2.,10)[1:]
    zs= numpy.linspace(-0.2,0.2,20)
    for r in rs:
        for z in zs:
            assert numpy.fabs((sp.Rzderiv(r,z)-kp.Rzderiv(r,z))/kp.Rzderiv(r,z)) < 10.**-4.*(1.+19.*(numpy.fabs(z) < 0.05)), 'RZPot interpolation of Rzderiv w/ InterpSnapShotPotential of KeplerPotential fails at (R,z) = (%g,%g) by %g' % (r,z,numpy.fabs((sp.Rzderiv(r,z)-kp.Rzderiv(r,z))/kp.Rzderiv(r,z)))
    return None

def setup():
    global f, h
    f = pynbody.new(dm=1000, star=500, gas=500, order='gas,dm,star')
    f['pos'] = np.random.normal(scale=1.0, size=f['pos'].shape)
    f['vel'] = np.random.normal(scale=1.0, size=f['vel'].shape)
    f['mass'] = np.random.uniform(1.0, 10.0, size=f['mass'].shape)
    f.gas['rho'] = np.ones(500, dtype=float)

def test_float_kd():
    f = pynbody.load("testdata/test_g2_snap")
    del f.properties['boxsize']

    assert f.dm['mass'].dtype==f.dm['pos'].dtype==np.float32
    assert f.dm['smooth'].dtype==np.float32

    # make double copy
    g = pynbody.new(len(f.dm))
    g.dm['pos']=f.dm['pos']
    g.dm['mass']=f.dm['mass']

    assert g.dm['mass'].dtype==g.dm['pos'].dtype==g.dm['smooth'].dtype==np.float64

    # check smoothing lengths agree (they have been calculated differently
    # using floating/double routines)

    npt.assert_allclose(f.dm['smooth'],g.dm['smooth'],rtol=1e-4)
    npt.assert_allclose(f.dm['rho'],g.dm['rho'],rtol=1e-4)

    # check all combinations of float/double smoothing
    double_ar = np.ones(len(f.dm),dtype=np.float64)
    float_ar = np.ones(len(f.dm),dtype=np.float32)

    double_double = g.dm.kdtree.sph_mean(double_ar,32)
    double_float = g.dm.kdtree.sph_mean(float_ar,32)
    float_double = f.dm.kdtree.sph_mean(double_ar,32)
    float_float = f.dm.kdtree.sph_mean(float_ar,32)

    # take double-double as 'gold standard' (though of course if any of these
    # fail it could also be a problem with the double-double case)

    npt.assert_allclose(double_double,double_float,rtol=1e-4)
    npt.assert_allclose(double_double,float_double,rtol=1e-4)
    npt.assert_allclose(double_double,float_float,rtol=1e-4)

def test_interpsnapshotKeplerPotential_normalize_units():
    # Set up a snapshot with just one unit mass at the origin
    s= pynbody.new(star=1)
    s['mass']= 4.
    s['eps']= 0.
    s['pos'].units= 'kpc'
    s['vel'].units= 'km s**-1'
    sp= potential.InterpSnapshotRZPotential(s,
                                          rgrid=(0.01,3.,201),
                                          zgrid=(0.,0.2,201),
                                          logR=False,
                                          interpPot=True,
                                          zsym=True)
    #Currently unnormalized
    assert numpy.fabs(sp.Rforce(1.,0.)+4.) < 10.**-7., "InterpSnapShotPotential that is assumed to be unnormalized doesn't behave as expected"
    # Normalize
    sp.normalize(R0=1.)
    assert numpy.fabs(sp.Rforce(1.,0.)+1.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    # De normalize
    sp.denormalize()
    assert numpy.fabs(sp.Rforce(1.,0.)+4.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    # Also test when R0 =/= 1
    sp.normalize(R0=2.)
    assert numpy.fabs(sp.Rforce(1.,0.)+1.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    # De normalize
    sp.denormalize()
    assert numpy.fabs(sp.Rforce(1.,0.)+4.) < 10.**-7., "InterpSnapShotPotential that is assumed to be normalized doesn't behave as expected"
    return None

 def __init__(self, **kwargs):
     self.__kwargs = kwargs
     self.__profile = kwargs.get('profile', density_profiles.alphabetagamma)
     self.__drhodr = kwargs.get('drhodr', density_profiles.dalphabetagammadr)
     self.__d2rhodr2 = kwargs.get('d2rhodr2', density_profiles.d2alphabetagammadr2)
     self.__pars = kwargs.get('pars', {'alpha': 1., 'beta': 3., 'gamma': 1.,
         'c': 10., 'factor': 0.1})
     if self.__profile == density_profiles.alphabetagamma and self.__pars['beta'] <= 3.:
         if 'factor' not in self.__pars.keys(): self.__pars['factor'] = 0.1
     self.__m_vir = kwargs.get('m_vir', '1e12 Msol')
     self.__m_vir = units.Unit(self.__m_vir)
     self.__h = kwargs.get('h', 0.7)
     self.__overden = kwargs.get('overden', 200.)
     self.__r_vir = tools.calc_r_vir(self.__m_vir, self.__h, self.__overden)
     self.__r_s = self.__r_vir/self.__pars['c']
     self.__n_particles = int(kwargs.get('n_particles', 1e5))
     self.__logxmax_rho = np.log10(self.__pars['c']) + 2.
     # Make sure to sample well inside the gravitational softening
     self.__logxmin_rho = self.__logxmax_rho - .5*np.log10(self.__n_particles) - 3.
     self.__logxmin_dist_func = kwargs.get('logxmin_dist_func', -3.)
     self.__logxmax_dist_func = kwargs.get('logxmax_dist_func', 14.)
     self.__n_sample_rho = int(kwargs.get('n_sample_rho', 1e4))
     self.__n_sample_dist_func = int(kwargs.get('n_sample_dist_func', 1e2))
     self.__n_sample_dist_func_rho = int(kwargs.get('n_sample_dist_func_rho', 1e4))
     self.__random_seed = kwargs.get('random_seed', 4)
     if 'prng' in kwargs.keys():
         self.__prng = kwargs['prng']
     else:
         self.__prng = np.random.RandomState(self.__random_seed)
     self.__spline_order = kwargs.get('spline_order', 3)
     self.__progress_bar = kwargs.get('progress_bar', False)
     self.__no_bulk_vel = kwargs.get('no_bulk_vel', True)
     self.__x_rho = np.logspace(self.__logxmin_rho, self.__logxmax_rho, self.__n_sample_rho)
     self.__f_bary = kwargs.get('f_bary', 0.1)
     self.__mu = kwargs.get('mu', 1.3)
     self.__spin_parameter = kwargs.get('spin_parameter', 0.04)
     self.__rot_balanced = kwargs.get('rot_balanced', False)
     # Different gas profiles are not yet implemented or successfully tested
     self.__gas_profile = self.__profile
     self.__gas_pars = self.__pars
     #self.__gas_profile = kwargs.get('gas_profile', density_profiles.alphabetagamma)
     #self.__gas_pars = kwargs.get('gas_pars', {'alpha': 1., 'beta': 3., 'gamma': 1.,
     #   'c': 10., 'factor': 0.1})
     self.__r_s_gas = self.__r_vir/self.__gas_pars['c']
     #self.__vel_prof = kwargs.get('vel_prof', None)
     self.__vel_pars = kwargs.get('vel_pars', {'rs_v': array.SimArray(1., 'kpc'),
         'c': self.__pars['c'], 'prefac': 1., 'factor': 1.})
     self.__n_gas_particles = int(kwargs.get('n_gas_particles', self.__n_particles))
     self.__ang_mom_prof = kwargs.get('ang_mom_prof', am_profiles.bullock_prof)
     self.__ang_mom_pars = kwargs.get('ang_mom_pars', {'mu': self.__mu})
     self.__fname = kwargs.get('fname', 'halo.out')
     # Careful here: as of now, only the output as tipsy files has been successfully tested
     self.__type = {'gadget': gadget.GadgetSnap,
                     'grafic': grafic.GrafICSnap,
                     'nchilada': nchilada.NchiladaSnap,
                     'ramses': ramses.RamsesSnap,
                     'tipsy': tipsy.TipsySnap}[kwargs.get('type', 'tipsy')]
     self.sim = new(dm=self.__n_particles, gas=self.__n_gas_particles)
     self.sim.physical_units()

def test_family_filter():
    f = pynbody.new(dm=100,gas=100)
    f_dm = f.dm
    f_dm_filter = f[pynbody.filt.FamilyFilter(pynbody.family.dm)]
    f_gas = f.gas
    f_gas_filter = f[pynbody.filt.FamilyFilter(pynbody.family.gas)]
    assert (f_dm.get_index_list(f) == f_dm_filter.get_index_list(f)).all()
    assert (f_gas.get_index_list(f) == f_gas_filter.get_index_list(f)).all()

def test_one_family_promotion():
    fx = pynbody.new(dm=10)
    fx.dm['bla'] = np.arange(10)
    # should have been made as a full-simulation array

    assert 'bla' in fx.keys()
    fx['bla']
    del fx

def setup():
    global f
    f = pynbody.new(1000)
    f['pos'] = np.random.normal(scale=1.0, size=f['pos'].shape)
    f['vel'] = np.random.normal(scale=1.0, size=f['vel'].shape)
    f['mass'] = np.random.uniform(1.0, 10.0, size=f['mass'].shape)
    f['pos'].units = 'kpc'
    f['vel'].units = 'km s^-1'
    f['mass'].units = 'Msol'

def test_aform_saturation():
    """Test that NaN is returned when tform cannot be calculated from aform"""
    ipoints = pynbody.analysis.cosmology._interp_points
    f = pynbody.new(ipoints + 1)
    f['aform'] = np.linspace(0.0,1.1,ipoints+1)-0.05
    tf = f['tform'][::100]
    assert tf[0]!=tf[0] # nan outside range
    assert tf[-1]!=tf[-1] # nan outside range
    assert (tf[1:-1]==tf[1:-1]).all() # no nans inside range