support for simtk units (needs tests)

dwhswenson · Dec 23, 2020 · Dec 1, 2020 · Dec 3, 2020 · Dec 3, 2020 · Dec 3, 2020
commit 1a2ae80626904f212d0fb736623989a605e185dd
diff --git a/openpathsampling/experimental/simstore/class_lookup.py b/openpathsampling/experimental/simstore/class_lookup.py
@@ -14,6 +14,14 @@ def __init__(self, check_method):
         self._true_set = set()
         self._false_set = set()
 
+    def force_true(self, cls):
+        self._false_set.discard(cls)
+        self._true_set.add(cls)
+
+    def force_false(self, cls):
+        self._true_set.discard(cls)
+        self._false_set.add(cls)
+
     def __call__(self, obj):
         key = obj.__class__
         if key in self._false_set:

diff --git a/openpathsampling/experimental/storage/ops_storage.py b/openpathsampling/experimental/storage/ops_storage.py
@@ -1,6 +1,8 @@
 from ..simstore import storage
 from ..simstore import sql_backend
 
+from ..simstore.attribute_handlers import DEFAULT_HANDLERS
+
 from ..simstore.serialization_helpers import to_json_obj as json_serializer
 from ..simstore.serialization_helpers import from_json_obj as deserialize_sim
 from ..simstore.serialization_helpers import import_class
@@ -36,6 +38,8 @@
 from . import snapshots
 from .snapshots_table import SnapshotsTable
 
+from .simtk_unit import simtk_quantity_codec, SimtkQuantityHandler
+
 import logging
 logger = logging.getLogger(__name__)
 
@@ -65,7 +69,9 @@
 ops_schema_sql_metadata = {}
 
 # this defines the simulation object serializer for OPS
-CODECS = [numpy_codec, bytes_codec, uuid_object_codec]
+CODECS = [numpy_codec, bytes_codec, uuid_object_codec, simtk_quantity_codec]
+
+HANDLERS = DEFAULT_HANDLERS + [SimtkQuantityHandler]
 
 UNSAFE_CODECS = CODECS + [CallableCodec()]
 SAFE_CODECS = CODECS + [CallableCodec({'safemode': True})]
@@ -74,11 +80,12 @@ class MoveChangeDeserializer(SchemaDeserializer):
     # in general, I think it would be better to reorg MoveChange to only be
     # one class, but this is aimed at fixing problems with reloading
     # MoveChange objects
-    def __init__(self, schema, table):
+    def __init__(self, schema, table, handlers):
         super(MoveChangeDeserializer, self).__init__(
             schema=schema,
             table=table,
-            cls=None
+            cls=None,
+            handlers=[]
         )
 
     def __call__(self, uuid, table_dct, cache_list):
@@ -202,7 +209,8 @@ def _build_ops_serializer(schema, safe_codecs, unsafe_codecs):
             ClassInfo(table='move_changes', cls=paths.MoveChange,
                       deserializer=MoveChangeDeserializer(
                           schema=schema,
-                          table='move_changes'
+                          table='move_changes',
+                          handlers=HANDLERS
                       )),
             ClassInfo(table='steps', cls=paths.MCStep),
             ClassInfo(table='details', cls=paths.Details,
@@ -222,7 +230,7 @@ def _build_ops_serializer(schema, safe_codecs, unsafe_codecs):
     )
 
     for info in ops_class_info.class_info_list:
-        info.set_defaults(schema)
+        info.set_defaults(schema, HANDLERS)
 
     return ops_class_info
 

diff --git a/openpathsampling/experimental/storage/simtk_unit.py b/openpathsampling/experimental/storage/simtk_unit.py
@@ -0,0 +1,95 @@
+from openpathsampling.experimental.simstore.custom_json import JSONCodec
+from openpathsampling.experimental.simstore.attribute_handlers import (
+    AttributeHandler, NDArrayHandler
+)
+import re
+
+try:
+    import simtk.unit
+except ImportError:
+    HAS_SIMTK = False
+else:
+    HAS_SIMTK = True
+    # note that we need to force simstore to not treat this as an iterable
+    from openpathsampling.experimental.simstore.class_lookup import \
+        is_storage_iterable
+    is_storage_iterable.force_false(simtk.unit.Quantity)
+
+### JSON SERIALIZATION ###################################################
+
+def unit_to_dict(obj):
+    return {p.name: int(power)
+            for p, power in obj.iter_base_or_scaled_units()}
+
+def unit_from_dict(dct):
+    unit = simtk.unit.Unit({})
+    for u_name, u_power in dct.items():
+        unit *= getattr(simtk.unit, u_name) ** u_power
+    return unit
+
+def quantity_to_dict(obj):
+    return {
+            'value': obj.value,
+            '__simtk_unit__': unit_to_dict(obj.unit)}
+
+def quantity_from_dict(dct):
+    unit = unit_from_dict(dct['__simtk_unit__'])
+    return dct['value'] * unit
+
+
+if HAS_SIMTK:
+    simtk_quantity_codec = JSONCodec(
+        cls=simtk.unit.Quantity,
+        to_dict=quantity_to_dict,
+        from_dict=quantity_to_dict,
+        is_my_dict=lambda x: '__simtk_unit__' in x
+    )
+
+### DIRECT SERIALIZATION #################################################
+_simtk_re = re.compile("simtk\((.*)\)\*((ndarray|float).*)")
+
+
+def simtk_unit_from_string(unit_str):
+    # TODO: add safety checks; parse the AST and ensure that all attributes
+    # are of `unit` and that all operations are mul/div/pow
+    from simtk import unit
+    return eval(unit_str, {'unit': unit})
+
+
+class SimtkQuantityHandler(AttributeHandler):
+    # NOTE: only supports ndarrays and floats for now
+    def __init__(self, type_info):
+        super().__init__(type_info)
+        self.unit_str, self.wrapped_type = type_info
+        self.unit = simtk_unit_from_string(self.unit_str)
+        if self.wrapped_type == 'float':
+            self.inner_serialize = lambda x: x
+            self.inner_deserialize = lambda x, _: x
+            self.backend_type = 'float'
+            self.type_size = None
+        else:
+            np_handler = NDArrayHandler.from_type_string(self.wrapped_type)
+            self.inner_serialize = np_handler.serialize
+            self.inner_deserialize = np_handler.deserialize
+            self.backend_type = np_handler.backend_type
+            self.type_size = np_handler.type_size
+
+    @staticmethod
+    def is_my_type(type_str):
+        m_simtk_quantity = _simtk_re.match(type_str)
+        if m_simtk_quantity:
+            unit = m_simtk_quantity.group(1)
+            wrapped_type = m_simtk_quantity.group(2)
+            return unit, wrapped_type
+
+    def serialize(self, obj):
+        unwrapped = obj.value_in_unit(self.unit)
+        return self.inner_serialize(unwrapped)
+
+    def deserialize(self, data, caches=None):
+        unwrapped = self.inner_deserialize(data, caches)
+        return self.unit * unwrapped
+
+
+
+
diff --git a/openpathsampling/experimental/storage/snapshots.py b/openpathsampling/experimental/storage/snapshots.py
@@ -96,7 +96,7 @@ def snapshot_registration_info(snapshot_instance, snapshot_number):
     attr_infos = []
     for table in [tbl for tbl in schema.keys() if tbl != 'snapshot']:
         obj = getattr(snapshot_instance, table)
-        attr_infos.append(ClassInfo(table=real_table,
+        attr_infos.append(ClassInfo(table=real_table[table],
                                     cls=obj.__class__,
                                     lookup_result=(engine_uuid,
                                                    obj.__class__)))