-
Notifications
You must be signed in to change notification settings - Fork 0
/
dark_loader.py
72 lines (55 loc) · 2.45 KB
/
dark_loader.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import yt
import numpy as np
from yt.utilities.sdf import load_sdf
from super_data import SuperData
from darksky_catalog import darksky
prefix = "http://darksky.slac.stanford.edu/simulations/"
def load_ds14_a(midx=None, bbox=None, a=1.0):
if midx == True:
midx = 10
ds = darksky['ds14_a'].load(bounding_box = bbox, midx = midx)
# the c code chokes on data types without these conversions
ds.domain_left_edge = ds.domain_left_edge.astype(np.float64)
ds.domain_right_edge = ds.domain_right_edge.astype(np.float64)
ds.domain_width = ds.domain_width.astype(np.float64)
ds.domain_center = ds.domain_center.astype(np.float64)
return ds
def load_ds14_a_halos(midx=7, bbox=None, a=1.0):
if a != 1.0:
raise RuntimeError("Only a=1.0 allowed at this time.")
ds = darksky['ds14_a']['halos_a_%0.4f' % a].load(
bounding_box = bbox, midx = midx)
return ds
def get_halo_params(nth_most_massive):
halo = darksky['ds14_a']['filtered_1e15_a_1.0000'].get_halo(nth_most_massive)
return halo
def get_halo_bounds(nth_most_massive, pdf, hdf):
particles_midx = pdf + ".midx10"
pds = yt.load(pdf, midx_filename=particles_midx)
halos_midx = hdf + ".midx7"
hds = yt.load(hdf, midx_filename=halos_midx)
halo = get_halo_params(nth_most_massive)
sds = SuperData(pds, hds)
hx, hy, hz, hr = halo['x'], halo['y'], halo['z'], halo['r200b']
px, py, pz = sds.particle_position(sds.halo_dataset.arr([hx, hy, hz], 'code_length')).d
pr = sds.conv_arr(sds.halo_dataset.arr([hr], 'code_length'), sds.full_particle_dataset).d
center = np.array([px, py, pz])
left = center - 2*pr
right = center + 2*pr
return halo, center, left, right, pr[0]
def load_halo(pdf, hdf, nth_most_massive):
halo, center, left, right, radius = get_halo_bounds(nth_most_massive, pdf, hdf)
particles_midx = pdf + ".midx10"
halos_midx = hdf + ".midx7"
bbox = np.array([left, right]).T
pds = yt.load(pdf,
midx_filename=particles_midx,
bounding_box = bbox,
n_ref=64, over_refine_factor=2,
)
# the c code chokes on data types without these conversions
pds.domain_left_edge = pds.domain_left_edge.astype(np.float64)
pds.domain_right_edge = pds.domain_right_edge.astype(np.float64)
pds.domain_width = pds.domain_width.astype(np.float64)
pds.domain_center = pds.domain_center.astype(np.float64)
return pds, halo, center, radius