black

examples/simple_usa_statelevel/model_input/my_initial_conditions.py

@@ -4,11 +4,10 @@
 
 class InitialConditions(gempyor.initial_conditions.InitialConditions):
     def get_from_config(self, sim_id: int, modinf) -> np.ndarray:
-
         y0 = np.zeros((modinf.compartments.compartments.shape[0], modinf.nsubpops))
         S_idx = modinf.compartments.get_comp_idx({"infection_stage": "S"})
         I_idx = modinf.compartments.get_comp_idx({"infection_stage": "I"})
-        prop_inf = .005 #np.random.uniform(low=0, high=0.01, size=modinf.nsubpops)
+        prop_inf = 0.005  # np.random.uniform(low=0, high=0.01, size=modinf.nsubpops)
         # TODO cause this is another example
         y0[S_idx, :] = modinf.subpop_pop * (1 - prop_inf)
         y0[I_idx, :] = modinf.subpop_pop * prop_inf

flepimop/gempyor_pkg/src/gempyor/calibrate.py

@@ -15,6 +15,7 @@
 # disable  operations using the MKL linear algebra.
 os.environ["OMP_NUM_THREADS"] = "1"
 
+
 @click.command()
 @click.option(
     "-c",

flepimop/gempyor_pkg/src/gempyor/inference.py

@@ -17,25 +17,28 @@
 # TODO: should be able to draw e.g from an initial condition folder buuut keep the draw as a blob
 # so it is saved by emcee, so I can build a posterio
 
-def emcee_logprob(proposal, modinf, inferpar, loss, static_sim_arguments, save=False, silent = True):
+
+def emcee_logprob(proposal, modinf, inferpar, loss, static_sim_arguments, save=False, silent=True):
     if not inferpar.check_in_bound(proposal=proposal):
         if not silent:
             print("OUT OF BOUND!!")
         return -np.inf
-
-    snpi_df_mod, hnpi_df_mod = inferpar.inject_proposal(proposal=proposal, snpi_df = static_sim_arguments["snpi_df_ref"], hnpi_df = static_sim_arguments["hnpi_df_ref"])
-
+
+    snpi_df_mod, hnpi_df_mod = inferpar.inject_proposal(
+        proposal=proposal, snpi_df=static_sim_arguments["snpi_df_ref"], hnpi_df=static_sim_arguments["hnpi_df_ref"]
+    )
+
     ss = copy.deepcopy(static_sim_arguments)
     ss["snpi_df_in"] = snpi_df_mod
     ss["hnpi_df_in"] = hnpi_df_mod
     del ss["snpi_df_ref"]
     del ss["hnpi_df_ref"]
 
-
     outcomes_df = simulation_atomic(**ss, modinf=modinf, save=save)
 
     ll_total, logloss, regularizations = loss.compute_logloss(model_df=outcomes_df, modinf=modinf)
-    if not silent: print(f"llik is {ll_total}")
+    if not silent:
+        print(f"llik is {ll_total}")
 
     return ll_total
 
@@ -55,7 +58,7 @@ def simulation_atomic(
     seeding_data,
     seeding_amounts,
     outcomes_parameters,
-    save=False
+    save=False,
 ):
     # We need to reseed because subprocess inherit of the same random generator state.
     np.random.seed(int.from_bytes(os.urandom(4), byteorder="little"))
@@ -176,7 +179,7 @@ def get_static_arguments(modinf):
         coords = dict(
             date=pd.date_range(modinf.ti, modinf.tf, freq="D"),
             **compartment_coords,
-            subpop=modinf.subpop_struct.subpop_names
+            subpop=modinf.subpop_struct.subpop_names,
         )
 
         zeros = np.zeros((len(coords["date"]), len(coords["mc_name"][1]), len(coords["subpop"])))

flepimop/gempyor_pkg/src/gempyor/inference_parameter.py

@@ -105,10 +105,11 @@ def print_summary(self):
 
     def __str__(self) -> str:
         from collections import Counter
+
         this_str = f"InferenceParameters: with {self.get_dim()} parameters: \n"
         for key, value in Counter(self.ptypes).items():
             this_str += f"    {key}: {value} parameters\n"
-        
+
         return this_str
 
     def get_dim(self):
@@ -160,7 +161,12 @@ def hit_ubs(self, proposal) -> np.ndarray:
         """
         return np.array((proposal > self.ubs))
 
-    def inject_proposal(self, proposal, snpi_df=None, hnpi_df=None,):
+    def inject_proposal(
+        self,
+        proposal,
+        snpi_df=None,
+        hnpi_df=None,
+    ):
         """
         Injects the proposal into model inputs, at the right place.
 

flepimop/gempyor_pkg/src/gempyor/logloss.py

@@ -24,8 +24,12 @@ def __init__(self, inference_config: confuse.ConfigView, modinf, data_dir: str =
         #  self.gt_xr = xr.Dataset.from_dataframe(self.gt.reset_index().set_index(["date","subpop"]))
         # then some NA were created if some dates where present in some gt but no other.
         # FIXME THIS IS FUNDAMENTALLY WRONG, especially as groundtruth resample by statistic !!!!
-
-        self.gt = pd.read_csv(os.path.join(data_dir, inference_config["gt_data_path"].get()), converters={"subpop": lambda x: str(x)}, skipinitialspace=True) # TODO: use read_df
+
+        self.gt = pd.read_csv(
+            os.path.join(data_dir, inference_config["gt_data_path"].get()),
+            converters={"subpop": lambda x: str(x)},
+            skipinitialspace=True,
+        )  # TODO: use read_df
         self.gt["date"] = pd.to_datetime(self.gt["date"])
         self.gt = self.gt.set_index("date")
 

flepimop/gempyor_pkg/src/gempyor/parameters.py

@@ -106,10 +106,10 @@ def __init__(
             else:
                 self.pdata[pn]["stacked_modifier_method"] = "product"
                 logging.debug(f"No 'stacked_modifier_method' for parameter {pn}, assuming multiplicative NPIs")
-            
+
             if self.pconfig[pn]["rolling_mean_windows"].exists():
                 self.pdata[pn]["rolling_mean_windows"] = self.pconfig[pn]["rolling_mean_windows"].get()
-            
+
             self.stacked_modifier_method[self.pdata[pn]["stacked_modifier_method"]].append(pn.lower())
 
         logging.debug(f"We have {self.npar} parameter: {self.pnames}")
@@ -196,14 +196,16 @@ def parameters_reduce(self, p_draw: ndarray, npi: object) -> ndarray:
         """
         p_reduced = copy.deepcopy(p_draw)
         if npi is not None:
-            for idx, pn in enumerate(self.pnames):  
+            for idx, pn in enumerate(self.pnames):
                 p_reduced[idx] = NPI.reduce_parameter(
                     parameter=p_draw[idx],
                     modification=npi.getReduction(pn.lower()),
                     method=self.pdata[pn]["stacked_modifier_method"],
                 )
                 # apply rolling mean if specified
                 if "rolling_mean_windows" in self.pdata[pn]:
-                    p_reduced[idx] = utils.rolling_mean_pad(data = p_reduced[idx], window=self.pdata[pn]["rolling_mean_windows"])
+                    p_reduced[idx] = utils.rolling_mean_pad(
+                        data=p_reduced[idx], window=self.pdata[pn]["rolling_mean_windows"]
+                    )
 
         return p_reduced

flepimop/gempyor_pkg/src/gempyor/statistics.py

@@ -111,8 +111,8 @@ def llik(self, model_data: xr.DataArray, gt_data: xr.DataArray):
                 x, loc=loc, scale=self.params.get("scale", scale)
             ),  # wrong:
             "norm_cov": lambda x, loc, scale: scipy.stats.norm.logpdf(
-                x, loc=loc, scale=scale*loc.where(loc > 5, 5)
-                ), # TODO: check, that it's really the loc
+                x, loc=loc, scale=scale * loc.where(loc > 5, 5)
+            ),  # TODO: check, that it's really the loc
             "nbinom": lambda x, n, p: scipy.stats.nbinom.logpmf(x, n=self.params.get("n"), p=model_data),
             "rmse": lambda x, y: -np.log(np.nansum(np.sqrt((x - y) ** 2))),
             "absolute_error": lambda x, y: -np.log(np.nansum(np.abs(x - y))),
@@ -122,11 +122,11 @@ def llik(self, model_data: xr.DataArray, gt_data: xr.DataArray):
         if self.dist in ["pois", "nbinom"]:
             model_data = model_data.astype(int)
             gt_data = gt_data.astype(int)
-        
+
         if self.zero_to_one:
             # so confusing, wish I had not used xarray to do model_data[model_data==0]=1
-            model_data=model_data.where(model_data != 0, 1)
-            gt_data=gt_data.where(gt_data != 0, 1)
+            model_data = model_data.where(model_data != 0, 1)
+            gt_data = gt_data.where(gt_data != 0, 1)
 
         # Use stored parameters in the distribution function call
         likelihood = dist_map[self.dist](gt_data, model_data, **self.params)

flepimop/gempyor_pkg/src/gempyor/steps_rk4.py

@@ -38,7 +38,7 @@ def rk4_integration(
     population,  # 15
     stochastic_p,  # 16
     method="rk4",
-    silent = False
+    silent=False
 ):
     states = np.zeros((ndays, ncompartments, nspatial_nodes))
     states_daily_incid = np.zeros((ndays, ncompartments, nspatial_nodes))

flepimop/gempyor_pkg/src/gempyor/subpopulation_structure.py

@@ -8,13 +8,13 @@
 
 logger = logging.getLogger(__name__)
 
-subpop_pop_key="population"
-subpop_names_key="subpop"
+subpop_pop_key = "population"
+subpop_names_key = "subpop"
 
 
 class SubpopulationStructure:
     def __init__(self, *, setup_name, subpop_config, path_prefix=pathlib.Path(".")):
-        """ Important attributes:
+        """Important attributes:
         - self.setup_name: Name of the setup
         - self.data: DataFrame with subpopulations and populations
         - self.nsubpops: Number of subpopulations
@@ -23,7 +23,7 @@ def __init__(self, *, setup_name, subpop_config, path_prefix=pathlib.Path(".")):
         - self.mobility: Mobility matrix
         """
 
-        geodata_file=path_prefix / subpop_config["geodata"].get()
+        geodata_file = path_prefix / subpop_config["geodata"].get()
 
         self.setup_name = setup_name
         self.data = pd.read_csv(
@@ -50,7 +50,7 @@ def __init__(self, *, setup_name, subpop_config, path_prefix=pathlib.Path(".")):
             raise ValueError(f"There are duplicate subpop_names in geodata.")
 
         if subpop_config["mobility"].exists():
-            mobility_file= path_prefix / subpop_config["mobility"].get()
+            mobility_file = path_prefix / subpop_config["mobility"].get()
             mobility_file = pathlib.Path(mobility_file)
             if mobility_file.suffix == ".txt":
                 print("Mobility files as matrices are not recommended. Please switch soon to long form csv files.")
@@ -130,8 +130,3 @@ def __init__(self, *, setup_name, subpop_config, path_prefix=pathlib.Path(".")):
             self.nsubpops = len(self.data)
             # TODO: this needs to be tested
             self.mobility = self.mobility[selected_subpop_indices][:, selected_subpop_indices]
-
-
-
-
-

flepimop/gempyor_pkg/src/gempyor/utils.py

@@ -273,20 +273,19 @@ def list_filenames(folder: str = ".", filters: list = []) -> list:
 
 
 def rolling_mean_pad(data, window):
-  """
-  Calculates rolling mean with centered window and pads the edges.
-
-  Args:
-      data: A NumPy array.
-      window: The window size for the rolling mean.
+    """
+    Calculates rolling mean with centered window and pads the edges.
 
-  Returns:
-      A NumPy array with the padded rolling mean.
-  """
-  padding_size = (window - 1) // 2
-  padded_data = np.pad(data, padding_size, mode='edge')
-  return np.convolve(padded_data, np.ones(window) / window, mode='valid')
+    Args:
+        data: A NumPy array.
+        window: The window size for the rolling mean.
 
+    Returns:
+        A NumPy array with the padded rolling mean.
+    """
+    padding_size = (window - 1) // 2
+    padded_data = np.pad(data, padding_size, mode="edge")
+    return np.convolve(padded_data, np.ones(window) / window, mode="valid")
 
 
 def print_disk_diagnosis():