Merge pull request #12 from paulromano/unit-improvements

Several improvements for unit conversion

doc/source/conf.py

@@ -48,7 +48,8 @@
     'python': ('https://docs.python.org/3', None),
     'numpy': ('https://numpy.org/doc/stable/', None),
     'openmc': ('https://docs.openmc.org/en/stable/', None),
-    'h5py': ('https://docs.h5py.org/en/stable', None)
+    'h5py': ('https://docs.h5py.org/en/stable', None),
+    'astropy': ('https://docs.astropy.org/en/stable/', None)
 }
 
 # -- Options for HTML output -------------------------------------------------

doc/source/reference/index.rst

@@ -12,7 +12,9 @@ API Reference
    watts.Parameters
    watts.Plugin
    watts.PluginOpenMC
-   watts.PluginSAM
+   watts.PluginMOOSE
+   watts.PluginPyARC
    watts.Results
    watts.ResultsOpenMC
-   watts.ResultsSAM
+   watts.ResultsMOOSE
+   watts.ResultsPyARC

doc/source/user/usage.rst

@@ -48,31 +48,35 @@ workflows.
 Units
 ~~~~~
 
-To handle codes that use different unit systems, WATTS relies on `Astropy <https://www.astropy.org>`_ to perform unit conversion on parameters to ensure that the correct units are used for each code. For instance, MOOSE-based codes use the SI units while OpenMC uses the CGS units. With the built-in unit-conversion capability, a parameter needs only to be set once in any unit system and WATTS can automatically convert it to the correct unit for different codes. To use the unit-conversion capability, parameters need to be set to the ``astropy.units.quantity.Quantity`` class 
-as follows::
-
-    from astropy import units as u
-
-    Quantity = u.Quantity
+To handle codes that use different unit systems, WATTS relies on the
+:class:`~astropy.units.Quantity` class from :mod:`astropy.units` to perform unit
+conversion on parameters to ensure that the correct units are used for each
+code. For instance, MOOSE-based codes use the SI units while OpenMC uses the CGS
+units. With the built-in unit-conversion capability, a parameter needs only to
+be set once in any unit system and WATTS can automatically convert it to the
+correct unit for different codes. To use the unit-conversion capability,
+parameters need to be set using the :class:`~astropy.units.Quantity` class as
+follows::
 
-    params['control_pin_rad'] = Quantity(9.9, "mm")
-    params['He_inlet_temp'] = Quantity(600, "Celsius")
-    params['He_cp'] = Quantity(4.9184126, "BTU/(kg*K)")
+    from astropy.units import Quantity
 
-with the format of ``Quantity(<value>, <current unit>)``. Imperial units can also be enabled as 
-follows::
+    params['radius'] = Quantity(9.9, "mm")
+    params['inlet_temperature'] = Quantity(600, "Celsius")
+    params['c_p'] = Quantity(4.9184126, "BTU/(kg*K)")
 
-    u.imperial.enable()
+with the format of ``Quantity(value, unit)``.
 
 Plugins
 +++++++
 
-Using a particular code within WATTS requires a "plugin" that controls input file
-generation, execution, and post-processing. Three plugin classes,
-:class:`~watts.PluginMOOSE`, :class:`~watts.PluginOpenMC`, and :class:`~watts.PluginPyARC`, have already been added to WATTS and are available for your use.
+Using a particular code within WATTS requires a "plugin" that controls input
+file generation, execution, and post-processing. Three plugin classes,
+:class:`~watts.PluginMOOSE`, :class:`~watts.PluginOpenMC`, and
+:class:`~watts.PluginPyARC`, have already been added to WATTS and are available
+for your use.
 
 MOOSE Plugin
-~~~~~~~~~~
+~~~~~~~~~~~~
 
 The :class:`~watts.PluginMOOSE` class enables MOOSE simulations using a
 templated input file. This is demonstrated here for a SAM application, but other
@@ -96,7 +100,8 @@ follows:
         Tsolid_sf = 1e-3
     []
 
-If the templated input file is ``sam_template.inp``, the SAM code will rely the general MOOSE plugin that can be created as::
+If the templated input file is ``sam_template.inp``, the SAM code will rely on
+the general MOOSE plugin that can be created as::
 
     moose_plugin = watts.PluginMOOSE('sam_template.inp')
 
@@ -137,7 +142,7 @@ OpenMC Plugin
 ~~~~~~~~~~~~~
 
 The :class:`~watts.PluginOpenMC` class handles OpenMC execution in a similar
-manner to the :class:`~watts.PluginSAM` class for SAM. However, for OpenMC,
+manner to the :class:`~watts.PluginMOOSE` class for MOOSE. However, for OpenMC,
 inputs are generated programmatically through the OpenMC Python API. Instead of
 writing a text template, for the OpenMC plugin you need to write a function that
 accepts an instance of :class:`~watts.Parameters` and generates the necessary
@@ -183,7 +188,7 @@ PyARC Plugin
 ~~~~~~~~~~~~~
 
 The :class:`~watts.PluginPyARC` class handles PyARC execution in a similar
-manner to the :class:`~watts.PluginSAM` class for SAM. PyARC use text-based
+manner to the :class:`~watts.PluginMOOSE` class for MOOSE. PyARC use text-based
 input files which can be templated as follows:
 
 .. code-block:: jinja
@@ -194,7 +199,8 @@ input files which can be templated as follows:
         plane ( z10 ) { z = {{ assembly_length }} }
     }
 
-If the templated input file is `pyarc_template`, then the PyARC plugin can be instantiated with following command line::
+If the templated input file is `pyarc_template`, then the PyARC plugin can be
+instantiated with following command line::
 
     pyarc_plugin = watts.PluginPyARC('pyarc_template', show_stdout=True, extra_inputs=['lumped_test5.son'])
 
@@ -248,9 +254,9 @@ k-effective value at the end of the simulation:
     >>> results.keff
     1.0026170700986219+/-0.003342785895893627
 
-For SAM, the :class:`~watts.ResultsMOOSE` class
-provides a :attr:`~watts.ResultsMOOSE.csv_data` attribute that gathers the
-results from every CSV files generated by MOOSE applications (such as SAM or BISON)::
+For MOOSE, the :class:`~watts.ResultsMOOSE` class provides a
+:attr:`~watts.ResultsMOOSE.csv_data` attribute that gathers the results from
+every CSV files generated by MOOSE applications (such as SAM or BISON)::
 
     moose_result = moose_plugin.workflow(params)
     for key in moose_result.csv_data:

examples/example1a_SAM/example1a.py

@@ -4,11 +4,8 @@
 from math import cos, pi
 import os
 import watts
-from astropy import units as u
+from astropy.units import Quantity
 
-# Uses Astropy for unit conversion
-u.imperial.enable()    # Enable imperial units
-Quantity = u.Quantity
 
 params = watts.Parameters()
 
@@ -30,7 +27,7 @@
 params['ax_ref'] = 20 # cm
 params['num_cool_pins'] = 1*6+2*6+6*2/2
 params['num_fuel_pins'] = 6+6+6+3*6+2*6/2+6/3
-params['Height_FC'] = Quantity(2000, "mm") # Automatically converts to 'm' for MOOSE and 'cm' for openmc 
+params['Height_FC'] = Quantity(2000, "mm") # Automatically converts to 'm' for MOOSE and 'cm' for openmc
 params['Lattice_pitch'] = 2.0
 params['FuelPin_rad'] = 0.90 # cm
 params['cool_hole_rad'] = 0.60 # cm
@@ -56,4 +53,4 @@
 for key in moose_result.csv_data:
     print(key, moose_result.csv_data[key])
 print(moose_result.inputs)
-print(moose_result.outputs)
+print(moose_result.outputs)

examples/example1c_PyARC/example1c.py

@@ -1,23 +1,15 @@
 # SPDX-FileCopyrightText: 2022 UChicago Argonne, LLC
 # SPDX-License-Identifier: MIT
 
-import os
 import watts
-from astropy import units as u
-
-# Uses Astropy for unit conversion
-u.imperial.enable()    # Enable imperial units
-Quantity = u.Quantity
-params = watts.Parameters()
+from astropy.units import Quantity
 
 # TH params
-
+params = watts.Parameters()
 params['assembly_pitch'] = Quantity(20, "cm")  # 20e-2  m
 params['assembly_length'] = Quantity(13, "cm")  # 0.13 m
 params['temp'] = Quantity(26.85, "Celsius")  # 300 K
 
-
-
 params.show_summary(show_metadata=False, sort_by='key')
 
 # PyARC Workflow

examples/example2_SAM_OpenMC/example2.py

@@ -6,13 +6,9 @@
 import watts
 from statistics import mean
 from openmc_template import build_openmc_model
-from astropy import units as u
+from astropy.units import Quantity
 
 
-# Uses Astropy for unit conversion
-u.imperial.enable()    # Enable imperial units
-Quantity = u.Quantity
-
 params = watts.Parameters()
 
 # TH params
@@ -91,4 +87,4 @@
 for i, power_frac in enumerate(power_fractions):
     params[f'Init_P_{i+1}'] = power_frac
 
-params.show_summary(show_metadata=True, sort_by='time')
+params.show_summary(show_metadata=True, sort_by='time')

src/watts/__init__.py

@@ -8,3 +8,6 @@
 from .template import *
 from .parameters import *
 from .database import *
+
+# This allows a user to write watts.Quantity
+from astropy.units import Quantity

src/watts/parameters.py

@@ -1,29 +1,69 @@
 # SPDX-FileCopyrightText: 2022 UChicago Argonne, LLC
 # SPDX-License-Identifier: MIT
 
-import textwrap
+from __future__ import annotations
+import copy
 from collections import namedtuple
 from collections.abc import MutableMapping, Mapping, Iterable
 from datetime import datetime
 from getpass import getuser
+import textwrap
 from typing import Any, Union
 from warnings import warn
 
+import astropy.units as u
 import h5py
 from prettytable import PrettyTable
 
 
+# Enable imperial units
+u.imperial.enable()
+
+# Normally saving parameters to HDF5 is as simple as calling:
+#
+#   obj.create_dataset(key, data=value)
+#
+# However, in some cases we need to transform the value before writing or add
+# extra metadata in the dataset. To do this, we setup a mapping of Python types
+# to functions that create a dataset.
+
+def _generate_save_func(dtype):
+    def make_dataset(obj, key, value):
+        dataset = obj.create_dataset(key, data=dtype(value))
+        return dataset
+    return make_dataset
+
+_default_save_func = _generate_save_func(lambda x: x)
+
+def _quantity_save_func(obj, key, value):
+    dataset = _default_save_func(obj, key, value)
+    dataset.attrs['unit'] = str(value.unit)
+    return dataset
+
 _SAVE_FUNCS = {
-    'set': list
+    'set': _generate_save_func(list),
+    'Quantity': _quantity_save_func
 }
 
+# In an HDF5 file, all iterable objects just appear as plain arrays (represented
+# by h5py as numpy arrays). To "round trip" data correctly, we again setup a
+# mapping of Python types to functions that load data out of a datset and
+# perform any transformation needed.
+
+def _generate_load_func(dtype):
+    return lambda obj, value: dtype(value)
+
+def _quantity_load_func(obj, value):
+    return u.Quantity(value, obj.attrs['unit'])
+
 _LOAD_FUNCS = {
-    'tuple': tuple,
-    'list': list,
-    'set': set,
-    'float': float,
-    'int': int,
-    'bool': bool
+    'tuple': _generate_load_func(tuple),
+    'list': _generate_load_func(list),
+    'set': _generate_load_func(set),
+    'float': _generate_load_func(float),
+    'int': _generate_load_func(int),
+    'bool': _generate_load_func(bool),
+    'Quantity': _quantity_load_func
 }
 
 ParametersMetadata = namedtuple('ParametersMetadata', ['user', 'time'])
@@ -203,10 +243,8 @@ def add_metadata(obj, metadata):
             else:
                 # Convert type if necessary. If the type is not listed, return a
                 # "null" function that just returns the original value
-                func = _SAVE_FUNCS.get(type(value).__name__, lambda x: x)
-                file_value = func(value)
-
-                dset = h5_obj.create_dataset(key, data=file_value)
+                func = _SAVE_FUNCS.get(type(value).__name__, _default_save_func)
+                dset = func(h5_obj, key, value)
                 dset.attrs['type'] = type(value).__name__
                 if isinstance(mapping, type(self)):
                     add_metadata(dset, self._metadata[key])
@@ -247,8 +285,8 @@ def metadata_from_obj(obj):
 
                 # Convert type if indicated. If the type is not listed, return a
                 # "null" function that just returns the original value
-                func = _LOAD_FUNCS.get(obj.attrs['type'], lambda x: x)
-                mapping[key] = func(value)
+                func = _LOAD_FUNCS.get(obj.attrs['type'], lambda obj, x: x)
+                mapping[key] = func(obj, value)
 
                 if root:
                     self._metadata[key] = metadata_from_obj(obj)
@@ -277,7 +315,7 @@ def load(self, filename_or_obj: Union[str, h5py.Group]):
             self._load_mapping(self, filename_or_obj)
 
     @classmethod
-    def from_hdf5(cls, filename_or_obj: Union[str, h5py.Group]):
+    def from_hdf5(cls, filename_or_obj: Union[str, h5py.Group]) -> Parameters:
         """Return parameters from HDF5 file/group
 
         Parameters
@@ -287,4 +325,34 @@ def from_hdf5(cls, filename_or_obj: Union[str, h5py.Group]):
         """
         params = cls()
         params.load(filename_or_obj)
-        return params
+        return params
+
+    def convert_units(self, system: str = 'si', temperature: str = 'K',
+                     inplace: bool = False) -> Parameters:
+        """Perform unit conversion
+
+        Parameters
+        ----------
+        system
+            Desired unit system: 'si' or 'cgs'
+        temperature
+            Desired unit for temperature conversions
+        inplace
+            Whether to modify the parameters (True) or return a copy (False)
+
+        Returns
+        -------
+        A :class:`Parameters` instance with converted units
+        """
+        params = self if inplace else copy.deepcopy(self)
+
+        for key, value in params.items():
+            if isinstance(value, u.Quantity):
+                # Unit conversion for temperature needs to be done separately because
+                # astropy uses a different method to convert temperature.
+                if value.unit.physical_type == 'temperature':
+                    params[key] = value.to(temperature, equivalencies=u.temperature()).value
+                else:
+                    params[key] = getattr(value, system).value
+
+        return params