minor adjustments to black formatting

sodym/export/sankey.py

@@ -14,17 +14,21 @@ class PlotlySankeyPlotter(CustomNameDisplayer, PydanticBaseModel):
     model_config = ConfigDict(arbitrary_types_allowed=True, extra="allow", protected_namespaces=())
 
     mfa: MFASystem
-    slice_dict: Optional[
-        dict
-    ] = {}  # for selection of a subset of the data; all other dimensions are summed over
+    """MFA system to visualize."""
+    slice_dict: Optional[dict] = {}
+    """for selection of a subset of the data; all other dimensions are summed over"""
     split_flows_by: Optional[str] = None
+    """dimension name to split and color flows by; if None, all flows are colored the same"""
     color_scheme: Optional[str] = "blueish"
+    """used if split_flows_by is not None and splitting dimension is in flow."""
     node_color: Optional[str] = "gray"
-    flow_color: Optional[
-        str
-    ] = "hsl(230,20,70)"  # used if split_flows_by is None or splitting dimension not in flow
+    """color of the nodes (processes and stocks)"""
+    flow_color: Optional[str] = "hsl(230,20,70)"
+    """used if split_flows_by is None or splitting dimension not in flow."""
     exclude_processes: Optional[list[str]] = ["sysenv"]
+    """processes that won't show up in the plot; neither will flows to and from them"""
     exclude_flows: Optional[list[str]] = []
+    """flows that won't show up in the plot"""
     __pydantic_extra__: dict[str, Any]
 
     @model_validator(mode="after")

sodym/stocks.py

@@ -135,9 +135,10 @@ def compute_outflow_by_cohort(self, stock_by_cohort) -> np.ndarray:
         """Compute outflow by cohort from changes in the stock by cohort and the known inflow."""
         outflow_by_cohort = np.zeros(self.shape_cohort)
         outflow_by_cohort[1:, :, ...] = -np.diff(stock_by_cohort, axis=0)
+        # allow for outflow in year 0 already
         outflow_by_cohort[self.t_diag_indices] = self.inflow.values - np.moveaxis(
             stock_by_cohort.diagonal(0, 0, 1), -1, 0
-        )  # allow for outflow in year 0 already
+        )
         return outflow_by_cohort
 
 
@@ -171,24 +172,25 @@ def compute_inflow_and_outflow(self, do_correct_negative_inflow=False) -> tuple[
         # construct the sf of a product of cohort tc remaining in the stock in year t
         # First year:
         inflow[0, ...] = np.where(sf[0, 0, ...] != 0.0, self.stock.values[0] / sf[0, 0], 0.0)
-        stock_by_cohort[:, 0, ...] = (
-            inflow[0, ...] * sf[:, 0, ...]
-        )  # Future decay of age-cohort of year 0.
+        # Future decay of age-cohort of year 0.
+        stock_by_cohort[:, 0, ...] = (inflow[0, ...] * sf[:, 0, ...])
         outflow_by_cohort[0, 0, ...] = inflow[0, ...] - stock_by_cohort[0, 0, ...]
         # all other years:
         for m in range(1, self.n_t):  # for all years m, starting in second year
             # 1) Compute outflow from previous age-cohorts up to m-1
+            # outflow table is filled row-wise, for each year m.
             outflow_by_cohort[m, 0:m, ...] = (
                 stock_by_cohort[m - 1, 0:m, ...] - stock_by_cohort[m, 0:m, ...]
-            )  # outflow table is filled row-wise, for each year m.
+                )
             # 2) Determine inflow from mass balance:
             if not do_correct_negative_inflow:  # if no correction for negative inflows is made
+                # allow for outflow during first year by rescaling with 1/sf[m,m]
                 inflow[m, ...] = np.where(
                     sf[m, m, ...] != 0.0,
                     (self.stock.values[m, ...] - stock_by_cohort[m, :, ...].sum(axis=0))
                     / sf[m, m, ...],
                     0.0,
-                )  # allow for outflow during first year by rescaling with 1/sf[m,m]
+                )
                 # 3) Add new inflow to stock and determine future decay of new age-cohort
                 stock_by_cohort[m::, m, ...] = inflow[m, ...] * sf[m::, m, ...]
                 outflow_by_cohort[m, m, ...] = inflow[m, ...] * (1 - sf[m, m, ...])
@@ -201,9 +203,8 @@ def compute_inflow_and_outflow(self, do_correct_negative_inflow=False) -> tuple[
                 inflow_test = self.stock.values[m, ...] - stock_by_cohort[m, :, ...].sum(axis=0)
                 if inflow_test < 0:  # if stock-driven model would yield negative inflow
                     delta = -1 * inflow_test  # Delta > 0!
-                    inflow[
-                        m, ...
-                    ] = 0  # Set inflow to 0 and distribute mass balance gap onto remaining cohorts:
+                    # Set inflow to 0 and distribute mass balance gap onto remaining cohorts:
+                    inflow[m, ...] = 0
                     delta_percent = np.where(
                         stock_by_cohort[m, :, ...].sum(axis=0) != 0,
                         delta / stock_by_cohort[m, :, ...].sum(axis=0),
@@ -227,11 +228,10 @@ def compute_inflow_and_outflow(self, do_correct_negative_inflow=False) -> tuple[
                     )
                 else:  # If no negative inflow would occur
                     inflow[m, ...] = np.where(
-                        sf[m, m, ...] != 0,  # Else, inflow is 0.
+                        sf[m, m, ...] != 0,
+                        # allow for outflow during first year by rescaling with 1/sf[m,m]
                         (self.stock.values[m, ...] - stock_by_cohort[m, :, ...].sum(axis=0))
-                        / sf[
-                            m, m, ...
-                        ],  # allow for outflow during first year by rescaling with 1/sf[m,m]
+                        / sf[m, m, ...],
                         0.0,
                     )
                     # Add new inflow to stock and determine future decay of new age-cohort

sodym/survival_functions.py

@@ -161,18 +161,16 @@ def survival_function_by_year_id(self, m, lifetime_mean, lifetime_std):
             lifetime_mean[m, ...]
             / np.sqrt(
                 1
-                + lifetime_mean[m, ...]
-                * lifetime_mean[m, ...]
-                / (lifetime_std[m, ...] * lifetime_std[m, ...])
+                + (lifetime_mean[m, ...] * lifetime_mean[m, ...]
+                / (lifetime_std[m, ...] * lifetime_std[m, ...]))
             )
         )
         # calculate parameter sigma of underlying normal distribution
         sg_ln = np.sqrt(
             np.log(
                 1
-                + lifetime_mean[m, ...]
-                * lifetime_mean[m, ...]
-                / (lifetime_std[m, ...] * lifetime_std[m, ...])
+                + (lifetime_mean[m, ...] * lifetime_mean[m, ...]
+                / (lifetime_std[m, ...] * lifetime_std[m, ...]))
             )
         )
         # compute survial function