Optimize ram usage

Infer_1Photon.py

@@ -81,31 +81,35 @@ def run_inference(model, data_loader, conf):
 
 def accumulate_hits(predictions: np.ndarray, reference_points: np.ndarray, 
                     binning_factor: int, number_of_subframes: int = 20):
-    ml_super_frames = np.zeros((number_of_subframes, NY * binning_factor, NX * binning_factor))
+    ml_super_frames = np.zeros((number_of_subframes, NY * binning_factor, NX * binning_factor), dtype=np.uint32)
-    count_frame = np.zeros((NY, NX))
+    count_frame = np.zeros((NY, NX), dtype=np.uint32)
-    subpixel_dist = np.zeros((binning_factor, binning_factor))
+    subpixel_dist = np.zeros((binning_factor, binning_factor), dtype=np.uint32)
     
-    # absolute coordinate = predicted subpixel + reference point
+    chunk_size = len(predictions) // number_of_subframes
-    absolute_positions = predictions + reference_points[:, :2]
-    
-    # super resolution frames (binning)
-    hit_x = np.floor(absolute_positions[:, 0] * binning_factor).astype(int)
-    hit_y = np.floor(absolute_positions[:, 1] * binning_factor).astype(int)
     
     for i in range(number_of_subframes):
-        start_idx = i * len(predictions) // number_of_subframes
+        start_idx = i * chunk_size
-        end_idx = min((i + 1) * len(predictions) // number_of_subframes, len(predictions))
+        end_idx = (i + 1) * chunk_size if i < number_of_subframes - 1 else len(predictions)
-        np.add.at(ml_super_frames[i], (hit_y[start_idx:end_idx], hit_x[start_idx:end_idx]), 1)
         
-    # count frame (by reference point pixel index)
+        pred_chunk = predictions[start_idx:end_idx]
-    ref_x = (reference_points[:, 0] + 1).astype(int)  # reference point is lower-left corner, +1 to get pixel index
+        ref_chunk = reference_points[start_idx:end_idx]
-    ref_y = (reference_points[:, 1] + 1).astype(int)
-    np.add.at(count_frame, (ref_y, ref_x), 1)
         
-    # subpixel distribution
+        abs_pos_chunk = pred_chunk + ref_chunk[:, :2]
-    sub_x = np.floor((absolute_positions[:, 0] % 1) * binning_factor).astype(int)
+        
-    sub_y = np.floor((absolute_positions[:, 1] % 1) * binning_factor).astype(int)
+        # --- Super resolution frames ---
-    np.add.at(subpixel_dist, (sub_y, sub_x), 1)
+        hit_x = np.floor(abs_pos_chunk[:, 0] * binning_factor).astype(np.int32)
+        hit_y = np.floor(abs_pos_chunk[:, 1] * binning_factor).astype(np.int32)
+        np.add.at(ml_super_frames[i], (hit_y, hit_x), 1)
+        
+        # --- Count frame ---
+        ref_x = (ref_chunk[:, 0] + 1).astype(np.int32)
+        ref_y = (ref_chunk[:, 1] + 1).astype(np.int32)
+        np.add.at(count_frame, (ref_y, ref_x), 1)
+        
+        # --- Subpixel distribution ---
+        sub_x = np.floor((abs_pos_chunk[:, 0] % 1) * binning_factor).astype(np.int32)
+        sub_y = np.floor((abs_pos_chunk[:, 1] % 1) * binning_factor).astype(np.int32)
+        np.add.at(subpixel_dist, (sub_y, sub_x), 1)
         
     return ml_super_frames, count_frame, subpixel_dist
 

Infer_3Photon.py

@@ -3,6 +3,8 @@ sys.path.append('./src')
 import matplotlib
 matplotlib.use('Agg')
 
+from etaInterpolationFromPoints import interpolate_eta_from_points
+
 from pathlib import Path
 from omegaconf import OmegaConf
 import torch
@@ -47,11 +49,7 @@ def apply_inverse_transforms(predictions: torch.Tensor, numberOfAugOps: int) ->
     return corrected.mean(dim=1)
 
 def prepare_output_folder(conf):
-    if conf.data.normalize:
+    output_base = Path(conf.experiment.output_base) / conf.experiment.name / (f'3ph_{conf.model.experiment_name}') / (f'augX{conf.inference.num_aug_ops}')
-        normalize_suffix = '_normalized'
-    else:
-        normalize_suffix = ''
-    output_base = Path(conf.experiment.output_base) / conf.experiment.name / conf.model.experiment_name / f'augX{conf.inference.num_aug_ops}{normalize_suffix}'
     output_base.mkdir(parents=True, exist_ok=True)
     OmegaConf.save(conf, output_base / 'config.yaml')
     return output_base
@@ -75,45 +73,52 @@ def run_inference(model, data_loader, conf):
     all_predictions = torch.cat(all_predictions, dim=0)
     # all_predictions = apply_inverse_transforms(all_predictions, conf.inference.num_aug_ops)
     all_predictions += torch.tensor([NSIZE/2., NSIZE/2.]).unsqueeze(0)  # adjust back to original coordinate system
-    print(f'mean x = {torch.mean(all_predictions[:, 0])}, std x = {torch.std(all_predictions[:, 0])}')
+    print(f'[Inference]: mean x = {torch.mean(all_predictions[:, 0])}, std x = {torch.std(all_predictions[:, 0])}')
-    print(f'mean y = {torch.mean(all_predictions[:, 1])}, std y = {torch.std(all_predictions[:, 1])}')
+    print(f'[Inference]: mean y = {torch.mean(all_predictions[:, 1])}, std y = {torch.std(all_predictions[:, 1])}')
     referencePoints = data_loader.dataset.referencePoint ### the lower-left corner of the cluster in absolute coordinate
     referencePoints = np.repeat(referencePoints, 3, axis=0) ### duplicate reference points for 3-photon clusters
     return all_predictions.numpy(), referencePoints
 
 def accumulate_hits(predictions: np.ndarray, reference_points: np.ndarray, 
-                    binning_factor: int):
+                    binning_factor: int, number_of_subframes: int = 20):
-    ### ret 
+    ml_super_frames = np.zeros((number_of_subframes, NY * binning_factor, NX * binning_factor))
-    ml_super_frame = np.zeros((NY*binning_factor, NX*binning_factor), dtype=np.int32)
     count_frame = np.zeros((NY, NX), dtype=np.int32)
     subpixel_dist = np.zeros((binning_factor, binning_factor), dtype=np.int32)
 
     ### absolute coordinate = predicted subpixel + reference point
     absolute_positions = predictions + reference_points
 
-    hit_x_superpixel_idx = np.floor(absolute_positions[:, 0] * binning_factor).astype(int)
-    hit_x_superpixel_idx = np.clip(hit_x_superpixel_idx, 0, NX*binning_factor-1)
-    hit_y_superpixel_idx = np.floor(absolute_positions[:, 1] * binning_factor).astype(int)
-    hit_y_superpixel_idx = np.clip(hit_y_superpixel_idx, 0, NY*binning_factor-1)
-    np.add.at(ml_super_frame, (hit_y_superpixel_idx, hit_x_superpixel_idx), 1)
 
-    hit_x_pixel_idx = np.floor(absolute_positions[:, 0]).astype(int)
+    # super resolution frames (binning)
-    hit_x_pixel_idx = np.clip(hit_x_pixel_idx, 0, NX-1)
+    hit_x = np.floor(absolute_positions[:, 0] * binning_factor).astype(int)
-    hit_y_pixel_idx = np.floor(absolute_positions[:, 1]).astype(int)
+    hit_y = np.floor(absolute_positions[:, 1] * binning_factor).astype(int)
-    hit_y_pixel_idx = np.clip(hit_y_pixel_idx, 0, NY-1)
-    np.add.at(count_frame, (hit_y_pixel_idx, hit_x_pixel_idx), 1)
     
-    subpixel_x_idx = np.floor((absolute_positions[:, 0] % 1) * binning_factor).astype(int)
+    for i in range(number_of_subframes):
-    subpixel_y_idx = np.floor((absolute_positions[:, 1] % 1) * binning_factor).astype(int)
+        start_idx = i * len(predictions) // number_of_subframes
-    np.add.at(subpixel_dist, (subpixel_y_idx, subpixel_x_idx), 1)
+        end_idx = min((i + 1) * len(predictions) // number_of_subframes, len(predictions))
+        np.add.at(ml_super_frames[i], (hit_y[start_idx:end_idx], hit_x[start_idx:end_idx]), 1)
     
-    return ml_super_frame, count_frame, subpixel_dist
+    # count frame (by reference point pixel index)
+    ref_x = (reference_points[:, 0] + 1).astype(int)  # reference point is lower-left corner, +1 to get pixel index
+    ref_y = (reference_points[:, 1] + 1).astype(int)
+    np.add.at(count_frame, (ref_y, ref_x), 1)
     
-def save_results(ml_super_frame, count_frame, subpixel_dist, 
+    # subpixel distribution
+    sub_x = np.floor((absolute_positions[:, 0] % 1) * binning_factor).astype(int)
+    sub_y = np.floor((absolute_positions[:, 1] % 1) * binning_factor).astype(int)
+    np.add.at(subpixel_dist, (sub_y, sub_x), 1)
+    
+    return ml_super_frames, count_frame, subpixel_dist
+
+def save_results(ml_super_frames, ml_super_frames_eta,
+                count_frame, subpixel_dist, subpixel_dist_eta,
                 roi: list, binning_factor: int, output_dir: Path):
     x_st, x_ed, y_st, y_ed = roi
 
     # 1. super-resolution frame
+    np.save(output_dir / '3Photon_ML_superFrames.npy', ml_super_frames)
+    
+    ml_super_frame = np.sum(ml_super_frames, axis=0)
     plt.figure(figsize=(8, 8))
     plt.imshow(ml_super_frame[y_st*binning_factor:y_ed*binning_factor, x_st*binning_factor:x_ed*binning_factor], origin='lower', extent=[x_st, x_ed, y_st, y_ed])
     plt.colorbar(label='Counts')
@@ -124,7 +129,21 @@ def save_results(ml_super_frame, count_frame, subpixel_dist,
     plt.clf()
     np.save(output_dir / '3Photon_ML_superFrame.npy', ml_super_frame)
 
-    # 2. count frame
+    # 2. super-resolution frame with eta interpolation
+    np.save(output_dir / '3Photon_ML_superFrames_etaInterpolated.npy', ml_super_frames_eta)
+
+    ml_super_frame_eta = np.sum(ml_super_frames_eta, axis=0)
+    plt.figure(figsize=(8, 8))
+    plt.imshow(ml_super_frame_eta[y_st*binning_factor:y_ed*binning_factor, x_st*binning_factor:x_ed*binning_factor], origin='lower', extent=[x_st, x_ed, y_st, y_ed])
+    plt.colorbar(label='Counts')
+    plt.title('ML Super-Resolution Frame with Eta Interpolation')
+    plt.xlabel('X (pixel)')
+    plt.ylabel('Y (pixel)')
+    plt.savefig(output_dir / '3Photon_ML_superFrame_etaInterpolated.png', dpi=300, bbox_inches='tight')
+    plt.clf()
+    np.save(output_dir / '3Photon_ML_superFrame_etaInterpolated.npy', ml_super_frame_eta)
+
+    # 3. count frame
     plt.imshow(count_frame[y_st:y_ed, x_st:x_ed], origin='lower', extent=[x_st, x_ed, y_st, y_ed])
     plt.colorbar(label='Counts')
     plt.title('Photon Count Frame')
@@ -134,7 +153,7 @@ def save_results(ml_super_frame, count_frame, subpixel_dist,
     plt.clf()
     np.save(output_dir / '3Photon_count_Frame.npy', count_frame)
 
-    # 3. subpixel distribution
+    # 4. subpixel distribution
     plt.imshow(subpixel_dist, origin='lower', extent=[0, 1, 0, 1])
     plt.colorbar(label='Counts')
     plt.title('Subpixel Distribution')
@@ -146,7 +165,19 @@ def save_results(ml_super_frame, count_frame, subpixel_dist,
     std, mean = np.std(subpixel_dist), np.mean(subpixel_dist)
     print(f"[Plotting]: Sub-pixel distribution: RMS/Mean: {std/mean:.4f}, expected  value = {1/np.sqrt(mean):.4f} for uniform distribution")
 
-    print(f"Results saved to: {output_dir}")
+    # 5. subpixel distribution with eta interpolation
+    plt.imshow(subpixel_dist_eta, origin='lower', extent=[0, 1, 0, 1])
+    plt.colorbar(label='Counts')
+    plt.title('Subpixel Distribution with Eta Interpolation')
+    plt.xlabel('Subpixel X')
+    plt.ylabel('Subpixel Y')
+    plt.savefig(output_dir / '3Photon_subpixel_Distribution_etaInterpolated.png', dpi=300, bbox_inches='tight')
+    plt.close()
+    np.save(output_dir / '3Photon_subpixel_Distribution_etaInterpolated.npy', subpixel_dist_eta)
+    std_eta, mean_eta = np.std(subpixel_dist_eta), np.mean(subpixel_dist_eta)
+    print(f"[Plotting]: Sub-pixel distribution with eta interpolation: RMS/Mean: {std_eta/mean_eta:.4f}, expected  value = {1/np.sqrt(mean_eta):.4f} for uniform distribution")
+
+    print(f"[Plotting]: Results saved to: {output_dir}")
 
 if __name__ == "__main__":
     ### output folder preparation
@@ -166,6 +197,9 @@ if __name__ == "__main__":
     flag_normalize = conf.data.normalize
     nChunks = np.ceil(len(files_list) / conf.inference.chunk_size).astype(int)
 
+    list_of_predictions = []
+    list_of_reference_points = []
+
     ml_super_frame = np.zeros((NY*BinningFactor, NX*BinningFactor), dtype=np.int32)
     count_frame = np.zeros((NY, NX), dtype=np.int32)
     subpixel_dist = np.zeros((BinningFactor, BinningFactor), dtype=np.int32)
@@ -191,8 +225,30 @@ if __name__ == "__main__":
         )
 
         predictions, reference_points = run_inference(model, dataLoader, conf)
-        ml_super_frame_chunk, count_frame_chunk, subpixel_dist_chunk = accumulate_hits(predictions, reference_points, BinningFactor)
+        list_of_predictions.append(predictions)
-        ml_super_frame += ml_super_frame_chunk
+        list_of_reference_points.append(reference_points)
-        count_frame += count_frame_chunk
+
-        subpixel_dist += subpixel_dist_chunk
+        del dataset, dataLoader
-    save_results(ml_super_frame, count_frame, subpixel_dist, roi, BinningFactor, output_dir)
+        torch.cuda.empty_cache()
+        
+    predictions = np.concatenate(list_of_predictions, axis=0)
+    ref_points = np.concatenate(list_of_reference_points, axis=0)
+
+    del list_of_predictions, list_of_reference_points
+
+    ml_super_frames, count_frame, subpixel_dist = accumulate_hits(
+        predictions, ref_points, binning_factor=BinningFactor
+    )
+    
+    print('[Main]: Applying eta interpolation to predictions...')
+    predictions_eta_interpolated = interpolate_eta_from_points(predictions)
+    ml_super_frames_eta, count_frame_eta, subpixel_dist_eta = accumulate_hits(
+        predictions_eta_interpolated, ref_points, binning_factor=BinningFactor
+    )
+
+
+    save_results(ml_super_frames, ml_super_frames_eta,
+                    count_frame, subpixel_dist, subpixel_dist_eta,
+                    roi=roi, binning_factor=BinningFactor, 
+                    output_dir=output_dir
+                )