perf: cap PyTorch threads, bump fix_illumination forks, skip intermediate pyramids

scripts/linum_crop_3d_mosaic_below_interface.py

@@ -55,6 +55,7 @@ def _build_arg_parser() -> argparse.ArgumentParser:
         "to finding the interface*. Original values will\n"
         "remain in output clipped volume (range [0-100]).",
     )
+    p.add_argument("--n_levels", type=int, default=5, help="Number of levels in pyramid representation. [%(default)s]")
     return p
 
 
@@ -101,7 +102,7 @@ def main() -> None:
 
     crop_dask = da.from_array(vol_crop, chunks=vol.chunks)
     # Save cropped volume as OME-Zarr
-    save_omezarr(crop_dask, output_path, voxel_size=res, chunks=vol.chunks)
+    save_omezarr(crop_dask, output_path, voxel_size=res, chunks=vol.chunks, n_levels=args.n_levels)
 
     # Collect metrics using helper function
     original_shape = vol.shape

scripts/linum_fix_illumination_3d.py

@@ -42,6 +42,7 @@ def _build_arg_parser() -> argparse.ArgumentParser:
         type=float,
         help="Values above this percentile will be clipped when\nestimating the flatfield (inside range [0-100]).",
     )
+    p.add_argument("--n_levels", type=int, default=5, help="Number of levels in pyramid representation. [%(default)s]")
     p.add_argument(
         "--use_gpu",
         default=True,
@@ -57,9 +58,12 @@ def _build_arg_parser() -> argparse.ArgumentParser:
 
 def process_tile(params: dict) -> tuple:
     """Process a tile and add it to the output mosaic."""
-    from linumpy.config.threads import apply_threadpool_limits
+    from linumpy.config.threads import configure_all_libraries
 
-    apply_threadpool_limits()
+    # configure_all_libraries() also caps PyTorch intra-/inter-op threads
+    # (used by BaSiCPy). apply_threadpool_limits() alone misses torch and
+    # lets it default to ~half the host cores, oversubscribing the node.
+    configure_all_libraries()
 
     file = params["slice_file"]
     z = params["z"]
@@ -182,7 +186,7 @@ def main() -> None:
         print(f"Minimum value in the output volume is {min_value}. Clipping at 0.")
         out_dask = da.clip(out_dask, 0.0, None)
 
-    save_omezarr(out_dask, output_zarr, voxel_size=resolution, chunks=vol.chunks)
+    save_omezarr(out_dask, output_zarr, voxel_size=resolution, chunks=vol.chunks, n_levels=args.n_levels)
     tmp_dir.cleanup()
 
 

scripts/linum_normalize_intensities_per_slice.py

@@ -39,6 +39,7 @@ def _build_arg_parser() -> argparse.ArgumentParser:
     p.add_argument("--sigma", type=float, default=1.0, help="Smoothing sigma for estimating the agarose mask. [%(default)s]")
     p.add_argument("--use_gpu", default=True, action=argparse.BooleanOptionalAction, help="Use GPU acceleration if available.")
     p.add_argument("--verbose", action="store_true", help="Print GPU information.")
+    p.add_argument("--n_levels", type=int, default=3, help="Number of levels in pyramid representation. [%(default)s]")
     return p
 
 
@@ -71,7 +72,7 @@ def main() -> None:
 
     vol_normalized, background_thresholds = normalize_volume(vol_data, agarose_mask, args.percentile_max)
 
-    save_omezarr(da.from_array(vol_normalized), args.out_image, res, n_levels=3)
+    save_omezarr(da.from_array(vol_normalized), args.out_image, res, n_levels=args.n_levels)
 
     collect_normalization_metrics(
         vol_normalized=vol_normalized,

scripts/linum_stitch_3d_refined.py

@@ -84,6 +84,7 @@ def _build_arg_parser() -> argparse.ArgumentParser:
         "--output_refinements", type=str, default=None, help="Output JSON file to save computed refinements for analysis."
     )
     p.add_argument("--overwrite", "-f", action="store_true", help="Overwrite output if it exists.")
+    p.add_argument("--n_levels", type=int, default=3, help="Number of levels in pyramid representation. [%(default)s]")
     return p
 
 
@@ -292,7 +293,7 @@ def main() -> None:
 
     from linumpy.io.zarr import save_omezarr
 
-    save_omezarr(da.from_array(output), output_file, resolution, n_levels=3)
+    save_omezarr(da.from_array(output), output_file, resolution, n_levels=args.n_levels)
 
     # Collect metrics
     from linumpy.metrics import PipelineMetrics

workflows/reconst_3d/nextflow.config

@@ -391,9 +391,8 @@ process {
 
     withName: "resample_mosaic_grid" {
         scratch = false
-        // Parallel resampling on GPU. Peak ~9 GB/fork (Gaussian intermediate at
-        // 2x tile size). With 2x A6000 (48 GB each), 6 forks = 3/GPU = ~27 GB,
-        // leaving headroom for fix_illumination running concurrently.
+        // Measured: ~1 GB GPU memory per fork. IO (RAID at 100%) is the
+        // gating factor at 6 forks, not GPU. Bumping higher won't help.
         maxForks = params.use_gpu ? 6 : null
     }
 
@@ -405,8 +404,9 @@ process {
     }
 
     withName: "fix_illumination" {
-        // Limit to 1 parallel instance - BaSiCPy/PyTorch is memory-intensive
-        maxForks = params.use_gpu ? 1 : null
+        // Measured: BaSiCPy/PyTorch uses ~374 MiB per fork on this dataset.
+        // 4 forks ≈ 1.5 GB, well under per-GPU capacity (49 GB).
+        maxForks = params.use_gpu ? 4 : null
         // Don't set CUDA_VISIBLE_DEVICES - let linumpy.gpu auto-select GPU with most free memory
     }
 

workflows/reconst_3d/soct_3d_reconst.nf

@@ -128,7 +128,7 @@ process resample_mosaic_grid {
     def gpu_flag = params.use_gpu ? "--use_gpu" : "--no-use_gpu"
     """
     linum_resample_mosaic_grid.py ${mosaic_grid} "mosaic_grid_z${slice_id}_resampled.ome.zarr" \
-        -r ${params.resolution} ${gpu_flag} -v
+        -r ${params.resolution} ${gpu_flag} --n_levels 0 -v
     """
 
     stub:
@@ -147,7 +147,8 @@ process fix_focal_curvature {
     script:
     def gpu_flag = params.use_gpu ? "--use_gpu" : "--no-use_gpu"
     """
-    linum_detect_focal_curvature.py ${mosaic_grid} "mosaic_grid_z${slice_id}_focal_fix.ome.zarr" ${gpu_flag}
+    linum_detect_focal_curvature.py ${mosaic_grid} "mosaic_grid_z${slice_id}_focal_fix.ome.zarr" \\
+        --n_levels 0 ${gpu_flag}
     """
 
     stub:
@@ -170,7 +171,7 @@ process fix_illumination {
     """
     linum_fix_illumination_3d.py ${mosaic_grid} "mosaic_grid_z${slice_id}_illum_fix.ome.zarr" \
         --n_processes ${params.processes} \
-        --percentile_max ${params.clip_percentile_upper} ${gpu_flag}
+        --percentile_max ${params.clip_percentile_upper} ${gpu_flag} --n_levels 0
     """
 
     stub:
@@ -244,6 +245,7 @@ process stitch_3d_with_refinement {
         --refinement_mode blend_shift \
         --max_refinement_px ${params.max_blend_refinement_px} \
         ${transform_arg} \
+        --n_levels 0 \
         -f
     """
 
@@ -291,7 +293,7 @@ process beam_profile_correction {
     script:
     """
     linum_compensate_psf_model_free.py ${slice_3d} "slice_z${slice_id}_axial_corr.ome.zarr" \
-        --percentile_max ${params.clip_percentile_upper}
+        --percentile_max ${params.clip_percentile_upper} --n_levels 0
     """
 
     stub:
@@ -315,7 +317,7 @@ process crop_interface {
     linum_crop_3d_mosaic_below_interface.py ${image} "slice_z${slice_id}_crop_interface.ome.zarr" \
         --depth ${params.crop_interface_out_depth} \
         --crop_before_interface \
-        --percentile_max ${params.clip_percentile_upper}
+        --percentile_max ${params.clip_percentile_upper} --n_levels 0
     """
 
     stub:
@@ -338,7 +340,7 @@ process normalize {
     def gpu_flag = params.use_gpu ? "--use_gpu" : "--no-use_gpu"
     """
     linum_normalize_intensities_per_slice.py ${image} "slice_z${slice_id}_normalize.ome.zarr" \
-        --percentile_max ${params.clip_percentile_upper} ${gpu_flag}
+        --percentile_max ${params.clip_percentile_upper} ${gpu_flag} --n_levels 0
     """
 
     stub: