Add opt-in build_knn flag to skip the unused kNN graph

rite_weight/rite_weight.py

@@ -6,29 +6,38 @@
 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
         unique_rows, unique_indices = np.unique(trans_from_feat_rounded, axis=0, return_index=True)
         # 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