New object oriented version using HDF5
This commit is contained in:
Noe Brucy
@@ -0,0 +1,276 @@
|
||||
# coding: utf-8
|
||||
|
||||
import sys
|
||||
import os
|
||||
|
||||
import tables
|
||||
import pymses
|
||||
import numpy as np
|
||||
from numpy.polynomial.polynomial import polyfit
|
||||
from pymses.sources.ramses import output
|
||||
from pymses.sources.hop.file_formats import *
|
||||
from pymses.analysis import Camera, raytracing, slicing, splatting
|
||||
from pymses.filters import CellsToPoints
|
||||
from pymses.analysis import ScalarOperator, FractionOperator, MaxLevelOperator
|
||||
|
||||
from pp_params import *
|
||||
|
||||
class Rule:
|
||||
|
||||
def __init__(self, process, description, group='', dependencies=[], axes=['x', 'y', 'z'],
|
||||
is_valid=lambda save, ax:True):
|
||||
self.process_fn = process
|
||||
self.dependencies = dependencies
|
||||
self.is_valid_add = is_valid
|
||||
self.group = group
|
||||
self.axes = axes
|
||||
self.description = description
|
||||
|
||||
def process(self, ax_los):
|
||||
return self.process_fn(ax_los)
|
||||
|
||||
def is_valid(self, save, ax):
|
||||
valid = True
|
||||
for dep in self.dependencies:
|
||||
valid = valid and self.group + '/' + dep + '_' + ax in save
|
||||
return ax in self.axes and valid and self.is_valid_add(save, ax)
|
||||
|
||||
class PostProcessor:
|
||||
"""
|
||||
This class enable to compute and save derived quantities from the raw output
|
||||
"""
|
||||
|
||||
# Axes information
|
||||
_ax_nb = {'x' : 0, 'y' : 1, 'z' : 2} # Number of each axes
|
||||
_axes_h = {'x' :'y', 'y' : 'x', 'z' : 'x'} # Associated horizontal axe
|
||||
_axes_v = {'x' : 'z', 'y' : 'z', 'z' : 'y'} # Associated vertical axe
|
||||
|
||||
G = 1. # Gravitational constant
|
||||
|
||||
|
||||
def __init__(self, path, num, path_out=None, pp_params=Params()):
|
||||
"""
|
||||
Creates the basic structures needed for the outputs
|
||||
"""
|
||||
|
||||
# TODO : Make possible to load the HDF5 file even without the original file
|
||||
self.pp_params = pp_params
|
||||
|
||||
# Determining output directory
|
||||
if (path_out is None):
|
||||
path_out = path
|
||||
|
||||
# Open outfile
|
||||
if not pp_params.out.tag == '':
|
||||
tag_name = '_' + pp_params.out.tag
|
||||
else :
|
||||
tag_name = ''
|
||||
|
||||
self.filename = (path_out + '/postproc_' +
|
||||
tag_name + format(num,'05') + '.h5')
|
||||
self.save = tables.open_file(self.filename, mode="a",
|
||||
title=os.path.basename(path) + format(num,'05'))
|
||||
|
||||
# Ramses Output
|
||||
self._ro = pymses.RamsesOutput(path, num, order=pp_params.pymses.order)
|
||||
self._amr = self._ro.amr_source(["rho","vel","P"])
|
||||
|
||||
# Density operator
|
||||
self._rho_op = ScalarOperator(lambda dset: dset["rho"], self._ro.info["unit_density"])
|
||||
|
||||
# Density ray tracer
|
||||
if(pp_params.pymses.fft):
|
||||
self._rt = splatting.SplatterProcessor(self._amr, self._ro.info, self._rho_op)
|
||||
else:
|
||||
self._rt = raytracing.RayTracer(self._amr, self._ro.info, self._rho_op)
|
||||
|
||||
# Set the extend of the image
|
||||
self._radius = 0.5 / pp_params.out.zoom
|
||||
self._lbox = self._ro.info['boxlen']
|
||||
center = pp_params.out.center
|
||||
im_extent = [(- self._radius + center[0]) * self._lbox,
|
||||
( self._radius + center[0]) * self._lbox,
|
||||
(- self._radius + center[1]) * self._lbox,
|
||||
( self._radius + center[1]) * self._lbox]
|
||||
|
||||
# Get time
|
||||
time = self._ro.info['time'] # time in codeunits
|
||||
|
||||
# Set post processing attributes
|
||||
self.save.root._v_attrs.num = num
|
||||
self.save.root._v_attrs.lbox = self._lbox
|
||||
self.save.root._v_attrs.time = time
|
||||
|
||||
self.save.root.maps._v_attrs.center = center
|
||||
self.save.root.maps._v_attrs.radius = self._radius
|
||||
self.save.root.maps._v_attrs.im_extent = im_extent
|
||||
|
||||
|
||||
# Initialize cameras
|
||||
self._cam = dict()
|
||||
for ax_los in self._ax_nb : # los = line of sight
|
||||
ax_h = self._axes_h[ax_los]
|
||||
ax_v = self._axes_v[ax_los]
|
||||
|
||||
self._cam[ax_los] = Camera(center=pp_params.out.center,
|
||||
line_of_sight_axis=ax_los,
|
||||
region_size=[2.*self._radius, 2.*self._radius],
|
||||
distance=self._radius,
|
||||
far_cut_depth=self._radius,
|
||||
up_vector=ax_v,
|
||||
map_max_size=pp_params.out.map_size)
|
||||
self.save.close()
|
||||
self.def_rules()
|
||||
|
||||
def process(self, to_process_list, axes, overwrite=False):
|
||||
"""
|
||||
Render the data in to_process_list and save them
|
||||
"""
|
||||
self.save = tables.open_file(self.filename, mode="a")
|
||||
for name in to_process_list:
|
||||
if name in self.rules:
|
||||
rule = self.rules[name]
|
||||
for ax_los in axes:
|
||||
# Solve dependencies
|
||||
for dep in rule.dependencies:
|
||||
if dep in self.rules:
|
||||
rule_dep = self.rules[dep]
|
||||
self._process_rule(dep, rule_dep, ax_los, overwrite)
|
||||
else:
|
||||
print("ERROR: Dependency {} for {} is unknown".format(dep, name))
|
||||
# Process rule
|
||||
self._process_rule(name, rule, ax_los, overwrite)
|
||||
else:
|
||||
print("ERROR: {} is unknown".format(name))
|
||||
self.save.close()
|
||||
|
||||
def _process_rule(self, name, rule, ax_los, overwrite):
|
||||
name_full = rule.group + '/' + name + '_' + ax_los
|
||||
if rule.is_valid(self.save, ax_los):
|
||||
if overwrite or not name_full in self.save:
|
||||
data = rule.process(ax_los)
|
||||
self._save_data(name_full, data, rule.description)
|
||||
else:
|
||||
print("Data for {} is already computed, skipping...".format(name_full))
|
||||
else:
|
||||
print("ERROR: {} is not valid in this context".format(name_full))
|
||||
|
||||
|
||||
def _save_data(self, name_full, data, description):
|
||||
"""
|
||||
Save data in the HDF5 structure, overwrite if necessary
|
||||
"""
|
||||
if name_full in self.save:
|
||||
node = self.save.get_node(name_full)
|
||||
del node
|
||||
self.save.create_array(os.path.dirname(name_full), os.path.basename(name_full),
|
||||
data, description, createparents=True)
|
||||
|
||||
def _coldens(self, ax_los):
|
||||
datamap = self._rt.process(self._cam[ax_los], surf_qty=True)
|
||||
return datamap.map.T * self._lbox
|
||||
|
||||
def _rho(self, ax_los):
|
||||
datamap_rho = slicing.SliceMap(self._amr, self._cam[ax_los], self._rho_op, z=0.)
|
||||
return (datamap_rho.map).T
|
||||
|
||||
def _speed_h(self, ax_los):
|
||||
vh_op = ScalarOperator(lambda dset: dset["vel"][:, self._ax_nb[self._axes_h[ax_los]]],
|
||||
self._ro.info["unit_velocity"])
|
||||
dmap_vh = slicing.SliceMap(self._amr, self._cam[ax_los], vh_op, z=0.).map.T
|
||||
return dmap_vh
|
||||
|
||||
def _speed_v(self, ax_los):
|
||||
vv_op = ScalarOperator(lambda dset: dset["vel"][:, self._ax_nb[self._axes_v[ax_los]]],
|
||||
self._ro.info["unit_velocity"])
|
||||
dmap_vv = slicing.SliceMap(self._amr, self._cam[ax_los], vv_op, z=0.).map.T
|
||||
return dmap_vv
|
||||
|
||||
def _temperature(self, ax_los):
|
||||
P_op = ScalarOperator(lambda dset: dset["P"], self._ro.info["unit_pressure"])
|
||||
dmap_P = (slicing.SliceMap(self._amr, self._cam[ax_los], P_op, z=0.)).map.T
|
||||
dmap_rho = self.save.get_node("/maps/rho_{}".format(ax_los)).read()
|
||||
return dmap_P/dmap_rho
|
||||
|
||||
def _levels(self, ax_los):
|
||||
self._amr.set_read_levelmax(20)
|
||||
level_op = MaxLevelOperator()
|
||||
rt_level = raytracing.RayTracer(self._amr, self._ro.info, level_op)
|
||||
datamap = rt_level.process(self._cam[ax_los], surf_qty=True)
|
||||
return datamap.map.T
|
||||
|
||||
def _jeans(self, ax_los):
|
||||
dmap_T = self.save.get_node('/maps/T_' + ax_los).read()
|
||||
dmap_rho = self.save.get_node('/maps/rho_' + ax_los).read()
|
||||
dmap_jeans = np.sqrt(np.pi * dmap_T / dmap_rho)
|
||||
return dmap_jeans
|
||||
|
||||
def _jeans_ratio(self, ax_los):
|
||||
dmap_jeans = self.save.get_node('/maps/jeans_' + ax_los).read()
|
||||
dmap_levels = self.save.get_node('/maps/levels_' + ax_los).read()
|
||||
dmap_jeans_ratio = dmap_jeans * 2**(dmap_levels)
|
||||
return dmap_jeans_ratio
|
||||
|
||||
def _toomreQ_disk(self, ax_los):
|
||||
"""
|
||||
Compute the Toomre Q parameter in a Keplerian disk
|
||||
"""
|
||||
|
||||
# Operator to compute the angular speed times rho
|
||||
def omega_rho_func(dset):
|
||||
pos = dset.get_cell_centers()
|
||||
pos = pos - (self.pp_params.disk.pos_star / self._lbox)
|
||||
xx = pos[:, :, 0]
|
||||
yy = pos[:, :, 1]
|
||||
rc = np.sqrt(xx**2 + yy**2) # cylindrical radius
|
||||
vx = dset["vel"][:, :, 0]
|
||||
vy = dset["vel"][:, :, 1]
|
||||
omega_rho = (1. / rc**2)
|
||||
omega_rho = omega_rho * dset["rho"]
|
||||
vyx = vy * xx
|
||||
vxy = vx * yy
|
||||
omega_rho = omega_rho * (vyx - vxy)
|
||||
return omega_rho
|
||||
|
||||
# Operator to compute the angular speed
|
||||
omega_op = FractionOperator(omega_rho_func, lambda dset: dset["rho"],
|
||||
1. / self._ro.info["unit_time"])
|
||||
|
||||
# Operator to compute the sound speed
|
||||
cs_op = FractionOperator(lambda dset: dset["P"],
|
||||
lambda dset: dset["rho"], self._ro.info["unit_velocity"])
|
||||
|
||||
# Ray tracer for the angular speed
|
||||
rt_omega = raytracing.RayTracer(self._amr, self._ro.info, omega_op)
|
||||
|
||||
# Ray tracer for the sound speed
|
||||
if self.pp_params.pymses.fft:
|
||||
rt_cs = splatting.SplatterProcessor(self._amr, ro.info, cs_op, surf_qty=False)
|
||||
else :
|
||||
rt_cs = raytracing.RayTracer(self._amr, self._ro.info, cs_op)
|
||||
|
||||
dmap_omega = rt_omega.process(self._cam[ax_los])
|
||||
dmap_cs = rt_cs.process(self._cam[ax_los])
|
||||
dmap_col = self.save.root.maps.coldens_z.read()
|
||||
map_Q = (self._lbox * dmap_cs.map.T) * dmap_omega.map.T / (np.pi * self.G * dmap_col)
|
||||
|
||||
return map_Q
|
||||
|
||||
def def_rules(self):
|
||||
self.rules = {
|
||||
'coldens' : Rule(self._coldens, "Column density", '/maps'),
|
||||
'rho' : Rule(self._rho, "Density slice", '/maps'),
|
||||
'speed_h' : Rule(self._speed_h, "Horizontal speed slice wrt the line of sight", '/maps'),
|
||||
'speed_v' : Rule(self._speed_v, "Vertical speed slice wrt the line of sight", '/maps'),
|
||||
'T' : Rule(self._temperature, "Temperature slice", '/maps', dependencies=['rho']),
|
||||
'levels' : Rule(self._levels, "Max level within line of sight", '/maps'),
|
||||
'jeans' : Rule(self._jeans, "Jeans lenght slice", '/maps', dependencies=['rho', 'T']),
|
||||
'jeans_ratio' : Rule(self._jeans_ratio, "Jeans' lenght divided by the max resolution",
|
||||
'/maps', dependencies=['jeans', 'levels']),
|
||||
'Q' : Rule(self._toomreQ_disk, "Toomre Q parameter for a Keplerian disk", '/maps',
|
||||
dependencies=['coldens'], axes=['z'],
|
||||
is_valid=lambda save, axe: self.pp_params.disk.on)
|
||||
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user