diff --git a/scripts/linum_crop_3d_mosaic_below_interface.py b/scripts/linum_crop_3d_mosaic_below_interface.py index 0fbcf2c9..ecdfdb44 100644 --- a/scripts/linum_crop_3d_mosaic_below_interface.py +++ b/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 diff --git a/scripts/linum_fix_illumination_3d.py b/scripts/linum_fix_illumination_3d.py index 7063c375..b55174e5 100644 --- a/scripts/linum_fix_illumination_3d.py +++ b/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() diff --git a/scripts/linum_normalize_intensities_per_slice.py b/scripts/linum_normalize_intensities_per_slice.py index 1630dd19..8b5503ca 100644 --- a/scripts/linum_normalize_intensities_per_slice.py +++ b/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, diff --git a/scripts/linum_stitch_3d_refined.py b/scripts/linum_stitch_3d_refined.py index a4490068..dfc859f8 100644 --- a/scripts/linum_stitch_3d_refined.py +++ b/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 diff --git a/workflows/reconst_3d/nextflow.config b/workflows/reconst_3d/nextflow.config index f7cafa5b..3765a1ab 100644 --- a/workflows/reconst_3d/nextflow.config +++ b/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 } diff --git a/workflows/reconst_3d/soct_3d_reconst.nf b/workflows/reconst_3d/soct_3d_reconst.nf index 3dbac57b..529afe1d 100644 --- a/workflows/reconst_3d/soct_3d_reconst.nf +++ b/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: