From ba5f6a62d50787bb0eaa2360fb713f0b5d5a30c3 Mon Sep 17 00:00:00 2001 From: IsabelThompson97 Date: Sun, 5 Jul 2026 01:11:33 -0400 Subject: [PATCH] Add opt-in build_knn flag to skip the unused kNN graph rite_weight_model always builds a NearestNeighbors graph in __init__ (fit + full N-point kneighbors query, held for the object's lifetime). On the default smoothing_factor_nn = 0 path that graph is built and never read, which is the dominant construction cost at large N. Add build_knn (default True, so existing behavior is unchanged). Pass build_knn=False to skip the graph when kNN smoothing is off; smooth_weights_knn then raises a clear error if smoothing is later requested, so the build-time and run-time switches cannot silently disagree. --- rite_weight/rite_weight.py | 39 ++++++++++++++++++++++++++++---------- 1 file changed, 29 insertions(+), 10 deletions(-) diff --git a/rite_weight/rite_weight.py b/rite_weight/rite_weight.py index 8d1b485..b4fcf52 100644 --- a/rite_weight/rite_weight.py +++ b/rite_weight/rite_weight.py @@ -6,12 +6,12 @@ from rite_weight.reweighting_utils import reweights_calculate class rite_weight_model: - def __init__(self, config, knn_k=10): + def __init__(self, config, knn_k=10, build_knn=True): # Load input arrays and other parameters self.trans_from_feat = np.load(config['trans_from_Rfeat_path']) self.trans_to_feat = np.load(config['trans_to_Rfeat_path']) self.thero_weight_ref = 0.0 #Change to very small value bcz of which NAN could be there in Tmatrix - + # Round to 3 decimals trans_from_feat_rounded = np.round(self.trans_from_feat, decimals=3) # Extract unique rows and their indices @@ -19,16 +19,25 @@ def __init__(self, config, knn_k=10): # Store for use in other methods self.ref_segs_feat = unique_rows self.ref_segs_indices = unique_indices - - # Build nearest-neighbor graph once + + # Build the nearest-neighbor graph once -- but only when kNN smoothing will + # actually be used (smoothing_factor_nn > 0). The fit + full N-point kneighbors + # query is the dominant construction cost and the resulting (N, knn_k) index + # array is held for the object's lifetime; on the smoothing_factor_nn = 0 path + # it is built and never read, so build_knn=False lets that (common) case skip it. + # Default True preserves the original unconditional behavior. self.knn_k = knn_k - self.nn = NearestNeighbors(n_neighbors=knn_k + 1, metric="euclidean") - self.nn.fit(self.trans_from_feat) + if build_knn: + self.nn = NearestNeighbors(n_neighbors=knn_k + 1, metric="euclidean") + self.nn.fit(self.trans_from_feat) - # Excludes self as first neighbor - distances, indices = self.nn.kneighbors(self.trans_from_feat) - #self.knn_distances = distances[:, 1:] - self.knn_indices = indices[:, 1:] + # Excludes self as first neighbor + distances, indices = self.nn.kneighbors(self.trans_from_feat) + #self.knn_distances = distances[:, 1:] + self.knn_indices = indices[:, 1:] + else: + self.nn = None + self.knn_indices = None def rite_weight_iter(self,weights_started,clusters,total_iterations,weightout_fq,output_file='Reweights',learning_rate=1,smoothing_factor=0.01,smoothing_factor_nn=0.00,seed=42, save_last=10): @@ -126,6 +135,16 @@ def smooth_weights_knn(self, weights, smoothing_factor_nn=0.01): if smoothing_factor_nn <= 0: return weights + # Smoothing was requested but the kNN graph was skipped at construction + # (build_knn=False). Fail loudly rather than dereference an absent graph, so + # the build-time and run-time switches cannot silently disagree. + if self.knn_indices is None: + raise RuntimeError( + "smooth_weights_knn called with smoothing_factor_nn > 0 but the kNN " + "graph was not built (build_knn=False). Construct " + "rite_weight_model(..., build_knn=True) to enable kNN smoothing." + ) + neighbor_avg = np.mean(weights[self.knn_indices], axis=1) smoothed = (1.0 - smoothing_factor_nn) * weights + smoothing_factor_nn * neighbor_avg