Add filaments postproc, improve units detection, add automatic map rules, add selection

baseprocessor.py

@@ -4,8 +4,10 @@ import sys
 import os
 import glob as glob
 import copy
-
+import time
 import tables
+from tables import HDF5ExtError
+
 import pymses
 import numpy as np
 from numpy.polynomial.polynomial import polyfit
@@ -211,7 +213,12 @@ class HDF5Container(BaseProcessor):
 
     def open(self):
         if not self.opened:
-            self.save = tables.open_file(self.filename, mode="a")
+            try:
+                self.save = tables.open_file(self.filename, mode="a")
+            except HDF5ExtError:
+                # Wait a bit if the lock was not still released
+                time.sleep(3)
+                self.save = tables.open_file(self.filename, mode="a")
             self.opened = True
 
     def close(self):
@@ -253,10 +260,52 @@ class HDF5Container(BaseProcessor):
             self.close()
         return value
 
+    def _get_units(self, unit, data=None):
+        """
+        Get real units from info files
+        unit is either:
+        1. An instance of cst.Unit (pymses unit class)
+        2. A string beginning by "unit_", referring to a code unit,
+           available in self.info
+        3. A dict {unit1 : exp1, unit2: exp2, ...} with unitX as 2.
+           and expX a float, referring to the compound unit
+           unit1**exp1 * unit2**exp2
+        4. A dict {key: unit, ...} where key is a field name (eg. 'time', or 'mass')
+           and unit the corresponding unit (on one on the above format)
+
+        Returns:
+        1-3. : a cst.Unit instance
+        4.   : a dict {key: unit, ...} with same key as input and unit being cst.Unit instances
+        """
+        if isinstance(unit, cst.Unit):
+            return unit
+        if isinstance(unit, str) and unit[:5] == "unit_":
+            res = self.info[unit]
+            if unit == "unit_length":
+                res = res / self.info["boxlen"]
+            return res
+        if list(unit)[0][:5] == "unit_":
+            new_unit = cst.none
+            for base_unit_str in unit:
+                expo = unit[base_unit_str]
+                base_unit = self._get_units(base_unit_str)
+                new_unit = new_unit * base_unit ** expo
+            return new_unit
+        if (not data is None) and isinstance(data, dict) and list(unit)[0] in data:
+            for key in unit:
+                unit[key] = self._get_units(unit[key])
+            return unit
+
+        else:
+            raise ValueError("Invalid unit")
+
     def _save_data(self, name_full, data, description, unit):
         """
         Save data in the HDF5 structure, overwrite if necessary
         """
+
+        unit = self._get_units(unit, data=data)
+
         if name_full in self.save:
             self.save.remove_node(name_full, recursive=True)
 
@@ -285,6 +334,7 @@ class HDF5Container(BaseProcessor):
                 self.save.get_node(name_full)._v_attrs.unit = unit
 
             for key in data:
+                key = str(key)
                 if isinstance(description, dict):
                     if isinstance(unit, dict):
                         self._save_data(
@@ -314,6 +364,7 @@ class HDF5Container(BaseProcessor):
 
         if not attrs is None:
             for key in attrs:
+                key = str(key)
                 self.save.get_node(name_full)._v_attrs[key] = attrs[key]
 
     def set_value(self, node_name, data, description, unit):
@@ -412,3 +463,11 @@ class HDF5Container(BaseProcessor):
 
 def simple_getter(name, dset):
     return dset[name]
+
+
+def vect_getter(name, i, dset):
+    return dset[name][:, i]
+
+
+def norm_getter(name, dset):
+    return np.sqrt(np.sum(dset[name] ** 2, axis=1))