Refactoring, split into more files. Add more personalisation
This commit is contained in:
Noe Brucy
@@ -0,0 +1,376 @@
|
||||
f # coding: utf-8
|
||||
|
||||
import sys
|
||||
import os
|
||||
import glob as glob
|
||||
|
||||
|
||||
import tables
|
||||
import pymses
|
||||
import numpy as np
|
||||
from numpy.polynomial.polynomial import polyfit
|
||||
from scipy.stats import linregress
|
||||
|
||||
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
|
||||
import subprocess
|
||||
|
||||
import module_extract as me
|
||||
|
||||
from mypool import MyPool as Pool
|
||||
from functools import partial
|
||||
from abc import ABCMeta, abstractmethod
|
||||
import bunch
|
||||
|
||||
from run_selector import *
|
||||
from units import *
|
||||
|
||||
|
||||
class Rule:
|
||||
def __init__(
|
||||
self,
|
||||
postproc,
|
||||
process,
|
||||
description="",
|
||||
group="",
|
||||
dependencies=[],
|
||||
is_valid=lambda arg: True,
|
||||
kind="classic",
|
||||
unit=cst.none,
|
||||
):
|
||||
self.postproc = postproc
|
||||
self.process_fn = process
|
||||
self.dependencies = dependencies
|
||||
self.is_valid_add = is_valid
|
||||
self.group = group
|
||||
self.description = description
|
||||
self.unit = unit
|
||||
self.kind = kind
|
||||
|
||||
def process(self, arg, **kwargs):
|
||||
if not arg is None:
|
||||
return self.process_fn(arg, **kwargs)
|
||||
else:
|
||||
return self.process_fn(**kwargs)
|
||||
|
||||
def is_valid(self, arg):
|
||||
# save = self.postproc.save
|
||||
# valid = True
|
||||
# for dep in self.dependencies:
|
||||
# if dep in self.postproc.rules:
|
||||
# rule_dep = self.postproc.rules[dep]
|
||||
# if not arg is None:
|
||||
# valid = valid and rule_dep.group + '/' + dep + '_' + str(arg) in save
|
||||
# else:
|
||||
# valid = valid and rule_dep.group + '/' + dep in save
|
||||
# return valid and self.is_valid_add(arg)
|
||||
return self.is_valid_add(arg)
|
||||
|
||||
|
||||
class BaseProcessor:
|
||||
"""
|
||||
Base class for processors, should not be instanciated
|
||||
"""
|
||||
|
||||
__metaclass__ = ABCMeta
|
||||
|
||||
log_id = ""
|
||||
rules = {}
|
||||
solve_self_dep = True
|
||||
|
||||
def __init__(self, path, path_out=None, pp_params=None, tag=None):
|
||||
if pp_params is None:
|
||||
self.pp_params = default_params()
|
||||
elif type(pp_params) == str:
|
||||
self.pp_params = load_params(pp_params)
|
||||
else:
|
||||
self.pp_params = pp_params
|
||||
|
||||
if tag is not None:
|
||||
self.pp_params.out.tag = tag
|
||||
|
||||
# Determining output directory
|
||||
if path_out is None:
|
||||
self.path_out = path
|
||||
else:
|
||||
self.path_out = path_out
|
||||
|
||||
def _log(self, string, status=""):
|
||||
if self.pp_params.process.verbose:
|
||||
if len(status) > 0:
|
||||
print(status + ": " + self.log_id + string)
|
||||
else:
|
||||
print(self.log_id + string)
|
||||
|
||||
def process(
|
||||
self,
|
||||
to_process_list,
|
||||
args=[None],
|
||||
overwrite=False,
|
||||
overwrite_dep=False,
|
||||
**kwargs
|
||||
):
|
||||
"""
|
||||
Render the data in to_process_list and save them
|
||||
"""
|
||||
|
||||
if type(to_process_list) == str:
|
||||
to_process_list = [to_process_list]
|
||||
|
||||
if type(args) == str or args is None:
|
||||
args = [args]
|
||||
|
||||
self.overwrite_dep = overwrite_dep
|
||||
self.just_done = [] # Computations done within this call
|
||||
|
||||
for name in to_process_list:
|
||||
if name in self.rules:
|
||||
rule = self.rules[name]
|
||||
for arg in args:
|
||||
self._solve_and_process_rule(name, rule, arg, overwrite, **kwargs)
|
||||
else:
|
||||
self._log(
|
||||
"{} is unknown, allowed rules are {}".format(
|
||||
name, self.rules.keys()
|
||||
),
|
||||
"ERROR",
|
||||
)
|
||||
|
||||
return self.just_done
|
||||
|
||||
def _solve_and_process_rule(self, name, rule, arg, overwrite=False, **kwargs):
|
||||
updated = self._solve_dependencies(name, rule, arg, overwrite, **kwargs)
|
||||
overwrite_rule = overwrite or updated
|
||||
self._process_rule(name, rule, arg, overwrite_rule, **kwargs)
|
||||
|
||||
def _solve_dependencies(self, name, rule, arg, overwrite=False, **kwargs):
|
||||
|
||||
self.done_before_dep = len(self.just_done)
|
||||
|
||||
# Solve dependencies
|
||||
for dep in rule.dependencies:
|
||||
try:
|
||||
dep_arg = rule.dependencies[dep]
|
||||
except:
|
||||
dep_arg = arg
|
||||
|
||||
if dep_arg == "__parent__":
|
||||
dep_arg = arg
|
||||
|
||||
if self.solve_self_dep and dep in self.rules:
|
||||
rule_dep = self.rules[dep]
|
||||
self._solve_and_process_rule(dep, rule_dep, dep_arg, self.overwrite_dep)
|
||||
else:
|
||||
self._not_self_dep(name, dep, dep_arg, self.overwrite_dep, **kwargs)
|
||||
|
||||
# Whether dependencies where updated
|
||||
return len(self.just_done) > self.done_before_dep
|
||||
|
||||
def _not_self_dep(self, name, dep, dep_arg, overwrite, **kwargs):
|
||||
self._log("Dependency {} for {} is unknown".format(dep, name), "ERROR")
|
||||
|
||||
def _needs_computation(self, overwrite, name_full):
|
||||
return overwrite
|
||||
|
||||
def _process_rule(self, name, rule, arg, overwrite=False, **kwargs):
|
||||
if not arg is None:
|
||||
name_full = rule.group + "/" + name + "_" + str(arg)
|
||||
else:
|
||||
name_full = rule.group + "/" + name
|
||||
|
||||
if rule.is_valid(arg):
|
||||
if not name_full in self.just_done:
|
||||
if self._needs_computation(overwrite, name_full):
|
||||
self._log("Processing {}".format(name_full))
|
||||
data = rule.process(arg, **kwargs)
|
||||
self._save_data(name_full, data, rule.description, rule.unit)
|
||||
self._log("Data for {} computed".format(name_full), "SUCCESS")
|
||||
self.just_done.append(name_full)
|
||||
else:
|
||||
self._log(
|
||||
"Data for {} is already computed, skipping...".format(name_full)
|
||||
)
|
||||
else:
|
||||
self._log("{} is not valid in this context".format(name_full), "ERROR")
|
||||
|
||||
def def_rules(self):
|
||||
for rule in self.rules:
|
||||
setattr(self, rule, partial(self.process, rule))
|
||||
|
||||
|
||||
class HDF5Container(BaseProcessor):
|
||||
|
||||
filename = ""
|
||||
save = None
|
||||
opened = False
|
||||
|
||||
def open(self):
|
||||
if not self.opened:
|
||||
self.save = tables.open_file(self.filename, mode="a")
|
||||
self.opened = True
|
||||
|
||||
def close(self):
|
||||
if self.opened:
|
||||
self.save.close()
|
||||
self.opened = False
|
||||
|
||||
def _needs_computation(self, overwrite, name_full):
|
||||
return overwrite or not (name_full in self.save)
|
||||
|
||||
def _process_rule(self, name, rule, arg, overwrite, **kwargs):
|
||||
self.open()
|
||||
try:
|
||||
super(HDF5Container, self)._process_rule(
|
||||
name, rule, arg, overwrite, **kwargs
|
||||
)
|
||||
finally:
|
||||
self.close()
|
||||
|
||||
def get_value(self, node_name):
|
||||
self.open()
|
||||
try:
|
||||
node = self.save.get_node(node_name)
|
||||
if node._v_attrs.CLASS == "GROUP":
|
||||
value = {}
|
||||
for child_name in node._v_children:
|
||||
value[child_name] = self.get_value(node_name + "/" + child_name)
|
||||
else:
|
||||
value = node.read()
|
||||
finally:
|
||||
self.close()
|
||||
return value
|
||||
|
||||
def _save_data(self, name_full, data, description, unit):
|
||||
"""
|
||||
Save data in the HDF5 structure, overwrite if necessary
|
||||
"""
|
||||
if name_full in self.save:
|
||||
self.save.remove_node(name_full, recursive=True)
|
||||
|
||||
attrs = None
|
||||
if isinstance(data, tuple):
|
||||
attrs = data[1]
|
||||
data = data[0]
|
||||
|
||||
if isinstance(data, dict):
|
||||
if type(description) == str:
|
||||
self.save.create_group(
|
||||
os.path.dirname(name_full),
|
||||
os.path.basename(name_full),
|
||||
description,
|
||||
createparents=True,
|
||||
)
|
||||
else:
|
||||
self.save.create_group(
|
||||
os.path.dirname(name_full),
|
||||
os.path.basename(name_full),
|
||||
"",
|
||||
createparents=True,
|
||||
)
|
||||
|
||||
if not isinstance(unit, dict):
|
||||
self.save.get_node(name_full)._v_attrs.unit = unit
|
||||
|
||||
for key in data:
|
||||
if isinstance(description, dict):
|
||||
if isinstance(unit, dict):
|
||||
self._save_data(
|
||||
name_full + "/" + key,
|
||||
data[key],
|
||||
description[key],
|
||||
unit[key],
|
||||
)
|
||||
else:
|
||||
self._save_data(
|
||||
name_full + "/" + key, data[key], description[key], unit
|
||||
)
|
||||
else:
|
||||
if isinstance(unit, dict):
|
||||
self._save_data(name_full + "/" + key, data[key], "", unit[key])
|
||||
else:
|
||||
self._save_data(name_full + "/" + key, data[key], "", unit)
|
||||
else:
|
||||
self.save.create_array(
|
||||
os.path.dirname(name_full),
|
||||
os.path.basename(name_full),
|
||||
data,
|
||||
description,
|
||||
createparents=True,
|
||||
)
|
||||
self.save.get_node(name_full).attrs.unit = unit
|
||||
|
||||
if not attrs is None:
|
||||
self.save.get_node(name_full).attrs.update(attrs)
|
||||
|
||||
def set_value(self, node_name, data, description, unit):
|
||||
self.open()
|
||||
try:
|
||||
self._save_data(node_name, data, description, unit)
|
||||
finally:
|
||||
self.close()
|
||||
|
||||
def get_attribute(self, node_name, attr_name):
|
||||
self.open()
|
||||
try:
|
||||
node = self.save.get_node(node_name)
|
||||
attr = node._v_attrs[attr_name]
|
||||
finally:
|
||||
self.close()
|
||||
return attr
|
||||
|
||||
def _transform(self, name, transform_fn, group="/maps", **kwargs):
|
||||
src = self.save.get_node(group + "/" + name).read()
|
||||
return transform_fn(src, **kwargs)
|
||||
|
||||
def _gen_rule_transform(
|
||||
self,
|
||||
rule_src_name,
|
||||
transform_fn,
|
||||
transform_name,
|
||||
subarray_name=None,
|
||||
group=None,
|
||||
):
|
||||
|
||||
rule_src = self.rules[rule_src_name]
|
||||
|
||||
if subarray_name is None:
|
||||
src_name = rule_src_name
|
||||
group_src = rule_src.group
|
||||
unit = rule_src.unit
|
||||
description = rule_src.description
|
||||
else:
|
||||
src_name = subarray_name
|
||||
group_src = rule_src.group + "/" + rule_src_name
|
||||
unit = rule_src.unit[subarray_name]
|
||||
description = rule_src.description[subarray_name]
|
||||
|
||||
def fn(arg=None, **kwargs):
|
||||
if arg is None:
|
||||
return self._transform(
|
||||
src_name, transform_fn, group=group_src, **kwargs
|
||||
)
|
||||
else:
|
||||
return self._transform(
|
||||
src_name + "_" + str(arg), transform_fn, group=group_src, **kwargs
|
||||
)
|
||||
|
||||
if group is None:
|
||||
group = group_src
|
||||
|
||||
name = transform_name + "_" + rule_src_name
|
||||
|
||||
self.rules[name] = Rule(
|
||||
self,
|
||||
fn,
|
||||
group=group,
|
||||
unit=unit,
|
||||
description=description,
|
||||
dependencies=[rule_src_name],
|
||||
)
|
||||
|
||||
|
||||
def simple_getter(name, dset):
|
||||
return dset[name]
|
||||
Reference in New Issue
Block a user