Update process_multiSample.py

mNSF/process_multiSample.py

@@ -402,7 +402,56 @@ def interpret_npf_v3(list_fit,list_X,S=10,**kwargs):
   return interpret_nonneg(np.exp(Fhat_c),list_fit[0].W.numpy(),sort=False,**kwargs)
 
 
-
+def reorder_spatial_factors(factors, list_D, list_X_original):
+    """
+    Reorder factors to match the original spatial coordinates order.
+    
+    Args:
+        factors: numpy array of factors from mNSF analysis
+        list_D: list of chunked data dictionaries
+        list_X_original: list of original spatial coordinate dataframes
+    
+    Returns:
+        list of reordered factor matrices, one per sample
+    """
+    reordered_factors = []
+    current_idx = 0
+
+    for sample_idx in range(len(list_X_original)):
+        # Get original number of spots for this sample
+        n_spots_original = list_X_original[sample_idx].shape[0]
+
+        # Count total spots in chunks for this sample
+        n_chunks_per_sample = sum(1 for d in list_D[current_idx:] 
+                                if d['sample_id'][0] == sample_idx)
+
+        # Collect all spot indices and factors for this sample's chunks
+        sample_indices = []
+        sample_factors = []
+
+        for chunk_idx in range(current_idx, current_idx + n_chunks_per_sample):
+            chunk_spots = list_D[chunk_idx]['X']
+            # Get original indices by matching coordinates
+            for idx, spot in chunk_spots.iterrows():
+                orig_idx = list_X_original[sample_idx][
+                    (list_X_original[sample_idx].iloc[:,0] == spot.iloc[0]) & 
+                    (list_X_original[sample_idx].iloc[:,1] == spot.iloc[1])
+                ].index[0]
+                sample_indices.append(orig_idx)
+                sample_factors.append(factors[idx])
+
+        # Convert to numpy arrays
+        sample_indices = np.array(sample_indices)
+        sample_factors = np.array(sample_factors)
+
+        # Create reordered factor matrix
+        reordered = np.zeros((n_spots_original, factors.shape[1]))
+        reordered[sample_indices] = sample_factors
+
+        reordered_factors.append(reordered)
+        current_idx += n_chunks_per_sample
+
+    return reordered_factors
 
 def interpret_nonneg(factors,loadings,lda_mode=False,sort=False):