Extracted make_quantiles from convert_to_output

src/pyciemss/utils/interface_utils.py

@@ -6,6 +6,7 @@
 from typing import Dict, Optional, Iterable, Callable
 
 
+
 def convert_to_output_format(
     samples: Dict[str, torch.Tensor],
     timepoints: Iterable[float],
@@ -97,49 +98,64 @@ def convert_to_output_format(
         result = result.assign(**{f"timepoint_{time_unit}": all_timepoints})
 
     if quantiles:
-        key_list = ["timepoint_id", "inc_cum", "output", "type", "quantile", "value"]
-        q = {k: [] for k in key_list}
-        if alpha_qs is None:
-            alpha_qs = np.linspace(0, 1, num_quantiles)
-            alpha_qs[0] = 0.01
-            alpha_qs[-1] = 0.99
-        else:
-            num_quantiles = len(alpha_qs)
-
-        # Solution (state variables)
-        for k, v in pyciemss_results["states"].items():
-            q_vals = np.quantile(v, alpha_qs, axis=0)
-            k = k.replace("_sol","")
-            if stacking_order == "timepoints":
-                # Keeping timepoints together
-                q["timepoint_id"].extend(list(np.repeat(np.array(range(num_timepoints)), num_quantiles)))
-                q["output"].extend([k]*num_timepoints*num_quantiles)
-                q["type"].extend(["quantile"]*num_timepoints*num_quantiles)
-                q["quantile"].extend(list(np.tile(alpha_qs, num_timepoints)))
-                q["value"].extend(list(np.squeeze(q_vals.T.reshape((num_timepoints * num_quantiles, 1)))))
-            elif stacking_order == "quantiles":
-                # Keeping quantiles together
-                q["timepoint_id"].extend(list(np.tile(np.array(range(num_timepoints)), num_quantiles)))
-                q["output"].extend([k]*num_timepoints*num_quantiles)
-                q["type"].extend(["quantile"]*num_timepoints*num_quantiles)
-                q["quantile"].extend(list(np.repeat(alpha_qs, num_timepoints)))
-                q["value"].extend(list(np.squeeze(q_vals.reshape((num_timepoints * num_quantiles, 1)))))
-            else:
-                raise Exception("Incorrect input for stacking_order.")
-        q["inc_cum"].extend(["inc"]*num_timepoints*num_quantiles*len(pyciemss_results["states"].items()))
-        result_q = pd.DataFrame(q)
-        if time_unit is not None:
-            all_timepoints = result_q["timepoint_id"].map(lambda v: timepoints[v])
-            result_q = result_q.assign(**{f"number_{time_unit}": all_timepoints})   
-            result_q = result_q[["timepoint_id", f"number_{time_unit}", "inc_cum", "output", "type", "quantile", "value"]]
-            if train_end_point is None:
-                result_q["Forecast_Backcast"] = "Forecast"
-            else:
-                result_q["Forecast_Backcast"] = np.where(result_q[f"number_{time_unit}"] > train_end_point, "Forecast", "Backcast")
+        result_q = make_quantiles(pyciemss_results, timepoints, alpha_qs, num_quantiles,
+                                  time_unit=time_unit, stacking_order=stacking_order, train_end_point=train_end_point)
         return result, result_q
     else:
         return result
 
+def make_quantiles(
+    pyciemss_results: dict[str,dict[str, torch.tensor]],
+    timepoints: Iterable[float],
+    alpha_qs: Optional[Iterable[float]] = [0.01, 0.025, 0.05, 0.1, 0.15, 0.2, 0.25, 0.3, 0.35, 0.4, 0.45, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95, 0.975, 0.99],
+    num_quantiles: Optional[int] = 0,
+    *,
+    time_unit: Optional[str] = "(unknown)",
+    stacking_order: Optional[str] = "timepoints",
+    train_end_point: Optional[float] = None,
+) -> pd.DataFrame:
+     """Make quantiles for each timepoint"""
+     num_samples, num_timepoints = next(iter(pyciemss_results["states"].values())).shape    
+     key_list = ["timepoint_id", "inc_cum", "output", "type", "quantile", "value"]
+     q = {k: [] for k in key_list}
+     if alpha_qs is None:
+         alpha_qs = np.linspace(0, 1, num_quantiles)
+         alpha_qs[0] = 0.01
+         alpha_qs[-1] = 0.99
+     else:
+         num_quantiles = len(alpha_qs)
+
+     # Solution (state variables)
+     for k, v in pyciemss_results["states"].items():
+         q_vals = np.quantile(v, alpha_qs, axis=0)
+         k = k.replace("_sol","")
+         if stacking_order == "timepoints":
+             # Keeping timepoints together
+             q["timepoint_id"].extend(list(np.repeat(np.array(range(num_timepoints)), num_quantiles)))
+             q["output"].extend([k]*num_timepoints*num_quantiles)
+             q["type"].extend(["quantile"]*num_timepoints*num_quantiles)
+             q["quantile"].extend(list(np.tile(alpha_qs, num_timepoints)))
+             q["value"].extend(list(np.squeeze(q_vals.T.reshape((num_timepoints * num_quantiles, 1)))))
+         elif stacking_order == "quantiles":
+             # Keeping quantiles together
+             q["timepoint_id"].extend(list(np.tile(np.array(range(num_timepoints)), num_quantiles)))
+             q["output"].extend([k]*num_timepoints*num_quantiles)
+             q["type"].extend(["quantile"]*num_timepoints*num_quantiles)
+             q["quantile"].extend(list(np.repeat(alpha_qs, num_timepoints)))
+             q["value"].extend(list(np.squeeze(q_vals.reshape((num_timepoints * num_quantiles, 1)))))
+         else:
+             raise Exception("Incorrect input for stacking_order.")
+     q["inc_cum"].extend(["inc"]*num_timepoints*num_quantiles*len(pyciemss_results["states"].items()))
+     result_q = pd.DataFrame(q)
+     if time_unit is not None:
+         all_timepoints = result_q["timepoint_id"].map(lambda v: timepoints[v])
+         result_q = result_q.assign(**{f"number_{time_unit}": all_timepoints})   
+         result_q = result_q[["timepoint_id", f"number_{time_unit}", "inc_cum", "output", "type", "quantile", "value"]]
+         if train_end_point is None:
+             result_q["Forecast_Backcast"] = "Forecast"
+         else:
+             result_q["Forecast_Backcast"] = np.where(result_q[f"number_{time_unit}"] > train_end_point, "Forecast", "Backcast")
+     return result_q
 
 def csv_to_list(filename):
     result = []

test/test_utils/expected_output_quantiles_format.csv

@@ -0,0 +1,61 @@
+timepoint_id,number_FancyUnit,inc_cum,output,type,quantile,value,Forecast_Backcast
+0,0.5,inc,Infected,quantile,0.01,0.9459380221366882,Forecast
+0,0.5,inc,Infected,quantile,0.05,0.946090018749237,Forecast
+0,0.5,inc,Infected,quantile,0.95,0.9489700138568878,Forecast
+0,0.5,inc,Infected,quantile,0.99,0.9490740168094635,Forecast
+1,1.0,inc,Infected,quantile,0.01,0.8946759748458862,Forecast
+1,1.0,inc,Infected,quantile,0.05,0.8949799776077271,Forecast
+1,1.0,inc,Infected,quantile,0.95,0.9005599915981293,Forecast
+1,1.0,inc,Infected,quantile,0.99,0.9007519900798797,Forecast
+2,2.0,inc,Infected,quantile,0.01,3.5019580268859865,Forecast
+2,2.0,inc,Infected,quantile,0.05,3.5033900260925295,Forecast
+2,2.0,inc,Infected,quantile,0.95,3.5321900844573975,Forecast
+2,2.0,inc,Infected,quantile,0.99,3.5333180904388426,Forecast
+3,3.0,inc,Infected,quantile,0.01,22.44344825744629,Forecast
+3,3.0,inc,Infected,quantile,0.05,22.45244026184082,Forecast
+3,3.0,inc,Infected,quantile,0.95,22.633249473571777,Forecast
+3,3.0,inc,Infected,quantile,0.99,22.640329399108886,Forecast
+4,4.0,inc,Infected,quantile,0.01,131.32871185302736,Forecast
+4,4.0,inc,Infected,quantile,0.05,131.37476043701173,Forecast
+4,4.0,inc,Infected,quantile,0.95,132.3012237548828,Forecast
+4,4.0,inc,Infected,quantile,0.99,132.3375274658203,Forecast
+0,0.5,inc,Recovered,quantile,0.01,0.06373199805617333,Forecast
+0,0.5,inc,Recovered,quantile,0.05,0.06385999843478203,Forecast
+0,0.5,inc,Recovered,quantile,0.95,0.06602000147104263,Forecast
+0,0.5,inc,Recovered,quantile,0.99,0.06608400136232376,Forecast
+1,1.0,inc,Recovered,quantile,0.01,0.1242620012164116,Forecast
+1,1.0,inc,Recovered,quantile,0.05,0.12451000064611435,Forecast
+1,1.0,inc,Recovered,quantile,0.95,0.1284700006246567,Forecast
+1,1.0,inc,Recovered,quantile,0.99,0.12857400119304657,Forecast
+2,2.0,inc,Recovered,quantile,0.01,0.5759500181674957,Forecast
+2,2.0,inc,Recovered,quantile,0.05,0.5761500179767609,Forecast
+2,2.0,inc,Recovered,quantile,0.95,0.5786700069904327,Forecast
+2,2.0,inc,Recovered,quantile,0.99,0.5786940062046051,Forecast
+3,3.0,inc,Recovered,quantile,0.01,1.6707160544395447,Forecast
+3,3.0,inc,Recovered,quantile,0.05,1.671180045604706,Forecast
+3,3.0,inc,Recovered,quantile,0.95,1.6778399705886842,Forecast
+3,3.0,inc,Recovered,quantile,0.99,1.6779679727554322,Forecast
+4,4.0,inc,Recovered,quantile,0.01,7.9489042186737064,Forecast
+4,4.0,inc,Recovered,quantile,0.05,7.951720190048218,Forecast
+4,4.0,inc,Recovered,quantile,0.95,8.005719709396363,Forecast
+4,4.0,inc,Recovered,quantile,0.99,8.007703695297241,Forecast
+0,0.5,inc,Susceptible,quantile,0.01,999.9868225097656,Forecast
+0,0.5,inc,Susceptible,quantile,0.05,999.9868469238281,Forecast
+0,0.5,inc,Susceptible,quantile,0.95,999.987890625,Forecast
+0,0.5,inc,Susceptible,quantile,0.99,999.987958984375,Forecast
+1,1.0,inc,Susceptible,quantile,0.01,999.9744897460937,Forecast
+1,1.0,inc,Susceptible,quantile,0.05,999.9744995117187,Forecast
+1,1.0,inc,Susceptible,quantile,0.95,999.9764770507812,Forecast
+1,1.0,inc,Susceptible,quantile,0.99,999.9766430664063,Forecast
+2,2.0,inc,Susceptible,quantile,0.01,996.8907592773437,Forecast
+2,2.0,inc,Susceptible,quantile,0.05,996.8917846679688,Forecast
+2,2.0,inc,Susceptible,quantile,0.95,996.919140625,Forecast
+2,2.0,inc,Susceptible,quantile,0.99,996.920546875,Forecast
+3,3.0,inc,Susceptible,quantile,0.01,976.6816833496093,Forecast
+3,3.0,inc,Susceptible,quantile,0.05,976.6889343261719,Forecast
+3,3.0,inc,Susceptible,quantile,0.95,976.8760864257813,Forecast
+3,3.0,inc,Susceptible,quantile,0.99,976.8854711914063,Forecast
+4,4.0,inc,Susceptible,quantile,0.01,860.65490234375,Forecast
+4,4.0,inc,Susceptible,quantile,0.05,860.693212890625,Forecast
+4,4.0,inc,Susceptible,quantile,0.95,861.673193359375,Forecast
+4,4.0,inc,Susceptible,quantile,0.99,861.7219921875,Forecast

test/test_utils/test_interface_utils.py

@@ -1,6 +1,7 @@
 import unittest
 from pyciemss.utils.interface_utils import (
     convert_to_output_format,
+    make_quantiles,
     interventions_and_sampled_params_to_interval,
     assign_interventions_to_timepoints,
     csv_to_list,
@@ -68,11 +69,14 @@ def setUp(self):
     def test_convert_to_output_format(self):
         """Test convert_to_output_format."""
         expected_output = pd.read_csv("test/test_utils/expected_output_format.csv")
-        result = convert_to_output_format(
+        expected_output_quantiles = pd.read_csv("test/test_utils/expected_output_quantiles_format.csv")
+        result, result_q = convert_to_output_format(
             self.intervened_samples,
             self.timepoints,
             self.interventions,
             time_unit="FancyUnit",
+            quantiles = True,
+            alpha_qs = [0.01, 0.05, 0.95, 0.99],
         )
 
         self.assertTrue(
@@ -85,6 +89,13 @@ def test_convert_to_output_format(self):
             atol=1e-5,
             rtol=1e-5,
         )
+        assert_frame_equal(
+            expected_output_quantiles,
+            result_q,
+            check_exact=False,
+            atol=1e-5,
+            rtol=1e-5,
+        )
 
         result = convert_to_output_format(
             self.intervened_samples, self.timepoints, self.interventions