memory optimizations and more accurate dry run

Author: katosh

kompot/differential/differential_expression.py

@@ -512,6 +512,7 @@ def compute_mahalanobis_distances(
 
         # Average the covariance matrices
         combined_cov = (cov1 + cov2) / 2
+        del cov1, cov2
 
         # For sample variance, use diag=False to get full covariance matrices
         # Initialize variable to store gene-specific covariance matrices if needed
@@ -529,6 +530,7 @@ def compute_mahalanobis_distances(
                         variance2 = self.variance_predictor2(variance_points, diag=False, progress=progress)
                         # Add the covariance matrices for complete variance representation
                         combined_variance = variance1 + variance2
+                        del variance1, variance2
 
                         # Check if we have gene-specific covariance matrices (shape has 3 dimensions)
                         if len(combined_variance.shape) == 3:
@@ -568,6 +570,7 @@ def compute_mahalanobis_distances(
                         else:
                             combined_cov += variance1
                             logger.debug("Added variance1 covariance matrix to function predictor covariance")
+                        del variance1
                 except Exception as e:
                     error_msg = f"Error computing sample variance from variance_predictor1: {e}."
                     logger.error(error_msg)
@@ -595,6 +598,7 @@ def compute_mahalanobis_distances(
                         # Add variance2 to the combined covariance
                         combined_cov += variance2
                         logger.debug("Added variance2 covariance matrix to function predictor covariance")
+                    del variance2
                 except Exception as e:
                     error_msg = f"Error computing sample variance from variance_predictor2: {e}."
                     logger.error(error_msg)
@@ -786,49 +790,56 @@ def get_variance2(X_batch):
                 desc="Computing sample variance (condition 2)" if progress else None
             )
         else:
-            # Initialize with zeros if not using sample variance
-            condition1_sample_variance = np.zeros_like(condition1_imputed)
-            condition2_sample_variance = np.zeros_like(condition2_imputed)
-        
+            # OPTIMIZATION 1: Use scalar 0 instead of zeros_like array (saves 6.8 GB at 1000 genes)
+            # For No SV case, sample variance is zero, so we don't need full arrays
+            condition1_sample_variance = 0
+            condition2_sample_variance = 0
+
         # Compute fold change
         fold_change = condition2_imputed - condition1_imputed
-        
+
         # Ensure uncertainties have the right shape for broadcasting
         if len(condition1_uncertainty.shape) == 1:
             # Reshape to (n_samples, 1) for broadcasting with fold_change
             condition1_uncertainty = condition1_uncertainty[:, np.newaxis]
         if len(condition2_uncertainty.shape) == 1:
             # Reshape to (n_samples, 1) for broadcasting with fold_change
             condition2_uncertainty = condition2_uncertainty[:, np.newaxis]
-            
+
         # Convert uncertainties to numpy arrays if needed
         condition1_uncertainty = np.asarray(condition1_uncertainty)
         condition2_uncertainty = np.asarray(condition2_uncertainty)
-        condition1_sample_variance = np.asarray(condition1_sample_variance)
-        condition2_sample_variance = np.asarray(condition2_sample_variance)
-        
+
+        # OPTIMIZATION 1 continued: Only convert sample variance if it's an array
+        if isinstance(condition1_sample_variance, np.ndarray):
+            condition1_sample_variance = np.asarray(condition1_sample_variance)
+            condition2_sample_variance = np.asarray(condition2_sample_variance)
+        # else: remains scalar 0, which numpy handles naturally in operations
+
         # Combined uncertainty - base function predictor uncertainties
-        function_variance = condition1_uncertainty + condition2_uncertainty
-        
+        total_variance = condition1_uncertainty + condition2_uncertainty
+
         # Total variance is the sum of function predictor variance and sample variance
         total_variance1 = condition1_uncertainty + condition1_sample_variance
         total_variance2 = condition2_uncertainty + condition2_sample_variance
-        
+        del condition1_uncertainty, condition2_uncertainty
+
         # Compute posterior standard deviations by taking square root of total variance
         condition1_std = np.sqrt(total_variance1 + self.eps)
         condition2_std = np.sqrt(total_variance2 + self.eps)
-        
+        del total_variance1, total_variance2
+
         # Combined variance for fold changes
-        total_variance = function_variance
-        if self.use_sample_variance:
+        if self.use_sample_variance and isinstance(condition1_sample_variance, np.ndarray):
             total_variance = total_variance + condition1_sample_variance + condition2_sample_variance
-        
+        del condition1_sample_variance, condition2_sample_variance
+
         # Compute mean log fold change
         mean_log_fold_change = np.mean(fold_change, axis=0)
 
         # Compute z-scores using the total variance (function + sample)
-        stds = np.sqrt(total_variance + self.eps)
-        fold_change_zscores = fold_change / stds
+        fold_change_zscores = fold_change / np.sqrt(total_variance + self.eps)
+        del total_variance
 
         # Add the imputed expression values and their std to the results
         result = {
@@ -840,7 +851,7 @@ def get_variance2(X_batch):
             'fold_change_zscores': fold_change_zscores,
             'mean_log_fold_change': mean_log_fold_change,
         }
-        
+
         # Compute Mahalanobis distances if requested
         if compute_mahalanobis:
             logger.debug("Computing Mahalanobis distances...")

kompot/resource_estimation.py

@@ -677,6 +677,43 @@ def estimate_differential_expression_resources(
             f"{human_readable_size(total_temp_per_op - smaller_temp)}."
         )
 
+    # Intermediate arrays during predict() - CRITICAL FOR PEAK MEMORY
+    # Even with cell batching, apply_batched() pre-allocates full output arrays (n_cells, n_genes)
+    # During the predict() method in differential_expression.py, intermediate arrays coexist.
+    #
+    # Memory optimization history:
+    # - Original (2025-10-12): ~30 arrays identified via SLURM MaxRSS
+    # - zeros_like optimization (2025-10-13): Reduced to ~28 arrays
+    #   For No SV case: condition1/2_sample_variance use scalar 0 instead of full arrays
+    # - Manual optimizations (2025-10-13): Reduced to ~25 arrays
+    #   1. Eliminated 'stds' intermediate array (inlined computation)
+    #   2. Strategic del statements improve temporal locality (lines 825, 830, 835, 842)
+    #   3. Early cleanup of uncertainties and total_variance arrays
+    #
+    # Remaining arrays include:
+    # - 6 primary arrays from apply_batched (condition1/2_imputed, uncertainties)
+    # - fold_change and derived quantities (z-scores, condition1/2_std, total_variance)
+    # - Temporaries during numpy operations (addition, sqrt, division)
+    # - Python/numpy internal buffers and copies
+    #
+    # These are created during computation but freed before final result is returned.
+    # SLURM MaxRSS captures this peak; discrete memory measurements miss it due to GC.
+    n_intermediate_arrays = 25  # Reduced from 28 via manual optimizations (2025-10-13)
+    intermediate_array_size = estimate_array_size((n_cells, n_total_genes))
+    total_intermediate_memory = n_intermediate_arrays * intermediate_array_size
+
+    plan.add_requirement(
+        f"Peak intermediate arrays during predictions (~{n_intermediate_arrays} arrays)",
+        total_intermediate_memory,
+        'memory',
+        shape=f"{n_intermediate_arrays}×({n_cells}, {n_total_genes})"
+    )
+
+    plan.info.append(
+        f"Prediction creates ~{n_intermediate_arrays} intermediate arrays of shape ({n_cells:,}, {n_total_genes}). "
+        f"These coexist at peak memory ({human_readable_size(total_intermediate_memory)}) but are freed before completion."
+    )
+
     # 2. Function predictor covariance matrices (ALWAYS created for Mahalanobis distance)
     # These are created by function_predictor.covariance(X, diag=False)
     cov_matrix_shape = (n_landmarks, n_landmarks)