diff --git a/src/smif/convert/unit.py b/src/smif/convert/unit.py index dfdf3306d..dd346598b 100644 --- a/src/smif/convert/unit.py +++ b/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 diff --git a/src/smif/data_layer/data_handle.py b/src/smif/data_layer/data_handle.py index f81df5610..f7382f901 100644 --- a/src/smif/data_layer/data_handle.py +++ b/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 diff --git a/tests/convert/test_unit.py b/tests/convert/test_unit.py index 552194be7..dc420592f 100644 --- a/tests/convert/test_unit.py +++ b/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) diff --git a/tests/data_layer/test_data_handle.py b/tests/data_layer/test_data_handle.py index 7061691b6..a16d605b9 100644 --- a/tests/data_layer/test_data_handle.py +++ b/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