Written multi-tld tests

Transport-for-the-North · Dec 11, 2023 · 38e76b4 · 38e76b4
1 parent e86b8c4
commit 38e76b4
Show file tree

Hide file tree

Showing 3 changed files with 130 additions and 60 deletions.
diff --git a/src/caf/distribute/utils.py b/src/caf/distribute/utils.py
@@ -47,8 +47,9 @@ def infill_cost_matrix(
     np.ndarray: The input matrix with values infilled.
     """
     # TODO add multiple factors by area-type
-    mins = cost_matrix.min(axis=1) * diag_factor
-    np.fill_diagonal(cost_matrix, mins)
+    min_row = np.min(np.ma.masked_where(cost_matrix <= 0, cost_matrix), axis=1) * diag_factor
+    np.fill_diagonal(cost_matrix, min_row)
+    cost_matrix[cost_matrix > 1e+10] = zeros_infill
     cost_matrix[cost_matrix == 0] = zeros_infill
     return cost_matrix
 
@@ -67,7 +68,7 @@ def process_tlds(
 ) -> list[multi_area.MultiCostDistribution]:
     """
     Read in a dataframe of distributions by category and a lookup, and return
-    a list of distributions reeady to be passed to a multi area gravity model.
+    a list of distributions ready to be passed to a multi area gravity model.
     """
     tlds = tlds.set_index(cat_col)
     tld_lookup = tld_lookup.sort_values(lookup_zone_col)