From 877e34723f27c95f3b3e06e427f4204e3128c69f Mon Sep 17 00:00:00 2001
From: Patrick Ferris <patrick@sirref.org>
Date: Sat, 21 Oct 2023 21:39:29 +0100
Subject: [PATCH] Fix distributions

---
 methods/matching/find_pairs.py | 53 ++++++++++++++++++----------------
 1 file changed, 28 insertions(+), 25 deletions(-)

diff --git a/methods/matching/find_pairs.py b/methods/matching/find_pairs.py
index 74ce6e6..9348fe9 100644
--- a/methods/matching/find_pairs.py
+++ b/methods/matching/find_pairs.py
@@ -88,11 +88,31 @@ def find_match_iteration(
 
     s_subset_mask = make_s_subset_mask(m_dist_thresholded, k_dist_thresholded, m_dist_hard, k_dist_hard, required)
 
-    logging.info("Done make_s_subset_mask. s_subset_mask.shape: %s", s_subset_mask.shape)
+    s_subset_mask_true = np.zeros(m_set.shape[0], dtype=np.bool_)
+
+    # Randomly select 100 potential pixels per K, these pixels are similar
+    # to the particular k and so we should end up with similar distributions
+    # for the matching variables.
+    no_potentials = np.zeros(k_subset.shape[0], dtype=np.bool_)
+    for k in range(s_subset_mask.shape[0]):
+        masks = s_subset_mask[k]
+        indices = np.argwhere(masks)
+        np.random.shuffle(indices)
+        idx = indices[:100]
+        if len(idx) == 0:
+            no_potentials[k] = True
+        else:
+            for i in idx:
+                s_subset_mask_true[i[0]] = True
+
+    logging.info(f"Done make_s_subset_mask. s_subset_mask.shape: {s_subset_mask.shape}")
+    logging.info(f"Actual indexes {np.cumsum(s_subset_mask_true)}")
 
-    s_subset = m_set[s_subset_mask].reset_index()
+    s_subset = m_set[s_subset_mask_true]
+    potentials = np.invert(no_potentials)
 
-    logging.info("Finished preparing s_subset. shape: %s", s_subset.shape)
+    k_subset = k_subset[potentials]
+    logging.info(f"Finished preparing s_subset. shape: {s_subset.shape}")
 
     # Notes:
     # 1. Not all pixels may have matches
@@ -165,21 +185,10 @@ def find_match_iteration(
     logging.info("Finished find match iteration")
 
 @jit(nopython=True, fastmath=True, error_model="numpy")
-def make_s_subset_mask(
-    m_dist_thresholded: np.ndarray,
-    k_dist_thresholded: np.ndarray,
-    m_dist_hard: np.ndarray,
-    k_dist_hard: np.ndarray,
-    number_required: int,
-):
-    s_include = np.zeros((m_dist_thresholded.shape[0],), dtype=np.bool_)
-    # create an array that is the indexes of the rows in s_dist_thresholded and shuffle it
-    m_indexes = np.arange(m_dist_thresholded.shape[0])
-    np.random.shuffle(m_indexes)
-    found = 0
-
-    for position in range(m_dist_thresholded.shape[0]):
-        i = m_indexes[position]
+def make_s_subset_mask(m_dist_thresholded: np.ndarray, k_dist_thresholded: np.ndarray, m_dist_hard: np.ndarray, k_dist_hard: np.ndarray, n: int):
+    s_include = np.zeros((k_dist_thresholded.shape[0], m_dist_thresholded.shape[0]), dtype=np.bool_)
+
+    for i in range(m_dist_thresholded.shape[0]):
         m_row = m_dist_thresholded[i, :]
         m_hard = m_dist_hard[i]
 
@@ -203,13 +212,7 @@ def make_s_subset_mask(
                         should_include = False
 
             if should_include:
-                s_include[i] = True
-                break
-
-        if s_include[i]:
-            found += 1
-            if found >= number_required:
-                break
+                s_include[k][i] = True
 
     return s_include