diff --git a/coarsen_weights.py b/coarsen_weights.py index b84a22df1384e582ad4ce30716433d21521eda86..d61995466996bec0af91e03eafcbdab37a253e26 100755 --- a/coarsen_weights.py +++ b/coarsen_weights.py @@ -17,9 +17,14 @@ def coarsen_weight_file(inweights, outweights, nbisections=1): w = weights[["src_address", "dst_address", "remap_matrix"]].to_dataframe() w.dst_address = ((w.dst_address - 1) // coarsen) + 1 - w.remap_matrix = w.remap_matrix / coarsen - new_w = w.groupby(["src_address", "dst_address"]).sum().reset_index() + # Normalize weights by actual number of links to new destination address + new_w = w.set_index("dst_address") + new_w["remap_matrix"] = new_w["remap_matrix"] / new_w["remap_matrix"].groupby("dst_address").count() + new_w = new_w.reset_index() + + # Combine equivalent links into one + new_w = new_w.groupby(["src_address", "dst_address"]).sum().reset_index() new_w.index.name = "num_links" new_w = new_w.to_xarray()