def test_attr(comm):
fname = tempfile.mkdtemp()
x = BigFile(fname, create=True)
with x.create('.', dtype=None) as b:
b.attrs['int'] = 128
b.attrs['float'] = [128.0, 3, 4]
b.attrs['string'] = 'abcdefg'
b.attrs['complex'] = 128 + 128J
b.attrs['bool'] = True
b.attrs['arrayustring'] = numpy.array(u'unicode')
b.attrs['arraysstring'] = numpy.array('str')
with x.open('.') as b:
assert_equal(b.attrs['int'], 128)
assert_equal(b.attrs['float'], [128.0, 3, 4])
assert_equal(b.attrs['string'], 'abcdefg')
assert_equal(b.attrs['complex'], 128 + 128J)
assert_equal(b.attrs['bool'], True)
b.attrs['int'] = 30
b.attrs['float'] = [3, 4]
b.attrs['string'] = 'defg'
b.attrs['complex'] = 32 + 32J
b.attrs['bool'] = False
with x.open('.') as b:
assert_equal(b.attrs['int'], 30)
assert_equal(b.attrs['float'], [3, 4])
assert_equal(b.attrs['string'], 'defg')
assert_equal(b.attrs['complex'], 32 + 32J)
assert_equal(b.attrs['bool'], False)
shutil.rmtree(fname)
def test_attr(comm):
fname = tempfile.mkdtemp()
x = BigFile(fname, create=True)
with x.create('.', dtype=None) as b:
b.attrs['int'] = 128
b.attrs['float'] = [128.0, 3, 4]
b.attrs['string'] = 'abcdefg'
b.attrs['complex'] = 128 + 128J
with x.open('.') as b:
assert_equal(b.attrs['int'], 128)
assert_equal(b.attrs['float'], [128.0, 3, 4])
assert_equal(b.attrs['string'], 'abcdefg')
assert_equal(b.attrs['complex'], 128 + 128J)
b.attrs['int'] = 30
b.attrs['float'] = [3, 4]
b.attrs['string'] = 'defg'
b.attrs['complex'] = 32 + 32J
with x.open('.') as b:
assert_equal(b.attrs['int'], 30)
assert_equal(b.attrs['float'], [3, 4])
assert_equal(b.attrs['string'], 'defg')
assert_equal(b.attrs['complex'], 32 + 32J)
shutil.rmtree(fname)
def test_append(comm):
fname = tempfile.mkdtemp()
x = BigFile(fname, create=True)
name = 'f4'
d = numpy.dtype(('f4', 3))
numpy.random.seed(1234)
data = numpy.random.uniform(100000, size=(100, 3)).astype('f4')
# test creating
with x.create(name, Nfile=3, dtype=d, size=100) as b:
b.write(0, data)
b.append(data, Nfile=2)
with x.open(name) as bb:
assert bb.size == 200
assert b.size == 200
with x.open(name) as b:
assert b.Nfile == 5
assert_equal(b[:100], data)
assert_equal(b[100:], data)
assert b.size == 200
shutil.rmtree(fname)
def test_attr(comm):
fname = tempfile.mkdtemp()
x = BigFile(fname, create=True)
with x.create('.', dtype=None) as b:
b.attrs['int'] = 128
b.attrs['float'] = [128.0, 3, 4]
b.attrs['string'] = 'abcdefg'
b.attrs['complex'] = 128 + 128J
b.attrs['bool'] = True
b.attrs['arrayustring'] = numpy.array(u'unicode')
b.attrs['arraysstring'] = numpy.array('str')
with x.open('.') as b:
assert_equal(b.attrs['int'], 128)
assert_equal(b.attrs['float'], [128.0, 3, 4])
assert_equal(b.attrs['string'], 'abcdefg')
assert_equal(b.attrs['complex'], 128 + 128J)
assert_equal(b.attrs['bool'], True)
b.attrs['int'] = 30
b.attrs['float'] = [3, 4]
b.attrs['string'] = 'defg'
b.attrs['complex'] = 32 + 32J
b.attrs['bool'] = False
with x.open('.') as b:
assert_equal(b.attrs['int'], 30)
assert_equal(b.attrs['float'], [3, 4])
assert_equal(b.attrs['string'], 'defg')
assert_equal(b.attrs['complex'], 32 + 32J)
assert_equal(b.attrs['bool'], False)
shutil.rmtree(fname)
def test_grow(comm):
fname = tempfile.mkdtemp()
x = BigFile(fname, create=True)
d = numpy.dtype(('f4', 3))
numpy.random.seed(1234)
data = numpy.random.uniform(100000, size=(100, 3)).astype('f4')
# test creating
with x.create(d.str, Nfile=3, dtype=d, size=100) as b:
b.write(0, data)
b.grow(size=100, Nfile=2)
with x.open(d.str) as bb:
assert bb.size == 200
b.write(100, data)
assert b.size == 200
with x.open(d.str) as b:
assert b.Nfile == 5
assert_equal(b[:100], data)
assert_equal(b[100:], data)
assert b.size == 200
shutil.rmtree(fname)
def calc_mf_each_bf(dir_name: str, bf: BigFile):
""" Calculate MFs of a single bf file"""
header = bf.open('Header')
redshift = 1. / header.attrs['Time'][0] - 1.
bhmass = bf.open('5/BlackholeMass')[:] * TO_MSUN
halomass = bf.open('FOFGroups/Mass')[:] * TO_MSUN
starmass = bf.open('FOFGroups/MassByType')[:][:, 4] * TO_MSUN
halo_mf = mass_function(halomass, HALO_MIN, HALO_MAX, N_BIN, BOXSIZE)
star_mf = mass_function(starmass, STAR_MIN, STAR_MAX, N_BIN, BOXSIZE)
bh_mf = mass_function(bhmass, BH_MIN, BH_MAX, N_BIN, BOXSIZE)
print(' Saving MFs at z = %0.4f' % redshift)
np.save('{}halo_mf_%0.4f'.format(dir_name) %redshift, halo_mf)
np.save('{}star_mf_%0.4f'.format(dir_name) %redshift, star_mf)
np.save('{}bh_mf_%0.4f'.format(dir_name) %redshift, bh_mf)
def ReadIC(filename):
# this reads in a MP-Gadget3/GENIC format IC
# major thing is to scale vel by a0**1.5
file = BigFile(filename)
header = file.open('header')
BoxSize = header.attrs['BoxSize'][0]
a0 = header.attrs['Time'][0]
Ntot = file.open('1/ID').size
myslice = slice(
MPI.COMM_WORLD.rank * Ntot // MPI.COMM_WORLD.size,
(MPI.COMM_WORLD.rank + 1) * Ntot // MPI.COMM_WORLD.size,
)
P = dict()
P['Mass'] = header.attrs['MassTable'][1]
P['Position'] = file.open('1/Position')[myslice]
P['Velocity'] = file.open('1/Velocity')[myslice]
P['Velocity'] *= a0**1.5
P['ID'] = file.open('1/ID')[myslice]
return P, BoxSize, a0
def ReadIC(filename):
# this reads in a MP-Gadget3/GENIC format IC
# major thing is to scale vel by a0**1.5
file = BigFile(filename)
header = file.open('header')
BoxSize = header.attrs['BoxSize'][0]
a0 = header.attrs['Time'][0]
Ntot = file.open('1/ID').size
myslice = slice(
MPI.COMM_WORLD.rank * Ntot // MPI.COMM_WORLD.size,
(MPI.COMM_WORLD.rank + 1) * Ntot // MPI.COMM_WORLD.size,
)
P = dict()
P['Mass'] = header.attrs['MassTable'][1]
P['Position'] = file.open('1/Position')[myslice]
P['Velocity'] = file.open('1/Velocity')[myslice]
P['Velocity'] *= a0 ** 1.5
P['ID'] = file.open('1/ID')[myslice]
return P, BoxSize, a0
def append_mmbh_data(part: str, redshifts: List, mmbhmasss: List, mmbhids: List,
mmbhaccs: List, mmbhposs: List, mmbhvels: List, mmbht1s: List):
""" Append the most massive BHs' quantities. """
bf = BigFile(part)
header = bf.open('Header')
redshift = 1. / header.attrs['Time'][0] - 1.
bhmass = bf.open('5/BlackholeMass')[:] * TO_MSUN
no_blackhole = len(bhmass) == 0
if no_blackhole:
print(' No BH formed at z = %0.2f' % redshift)
return None
bhid = bf.open('5/ID')[:]
bhacc = bf.open('5/BlackholeAccretionRate')[:] * TO_MSUN_YEAR
bhpos = bf.open('5/Position')[:]
bhvel = bf.open('5/Velocity')[:]
mmbhmass = bhmass.max()
mmbhid = bhid[np.argmax(bhmass)]
mmbhacc = bhacc[np.argmax(bhmass)]
mmbhpos = bhpos[np.argmax(bhmass)]
mmbhvel = bhvel[np.argmax(bhmass)]
print(' Appending BH quantities at z = %0.4f' % redshift)
if GET_T1 & (redshift <= STARTZ) & (np.abs(redshift - round(redshift)) < DZ):
print(' -- calculating t1... at %0.4f' % redshift)
mmbht1 = [calc_t1(part, mmbhpos, nmesh) for nmesh in NMESHS]
else:
mmbht1 = [np.nan] * len(NMESHS)
redshifts.append(redshift)
mmbhmasss.append(mmbhmass)
mmbhids.append(mmbhid)
mmbhaccs.append(mmbhacc)
mmbhposs.append(mmbhpos)
mmbhvels.append(mmbhvel)
mmbht1s.append(mmbht1)
def append_merger_data(part: str, mergerid: np.ndarray, merger_datas: List):
""" Append the merger history of the most massive BH """
bf = BigFile(part)
header = bf.open('Header')
redshift = 1. / header.attrs['Time'][0] - 1.
bhmass = bf.open('5/BlackholeMass')[:] * TO_MSUN
bhid = bf.open('5/ID')[:]
bhacc = bf.open('5/BlackholeAccretionRate')[:] * TO_MSUN_YEAR
bhpos = bf.open('5/Position')[:]
bhvel = bf.open('5/Velocity')[:]
if len(bhmass) > 0:
mask_i = bhid == mergerid
if True in mask_i:
bhmass_i = bhmass[mask_i]
bhacc_i = bhacc[mask_i]
bhpos_i = bhpos[mask_i]
bhvel_i = bhvel[mask_i]
merger_data = [mergerid, redshift, bhmass_i[0], bhacc_i[0],
bhpos_i[0][0], bhpos_i[0][1], bhpos_i[0][2],
bhvel_i[0][0], bhvel_i[0][1], bhvel_i[0][2]]
merger_datas.append(merger_data)
def test_string(comm):
fname = tempfile.mkdtemp()
x = BigFile(fname, create=True)
# test creating
with x.create("Header", Nfile=1, dtype=None, size=128) as b:
b.attrs['v3'] = ['a', 'bb', 'ccc']
b.attrs['v32'] = [
['a', 'bb', 'ccc'],
['1', '22', '333'],]
b.attrs['s'] = 'abc'
b.attrs['l'] = 'a' * 65536
with x.open("Header") as b:
assert_equal(b.attrs['v3'], ['a', 'bb', 'ccc'])
assert_equal(b.attrs['v32'], ['a', 'bb', 'ccc', '1', '22', '333'])
assert_equal(b.attrs['s'], 'abc')
assert_equal(b.attrs['l'], 'a' * 65536)
shutil.rmtree(fname)
def test_string(comm):
fname = tempfile.mkdtemp()
x = BigFile(fname, create=True)
# test creating
with x.create("Header", Nfile=1, dtype=None, size=128) as b:
b.attrs['v3'] = ['a', 'bb', 'ccc']
b.attrs['v32'] = [
['a', 'bb', 'ccc'],
['1', '22', '333'],]
b.attrs['s'] = 'abc'
b.attrs['l'] = 'a' * 65536
with x.open("Header") as b:
assert_equal(b.attrs['v3'], ['a', 'bb', 'ccc'])
assert_equal(b.attrs['v32'], ['a', 'bb', 'ccc', '1', '22', '333'])
assert_equal(b.attrs['s'], 'abc')
assert_equal(b.attrs['l'], 'a' * 65536)
shutil.rmtree(fname)
print(tag)
#bigfile = BigFile(argv[1])
from argparse import ArgumentParser
ap = ArgumentParser()
ap.add_argument("config")
ap.add_argument("input")
ap.add_argument("output")
ns = ap.parse_args()
bigfile = BigFile(ns.input)
world = MPI.COMM_WORLD
if world.rank == 0:
attrs = bigfile.open("Header").attrs
HEADER = dict([(i, attrs[i]) for i in attrs])
else:
HEADER = None
HEADER = world.bcast(HEADER)
BoxSize = HEADER['BoxSize']
# set up some defaults
SMLFACTOR = 1.0
TilePadding = 256
M = [
[1, 1, 0],
[0, 1, 0],
[0, 0, 1]]
if __name__ == '__main__':
from sys import argv
# this will set the units to
#
# time: 980 Myear/h
# distance: 1 Kpc/h
# speed: 100 km/s
# mass: 1e10 Msun /h
DH = 3e5 / 100.
G = 43007.1
H0 = 0.1
Nmesh = int(argv[2])
file = BigFile(argv[1])
header = file.open('header')
BoxSize = header.attrs['BoxSize'][0]
a0 = header.attrs['Time'][0]
Ntot = file.open('1/ID').size
myslice = slice(
MPI.COMM_WORLD.rank * Ntot // MPI.COMM_WORLD.size,
(MPI.COMM_WORLD.rank + 1) * Ntot // MPI.COMM_WORLD.size,
)
P = lambda : None
P.Pos = file.open('1/Position')[myslice]
NumPart = len(P.Pos)
