Add more unit tests

Signed-off-by: Sylvain Leclerc <[email protected]>

src/antares/tsgen/ts_generator.py

@@ -66,14 +66,23 @@ def __init__(self, ts_count: int, days: int) -> None:
 
 def _column_powers(column: npt.NDArray[np.float_], width: int) -> npt.NDArray:
     """
-    Returns a matrix of given width where column[i] is the ith power of the input column.
+    Returns a matrix of given width where column[i] is the ith power of the input vector.
     """
     powers = np.arange(width)
     powers.shape = (1, len(powers))
     column.shape = (len(column), 1)
     return pow(column, powers)
 
 
+def _daily_to_hourly(daily_data: npt.NDArray) -> npt.NDArray:
+    """
+    Converts daily rows of a 2D array to hourly rows
+    """
+    if daily_data.ndim != 2:
+        raise ValueError("Daily data must be a 2D-array")
+    return np.repeat(daily_data, 24, axis=1)
+
+
 class ThermalDataGenerator:
     def __init__(self, rng: RNG = MersenneTwisterRNG(), days: int = 365) -> None:
         self.rng = rng
@@ -319,16 +328,16 @@ def generate_time_series(
 
                 # = storing output in output arrays =
                 if ts_index >= 0:  # drop the 2 first generated timeseries
-                    output.planned_outages[ts_index][day] = pure_planned_outages
-                    output.forced_outages[ts_index][day] = pure_forced_outages
-                    output.mixed_outages[ts_index][day] = mixed_outages
-                    output.planned_outage_durations[ts_index][day] = po_duration
-                    output.forced_outage_durations[ts_index][day] = fo_duration
-                    output.available_units[ts_index][day] = current_available_units
+                    output.planned_outages[ts_index, day] = pure_planned_outages
+                    output.forced_outages[ts_index, day] = pure_forced_outages
+                    output.mixed_outages[ts_index, day] = mixed_outages
+                    output.planned_outage_durations[ts_index, day] = po_duration
+                    output.forced_outage_durations[ts_index, day] = fo_duration
+                    output.available_units[ts_index, day] = current_available_units
 
                 now = (now + 1) % log_size
 
-        hourly_available_units = np.repeat(output.available_units, 24, axis=1)
+        hourly_available_units = _daily_to_hourly(output.available_units)
         hourly_modulation = np.tile(cluster.modulation, self.days)
         output.available_power = (
             hourly_available_units * cluster.nominal_power * hourly_modulation

tests/test_unit.py

@@ -12,14 +12,77 @@
 
 import numpy as np
 import numpy.testing as npt
+import pytest
 
 from antares.tsgen.ts_generator import (
     ProbabilityLaw,
     ThermalCluster,
     ThermalDataGenerator,
+    _column_powers,
+    _daily_to_hourly,
 )
 
 
+def test_daily_to_hourly():
+    daily = np.array([[1, 2]])
+    hourly = _daily_to_hourly(daily)
+    expected = [[1] * 24 + [2] * 24]
+    npt.assert_equal(hourly, expected)
+
+
+def test_elevate_to_power():
+    input = np.array([1, 0.5, 0.1])
+    powers = _column_powers(input, 3)
+    expected = np.array([[1, 1, 1], [1, 0.5, 0.25], [1, 0.1, 0.01]])
+    npt.assert_almost_equal(powers, expected, decimal=3)
+
+
+@pytest.fixture()
+def base_cluster_365_days():
+    days = 365
+    return ThermalCluster(
+        unit_count=10,
+        nominal_power=100,
+        modulation=np.ones(dtype=float, shape=24),
+        fo_law=ProbabilityLaw.UNIFORM,
+        fo_volatility=0,
+        po_law=ProbabilityLaw.UNIFORM,
+        po_volatility=0,
+        fo_duration=10 * np.ones(dtype=int, shape=days),
+        fo_rate=0.2 * np.ones(dtype=float, shape=days),
+        po_duration=10 * np.ones(dtype=int, shape=days),
+        po_rate=np.zeros(dtype=float, shape=days),
+        npo_min=np.zeros(dtype=int, shape=days),
+        npo_max=10 * np.ones(dtype=int, shape=days),
+    )
+
+
+def test_invalid_fo_rates(rng, base_cluster_365_days):
+    days = 365
+    cluster = base_cluster_365_days
+    cluster.fo_rate[12] = -0.2
+    cluster.fo_rate[10] = -0.1
+
+    with pytest.raises(
+        ValueError, match="forced failure rate is negative on days \[10, 12\]"
+    ):
+        generator = ThermalDataGenerator(rng=rng, days=days)
+        generator.generate_time_series(cluster, 1)
+
+
+def test_invalid_po_rates(rng, base_cluster_365_days):
+    days = 365
+    cluster = base_cluster_365_days
+    cluster.po_rate[12] = -0.2
+    cluster.po_rate[10] = -0.1
+
+    with pytest.raises(
+        ValueError, match="planned failure rate is negative on days \[10, 12\]"
+    ):
+        generator = ThermalDataGenerator(rng=rng, days=days)
+        generator.generate_time_series(cluster, 1)
+
+
 def test_forced_outages(rng):
     days = 365
     cluster = ThermalCluster(