Linting

src/caf/distribute/furness.py

@@ -164,6 +164,11 @@ def doubly_constrained_furness(
 @dataclass
 class PropsInput:
     """
+    Input to triply constrained furness.
+
+    This is returned from cost_to_prop
+    Parameters
+    ----------
     props: np.ndarray
         This is essentially a cost matrix, but costs are replaced by the percentage
         from the given cost band.
@@ -204,12 +209,14 @@ def cost_to_prop(costs: np.ndarray, bands: pd.DataFrame, val_col: str):
     return band_indices, bands[val_col].values
 
 
+# pylint: disable=too-many-locals
 def triply_constrained_furness(
     props: list[PropsInput],
     row_targets,
     col_targets,
     max_iters,
     mat_size: tuple[int, int],
+    init_mat: np.ndarray = None,
     tol=1e-5,
 ):
     """
@@ -241,31 +248,17 @@ def triply_constrained_furness(
     cur_rmse = np.inf
     iter_num = 0
     n_vals = len(row_targets)
-    # build seed
-    furnessed_mat = np.zeros(mat_size)
-    for distro in props:
-        furnessed_mat[distro.zones] = distro.props
-
-    # one row, col furness loop outside iterations
-    row_ach = np.sum(furnessed_mat, axis=1)
-    diff_factor = np.divide(
-        row_targets, row_ach, where=row_ach != 0, out=np.ones_like(row_targets, dtype=float)
-    )
-
-    furnessed_mat = np.multiply(furnessed_mat.T, diff_factor).T
-    # Adjust cols
-    col_ach = np.sum(furnessed_mat, axis=0)
-    diff_factor = np.divide(
-        col_targets,
-        col_ach,
-        where=col_ach != 0,
-        out=np.ones_like(col_targets, dtype=float),
-    )
-    furnessed_mat *= diff_factor
-    # Can return early if all 0 - probably shouldn't happen!
-    if row_targets.sum() == 0 or col_targets.sum() == 0:
-        warnings.warn("Furness given targets of 0. Returning all 0's")
-        return np.zeros_like(props), iter_num, np.inf
+    if init_mat is None:
+        # build seed
+        furnessed_mat = np.zeros(mat_size)
+        for distro in props:
+            furnessed_mat[distro.zones] = distro.props
+
+        furnessed_mat, _, _ = doubly_constrained_furness(
+            furnessed_mat, row_targets, col_targets, max_iters=10, warning=False
+        )
+    else:
+        furnessed_mat = init_mat
     for iter_num in range(1, max_iters):
         # first adjust to match cost bands; this is the 'third' constraint but
         # is done first as the other two need to be matched more closely
@@ -275,15 +268,19 @@ def triply_constrained_furness(
             for i in distro.prop_vals:
                 tot_demand = to_alter[distro.props == i].sum()
                 checker[i] = tot_demand
-            df = pd.DataFrame.from_dict(checker, orient="index").reset_index()
-            df.columns = ["target_prop", "demand"]
-            df["act_prop"] = df["demand"] / df["demand"].sum()
-            df["adj"] = df["target_prop"] / df["act_prop"]
-            df.fillna(0, inplace=True)
-            df.set_index("target_prop", inplace=True)
-            for i in df.index:
-                to_alter[distro.props == i] *= df.loc[i, "adj"]
+            # pylint: disable=unsupported-assignment-operation, unsubscriptable-object
+            # pylint error
+            checker_df = pd.DataFrame.from_dict(checker, orient="index").reset_index()
+            checker_df.columns = ["target_prop", "demand"]
+            checker_df["act_prop"] = checker_df["demand"] / checker_df["demand"].sum()
+            checker_df["adj"] = checker_df["target_prop"] / checker_df["act_prop"]
+            checker_df.fillna(0, inplace=True)
+            checker_df.set_index("target_prop", inplace=True)
+
+            for i in checker_df.index:
+                to_alter[distro.props == i] *= checker_df.loc[i, "adj"]
             furnessed_mat[distro.zones] = to_alter
+            # pylint:enable=unsupported-assignment-operation, unsubscriptable-object
         # Adjust rows
         row_ach = np.sum(furnessed_mat, axis=1)
         diff_factor = np.divide(
@@ -310,7 +307,7 @@ def triply_constrained_furness(
         if cur_rmse < tol:
             early_exit = True
             break
-    if not early_exit:
+    if not early_exit & iter_num >= max_iters:
         warnings.warn(
             f"The triply constrained furness exhausted its max "
             f"number of loops ({max_iters:d}), while achieving an RMSE "
@@ -319,3 +316,6 @@ def triply_constrained_furness(
         )
 
     return furnessed_mat
+
+
+# pylint: enable=too-many-locals

src/caf/distribute/gravity_model/core.py

@@ -111,13 +111,14 @@ class GravityModelCalibrateResults(GravityModelResults):
 
     class OutputYaml(BaseConfig):
         """Class for outputting some data from this class."""
+
         cost_params: dict[str, Any]
         cost_function: str
         matrix_total: float
         cost_convergence: float
 
     def save(self, out_dir: Path):
-        """Save method for class"""
+        """Save method for class."""
         out_dir.mkdir(parents=False, exist_ok=True)
         achieved = self.cost_distribution.df.copy()
         achieved["achieved_normalised_demand"] = (

src/caf/distribute/gravity_model/multi_area.py

@@ -3,7 +3,6 @@
 # Built-Ins
 import functools
 import logging
-import os
 from dataclasses import dataclass
 from pathlib import Path
 from typing import Any, Optional
@@ -124,6 +123,8 @@ class MultiCostDistribution:
     function_params: dict[str, float]
 
 
+# pylint: disable=too-many-instance-attributes
+# 16 fine here
 class MultiAreaGravityModelCalibrator(core.GravityModelBase):
     """
     A self-calibrating multi-area gravity model.
@@ -164,6 +165,7 @@ def __init__(
         params: Optional[MultiDistInput],
     ):
         super().__init__(cost_function=cost_function, cost_matrix=cost_matrix)
+
         self.row_targets = row_targets
         self.col_targets = col_targets
         if len(row_targets) != cost_matrix.shape[0]:
@@ -196,6 +198,20 @@ def __init__(
             self.furness_tol = params.furness_tolerance
             self.furness_jac = params.furness_jac
 
+    @property
+    def achieved_tripends(self) -> pd.DataFrame:
+        """
+        Return achieved trip-ends.
+
+        Simply sums achieved distribution over each axis.
+        """
+        return pd.DataFrame(
+            {
+                "origins": self.achieved_distribution.sum(axis=1),
+                "destinations": self.achieved_distribution.sum(axis=0),
+            }
+        )
+
     def process_tlds(self):
         """Get distributions in the right format for a multi-area gravity model."""
         dists = []
@@ -348,7 +364,7 @@ def calibrate(
         *args,
         update_params: bool = False,
         **kwargs,
-    ) -> GravityModelCalibrateResults:
+    ) -> dict[str, GravityModelCalibrateResults]:
         """Find the optimal parameters for self.cost_function.
 
         Optimal parameters are found using `scipy.optimize.least_squares`
@@ -647,8 +663,9 @@ def run(self, triply_constrain: bool = False, xamax: int = 2):
                 self.row_targets,
                 self.col_targets,
                 5000,
-                self.cost_matrix.shape,
-                0.01,
+                init_mat=self.achieved_distribution,
+                mat_size=(self.cost_matrix.shape[0], self.cost_matrix.shape[1]),
+                tol=0.01,
             )
 
             assert self.achieved_cost_dist is not None
@@ -696,6 +713,9 @@ def run(self, triply_constrain: bool = False, xamax: int = 2):
         return results
 
 
+# pylint: enable=too-many-instance-attributes
+
+
 def gravity_model(
     row_targets: pd.Series,
     col_targets: np.ndarray,