Merge pull request #327 from tomalrussell/fix/convert_custom_units

Fix/convert custom units

src/smif/convert/unit.py

@@ -1,7 +1,9 @@
 """Handles conversion between units used in the `SosModel`
 """
 from pint import DimensionalityError, UndefinedUnitError, UnitRegistry  # type: ignore
+
 from smif.convert.adaptor import Adaptor
+from smif.data_layer.data_handle import DataHandle
 
 
 class UnitAdaptor(Adaptor):
@@ -11,6 +13,13 @@ def __init__(self, name):
         self._register = UnitRegistry()
         super().__init__(name)
 
+    def before_model_run(self, data_handle: DataHandle):
+        """Register unit definitions in registry before model run
+        """
+        units = data_handle.read_unit_definitions()
+        for unit in units:
+            self._register.define(unit)
+
     def convert(self, data_array, to_spec, coefficients):
         data = data_array.data
         from_spec = data_array.spec

src/smif/data_layer/data_handle.py

@@ -9,7 +9,7 @@
 from copy import copy
 from logging import getLogger
 from types import MappingProxyType
-from typing import Dict
+from typing import Dict, List
 
 import numpy as np  # type: ignore
 from smif.data_layer.store import Store
@@ -567,6 +567,15 @@ def get_results(self, output_name, decision_iteration=None,
             decision_iteration
         )
 
+    def read_unit_definitions(self) -> List[str]:
+        """Read unit definitions
+
+        Returns
+        -------
+        list[str]
+        """
+        return self._store.read_unit_definitions()
+
     def read_coefficients(self, source_dim: str, destination_dim: str) -> np.ndarray:
         """Reads coefficients from the store
 

tests/convert/test_unit.py

@@ -36,3 +36,37 @@ def test_convert_unit():
     actual = data_handle.set_results.call_args[0][1]
     expected = np.array([1000], dtype=float)
     np.testing.assert_allclose(actual, expected)
+
+
+def test_convert_custom():
+    """Convert custom units
+    """
+    data_handle = Mock()
+    data = np.array([0.18346346], dtype=float)
+
+    from_spec = Spec(
+        name='test_variable',
+        dtype='float',
+        unit='mcm'
+    )
+
+    to_spec = Spec(
+        name='test_variable',
+        dtype='float',
+        unit='GW'
+    )
+
+    data_array = DataArray(from_spec, data)
+
+    data_handle.get_data = Mock(return_value=data_array)
+    data_handle.read_unit_definitions = Mock(return_value=['mcm = 10.901353 * GW'])
+
+    adaptor = UnitAdaptor('test-mcm-GW')
+    adaptor.add_input(from_spec)
+    adaptor.add_output(to_spec)
+    adaptor.before_model_run(data_handle)  # must have run before_model_run to register units
+    adaptor.simulate(data_handle)
+
+    actual = data_handle.set_results.call_args[0][1]
+    expected = np.array([2], dtype=float)
+    np.testing.assert_allclose(actual, expected)

tests/data_layer/test_data_handle.py

@@ -21,6 +21,11 @@ def mock_store(sample_dimensions, annual, get_sector_model, empty_store):
     """
     store = empty_store
 
+    store.write_unit_definitions([
+        'people',
+        'million people = 1000000 * people'
+    ])
+
     for dim in sample_dimensions:
         store.write_dimension(dim)
 
@@ -830,6 +835,17 @@ def test_read_coefficient_(self, mock_store, mock_model):
         actual = dh.read_coefficients('from_dim', 'to_dim')
         np.testing.assert_equal(actual, np.array([0, 0, 0, 0]))
 
+    def test_read_unit_definitions(self, mock_store, mock_model):
+        """Reading list of unit definition strings
+        """
+        store = mock_store
+        dh = DataHandle(store, 1, 2010, [2010], mock_model)
+        actual = dh.read_unit_definitions()
+        assert actual == [
+            'people',
+            'million people = 1000000 * people'
+        ]
+
 
 class TestResultsHandle:
     """Get results from any model