test: add unit tests

Author: frodehk

src/libecalc/presentation/yaml/domain/expression_time_series_variable.py

@@ -25,8 +25,6 @@ def __init__(
         self._regularity = regularity
         self._is_rate = is_rate
 
-        self._condition = self._time_series_expression.get_condition_mask()
-
     @property
     def name(self) -> str:
         return self._name
@@ -36,7 +34,7 @@ def is_rate(self) -> bool:
         return self._is_rate
 
     def get_values(self) -> list[float]:
-        values: np.ndarray = np.asarray(self._time_series_expression.get_evaluated_expressions(), dtype=np.float64)
+        values: np.ndarray = np.asarray(self._time_series_expression.get_masked_values(), dtype=np.float64)
         # If some of these are rates, we need to calculate stream day rate for use
         # Also take a copy of the calendar day rate and stream day rate for input to result object
 
@@ -46,7 +44,6 @@ def get_values(self) -> list[float]:
                 regularity=self._regularity.values,
             )
 
-        values = self._condition.apply(values)
         return values.tolist()
 
     def get_periods(self) -> Periods:

src/libecalc/presentation/yaml/domain/time_series_expression.py

@@ -54,6 +54,15 @@ def get_evaluated_expressions(self) -> list[float]:
             arr = np.atleast_1d(arr[0])  # Flattens (1, N) to (N,)
         return arr.tolist()
 
+    def get_masked_values(self) -> list[float]:
+        """
+        Returns the evaluated expressions with the condition mask applied.
+        """
+        values = np.asarray(self.get_evaluated_expressions(), dtype=np.float64)
+        mask = self.get_condition_mask()
+        masked = mask.apply(values)
+        return masked.tolist()
+
     def get_condition_mask(self) -> TimeSeriesMask:
         mask = self.expression_evaluator.evaluate(expression=self._condition) if self._condition is not None else None
         return TimeSeriesMask.from_evaluated_mask(mask)

tests/libecalc/presentation/yaml/domain/test_time_series_mask.py

@@ -0,0 +1,61 @@
+from datetime import datetime
+
+import numpy as np
+
+from libecalc.common.time_utils import Period
+from libecalc.domain.regularity import Regularity
+from libecalc.presentation.yaml.domain.time_series_mask import TimeSeriesMask
+from libecalc.presentation.yaml.domain.time_series_expression import TimeSeriesExpression
+from libecalc.presentation.yaml.domain.expression_time_series_variable import ExpressionTimeSeriesVariable
+
+
+periods = [
+    Period(start=datetime(2020, 1, 1), end=datetime(2021, 1, 1)),
+    Period(start=datetime(2021, 1, 1), end=datetime(2022, 1, 1)),
+    Period(start=datetime(2022, 1, 1), end=datetime(2023, 1, 1)),
+]
+
+
+def test_time_series_mask_none():
+    """Test that TimeSeriesMask with None mask returns the input array unchanged."""
+    mask = TimeSeriesMask.from_evaluated_mask(None)
+    arr = np.array([1, 2, 3])
+    np.testing.assert_array_equal(mask.apply(arr), arr)
+
+
+def test_time_series_mask_int_mask():
+    """Test that TimeSeriesMask with integer mask applies masking correctly."""
+    mask = TimeSeriesMask.from_evaluated_mask(np.array([1, 0, 1]))
+    arr = np.array([10, 20, 30])
+    np.testing.assert_array_equal(mask.apply(arr), [10, 0, 30])
+
+
+def test_time_series_mask_float_mask():
+    """Test that TimeSeriesMask with float mask applies masking (nonzero as 1, zero as 0)."""
+    mask = TimeSeriesMask.from_evaluated_mask(np.array([0.5, 0.0, -2.1]))
+    arr = np.array([7, 8, 9])
+    np.testing.assert_array_equal(mask.apply(arr), [7, 0, 9])
+
+
+def test_time_series_mask_condition_mask(expression_evaluator_factory):
+    """Test that TimeSeriesExpression uses condition mask correctly."""
+    evaluator = expression_evaluator_factory.from_periods(
+        periods, variables={"SIM1;GAS_PROD": [10.0, 5.0, 10.0], "TEST_VAR": [1.0, 2.0, 3.0]}
+    )
+
+    expr = TimeSeriesExpression(expressions="TEST_VAR", expression_evaluator=evaluator, condition="SIM1;GAS_PROD > 5")
+    # The condition "SIM1;GAS_PROD > 5" results in a mask [1, 0, 1], which is applied to the values of TEST_VAR
+    assert expr.get_masked_values() == [1, 0, 3]
+
+
+def test_expression_time_series_variable_get_values_with_mask(expression_evaluator_factory):
+    """Test that ExpressionTimeSeriesVariable applies mask to values as expected."""
+    evaluator = expression_evaluator_factory.from_periods(
+        periods, variables={"SIM1;GAS_PROD": [10.0, 5.0, 10.0], "TEST_VAR": [1.0, 2.0, 3.0]}
+    )
+    regularity = Regularity(expression_evaluator=evaluator, target_period=evaluator.get_period(), expression_input=1)
+    expr = TimeSeriesExpression(expressions="TEST_VAR", expression_evaluator=evaluator, condition="SIM1;GAS_PROD > 5")
+
+    var = ExpressionTimeSeriesVariable(name="test", time_series_expression=expr, regularity=regularity, is_rate=False)
+    # The mask [1, 0, 1], created from condition, is applied to [1.0, 2.0, 3.0], resulting in [1.0, 0.0, 3.0]
+    assert var.get_values() == [1, 0, 3]