[Units] Set YamlWriter UnitSystem directly

Cantera · Aug 19, 2021 · 5751633 · 5751633
1 parent d10d1f4
commit 5751633
Show file tree

Hide file tree

Showing 3 changed files with 31 additions and 3 deletions.
diff --git a/interfaces/cython/cantera/test/test_composite.py b/interfaces/cython/cantera/test/test_composite.py
@@ -580,7 +580,7 @@ def test_yaml_simple(self):
                              gas2.forward_rate_constants)
         self.assertArrayNear(gas.mix_diff_coeffs, gas2.mix_diff_coeffs)
 
-    def test_yaml_outunits(self):
+    def test_yaml_outunits1(self):
         gas = ct.Solution('h2o2.yaml')
         gas.TPX = 500, ct.one_atm, 'H2: 1.0, O2: 1.0'
         gas.equilibrate('HP')
@@ -604,6 +604,30 @@ def test_yaml_outunits(self):
                              gas2.forward_rate_constants)
         self.assertArrayNear(gas.mix_diff_coeffs, gas2.mix_diff_coeffs)
 
+    def test_yaml_outunits2(self):
+        gas = ct.Solution('h2o2.yaml')
+        gas.TPX = 500, ct.one_atm, 'H2: 1.0, O2: 1.0'
+        gas.equilibrate('HP')
+        gas.TP = 1500, ct.one_atm
+        units = {'length': 'cm', 'quantity': 'mol', 'energy': 'cal'}
+        system = ct.UnitSystem(units)
+        gas.write_yaml('h2o2-generated.yaml', units=system)
+        generated = utilities.load_yaml("h2o2-generated.yaml")
+        original = utilities.load_yaml(self.cantera_data_path / "h2o2.yaml")
+
+        for r1, r2 in zip(original['reactions'], generated['reactions']):
+            if 'rate-constant' in r1:
+                self.assertNear(r1['rate-constant']['A'], r2['rate-constant']['A'])
+                self.assertNear(r1['rate-constant']['Ea'], r2['rate-constant']['Ea'])
+
+        gas2 = ct.Solution("h2o2-generated.yaml")
+        self.assertArrayNear(gas.concentrations, gas2.concentrations)
+        self.assertArrayNear(gas.partial_molar_enthalpies,
+                             gas2.partial_molar_enthalpies)
+        self.assertArrayNear(gas.forward_rate_constants,
+                             gas2.forward_rate_constants)
+        self.assertArrayNear(gas.mix_diff_coeffs, gas2.mix_diff_coeffs)
+
     def test_yaml_surface(self):
         gas = ct.Solution('ptcombust.yaml', 'gas')
         surf = ct.Interface('ptcombust.yaml', 'Pt_surf', [gas])

diff --git a/interfaces/cython/cantera/units.pyx b/interfaces/cython/cantera/units.pyx
@@ -46,6 +46,7 @@ cdef class Units:
         units.units = CxxUnits(other)
         return units
 
+
 cdef class UnitSystem:
     """
     Unit conversion utility

diff --git a/interfaces/cython/cantera/yamlwriter.pyx b/interfaces/cython/cantera/yamlwriter.pyx
@@ -50,11 +50,14 @@ cdef class YamlWriter:
         will use Cantera's defaults.
 
         :param units:
-            A map where keys are dimensions (mass, length, time, quantity,
-            pressure, energy, activation-energy), and the values are
+            A UnitSystem object or map where keys are dimensions (mass, length, time,
+            quantity, pressure, energy, activation-energy), and the values are
             corresponding units such as kg, mm, s, kmol, Pa, cal, and eV.
         """
         def __set__(self, units):
+            if isinstance(units, UnitSystem):
+                defaults = UnitSystem().defaults
+                units = {k: v for k, v in units.defaults.items() if defaults[k] != v}
             cdef stdmap[string, string] cxxunits
             for dimension, unit in units.items():
                 cxxunits[stringify(dimension)] = stringify(unit)