def test_clump_tree_save():
tmpdir = tempfile.mkdtemp()
curdir = os.getcwd()
os.chdir(tmpdir)
ds = data_dir_load(i30)
data_source = ds.disk([0.5, 0.5, 0.5], [0.0, 0.0, 1.0], (8, "kpc"),
(1, "kpc"))
field = ("gas", "density")
step = 2.0
c_min = 10**np.floor(np.log10(data_source[field]).min())
c_max = 10**np.floor(np.log10(data_source[field]).max() + 1)
master_clump = Clump(data_source, field)
master_clump.add_info_item("center_of_mass")
master_clump.add_validator("min_cells", 20)
find_clumps(master_clump, c_min, c_max, step)
leaf_clumps = master_clump.leaves
fn = master_clump.save_as_dataset(fields=[
("gas", "density"),
("index", "x"),
("index", "y"),
("index", "z"),
("all", "particle_mass"),
])
ds2 = load(fn)
# compare clumps in the tree
t1 = [c for c in master_clump]
t2 = [c for c in ds2.tree]
mt1 = ds.arr([c.info["cell_mass"][1] for c in t1])
mt2 = ds2.arr([c["clump", "cell_mass"] for c in t2])
it1 = np.array(np.argsort(mt1).astype(int))
it2 = np.array(np.argsort(mt2).astype(int))
assert_array_equal(mt1[it1], mt2[it2])
for i1, i2 in zip(it1, it2):
ct1 = t1[i1]
ct2 = t2[i2]
assert_array_equal(ct1["gas", "density"], ct2["grid", "density"])
assert_array_equal(ct1["all", "particle_mass"], ct2["all",
"particle_mass"])
# compare leaf clumps
c1 = [c for c in leaf_clumps]
c2 = [c for c in ds2.leaves]
mc1 = ds.arr([c.info["cell_mass"][1] for c in c1])
mc2 = ds2.arr([c["clump", "cell_mass"] for c in c2])
ic1 = np.array(np.argsort(mc1).astype(int))
ic2 = np.array(np.argsort(mc2).astype(int))
assert_array_equal(mc1[ic1], mc2[ic2])
os.chdir(curdir)
shutil.rmtree(tmpdir)
class ClumpFinder:
def __init__(self, max_level, label="", file_cache=True):
if not os.path.exists(CUBE_DIR):
os.makedirs(CUBE_DIR)
if not os.path.exists(PLOT_DIR):
os.makedirs(PLOT_DIR)
if not os.path.exists(CLUMP_DIR):
os.makedirs(CLUMP_DIR)
self._cube_data = {}
self._ramses_ds = None
self._cube_ds = None
self._disk = None
self._master_clump = None
self._leaf_clumps = None
self._clump_quantities = None
self._molecular_clouds = None
self.max_level = int(max_level)
self.file_cache = file_cache
if label:
self.label = label
else:
self.label = max_level
@property
def ramses_ds(self):
if not self._ramses_ds:
self._ramses_ds = yt.load(RAMSES_INPUT_INFO)
return self._ramses_ds
def cube_data(self, sim_type):
if not sim_type in self._cube_data:
self._cube_data[sim_type] = RamsesData(
idir=RAMSES_INPUT_DIR,
sim_type=sim_type,
xmin=GALAXY_CENTRE[0] - CUBE_PADDING,
xmax=GALAXY_CENTRE[0] + CUBE_PADDING,
ymin=GALAXY_CENTRE[1] - CUBE_PADDING,
ymax=GALAXY_CENTRE[1] + CUBE_PADDING,
zmin=GALAXY_CENTRE[2] - CUBE_PADDING,
zmax=GALAXY_CENTRE[2] + CUBE_PADDING,
lmax=self.max_level,
save_dir=CUBE_DIR,
use_file_cache=self.file_cache,
)
return self._cube_data[sim_type]
@property
def cube_ds(self):
if not self._cube_ds:
self._cube_ds = yt.load_uniform_grid(
dict(
density=self.cube_data(SimTypes.DENSITY).cube,
velocity_x=self.cube_data(SimTypes.X_VELOCITY).cube,
velocity_y=self.cube_data(SimTypes.Y_VELOCITY).cube,
velocity_z=self.cube_data(SimTypes.Z_VELOCITY).cube,
pressure=self.cube_data(SimTypes.PRESSURE).cube,
),
self.cube_data(SimTypes.DENSITY).cube.shape,
# TODO: Fix scaling. Doesn't find many clumps with this enabled.
#length_unit=self.ramses_ds.length_unit/512,#3080*6.02,
)
return self._cube_ds
@property
def disk(self):
if not self._disk:
self._disk = self.cube_ds.disk(
GALAXY_CENTRE,
[0., 0., 1.],
(1, 'kpc'),
(0.5, 'kpc'),
)
return self._disk
@property
def master_clump(self):
if not self._master_clump:
clump_file = os.path.join(CLUMP_DIR,
'{}_clumps.h5'.format(self.max_level))
# TODO: Fix file format -- saved dataset loses attributes/isn't
# loaded as the right type
orig_file_cache = self.file_cache
self.file_cache = False
if self.file_cache and os.path.isfile(clump_file):
self._master_clump = yt.load(clump_file)
else:
self._master_clump = Clump(self.disk, ('gas', "density"))
find_clumps(
clump=self._master_clump,
min_val=self.disk["density"].min(),
max_val=self.disk["density"].max(),
d_clump=8.0, # Step size
)
if self.file_cache:
self._master_clump.save_as_dataset(clump_file, [
'density',
])
self.file_cache = orig_file_cache
return self._master_clump
@property
def leaf_clumps(self):
if not self._leaf_clumps:
self._leaf_clumps = self.master_clump.leaves
return self._leaf_clumps
@property
def clump_quantities(self):
if not self._clump_quantities:
self._clump_quantities = []
for clump in self.leaf_clumps:
self._clump_quantities.append({
'clump':
clump,
'volume':
clump.data.volume().to_value(),
'mass':
clump.data.quantities.total_mass().to_value()[0],
'velocity_x_mean':
clump.data['velocity_x'].mean(),
'velocity_y_mean':
clump.data['velocity_y'].mean(),
'velocity_z_mean':
clump.data['velocity_z'].mean(),
'velocity_x_var':
clump.data['velocity_x'].var(),
'velocity_y_var':
clump.data['velocity_y'].var(),
'velocity_z_var':
clump.data['velocity_z'].var(),
'pressure_mean':
clump.data['pressure'].mean(),
})
self._clump_quantities[-1]['density'] = (
self._clump_quantities[-1]['mass'] /
self._clump_quantities[-1]['volume'])
(
self._clump_quantities[-1]['bulk_velocity_0'],
self._clump_quantities[-1]['bulk_velocity_1'],
self._clump_quantities[-1]['bulk_velocity_2'],
) = clump.quantities.bulk_velocity().to_value()
return self._clump_quantities
@property
def molecular_clouds(self):
if not self._molecular_clouds:
self._molecular_clouds = [
cq for cq in self.clump_quantities
if cq['density'] >= CLOUD_DENSITY_THRESHOLD
]
return self._molecular_clouds
# Now find get our 'base' clump -- this one just covers the whole domain.
master_clump = Clump(data_source, field)
# Add a "validator" to weed out clumps with less than 20 cells.
# As many validators can be added as you want.
master_clump.add_validator("min_cells", 20)
# Calculate center of mass for all clumps.
master_clump.add_info_item("center_of_mass")
# Begin clump finding.
find_clumps(master_clump, c_min, c_max, step)
# Save the clump tree as a reloadable dataset
fn = master_clump.save_as_dataset(fields=[("gas",
"density"), ("all",
"particle_mass")])
# We can traverse the clump hierarchy to get a list of all of the 'leaf' clumps
leaf_clumps = master_clump.leaves
# Get total cell and particle masses for each leaf clump
leaf_masses = [leaf.quantities.total_mass() for leaf in leaf_clumps]
# If you'd like to visualize these clumps, a list of clumps can be supplied to
# the "clumps" callback on a plot. First, we create a projection plot:
prj = yt.ProjectionPlot(ds, 2, field, center="c", width=(20, "kpc"))
# Next we annotate the plot with contours on the borders of the clumps
prj.annotate_clumps(leaf_clumps)
# Now find get our 'base' clump -- this one just covers the whole domain.
master_clump = Clump(data_source, field)
# Add a "validator" to weed out clumps with less than 20 cells.
# As many validators can be added as you want.
master_clump.add_validator("min_cells", 20)
# Calculate center of mass for all clumps.
master_clump.add_info_item("center_of_mass")
# Begin clump finding.
find_clumps(master_clump, c_min, c_max, step)
# Save the clump tree as a reloadable dataset
fn = master_clump.save_as_dataset(fields=["density", "particle_mass"])
# We can traverse the clump hierarchy to get a list of all of the 'leaf' clumps
leaf_clumps = master_clump.leaves
# Get total cell and particle masses for each leaf clump
leaf_masses = [leaf.quantities.total_mass() for leaf in leaf_clumps]
# If you'd like to visualize these clumps, a list of clumps can be supplied to
# the "clumps" callback on a plot. First, we create a projection plot:
prj = yt.ProjectionPlot(ds, 2, field, center="c", width=(20, "kpc"))
# Next we annotate the plot with contours on the borders of the clumps
prj.annotate_clumps(leaf_clumps)
# Save the plot to disk.
output_dir = "clumps/"
ensure_dir(output_dir)
master_clump.add_validator("future_bound",
use_thermal_energy=True,
truncate=True,
include_cooling=True)
master_clump.add_info_item("center_of_mass")
master_clump.add_info_item("min_number_density")
master_clump.add_info_item("max_number_density")
master_clump.add_info_item("jeans_mass")
find_clumps(master_clump, c_min, c_max, step)
fn = master_clump.save_as_dataset(filename=output_dir,
fields=["density", "particle_mass"])
leaf_clumps = master_clump.leaves
pdir = os.path.join(output_dir, 'projections')
ensure_dir(pdir)
inner_radius = ds.quan(100, 'AU')
units = 'pc'
for i, sphere in enumerate(
iterate_center_of_mass(data_source,
inner_radius,
com_kwargs={
'use_gas': True,
'use_particles': True
})):
width = 2 * sphere.radius