recycle unused memory to prevent allocation #168

jacanchaplais · Nov 23, 2023 · 0ef34c4 · 0ef34c4
1 parent 8ca9b1e
commit 0ef34c4
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/graphicle/calculate.py b/graphicle/calculate.py
@@ -666,7 +666,7 @@ def _assignment_cost(
 ) -> base.DoubleVector:
     dist_matrix = _delta_R(rapidity_1, rapidity_2, xy_pol_1, xy_pol_2)
     num_partons = dist_matrix.shape[0]
-    pt_2_cache = np.abs(xy_pol_2)
+    pt_2_cache = np.abs(xy_pol_2, out=rapidity_2)
     var_pt_recip = 1.0 / np.var(pt_2_cache)
     for parton_idx in nb.prange(num_partons):
         row = dist_matrix[parton_idx, :]

diff --git a/graphicle/select.py b/graphicle/select.py
@@ -1228,6 +1228,7 @@ def arg_closest(
             focus._xy_pol,
             candidate._xy_pol,
         )
+    del candidate.rapidity  # inplace operation above invalidates cache
     _, idxs = opt.linear_sum_assignment(cost_matrix)
     return idxs.tolist()