From 5c7f4da25f388aeb0533220e486e518726a22689 Mon Sep 17 00:00:00 2001
From: MartinBelthle <martin.belthle@rte-france.com>
Date: Tue, 12 Nov 2024 14:34:42 +0100
Subject: [PATCH] fix(drawer): don't crash when rates are at 1 (#19)

---
 src/antares/tsgen/ts_generator.py | 14 ++++----------
 tests/test_ts_generator.py        | 14 ++++++++++++++
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/src/antares/tsgen/ts_generator.py b/src/antares/tsgen/ts_generator.py
index 438dff7..49df6da 100644
--- a/src/antares/tsgen/ts_generator.py
+++ b/src/antares/tsgen/ts_generator.py
@@ -236,14 +236,11 @@ def _categorize_outages(available_units: int, po_candidates: int, fo_candidates:
 
 class ForcedOutagesDrawer:
     def __init__(self, rng: RNG, unit_count: int, failure_rate: FloatArray):
-        days = len(failure_rate)
         self.rng = rng
         self.failure_rate = failure_rate
-        self.ff = np.zeros(days, dtype=float)  # ff = lf / (1 - lf)
-        a = np.zeros(shape=days, dtype=float)
         mask = failure_rate <= FAILURE_RATE_EQ_1
-        a[mask] = 1 - failure_rate[mask]
-        self.ff[mask] = failure_rate[mask] / a
+        a = np.where(mask, 1 - failure_rate, 0)
+        self.ff = np.where(mask, failure_rate / a, 0)  # ff = lf / (1 - lf)
         self.fpow = _column_powers(a, unit_count + 1)
 
     def draw(self, available_units: int, day: int) -> int:
@@ -269,14 +266,11 @@ def draw(self, available_units: int, day: int) -> int:
 
 class PlannedOutagesDrawer:
     def __init__(self, rng: RNG, unit_count: int, failure_rate: FloatArray):
-        days = len(failure_rate)
         self.rng = rng
         self.failure_rate = failure_rate
-        self.pp = np.zeros(days, dtype=float)  # pp = lp / (1 - lp)
-        a = np.zeros(shape=days, dtype=float)
         mask = failure_rate <= FAILURE_RATE_EQ_1
-        a[mask] = 1 - failure_rate[mask]
-        self.pp[mask] = failure_rate[mask] / a
+        a = np.where(mask, 1 - failure_rate, 0)
+        self.pp = np.where(mask, failure_rate / a, 0)  # pp = lp / (1 - lp)
         self.ppow = _column_powers(a, unit_count + 1)
 
     def draw(self, available_units: int, day: int, stock: int) -> Tuple[int, int]:
diff --git a/tests/test_ts_generator.py b/tests/test_ts_generator.py
index e388e4e..1a8baaa 100644
--- a/tests/test_ts_generator.py
+++ b/tests/test_ts_generator.py
@@ -12,6 +12,7 @@
 
 import csv
 
+import numpy as np
 import pytest
 
 from antares.tsgen.cluster_import import import_thermal_cluster
@@ -57,6 +58,19 @@ def test_one_unit_cluster(cluster_1, output_directory):
         writer.writerow(["PO rate :", round(true_por, 4), "FO rate :", round(true_for, 4)])
 
 
+def test_generation_with_fo_rate_at_1(cluster_1):
+    # Put FO_rate at 1 shouldn't raise an issue
+    cluster_1.outage_gen_params.fo_rate = np.ones(365)
+    TimeseriesGenerator().generate_time_series_for_clusters(cluster_1, 1)
+
+    # Reset
+    cluster_1.outage_gen_params.fo_rate = np.zeros(365)
+
+    # Put PO_rate at 1 shouldn't raise an issue
+    cluster_1.outage_gen_params.po_rate = np.ones(365)
+    TimeseriesGenerator().generate_time_series_for_clusters(cluster_1, 1)
+
+
 def test_hundred_unit_cluster(cluster_100, output_directory):
     ts_nb = 50