Python pynbody.load函数代码示例

def pbload(filename, paramname=None):
    """
    Loads a snapshot using pynbody.  Can load a single species by appending
    ::gas, ::star, or ::dm to the filename
    
    Parameters
    ----------
    
    filename : str
        Filename to load
    paramname : str
        (optional) .param file to use
    
    Returns
    -------
    
    sim : snapshot
        A pynbody snapshot
    """
    if '::' in filename:
        
        filename, species = filename.split('::')
        sim = pb.load(filename, paramname=paramname)
        sim = getattr(sim, species)
        
    else:
        
        sim = pb.load(filename, paramname=paramname)
        
    return sim

def test_array_metadata() :
    f1 = pynbody.load("testdata/test_out.tipsy")
    
    f1.gas['zog'] = np.ones(len(f1.gas))
    f1.gas['zog'].units = "Msol kpc^-1"
    f1.gas['zog'].write()

    f1['banana'] = np.ones(len(f1))*0.5
    f1['banana'].units = "kpc^3 Myr^-1"
    f1['banana'].write()

    del f1

    f1 = pynbody.load("testdata/test_out.tipsy")
    assert "banana" in f1.loadable_keys()
    assert "zog" not in f1.loadable_keys()
    assert "banana" in f1.gas.loadable_keys()
    assert "zog" in f1.gas.loadable_keys()

    try:
        f1.star["zog"] # -> KeyError
        assert False # Shouldn't have been able to load gas-only array zog
    except KeyError :
        pass

    f1.gas['zog']
    assert f1.gas['zog'][0]==1.0
    assert f1.gas['zog'].units == "Msol kpc^-1"

    f1.star['banana']
    f1.gas['banana']
    f1.dm['banana']
    assert f1['banana'].units=="kpc^3 Myr^-1"

def test_unit_persistence():
    f = pynbody.load("testdata/test_g2_snap")

    # f2 is the comparison case - just load the whole
    # position array and convert it, simple
    f2 = pynbody.load("testdata/test_g2_snap")
    f2['pos']
    f2.physical_units()

    f.gas['pos']
    f.physical_units()
    assert (f.gas['pos'] == f2.gas['pos']).all()

    # the following lazy-loads should lead to the data being
    # auto-converted
    f.dm['pos']
    assert (f.gas['pos'] == f2.gas['pos']).all()
    assert (f.dm['pos'] == f2.dm['pos']).all()

    # the final one is the tricky one because this will trigger
    # an array promotion and hence internally inconsistent units
    f.star['pos']

    assert (f.star['pos'] == f2.star['pos']).all()

    # also check it hasn't messed up the other bits of the array!
    assert (f.gas['pos'] == f2.gas['pos']).all()
    assert (f.dm['pos'] == f2.dm['pos']).all()

    assert (f['pos'] == f2['pos']).all()

def test_snapshot_update() :
    f1 = pynbody.load("testdata/test_out.tipsy")
    f1['pos']
    f1['pos'] = np.arange(0,len(f1)*3).reshape(len(f1),3)
    
    # convert units -- the array should get written out in simulation units
    f1.g['pos'].convert_units('Mpc')

    f1['pos'].write(overwrite=True)
    f1.gas['metals'] = np.ones(len(f1.gas))*123.
    f1.star['metals'] = np.ones(len(f1.star))*345.

    f1.gas['metals'].write(overwrite=True)
    del f1

    f2 = pynbody.load("testdata/test_out.tipsy")
    assert (f2['pos']==np.arange(0,len(f2)*3).reshape(len(f2),3)).all()
    assert (f2.gas['metals']==123.).all() # should have updated gas.metals
    assert not (f2.star['metals']==345.).any() # should not have written out changes to star.metals

    # this is a completion:
    f2.dm['metals'] = np.ones(len(f2.dm))*789.1

    # should now be a simulation-level array... write it out
    f2['metals'].write(overwrite=True)

    del f2['metals']

    
    f3 = pynbody.load("testdata/test_out.tipsy")
    assert (f3.dm['metals']==789.1).all()

def test_fam_sim():
    """Check that an array loaded as families is the same as one loaded as a simulation array"""
    snap2 = pynbody.load("testdata/test_g2_snap")
    snap3 = pynbody.load("testdata/test_g2_snap")
    snap3.gas["pos"]
    snap3.dm["pos"]
    snap3.star["pos"]
    assert((snap3["pos"] == snap2["pos"]).all())

def test_units_override():
    f = pynbody.load("testdata/test_g2_snap.0")
    assert_equals(f.filename, "testdata/test_g2_snap")  # final '.0' is stripped
    assert_equals(f['pos'].units, "kpc a h^-1")

    # In this case the unit override system is not effective because
    # the final ".1" is not stripped away in the filename:
    # the file `test_g2_snap.units` is not used
    f_no_unit_override = pynbody.load("testdata/test_g2_snap.1")
    assert_equals(f_no_unit_override.filename, "testdata/test_g2_snap.1")
    assert_equals(f_no_unit_override['pos'].units, "Mpc a h^-1")  # from default_config.ini

def test_indexing():
    f1 = pynbody.load("testdata/g15784.lr.01024")

    test_len = 12310
    for test_len in [100, 10000, 20000]:
        for i in range(5):
            subindex = np.random.permutation(np.arange(0, len(f1)))[:test_len]
            subindex.sort()
            f2 = pynbody.load("testdata/g15784.lr.01024", take=subindex)

            assert (f2['x'] == f1[subindex]['x']).all()
            assert (f2['iord'] == f1[subindex]['iord']).all()

def getAllDust(lowz,catfiles,Rv,A1600max=2.0,halonum=[1],filename='dust.pkl'):
	f = open(catfiles,'r')
	files = f.readlines()
	dustExt = {'halos':halonum,'Rhl':np.zeros((len(halonum),len(files)+1)),'RDhalf':np.zeros((len(halonum),len(files)+1)),'z':np.zeros(len(files)+1),'u':np.zeros((len(halonum),len(files)+1)),'b':np.zeros((len(halonum),len(files)+1)),'v':np.zeros((len(halonum),len(files)+1)),'r':np.zeros((len(halonum),len(files)+1)),'i':np.zeros((len(halonum),len(files)+1)),'j':np.zeros((len(halonum),len(files)+1)),'h':np.zeros((len(halonum),len(files)+1)),'k':np.zeros((len(halonum),len(files)+1))}
	slz = pynbody.load(lowz)
        lz = str(round(slz.properties['a']**-1 -1,3))
        catend = '.cat.z'+lz+'\n'
        for i in range(len(files)):
                print "calculating dust corrections for ", files[i].strip('\n')
                xx = files[i].find(catend)
                simname=files[i][0:xx]
                s = pynbody.load(simname)
                h = s.halos()
                dustExt['z'][i] = s.properties['a']**-1 -1
                s.physical_units()
                catf = open(files[i].strip('\n'))
                cat = pickle.load(catf)
                catf.close()
		for j in range(len(halonum)):
                        print "halo", halonum[j]
                        progs, = np.where(cat==halonum[j])
                        if len(progs)==0:
                                print "no progenitors found in this step!"
                                continue
                        main = progs[0]
                        print "progenitor", main
                        h1 = h[main]
			dust,Rhalf = dustCor(h1,s,Rv,A1600max=A1600max)
			dustExt['RDhalf'][j,i] = Rhalf
#			dustExt['Rhl'][j,i] = Rhl
			for key in dust.keys():
				dustExt[key][j,i] = dust[key]
		del(s)
                del(h)
                del(h1)
                del(cat)
                gc.collect()
	print "calculating values for final step"
        h = slz.halos()
        slz.physical_units()
        for j in range(len(halonum)):
                h1 = h[halonum[j]]
		dust,Rhalf= dustCor(h1,slz,Rv)
		dustExt['z'][i+1] = slz.properties['a']**-1 -1
		dustExt['RDhalf'][j,i+1] = Rhalf
		for key in dust.keys():
			dustExt[key][j,i+1] = dust[key]
	if filename:
                print "saving data..."
                f = open(filename,'wb')
                pickle.dump(dustExt,f)
                f.close()
	return dustExt

def test_partial_loading() :
    f_f = pynbody.load("testdata/nchilada_test/12M.00001")
    
    test_ptcls = [ 11634,  24181,  26275,  37336,  37795,  38040,  38280,  38327,
        38524,  39349,  46758,  48892,  52160,  53267,  53745,  68970,
        78073,  83777,  86865,  93492,  94596,  96567,  99713, 106100,
       107856, 111036, 111830, 112560, 115082, 117111, 117444, 117667,
       123604, 123665, 124911, 132957, 138551, 154869, 158919, 182131,
       184252, 190498, 197946, 198288, 204526, 221720, 226375, 226915,
       229959, 231778] # randomly generated sample
    
    f_p = pynbody.load("testdata/nchilada_test/12M.00001", take = test_ptcls )

    assert((f_p['pos']==f_f['pos'][test_ptcls]).all())

def haloCat(lowz,highz,nhalos=50):
	print "matching halo catalog from ", highz, "to ", lowz
	s1 = pynbody.load(highz)
	s2 = pynbody.load(lowz)
	if s1.properties['a'] > s2.properties['a']:
		print "uh oh! highz file must actually be at higher z!"
		return
	b = pynbody.bridge.OrderBridge(s1,s2)
	cat = b.match_catalog()
	filename = highz+'.cat.z'+str(round(s2.properties['a']**-1-1,3))
	f = open(filename,'wb')
	pickle.dump(cat,f)
	f.close()
	return

def test_per_particle_loading():
    """Tests that loading one family at a time results in the
    same final array as loading all at once. There are a number of
    subtelties in the gadget handler that could mess this up by loading
    the wrong data."""

    f_all = pynbody.load("testdata/test_g2_snap")
    f_part = pynbody.load("testdata/test_g2_snap")

    f_part.dm['pos']
    f_part.star['pos']
    f_part.gas['pos']

    assert (f_all['pos'] == f_part['pos']).all()

def test_mass_unit_sanity():
    """Picks up on problems with converting array units as
    mass array gets loaded (which is a combination of a derived
    array and a loaded array)"""

    f1 = pynbody.load("testdata/ramses_partial_output_00250")
    f1['mass']
    f1.physical_units()

    f2 = pynbody.load("testdata/ramses_partial_output_00250")
    f2.physical_units()
    f2['mass']

    np.testing.assert_allclose(f1.dm['mass'], f2.dm['mass'], atol=1e-5)

def test_array_update():
    f1 = pynbody.load("testdata/test_out.tipsy")

    f1['bla'] = np.zeros(len(f1))
    f1['bla'].units = 'km'
    f1['bla'].write()

    del(f1['bla'])

    f1['bla']

    f1.g['bla'] = 1
    f1.d['bla'] = 2
    f1.s['bla'] = 3

    # test the case where bla is a snapshot-level array

    try:
        f1.g['bla'].write()
        assert False  # should not be allowed to overwrite here
    except IOError:
        pass

    f1.write_array(
        'bla', [pynbody.family.gas, pynbody.family.dm], overwrite=True)

    del(f1['bla'])

    f1['bla']

    assert all(f1.g['bla'] == 1)
    assert all(f1.d['bla'] == 2)
    assert all(f1.s['bla'] == 0)

    # test the case where bla2 is a family-level array

    f1.g['bla2'] = np.zeros(len(f1.g))
    f1.g['bla2'].units = 'km'
    f1.s['bla2'] = np.ones(len(f1.s))

    f1.write_array('bla2', [pynbody.family.gas, pynbody.family.star])

    del(f1)

    f1 = pynbody.load("testdata/test_out.tipsy")

    assert all(f1.g['bla2'] == 0)
    assert all(f1.s['bla2'] == 1)

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 writeBHMark(simname,step,Name=None,iord=False,massrange=False):
	if not Name: f = open('BH.'+step+'.mark','w')
	else: f = open(Name,'w')
	s = pynbody.load(simname+'.'+step)
	f.write(str(len(s))+' '+str(len(s.gas))+' '+str(len(s.star))+'\n')
	if not iord: 
		if not massrange:
			bhind, = np.where(s.stars['tform']<0)
		else:
			if len(massrange) != 2:
				print "error massrange must be a length 2 tuple!"
				return
			bhind, = np.where((s.stars['tform']<0)&(s.stars['mass'].in_units('Msol')<massrange[1])&(s.stars['mass'].in_units('Msol')>massrange[0]))
	else:	
		bhind = np.array([])
		for ii in range(len(iord)):
			tmpind, = np.where(s.stars['iord']==iord[ii])
			if len(tmpind)==0: print "uh oh... iord ", iord[ii]," not found!"
			bhind = np.append(bhind,tmpind)
	bhindreal = bhind+len(s.dark)+len(s.gas)+1
	for ii in range(len(bhindreal)):
		f.write(str(bhindreal[ii])+'\n')
	f.close()
	del(s)
	return bhindreal