Add support for fixing scenarios

scenarioReducer/fast_forward.py

@@ -34,10 +34,14 @@ def __init__(self, initialSet, initProbs):
         if round(np.sum(initProbs),2) != 1:
             raise ValueError('Probs must sum to one')
 
-    def reduce(self,distance,n_scenarios: int = 1):
+    def reduce(self, distance, n_scenarios: int = 1, fixed_idx: list[int] = []):
         """
         reduces the initial set of scenarios
         """
+
+        if len(fixed_idx) > n_scenarios:
+            raise ValueError('Number of scenarios to keep must be less than total number of scenarios')
+
         indxR = [] #indeces of the reduced set
         probs_initial = self.initProbs.copy() 
         #### computation of the distance matrix
@@ -53,9 +57,15 @@ def reduce(self,distance,n_scenarios: int = 1):
         ##first indx
         u = np.nanargmin(zeta)
         indxR.append(u)
+
+        if fixed_idx:
+            for idx in fixed_idx:
+                if idx not in indxR:
+                    indxR.append(idx)
+
         ####
         ##Step i
-        for it in range(n_scenarios-1): #we already did the first
+        for it in range(n_scenarios-len(indxR)): #we already fixed scenarios in indxR
             #update the distance matrix
             dist_mtrx = np.minimum(dist_mtrx, dist_mtrx[u, :])
             probs_initial[indxR] = 0 #set zero chosen elements

tests/test_fast_forward.py

@@ -78,6 +78,46 @@ def test_fast_forward_stats(scenario_generator):
     q = [0.05, 0.5, 0.95]
     np.testing.assert_allclose(d_original.quantile(q), d_reduced.quantile(q), rtol=1e-2, atol=1e-3)
 
+def test_fixed_scenarios():
+    outcomes_original = np.array([[215., 216., 217., 218., 219., 220., 221., 222., 223., 224., 225.,
+        226., 227., 228., 229., 230., 231., 232., 233., 234., 235., 236.,
+        237., 238., 239., 240., 241., 242., 243., 244., 245., 246., 247.,
+        248., 249., 250., 251., 252., 253., 254., 255., 256., 257., 258.,
+        259., 260., 261., 262., 263., 264., 265., 266., 267., 268., 269.,
+        270., 271., 272., 273., 274., 275., 276.]])
+
+    probabilities_original = np.array([0.00419851, 0.00476228, 0.00537685, 0.00604288, 0.0067604,
+           0.00752875, 0.00834648, 0.00921134, 0.01012024, 0.01106919,
+           0.0120533, 0.01306683, 0.01410319, 0.01515497, 0.01621409,
+           0.01727179, 0.01831885, 0.01934564, 0.02034229, 0.02129888,
+           0.02220556, 0.02305274, 0.02383126, 0.02453255, 0.02514882,
+           0.02567315, 0.02609966, 0.02642362, 0.02664151, 0.02675111,
+           0.02675152, 0.02664319, 0.02642786, 0.02610857, 0.02568955,
+           0.02517615, 0.02457471, 0.02389245, 0.02313731, 0.02231783,
+           0.02144295, 0.02052189, 0.01956397, 0.01857848, 0.01757451,
+           0.01656085, 0.01554586, 0.01453739, 0.01354264, 0.01256818,
+           0.01161981, 0.01070263, 0.00982091, 0.00897821, 0.00817731,
+           0.00742026, 0.00670846, 0.00604264, 0.00542297, 0.00484908,
+           0.00432015, 0.00383498])
+
+    FFreducer = Fast_forward(initialSet=outcomes_original, initProbs=probabilities_original)
+
+    test_cases_fixed_idx = [
+        [1,2,3],
+        [5],
+        [4,10],
+        [0, len(outcomes_original)-1],
+        [0,1,2,3,4,5,6,7,8,9,10]
+    ]
+
+    for fixed_idx in test_cases_fixed_idx:
+        outcomes_reduced, probabilities_reduced = FFreducer.reduce(
+            np.inf, 10, fixed_idx=fixed_idx
+        )
+        assert np.isclose(probabilities_reduced.sum(), 1.0)
+        for idx in fixed_idx:
+            assert outcomes_original[idx] in outcomes_reduced
+
 
 def print_comparative_stats(d1, d2):
     compare_list = {