diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 321f3ad..91ecc2f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -75,7 +75,7 @@ jobs: strategy: matrix: - python-version: ["3.10", "3.11", "3.12"] # Ensure versions are strings + python-version: ["3.11", "3.12", "3.13"] steps: # Set debug environment variable @@ -100,4 +100,4 @@ jobs: # Run Nox sessions - name: Run Nox - run: nox --python ${{ matrix.python-version }} \ No newline at end of file + run: nox --python ${{ matrix.python-version }} diff --git a/docs/advanced/index.md b/docs/advanced/index.md index 74c86ea..450b421 100644 --- a/docs/advanced/index.md +++ b/docs/advanced/index.md @@ -9,7 +9,6 @@ Pages: - [Python API](./python-api.md): Use `cardiotensor` as a Python library for structure tensor analysis, angle computation, tractography, and plotting. - [Transmural Analysis](./cardio-analysis.md): Extract transmural profiles from angle and FA maps with `cardio-analysis`. -- [Streamlines](./streamlines.md): Generate and visualize 3D streamlines with `cardio-generate-streamlines` and `cardio-visualize-streamlines`. - [SLURM Launcher](./slurm-launcher.md): Submit chunked `cardio-tensor` jobs on HPC clusters. See also: diff --git a/docs/advanced/slurm-launcher.md b/docs/advanced/slurm-launcher.md index 0111cdc..aef1d51 100644 --- a/docs/advanced/slurm-launcher.md +++ b/docs/advanced/slurm-launcher.md @@ -42,6 +42,15 @@ Other useful flags: - `--no_monitor`: submit jobs and return immediately (do not watch output progress). - `--dry_run`: generate the `.slurm` script and print `sbatch` command without submitting. +When `LOW_MEMORY=True`, the launcher estimates the padded tensor, mask, worker, +and post-processing peak from `--mem_gb`. If the requested `--chunk_size` is +unsafe, it is reduced before the array is submitted. Low-memory tensor files +default to `OUTPUT_PATH/.cardiotensor_scratch`, not node `/tmp`. Set +`LOW_MEMORY_DIR` only when another data-processing filesystem should be used. +The generated job prints the selected tensor scratch root and its free capacity. +After an OOM or cancelled job, remove any orphaned `cardiotensor_tensor_*` +directories from that scratch root. + Run help for the full list: ```bash diff --git a/docs/advanced/streamlines.md b/docs/advanced/streamlines.md deleted file mode 100644 index 07004da..0000000 --- a/docs/advanced/streamlines.md +++ /dev/null @@ -1,112 +0,0 @@ -# Streamlines - -Use this page to generate and inspect 3D streamlines with: - -- `cardio-generate-streamlines` -- `cardio-visualize-streamlines` - -## Before You Start - -- In the config file, set `WRITE_VECTORS = True` and `TEST = False`. -- Run `cardio-tensor` on your full dataset. -- Make sure these outputs exist in `OUTPUT_PATH`: - - `eigen_vec/` - - `FA/` - - angle folders (`HA`/`IA` or `AZ`/`EL`, depending on `ANGLE_MODE`) - -!!! note - - It is useful to provide a `MASK_PATH` to avoid placing seeds and tracing streamline outside the sample, even though the FA threshold should avoid this. - - -## Generate Streamlines - -Basic example: - -```console -$ cardio-generate-streamlines ./parameters_example.conf -``` - -Useful options: - -- `--seeds`: number of seeds (default: 20000). -- `--fa-seed-min`: minimum FA to place seeds (default: 0.2). -- `--fa-threshold`: minimum FA to continue tracking (default: 0.1). -- `--angle`: maximum turning angle (degrees - default: 60.0). -- `--step`: integration step length (voxels - default: 0.5). -- `--min-len`: minimum streamline length (points - default: 10). -- `--bin`: downsampling factor for faster processing. -- `--start-z/--end-z` (and `x`, `y` variants): process only a sub-volume. - -Generated files: - -- `OUTPUT_PATH/streamlines.trk` - - -!!! note - - Use `--bin` if the dataset is too big to fit in RAM. This will bin the eigenvectors and the maps. - -## Visualize Streamlines - -```console -$ cardio-visualize-streamlines ./parameters_example.conf --downsample 5 -``` - -Useful options: - -- `--color-by`: `auto`, `elevation`, or stored fields like `HA`, `IA`, `AZ`, `EL`. -- `--list-color-by`: list available fields in the `.trk`. -- `--line-width`: tube width. -- `--subsample`: keep every Nth streamline. -- `--min-length`: filter short streamlines. -- `--crop-x`, `--crop-y`, `--crop-z`: spatial crop ranges. -- `--screenshot`: save a PNG snapshot. - -!!! note - - Use `--downsample` to keep fewer points per streamline and improve rendering performance on large datasets. - -
-![Streamline visualization](../assets/images/streamline_all_advanced.png){ width="75%" } -
Example 3D streamline rendering in a human heart.
-
- -The dataset used as an example here can be find at [DOI](https://doi.org/10.15151/esrf-dc-1634390196) - - -## Keybindings - -In the interactive viewer: - -- `O`: toggle clipping plane on/off. -- `H`: show/hide clipping plane gizmo. -- `I`: flip clipping plane normal. -- `R`: reset clipping plane to center. -- `+` / `-`: increase/decrease streamline thickness. -- `B`: toggle background color (black/white). -- `S`: show/hide scale bar. -- `P`: save a high-resolution PNG snapshot. - -!!! Note - - During interaction (rotate/pan/zoom), the viewer switches to a low-resolution for better responsiveness, then automatically returns to full resolution when interaction ends. - -
-![Streamline visualization](../assets/images/streamline_crop_advanced.png){ width="75%" } -
Example 3D streamline rendering in a human heart, cropped using the clipping plane.
-
- - -## Quick Tuning Workflow - -1. Start with `--seeds 10000`. -2. Check global coverage and continuity. -3. Increase `--seeds` if regions are undersampled. -4. Increase `--fa-threshold` or decrease `--angle` to reduce implausible curves. -5. Use `--subsample` and `--downsample` for faster visual inspection. - -!!! warning - - Interpretation depends strongly on tracking parameters. Compare conditions only when parameters, such as seeding and thresholds, are similar. - diff --git a/docs/getting-technical/angles.md b/docs/getting-technical/angles.md index 724b3c5..a4cc972 100644 --- a/docs/getting-technical/angles.md +++ b/docs/getting-technical/angles.md @@ -1,60 +1,53 @@ # Angle Definitions -This page explains how helical and intrusion angles are calculated from the 3D eigenvector field derived by `cardiotensor`. +Cardiotensor computes helical angle (HA) and intrusion angle (IA) from the local myocyte-axis eigenvector. -By default, Cardiotensor reports **unprojected 3D angles**. The primary eigenvector is not first flattened onto a 2D plane before the angle is measured, because projection discards one component of the vector and can bias the resulting angle. Projection-based angles are available only for legacy comparison. +By default, these are **unprojected 3D angles**. The vector is not flattened onto a 2D plane before measuring HA or IA, because projection removes one component of the vector and can bias the result. -## Coordinate System +## Local Components -A transformation to a cylindrical coordinate system is defined for each voxel based on an approximation of the left ventricle (LV) centerline. +Cardiotensor expresses the vector in a centerline-based cylindrical frame for the left ventricle: -- **Radial (r)**: outward from the LV center -- **Circumferential (θ)**: tangential around the ventricle -- **Longitudinal (z)**: base to apex direction +- **R**: radial component, outward from the LV centerline +- **C**: circumferential component, tangent around the ventricle +- **L**: longitudinal component, along the LV axis -To compute local fiber angles consistently, all eigenvectors are first rotated into this cylindrical coordinate frame. This alignment is performed using the Rodrigues rotation formula, which computes the minimal-angle rotation that maps the global reference axis (here the z-axis) onto the local longitudinal axis at each point. This allows a robust comparison of orientations across the myocardium. +This is a practical LV cylindrical approximation. It should not be read as a full reproduction of anatomically corrected methods that explicitly model epicardial curvature. -## Helical Angle (HA) +## Helical Angle -The helical angle is defined as the angle between the primary myocyte-axis eigenvector \\( \vec{v}_1 \\) and the local circumferential plane. +The helical angle measures how far the vector points out of the local horizontal, or short-axis, plane. That plane is spanned by the radial and circumferential directions. -In local cylindrical coordinates, with radial component \\(R\\), circumferential component \\(C\\), and longitudinal component \\(L\\), the unprojected helical angle is: +```text +HA = atan2(L, sqrt(R^2 + C^2)) +``` -\\[ -\mathrm{HA} = \arctan2\left(L, \sqrt{R^2 + C^2}\right) -\\] +Typical values are around -60 degrees at the epicardium, 0 degrees in the mid-wall, and +60 degrees at the endocardium. -It captures the transmural variation of fiber orientation from epicardium to endocardium. +## Intrusion Angle -Typical pattern: -- ~−60° at epicardium -- ~0° in mid-wall -- ~+60° at endocardium +The intrusion angle measures radial deviation from the local tangential plane. That plane is spanned by the circumferential and longitudinal directions. -## Intrusion Angle (IA) +```text +IA = atan2(R, sqrt(C^2 + L^2)) +``` -The intrusion angle is the angle between the primary myocyte-axis eigenvector \\( \vec{v}_1 \\) and the **tangential plane** (longitudinal + circumferential). +## Projected Legacy Angles -Using the same local components, the unprojected intrusion angle is: +Set `PROJECTED_ANGLES = True` only when comparing with legacy projection-based maps: -\\[ -\mathrm{IA} = \arctan2\left(R, \sqrt{C^2 + L^2}\right) -\\] +```text +HA_projected = atan2(L, C) +IA_projected = atan2(R, C) +``` -It captures radial deviation of fiber aggregates and can help identify wall thickening or microstructural disruptions. +`IA_projected` corresponds to the projected transverse angle terminology often used in the literature. -## Projection Bias - -Conventional projected angles are computed after removing one vector component, for example \\(\arctan2(L, C)\\) for projected helical angle or \\(\arctan2(R, C)\\) for projected intrusion angle. These projected quantities can differ from the true 3D orientation when the discarded component is large. - -Set `PROJECTED_ANGLES = True` only when you need legacy projected `HA_projected` and `IA_projected` maps for comparison with literature. - -This convention follows the projection-error discussion in Agger et al., "Anatomically correct assessment of the orientation of the cardiomyocytes using diffusion tensor imaging", *NMR in Biomedicine* (2020), https://doi.org/10.1002/nbm.4205. +Projection-based angles can differ from the true 3D orientation when the discarded component is large. This projection bias is discussed in Agger et al., "Anatomically correct assessment of the orientation of the cardiomyocytes using diffusion tensor imaging", *NMR in Biomedicine* (2020), https://doi.org/10.1002/nbm.4205. ## Angle Ranges Both angles are reported in degrees: -- **HA**: −90° to +90° -- **IA**: −90° to +90° -Angles are defined in a left-handed cylindrical coordinate system aligned to the LV. +- **HA**: -90 to +90 degrees +- **IA**: -90 to +90 degrees diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index adfeef3..f6c6a4c 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -36,7 +36,7 @@ RHO = 0.6 TRUNCATE = 4 # Padding to avoid border artifacts -# Default value is TRUNCATE * RHO + 0.5 +# By default this includes the support of both Gaussian filtering stages. # VERTICAL_PADDING = 10 # Number of slices to load into memory at a time during processing. @@ -46,10 +46,22 @@ N_CHUNK = 20 # Enable GPU computation during the structure tensor calculation (True/False) USE_GPU = True -# Whether to save the orientation vectors (as .npy) (True/False) +# Store temporary eigenvalue/eigenvector volumes in memory-mapped files. +# Enable this when tensor outputs do not comfortably fit in RAM. Temporary +# files default to OUTPUT_PATH/.cardiotensor_scratch, never the system /tmp. +# Unsafe N_CHUNK values are reduced automatically from the available memory. +LOW_MEMORY = False +# Optional explicit scratch root on your data-processing filesystem. +# LOW_MEMORY_DIR = /path/to/data-proposal/cardiotensor_scratch + +# Whether to save the orientation vectors (True/False) # Use for 3D vector/fiber visualisation WRITE_VECTORS = True +# Vector storage format: npy or zarr +# Use zarr for compressed, chunked storage of very large volumes. +VECTOR_FORMAT = zarr + # Specify the processing direction: # - True: Process slices from the beginning (0) to the end. # - False: Process slices from the end to the beginning. @@ -107,6 +119,8 @@ OUTPUT_TYPE = 8bit !!! note Modify the configuration file as needed to fit your dataset. + Relative dataset, mask, and output paths are resolved from the directory + containing the configuration file. --- @@ -127,14 +141,34 @@ OUTPUT_TYPE = 8bit - **`SIGMA`**: Noise scale before gradient computation. Helps reduce noise while preserving structures. - **`RHO`**: Integration scale for smoothing tensor components. Larger values yield smoother orientation fields. - **`TRUNCATE`**: Multiple of RHO for defining the gradient filter kernel size. -- **`VERTICAL_PADDING (optional)`**: Extra padding (in voxels) to avoid edge artifacts. If not set, defaults to TRUNCATE * RHO + 0.5. +- **`VERTICAL_PADDING (optional)`**: Extra padding (in voxels) to avoid edge + artifacts. If omitted, Cardiotensor includes the support radius of both the + `SIGMA` and `RHO` Gaussian filters. - **`N_CHUNK`**: Number of slices to process simultaneously. - **`USE_GPU`**: Enable GPU computation (requires CuPy). -- **`WRITE_VECTORS`**: Save orientation 3rd vector as `.npy` files. +- **`LOW_MEMORY`**: Store temporary tensor outputs on disk using memory maps. + This reduces RAM use but can be slower unless the temporary directory is on + fast node-local storage. Cardiotensor treats `N_CHUNK` as an upper bound and + reduces it when the padded chunk would exceed a conservative cgroup-aware + memory estimate. +- **`LOW_MEMORY_DIR (optional)`**: Scratch root for tensor memory maps. If + omitted, CardioTensor uses `OUTPUT_PATH/.cardiotensor_scratch`; it never + falls back to the system `/tmp`. The required and available scratch + capacities are printed before calculation. Temporary task directories are + removed after normal completion. Because OOM and `SIGKILL` cannot run + cleanup code, inspect this scratch root after failed jobs and remove orphaned + `cardiotensor_tensor_*` directories. +- **`WRITE_VECTORS`**: Save the orientation vector field. +- **`VECTOR_FORMAT`**: Vector storage backend: + - `npy`: one `float32` file per Z slice in `eigen_vec/`. + - `zarr` (default): compressed, chunked and sharded storage in `eigen_vec.zarr/`. + This is recommended for very large volumes and supports safe slice-level + restart tracking. !!! warning - Orientation vectors are saved in `float32` format and may consume significant disk space. + Orientation vectors are saved in `float32` format. Zarr uses Blosc Zstd + compression with bit-shuffle; zero-filled masked regions generally compress well. - **`REVERSE`**: Process volume from end to start if set to `True`. diff --git a/examples/parameters_example.conf b/examples/parameters_example.conf index 811cc20..5935892 100644 --- a/examples/parameters_example.conf +++ b/examples/parameters_example.conf @@ -24,7 +24,7 @@ RHO = 0.6 TRUNCATE = 4 # Padding to avoid border artifacts -# Default value is TRUNCATE * RHO + 0.5 +# By default this includes the support of both Gaussian filtering stages. # VERTICAL_PADDING = 10 # Number of slices to load into memory at a time during processing. @@ -34,10 +34,22 @@ N_CHUNK = 20 # Enable GPU computation during the structure tensor calculation (True/False) USE_GPU = True -# Whether to save the orientation vectors (as .npy) (True/False) +# Store temporary tensor outputs as memory-mapped files. +# By default they use OUTPUT_PATH/.cardiotensor_scratch, never system /tmp. +# N_CHUNK is automatically reduced when its conservative peak exceeds the job limit. +LOW_MEMORY = False +# Optional explicit scratch root on your data-processing filesystem. +# LOW_MEMORY_DIR = /path/to/data-proposal/cardiotensor_scratch + +# Whether to save the orientation vectors (True/False) # Use for 3D vector/fiber visualisation WRITE_VECTORS = True +# Vector storage format: npy or zarr (default) +# Choose npy for one file per slice, or zarr for compressed, chunked storage that +# is better suited to very large volumes and slice-by-slice processing. +VECTOR_FORMAT = zarr + # Specify the processing direction: # - True: Process slices from the beginning (0) to the end. # - False: Process slices from the end to the beginning. @@ -53,8 +65,8 @@ WRITE_ANGLES = True # az_el → Azimuth / Elevation (generic vector orientation in 3D) ANGLE_MODE = ha_ia -# Coordinates of points along the left ventricle axis. -# The first point should be coordinates of the mitral valve point in the volume ([X, Y, Z]) +# Coordinates of points along the left ventricle axis. +# The first point should be coordinates of the mitral valve point in the volume ([X, Y, Z]) # The last point should be coordinates of the apex point in the volume ([X, Y, Z]) # Intermediate points will be interpolated to create a curved centre line. AXIS_POINTS = [104,110,116], [41,87,210], [68,95,162] diff --git a/mkdocs.yml b/mkdocs.yml index ec03d0b..66fb44d 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -34,7 +34,6 @@ nav: - advanced/index.md - Python API: advanced/python-api.md - Transmural Analysis: advanced/cardio-analysis.md - - Streamlines: advanced/streamlines.md - SLURM Launcher: advanced/slurm-launcher.md - Reference: - reference/index.md diff --git a/noxfile.py b/noxfile.py index 4f63169..80d046b 100644 --- a/noxfile.py +++ b/noxfile.py @@ -4,7 +4,7 @@ # nox.options.reuse_existing_virtualenvs = True # Define default Python versions for testing -PYTHON_VERSIONS = ["3.10", "3.11", "3.12"] +PYTHON_VERSIONS = ["3.11", "3.12", "3.13"] # Locations to check for linting, formatting, and type checking PACKAGE_LOCATIONS = ["src", "tests", "noxfile.py", "pyproject.toml"] diff --git a/pyproject.toml b/pyproject.toml index 2616710..d8a187a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,7 +11,6 @@ classifiers = [ "Operating System :: OS Independent", "Programming Language :: Python :: 3", "Programming Language :: Python :: 3 :: Only", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", "Programming Language :: Python :: 3.13", @@ -22,9 +21,7 @@ classifiers = [ ] dependencies = [ "PyQt5", - "SimpleITK", "alive_progress", - "dask", "dipy", "fury", "glymur", @@ -38,16 +35,17 @@ dependencies = [ "psutil", "pyvista", "scikit-image", - "structure-tensor>=0.3.0", + "structure-tensor>=0.3.4", "tifffile", "tqdm", + "zarr>=3,<4", ] description = "Toolkit designed for quantifying and visualising 3D cardiomyocytes orientations in heart images" keywords = ["cardiotensor, cardiomyocytes, heart, orientation, structure tensor, image processing"] license = {file = "LICENSE"} name = "cardiotensor" readme = "README.md" -requires-python = ">=3.10" +requires-python = ">=3.11" version = "1.1.9" [project.optional-dependencies] @@ -121,7 +119,7 @@ testpaths = [ [tool.ruff] fix = true force-exclude = true -target-version = "py312" +target-version = "py311" lint.select = ["F", "I", "UP"] lint.mccabe.max-complexity = 10 lint.pep8-naming.classmethod-decorators = ["classmethod", "pydantic.validator"] diff --git a/src/cardiotensor/analysis/gui_analysis_tool.py b/src/cardiotensor/analysis/gui_analysis_tool.py index f8fc0fb..19c8d77 100644 --- a/src/cardiotensor/analysis/gui_analysis_tool.py +++ b/src/cardiotensor/analysis/gui_analysis_tool.py @@ -1,4 +1,3 @@ -import sys from pathlib import Path import cv2 @@ -105,9 +104,11 @@ def plot_label_and_limits(mode: str): if m == "FA": return "Fractional Anisotropy", 0.0, 1.0 if m in {"HA", "IA", "EL"}: - name = {"HA": "Helical Angle", "IA": "Intrusion Angle", "EL": "Elevation Angle"}[ - m - ] + name = { + "HA": "Helical Angle", + "IA": "Intrusion Angle", + "EL": "Elevation Angle", + }[m] return f"{name} (°)", -90.0, 90.0 if m == "AZ": return "Azimuth (°)", 0.0, 360.0 @@ -158,12 +159,12 @@ def __init__( # Discover available modes, select default self.available_modes = discover_modes(self.output_path) if not self.available_modes: - sys.exit( + raise FileNotFoundError( f"No angle or FA folders found in {self.output_path}. Expected one of HA, IA, EL, AZ, FA." ) self.image_mode = (image_mode or self.available_modes[0]).upper() if self.image_mode not in self.available_modes: - sys.exit( + raise ValueError( f"Requested mode {self.image_mode} not found. Available: {self.available_modes}" ) diff --git a/src/cardiotensor/launcher/slurm_launcher.py b/src/cardiotensor/launcher/slurm_launcher.py index 43e4847..2897861 100644 --- a/src/cardiotensor/launcher/slurm_launcher.py +++ b/src/cardiotensor/launcher/slurm_launcher.py @@ -4,12 +4,20 @@ import re import shlex import subprocess -import sys import time from datetime import datetime from pathlib import Path +from cardiotensor.orientation.orientation_computation_pipeline import ( + _safe_low_memory_chunk_size, +) from cardiotensor.utils.DataReader import DataReader +from cardiotensor.utils.image_io import ( + initialize_zarr_vector_field, + normalize_vector_format, + open_zarr_vector_field, + vector_field_path, +) from cardiotensor.utils.utils import read_conf_file @@ -90,6 +98,7 @@ def submit_job_to_slurm( echo SLURM_NTASKS_PER_NODE: ${{SLURM_NTASKS_PER_NODE:-}} echo SLURM_MEM_PER_CPU: ${{SLURM_MEM_PER_CPU:-}} echo SLURM_MEM_PER_NODE: ${{SLURM_MEM_PER_NODE:-}} +echo TMPDIR: ${{TMPDIR:-/tmp}} echo ------------------------------------------------------ IMAGES_PER_JOB={chunk_size} @@ -134,12 +143,10 @@ def submit_job_to_slurm( check=True, ) except subprocess.CalledProcessError as exc: - print(f"⚠️ Failed to submit SLURM script {job_filename}", flush=True) - if exc.stdout: - print(exc.stdout, flush=True) - if exc.stderr: - print(exc.stderr, flush=True) - sys.exit(1) + details = (exc.stderr or exc.stdout or str(exc)).strip() + raise RuntimeError( + f"Failed to submit SLURM script {job_filename}: {details}" + ) from exc stdout = result.stdout.strip() match = re.search(r"Submitted batch job\s+(\d+)", stdout) @@ -186,11 +193,7 @@ def slurm_launcher( dry_run : bool If True, generate scripts but do not submit jobs. """ - try: - params = read_conf_file(conf_file_path) - except Exception as err: - print(f"⚠️ Error reading parameter file '{conf_file_path}': {err}", flush=True) - sys.exit(1) + params = read_conf_file(conf_file_path) volume_path = params.get("IMAGES_PATH", "") output_dir = params.get("OUTPUT_PATH", "./output") @@ -198,10 +201,13 @@ def slurm_launcher( angle_mode = str(params.get("ANGLE_MODE", "ha_ia")).strip().lower() write_angles = bool(params.get("WRITE_ANGLES", True)) write_vectors = bool(params.get("WRITE_VECTORS", False)) + vector_format = normalize_vector_format(params.get("VECTOR_FORMAT", "zarr")) + projected = bool(params.get("PROJECTED_ANGLES", False)) + low_memory = bool(params.get("LOW_MEMORY", False)) is_test = bool(params.get("TEST", False)) if is_test: - sys.exit( + raise ValueError( "Test mode is enabled. Disable TEST in the configuration for SLURM runs." ) @@ -232,6 +238,8 @@ def slurm_launcher( total_slices = int(data_reader.shape[0]) if total_slices <= 0: raise ValueError(f"Dataset at {volume_path} contains no slices.") + if write_vectors and vector_format == "zarr" and not dry_run: + initialize_zarr_vector_field(output_dir, tuple(data_reader.shape)) first = max(0, int(start_index)) last_exclusive = total_slices if end_index is None else int(end_index) @@ -244,6 +252,48 @@ def slurm_launcher( ) window_len = last_exclusive - first + if low_memory: + sigma = float(params.get("SIGMA", 1.0)) + rho = float(params.get("RHO", 3.0)) + truncate = float(params.get("TRUNCATE", 4.0)) + vertical_padding = params.get("VERTICAL_PADDING", None) + padding = math.ceil( + vertical_padding + if vertical_padding is not None + else int(sigma * truncate + 0.5) + int(rho * truncate + 0.5) + ) + requested_chunk = min(chunk_size, window_len) + safe_chunk, requested_peak, safe_peak = _safe_low_memory_chunk_size( + requested_chunk, + total_slices, + data_reader.shape[-2], + data_reader.shape[-1], + data_reader.dtype, + padding=padding, + include_eigenvalues=write_angles, + has_mask=params.get("MASK_PATH") is not None, + write_angles=write_angles, + sigma=sigma, + rho=rho, + truncate=truncate, + available_memory_bytes=mem_gb * 1024**3, + cpu_count=cpus_per_task, + ) + if safe_chunk < requested_chunk: + print( + "LOW_MEMORY adjusted SLURM chunk_size from " + f"{requested_chunk} to {safe_chunk} for {mem_gb} GiB: " + f"conservative peak {requested_peak / 1024**3:.2f} GiB -> " + f"{safe_peak / 1024**3:.2f} GiB.", + flush=True, + ) + chunk_size = safe_chunk + else: + print( + f"LOW_MEMORY validated chunk_size={chunk_size}: conservative peak " + f"{safe_peak / 1024**3:.2f} GiB for {mem_gb} GiB requested.", + flush=True, + ) print( f"Processing slice window [{first}, {last_exclusive}) " f"(len={window_len}) out of 0..{total_slices}", @@ -269,6 +319,31 @@ def build_intervals( return out intervals = build_intervals(first, last_exclusive, chunk_size) + all_intervals = intervals + intervals = [ + interval + for interval in all_intervals + if not is_chunk_done( + output_dir, + *interval, + output_format=output_format, + angle_mode=angle_mode, + write_angles=write_angles, + write_vectors=write_vectors, + vector_format=vector_format, + projected=projected, + ) + ] + skipped_jobs = len(all_intervals) - len(intervals) + print( + f"Output check: {skipped_jobs} complete job(s) skipped, " + f"{len(intervals)} job(s) needed", + flush=True, + ) + if not intervals: + print("All requested outputs already exist. Nothing to submit.", flush=True) + return + n_jobs_total = len(intervals) print( f"Splitting data into {n_jobs_total} jobs of up to {chunk_size} slices each", @@ -276,10 +351,16 @@ def build_intervals( ) max_tasks_per_array = 999 - batched = [ - intervals[i : i + max_tasks_per_array] - for i in range(0, n_jobs_total, max_tasks_per_array) - ] + batched: list[list[tuple[int, int]]] = [] + for interval in intervals: + if ( + not batched + or len(batched[-1]) >= max_tasks_per_array + or batched[-1][-1][1] != interval[0] + ): + batched.append([interval]) + else: + batched[-1].append(interval) print( f"Launching {len(batched)} array batch(es) " f"(tasks per batch: {[len(batch) for batch in batched]})", @@ -324,6 +405,8 @@ def build_intervals( angle_mode=angle_mode, write_angles=write_angles, write_vectors=write_vectors, + vector_format=vector_format, + projected=projected, ) elapsed = time.time() - start_t @@ -371,32 +454,29 @@ def monitor_job_output( angle_mode: str = "ha_ia", write_angles: bool = True, write_vectors: bool = False, + vector_format: str = "zarr", poll_interval_sec: int = 60, + projected: bool = False, ) -> None: """ - Monitor output progress for the requested index range. + Monitor every requested output for the requested index range. - The monitor tracks one representative output per slice: - - HA or AZ image when `write_angles=True` - - eigen_vec when only vectors are written + Completion requires both angle outputs, FA, and vectors when enabled. """ total_images = end_index_exclusive - start_index if total_images <= 0: return - if write_angles: - if angle_mode == "az_el": - mode_prefix = "AZ" - else: - mode_prefix = "HA" - folder = os.path.join(output_directory, mode_prefix) - extension = output_format - prefix = mode_prefix - elif write_vectors: - folder = os.path.join(output_directory, "eigen_vec") - extension = "npy" - prefix = "eigen_vec" + vector_format = normalize_vector_format(vector_format) + mode = angle_mode.strip().lower() + if mode == "az_el": + angle_names = ("AZ", "EL") + elif mode == "ha_ia": + angle_names = ("HA_projected", "IA_projected") if projected else ("HA", "IA") else: + raise ValueError("ANGLE_MODE must be 'ha_ia' or 'az_el'") + + if not write_angles and not write_vectors: print( "No angle/vector outputs requested; skipping monitor (nothing to track).", flush=True, @@ -404,36 +484,68 @@ def monitor_job_output( return print( - f"Monitoring outputs in {folder} for range [{start_index}, {end_index_exclusive})", + f"Monitoring all requested outputs for range " + f"[{start_index}, {end_index_exclusive})", flush=True, ) + def count_completed() -> dict[str, int]: + counts = {} + if write_angles: + for name in (*angle_names, "FA"): + counts[name] = _count_files_in_range( + os.path.join(output_directory, name), + name, + output_format, + start_index, + end_index_exclusive, + ) + + if write_vectors: + if vector_format == "zarr": + store = open_zarr_vector_field(output_directory) + counts["Zarr vectors"] = int( + store.completed_range(start_index, end_index_exclusive).sum() + ) + else: + counts["NPY vectors"] = _count_files_in_range( + str(vector_field_path(output_directory, "npy")), + "eigen_vec", + "npy", + start_index, + end_index_exclusive, + ) + return counts + start_t = time.time() - prev_count = _count_files_in_range( - folder, prefix, extension, start_index, end_index_exclusive - ) + previous_total = None while True: - current_count = _count_files_in_range( - folder, prefix, extension, start_index, end_index_exclusive + counts = count_completed() + print( + " | ".join( + f"{name}: {min(count, total_images)}/{total_images}" + for name, count in counts.items() + ), + flush=True, ) - processed = current_count - remaining = max(total_images - processed, 0) - print(f"{processed}/{total_images} processed", flush=True) - if processed >= total_images: + if all(count >= total_images for count in counts.values()): break - delta = current_count - prev_count - if delta > 0: + current_total = sum(counts.values()) + target_total = total_images * len(counts) + remaining = max(target_total - current_total, 0) + delta = 0 if previous_total is None else max(current_total - previous_total, 0) + if delta: rate_per_min = delta * (60.0 / max(poll_interval_sec, 1)) eta_min = remaining / rate_per_min if rate_per_min > 0 else float("inf") print( - f"{delta} new files in last {poll_interval_sec}s. " + f"{delta} new outputs in last {poll_interval_sec}s. " f"Approx. {eta_min:.2f} minutes remaining.", flush=True, ) - prev_count = current_count + previous_total = current_total print(f"Elapsed time (s): {time.time() - start_t:.1f}", flush=True) print(f"Waiting {poll_interval_sec} seconds...\n", flush=True) time.sleep(poll_interval_sec) @@ -445,23 +557,47 @@ def is_chunk_done( end: int, output_format: str = "jp2", angle_mode: str = "ha_ia", + write_angles: bool = True, + write_vectors: bool = False, + vector_format: str = "zarr", + projected: bool = False, ) -> bool: - """ - Check if all output files for a given chunk [start, end) are already present. - """ + """Return True when every requested output exists for [start, end).""" ext = output_format.lstrip(".") mode = angle_mode.lower().strip() if mode == "az_el": angle1, angle2 = "AZ", "EL" else: - angle1, angle2 = "HA", "IA" + angle1, angle2 = ( + ("HA_projected", "IA_projected") if projected else ("HA", "IA") + ) + + vector_format = normalize_vector_format(vector_format) + zarr_completed = None + if write_vectors and vector_format == "zarr": + try: + zarr_completed = open_zarr_vector_field(output_dir).completed_range( + start, end + ) + except (FileNotFoundError, KeyError, ValueError): + return False for idx in range(start, end): - expected_files = [ - f"{output_dir}/{angle1}/{angle1}_{idx:06d}.{ext}", - f"{output_dir}/{angle2}/{angle2}_{idx:06d}.{ext}", - f"{output_dir}/FA/FA_{idx:06d}.{ext}", - ] + expected_files = [] + if write_angles: + expected_files.extend( + [ + f"{output_dir}/{angle1}/{angle1}_{idx:06d}.{ext}", + f"{output_dir}/{angle2}/{angle2}_{idx:06d}.{ext}", + f"{output_dir}/FA/FA_{idx:06d}.{ext}", + ] + ) + if write_vectors and vector_format == "npy": + expected_files.append( + str(vector_field_path(output_dir, "npy") / f"eigen_vec_{idx:06d}.npy") + ) if not all(os.path.exists(path) for path in expected_files): return False + if zarr_completed is not None and not zarr_completed[idx - start]: + return False return True diff --git a/src/cardiotensor/orientation/orientation_computation_functions.py b/src/cardiotensor/orientation/orientation_computation_functions.py index 4341d8b..667f98a 100644 --- a/src/cardiotensor/orientation/orientation_computation_functions.py +++ b/src/cardiotensor/orientation/orientation_computation_functions.py @@ -1,5 +1,7 @@ +import mmap import os import warnings +from pathlib import Path import matplotlib.pyplot as plt import numpy as np @@ -140,8 +142,10 @@ def calculate_structure_tensor( devices: list[str] | None = None, block_size: int = 200, use_gpu: bool = False, - dtype: type = np.float32, # Default to np.float64 -) -> tuple[np.ndarray, np.ndarray]: + dtype: type = np.float32, + return_eigenvalues: bool = True, + memmap_dir: str | os.PathLike[str] | None = None, +) -> tuple[np.ndarray | None, np.ndarray]: """ Calculates the structure tensor of a volume. @@ -152,9 +156,13 @@ def calculate_structure_tensor( devices (Optional[list[str]]): List of devices for parallel processing (e.g., ['cpu', 'cuda:0']). block_size (int): Size of the blocks for processing. Default is 200. use_gpu (bool): If True, uses GPU for calculations. Default is False. + dtype: Output dtype for eigenvalues and eigenvectors. + return_eigenvalues: If False, do not allocate the eigenvalue volume. + memmap_dir: Optional scratch directory for disk-backed tensor outputs. Returns: - tuple[np.ndarray, np.ndarray]: Eigenvalues and eigenvectors of the structure tensor. + tuple[np.ndarray | None, np.ndarray]: Eigenvalues, when requested, and + eigenvectors of the structure tensor. """ # Filter or ignore specific warnings warnings.filterwarnings("ignore", category=RuntimeWarning) @@ -198,8 +206,45 @@ class TqdmTotal(tqdm): def update_with_total(self, n=1, total=None): if total is not None: self.total = total + if ( + memmap_dir is not None + and total is not None + and (n % 16 == 0 or n == total) + ): + for output in (eigenvector_output, eigenvalue_output): + if isinstance(output, np.memmap): + output.flush() + mapped = getattr(output, "_mmap", None) + if ( + n == total + and mapped is not None + and hasattr(mapped, "madvise") + and hasattr(mmap, "MADV_DONTNEED") + ): + mapped.madvise(mmap.MADV_DONTNEED) return self.update(1) + eigenvector_output: type | np.memmap = dtype + eigenvalue_output: type | np.memmap | None = dtype if return_eigenvalues else None + if memmap_dir is not None: + scratch_dir = Path(memmap_dir) + scratch_dir.mkdir(parents=True, exist_ok=True) + output_shape = (3, *volume.shape) + eigenvector_output = np.lib.format.open_memmap( + scratch_dir / "eigenvectors.npy", + mode="w+", + dtype=dtype, + shape=output_shape, + ) + if return_eigenvalues: + eigenvalue_output = np.lib.format.open_memmap( + scratch_dir / "eigenvalues.npy", + mode="w+", + dtype=dtype, + shape=output_shape, + ) + print("--- Flushing memory-mapped tensor outputs every 16 blocks") + def run_structure_tensor(selected_devices: list[str]): with TqdmTotal(desc="Computing structure tensors", unit="block") as t: return parallel_structure_tensor_analysis( @@ -210,8 +255,8 @@ def run_structure_tensor(selected_devices: list[str]): block_size=block_size, truncate=truncate, structure_tensor=None, - eigenvectors=dtype, - eigenvalues=dtype, + eigenvectors=eigenvector_output, + eigenvalues=eigenvalue_output, progress_callback_fn=t.update_with_total, ) @@ -230,6 +275,8 @@ def run_structure_tensor(selected_devices: list[str]): print("Structure tensor computation completed\n") # vec has shape =(3,z,y,x) in the order of (x,y,z) + if vec is None: + raise RuntimeError("Structure tensor calculation did not return eigenvectors") return val, vec @@ -240,6 +287,8 @@ def remove_padding( vec: np.ndarray, padding_start: int, padding_end: int, + *, + copy: bool = False, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """ Removes padding from the volume, eigenvalues, and eigenvectors. @@ -250,14 +299,19 @@ def remove_padding( vec (np.ndarray): The eigenvectors. padding_start (int): Padding at the start to remove. padding_end (int): Padding at the end to remove. + copy: Return compact copies instead of views into the padded arrays. Returns: - Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: Adjusted data without padding. + Tuple[np.ndarray, np.ndarray, np.ndarray]: Adjusted data without padding. """ array_end = vec.shape[1] - padding_end volume = volume[padding_start:array_end, :, :] vec = vec[:, padding_start:array_end, :, :] val = val[:, padding_start:array_end, :, :] + if copy: + volume = volume.copy(order="C") + vec = vec.copy(order="C") + val = val.copy(order="C") return volume, val, vec @@ -327,21 +381,10 @@ def rotate_vectors_to_new_axis( + np.dot(skew_matrix, skew_matrix) * ((1.0 - cos_angle) / (axis_length**2)) ) - # Flatten the vector field to shape (3, N) - vec_reshaped = np.reshape(vector_field_slice, (3, -1)) - - # Normalize safely - norms = np.linalg.norm(vec_reshaped, axis=0) - nonzero_mask = norms > 0 - vec_reshaped[:, nonzero_mask] /= norms[nonzero_mask] - - # Rotate - rotated_vecs = np.dot(rotation_matrix, vec_reshaped) - - # Reshape back - rotated_vecs = rotated_vecs.reshape(vector_field_slice.shape) - - return rotated_vecs + vector_dtype = np.result_type(vector_field_slice.dtype, np.float32) + rotation_matrix = rotation_matrix.astype(vector_dtype, copy=False) + vectors = np.asarray(vector_field_slice, dtype=vector_dtype) + return np.einsum("ij,jyx->iyx", rotation_matrix, vectors, optimize=True) def orient_vectors_y_positive(vector_field_slice: np.ndarray) -> np.ndarray: @@ -398,42 +441,48 @@ def compute_helical_and_intrusion_angles( Returns: Tuple[np.ndarray, np.ndarray]: Helical and intrusion angle arrays. """ - center = center_point[0:2] # Replace with actual values rows, cols = vector_field_2d.shape[1:3] - - reshaped_vector_field = np.reshape(vector_field_2d, (3, -1)) - - center_x, center_y = center[0], center[1] - - X, Y = np.meshgrid(np.arange(cols) - center_x, np.arange(rows) - center_y) - - theta = -np.arctan2(Y.flatten(), X.flatten()) - cos_angle = np.cos(theta) - sin_angle = np.sin(theta) - - # Change coordinate system to cylindrical - rotated_vector_field = np.copy(reshaped_vector_field) - rotated_vector_field[0, :] = ( - cos_angle * reshaped_vector_field[0, :] - - sin_angle * reshaped_vector_field[1, :] + vector_dtype = np.result_type(vector_field_2d.dtype, np.float32) + vectors = np.asarray(vector_field_2d, dtype=vector_dtype) + vx, vy, vz = vectors + + center_x, center_y = center_point[:2] + x = np.arange(cols, dtype=vector_dtype) - vector_dtype.type(center_x) + y = np.arange(rows, dtype=vector_dtype) - vector_dtype.type(center_y) + x = x[None, :] + y = y[:, None] + + # Compute the cylindrical components directly. This avoids full-size X/Y, + # theta, sine, cosine, and three-component intermediate arrays. + radius = np.hypot(x, y) + radial_component = np.multiply(vx, x) + radial_component += vy * y + circumferential_component = np.multiply(vy, x) + circumferential_component -= vx * y + + away_from_center = radius > 0 + np.divide( + radial_component, + radius, + out=radial_component, + where=away_from_center, ) - rotated_vector_field[1, :] = ( - sin_angle * reshaped_vector_field[0, :] - + cos_angle * reshaped_vector_field[1, :] + np.divide( + circumferential_component, + radius, + out=circumferential_component, + where=away_from_center, ) + at_center = ~away_from_center + radial_component[at_center] = vx[at_center] + circumferential_component[at_center] = vy[at_center] + del radius, away_from_center, at_center - # Reshape rotated vector field to original image dimensions - reshaped_rotated_vector_field = np.zeros((3, rows, cols)) - for i in range(3): - reshaped_rotated_vector_field[i] = rotated_vector_field[i].reshape(rows, cols) - - reshaped_rotated_vector_field = orient_vectors_y_positive( - reshaped_rotated_vector_field - ) - - radial_component = reshaped_rotated_vector_field[0, :, :] - circumferential_component = reshaped_rotated_vector_field[1, :, :] - longitudinal_component = reshaped_rotated_vector_field[2, :, :] + longitudinal_component = np.array(vz, copy=True) + flip_mask = circumferential_component < 0 + radial_component[flip_mask] *= -1 + circumferential_component[flip_mask] *= -1 + longitudinal_component[flip_mask] *= -1 # Calculate helical and intrusion angles if projected: @@ -448,8 +497,8 @@ def compute_helical_and_intrusion_angles( radial_component, np.hypot(circumferential_component, longitudinal_component), ) - helical_angle = np.rad2deg(helical_angle) - intrusion_angle = np.rad2deg(intrusion_angle) + np.rad2deg(helical_angle, out=helical_angle) + np.rad2deg(intrusion_angle, out=intrusion_angle) return helical_angle, intrusion_angle @@ -461,13 +510,15 @@ def compute_azimuth_and_elevation( Azimuth = angle in XY plane from +X toward +Y, in degrees [-180, 180] Elevation = unsigned angle from XY plane toward +Z, in degrees [0, 90] """ - oriented = orient_vectors_z_positive(vector_field_2d) - vx = oriented[0, :, :] - vy = oriented[1, :, :] - vz = oriented[2, :, :] - - az = np.rad2deg(np.arctan2(vy, vx)) # [-180, 180] - el = np.rad2deg(np.arctan2(vz, np.hypot(vx, vy))) # [0, 90] + vx, vy, vz = vector_field_2d + az = np.arctan2(vy, vx) + flip_mask = vz < 0 + az[flip_mask] = np.arctan2(-vy[flip_mask], -vx[flip_mask]) + np.rad2deg(az, out=az) + + horizontal = np.hypot(vx, vy) + el = np.arctan2(np.abs(vz), horizontal) + np.rad2deg(el, out=el) return az, el diff --git a/src/cardiotensor/orientation/orientation_computation_pipeline.py b/src/cardiotensor/orientation/orientation_computation_pipeline.py index 65114c8..7a27905 100644 --- a/src/cardiotensor/orientation/orientation_computation_pipeline.py +++ b/src/cardiotensor/orientation/orientation_computation_pipeline.py @@ -1,10 +1,13 @@ +import mmap import math import os -import sys +import shutil import time import warnings from collections.abc import Sequence from multiprocessing.pool import ThreadPool +from pathlib import Path +from tempfile import TemporaryDirectory import numpy as np from alive_progress import alive_bar @@ -21,13 +24,22 @@ compute_helical_and_intrusion_angles, interpolate_points, plot_images, - remove_padding, rotate_vectors_to_new_axis, write_images, ) from cardiotensor.utils.DataReader import DataReader -from cardiotensor.utils.image_io import write_vector_field -from cardiotensor.utils.utils import get_available_cpu_count, remove_corrupted_files +from cardiotensor.utils.image_io import ( + ZarrVectorFieldStore, + initialize_zarr_vector_field, + normalize_vector_format, + vector_field_path, + write_vector_field, +) +from cardiotensor.utils.utils import ( + get_available_cpu_count, + get_available_memory_bytes, + remove_corrupted_files, +) # --- small helpers --- @@ -49,6 +61,244 @@ def _resolve_colormap(colormap: str | None): ) from err +def _centerline_neighborhood( + center_line: np.ndarray, global_slice_idx: int, buffer: int = 5 +) -> np.ndarray: + """Return a chunk-independent centerline window around one global slice.""" + if not 0 <= global_slice_idx < len(center_line): + raise IndexError( + f"global_slice_idx={global_slice_idx} is outside centerline length " + f"{len(center_line)}" + ) + start = max(0, global_slice_idx - buffer) + end = min(len(center_line), global_slice_idx + buffer + 1) + return center_line[start:end] + + +def _normalize_vectors_in_place(vector_field: np.ndarray) -> np.ndarray: + """Normalize one Z slice at a time, preserving zero-filled masked voxels.""" + for z in range(vector_field.shape[1]): + vector_slice = vector_field[:, z] + norms = np.linalg.norm(vector_slice, axis=0) + np.divide(vector_slice, norms[None], out=vector_slice, where=norms[None] > 0) + if isinstance(vector_field, np.memmap) and ( + (z + 1) % 4 == 0 or z + 1 == vector_field.shape[1] + ): + vector_field.flush() + mapped = getattr(vector_field, "_mmap", None) + if ( + mapped is not None + and hasattr(mapped, "madvise") + and hasattr(mmap, "MADV_DONTNEED") + ): + mapped.madvise(mmap.MADV_DONTNEED) + return vector_field + + +def _estimate_low_memory_peak_bytes( + volume_shape: tuple[int, int, int], + input_dtype: np.dtype | type, + *, + include_eigenvalues: bool, + has_mask: bool, + write_angles: bool, + sigma: float, + rho: float, + truncate: float, + cpu_count: int | None = None, + include_loaded_inputs: bool = True, +) -> int: + """Conservative peak estimate for a disk-backed tensor calculation.""" + depth, height, width = volume_shape + voxels = math.prod(volume_shape) + output_components = 6 if include_eigenvalues else 3 + mapped_outputs = voxels * output_components * np.dtype(np.float32).itemsize + input_bytes = voxels * np.dtype(input_dtype).itemsize + mask_bytes = voxels if has_mask else 0 + + # structure-tensor uses float64 blocks with overlap and several component arrays. + block_size = 200 + radius = 2 * int(max(sigma, rho) * truncate + 0.5) + block_shape = ( + min(depth, block_size + 2 * radius), + min(height, block_size + 2 * radius), + min(width, block_size + 2 * radius), + ) + block_count = math.prod(math.ceil(size / block_size) for size in volume_shape) + workers = min( + max(1, get_available_cpu_count() if cpu_count is None else cpu_count), + block_count, + ) + worker_temporaries = workers * math.prod(block_shape) * 18 * 8 + + # At least one angle-writing worker must fit. Vector-only writing is much lighter. + slice_work = height * width * (128 if write_angles else 16) + reserve = 512 * 1024**2 + + if include_loaded_inputs: + # Mask resizing can temporarily hold the resized mask and its fitted copy. + load_phase = input_bytes + 2 * mask_bytes + tensor_phase = mapped_outputs + input_bytes + mask_bytes + worker_temporaries + post_phase = mapped_outputs + mask_bytes + slice_work + else: + # The runtime call is made after input/mask allocations, so available memory + # already excludes them. Estimate only additional allocations from this point. + load_phase = 0 + tensor_phase = mapped_outputs + worker_temporaries + post_phase = mapped_outputs + slice_work + + return max(load_phase, tensor_phase, post_phase) + reserve + + +def _safe_low_memory_chunk_size( + requested_chunk: int, + total_slices: int, + height: int, + width: int, + input_dtype: np.dtype | type, + *, + padding: int, + include_eigenvalues: bool, + has_mask: bool, + write_angles: bool, + sigma: float, + rho: float, + truncate: float, + available_memory_bytes: int | None = None, + cpu_count: int | None = None, +) -> tuple[int, int, int]: + """Return safe payload slices plus requested/safe peak byte estimates.""" + if requested_chunk <= 0: + raise ValueError("N_CHUNK must be > 0") + available = ( + get_available_memory_bytes() + if available_memory_bytes is None + else available_memory_bytes + ) + budget = int(available * 0.82) + requested_chunk = min(requested_chunk, total_slices) + + def estimate(payload: int) -> int: + padded_depth = min(total_slices, payload + 2 * padding) + return _estimate_low_memory_peak_bytes( + (padded_depth, height, width), + input_dtype, + include_eigenvalues=include_eigenvalues, + has_mask=has_mask, + write_angles=write_angles, + sigma=sigma, + rho=rho, + truncate=truncate, + cpu_count=cpu_count, + ) + + requested_peak = estimate(requested_chunk) + safe_chunk = requested_chunk + safe_peak = requested_peak + while safe_chunk > 1 and safe_peak > budget: + safe_chunk -= 1 + safe_peak = estimate(safe_chunk) + + if safe_peak > budget: + raise MemoryError( + "Even a one-slice low-memory chunk exceeds the conservative memory " + f"budget ({safe_peak / 1024**3:.2f} GiB estimated, " + f"{available / 1024**3:.2f} GiB available). Request more memory." + ) + return safe_chunk, requested_peak, safe_peak + + +def _memory_aware_worker_count( + num_slices: int, + height: int, + width: int, + write_angles: bool, + *, + available_memory_bytes: int | None = None, + cpu_count: int | None = None, +) -> int: + """Choose a conservative worker count from slice size and free memory.""" + if num_slices <= 0: + return 1 + available_memory_bytes = ( + get_available_memory_bytes() + if available_memory_bytes is None + else available_memory_bytes + ) + cpu_count = get_available_cpu_count() if cpu_count is None else cpu_count + + # Angles create rotated vectors, three outputs, and temporary scalar arrays. + # Vector-only workers mainly need codec and file-write buffers. + bytes_per_pixel = 128 if write_angles else 12 + bytes_per_worker = max(1, height * width * bytes_per_pixel) + memory_workers = max(1, (available_memory_bytes // 2) // bytes_per_worker) + return max(1, min(num_slices, cpu_count, 32, memory_workers)) + + +def _ensure_structure_tensor_memory( + volume_shape: tuple[int, ...], + *, + include_eigenvalues: bool = True, + low_memory: bool = False, + available_memory_bytes: int | None = None, + input_dtype: np.dtype | type = np.float32, + has_mask: bool = False, + write_angles: bool = False, + sigma: float = 1.0, + rho: float = 3.0, + truncate: float = 4.0, +) -> None: + """Fail before tensor allocation when a job cannot hold its output arrays.""" + voxel_count = math.prod(volume_shape) + output_components = 6 if include_eigenvalues else 3 + output_bytes = voxel_count * output_components * np.dtype(np.float32).itemsize + available_memory_bytes = ( + get_available_memory_bytes() + if available_memory_bytes is None + else available_memory_bytes + ) + + output_gib = output_bytes / 1024**3 + available_gib = available_memory_bytes / 1024**3 + print( + f"Structure tensor outputs need {output_gib:.2f} GiB; " + f"{available_gib:.2f} GiB is currently available" + ) + + if low_memory: + additional_bytes = _estimate_low_memory_peak_bytes( + tuple(volume_shape), + input_dtype, + include_eigenvalues=include_eigenvalues, + has_mask=has_mask, + write_angles=write_angles, + sigma=sigma, + rho=rho, + truncate=truncate, + include_loaded_inputs=False, + ) + print( + "Low-memory mode: tensor outputs will use temporary memory-mapped files; " + f"conservative additional working set is {additional_bytes / 1024**3:.2f} GiB" + ) + if additional_bytes > available_memory_bytes: + raise MemoryError( + f"This padded low-memory chunk still needs approximately " + f"{additional_bytes / 1024**3:.2f} GiB beyond currently resident " + f"inputs, but only {available_gib:.2f} GiB is available. Reduce " + "N_CHUNK (or --chunk_size for SLURM) or request more memory." + ) + return + + # Keep 30% for block temporaries, Python, codecs, and image-writing buffers. + if output_bytes > available_memory_bytes * 0.7: + raise MemoryError( + f"Structure tensor outputs need {output_gib:.2f} GiB, leaving " + f"insufficient working memory from {available_gib:.2f} GiB available. " + "Reduce N_CHUNK (or --chunk_size for SLURM) or request more memory." + ) + + def check_already_processed( output_dir: str, start_index: int, @@ -56,6 +306,8 @@ def check_already_processed( write_vectors: bool, write_angles: bool, output_format: str, + vector_format: str = "zarr", + zarr_store: ZarrVectorFieldStore | None = None, angle_names: tuple[str, str] = ("HA", "IA"), fa_name: str = "FA", extra_expected: Sequence[str] | None = None, @@ -72,7 +324,7 @@ def check_already_processed( end_index : int Last global slice index to check (exclusive). write_vectors : bool - If True, expect eigenvector .npy files (e.g., eigen_vec_{idx:06d}.npy). + If True, expect completed vector output in the selected backend. write_angles : bool If True, expect angle images for angle_names[0], angle_names[1], and FA. output_format : str @@ -105,6 +357,13 @@ def check_already_processed( # Prepare optional extras extra_expected = tuple(extra_expected or ()) + vector_format = normalize_vector_format(vector_format) + zarr_completed = None + if write_vectors and vector_format == "zarr": + if zarr_store is None: + print(f"Missing Zarr vector store in {output_dir}") + return False + zarr_completed = zarr_store.completed_range(start_index, end_index) for idx in range(start_index, end_index): expected_files = [] @@ -117,20 +376,34 @@ def check_already_processed( os.path.join(output_dir, fa_name, f"{fa_name}_{idx:06d}.{ext}"), ] - if write_vectors: + if write_vectors and vector_format == "npy": expected_files.append( - os.path.join(output_dir, "eigen_vec", f"eigen_vec_{idx:06d}.npy") + str(vector_field_path(output_dir, "npy") / f"eigen_vec_{idx:06d}.npy") ) + elif write_vectors and not zarr_completed[idx - start_index]: + print(f"Missing or incomplete Zarr vector output for slice {idx:06d}") + return False # User-specified extras, if any for tmpl in extra_expected: expected_files.append(tmpl.format(idx=f"{idx:06d}", ext=ext)) # Remove small/corrupted files before checking (function defined elsewhere) - remove_corrupted_files(expected_files) + corrupted_files = remove_corrupted_files(expected_files) + if corrupted_files: + print( + f"Corrupted output file(s) found for slice {idx:06d}: " + + ", ".join(corrupted_files) + ) # If any file is missing, we need to process - if not all(os.path.exists(p) for p in expected_files): + missing_files = [path for path in expected_files if not os.path.exists(path)] + if missing_files: + if not corrupted_files: + print( + f"Missing output file(s) for slice {idx:06d}: " + + ", ".join(missing_files) + ) return False print(f"Checking already processed files: all expected files exist in {output_dir}") @@ -162,6 +435,9 @@ def compute_orientation( colormap_angle1: str | None = None, colormap_angle2: str | None = None, projected: bool = False, + vector_format: str = "zarr", + low_memory: bool = False, + low_memory_dir: str | os.PathLike[str] | None = None, ) -> None: """ Compute the orientation for a volume dataset. @@ -178,6 +454,10 @@ def compute_orientation( axis_points: 3D points defining LV axis for cylindrical coordinates. vertical_padding: Padding slices for tensor computation. write_vectors: Whether to save eigenvectors. Ignored in test mode. + vector_format: Vector storage backend, either "npy" or "zarr". + low_memory: Store temporary tensor outputs as memory-mapped files. + low_memory_dir: Optional scratch root. Defaults to + OUTPUT_PATH/.cardiotensor_scratch, never the system /tmp directory. write_angles: Whether to save HA/IA/FA maps. projected: If True in ha_ia mode, write projected HA/IA legacy maps. use_gpu: Use GPU acceleration for tensor computation. @@ -194,11 +474,12 @@ def compute_orientation( # --- Sanity checks --- if sigma > rho: raise ValueError("sigma must be <= rho") + if not write_vectors and not write_angles: + raise ValueError("At least one of write_vectors or write_angles must be True") + vector_format = normalize_vector_format(vector_format) if angle_mode.lower() == "ha_ia": - angle_names = ( - ("HA_projected", "IA_projected") if projected else ("HA", "IA") - ) + angle_names = ("HA_projected", "IA_projected") if projected else ("HA", "IA") elif angle_mode.lower() == "az_el": angle_names = ("AZ", "EL") else: @@ -214,11 +495,10 @@ def compute_orientation( ) write_vectors = False - projected_status = ( - projected if angle_mode.lower().strip() == "ha_ia" else "[n/a]" - ) + projected_status = projected if angle_mode.lower().strip() == "ha_ia" else "[n/a]" - print(f""" + print( + f""" Parameters: - Volume path: {volume_path} - Mask path: {mask_path or "[None]"} @@ -231,13 +511,16 @@ def compute_orientation( - Angle mode: {angle_mode} -> {angle_names[0]}, {angle_names[1]} - Projected HA/IA:{projected_status} - Write vectors: {write_vectors} + - Vector format: {vector_format} + - Low memory: {low_memory} - Use GPU: {use_gpu} - Test mode: {is_test} - Show quiver: {show_quiver} - Colormap: {colormap or "[default]"} - Colormap angle1:{colormap_angle1 or "[default]"} - Colormap angle2:{colormap_angle2 or "[default]"} - """) + """ + ) print("\n" + "-" * 40) print("READING VOLUME INFORMATION") @@ -246,6 +529,11 @@ def compute_orientation( print(f"Volume path: {volume_path}") data_reader = DataReader(volume_path) + vector_store = None + if write_vectors and vector_format == "zarr": + vector_store = initialize_zarr_vector_field( + output_dir, tuple(data_reader.shape) + ) if end_index is None: end_index = data_reader.shape[0] @@ -262,6 +550,8 @@ def compute_orientation( write_vectors, write_angles, output_format, + vector_format=vector_format, + zarr_store=vector_store, angle_names=angle_names, ) and not is_test @@ -277,7 +567,7 @@ def compute_orientation( print("CALCULATE PADDING START AND ENDING INDEXES\n") if vertical_padding is None: - vertical_padding = truncate * rho + 0.5 + vertical_padding = int(sigma * truncate + 0.5) + int(rho * truncate + 0.5) padding_start = padding_end = math.ceil(vertical_padding) if not is_test: @@ -286,8 +576,13 @@ def compute_orientation( if padding_end > (data_reader.shape[0] - end_index): padding_end = data_reader.shape[0] - end_index if is_test: - if n_slice_test > data_reader.shape[0]: - sys.exit("Error: n_slice_test > number of images") + if n_slice_test is None: + raise ValueError("n_slice_test is required in test mode") + if not 0 <= n_slice_test < data_reader.shape[0]: + raise ValueError( + f"n_slice_test={n_slice_test} is outside volume bounds " + f"[0, {data_reader.shape[0]})" + ) print(f"Padding start, Padding end : {padding_start}, {padding_end}") start_index_padded, end_index_padded = adjust_start_end_index( @@ -305,10 +600,9 @@ def compute_orientation( print("\n---------------------------------") print("LOAD DATASET\n") - volume = data_reader.load_volume(start_index_padded, end_index_padded).astype( - "float32" - ) + volume = data_reader.load_volume(start_index_padded, end_index_padded) print(f"Loaded volume shape {volume.shape}") + invalid_mask = None if mask_path is not None: print("\n---------------------------------") print("LOAD MASK\n") @@ -318,43 +612,132 @@ def compute_orientation( start_index_padded, end_index_padded, unbinned_shape=data_reader.shape ) - assert mask.shape == volume.shape, ( - f"Mask shape {mask.shape} does not match volume shape {volume.shape}" - ) + if mask.shape != volume.shape: + raise ValueError( + f"Mask shape {mask.shape} does not match volume shape {volume.shape}" + ) - volume[mask == 0] = 0 + invalid_mask = mask == 0 + volume[invalid_mask] = 0 + del mask print("\n" + "-" * 40) print("CALCULATING STRUCTURE TENSOR") print("-" * 40 + "\n") + need_eigenvalues = write_angles or is_test + _ensure_structure_tensor_memory( + tuple(volume.shape), + include_eigenvalues=need_eigenvalues, + low_memory=low_memory, + input_dtype=volume.dtype, + has_mask=invalid_mask is not None, + write_angles=write_angles, + sigma=sigma, + rho=rho, + truncate=truncate, + ) + tensor_scratch = None + if low_memory: + scratch_root = Path( + low_memory_dir or Path(output_dir) / ".cardiotensor_scratch" + ).expanduser().resolve() + scratch_root.mkdir(parents=True, exist_ok=True) + output_components = 6 if need_eigenvalues else 3 + required_scratch = ( + math.prod(volume.shape) * output_components * np.dtype(np.float32).itemsize + ) + free_scratch = shutil.disk_usage(scratch_root).free + print( + f"Tensor scratch root: {scratch_root} | required: " + f"{required_scratch / 1024**3:.2f} GiB | free: " + f"{free_scratch / 1024**3:.2f} GiB" + ) + if required_scratch > free_scratch * 0.9: + raise OSError( + f"Low-memory tensor outputs need {required_scratch / 1024**3:.2f} " + f"GiB in {scratch_root}, but only {free_scratch / 1024**3:.2f} " + "GiB is free. Choose LOW_MEMORY_DIR with more space." + ) + tensor_scratch = TemporaryDirectory( + prefix="cardiotensor_tensor_", dir=scratch_root + ) + print(f"Tensor scratch directory: {tensor_scratch.name}") t1 = time.perf_counter() # start time val, vec = calculate_structure_tensor( - volume, sigma, rho, truncate=truncate, use_gpu=use_gpu + volume, + sigma, + rho, + truncate=truncate, + use_gpu=use_gpu, + return_eigenvalues=need_eigenvalues, + memmap_dir=None if tensor_scratch is None else tensor_scratch.name, ) print(f"Vector field shape: {vec.shape}") + if not is_test: + # The intensity data is no longer needed once tensor calculation ends. + del volume + volume = None + + array_end = vec.shape[1] - padding_end + interior = slice(padding_start, array_end) + if low_memory: + vec = vec[:, interior, :, :] + if val is not None: + val = val[:, interior, :, :] + else: + cropped_vec = vec[:, interior, :, :].copy(order="C") + del vec + vec = cropped_vec + if val is not None: + cropped_val = val[:, interior, :, :].copy(order="C") + del val + val = cropped_val - if mask_path is not None: - print("Applying mask to tensors and vectors...") - - volume[mask == 0] = np.nan - val[0, :, :, :][mask == 0] = np.nan - val[1, :, :, :][mask == 0] = np.nan - val[2, :, :, :][mask == 0] = np.nan - vec[0, :, :, :][mask == 0] = np.nan - vec[1, :, :, :][mask == 0] = np.nan - vec[2, :, :, :][mask == 0] = np.nan + if is_test: + cropped_volume = volume[interior, :, :].copy(order="C") + del volume + volume = cropped_volume + print(f"Vector shape after removing padding: {vec.shape}") + if invalid_mask is not None: + mask_end = invalid_mask.shape[0] - padding_end + invalid_mask = invalid_mask[padding_start:mask_end].copy(order="C") + print("Applying mask to cropped tensors and vectors...") + if val is not None: + val[:, invalid_mask] = 0 + vec[:, invalid_mask] = 0 + for mapped_output in (val, vec): + if isinstance(mapped_output, np.memmap): + mapped_output.flush() + mapped = getattr(mapped_output, "_mmap", None) + if ( + mapped is not None + and hasattr(mapped, "madvise") + and hasattr(mmap, "MADV_DONTNEED") + ): + mapped.madvise(mmap.MADV_DONTNEED) print("Masking complete") + else: + invalid_mask = None - del mask - - volume, val, vec = remove_padding(volume, val, vec, padding_start, padding_end) - print(f"Vector shape after removing padding: {vec.shape}") - - center_line = center_line[start_index_padded:end_index_padded] + _normalize_vectors_in_place(vec) - vec = vec / np.linalg.norm(vec, axis=0) + if write_vectors and vector_format == "zarr": + vector_completed = vector_store.completed_range( + start_index, start_index + vec.shape[1] + ) + elif write_vectors: + npy_dir = vector_field_path(output_dir, "npy") + vector_completed = np.array( + [ + (npy_dir / f"eigen_vec_{idx:06d}.npy").exists() + for idx in range(start_index, start_index + vec.shape[1]) + ], + dtype=bool, + ) + else: + vector_completed = np.zeros(vec.shape[1], dtype=bool) t2 = time.perf_counter() # stop time print(f"finished calculating structure tensors in {t2 - t1} seconds") @@ -365,10 +748,14 @@ def compute_orientation( if not is_test: num_slices = vec.shape[1] - available_cpus = get_available_cpu_count() - num_workers = min(available_cpus, 32) # Avoid too many concurrent writers + num_workers = _memory_aware_worker_count( + num_slices, + height=vec.shape[2], + width=vec.shape[3], + write_angles=write_angles, + ) - print(f"Using {num_workers} threads") + print(f"Using {num_workers} memory-aware slice worker(s)") def update_bar(_): """Callback to tick the progress bar after each finished task.""" @@ -380,14 +767,15 @@ def update_bar(_): ) as bar: results = [] for z in range(num_slices): + global_slice_idx = start_index + z result = pool.apply_async( compute_slice_angles_and_anisotropy, ( z, vec[:, z, :, :], - volume[z, :, :], - np.around(center_line[z]), - val[:, z, :, :], + None, + np.around(center_line[global_slice_idx]), + None if val is None else val[:, z, :, :], center_line, output_dir, output_format, @@ -402,6 +790,11 @@ def update_bar(_): colormap_angle1, colormap_angle2, projected, + global_slice_idx, + None if invalid_mask is None else invalid_mask[z], + vector_format, + vector_store, + bool(vector_completed[z]), ), callback=update_bar, ) @@ -431,12 +824,14 @@ def update_bar(_): ) z = 0 + global_slice_idx = n_slice_test compute_slice_angles_and_anisotropy( z, vec[:, z, :, :], + # Test mode is the only path that retains the intensity image. volume[z, :, :], - np.around(center_line[z]), - val[:, z, :, :], + np.around(center_line[global_slice_idx]), + None if val is None else val[:, z, :, :], center_line, output_dir, output_format, @@ -451,6 +846,11 @@ def update_bar(_): colormap_angle1, colormap_angle2, projected, + global_slice_idx, + None if invalid_mask is None else invalid_mask[z], + vector_format, + vector_store, + bool(vector_completed[z]), ) if is_test: @@ -459,15 +859,23 @@ def update_bar(_): end_index_local = start_index + vec.shape[1] print(f"\nFinished processing slices {start_index} to {end_index_local}") print("---------------------------------\n\n") + + if tensor_scratch is not None: + if isinstance(vec, np.memmap): + vec.flush() + if isinstance(val, np.memmap): + val.flush() + del vec, val + tensor_scratch.cleanup() return def compute_slice_angles_and_anisotropy( z: int, vector_field_slice: np.ndarray, - img_slice: np.ndarray, + img_slice: np.ndarray | None, center_point: np.ndarray, - eigen_val_slice: np.ndarray, + eigen_val_slice: np.ndarray | None, center_line: np.ndarray, output_dir: str, output_format: str = "jp2", @@ -482,6 +890,11 @@ def compute_slice_angles_and_anisotropy( colormap_angle1: str | None = None, colormap_angle2: str | None = None, projected: bool = False, + global_slice_idx: int | None = None, + invalid_mask_slice: np.ndarray | None = None, + vector_format: str = "zarr", + zarr_store: ZarrVectorFieldStore | None = None, + vector_already_written: bool = False, ) -> tuple[float, float, bool]: """ Compute either HA/IA or Azimuth/Elevation plus FA for a single slice, @@ -490,9 +903,7 @@ def compute_slice_angles_and_anisotropy( # Decide angle labels and ranges based on mode mode = angle_mode.lower().strip() if mode == "ha_ia": - angle_names = ( - ("HA_projected", "IA_projected") if projected else ("HA", "IA") - ) + angle_names = ("HA_projected", "IA_projected") if projected else ("HA", "IA") angle_ranges = ((-90.0, 90.0), (-90.0, 90.0)) elif mode == "az_el": angle_names = ("AZ", "EL") @@ -501,7 +912,7 @@ def compute_slice_angles_and_anisotropy( raise ValueError("ANGLE_MODE must be 'ha_ia' or 'az_el'") ext = output_format.lstrip(".") - idx = start_index + z + idx = start_index + z if global_slice_idx is None else global_slice_idx shared_colormap = _resolve_colormap(colormap) colormap_angle = _resolve_colormap(colormap_angle1) or shared_colormap colormap_angle2 = _resolve_colormap(colormap_angle2) or shared_colormap @@ -510,7 +921,7 @@ def compute_slice_angles_and_anisotropy( colormap_angle = helix_angle_cmap if colormap_angle2 is None: colormap_angle2 = plt.cm.viridis - rows, cols = img_slice.shape[:2] + rows, cols = vector_field_slice.shape[1:] center_x, center_y = float(center_point[0]), float(center_point[1]) if not (0 <= center_x < cols and 0 <= center_y < rows): print( @@ -527,34 +938,27 @@ def compute_slice_angles_and_anisotropy( os.path.join(output_dir, a2, f"{a2}_{idx:06d}.{ext}"), os.path.join(output_dir, "FA", f"FA_{idx:06d}.{ext}"), ] - if write_vectors: + if write_vectors and vector_format == "npy": expected_paths.append( - os.path.join(output_dir, "eigen_vec", f"eigen_vec_{idx:06d}.npy") + str(vector_field_path(output_dir, "npy") / f"eigen_vec_{idx:06d}.npy") ) # Skip if all outputs are already present and we are not in test mode if ( not is_test - and expected_paths + and (write_angles or write_vectors) and all(os.path.exists(p) for p in expected_paths) + and (not write_vectors or vector_already_written) ): return 0.0, 0.0, True compute_t0 = time.perf_counter() - # Build a small window around the slice index to estimate the local axis direction - buffer = 5 - if z < buffer: - VEC_PTS = center_line[: min(z + buffer, len(center_line))] - elif z >= len(center_line) - buffer: - VEC_PTS = center_line[max(z - buffer, 0) :] - else: - VEC_PTS = center_line[z - buffer : z + buffer] - - center_vec = calculate_center_vector(VEC_PTS) - # Compute FA and the chosen angle pair if write_angles or is_test: + if eigen_val_slice is None: + raise ValueError("Eigenvalues are required for angle and FA calculation") + center_vec = calculate_center_vector(_centerline_neighborhood(center_line, idx)) img_FA = compute_fraction_anisotropy(eigen_val_slice) vector_field_slice_rotated = rotate_vectors_to_new_axis( vector_field_slice, center_vec @@ -569,11 +973,18 @@ def compute_slice_angles_and_anisotropy( vector_field_slice_rotated ) + if invalid_mask_slice is not None: + img_angle1[invalid_mask_slice] = np.nan + img_angle2[invalid_mask_slice] = np.nan + img_FA[invalid_mask_slice] = np.nan + compute_time = time.perf_counter() - compute_t0 write_t0 = time.perf_counter() # Test mode: visualize a 2x2 figure and write to test subfolder if is_test: + if img_slice is None: + raise ValueError("The intensity image is required in test mode") if mode == "ha_ia": titles = ( ("Projected Helical Angle", "Projected Intrusion Angle") @@ -602,11 +1013,11 @@ def compute_slice_angles_and_anisotropy( img_angle1, img_angle2, img_FA, - start_index, + idx, test_output_dir, ext, output_type, - z, + 0, colormap_angle=colormap_angle, colormap_angle2=colormap_angle2, angle_names=angle_names, @@ -620,17 +1031,24 @@ def compute_slice_angles_and_anisotropy( img_angle1, img_angle2, img_FA, - start_index, + idx, output_dir, ext, output_type, - z, + 0, colormap_angle=colormap_angle, colormap_angle2=colormap_angle2, angle_names=angle_names, angle_ranges=angle_ranges, ) - if write_vectors: - write_vector_field(vector_field_slice, start_index, output_dir, z) + if write_vectors and not vector_already_written: + write_vector_field( + vector_field_slice, + idx, + output_dir, + 0, + vector_format=vector_format, + zarr_store=zarr_store, + ) return compute_time, time.perf_counter() - write_t0, False diff --git a/src/cardiotensor/scripts/analysis_streamlines.py b/src/cardiotensor/scripts/analysis_streamlines.py index 8b6b7c6..f798022 100644 --- a/src/cardiotensor/scripts/analysis_streamlines.py +++ b/src/cardiotensor/scripts/analysis_streamlines.py @@ -16,20 +16,15 @@ from __future__ import annotations import argparse +import csv +import itertools from pathlib import Path +from typing import Iterator import matplotlib as mpl import matplotlib.pyplot as plt import numpy as np -try: - from tqdm import tqdm -except Exception: - - def tqdm(it, **kwargs): # fallback - return it - - mpl.rcParams.update( { "savefig.dpi": 300, @@ -56,36 +51,36 @@ def tqdm(it, **kwargs): # fallback # ---------- helpers ---------- -def load_streamlines(path: Path, key: str | None = None) -> list[np.ndarray]: +def iter_streamlines(path: Path, key: str | None = None) -> Iterator[np.ndarray]: """ - Load streamlines from .trk or .npz/.npy object arrays. - Always returns list of streamlines in (x,y,z) coordinates. + Yield streamline geometry without loading unused TRK per-point fields. + + Nibabel applies the TRK affine lazily, so yielded TRK coordinates are in + the physical space encoded in the file header. Legacy CardioTensor TRKs + used an identity affine and can still be scaled with --voxel-size. """ suffix = path.suffix.lower() if suffix == ".trk": - from cardiotensor.utils.streamlines_io_utils import load_trk_streamlines - - streamlines_xyz, attrs = load_trk_streamlines(path) - return [np.asarray(sl, dtype=np.float32) for sl in streamlines_xyz] + import nibabel as nib - elif suffix == ".npz": + tractogram = nib.streamlines.load(str(path), lazy_load=True).tractogram + source = tractogram.streamlines + elif suffix in {".npz", ".npy"}: data = np.load(path, allow_pickle=True) - k = key if key is not None else "streamlines" - arr = data[k] - - elif suffix == ".npy": - arr = np.load(path, allow_pickle=True) - + source = data[key or "streamlines"] if suffix == ".npz" else data else: raise ValueError(f"Unsupported file type: {path.suffix}") - if isinstance(arr, np.ndarray) and arr.dtype == object: - streamlines = [np.asarray(sl) for sl in arr.tolist()] - else: - streamlines = [np.asarray(sl) for sl in arr] + for streamline in source: + points = np.asarray(streamline, dtype=np.float32) + if points.ndim == 2 and points.shape[1] == 3 and len(points) >= 2: + yield points - return [sl for sl in streamlines if sl.ndim == 2 and sl.shape[1] == 3] + +def load_streamlines(path: Path, key: str | None = None) -> list[np.ndarray]: + """Compatibility wrapper for callers that explicitly need a list.""" + return list(iter_streamlines(path, key)) VoxelSize = float | tuple[float, float, float] @@ -101,6 +96,26 @@ def _scale_phys(P: np.ndarray, voxel_size: VoxelSize) -> np.ndarray: return Q +def resample_streamline(P: np.ndarray, spacing: float) -> np.ndarray: + """Linearly resample a polyline at approximately uniform arc-length spacing.""" + if len(P) < 2 or spacing <= 0: + return P + + segment_lengths = np.linalg.norm(np.diff(P, axis=0), axis=1) + cumulative = np.concatenate(([0.0], np.cumsum(segment_lengths))) + keep = np.concatenate(([True], np.diff(cumulative) > 0)) + cumulative = cumulative[keep] + P = P[keep] + if len(P) < 2 or cumulative[-1] <= 0: + return P + + n_points = max(2, int(np.ceil(cumulative[-1] / spacing)) + 1) + distances = np.linspace(0.0, cumulative[-1], n_points) + return np.column_stack( + [np.interp(distances, cumulative, P[:, axis]) for axis in range(3)] + ) + + def streamline_length(sl_xyz: np.ndarray, voxel_size: VoxelSize) -> float: P = _scale_phys(sl_xyz, voxel_size) diffs = np.diff(P, axis=0) @@ -113,6 +128,7 @@ def chord_length(sl_xyz: np.ndarray, voxel_size: VoxelSize) -> float: def curvature_discrete(sl_xyz: np.ndarray, voxel_size: VoxelSize) -> np.ndarray: + """Menger curvature through each consecutive triplet of points.""" if len(sl_xyz) < 3: return np.zeros((0,), dtype=np.float32) P = _scale_phys(sl_xyz, voxel_size) @@ -125,16 +141,17 @@ def curvature_discrete(sl_xyz: np.ndarray, voxel_size: VoxelSize) -> np.ndarray: ac = np.linalg.norm(AC, axis=1) denom = ab * bc * ac with np.errstate(divide="ignore", invalid="ignore"): - kappa = np.where(denom > 0, (area2) / denom, 0.0) + kappa = np.where(denom > 0, 2.0 * area2 / denom, 0.0) return np.nan_to_num(kappa).astype(np.float32) def percentile_bounds( values_list: list[np.ndarray], p_lo: float, p_hi: float ) -> tuple[float, float]: - vals = np.concatenate( - [v[np.isfinite(v)] for v in values_list if v.size > 0], axis=0 - ) + finite_values = [v[np.isfinite(v)] for v in values_list if v.size > 0] + if not finite_values: + return 0.0, 1.0 + vals = np.concatenate(finite_values, axis=0) if vals.size == 0: return 0.0, 1.0 lo = float(np.percentile(vals, p_lo)) @@ -147,9 +164,10 @@ def percentile_bounds( def gaussian_kde_1d( samples: np.ndarray, grid: np.ndarray, bandwidth: float | None = None ) -> np.ndarray: + """Approximate a Gaussian KDE using a fixed-memory smoothed histogram.""" x = samples[np.isfinite(samples)].astype(np.float64) n = len(x) - if n < 2: + if n < 2 or len(grid) < 2: return np.zeros_like(grid) std = x.std(ddof=1) if bandwidth is None: @@ -160,9 +178,21 @@ def gaussian_kde_1d( ) else: h = float(bandwidth) - inv = 1.0 / (h * np.sqrt(2.0 * np.pi)) - Z = (grid[:, None] - x[None, :]) / h - return inv * np.exp(-0.5 * Z * Z).mean(axis=1) + step = float(np.median(np.diff(grid))) + if not np.isfinite(h) or h <= 0 or not np.isfinite(step) or step <= 0: + return np.zeros_like(grid) + + edges = np.concatenate((grid - step / 2.0, [grid[-1] + step / 2.0])) + counts = np.histogram(x, bins=edges)[0].astype(np.float64) + sigma_bins = h / step + radius = max(1, int(np.ceil(4.0 * sigma_bins))) + offsets = np.arange(-radius, radius + 1, dtype=np.float64) + kernel = np.exp(-0.5 * (offsets / sigma_bins) ** 2) + kernel /= kernel.sum() + smoothed = np.convolve( + np.pad(counts, radius, mode="constant"), kernel, mode="same" + )[radius:-radius] + return smoothed / (n * step) def make_grid(lo: float, hi: float, n_points: int = 512) -> np.ndarray: @@ -182,14 +212,14 @@ def normalize_curve(y: np.ndarray, do_norm: bool) -> np.ndarray: def plot_hist(ax, data, bins, xlabel, normalize, label): - counts, edges = np.histogram(data, bins=bins) - y = normalize_curve(counts, normalize) - - # Use ax.hist with black borders + counts, _ = np.histogram(data, bins=bins) + weights = None + if normalize and counts.size and counts.max() > 0: + weights = np.full(data.shape, 1.0 / counts.max(), dtype=float) ax.hist( data, bins=bins, - weights=y / counts if normalize and counts.max() > 0 else None, + weights=weights, alpha=0.4, label=label, edgecolor="black", @@ -226,38 +256,603 @@ def plot_kde(ax, data, grid, xlabel, bandwidth_mult, normalize, label): def compute_metrics( - path: Path, key: str | None, voxel_size: VoxelSize, min_points: int + path: Path, + key: str | None, + voxel_size: VoxelSize, + min_points: int, + resample_spacing: float | None = None, ): print(f"Load {path}") - sls = load_streamlines(path, key) - - # 🔹 Quick test: print size (min/max in pixels) - if sls: - all_pts = np.vstack(sls) - min_vals = all_pts.min(axis=0) - max_vals = all_pts.max(axis=0) - size = max_vals - min_vals - print( - f"Space size (pixels): {size} (x={size[0]:.1f}, y={size[1]:.1f}, z={size[2]:.1f})" + iterator = iter_streamlines(path, key) + sample = list(itertools.islice(iterator, 512)) + if not sample: + return { + "length": np.array([]), + "mean_curvature": np.array([]), + "tortuosity": np.array([]), + "total_streamlines": 0, + "eligible_streamlines": 0, + "space_size": np.zeros(3), + "space_diagonal": 0.0, + } + + if resample_spacing is None: + sampled_lengths = [] + for sl in sample: + P = _scale_phys(sl, voxel_size) + lengths = np.linalg.norm(np.diff(P, axis=0), axis=1) + sampled_lengths.append(lengths[np.isfinite(lengths) & (lengths > 0)]) + valid_lengths = [values for values in sampled_lengths if values.size] + resample_spacing = ( + float(np.median(np.concatenate(valid_lengths))) if valid_lengths else 0.0 ) + if resample_spacing <= 0: + raise ValueError("resample_spacing must be positive") + print(f"Uniform curvature sampling interval: {resample_spacing:g}") lengths, mean_curv, torts = [], [], [] - for sl in sls: + min_vals = np.full(3, np.inf) + max_vals = np.full(3, -np.inf) + total_streamlines = 0 + for sl in itertools.chain(sample, iterator): + P = _scale_phys(sl, voxel_size) + total_streamlines += 1 + min_vals = np.minimum(min_vals, np.min(P, axis=0)) + max_vals = np.maximum(max_vals, np.max(P, axis=0)) if sl.shape[0] < min_points: continue - L = streamline_length(sl, voxel_size) - D = chord_length(sl, voxel_size) - K = curvature_discrete(sl, voxel_size) + L = float(np.linalg.norm(np.diff(P, axis=0), axis=1).sum()) + D = float(np.linalg.norm(P[-1] - P[0])) + uniform_points = resample_streamline(P, resample_spacing) + K = curvature_discrete(uniform_points, 1.0) tort = (L / D - 1.0) if D > 0 else np.nan lengths.append(L) mean_curv.append(K.mean() if K.size else 0.0) torts.append(tort) + size = max_vals - min_vals + print( + "Coordinate-space size: " + f"{size} (x={size[0]:.3g}, y={size[1]:.3g}, z={size[2]:.3g})" + ) + if ( + path.suffix.lower() == ".trk" + and np.allclose(np.asarray(voxel_size), 1.0) + and np.max(size) > 1000 + ): + print( + "WARNING: This looks like a legacy TRK stored in voxel coordinates. " + "QuickBundles thresholds would use voxels, not millimetres. Rerun with " + "--voxel-size MICROMETRES/1000 (for example, 0.016495 for 16.495 µm).", + flush=True, + ) return dict( length=np.array(lengths), mean_curvature=np.array(mean_curv), tortuosity=np.array(torts), + total_streamlines=total_streamlines, + eligible_streamlines=len(lengths), + space_size=size, + space_diagonal=float(np.linalg.norm(size)), + ) + + +def run_quickbundles( + path: Path, + key: str | None, + label: str, + outdir: Path, + voxel_size: VoxelSize, + min_points: int, + thresholds: list[float], + cluster_points: int, + max_streamlines: int, + seed: int, + cluster_min_size: int = 1, + max_clusters: int | None = None, +) -> dict: + """Cluster a reproducible sample and save its centroids and members.""" + if max_clusters is not None and max_clusters < 1: + raise ValueError("max_clusters must be at least 1") + from dipy.segment.clustering import QuickBundles, QuickBundlesX + from dipy.segment.featurespeed import ResampleFeature + from dipy.segment.metricspeed import AveragePointwiseEuclideanMetric + + feature = ResampleFeature(nb_points=cluster_points) + rng = np.random.default_rng(seed) + reservoir: list[tuple[int, np.ndarray]] = [] + eligible = 0 + limit = None if max_streamlines == 0 else max_streamlines + method = "QuickBundles" if len(thresholds) == 1 else "QuickBundlesX" + + print(f"\nPreparing {method} for {path}", flush=True) + print( + f" thresholds: {thresholds} | points per streamline: {cluster_points}", + flush=True, + ) + if limit is None: + print(" sampling: all eligible streamlines", flush=True) + else: + print( + f" sampling: reproducible reservoir of at most {limit:,} " + f"streamlines (seed={seed})", + flush=True, + ) + + for source_index, streamline in enumerate(iter_streamlines(path, key)): + if source_index and source_index % 100_000 == 0: + print( + f" scanned {source_index:,} | eligible {eligible:,} | " + f"selected {len(reservoir):,}", + flush=True, + ) + if len(streamline) < min_points: + continue + if not np.all(np.isfinite(streamline)): + continue + + eligible += 1 + replacement = None + if limit is not None and len(reservoir) >= limit: + replacement = int(rng.integers(eligible)) + if replacement >= limit: + continue + + points = feature.extract(_scale_phys(streamline, voxel_size)).astype( + np.float32, copy=False + ) + item = (source_index, points) + if replacement is None: + reservoir.append(item) + else: + reservoir[replacement] = item + + if not reservoir: + raise ValueError(f"No valid streamlines available for clustering in {path}") + + # QuickBundles is order-dependent. Restore source order after sampling so + # that the fixed random seed gives reproducible results. + reservoir.sort(key=lambda item: item[0]) + source_indices = np.asarray([item[0] for item in reservoir], dtype=np.int64) + streamlines = [item[1] for item in reservoir] + print( + f"Sampling complete: selected {len(streamlines):,} of " + f"{eligible:,} eligible streamlines", + flush=True, + ) + + metric = AveragePointwiseEuclideanMetric(feature) + print(f"Running {method}...", flush=True) + if len(thresholds) == 1: + clusters = QuickBundles(thresholds[0], metric=metric).cluster(streamlines) + else: + tree = QuickBundlesX(thresholds, metric=metric).cluster(streamlines) + clusters = tree.get_clusters(len(thresholds)) + + ranked_clusters = sorted(clusters, key=len, reverse=True) + centroids = [ + np.asarray(cluster.centroid, dtype=np.float32) for cluster in ranked_clusters + ] + sizes = np.asarray([len(cluster) for cluster in ranked_clusters], dtype=np.int32) + print(f"Clustering complete: {len(ranked_clusters):,} clusters", flush=True) + if len(sizes): + largest = ", ".join(f"{size:,}" for size in sizes[:5]) + print(f" five largest cluster sizes: {largest}", flush=True) + + eligible_cluster_ids = np.flatnonzero(sizes >= cluster_min_size) + if not len(eligible_cluster_ids): + raise ValueError( + f"No cluster has at least {cluster_min_size} streamlines; " + f"the largest has {int(sizes.max())}" + ) + kept_ids = eligible_cluster_ids[:max_clusters] + saved_cluster_mask = np.zeros(len(sizes), dtype=bool) + saved_cluster_mask[kept_ids] = True + kept_centroids = [centroids[cluster_id] for cluster_id in kept_ids] + kept_sizes = sizes[kept_ids] + kept_streamlines = int(kept_sizes.sum()) + print( + f"Keeping {len(kept_ids):,}/{len(eligible_cluster_ids):,} clusters with " + f"at least {cluster_min_size:,} members; they contain " + f"{kept_streamlines:,}/{len(streamlines):,} sampled streamlines " + f"({100.0 * kept_streamlines / len(streamlines):.1f}%)", + flush=True, + ) + assignments = np.full(len(streamlines), -1, dtype=np.int32) + for cluster_id, cluster in enumerate(ranked_clusters): + assignments[np.asarray(cluster.indices, dtype=np.int64)] = cluster_id + if np.any(assignments < 0): + raise RuntimeError("QuickBundles did not assign every sampled streamline") + + kept_member_indices = np.flatnonzero(saved_cluster_mask[assignments]) + member_clusters = { + int(source_indices[index]): int(assignments[index]) + for index in kept_member_indices + } + + safe_label = "".join( + char if char.isalnum() or char in "-_." else "_" for char in label + ).strip("._") or path.stem + prefix = outdir / f"{safe_label}_quickbundles" + csv_path = prefix.with_name(f"{prefix.name}_clusters.csv") + trk_path = prefix.with_name(f"{prefix.name}_centroids.trk") + member_trk_path = prefix.with_name(f"{prefix.name}_members.trk") + membership_path = prefix.with_name(f"{prefix.name}_membership.npz") + + print("Writing clustering results...", flush=True) + with csv_path.open("w", newline="") as output: + writer = csv.DictWriter( + output, + fieldnames=( + "cluster_id", + "streamline_count", + "percentage", + "centroid_length", + "method", + "thresholds", + "saved_centroid", + "sampled_streamlines", + "eligible_streamlines", + ), + ) + writer.writeheader() + for cluster_id, (centroid, size) in enumerate(zip(centroids, sizes)): + centroid_length = float( + np.linalg.norm(np.diff(centroid, axis=0), axis=1).sum() + ) + writer.writerow( + { + "cluster_id": cluster_id, + "streamline_count": int(size), + "percentage": 100.0 * int(size) / len(streamlines), + "centroid_length": centroid_length, + "method": method, + "thresholds": ";".join(map(str, thresholds)), + "sampled_streamlines": len(streamlines), + "eligible_streamlines": eligible, + "saved_centroid": bool(saved_cluster_mask[cluster_id]), + } + ) + + np.savez_compressed( + membership_path, + source_streamline_index=source_indices, + cluster_id=assignments, + ) + + import nibabel as nib + + centroid_tractogram = nib.streamlines.Tractogram( + kept_centroids, + data_per_point={ + "cluster_id": [ + np.full((len(centroid), 1), cluster_id, dtype=np.float32) + for cluster_id, centroid in zip(kept_ids, kept_centroids) + ] + }, + data_per_streamline={ + "cluster_size": kept_sizes.astype(np.float32)[:, None], + }, + affine_to_rasmm=np.eye(4), + ) + header = None + if path.suffix.lower() == ".trk": + header = nib.streamlines.load(str(path), lazy_load=True).header.copy() + nib.streamlines.save( + nib.streamlines.TrkFile( + centroid_tractogram, + header=header.copy() if header is not None else None, + ), + str(trk_path), + ) + + from nibabel.streamlines.tractogram import TractogramItem + + def iter_cluster_members(): + written = 0 + for source_index, streamline in enumerate(iter_streamlines(path, key)): + cluster_id = member_clusters.get(source_index) + if cluster_id is None: + continue + points = _scale_phys(streamline, voxel_size).astype( + np.float32, copy=False + ) + written += 1 + if written % 10_000 == 0: + print(f" wrote {written:,} cluster members", flush=True) + yield TractogramItem( + points, + {"cluster_size": np.asarray([sizes[cluster_id]], dtype=np.float32)}, + { + "cluster_id": np.full( + (len(points), 1), cluster_id, dtype=np.float32 + ) + }, + ) + + print( + f"Writing {len(member_clusters):,} full-resolution cluster members...", + flush=True, + ) + member_tractogram = nib.streamlines.LazyTractogram.from_data_func( + iter_cluster_members + ) + member_tractogram.affine_to_rasmm = np.eye(4) + nib.streamlines.save( + nib.streamlines.TrkFile( + member_tractogram, + header=header.copy() if header is not None else None, + ), + str(member_trk_path), + ) + + print(f"Cluster summary: {csv_path}", flush=True) + print(f"Cluster membership: {membership_path}", flush=True) + print(f"Cluster centroids: {trk_path}", flush=True) + print( + f"Cluster members: {member_trk_path} " + f"({len(member_clusters):,} streamlines)", + flush=True, + ) + return { + "method": method, + "eligible_streamlines": eligible, + "sampled_streamlines": len(streamlines), + "clusters": len(centroids), + "csv": csv_path, + "membership": membership_path, + "centroids": trk_path, + "members": member_trk_path, + "saved_clusters": len(kept_centroids), + "cluster_sizes": sizes, + } + + +def write_cohort_report( + results: list[dict], + outdir: Path, + input_paths: list[Path] | None = None, + cluster_results: list[dict] | None = None, + max_samples: int = 20_000, +) -> None: + """Write equal-sampled, heart-level comparison tables and figures.""" + if len(results) < 2: + return + + print("\nCreating cohort comparison report...", flush=True) + rng = np.random.default_rng(0) + prepared = [] + for result in results: + diagonal = float(result.get("space_diagonal", 0.0)) + if diagonal <= 0: + raise ValueError(f"Cannot normalize {result['label']}: zero spatial extent") + values = np.column_stack( + ( + result["length"] / diagonal, + result["mean_curvature"], + result["tortuosity"], + ) + ) + values = values[np.all(np.isfinite(values), axis=1)] + if not len(values): + raise ValueError(f"No finite streamline metrics for {result['label']}") + prepared.append(values) + + sample_count = min(len(values) for values in prepared) + if max_samples: + sample_count = min(sample_count, max_samples) + samples = [] + for values in prepared: + indices = rng.choice(len(values), size=sample_count, replace=False) + samples.append(values[indices]) + print( + f" using {sample_count:,} reproducibly sampled streamlines per heart", + flush=True, + ) + + metric_specs = ( + (0, "normalized_length", "Normalized streamline length"), + (1, "mean_curvature", "Mean curvature"), + (2, "tortuosity", "Tortuosity"), + ) + labels = [str(result["label"]) for result in results] + paths = input_paths or [Path("")] * len(results) + + summary_path = outdir / "heart_summary.csv" + summary_rows = [] + for index, (result, path) in enumerate(zip(results, paths)): + diagonal = float(result["space_diagonal"]) + full_metrics = { + "length": np.asarray(result["length"], dtype=float), + "normalized_length": np.asarray(result["length"], dtype=float) + / diagonal, + "mean_curvature": np.asarray(result["mean_curvature"], dtype=float), + "tortuosity": np.asarray(result["tortuosity"], dtype=float), + } + row = { + "heart": result["label"], + "input": str(path), + "total_streamlines": result.get( + "total_streamlines", len(result["length"]) + ), + "eligible_streamlines": result.get( + "eligible_streamlines", len(result["length"]) + ), + "space_x": float(result["space_size"][0]), + "space_y": float(result["space_size"][1]), + "space_z": float(result["space_size"][2]), + "space_diagonal": diagonal, + } + for name, values in full_metrics.items(): + values = values[np.isfinite(values)] + q1, median, q3 = np.percentile(values, [25, 50, 75]) + row.update( + { + f"{name}_mean": float(np.mean(values)), + f"{name}_q1": float(q1), + f"{name}_median": float(median), + f"{name}_q3": float(q3), + f"{name}_iqr": float(q3 - q1), + } + ) + + if cluster_results: + clustering = cluster_results[index] + sizes = np.asarray(clustering["cluster_sizes"], dtype=float) + fractions = sizes / sizes.sum() + entropy = float(-np.sum(fractions * np.log(fractions))) + row.update( + { + "quickbundles_method": clustering["method"], + "quickbundles_sampled_streamlines": clustering[ + "sampled_streamlines" + ], + "quickbundles_clusters": clustering["clusters"], + "quickbundles_top10_fraction": float(fractions[:10].sum()), + "quickbundles_entropy": entropy, + "quickbundles_normalized_entropy": ( + entropy / np.log(len(sizes)) if len(sizes) > 1 else 0.0 + ), + "quickbundles_effective_clusters": float(np.exp(entropy)), + } + ) + summary_rows.append(row) + + summary_path = outdir / "heart_summary.csv" + with summary_path.open("w", newline="") as output: + writer = csv.DictWriter(output, fieldnames=list(summary_rows[0])) + writer.writeheader() + writer.writerows(summary_rows) + print(f" heart-level summary: {summary_path}", flush=True) + + colors = plt.get_cmap("tab10")(np.linspace(0, 1, len(results))) + fig, axes = plt.subplots(1, 3, figsize=(16, 4.8)) + for metric_index, _, title in metric_specs: + ax = axes[metric_index] + for label, color, values in zip(labels, colors, samples): + sorted_values = np.sort(values[:, metric_index]) + probability = np.arange(1, len(sorted_values) + 1) / len(sorted_values) + ax.plot(sorted_values, probability, label=label, color=color, linewidth=2) + ax.set_xlabel(title, fontweight="bold") + ax.set_ylabel("Cumulative fraction" if metric_index == 0 else "") + ax.set_ylim(0, 1) + ax.grid(alpha=0.2) + axes[-1].legend(bbox_to_anchor=(1.02, 1), loc="upper left") + fig.tight_layout() + for suffix in ("png", "pdf"): + fig.savefig(outdir / f"geometry_distributions.{suffix}", bbox_inches="tight") + plt.close(fig) + print( + f" geometry distributions: {outdir / 'geometry_distributions.png'}", + flush=True, ) + fig, axes = plt.subplots(1, 3, figsize=(16, 4.8), sharex=True) + x = np.arange(len(results)) + for metric_index, _, title in metric_specs: + medians, lower, upper = [], [], [] + for values in prepared: + q1, median, q3 = np.percentile(values[:, metric_index], [25, 50, 75]) + medians.append(median) + lower.append(median - q1) + upper.append(q3 - median) + axes[metric_index].errorbar( + x, + medians, + yerr=np.asarray([lower, upper]), + fmt="o", + capsize=4, + color="black", + ) + axes[metric_index].scatter(x, medians, c=colors, s=55, zorder=3) + axes[metric_index].set_title(title, fontweight="bold") + axes[metric_index].set_xticks(x, labels, rotation=45, ha="right") + axes[metric_index].set_ylabel("Median and interquartile range") + axes[metric_index].grid(axis="y", alpha=0.2) + fig.tight_layout() + for suffix in ("png", "pdf"): + fig.savefig(outdir / f"heart_metric_summary.{suffix}", bbox_inches="tight") + plt.close(fig) + print(f" metric summary: {outdir / 'heart_metric_summary.png'}", flush=True) + + distance = np.zeros((len(results), len(results)), dtype=float) + metric_distances = [] + for metric_index, _, _ in metric_specs: + matrix = np.zeros_like(distance) + pooled = np.concatenate([values[:, metric_index] for values in samples]) + q1, q3 = np.percentile(pooled, [25, 75]) + scale = q3 - q1 + if scale <= 0: + scale = np.std(pooled) or 1.0 + sorted_values = [np.sort(values[:, metric_index]) for values in samples] + for left in range(len(results)): + for right in range(left + 1, len(results)): + value = float( + np.mean(np.abs(sorted_values[left] - sorted_values[right])) + ) + matrix[left, right] = matrix[right, left] = value / scale + metric_distances.append(matrix) + distance = np.mean(metric_distances, axis=0) + + matrix_path = outdir / "heart_distance_matrix.csv" + with matrix_path.open("w", newline="") as output: + writer = csv.writer(output) + writer.writerow(["heart", *labels]) + for label, row in zip(labels, distance): + writer.writerow([label, *row]) + + fig, ax = plt.subplots(figsize=(7.5, 6.5)) + image = ax.imshow(distance, cmap="magma", vmin=0) + ax.set_xticks(np.arange(len(labels)), labels, rotation=45, ha="right") + ax.set_yticks(np.arange(len(labels)), labels) + ax.set_title("Combined streamline-distribution distance", fontweight="bold") + colorbar = fig.colorbar(image, ax=ax) + colorbar.set_label("Mean Wasserstein distance / pooled IQR") + fig.tight_layout() + for suffix in ("png", "pdf"): + fig.savefig(outdir / f"heart_distance_heatmap.{suffix}", bbox_inches="tight") + plt.close(fig) + print(f" distance matrix: {matrix_path}", flush=True) + print( + f" distance heatmap: {outdir / 'heart_distance_heatmap.png'}", + flush=True, + ) + + if cluster_results: + fig, axes = plt.subplots(1, 2, figsize=(12, 4.8)) + for label, color, clustering in zip(labels, colors, cluster_results): + sizes = np.asarray(clustering["cluster_sizes"], dtype=float) + fractions = sizes / sizes.sum() + ranks = np.arange(1, len(sizes) + 1) + axes[0].plot(ranks, 100 * fractions, label=label, color=color) + axes[1].plot( + ranks, 100 * np.cumsum(fractions), label=label, color=color + ) + axes[0].set_xscale("log") + axes[0].set_yscale("log") + axes[0].set_xlabel("Cluster rank") + axes[0].set_ylabel("Streamlines in cluster (%)") + axes[0].set_title("Cluster rank-abundance", fontweight="bold") + axes[1].set_xscale("log") + axes[1].set_xlabel("Number of largest clusters") + axes[1].set_ylabel("Cumulative streamlines (%)") + axes[1].set_ylim(0, 100) + axes[1].set_title("Cumulative cluster coverage", fontweight="bold") + for ax in axes: + ax.grid(alpha=0.2) + axes[1].legend(bbox_to_anchor=(1.02, 1), loc="upper left") + fig.tight_layout() + for suffix in ("png", "pdf"): + fig.savefig( + outdir / f"cluster_rank_abundance.{suffix}", bbox_inches="tight" + ) + plt.close(fig) + print( + f" cluster rank-abundance: {outdir / 'cluster_rank_abundance.png'}", + flush=True, + ) + else: + print(" cluster rank-abundance skipped (use --quickbundles to create it)") + # ---------- CLI ---------- @@ -267,7 +862,73 @@ def script(): ap.add_argument("inputs", nargs="+", type=Path) ap.add_argument("--labels", nargs="+", type=str) ap.add_argument("--key", type=str, default=None) - ap.add_argument("--voxel-size", type=float, default=1.0) + ap.add_argument( + "--voxel-size", + type=float, + default=1.0, + help=( + "Coordinate multiplier for legacy TRK/NPY/NPZ inputs. Keep 1 for " + "TRKs whose header already contains physical spacing" + ), + ) + ap.add_argument( + "--resample-spacing", + type=float, + default=None, + help=( + "Physical spacing used for curvature. By default, estimate one common " + "spacing per input from its first 512 streamlines" + ), + ) + ap.add_argument( + "--quickbundles", + type=float, + nargs="+", + metavar="THRESHOLD", + help=( + "Cluster streamlines in physical units. One threshold uses QuickBundles; " + "multiple decreasing thresholds use QuickBundlesX" + ), + ) + ap.add_argument( + "--cluster-points", + type=int, + default=12, + help="Points used by the orientation-invariant clustering metric (default: 12)", + ) + ap.add_argument( + "--cluster-max-streamlines", + type=int, + default=50_000, + help=( + "Maximum reproducible sample used for clustering (default: 50000; " + "use 0 for all streamlines)" + ), + ) + ap.add_argument( + "--cluster-min-size", + type=int, + default=1, + help=( + "Write centroid and member streamlines only for clusters with at " + "least this many sampled members" + ), + ) + ap.add_argument( + "--max-clusters", + type=int, + default=None, + help=( + "Write only the N largest clusters that pass --cluster-min-size; " + "the CSV and membership NPZ still contain all clusters" + ), + ) + ap.add_argument( + "--cluster-seed", + type=int, + default=0, + help="Random seed for reservoir sampling (default: 0)", + ) ap.add_argument("--min-points", type=int, default=2) ap.add_argument("--bins", type=int, default=60) ap.add_argument("--kde-points", type=int, default=512) @@ -283,6 +944,34 @@ def script(): args = ap.parse_args() labels = args.labels or [p.stem for p in args.inputs] + if len(labels) != len(args.inputs): + ap.error("--labels must contain exactly one label per input") + if len(set(labels)) != len(labels): + ap.error("--labels must be unique for cohort comparison") + if args.voxel_size <= 0: + ap.error("--voxel-size must be positive") + if args.min_points < 2: + ap.error("--min-points must be at least 2") + if args.resample_spacing is not None and args.resample_spacing <= 0: + ap.error("--resample-spacing must be positive") + if args.cluster_points < 2: + ap.error("--cluster-points must be at least 2") + if args.cluster_max_streamlines < 0: + ap.error("--cluster-max-streamlines cannot be negative") + if args.cluster_min_size < 1: + ap.error("--cluster-min-size must be at least 1") + if args.max_clusters is not None and args.max_clusters < 1: + ap.error("--max-clusters must be at least 1") + if args.cluster_seed < 0: + ap.error("--cluster-seed cannot be negative") + if args.quickbundles: + if any(threshold <= 0 for threshold in args.quickbundles): + ap.error("--quickbundles thresholds must be positive") + if len(args.quickbundles) > 1 and any( + coarse <= fine + for coarse, fine in zip(args.quickbundles, args.quickbundles[1:]) + ): + ap.error("QuickBundlesX thresholds must be strictly decreasing") args.outdir.mkdir(parents=True, exist_ok=True) @@ -292,10 +981,41 @@ def script(): print("") results = [] + cluster_results = [] for path, label in zip(args.inputs, labels): - metrics = compute_metrics(path, args.key, args.voxel_size, args.min_points) + metrics = compute_metrics( + path, + args.key, + args.voxel_size, + args.min_points, + args.resample_spacing, + ) metrics["label"] = label results.append(metrics) + if args.quickbundles: + cluster_results.append( + run_quickbundles( + path=path, + key=args.key, + label=label, + outdir=args.outdir, + voxel_size=args.voxel_size, + min_points=args.min_points, + thresholds=args.quickbundles, + cluster_points=args.cluster_points, + max_streamlines=args.cluster_max_streamlines, + seed=args.cluster_seed, + cluster_min_size=args.cluster_min_size, + max_clusters=args.max_clusters, + ) + ) + + write_cohort_report( + results, + args.outdir, + input_paths=args.inputs, + cluster_results=cluster_results or None, + ) # Percentile bounds arr_len = [m["length"] for m in results if m["length"].size > 0] @@ -345,7 +1065,7 @@ def script(): fig.tight_layout() fig.savefig(args.outdir / f"{metric}.png", bbox_inches="tight") fig.savefig(args.outdir / f"{metric}.pdf", bbox_inches="tight") - plt.show() + plt.close(fig) if __name__ == "__main__": diff --git a/src/cardiotensor/scripts/compute_orientation.py b/src/cardiotensor/scripts/compute_orientation.py index 7c730a4..34873e1 100644 --- a/src/cardiotensor/scripts/compute_orientation.py +++ b/src/cardiotensor/scripts/compute_orientation.py @@ -1,8 +1,10 @@ import argparse +import math import sys import time from cardiotensor.orientation.orientation_computation_pipeline import ( + _safe_low_memory_chunk_size, compute_orientation, ) from cardiotensor.utils.DataReader import DataReader @@ -106,9 +108,12 @@ def script() -> None: axis_points = params.get("AXIS_POINTS", None) vertical_padding = params.get("VERTICAL_PADDING", None) write_vectors = params.get("WRITE_VECTORS", False) + vector_format = params.get("VECTOR_FORMAT", "zarr") write_angles = params.get("WRITE_ANGLES", True) angle_mode = params.get("ANGLE_MODE", "ha_ia") use_gpu = params.get("USE_GPU", True) + low_memory = params.get("LOW_MEMORY", False) + low_memory_dir = params.get("LOW_MEMORY_DIR", None) is_test = force_test or params.get("TEST", False) n_slice_test = params.get("N_SLICE_TEST", None) show_quiver = params.get("SHOW_QUIVER", True) @@ -148,9 +153,12 @@ def script() -> None: axis_points=axis_points, vertical_padding=vertical_padding, write_vectors=write_vectors, + vector_format=vector_format, angle_mode=angle_mode, write_angles=write_angles, use_gpu=use_gpu, + low_memory=low_memory, + low_memory_dir=low_memory_dir, is_test=is_test, n_slice_test=n_slice_test, show_quiver=show_quiver, @@ -164,6 +172,41 @@ def script() -> None: print(f"--- {time.time() - t0:.1f} seconds (TEST mode) ---") return + if low_memory: + padding = math.ceil( + vertical_padding + if vertical_padding is not None + else int(sigma * truncate + 0.5) + int(rho * truncate + 0.5) + ) + requested_chunk = min(n_chunk, end_index - start_index) + safe_chunk, requested_peak, safe_peak = _safe_low_memory_chunk_size( + requested_chunk, + total_slices, + data_reader.shape[-2], + data_reader.shape[-1], + data_reader.dtype, + padding=padding, + include_eigenvalues=write_angles, + has_mask=mask_path is not None, + write_angles=write_angles, + sigma=sigma, + rho=rho, + truncate=truncate, + ) + if safe_chunk < requested_chunk: + print( + "⚠️ LOW_MEMORY adjusted N_CHUNK from " + f"{requested_chunk} to {safe_chunk}: conservative peak " + f"{requested_peak / 1024**3:.2f} GiB -> " + f"{safe_peak / 1024**3:.2f} GiB." + ) + else: + print( + f"LOW_MEMORY chunk check: N_CHUNK={safe_chunk}, " + f"conservative peak={safe_peak / 1024**3:.2f} GiB." + ) + n_chunk = safe_chunk + # --- Build Chunks --- chunks: list[tuple[int, int]] = [] if reverse: @@ -196,9 +239,12 @@ def script() -> None: axis_points=axis_points, vertical_padding=vertical_padding, write_vectors=write_vectors, + vector_format=vector_format, write_angles=write_angles, angle_mode=angle_mode, use_gpu=use_gpu, + low_memory=low_memory, + low_memory_dir=low_memory_dir, is_test=is_test, n_slice_test=n_slice_test, show_quiver=show_quiver, diff --git a/src/cardiotensor/scripts/convert_trk_to_am.py b/src/cardiotensor/scripts/convert_trk_to_am.py index 1239cc1..aa877b1 100644 --- a/src/cardiotensor/scripts/convert_trk_to_am.py +++ b/src/cardiotensor/scripts/convert_trk_to_am.py @@ -67,6 +67,21 @@ def script(): streamlines_xyz, attrs = load_trk_streamlines( inp ) # dict like {"HA":[...], "IA":[...], ...} + keep = [ + index + for index, streamline in enumerate(streamlines_xyz) + if len(streamline) >= 2 + ] + removed = len(streamlines_xyz) - len(keep) + if removed: + print(f"Skipping {removed} streamline(s) with fewer than 2 points") + streamlines_xyz = [streamlines_xyz[index] for index in keep] + attrs = { + name: [values[index] for index in keep] for name, values in attrs.items() + } + if not streamlines_xyz: + raise ValueError("No streamlines with at least 2 points to export") + attrs_deg = normalize_attrs_to_degrees( attrs ) # cast to float32 degrees where needed diff --git a/src/cardiotensor/scripts/create_movie.py b/src/cardiotensor/scripts/create_movie.py new file mode 100644 index 0000000..201ea4e --- /dev/null +++ b/src/cardiotensor/scripts/create_movie.py @@ -0,0 +1,269 @@ +#!/usr/bin/env python3 +"""Create an RGB movie from a scalar image stack.""" + +from __future__ import annotations + +import argparse +import sys +from collections.abc import Iterator +from pathlib import Path + +import cv2 +import matplotlib.pyplot as plt +import numpy as np +from alive_progress import alive_bar + +from cardiotensor.colormaps.helix_angle import helix_angle_cmap +from cardiotensor.utils.DataReader import DataReader +from cardiotensor.utils.image_io import read_image_file +from cardiotensor.utils.utils import read_conf_file + +ANGLE_LIMITS = { + "HA": (-90.0, 90.0), + "IA": (-90.0, 90.0), + "EL": (-90.0, 90.0), + "AZ": (0.0, 360.0), + "FA": (0.0, 1.0), +} + + +def _resolve_stack_path(input_path: Path, volume: str) -> Path: + input_path = input_path.resolve() + volume = volume.upper() + + if input_path.suffix.lower() == ".conf": + params = read_conf_file(str(input_path)) + return (Path(params.get("OUTPUT_PATH", "./output")) / volume).resolve() + + if not input_path.exists(): + raise FileNotFoundError(f"Input path does not exist: {input_path}") + + if input_path.is_dir() and (input_path / volume).is_dir(): + return (input_path / volume).resolve() + + return input_path + + +def _iter_slices( + reader: DataReader, start: int, end: int, step: int +) -> Iterator[np.ndarray]: + if reader.volume_info["stack"]: + files = reader.volume_info["file_list"] + for path in files[start:end:step]: + frame = read_image_file(path) + if frame.ndim == 3 and frame.shape[2] == 1: + frame = frame[:, :, 0] + yield frame + return + + for z in range(start, end, step): + yield reader.load_volume(start_index=z, end_index=z + 1)[0] + + +def _normalise_scalar( + frame: np.ndarray, + *, + volume: str, + vmin: float | None, + vmax: float | None, +) -> np.ndarray: + original_dtype = frame.dtype + frame = frame.astype(np.float32, copy=False) + finite = np.isfinite(frame) + + if vmin is None or vmax is None: + if original_dtype == np.uint8 or ( + finite.any() and float(np.nanmax(frame)) > 1.5 + ): + norm = frame / 255.0 + return np.clip(np.nan_to_num(norm, nan=0.0), 0.0, 1.0) + + default_limits = ANGLE_LIMITS.get(volume.upper()) + if default_limits is not None: + auto_vmin, auto_vmax = default_limits + elif finite.any(): + auto_vmin = float(np.nanpercentile(frame[finite], 1.0)) + auto_vmax = float(np.nanpercentile(frame[finite], 99.0)) + else: + auto_vmin, auto_vmax = 0.0, 1.0 + + vmin = auto_vmin if vmin is None else vmin + vmax = auto_vmax if vmax is None else vmax + + if vmax <= vmin: + return np.zeros(frame.shape, dtype=np.float32) + + norm = (frame - vmin) / (vmax - vmin) + return np.clip(np.nan_to_num(norm, nan=0.0, posinf=1.0, neginf=0.0), 0.0, 1.0) + + +def _to_rgb_frame( + frame: np.ndarray, + *, + volume: str, + colormap: str, + vmin: float | None, + vmax: float | None, +) -> np.ndarray: + if frame.ndim == 3 and frame.shape[2] in {3, 4}: + rgb = frame[:, :, :3] + if rgb.dtype != np.uint8: + rgb = _normalise_scalar(rgb, volume=volume, vmin=vmin, vmax=vmax) * 255 + return rgb.astype(np.uint8) + + norm = _normalise_scalar(frame, volume=volume, vmin=vmin, vmax=vmax) + cmap = helix_angle_cmap if colormap == "helix_angle" else plt.get_cmap(colormap) + return (cmap(norm)[:, :, :3] * 255).astype(np.uint8) + + +def _resize_frame( + frame: np.ndarray, + *, + width: int | None, + height: int | None, + scale: float, +) -> np.ndarray: + if scale <= 0: + raise ValueError("--scale must be > 0") + + src_h, src_w = frame.shape[:2] + dst_w = int(round(src_w * scale)) if width is None else width + dst_h = int(round(src_h * scale)) if height is None else height + + if width is None and height is not None: + dst_w = int(round(src_w * (height / src_h))) + elif height is None and width is not None: + dst_h = int(round(src_h * (width / src_w))) + + if (dst_w, dst_h) == (src_w, src_h): + return frame + + return cv2.resize(frame, (dst_w, dst_h), interpolation=cv2.INTER_AREA) + + +def create_movie( + input_path: Path, + *, + volume: str = "HA", + output: Path | None = None, + fps: float = 30.0, + start: int = 0, + end: int | None = None, + step: int = 1, + width: int | None = None, + height: int | None = None, + scale: float = 1.0, + colormap: str = "helix_angle", + vmin: float | None = None, + vmax: float | None = None, +) -> Path: + stack_path = _resolve_stack_path(input_path, volume) + reader = DataReader(stack_path) + + if len(reader.shape) not in {3, 4}: + raise ValueError( + f"Expected a 3D stack, got shape {reader.shape} at {stack_path}" + ) + + n_slices = reader.shape[0] + start = max(0, start) + end = n_slices if end is None else min(end, n_slices) + if step < 1: + raise ValueError("--step must be >= 1") + if start >= end: + raise ValueError(f"Empty slice range: start={start}, end={end}") + + if output is None: + output = stack_path.parent / f"{stack_path.name}_movie.mp4" + output = output.resolve() + output.parent.mkdir(parents=True, exist_ok=True) + + first_slice = next(_iter_slices(reader, start, start + 1, 1)) + first_rgb = _to_rgb_frame( + first_slice, volume=volume, colormap=colormap, vmin=vmin, vmax=vmax + ) + first_rgb = _resize_frame(first_rgb, width=width, height=height, scale=scale) + out_h, out_w = first_rgb.shape[:2] + + fourcc = cv2.VideoWriter_fourcc(*"mp4v") + writer = cv2.VideoWriter(str(output), fourcc, fps, (out_w, out_h), isColor=True) + if not writer.isOpened(): + raise RuntimeError(f"Could not open video writer for {output}") + + print(f"Creating RGB movie from {stack_path}") + print(f"Slices: {start}:{end}:{step} | Size: {out_w}x{out_h} | FPS: {fps:g}") + try: + total_frames = len(range(start, end, step)) + with alive_bar(total_frames, title="Writing movie", length=40) as bar: + for frame in _iter_slices(reader, start, end, step): + rgb = _to_rgb_frame( + frame, volume=volume, colormap=colormap, vmin=vmin, vmax=vmax + ) + rgb = _resize_frame(rgb, width=width, height=height, scale=scale) + writer.write(cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)) + bar() + finally: + writer.release() + + print(f"Movie created: {output}") + return output + + +def script() -> None: + parser = argparse.ArgumentParser( + description="Create an RGB MP4 movie from a scalar stack such as HA.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "input", + type=Path, + help="A .conf file, an OUTPUT_PATH directory, or a direct stack directory/file.", + ) + parser.add_argument( + "--volume", + default="HA", + help="Subfolder to use when input is a .conf or output directory.", + ) + parser.add_argument("--output", type=Path, default=None, help="Output movie path.") + parser.add_argument("--fps", type=float, default=30.0, help="Frames per second.") + parser.add_argument("--start", type=int, default=0, help="First slice index.") + parser.add_argument("--end", type=int, default=None, help="Stop slice index.") + parser.add_argument("--step", type=int, default=1, help="Slice step.") + parser.add_argument( + "--width", type=int, default=None, help="Output width in pixels." + ) + parser.add_argument( + "--height", type=int, default=None, help="Output height in pixels." + ) + parser.add_argument("--scale", type=float, default=1.0, help="Output scale factor.") + parser.add_argument( + "--colormap", + default="helix_angle", + help="Matplotlib colormap name, or 'helix_angle' for the project HA map.", + ) + parser.add_argument("--vmin", type=float, default=None, help="Scalar value for 0.") + parser.add_argument("--vmax", type=float, default=None, help="Scalar value for 1.") + args = parser.parse_args() + + try: + create_movie( + args.input, + volume=args.volume, + output=args.output, + fps=args.fps, + start=args.start, + end=args.end, + step=args.step, + width=args.width, + height=args.height, + scale=args.scale, + colormap=args.colormap, + vmin=args.vmin, + vmax=args.vmax, + ) + except Exception as exc: + sys.exit(f"Error: {exc}") + + +if __name__ == "__main__": + script() diff --git a/src/cardiotensor/scripts/generate_streamlines.py b/src/cardiotensor/scripts/generate_streamlines.py index 0115588..1b19cf0 100644 --- a/src/cardiotensor/scripts/generate_streamlines.py +++ b/src/cardiotensor/scripts/generate_streamlines.py @@ -17,6 +17,7 @@ from cardiotensor.tractography.generate_streamlines import ( generate_streamlines_from_params, ) +from cardiotensor.utils.image_io import vector_field_path from cardiotensor.utils.utils import read_conf_file @@ -35,6 +36,9 @@ def script() -> None: parser.add_argument("--end-x", type=int, default=None, help="End slice index in X") parser.add_argument("--bin", type=int, default=1, help="Downsampling factor") parser.add_argument("--seeds", type=int, default=20000, help="Number of seeds") + parser.add_argument( + "--random-seed", type=int, default=0, help="Random seed for seed selection" + ) parser.add_argument( "--fa-seed-min", type=float, default=0.2, help="Min FA for seeding" ) @@ -47,7 +51,10 @@ def script() -> None: "--angle", type=float, default=60.0, help="Max turning angle in degrees" ) parser.add_argument( - "--min-len", type=int, default=10, help="Minimum streamline length in points" + "--min-len", + type=int, + default=10, + help="Minimum streamline length in points (values below 2 are raised to 2)", ) # Base angle folder selection for discovery @@ -77,10 +84,16 @@ def script() -> None: # Resolve standard output structure output_dir = Path(params.get("OUTPUT_PATH", "./output")) - vector_field_dir = output_dir / "eigen_vec" + vector_field_dir = vector_field_path( + output_dir, params.get("VECTOR_FORMAT", "zarr") + ) fa_dir = output_dir / "FA" angle_dir = output_dir / angle_folder mask_path = params.get("MASK_PATH", None) + voxel_size_mm = float(params.get("VOXEL_SIZE", 1000.0)) / 1000.0 + if voxel_size_mm <= 0: + print("VOXEL_SIZE must be positive.") + sys.exit(2) # Call library generate_streamlines_from_params( @@ -99,7 +112,9 @@ def script() -> None: max_steps=args.max_steps, angle_threshold=args.angle, min_length_pts=args.min_len, + voxel_sizes_zyx=(voxel_size_mm,) * 3, angle_mode=angle_mode, + random_seed=args.random_seed, ) print(f"Done! Streamlines generated and saved to {output_dir}.") diff --git a/src/cardiotensor/scripts/slurm_launcher.py b/src/cardiotensor/scripts/slurm_launcher.py index 2eb02c9..c9f87f1 100644 --- a/src/cardiotensor/scripts/slurm_launcher.py +++ b/src/cardiotensor/scripts/slurm_launcher.py @@ -92,21 +92,24 @@ def script() -> None: args = parser.parse_args() conf_file_path = args.conf_file_path - slurm_launcher( - conf_file_path, - start_index=args.start_index, - end_index=args.end_index, - chunk_size=args.chunk_size, - partition=args.partition, - time_limit=args.time_limit, - cpus_per_task=args.cpus_per_task, - mem_gb=args.mem_gb, - array_parallel=args.array_parallel, - log_dir=args.log_dir, - submit_dir=args.submit_dir, - monitor=not args.no_monitor, - dry_run=args.dry_run, - ) + try: + slurm_launcher( + conf_file_path, + start_index=args.start_index, + end_index=args.end_index, + chunk_size=args.chunk_size, + partition=args.partition, + time_limit=args.time_limit, + cpus_per_task=args.cpus_per_task, + mem_gb=args.mem_gb, + array_parallel=args.array_parallel, + log_dir=args.log_dir, + submit_dir=args.submit_dir, + monitor=not args.no_monitor, + dry_run=args.dry_run, + ) + except Exception as err: + parser.exit(1, f"Error: {err}\n") if __name__ == "__main__": diff --git a/src/cardiotensor/scripts/visualize_streamlines.py b/src/cardiotensor/scripts/visualize_streamlines.py index 1340149..784759f 100644 --- a/src/cardiotensor/scripts/visualize_streamlines.py +++ b/src/cardiotensor/scripts/visualize_streamlines.py @@ -15,6 +15,8 @@ """ import argparse +import json +import random import sys from pathlib import Path @@ -28,7 +30,9 @@ def _discover_trk_color_fields(trk_path: Path) -> list[str]: """Return list of available per-point scalar keys in TRK, case-preserving.""" - obj = nib.streamlines.load(str(trk_path)) + # Only the TRK header is needed here. A normal load would read the whole + # tractogram once for discovery and then a second time for visualization. + obj = nib.streamlines.load(str(trk_path), lazy_load=True) tg = obj.tractogram dpp = getattr(tg, "data_per_point", None) if not dpp: @@ -56,6 +60,55 @@ def _resolve_streamlines_path(input_path: Path) -> Path: sys.exit(1) +def _apply_session_settings(args, settings: dict, provided_options: set[str]) -> None: + """Use saved values as defaults while keeping explicit CLI options.""" + option_map = { + "color_by": ("color_by", "--color-by"), + "line_width": ("line_width", "--line-width"), + "subsample_factor": ("subsample", "--subsample"), + "filter_min_len": ("min_length", "--min-length"), + "downsample_factor": ("downsample_factor", "--downsample-factor"), + "max_streamlines": ("max_streamlines", "--max-streamlines"), + "top_clusters": ("top_clusters", "--top-clusters"), + "spline_subdiv": ("spline_subdiv", "--spline-subdiv"), + "colormap": ("colormap", "--colormap"), + "backend": ("backend", "--backend"), + "mode": ("mode", "--mode"), + "background_color": ("background_color", "--background-color"), + "tube_sides": ("tube_sides", "--tube-sides"), + "opacity": ("opacity", "--opacity"), + "quality": ("quality", "--quality"), + "random_seed": ("random_seed", "--random-seed"), + } + for key, (attribute, option) in option_map.items(): + if option not in provided_options and key in settings: + setattr(args, attribute, settings[key]) + + if "--line-width" not in provided_options and "tube_thickness" in settings: + args.line_width = float(settings["tube_thickness"]) + + if not {"--crop-x", "--crop-y", "--crop-z"} & provided_options: + crop_bounds = settings.get("crop_bounds") + if crop_bounds: + args.crop_x, args.crop_y, args.crop_z = crop_bounds + + window_size = settings.get("window_size") + if window_size: + if "--width" not in provided_options: + args.width = int(window_size[0]) + if "--height" not in provided_options: + args.height = int(window_size[1]) + + if "--hide-axes" not in provided_options and "show_axes" in settings: + args.hide_axes = not bool(settings["show_axes"]) + if not {"--show-bounds", "--hide-bounds"} & provided_options: + args.show_bounds = bool(settings.get("show_bounds", False)) + args.hide_bounds = False + if not {"--shadows", "--no-shadows"} & provided_options: + args.shadows = bool(settings.get("shadows", False)) + args.no_shadows = False + + def script() -> None: parser = argparse.ArgumentParser( description="Visualize cardiac streamlines from .trk, color by elevation or stored per-point fields", @@ -77,6 +130,22 @@ def script() -> None: action="store_true", help="List available color-by options from the .trk and exit", ) + parser.add_argument( + "--session", + type=Path, + default=None, + help=( + "FURY session JSON to load and update. When omitted, use " + "streamlines_session.json beside the TRK. Press Ctrl+S in the viewer " + "to save the current camera, zoom, crop gizmos and controls" + ), + ) + parser.add_argument( + "--random-seed", + type=int, + default=None, + help="Seed used for reproducible streamline subsampling", + ) parser.add_argument( "--line-width", type=float, @@ -107,6 +176,15 @@ def script() -> None: default=None, help="Maximum number of streamlines to render after filtering", ) + parser.add_argument( + "--top-clusters", + type=int, + default=None, + help=( + "Keep all streamlines from the N largest clusters; requires " + "cluster_id metadata" + ), + ) parser.add_argument( "--crop-x", nargs=2, @@ -164,8 +242,11 @@ def script() -> None: parser.add_argument( "--spline-subdiv", type=int, - default=2, - help="Spline/display subdivisions for smoother streamlines (higher = smoother, heavier)", + default=None, + help=( + "Number of cardinal-spline segments per streamline; 1 disables " + "smoothing (quality preset default when omitted)" + ), ) parser.add_argument("--width", type=int, default=800, help="Window width in pixels") parser.add_argument( @@ -198,14 +279,20 @@ def script() -> None: parser.add_argument( "--tube-sides", type=int, - default=9, - help="Number of sides for tube geometry", + default=None, + help="Number of sides for tube geometry (quality preset default when omitted)", + ) + parser.add_argument( + "--quality", + choices=("interactive", "publication"), + default="interactive", + help="FURY rendering preset; individual options still override it", ) parser.add_argument( "--opacity", type=float, default=1.0, - help="Streamline opacity for the PyVista backend", + help="Streamline opacity", ) parser.add_argument( "--shadows", @@ -233,11 +320,77 @@ def script() -> None: help=argparse.SUPPRESS, ) + provided_options = { + token.split("=", 1)[0] for token in sys.argv[1:] if token.startswith("--") + } args = parser.parse_args() # Resolve .trk trk_path = _resolve_streamlines_path(args.input_path) + restore_session = args.session is not None + if args.session is None: + args.session = ( + trk_path.resolve().parent / "streamlines_session.json" + ).resolve() + restore_session = False + if args.session.exists() and not args.list_color_by: + if sys.stdin.isatty(): + answer = ( + input( + f"Found default visualization session {args.session}. " + "Restart from it? [Y/n] " + ) + .strip() + .lower() + ) + restore_session = answer in {"", "y", "yes"} + else: + print( + f"Found default visualization session {args.session}, but input " + "is not interactive. Starting a fresh view; pass --session " + "explicitly to restore it." + ) + elif not args.list_color_by: + print(f"Visualization session will be saved by default to {args.session}") + else: + args.session = args.session.expanduser().resolve() + + if not args.list_color_by: + if args.session.exists() and restore_session: + try: + session_data = json.loads(args.session.read_text()) + saved_file = session_data.get("streamlines_file") + saved_size = session_data.get("streamlines_size") + if saved_file and Path(saved_file).resolve() != trk_path.resolve(): + print( + "Warning: session was created for a different TRK path: " + f"{saved_file}" + ) + if saved_size and int(saved_size) != trk_path.stat().st_size: + print("Warning: session TRK size differs from the current file.") + _apply_session_settings( + args, session_data.get("settings", {}), provided_options + ) + print(f"Loaded visualization settings from {args.session}") + except (OSError, ValueError, TypeError, json.JSONDecodeError) as err: + print(f"Warning: could not load session {args.session}: {err}") + restore_session = False + elif args.session.exists(): + print(f"Starting a fresh visualization session at {args.session}") + elif "--session" in provided_options: + print(f"New visualization session will be saved to {args.session}") + + if args.top_clusters is not None and args.top_clusters < 1: + parser.error("--top-clusters must be at least 1") + + if args.random_seed is None: + args.random_seed = random.SystemRandom().randrange(2**63) + if args.spline_subdiv is None: + args.spline_subdiv = 1 if args.quality == "interactive" else 2 + if args.tube_sides is None: + args.tube_sides = 6 if args.quality == "interactive" else 12 + # List available color-by and exit if requested available_dpp = _discover_trk_color_fields(trk_path) computed_color_options = ["elevation", "azimuth", "az", "el"] @@ -249,8 +402,10 @@ def script() -> None: return # Determine color-by - color_by = args.color_by.strip().lower() - dpp_lower = {k.lower(): k for k in available_dpp} # map lower to original key + color_by = args.color_by.strip().lower().replace("-", "_") + dpp_lower = { + key.lower().replace("-", "_"): key for key in available_dpp + } if color_by == "auto": if "ha" in dpp_lower: @@ -298,6 +453,30 @@ def script() -> None: # --video implies off-screen is_interactive = not args.no_interactive and not args.video + session_settings = { + "color_by": selected_color_by, + "line_width": args.line_width, + "subsample_factor": args.subsample, + "filter_min_len": args.min_length, + "downsample_factor": args.downsample_factor, + "max_streamlines": args.max_streamlines, + "top_clusters": args.top_clusters, + "crop_bounds": crop_bounds, + "window_size": [args.width, args.height], + "colormap": args.colormap, + "spline_subdiv": args.spline_subdiv, + "backend": args.backend, + "mode": args.mode, + "background_color": args.background_color, + "tube_sides": args.tube_sides, + "opacity": args.opacity, + "quality": args.quality, + "show_axes": not args.hide_axes, + "show_bounds": args.show_bounds and not args.hide_bounds, + "shadows": args.shadows and not args.no_shadows, + "random_seed": args.random_seed, + } + # Call the visualizer visualize_streamlines( streamlines_file=trk_path, @@ -307,6 +486,7 @@ def script() -> None: filter_min_len=args.min_length, downsample_factor=args.downsample_factor, max_streamlines=args.max_streamlines, + top_clusters=args.top_clusters, crop_bounds=crop_bounds, interactive=is_interactive, screenshot_path=args.screenshot, @@ -324,6 +504,11 @@ def script() -> None: pyvista_show_axes=not args.hide_axes, pyvista_show_bounds=args.show_bounds and not args.hide_bounds, pyvista_shadows=args.shadows and not args.no_shadows, + random_seed=args.random_seed, + session_path=args.session, + restore_session=restore_session, + session_settings=session_settings, + fury_quality=args.quality, ) diff --git a/src/cardiotensor/scripts/visualize_vector_field.py b/src/cardiotensor/scripts/visualize_vector_field.py index 002c069..7050af9 100644 --- a/src/cardiotensor/scripts/visualize_vector_field.py +++ b/src/cardiotensor/scripts/visualize_vector_field.py @@ -17,6 +17,7 @@ import matplotlib.pyplot as plt from cardiotensor.colormaps.helix_angle import helix_angle_cmap +from cardiotensor.utils.image_io import vector_field_path from cardiotensor.utils.utils import read_conf_file from cardiotensor.visualization.vector_field import visualize_vector_field @@ -24,7 +25,7 @@ def _discover_color_folders(output_dir: Path) -> list[str]: candidates = [] for path in sorted(output_dir.iterdir()): - if not path.is_dir() or path.name == "eigen_vec": + if not path.is_dir() or path.name in {"eigen_vec", "eigen_vec.zarr"}: continue try: if any( @@ -143,7 +144,9 @@ def script(): # Call visualization function (handles both plotting and VTK export) visualize_vector_field( - vector_field_path=output_dir / "eigen_vec", + vector_field_path=vector_field_path( + output_dir, params.get("VECTOR_FORMAT", "zarr") + ), color_volume_path=output_dir / color_volume_dir, mask_path=mask_path, downsample=args.downsample, diff --git a/src/cardiotensor/tractography/generate_streamlines.py b/src/cardiotensor/tractography/generate_streamlines.py index 34fc176..bb5d751 100644 --- a/src/cardiotensor/tractography/generate_streamlines.py +++ b/src/cardiotensor/tractography/generate_streamlines.py @@ -11,6 +11,36 @@ from cardiotensor.utils.am_utils import write_spatialgraph_am from cardiotensor.utils.DataReader import DataReader from cardiotensor.utils.downsampling import downsample_vector_volume, downsample_volume +from cardiotensor.utils.image_io import open_zarr_vector_field + + +def _fa_to_unit_interval(volume: np.ndarray) -> np.ndarray: + """Return FA as float32 in [0, 1] from scientific or integer image data.""" + source = np.asarray(volume) + fa = source.astype(np.float32, copy=False) + + if np.issubdtype(source.dtype, np.integer): + fa = fa / float(np.iinfo(source.dtype).max) + elif fa.size and np.nanmax(fa) > 1.0: + # Legacy float arrays may contain the pipeline's 0..255 image encoding. + fa = fa / 255.0 + + return np.clip(fa, 0.0, 1.0) + + +def _select_seed_points( + seed_mask: np.ndarray, num_seeds: int, random_seed: int +) -> np.ndarray: + """Select reproducible seed coordinates without storing all 3D coordinates.""" + if num_seeds <= 0: + raise ValueError("num_seeds must be > 0") + valid_flat = np.flatnonzero(np.asarray(seed_mask).ravel()) + if valid_flat.size == 0: + raise RuntimeError("No voxels above FA seed threshold") + if valid_flat.size > num_seeds: + rng = np.random.default_rng(random_seed) + valid_flat = rng.choice(valid_flat, size=num_seeds, replace=False) + return np.column_stack(np.unravel_index(valid_flat, seed_mask.shape)) @njit(cache=True) @@ -356,6 +386,13 @@ def generate_streamlines_from_vector_field( min_length_pts: int = 10, bidirectional: bool = True, ) -> list[list[tuple[float, float, float]]]: + if min_length_pts < 2: + print( + f"Minimum streamline length raised from {min_length_pts} to 2 points " + "because exported streamlines must contain an edge." + ) + min_length_pts = 2 + all_streamlines: list[list[tuple[float, float, float]]] = [] with alive_bar(len(seed_points), title="Tracing Streamlines") as bar: for zi, yi, xi in seed_points: @@ -475,6 +512,7 @@ def generate_streamlines_from_params( voxel_sizes_zyx: tuple[float, float, float] = (1.0, 1.0, 1.0), save_trk_file: bool = True, angle_mode: str = "ha_ia", + random_seed: int = 0, ) -> None: """ Generate streamlines from the eigenvector field, then export: @@ -543,10 +581,34 @@ def generate_streamlines_from_params( end_y = full_shape[2] if end_y is None else end_y end_x = full_shape[3] if end_x is None else end_x + incomplete = [] + if getattr(vec_probe, "volume_info", {}).get("type") == "zarr": + vector_store = open_zarr_vector_field(vector_field_dir.parent) + completed = np.asarray(vector_store.completed[:], dtype=bool) + completed_count = int(completed.sum()) + if completed_count != full_shape[1]: + incomplete.append( + f"vectors: {completed_count}/{full_shape[1]} completed slices" + ) + + scalar_paths = {"FA": fa_dir, **discovered} + for name, path in scalar_paths.items(): + scalar_depth = DataReader(path).shape[0] + if scalar_depth != full_shape[1]: + incomplete.append(f"{name}: {scalar_depth}/{full_shape[1]} output slices") + + if incomplete: + details = "\n - ".join(incomplete) + raise RuntimeError( + "Orientation outputs are incomplete. Wait for all cardio-tensor/SLURM " + f"jobs to finish before generating streamlines:\n - {details}" + ) + # Binning if bin_factor > 1: - downsample_vector_volume(vector_field_dir, bin_factor, output_dir) - vec_load_dir = output_dir / f"bin{bin_factor}" / vector_field_dir.name + vec_load_dir = downsample_vector_volume( + vector_field_dir, bin_factor, output_dir, mask_path=mask_path + ) downsample_volume(fa_dir, bin_factor, output_dir, subfolder="FA", out_ext="tif") fa_load_dir = output_dir / f"bin{bin_factor}" / "FA" @@ -573,14 +635,20 @@ def generate_streamlines_from_params( # Load vector field print("Loading vector field") vec_reader = DataReader(vec_load_dir) - vector_field = vec_reader.load_volume(start_index=start_z_b, end_index=end_z_b)[ - :, :, start_y_b:end_y_b, start_x_b:end_x_b - ] + vector_field = vec_reader.load_region( + start_index=start_z_b, + end_index=end_z_b, + start_y=start_y_b, + end_y=end_y_b, + start_x=start_x_b, + end_x=end_x_b, + ) if vector_field.ndim == 4 and vector_field.shape[-1] == 3: print("Reordering vector field axes") vector_field = np.moveaxis(vector_field, -1, 0) # Mask + mask = None if mask_path: print("Loading mask") mask_reader = DataReader(mask_path) @@ -599,20 +667,14 @@ def generate_streamlines_from_params( start_index=start_z_b, end_index=end_z_b ) fa_volume = fa_volume[:, start_y_b:end_y_b, start_x_b:end_x_b] + fa_volume = _fa_to_unit_interval(fa_volume) # Seeds print("Selecting seeds") - seed_mask = fa_volume > (fa_seed_min * 255) - valid_indices = np.argwhere(seed_mask) - if valid_indices.size == 0: - raise RuntimeError("No voxels above FA seed threshold") - chosen = ( - valid_indices - if len(valid_indices) <= num_seeds - else valid_indices[ - np.random.choice(valid_indices.shape[0], num_seeds, replace=False) - ] - ) + seed_mask = fa_volume > fa_seed_min + if mask is not None: + seed_mask &= mask > 0 + chosen = _select_seed_points(seed_mask, num_seeds, random_seed) # Streamlines streamlines = generate_streamlines_from_vector_field( diff --git a/src/cardiotensor/utils/DataReader.py b/src/cardiotensor/utils/DataReader.py index 492cf21..9c38808 100644 --- a/src/cardiotensor/utils/DataReader.py +++ b/src/cardiotensor/utils/DataReader.py @@ -1,17 +1,32 @@ -import sys +from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from os import PathLike from pathlib import Path +from tempfile import TemporaryDirectory from typing import Any import numpy as np -import psutil -import SimpleITK as sitk +import tifffile +import zarr from alive_progress import alive_bar -from dask import compute, delayed from skimage.measure import block_reduce -from cardiotensor.utils.image_io import read_image_file -from cardiotensor.utils.utils import get_available_cpu_count +from cardiotensor.utils.image_io import _validate_zarr_vector_group, read_image_file +from cardiotensor.utils.utils import ( + get_available_cpu_count, + get_available_memory_bytes, +) + + +_MHD_DTYPES = { + "MET_FLOAT": np.float32, + "MET_DOUBLE": np.float64, + "MET_CHAR": np.int8, + "MET_UCHAR": np.uint8, + "MET_SHORT": np.int16, + "MET_USHORT": np.uint16, + "MET_INT": np.int32, + "MET_UINT": np.uint32, +} # --------------------------- @@ -122,8 +137,17 @@ def _get_volume_info(self) -> dict: if not self.path.exists(): raise ValueError(f"The path does not exist: {self.path}") - # Case 1: Directory of images - if self.path.is_dir(): + # Case 1: Cardiotensor Zarr vector field + if self.path.is_dir() and (self.path / "zarr.json").is_file(): + group = zarr.open_group(store=str(self.path), mode="r") + _validate_zarr_vector_group(group) + vectors = group["vectors"] + volume_info["type"] = "zarr" + volume_info["shape"] = tuple(vectors.shape) + volume_info["dtype"] = np.dtype(vectors.dtype) + + # Case 2: Directory of images + elif self.path.is_dir(): volume_info["stack"] = True image_files = { ext: sorted(self.path.glob(f"*.{ext}")) @@ -168,20 +192,41 @@ def _get_volume_info(self) -> dict: first_image.shape[1], ) - # Case 2: Single MHD file - elif self.path.is_file() and self.path.suffix == ".mhd": + # Case 3: Single MHD file + elif self.path.is_file() and self.path.suffix.lower() == ".mhd": volume_info["type"] = "mhd" - img = sitk.ReadImage(str(self.path)) - arr = sitk.GetArrayFromImage(img) # Z, Y, X - volume_info["shape"] = arr.shape - volume_info["dtype"] = arr.dtype - - # Case 3: Single TIFF file storing a stack + metadata = _read_mhd(self.path) + if int(metadata.get("NDims", 0)) != 3: + raise ValueError("Only 3D MHD volumes are supported") + try: + dtype = np.dtype(_MHD_DTYPES[metadata["ElementType"]]) + except KeyError as exc: + raise NotImplementedError( + f"ElementType {metadata.get('ElementType')} is not supported" + ) from exc + byte_order = ">" if metadata.get("BinaryDataByteOrderMSB", False) else "<" + dtype = dtype.newbyteorder(byte_order) + shape = tuple(reversed(metadata["DimSize"])) + channels = int(metadata.get("ElementNumberOfChannels", 1)) + volume_info["shape"] = shape if channels == 1 else (*shape, channels) + volume_info["dtype"] = dtype + + # Case 4: Single TIFF file storing a stack elif self.path.is_file() and self.path.suffix.lower() in {".tif", ".tiff"}: volume_info["type"] = "tiff" - arr = _normalize_single_tiff_volume(self._custom_image_reader(self.path)) - volume_info["shape"] = arr.shape - volume_info["dtype"] = arr.dtype + with tifffile.TiffFile(self.path) as tif: + series = tif.series[0] + shape = tuple(series.shape) + dtype = np.dtype(series.dtype) + if len(shape) == 2: + shape = (1, *shape) + elif len(shape) != 3 or shape[-1] in {3, 4}: + raise ValueError( + f"Unsupported TIFF volume shape {shape}. " + "Expected a 2D image or 3D grayscale stack." + ) + volume_info["shape"] = shape + volume_info["dtype"] = dtype else: raise ValueError(f"Unsupported volume type for path: {self.path}") @@ -193,6 +238,7 @@ def load_volume( start_index: int = 0, end_index: int | None = None, unbinned_shape: tuple[int, int, int] | None = None, + show_progress: bool = True, ) -> np.ndarray: """ Loads the volume and resizes it to unbinned_shape if provided, using fast @@ -204,12 +250,13 @@ def load_volume( start_index (int): Start index for slicing (for stacks). end_index (int): End index for slicing (for stacks). If None, loads the entire stack. unbinned_shape (tuple): Desired shape (Z, Y, X). If None, no resizing is done. + show_progress: Display the stack-loading progress bar. Returns: np.ndarray: Loaded volume. """ if end_index is None: - end_index = self.shape[0] + end_index = self.shape[1] if len(self.shape) == 4 else self.shape[0] # Check memory available is enough effective_shape = list(self.shape) @@ -217,14 +264,17 @@ def load_volume( effective_shape[0] = end_index - start_index elif len(effective_shape) == 4: effective_shape[1] = end_index - start_index - self.check_memory_requirement(tuple(effective_shape), self.dtype) + self.check_memory_requirement( + tuple(effective_shape), self.dtype, verbose=show_progress + ) # Decide if resize is needed need_resize = False if unbinned_shape is not None and self.shape != unbinned_shape: need_resize = True zoom_factors = tuple(u / s for u, s in zip(unbinned_shape, self.shape)) - print(f"Resample factors: {zoom_factors}") + if show_progress: + print(f"Resample factors: {zoom_factors}") else: zoom_factors = (1.0, 1.0, 1.0) @@ -234,27 +284,65 @@ def load_volume( start_index = max(start_index, 0) end_index = int(end_index_ini / zoom_factors[0]) + 1 end_index = min(end_index, self.shape[0]) - print(f"Volume start index padded: {start_index} - end: {end_index}") + if show_progress: + print(f"Volume start index padded: {start_index} - end: {end_index}") # Load volume from stack or mhd if not self.volume_info["stack"]: - if self.volume_info["type"] == "mhd": - volume, _ = _load_raw_data_with_mhd(self.path) - volume = volume[start_index:end_index, :, :] - elif self.volume_info["type"] == "tiff": - volume = _normalize_single_tiff_volume( - self._custom_image_reader(self.path) + if self.volume_info["type"] == "zarr": + group = zarr.open_group(store=str(self.path), mode="r") + volume = np.asarray(group["vectors"][:, start_index:end_index]) + elif self.volume_info["type"] == "mhd": + volume, _ = _load_raw_data_with_mhd( + self.path, start_index=start_index, end_index=end_index ) - volume = volume[start_index:end_index, :, :] + elif self.volume_info["type"] == "tiff": + with tifffile.TiffFile(self.path) as tif: + series = tif.series[0] + if len(series.pages) == self.shape[0]: + volume = np.array( + _normalize_single_tiff_volume( + tif.asarray( + key=slice(start_index, end_index), series=series + ) + ), + copy=True, + order="C", + ) + else: + scratch_root = self.path.parent / ".cardiotensor_scratch" + scratch_root.mkdir(parents=True, exist_ok=True) + with TemporaryDirectory( + prefix="cardiotensor_tiff_", dir=scratch_root + ) as tmpdir: + mapped = series.asarray( + out=str(Path(tmpdir) / "volume.memmap") + ) + try: + volume = np.array( + _normalize_single_tiff_volume( + mapped[start_index:end_index] + ), + copy=True, + order="C", + ) + finally: + mapped_file = getattr(mapped, "_mmap", None) + if mapped_file is not None: + mapped_file.close() else: raise ValueError(f"Unsupported volume type for path: {self.path}") else: volume = self._load_image_stack( - self.volume_info["file_list"], start_index, end_index + self.volume_info["file_list"], + start_index, + end_index, + show_progress=show_progress, ) if need_resize: - print("Resizing with integer-only resampling...") + if show_progress: + print("Resizing with integer-only resampling...") _, y1, x1 = unbinned_shape z1 = end_index_ini - start_index_ini @@ -318,16 +406,68 @@ def _fmt(k, d): return f"/{d}" return "x1" - print( - f"Applied integer resampling: Z {_fmt(kz, dz)}, Y {_fmt(ky, dy)}, X {_fmt(kx, dx)} -> resulting shape {volume.shape}" - ) + if show_progress: + print( + f"Applied integer resampling: Z {_fmt(kz, dz)}, " + f"Y {_fmt(ky, dy)}, X {_fmt(kx, dx)} -> " + f"resulting shape {volume.shape}" + ) # Enforce exact shape - print(f"Fitting to exact unbinned shape: {unbinned_shape}") + if show_progress: + print(f"Fitting to exact unbinned shape: {unbinned_shape}") volume = _fit(volume, (z1, y1, x1), pad_value=0) return volume + def load_region( + self, + start_index: int = 0, + end_index: int | None = None, + start_y: int = 0, + end_y: int | None = None, + start_x: int = 0, + end_x: int | None = None, + ) -> np.ndarray: + """Load a Z/Y/X region, using direct chunk reads for Zarr vectors.""" + is_vector = len(self.shape) == 4 and self.shape[0] == 3 + depth = self.shape[1] if is_vector else self.shape[0] + height = self.shape[-2] + width = self.shape[-1] + end_index = depth if end_index is None else end_index + + if not 0 <= start_index < end_index <= depth: + raise ValueError( + f"Invalid Z range [{start_index}, {end_index}) for depth {depth}" + ) + + y_slice = slice(start_y, end_y) + x_slice = slice(start_x, end_x) + y_start, y_stop, y_step = y_slice.indices(height) + x_start, x_stop, x_step = x_slice.indices(width) + if y_step != 1 or x_step != 1 or y_start >= y_stop or x_start >= x_stop: + raise ValueError("Y and X regions must be non-empty contiguous ranges") + + if self.volume_info["type"] == "zarr": + region_shape = ( + 3, + end_index - start_index, + y_stop - y_start, + x_stop - x_start, + ) + self.check_memory_requirement(region_shape, self.dtype) + group = zarr.open_group(store=str(self.path), mode="r") + return np.asarray( + group["vectors"][ + :, start_index:end_index, y_start:y_stop, x_start:x_stop + ] + ) + + volume = self.load_volume(start_index=start_index, end_index=end_index) + if is_vector: + return volume[:, :, y_start:y_stop, x_start:x_stop] + return volume[:, y_start:y_stop, x_start:x_stop] + def _custom_image_reader(self, file_path: Path) -> np.ndarray: """ Reads an image from the given file path into a NumPy array. @@ -335,12 +475,16 @@ def _custom_image_reader(self, file_path: Path) -> np.ndarray: return read_image_file(file_path) def _load_image_stack( - self, file_list: list[Path], start_index: int, end_index: int + self, + file_list: list[Path], + start_index: int, + end_index: int, + show_progress: bool = True, ) -> np.ndarray: """ Efficiently loads a stack of images into a 3D (or 4D for vector .npy) NumPy array - by preallocating the output and filling it in place. Uses batched dask - multiprocess reads and keeps slice order via indices. + by preallocating the output and filling it in place. The number of + in-flight reads is bounded so loaded slices are not duplicated in memory. """ if end_index == 0: @@ -409,30 +553,54 @@ def _assign(z_idx: int, arr: np.ndarray): ) volume[z_idx, :, :] = arr - scheduler = "threads" - max_workers = min(32, get_available_cpu_count(default=8)) + # Keep concurrent decoded slices within a small, predictable buffer. + # A single very large slice still gets one reader. + read_buffer_bytes = 2 * 1024**3 + memory_workers = max(1, read_buffer_bytes // max(first.nbytes, 1)) + max_workers = min( + 32, + get_available_cpu_count(default=8), + memory_workers, + ) # Fill the rest with a progress bar - with alive_bar(total_files, title="Loading Volume", length=40) as bar: + with alive_bar( + total_files, + title="Loading Volume", + length=40, + disable=not show_progress, + ) as bar: # We already placed index 0 bar() # account for the first already loaded slice if start_fill_idx < total_files: - delayed_reads = [ - delayed(read_image_file)(path) for path in paths[start_fill_idx:] - ] - arrays = compute( - *delayed_reads, - scheduler=scheduler, - num_workers=max_workers, - ) - for z_idx, arr in enumerate(arrays, start=start_fill_idx): - _assign(z_idx, arr) - bar() + indexed_paths = iter(enumerate(paths[start_fill_idx:], start_fill_idx)) + with ThreadPoolExecutor(max_workers=max_workers) as executor: + pending = {} + for _ in range(max_workers): + try: + z_idx, path = next(indexed_paths) + except StopIteration: + break + pending[executor.submit(read_image_file, path)] = z_idx + + while pending: + done, _ = wait(pending, return_when=FIRST_COMPLETED) + for future in done: + z_idx = pending.pop(future) + _assign(z_idx, future.result()) + bar() + try: + next_idx, next_path = next(indexed_paths) + except StopIteration: + continue + pending[executor.submit(read_image_file, next_path)] = ( + next_idx + ) return volume - def check_memory_requirement(self, shape, dtype, safety_factor=0.8): + def check_memory_requirement(self, shape, dtype, safety_factor=0.8, verbose=True): """ Check if the dataset can fit in available memory. @@ -446,15 +614,19 @@ def check_memory_requirement(self, shape, dtype, safety_factor=0.8): size_gb = n_bytes / (1024**3) # Check available memory - available_gb = psutil.virtual_memory().available / (1024**3) + available_gb = get_available_memory_bytes() / (1024**3) - print( - f"Dataset size: {size_gb:.2f} GB | Available memory: {available_gb:.2f} GB" - ) + if verbose: + print( + f"Dataset size: {size_gb:.2f} GB | " + f"Available memory: {available_gb:.2f} GB" + ) if size_gb > available_gb * safety_factor: - print("❌ Dataset is too large to safely load into memory.") - sys.exit(1) + raise MemoryError( + f"Dataset requires {size_gb:.2f} GB but only " + f"{available_gb:.2f} GB is available" + ) def _read_mhd(filename: PathLike[str]) -> dict[str, Any]: @@ -524,6 +696,8 @@ def _read_mhd(filename: PathLike[str]) -> dict[str, Any]: def _load_raw_data_with_mhd( filename: PathLike[str], + start_index: int = 0, + end_index: int | None = None, ) -> tuple[np.ndarray, dict[str, Any]]: """ Load a MHD file @@ -532,59 +706,41 @@ def _load_raw_data_with_mhd( :returns: tuple with raw data and dictionary of meta data """ meta_dict = _read_mhd(filename) - dim = int(meta_dict["NDims"]) - if "ElementNumberOfChannels" in meta_dict: - element_channels = int(meta_dict["ElementNumberOfChannels"]) - else: - element_channels = 1 - - if meta_dict["ElementType"] == "MET_FLOAT": - np_type = np.float32 - elif meta_dict["ElementType"] == "MET_DOUBLE": - np_type = np.float64 - elif meta_dict["ElementType"] == "MET_CHAR": - np_type = np.byte - elif meta_dict["ElementType"] == "MET_UCHAR": - np_type = np.ubyte - elif meta_dict["ElementType"] == "MET_SHORT": - np_type = np.int16 - elif meta_dict["ElementType"] == "MET_USHORT": - np_type = np.ushort - elif meta_dict["ElementType"] == "MET_INT": - np_type = np.int32 - elif meta_dict["ElementType"] == "MET_UINT": - np_type = np.uint32 - else: + if int(meta_dict.get("NDims", 0)) != 3: + raise ValueError("Only 3D MHD volumes are supported") + if meta_dict.get("CompressedData", False): + raise NotImplementedError("Compressed MHD data cannot be memory-mapped") + if not meta_dict.get("BinaryData", True): + raise NotImplementedError("ASCII MHD data is not supported") + + try: + np_type = np.dtype(_MHD_DTYPES[meta_dict["ElementType"]]) + except KeyError as exc: raise NotImplementedError( - "ElementType " + meta_dict["ElementType"] + " not understood." - ) - arr = list(meta_dict["DimSize"]) - - volume = np.prod(arr[0 : dim - 1]) + f"ElementType {meta_dict.get('ElementType')} is not supported" + ) from exc + byte_order = ">" if meta_dict.get("BinaryDataByteOrderMSB", False) else "<" + np_type = np_type.newbyteorder(byte_order) + shape = tuple(reversed(meta_dict["DimSize"])) + element_channels = int(meta_dict.get("ElementNumberOfChannels", 1)) + if element_channels > 1: + shape = (*shape, element_channels) pwd = Path(filename).parents[0].resolve() data_file = Path(meta_dict["ElementDataFile"]) + if str(data_file).upper() == "LOCAL": + raise NotImplementedError("MHD files with embedded LOCAL data are unsupported") if not data_file.is_absolute(): data_file = pwd / data_file - shape = (arr[dim - 1], volume, element_channels) - with open(data_file, "rb") as f: - data = np.fromfile(f, count=np.prod(shape), dtype=np_type) - data.shape = shape - - # Adjust byte order in numpy array to match default system byte order - if "BinaryDataByteOrderMSB" in meta_dict: - sys_byteorder_msb = sys.byteorder == "big" - file_byteorder_ms = meta_dict["BinaryDataByteOrderMSB"] - if sys_byteorder_msb != file_byteorder_ms: - data = data.byteswap() - - # Begin 3D fix - # arr.reverse() - if element_channels > 1: - data = data.reshape(arr + [element_channels]) - else: - data = data.reshape(arr) - # End 3D fix + depth = shape[0] + end_index = depth if end_index is None else end_index + if not 0 <= start_index < end_index <= depth: + raise ValueError( + f"Invalid MHD Z range [{start_index}, {end_index}) for depth {depth}" + ) - return (data, meta_dict) + # Copy-on-write keeps reads lazy and lets masking modify pages in memory + # without changing the source RAW file. + data = np.memmap(data_file, dtype=np_type, mode="c", shape=shape) + return data[start_index:end_index], meta_dict diff --git a/src/cardiotensor/utils/downsampling.py b/src/cardiotensor/utils/downsampling.py index 1f7cb3e..e3fb4ab 100644 --- a/src/cardiotensor/utils/downsampling.py +++ b/src/cardiotensor/utils/downsampling.py @@ -12,6 +12,60 @@ from cardiotensor.utils.utils import convert_to_8bit, get_available_cpu_count +def _downsample_vector_axes(vector_block: np.ndarray, bin_factor: int) -> np.ndarray: + """Downsample unoriented vectors without cancelling equivalent v/-v axes.""" + if vector_block.ndim != 4 or vector_block.shape[0] != 3: + raise ValueError("vector_block must have shape (3, Z, Y, X)") + if bin_factor <= 0: + raise ValueError("bin_factor must be > 0") + + vectors = np.asarray(vector_block, dtype=np.float32) + vectors = np.nan_to_num(vectors, copy=False) + _, depth, height, width = vectors.shape + out_h = math.ceil(height / bin_factor) + out_w = math.ceil(width / bin_factor) + pad_h = out_h * bin_factor - height + pad_w = out_w * bin_factor - width + if pad_h or pad_w: + vectors = np.pad(vectors, ((0, 0), (0, 0), (0, pad_h), (0, pad_w))) + + # Each output pixel is represented by a 3x3 orientation tensor: + # sum(v @ v.T). Its principal eigenvector is invariant to vector sign. + samples = vectors.reshape(3, depth, out_h, bin_factor, out_w, bin_factor).transpose( + 2, 4, 1, 3, 5, 0 + ) + samples = samples.reshape(out_h, out_w, -1, 3) + valid = np.any(samples, axis=(2, 3)) + axes = np.zeros((out_h, out_w, 3), dtype=np.float32) + if not np.any(valid): + return np.moveaxis(axes, -1, 0) + + all_valid = np.all(valid) + active_samples = samples if all_valid else samples[valid] + tensors = np.einsum( + "...ni,...nj->...ij", active_samples, active_samples, optimize=True + ) + eigenvalues, eigenvectors = np.linalg.eigh(tensors) + active_axes = eigenvectors[..., -1].astype(np.float32, copy=False) + + # Choose a deterministic polarity for storage. The physical orientation + # remains an axis, so this does not change the result. + eps = np.finfo(np.float32).eps + z_zero = np.abs(active_axes[..., 2]) <= eps + y_zero = np.abs(active_axes[..., 1]) <= eps + flip = active_axes[..., 2] < -eps + flip |= z_zero & (active_axes[..., 1] < -eps) + flip |= z_zero & y_zero & (active_axes[..., 0] < 0) + active_axes[flip] *= -1 + active_axes[eigenvalues[..., -1] <= eps] = 0 + if all_valid: + axes = active_axes + else: + axes[valid] = active_axes + + return np.moveaxis(axes, -1, 0) + + def process_vector_block( block: list[Path], bin_factor: int, @@ -40,34 +94,86 @@ def process_vector_block( return array = np.empty((3, len(block), h, w), dtype=np.float32) - bin_array = np.empty( - (3, math.ceil(h / bin_factor), math.ceil(w / bin_factor)), dtype=np.float32 - ) - for i, p in enumerate(block): array[:, i, :, :] = np.load(p) - array = array.mean(axis=1) - - block_size = (bin_factor, bin_factor) - for comp in range(3): - bin_array[comp] = block_reduce( - array[comp], block_size=block_size, func=np.mean - ) - - np.save(output_file, bin_array.astype(np.float32)) + bin_array = _downsample_vector_axes(array, bin_factor) + np.save(output_file, bin_array) except Exception as e: print(f"Error processing block {idx}: {e}") # print(f"Failed to process files: {[str(p) for p in block]}") - raise e + raise + + +def _process_vector_range( + vector_reader: DataReader, + start_index: int, + end_index: int, + bin_factor: int, + output_dir: Path, + output_index: int, + mask_reader: DataReader | None = None, + mask_y_indices: np.ndarray | None = None, + mask_x_indices: np.ndarray | None = None, + overwrite: bool = False, +) -> None: + """Load and downsample one Z range from any DataReader vector backend.""" + output_file = output_dir / f"eigen_vec_{output_index:06d}.npy" + if output_file.exists() and not overwrite: + return + + mask = None + if mask_reader is not None: + vector_depth = vector_reader.shape[1] + mask_depth = mask_reader.shape[0] + mask_z_indices = ( + np.arange(start_index, end_index, dtype=np.int64) + * mask_depth + // vector_depth + ) + mask_start = int(mask_z_indices[0]) + mask = mask_reader.load_volume( + mask_start, + int(mask_z_indices[-1]) + 1, + show_progress=False, + ) + mask = np.take(mask, mask_y_indices, axis=1) + mask = np.take(mask, mask_x_indices, axis=2) + mask = mask[mask_z_indices - mask_start] + expected_shape = ( + end_index - start_index, + vector_reader.shape[2], + vector_reader.shape[3], + ) + if mask.shape != expected_shape: + raise ValueError( + f"Mask block shape {mask.shape} does not match vector block " + f"shape {expected_shape}" + ) + if not np.any(mask): + output_shape = ( + 3, + math.ceil(vector_reader.shape[2] / bin_factor), + math.ceil(vector_reader.shape[3] / bin_factor), + ) + np.save(output_file, np.zeros(output_shape, dtype=np.float32)) + return + + vector_block = vector_reader.load_volume( + start_index, end_index, show_progress=False + ) + if mask is not None: + vector_block[:, mask <= 0] = 0 + np.save(output_file, _downsample_vector_axes(vector_block, bin_factor)) def downsample_vector_volume( input_npy: Path, bin_factor: int, output_dir: Path, -) -> None: + mask_path: str | Path | None = None, +) -> Path: """ Downsamples a vector volume using a thread pool. @@ -75,21 +181,57 @@ def downsample_vector_volume( input_npy (Path): Path to the directory containing numpy files. bin_factor (int): Binning factor for downsampling. output_dir (Path): Path to the output directory. + mask_path: Optional mask applied before vector downsampling. """ bin_dir = output_dir / f"bin{bin_factor}" eig_out_dir = bin_dir / "eigen_vec" os.makedirs(eig_out_dir, exist_ok=True) - npy_list = sorted(input_npy.glob("*.npy")) - if len(npy_list) == 0: - return - - # Determine block count - blocks = [npy_list[i : i + bin_factor] for i in range(0, len(npy_list), bin_factor)] + input_npy = Path(input_npy) + mask_path = None if mask_path is None else Path(mask_path) + reader = DataReader(input_npy) + if len(reader.shape) != 4 or reader.shape[0] != 3: + raise ValueError( + f"Expected vector field shape (3, Z, Y, X), got {reader.shape}" + ) + total_slices = reader.shape[1] + mask_reader = None + mask_y_indices = None + mask_x_indices = None + if mask_path is not None: + mask_reader = DataReader(mask_path) + if len(mask_reader.shape) != 3: + raise ValueError(f"Expected a 3D mask, got shape {mask_reader.shape}") + _, vector_height, vector_width = reader.shape[1:] + _, mask_height, mask_width = mask_reader.shape + mask_y_indices = ( + np.arange(vector_height, dtype=np.int64) * mask_height // vector_height + ) + mask_x_indices = ( + np.arange(vector_width, dtype=np.int64) * mask_width // vector_width + ) + print( + f"Applying mask shape {mask_reader.shape} to vector field " + f"shape {reader.shape[1:]} before binning" + ) + blocks = [ + (start, min(start + bin_factor, total_slices)) + for start in range(0, total_slices, bin_factor) + ] total_blocks = len(blocks) - # Quick check: if all expected output files already exist, skip processing - all_exist = True + # Do not reuse vectors produced with a different mask. A missing marker also + # invalidates caches created by older Cardiotensor versions. + mask_source = ( + "mask-index-v3:none" + if mask_path is None + else f"mask-index-v3:{mask_path.resolve()}" + ) + cache_marker = eig_out_dir / ".mask_source" + cache_matches = ( + cache_marker.exists() and cache_marker.read_text().strip() == mask_source + ) + all_exist = cache_matches for idx in range(total_blocks): expected_file = eig_out_dir / f"eigen_vec_{idx:06d}.npy" if not expected_file.exists(): @@ -97,26 +239,36 @@ def downsample_vector_volume( break if all_exist: print("✅ Downsampled images for eigen_vec already exist. Skipping.") - return - - # Load dimensions from the first npy in each block - sample = np.load(npy_list[0]) - _, h, w = sample.shape + return eig_out_dir tasks = [ - (block, bin_factor, h, w, bin_dir, idx) for idx, block in enumerate(blocks) + ( + reader, + start, + end, + bin_factor, + eig_out_dir, + idx, + mask_reader, + mask_y_indices, + mask_x_indices, + not cache_matches, + ) + for idx, (start, end) in enumerate(blocks) ] with ThreadPool(processes=min(get_available_cpu_count(), 32)) as pool: with alive_bar(len(tasks), title="Downsampling vector volumes") as bar: results = [ pool.apply_async( - process_vector_block, args=task, callback=lambda _: bar() + _process_vector_range, args=task, callback=lambda _: bar() ) for task in tasks ] for result in results: - result.wait() + result.get() + cache_marker.write_text(mask_source + "\n") + return eig_out_dir def process_image_block( @@ -200,11 +352,15 @@ def downsample_volume( out_dir = bin_dir / subfolder out_dir.mkdir(parents=True, exist_ok=True) - # Early exit if all output files already exist + cache_key = f"scalar-v2:{Path(input_path).resolve()}:{reader.shape}:bin{bin_factor}" + cache_marker = out_dir / ".source_complete" + cache_matches = ( + cache_marker.exists() and cache_marker.read_text().strip() == cache_key + ) expected_files = [ out_dir / f"{subfolder}_{i:06d}.{out_ext}" for i in range(num_blocks) ] - if all(f.exists() for f in expected_files): + if cache_matches and all(f.exists() for f in expected_files): print(f"✅ Downsampled images for '{subfolder}' already exist. Skipping.") return @@ -220,7 +376,7 @@ def downsample_volume( tasks = [] for block_idx in range(num_blocks): out_file = out_dir / f"{subfolder}_{block_idx:06d}.{out_ext}" - if not out_file.exists(): + if not cache_matches or not out_file.exists(): tasks.append( (file_list, block_idx, bin_factor, H, W, out_file, min_value, max_value) ) @@ -240,4 +396,5 @@ def downsample_volume( for task in tasks ] for r in results: - r.wait() + r.get() + cache_marker.write_text(cache_key + "\n") diff --git a/src/cardiotensor/utils/image_io.py b/src/cardiotensor/utils/image_io.py index c382775..e47d210 100644 --- a/src/cardiotensor/utils/image_io.py +++ b/src/cardiotensor/utils/image_io.py @@ -1,11 +1,269 @@ +import math import os from collections.abc import Callable +from contextlib import contextmanager +from dataclasses import dataclass from pathlib import Path +from typing import Any import cv2 import glymur import numpy as np import tifffile +import zarr + +VECTOR_FORMATS = {"npy", "zarr"} +ZARR_VECTOR_FORMAT_VERSION = 1 +ZARR_VECTOR_CHUNK_SIZE = 512 +ZARR_VECTOR_SHARD_CHUNKS = 8 + + +def normalize_vector_format(vector_format: str) -> str: + """Validate and normalize a vector output format.""" + normalized = vector_format.lower().strip() + if normalized not in VECTOR_FORMATS: + choices = ", ".join(sorted(VECTOR_FORMATS)) + raise ValueError(f"VECTOR_FORMAT must be one of: {choices}") + return normalized + + +def vector_field_path(output_dir: str | Path, vector_format: str) -> Path: + """Return the vector-field path for one configured storage format.""" + vector_format = normalize_vector_format(vector_format) + name = "eigen_vec" if vector_format == "npy" else "eigen_vec.zarr" + return Path(output_dir) / name + + +def _zarr_vector_layout(height: int, width: int) -> tuple[tuple[int, ...], ...]: + """Return the fixed chunk and shard layout for a vector field.""" + chunk_y = min(ZARR_VECTOR_CHUNK_SIZE, height) + chunk_x = min(ZARR_VECTOR_CHUNK_SIZE, width) + shard_y = chunk_y * min(ZARR_VECTOR_SHARD_CHUNKS, math.ceil(height / chunk_y)) + shard_x = chunk_x * min(ZARR_VECTOR_SHARD_CHUNKS, math.ceil(width / chunk_x)) + return (3, 1, chunk_y, chunk_x), (3, 1, shard_y, shard_x) + + +def _codec_value(value: Any) -> Any: + """Return the plain value stored by a codec enum.""" + return getattr(value, "value", value) + + +def _validate_zarr_vector_group( + group: Any, + volume_shape: tuple[int, int, int] | None = None, + *, + allow_missing_version: bool = False, +) -> None: + """Validate arrays and metadata used by a Cardiotensor vector store.""" + if "vectors" not in group or "completed" not in group: + raise ValueError("Zarr vector store must contain 'vectors' and 'completed'") + + vectors = group["vectors"] + completed = group["completed"] + if len(vectors.shape) != 4 or vectors.shape[0] != 3: + raise ValueError( + f"Zarr vectors must have shape (3, Z, Y, X), got {vectors.shape}" + ) + + z_count, height, width = (int(size) for size in vectors.shape[1:]) + if volume_shape is not None and (z_count, height, width) != tuple(volume_shape): + raise ValueError( + f"Existing Zarr volume shape {(z_count, height, width)} does not match " + f"requested shape {tuple(volume_shape)}" + ) + if np.dtype(vectors.dtype) != np.dtype(np.float32): + raise ValueError(f"Zarr vector dtype must be float32, got {vectors.dtype}") + + expected_chunks, expected_shards = _zarr_vector_layout(height, width) + if tuple(vectors.chunks) != expected_chunks: + raise ValueError( + f"Zarr vector chunks {vectors.chunks} do not match {expected_chunks}" + ) + if vectors.shards is None or tuple(vectors.shards) != expected_shards: + raise ValueError( + f"Zarr vector shards {vectors.shards} do not match {expected_shards}" + ) + if vectors.fill_value != 0: + raise ValueError(f"Zarr vector fill value must be 0, got {vectors.fill_value}") + dimension_names = getattr(vectors.metadata, "dimension_names", None) + if dimension_names is None or tuple(dimension_names) != ( + "component", + "z", + "y", + "x", + ): + raise ValueError("Zarr vector dimensions must be ('component', 'z', 'y', 'x')") + + compressors = vectors.compressors + if len(compressors) != 1 or not isinstance(compressors[0], zarr.codecs.BloscCodec): + raise ValueError("Zarr vectors must use one Blosc compressor") + compressor = compressors[0] + if ( + _codec_value(compressor.cname) != "zstd" + or compressor.clevel != 3 + or _codec_value(compressor.shuffle) != "bitshuffle" + ): + raise ValueError("Zarr vectors must use Blosc Zstd level 3 with bit-shuffle") + + if tuple(completed.shape) != (z_count,): + raise ValueError( + f"Zarr completion shape {completed.shape} does not match {(z_count,)}" + ) + if tuple(completed.chunks) != (1,) or np.dtype(completed.dtype) != np.dtype(bool): + raise ValueError("Zarr completion markers must be bool with chunks=(1,)") + if completed.fill_value not in (False, 0): + raise ValueError("Zarr completion fill value must be False") + + expected_attrs = { + "cardiotensor_format": "vector_field", + "axis_order": ["component", "z", "y", "x"], + "components": ["x", "y", "z"], + "masked_fill_value": 0.0, + } + for name, expected in expected_attrs.items(): + if group.attrs.get(name) != expected: + raise ValueError( + f"Invalid Zarr metadata {name!r}: expected {expected!r}, " + f"got {group.attrs.get(name)!r}" + ) + version = group.attrs.get("cardiotensor_format_version") + if version is None and allow_missing_version: + return + if version != ZARR_VECTOR_FORMAT_VERSION: + raise ValueError( + f"Unsupported Zarr vector format version {version!r}; " + f"expected {ZARR_VECTOR_FORMAT_VERSION}" + ) + + +@contextmanager +def _zarr_metadata_lock(lock_path: Path): + """Serialize Zarr metadata creation across local or SLURM processes.""" + lock_path.parent.mkdir(parents=True, exist_ok=True) + with open(lock_path, "a+b") as lock_file: + try: + import fcntl + except ImportError: # pragma: no cover - Windows fallback + fcntl = None + + if fcntl is not None: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) + try: + yield + finally: + if fcntl is not None: + fcntl.flock(lock_file.fileno(), fcntl.LOCK_UN) + + +@dataclass(frozen=True) +class ZarrVectorFieldStore: + """Open arrays belonging to one Cardiotensor Zarr vector field.""" + + path: Path + vectors: Any + completed: Any + + def write_slice(self, vector_slice: np.ndarray, global_slice_idx: int) -> None: + expected_shape = (3, *self.vectors.shape[2:]) + if vector_slice.shape != expected_shape: + raise ValueError( + f"Vector slice shape {vector_slice.shape} does not match " + f"Zarr slice shape {expected_shape}" + ) + if not 0 <= global_slice_idx < self.vectors.shape[1]: + raise IndexError( + f"Vector slice {global_slice_idx} is outside Zarr bounds " + f"[0, {self.vectors.shape[1]})" + ) + + self.vectors[:, global_slice_idx, :, :] = np.asarray( + vector_slice, dtype=np.float32 + ) + # Completion is written last, so interrupted slices are recomputed. + self.completed[global_slice_idx] = True + + def completed_range(self, start_index: int, end_index: int) -> np.ndarray: + return np.asarray(self.completed[start_index:end_index], dtype=bool) + + +def initialize_zarr_vector_field( + output_dir: str | Path, + volume_shape: tuple[int, int, int], +) -> ZarrVectorFieldStore: + """Create or validate the sharded Zarr vector field for a full volume.""" + z_count, height, width = (int(size) for size in volume_shape) + if min(z_count, height, width) <= 0: + raise ValueError(f"Invalid vector volume shape: {volume_shape}") + + store_path = vector_field_path(output_dir, "zarr") + lock_path = Path(output_dir) / ".eigen_vec.zarr.lock" + vector_shape = (3, z_count, height, width) + chunks, shards = _zarr_vector_layout(height, width) + + with _zarr_metadata_lock(lock_path): + group = zarr.open_group( + store=str(store_path), + mode="a", + zarr_format=3, + ) + existing_store = "vectors" in group and "completed" in group + if existing_store: + _validate_zarr_vector_group(group, volume_shape, allow_missing_version=True) + + if "vectors" not in group: + vectors = group.create_array( + "vectors", + shape=vector_shape, + chunks=chunks, + shards=shards, + dtype="float32", + fill_value=0.0, + compressors=zarr.codecs.BloscCodec( + cname="zstd", + clevel=3, + shuffle=zarr.codecs.BloscShuffle.bitshuffle, + ), + dimension_names=("component", "z", "y", "x"), + config={"write_empty_chunks": False}, + ) + else: + vectors = group["vectors"] + + if "completed" not in group: + completed = group.create_array( + "completed", + shape=(z_count,), + chunks=(1,), + dtype="bool", + fill_value=False, + config={"write_empty_chunks": False}, + ) + else: + completed = group["completed"] + + group.attrs.update( + { + "cardiotensor_format": "vector_field", + "axis_order": ["component", "z", "y", "x"], + "components": ["x", "y", "z"], + "masked_fill_value": 0.0, + "cardiotensor_format_version": ZARR_VECTOR_FORMAT_VERSION, + } + ) + _validate_zarr_vector_group(group, volume_shape) + + return ZarrVectorFieldStore(store_path, vectors, completed) + + +def open_zarr_vector_field( + output_dir: str | Path, + mode: str = "r", +) -> ZarrVectorFieldStore: + """Open an existing Cardiotensor Zarr vector field.""" + store_path = vector_field_path(output_dir, "zarr") + group = zarr.open_group(store=str(store_path), mode=mode) + _validate_zarr_vector_group(group) + return ZarrVectorFieldStore(store_path, group["vectors"], group["completed"]) def read_image_file(file_path: str | Path) -> np.ndarray: @@ -92,12 +350,25 @@ def write_image(out_path: str | Path, data: np.ndarray) -> None: def write_vector_field( - vector_field_slice: np.ndarray, start_index: int, output_dir: str, slice_idx: int + vector_field_slice: np.ndarray, + start_index: int, + output_dir: str, + slice_idx: int, + vector_format: str = "npy", + zarr_store: ZarrVectorFieldStore | None = None, ) -> None: - """Save one vector field slice as eigen_vec/eigen_vec_XXXXXX.npy.""" - eig_dir = Path(output_dir) / "eigen_vec" + """Save one vector field slice using the configured vector backend.""" + vector_format = normalize_vector_format(vector_format) + global_slice_idx = start_index + slice_idx + if vector_format == "zarr": + if zarr_store is None: + raise ValueError("zarr_store is required when vector_format='zarr'") + zarr_store.write_slice(vector_field_slice, global_slice_idx) + return + + eig_dir = vector_field_path(output_dir, "npy") eig_dir.mkdir(parents=True, exist_ok=True) - out_path = eig_dir / f"eigen_vec_{(start_index + slice_idx):06d}.npy" + out_path = eig_dir / f"eigen_vec_{global_slice_idx:06d}.npy" def _writer(tmp_path: str) -> None: with open(tmp_path, "wb") as f: diff --git a/src/cardiotensor/utils/streamlines_io_utils.py b/src/cardiotensor/utils/streamlines_io_utils.py index e05bc8c..822d13a 100644 --- a/src/cardiotensor/utils/streamlines_io_utils.py +++ b/src/cardiotensor/utils/streamlines_io_utils.py @@ -48,7 +48,8 @@ def load_npz_streamlines( def load_trk_streamlines( p: Path, -) -> tuple[list[np.ndarray], dict[str, list[np.ndarray]]]: + include_per_streamline: bool = False, +): """ Load streamlines and all per-point fields from a TrackVis .trk file. Returns streamlines in (x, y, z) voxel/world space (as stored in the TRK), @@ -91,7 +92,22 @@ def load_trk_streamlines( ) per_point[name] = vals - return streamlines_xyz, per_point + if not include_per_streamline: + return streamlines_xyz, per_point + + per_streamline: dict[str, np.ndarray] = {} + dps = getattr(tg, "data_per_streamline", None) + if dps: + for name, values in dps.items(): + values = np.asarray(values) + if len(values) != len(streamlines_xyz): + raise ValueError( + f"data_per_streamline['{name}'] length {len(values)} does not " + f"match number of streamlines {len(streamlines_xyz)}" + ) + per_streamline[name] = values + + return streamlines_xyz, per_point, per_streamline def ha_to_degrees_per_streamline(ha_list: list[np.ndarray]) -> list[np.ndarray]: @@ -124,22 +140,35 @@ def normalize_attrs_to_degrees(attrs: dict | None) -> dict[str, list[np.ndarray] if not attrs: return {} + angle_ranges = { + "HA": (-90.0, 90.0), + "IA": (-90.0, 90.0), + "AZ": (-180.0, 180.0), + "EL": (0.0, 90.0), + } normalized = {} for key, arr_list in attrs.items(): + name = key.upper() + source_arrays = [np.asarray(arr) for arr in arr_list] + nonempty = [arr.reshape(-1) for arr in source_arrays if arr.size] + byte_scaled = any(arr.dtype == np.uint8 for arr in source_arrays) + if nonempty and name in angle_ranges and not byte_scaled: + field_min = min(float(np.nanmin(arr)) for arr in nonempty) + field_max = max(float(np.nanmax(arr)) for arr in nonempty) + _, physical_max = angle_ranges[name] + byte_scaled = ( + field_min >= 0.0 and field_max <= 255.0 and field_max > physical_max + ) + out_list = [] - for arr in arr_list: + for arr in source_arrays: a = np.asarray(arr, dtype=np.float32).reshape(-1) - - # Detect encoding: - # If input uses uint8 (0-255), convert -> degrees - # Otherwise assume already degrees - if a.dtype == np.uint8 or (np.nanmax(a) > 1.5 and np.nanmax(a) <= 255): - # Convert 0–255 → -90° to +90° convention - a = (a / 255.0) * 180.0 - 90.0 - + if byte_scaled and name in angle_ranges: + angle_min, angle_max = angle_ranges[name] + a = angle_min + (a / 255.0) * (angle_max - angle_min) out_list.append(a.astype(np.float32)) - normalized[key.upper()] = out_list + normalized[name] = out_list return normalized diff --git a/src/cardiotensor/utils/utils.py b/src/cardiotensor/utils/utils.py index 59ff4c1..8b26962 100644 --- a/src/cardiotensor/utils/utils.py +++ b/src/cardiotensor/utils/utils.py @@ -7,6 +7,7 @@ from typing import Any import numpy as np +import psutil def get_available_cpu_count(default: int = 1) -> int: @@ -24,6 +25,97 @@ def get_available_cpu_count(default: int = 1) -> int: return os.cpu_count() or default +def _memory_available_from_cgroup( + proc_cgroup_path: Path = Path("/proc/self/cgroup"), + cgroup_root: Path = Path("/sys/fs/cgroup"), +) -> int | None: + """Return unused bytes in the process memory cgroup, when available.""" + + def read_pair(directory: Path, limit_name: str, usage_name: str) -> int | None: + try: + limit_text = (directory / limit_name).read_text().strip() + usage = int((directory / usage_name).read_text().strip()) + except (OSError, ValueError): + return None + if limit_text == "max": + return None + try: + limit = int(limit_text) + except ValueError: + return None + return max(limit - usage, 0) + + try: + entries = proc_cgroup_path.read_text().splitlines() + except OSError: + entries = [] + + for entry in entries: + parts = entry.split(":", maxsplit=2) + if len(parts) != 3: + continue + _, controllers, relative_path = parts + relative_path = relative_path.lstrip("/") + + if not controllers: # cgroup v2 + available = read_pair( + cgroup_root / relative_path, "memory.max", "memory.current" + ) + if available is not None: + return available + elif "memory" in controllers.split(","): # cgroup v1 + for base in (cgroup_root / "memory", cgroup_root): + available = read_pair( + base / relative_path, + "memory.limit_in_bytes", + "memory.usage_in_bytes", + ) + if available is not None: + return available + + # Common when the process sees a cgroup namespace rooted at its own group. + available = read_pair(cgroup_root, "memory.max", "memory.current") + if available is not None: + return available + return read_pair( + cgroup_root / "memory", + "memory.limit_in_bytes", + "memory.usage_in_bytes", + ) + + +def _memory_available_from_slurm() -> int | None: + """Estimate unused allocation from SLURM variables when cgroups are hidden.""" + memory_mb = os.environ.get("SLURM_MEM_PER_NODE") + if memory_mb is None: + memory_per_cpu_mb = os.environ.get("SLURM_MEM_PER_CPU") + if memory_per_cpu_mb is not None: + try: + memory_mb = str( + int(memory_per_cpu_mb) * get_available_cpu_count(default=1) + ) + except ValueError: + return None + if memory_mb is None: + return None + try: + limit = int(memory_mb) * 1024**2 + except ValueError: + return None + return max(limit - psutil.Process().memory_info().rss, 0) + + +def get_available_memory_bytes() -> int: + """Return free memory visible to this process, respecting job limits.""" + host_available = int(psutil.virtual_memory().available) + job_available = _memory_available_from_cgroup() + if job_available is None: + job_available = _memory_available_from_slurm() + return ( + host_available if job_available is None else min(host_available, job_available) + ) + + def get_gpu_count() -> int: """Return the number of visible NVIDIA GPUs, or 0 if none are detected.""" @@ -100,8 +192,9 @@ def read_conf_file(file_path: str) -> dict[str, Any]: ValueError: If expected numerical or array values are incorrectly formatted. """ - file_path = str(Path(file_path).resolve()) # Ensure the file path is absolute - if not os.path.exists(file_path): + config_path = Path(file_path).resolve() + file_path = str(config_path) + if not config_path.exists(): raise FileNotFoundError(f"The configuration file {file_path} does not exist.") if not file_path.endswith(".conf"): @@ -110,6 +203,12 @@ def read_conf_file(file_path: str) -> dict[str, Any]: config = configparser.ConfigParser() config.read(file_path) + def resolve_path(value: str) -> str: + path = Path(value).expanduser() + if not path.is_absolute(): + path = config_path.parent / path + return str(path.resolve()) + def parse_coordinates(section: str, option: str, fallback: str = ""): """ Parses a coordinate string from a configuration file into a list of tuples. @@ -138,12 +237,12 @@ def parse_coordinates(section: str, option: str, fallback: str = ""): ) from e # Read the two paths - images_path = config.get("DATASET", "IMAGES_PATH").strip() + images_path = resolve_path(config.get("DATASET", "IMAGES_PATH").strip()) mask_path = config.get("DATASET", "MASK_PATH", fallback=None) if mask_path == "": mask_path = None if mask_path is not None: - mask_path = mask_path.strip() + mask_path = resolve_path(mask_path.strip()) # Existence check (file or directory) if not os.path.exists(images_path): @@ -155,15 +254,22 @@ def parse_coordinates(section: str, option: str, fallback: str = ""): output_dir = config.get("OUTPUT", "OUTPUT_PATH", fallback="").strip() if not output_dir: output_dir = "./output" # Default output directory if not specified - - return { + output_dir = resolve_path(output_dir) + low_memory_dir = config.get( + "STRUCTURE TENSOR CALCULATION", "LOW_MEMORY_DIR", fallback="" + ).strip() + low_memory_dir = ( + resolve_path(os.path.expandvars(low_memory_dir)) if low_memory_dir else None + ) + + params = { # DATASET "IMAGES_PATH": images_path, "MASK_PATH": mask_path, "VOXEL_SIZE": config.getfloat("DATASET", "VOXEL_SIZE", fallback=1.0), # STRUCTURE TENSOR CALCULATION - "SIGMA": config.getfloat("STRUCTURE TENSOR CALCULATION", "SIGMA", fallback=3.0), - "RHO": config.getfloat("STRUCTURE TENSOR CALCULATION", "RHO", fallback=1.0), + "SIGMA": config.getfloat("STRUCTURE TENSOR CALCULATION", "SIGMA", fallback=1.0), + "RHO": config.getfloat("STRUCTURE TENSOR CALCULATION", "RHO", fallback=3.0), "TRUNCATE": config.getfloat( "STRUCTURE TENSOR CALCULATION", "TRUNCATE", fallback=4.0 ), @@ -176,9 +282,18 @@ def parse_coordinates(section: str, option: str, fallback: str = ""): "USE_GPU": config.getboolean( "STRUCTURE TENSOR CALCULATION", "USE_GPU", fallback=False ), + "LOW_MEMORY": config.getboolean( + "STRUCTURE TENSOR CALCULATION", "LOW_MEMORY", fallback=False + ), + "LOW_MEMORY_DIR": low_memory_dir, "WRITE_VECTORS": config.getboolean( "STRUCTURE TENSOR CALCULATION", "WRITE_VECTORS", fallback=False ), + "VECTOR_FORMAT": config.get( + "STRUCTURE TENSOR CALCULATION", "VECTOR_FORMAT", fallback="zarr" + ) + .strip() + .lower(), "REVERSE": config.getboolean( "STRUCTURE TENSOR CALCULATION", "REVERSE", fallback=False ), @@ -209,11 +324,15 @@ def parse_coordinates(section: str, option: str, fallback: str = ""): "COLORMAP_ANGLE1": config.get("OUTPUT", "COLORMAP_ANGLE1", fallback=None), "COLORMAP_ANGLE2": config.get("OUTPUT", "COLORMAP_ANGLE2", fallback=None), } + if not params["WRITE_VECTORS"] and not params["WRITE_ANGLES"]: + raise ValueError("At least one of WRITE_VECTORS or WRITE_ANGLES must be True") + return params def remove_corrupted_files(file_paths, size_threshold=200): import warnings + removed_files = [] for file_path in file_paths: if os.path.exists(file_path) and os.path.getsize(file_path) < size_threshold: warnings.warn( @@ -222,6 +341,9 @@ def remove_corrupted_files(file_paths, size_threshold=200): stacklevel=2, ) os.remove(file_path) + removed_files.append(file_path) + + return removed_files def convert_to_8bit( @@ -244,12 +366,14 @@ def convert_to_8bit( Returns: np.ndarray: 8-bit converted image. """ - # Compute percentiles - minimum, maximum = np.nanpercentile(img, (perc_min, perc_max)) - - # Override percentiles with explicit min/max if provided if min_value is not None and max_value is not None: minimum, maximum = min_value, max_value + else: + minimum, maximum = np.nanpercentile(img, (perc_min, perc_max)) + if min_value is not None: + minimum = min_value + if max_value is not None: + maximum = max_value # Avoid division by zero if maximum == minimum: diff --git a/src/cardiotensor/visualization/__init__.py b/src/cardiotensor/visualization/__init__.py new file mode 100644 index 0000000..9431c78 --- /dev/null +++ b/src/cardiotensor/visualization/__init__.py @@ -0,0 +1,9 @@ +"""Public visualization API.""" + +from cardiotensor.visualization.streamlines import visualize_streamlines +from cardiotensor.visualization.vector_field import visualize_vector_field + +__all__ = [ + "visualize_streamlines", + "visualize_vector_field", +] diff --git a/src/cardiotensor/visualization/fury_plotting_streamlines.py b/src/cardiotensor/visualization/fury_plotting_streamlines.py index 7ff1d2a..dd1511a 100644 --- a/src/cardiotensor/visualization/fury_plotting_streamlines.py +++ b/src/cardiotensor/visualization/fury_plotting_streamlines.py @@ -3,6 +3,7 @@ import datetime import inspect +import json import random import tempfile from pathlib import Path @@ -11,7 +12,7 @@ import matplotlib.pyplot as plt import numpy as np import vtk -from fury import actor, window +from fury import actor, ui, window # --------------------------- # small utilities @@ -116,8 +117,15 @@ def _compute_streamline_bounds( if not streamlines_xyz: raise ValueError("No streamlines available to compute bounds.") - mins = np.min([sl.min(axis=0) for sl in streamlines_xyz], axis=0) - maxs = np.max([sl.max(axis=0) for sl in streamlines_xyz], axis=0) + mins = np.full(3, np.inf) + maxs = np.full(3, -np.inf) + for streamline in streamlines_xyz: + if len(streamline) == 0: + continue + mins = np.minimum(mins, np.min(streamline, axis=0)) + maxs = np.maximum(maxs, np.max(streamline, axis=0)) + if not np.all(np.isfinite(mins)): + raise ValueError("No non-empty streamlines available to compute bounds.") return mins, maxs @@ -194,6 +202,13 @@ def __init__( tube_sides: int = 9, color_range: tuple[float, float] | None = None, color_label: str = "Angle", + flat_values: np.ndarray | None = None, + session_path: str | Path | None = None, + restore_session: bool = True, + session_settings: dict | None = None, + streamlines_file: str | Path | None = None, + opacity: float = 1.0, + quality: str = "interactive", ): self.streamlines_xyz = streamlines_xyz self.color_values = color_values @@ -201,16 +216,48 @@ def __init__( self.window_size = window_size self.lut = lut self.color_label = color_label + self.session_path = ( + Path(session_path).expanduser().resolve() if session_path else None + ) + self.restore_session = bool(restore_session) + self.session_settings = dict(session_settings or {}) + self.streamlines_file = ( + Path(streamlines_file).resolve() if streamlines_file else None + ) + self.quality = quality.lower().strip() + if self.quality not in {"interactive", "publication"}: + raise ValueError("quality must be 'interactive' or 'publication'") # Scene self.scene = fury.window.Scene() self.current_bg = parse_background_color(background_color) self.scene.SetBackground(*self.current_bg) + self._setup_lighting() # thickness state - self.linewidth = max(1.0, float(line_width)) # used for both line and tube - self.spline_subdiv = spline_subdiv + self.linewidth = max(0.001, float(line_width)) # used for both line and tube + self.spline_subdiv = max(1, int(spline_subdiv)) + self.actor_spline_subdiv = ( + self.spline_subdiv if self.spline_subdiv > 1 else None + ) self.tube_sides = max(3, int(tube_sides)) + if self.quality == "publication": + self.material = { + "ambient": 0.30, + "diffuse": 0.80, + "specular": 0.25, + "opacity": float(np.clip(opacity, 0.05, 1.0)), + } + else: + self.material = { + "ambient": 0.45, + "diffuse": 0.70, + "specular": 0.12, + "opacity": float(np.clip(opacity, 0.05, 1.0)), + } + self.controls_panel = None + self.controls_visible = False + self.control_sliders = {} self.scale_bar = None self.scale_bar_on = False @@ -220,18 +267,25 @@ def __init__( # VTK/FURY objects self.showm: window.ShowManager | None = None + self._closing = False # clipped branch objects self.plane_rep = None self.plane_fn = None self.plane_widget = None + self.box_widget = None + self.box_rep = None + self.box_planes = vtk.vtkPlanes() + self.box_clipping_active = False self.mapper0 = None self.actor0 = None # precompute flat scalars for LUT mapping - self.flat_vals = np.concatenate( - [np.asarray(c).ravel() for c in self.color_values] - ).astype(np.float32) + if flat_values is None: + flat_values = np.concatenate( + [np.asarray(c).ravel() for c in self.color_values] + ) + self.flat_vals = np.asarray(flat_values, dtype=np.float32).reshape(-1) if color_range is None: self.vmin = float(np.nanmin(self.flat_vals)) self.vmax = float(np.nanmax(self.flat_vals)) @@ -240,8 +294,7 @@ def __init__( self.lut.SetRange(self.vmin, self.vmax) # bounds and center from NumPy - mins = np.min([sl.min(axis=0) for sl in self.streamlines_xyz], axis=0) - maxs = np.max([sl.max(axis=0) for sl in self.streamlines_xyz], axis=0) + mins, maxs = _compute_streamline_bounds(self.streamlines_xyz) self.center = (mins + maxs) / 2.0 self.bounds = [mins[0], maxs[0], mins[1], maxs[1], mins[2], maxs[2]] @@ -258,7 +311,7 @@ def _build_pipeline(self): self.streamlines_xyz, colors=self.flat_vals, linewidth=self.linewidth, - spline_subdiv=self.spline_subdiv, + spline_subdiv=self.actor_spline_subdiv, lookup_colormap=self.lut, tube_sides=self.tube_sides, **_supported_actor_kwargs(actor.streamtube, lod=False), @@ -268,6 +321,7 @@ def _build_pipeline(self): self.streamlines_xyz, colors=self.flat_vals, # scalars linewidth=self.linewidth, + spline_subdiv=self.actor_spline_subdiv, lookup_colormap=self.lut, **_supported_actor_kwargs(actor.line, lod=False), ) @@ -285,6 +339,7 @@ def _build_pipeline(self): **_supported_actor_kwargs(actor.line, lod=False), ) self.actor_fast.SetVisibility(False) + self.actor_fast.GetProperty().SetOpacity(self.material["opacity"]) self.scene.add(self.actor_fast) # clipping plane setup, start disabled @@ -308,10 +363,130 @@ def _build_pipeline(self): def _style_streamline_actor(self): prop = self.actor0.GetProperty() prop.SetInterpolationToPhong() - prop.SetAmbient(0.45) # raised from 0.1 → brighter unlit sides, less dark - prop.SetDiffuse(0.75) - prop.SetSpecular(0.15) - prop.SetSpecularPower(10) + prop.SetAmbient(self.material["ambient"]) + prop.SetDiffuse(self.material["diffuse"]) + prop.SetSpecular(self.material["specular"]) + prop.SetSpecularPower(20 if self.quality == "publication" else 10) + prop.SetOpacity(self.material["opacity"]) + + def _setup_lighting(self): + """Use stable camera-relative key, fill and rim lights.""" + try: + self.scene.RemoveAllLights() + headlight = vtk.vtkLight() + headlight.SetLightTypeToHeadlight() + headlight.SetIntensity(0.80) + + fill = vtk.vtkLight() + fill.SetLightTypeToCameraLight() + fill.SetPosition(-1.0, 1.0, 0.5) + fill.SetFocalPoint(0.0, 0.0, 0.0) + fill.SetIntensity(0.35) + + rim = vtk.vtkLight() + rim.SetLightTypeToCameraLight() + rim.SetPosition(1.0, -1.0, 0.75) + rim.SetFocalPoint(0.0, 0.0, 0.0) + rim.SetIntensity(0.25) + + self.lights = [headlight, fill, rim] + for light in self.lights: + self.scene.AddLight(light) + except Exception: + self.lights = [] + + def _apply_material(self): + self._style_streamline_actor() + if self.actor_fast is not None: + self.actor_fast.GetProperty().SetOpacity(self.material["opacity"]) + self._render_now() + + def _set_material_value(self, name: str, value: float): + self.material[name] = float(value) + self._apply_material() + + def _build_controls_panel(self): + panel = ui.Panel2D( + size=(280, 255), + position=(12, 12), + color=(0.08, 0.09, 0.12), + opacity=0.88, + has_border=True, + border_color=(0.55, 0.58, 0.65), + border_width=1, + ) + title = ui.TextBlock2D( + text="FURY material controls (L)", + font_size=16, + bold=True, + color=(1, 1, 1), + ) + panel.add_element(title, (12, 225)) + + specs = ( + ("ambient", "Ambient", 0.0, 1.0), + ("diffuse", "Diffuse", 0.0, 1.0), + ("specular", "Specular", 0.0, 1.0), + ("opacity", "Opacity", 0.05, 1.0), + ) + for row, (name, label_text, minimum, maximum) in enumerate(specs): + y = 185 - row * 52 + label = ui.TextBlock2D( + text=label_text, font_size=13, color=(0.92, 0.94, 1.0) + ) + slider = ui.LineSlider2D( + initial_value=self.material[name], + min_value=minimum, + max_value=maximum, + length=145, + line_width=3, + outer_radius=7, + font_size=12, + text_template="{value:.2f}", + ) + slider.on_change = ( + lambda current_slider, material_name=name: self._set_material_value( + material_name, current_slider.value + ) + ) + panel.add_element(label, (12, y)) + panel.add_element(slider, (180, y + 5), anchor="center") + self.control_sliders[name] = slider + + self.controls_panel = panel + self.scene.add(panel) + panel.set_visibility(False) + + def _toggle_controls_panel(self): + if self.controls_panel is None: + return + self.controls_visible = not self.controls_visible + self.controls_panel.set_visibility(self.controls_visible) + self._render_now() + print(f"Material controls {'ON' if self.controls_visible else 'OFF'}") + + def _update_overlay_colors(self): + luminance = ( + 0.2126 * self.current_bg[0] + + 0.7152 * self.current_bg[1] + + 0.0722 * self.current_bg[2] + ) + color = (1.0, 1.0, 1.0) if luminance < 0.5 else (0.0, 0.0, 0.0) + if getattr(self, "scalar_bar", None) is not None: + self.scalar_bar.GetTitleTextProperty().SetColor(*color) + self.scalar_bar.GetLabelTextProperty().SetColor(*color) + try: + self.scalar_bar.GetAnnotationTextProperty().SetColor(*color) + except Exception: + pass + if self.scale_bar is not None: + for axis_name in ("Bottom", "Left", "Right", "Top"): + try: + axis = getattr(self.scale_bar, f"Get{axis_name}Axis")() + axis.GetLabelTextProperty().SetColor(*color) + axis.GetTitleTextProperty().SetColor(*color) + except Exception: + pass def _add_origin_marker(self): bounds_size = np.array( @@ -347,9 +522,11 @@ def _add_origin_marker(self): ) def _add_scalar_bar(self): - self.scene.add( - fury.actor.scalar_bar(lookup_table=self.lut, title=self.color_label) + self.scalar_bar = fury.actor.scalar_bar( + lookup_table=self.lut, title=self.color_label ) + self.scene.add(self.scalar_bar) + self._update_overlay_colors() def _render_now(self): try: @@ -359,6 +536,237 @@ def _render_now(self): if self.showm is not None: self.showm.render() + def _save_screenshot(self, out_path: Path, scale: int = 1) -> tuple[int, int]: + """Capture the current viewer without creating another VTK window.""" + if self.showm is None or not hasattr(self.showm, "window"): + raise RuntimeError("The interactive viewer is not initialized.") + + render_window = self.showm.window + render_window.Render() + + capture = vtk.vtkWindowToImageFilter() + capture.SetInput(render_window) + capture.SetInputBufferTypeToRGB() + if scale > 1: + capture.SetScale(int(scale)) + capture.Update() + + writer = vtk.vtkPNGWriter() + writer.SetFileName(str(out_path)) + writer.SetInputConnection(capture.GetOutputPort()) + writer.Write() + + width, height, _ = capture.GetOutput().GetDimensions() + return int(width), int(height) + + def _save_session(self, announce: bool = True) -> Path: + if self.session_path is None: + self.session_path = (Path.cwd() / "streamlines_session.json").resolve() + + camera = self.scene.GetActiveCamera() + origin = [0.0, 0.0, 0.0] + normal = [1.0, 0.0, 0.0] + self.plane_rep.GetOrigin(origin) + self.plane_rep.GetNormal(normal) + + window_size = list(self.window_size) + if self.showm is not None and hasattr(self.showm, "window"): + window_size = [int(value) for value in self.showm.window.GetSize()] + + settings = dict(self.session_settings) + settings.update( + { + "line_width": float(self.linewidth), + "tube_thickness": float(self.linewidth), + "window_size": window_size, + "background_color": list(self.current_bg), + "opacity": float(self.material["opacity"]), + "quality": self.quality, + } + ) + box_transform = None + if self.box_rep is not None: + transform = vtk.vtkTransform() + self.box_rep.GetTransform(transform) + matrix = transform.GetMatrix() + box_transform = [ + matrix.GetElement(row, column) + for row in range(4) + for column in range(4) + ] + payload = { + "format": "cardiotensor-fury-session", + "version": 1, + "updated_at": datetime.datetime.now().isoformat(timespec="seconds"), + "streamlines_file": ( + str(self.streamlines_file) if self.streamlines_file else None + ), + "streamlines_size": ( + self.streamlines_file.stat().st_size + if self.streamlines_file and self.streamlines_file.exists() + else None + ), + "settings": settings, + "view": { + "camera": { + "position": list(camera.GetPosition()), + "focal_point": list(camera.GetFocalPoint()), + "view_up": list(camera.GetViewUp()), + "clipping_range": list(camera.GetClippingRange()), + "view_angle": float(camera.GetViewAngle()), + "parallel_projection": bool(camera.GetParallelProjection()), + "parallel_scale": float(camera.GetParallelScale()), + }, + "clipping_plane": { + "origin": origin, + "normal": normal, + "enabled": bool(self.clipping_active), + "gizmo_visible": bool( + self.plane_widget and self.plane_widget.GetEnabled() + ), + }, + "crop_box": { + "transform": box_transform, + "enabled": bool(self.box_clipping_active), + "gizmo_visible": bool( + self.box_widget and self.box_widget.GetEnabled() + ), + }, + "material": dict(self.material), + "controls_visible": bool(self.controls_visible), + "background_color": list(self.current_bg), + "scale_bar_visible": bool(self.scale_bar_on), + "window_size": window_size, + }, + } + + self.session_path.parent.mkdir(parents=True, exist_ok=True) + temporary_path = self.session_path.with_suffix( + self.session_path.suffix + ".tmp" + ) + temporary_path.write_text(json.dumps(payload, indent=2) + "\n") + temporary_path.replace(self.session_path) + self.session_settings = settings + if announce: + print(f"Saved FURY session to {self.session_path}") + return self.session_path + + def _restore_session(self) -> bool: + if self.session_path is None or not self.session_path.exists(): + return False + + try: + view = json.loads(self.session_path.read_text()).get("view", {}) + camera_state = view.get("camera", {}) + camera = self.scene.GetActiveCamera() + if camera_state.get("position"): + camera.SetPosition(*camera_state["position"]) + if camera_state.get("focal_point"): + camera.SetFocalPoint(*camera_state["focal_point"]) + if camera_state.get("view_up"): + camera.SetViewUp(*camera_state["view_up"]) + if camera_state.get("view_angle") is not None: + camera.SetViewAngle(float(camera_state["view_angle"])) + if camera_state.get("parallel_projection") is not None: + camera.SetParallelProjection(bool(camera_state["parallel_projection"])) + if camera_state.get("parallel_scale") is not None: + camera.SetParallelScale(float(camera_state["parallel_scale"])) + if camera_state.get("clipping_range"): + camera.SetClippingRange(*camera_state["clipping_range"]) + + plane = view.get("clipping_plane", {}) + if plane.get("origin"): + self.plane_rep.SetOrigin(*plane["origin"]) + if plane.get("normal"): + self.plane_rep.SetNormal(*plane["normal"]) + self.plane_rep.UpdatePlacement() + self._sync_plane_from_widget() + + self.clipping_active = bool(plane.get("enabled", False)) + + if self.plane_widget is not None: + if plane.get("gizmo_visible", False): + self.plane_widget.EnabledOn() + else: + self.plane_widget.EnabledOff() + + crop_box = view.get("crop_box", {}) + transform_values = crop_box.get("transform") + if getattr(self, "box_rep", None) is not None and transform_values: + matrix = vtk.vtkMatrix4x4() + for row in range(4): + for column in range(4): + matrix.SetElement( + row, column, transform_values[row * 4 + column] + ) + transform = vtk.vtkTransform() + transform.SetMatrix(matrix) + self.box_rep.SetTransform(transform) + self.box_clipping_active = bool(crop_box.get("enabled", False)) + if getattr(self, "box_widget", None) is not None: + if crop_box.get("gizmo_visible", False): + self.box_widget.EnabledOn() + else: + self.box_widget.EnabledOff() + self._apply_clipping_planes() + + saved_material = view.get("material", {}) + if saved_material: + for name in self.material: + if name in saved_material: + self.material[name] = float(saved_material[name]) + self._apply_material() + for name, slider in self.control_sliders.items(): + slider.value = self.material[name] + self.controls_visible = bool(view.get("controls_visible", False)) + if getattr(self, "controls_panel", None) is not None: + self.controls_panel.set_visibility(self.controls_visible) + + background = view.get("background_color") + if background: + self.current_bg = tuple(float(value) for value in background) + self.scene.SetBackground(*self.current_bg) + self._update_overlay_colors() + + show_scale_bar = bool(view.get("scale_bar_visible", True)) + if show_scale_bar and not self.scale_bar_on: + self.scene.add(self.scale_bar) + self.scale_bar_on = True + elif not show_scale_bar and self.scale_bar_on: + self.scene.rm(self.scale_bar) + self.scale_bar_on = False + + if camera_state.get("clipping_range"): + camera.SetClippingRange(*camera_state["clipping_range"]) + if self.showm is not None: + self.showm.render() + print(f"Restored FURY session from {self.session_path}") + return True + except (OSError, ValueError, TypeError, json.JSONDecodeError) as err: + print(f"Warning: could not restore session {self.session_path}: {err}") + return False + + def _autosave_session(self, *_): + if self.session_path is None: + return + try: + self._save_session(announce=False) + print(f"Updated FURY session: {self.session_path}") + except (OSError, ValueError, TypeError) as err: + print(f"Warning: could not update session {self.session_path}: {err}") + + def _close_window(self, *_): + """Save once, then restore VTK's normal window-close behavior.""" + if self._closing: + return + self._closing = True + self._autosave_session() + for widget in (self.plane_widget, self.box_widget): + if widget is not None: + widget.EnabledOff() + if self.showm is not None: + self.showm.exit() + def _sync_plane_from_widget(self, *_): origin = [0.0, 0.0, 0.0] normal = [1.0, 0.0, 0.0] @@ -366,6 +774,83 @@ def _sync_plane_from_widget(self, *_): self.plane_rep.GetNormal(normal) self.plane_fn.SetOrigin(origin) self.plane_fn.SetNormal(normal) + if self.clipping_active: + self._apply_clipping_planes() + + def _setup_box_widget(self): + self.box_rep = vtk.vtkBoxRepresentation() + self.box_rep.SetPlaceFactor(1.0) + self.box_rep.PlaceWidget(self.bounds) + try: + self.box_rep.GetOutlineProperty().SetColor(0.2, 0.8, 1.0) + self.box_rep.GetHandleProperty().SetColor(1.0, 0.75, 0.2) + except Exception: + pass + + self.box_widget = vtk.vtkBoxWidget2() + self.box_widget.SetRepresentation(self.box_rep) + self.box_widget.SetInteractor(self.showm.iren) + self.box_widget.RotationEnabledOn() + self.box_widget.EnabledOff() + for event in ( + vtk.vtkCommand.StartInteractionEvent, + vtk.vtkCommand.InteractionEvent, + vtk.vtkCommand.EndInteractionEvent, + ): + self.box_widget.AddObserver(event, self._sync_box_from_widget) + + def _sync_box_from_widget(self, *_): + if self.box_rep is not None: + self.box_rep.GetPlanes(self.box_planes) + if self.box_clipping_active: + self._apply_clipping_planes() + + def _apply_clipping_planes(self): + box_rep = getattr(self, "box_rep", None) + box_active = bool(getattr(self, "box_clipping_active", False)) + if box_rep is not None: + box_rep.GetPlanes(self.box_planes) + for current_actor in (self.actor0, self.actor_fast): + if current_actor is None: + continue + mapper = current_actor.GetMapper() + mapper.RemoveAllClippingPlanes() + if self.clipping_active: + mapper.AddClippingPlane(self.plane_fn) + if box_active: + for plane_index in range(self.box_planes.GetNumberOfPlanes()): + mapper.AddClippingPlane(self.box_planes.GetPlane(plane_index)) + + def _toggle_box_clipping(self): + self.box_clipping_active = not self.box_clipping_active + if self.box_widget is not None: + if self.box_clipping_active: + self.box_widget.EnabledOn() + else: + self.box_widget.EnabledOff() + self._apply_clipping_planes() + self._render_now() + state = "ON" if self.box_clipping_active else "OFF" + print(f"Rotatable crop box {state}") + + def _toggle_box_gizmo(self): + if self.box_widget is None: + return + if self.box_widget.GetEnabled(): + self.box_widget.EnabledOff() + print("Crop-box gizmo hidden; box clipping state unchanged") + else: + self.box_widget.EnabledOn() + print("Crop-box gizmo shown; drag faces or handles and rotate it") + self._render_now() + + def _reset_box(self): + if self.box_rep is None: + return + self.box_rep.PlaceWidget(self.bounds) + self._sync_box_from_widget() + self._render_now() + print("Crop box reset to the full streamline bounds") def _add_scale_bar(self): self.scale_bar = vtk.vtkLegendScaleActor() @@ -383,28 +868,12 @@ def _add_scale_bar(self): def _toggle_clipping(self): """Toggle clipping state.""" - if self.clipping_active: - # Deactivate clipping - self.actor0.GetMapper().RemoveAllClippingPlanes() - self.actor_fast.GetMapper().RemoveAllClippingPlanes() - self.clipping_active = False - print("Clipping OFF") - else: - # Activate clipping - self.actor0.GetMapper().AddClippingPlane(self.plane_fn) - self.actor_fast.GetMapper().AddClippingPlane(self.plane_fn) - self.clipping_active = True - print("Clipping ON") - + self.clipping_active = not self.clipping_active + self._apply_clipping_planes() + print(f"Clipping plane {'ON' if self.clipping_active else 'OFF'}") self._render_now() def _rebuild_unclipped_actor(self): - clipping_on = ( - self.mapper0.GetNumberOfClippingPlanes() > 0 - if self.mapper0 is not None - else False - ) - if self.actor0 is not None: try: self.scene.rm(self.actor0) @@ -416,7 +885,7 @@ def _rebuild_unclipped_actor(self): self.streamlines_xyz, colors=self.flat_vals, linewidth=self.linewidth, - spline_subdiv=self.spline_subdiv, + spline_subdiv=self.actor_spline_subdiv, tube_sides=self.tube_sides, lookup_colormap=self.lut, **_supported_actor_kwargs(actor.streamtube, lod=False), @@ -426,6 +895,7 @@ def _rebuild_unclipped_actor(self): self.streamlines_xyz, colors=self.flat_vals, linewidth=self.linewidth, + spline_subdiv=self.actor_spline_subdiv, lookup_colormap=self.lut, **_supported_actor_kwargs(actor.line, lod=False), ) @@ -434,19 +904,57 @@ def _rebuild_unclipped_actor(self): self._style_streamline_actor() self.mapper0 = self.actor0.GetMapper() - if clipping_on: - self.mapper0.RemoveAllClippingPlanes() - self.mapper0.AddClippingPlane(self.plane_fn) + self._apply_clipping_planes() # --------------------------- # key handling # --------------------------- def _on_keypress(self, obj, evt): key = obj.GetKeySym().lower() + control_pressed = ( + bool(obj.GetControlKey()) if hasattr(obj, "GetControlKey") else False + ) + shift_pressed = ( + bool(obj.GetShiftKey()) if hasattr(obj, "GetShiftKey") else False + ) + + if key == "s" and control_pressed: + try: + self._save_session() + except (OSError, ValueError, TypeError) as err: + print(f"Failed to save session: {err}") + return if key == "o": self._toggle_clipping() + elif key == "c": + self._toggle_box_clipping() + + elif key == "g": + self._toggle_box_gizmo() + + elif key == "x": + self._reset_box() + + elif key == "l": + self._toggle_controls_panel() + + elif key in ("1", "kp_1"): + self._set_camera_preset("front") + + elif key in ("2", "kp_2"): + self._set_camera_preset("side") + + elif key in ("3", "kp_3"): + self._set_camera_preset("top") + + elif key in ("4", "kp_4"): + self._set_camera_preset("isometric") + + elif key in ("5", "kp_5"): + self._toggle_projection() + elif key == "h": if self.plane_widget: currently_on = self.plane_widget.GetEnabled() @@ -471,6 +979,7 @@ def _on_keypress(self, obj, evt): else (1.0, 1.0, 1.0) ) self.scene.SetBackground(*self.current_bg) + self._update_overlay_colors() self._render_now() print(f"Background set to {self.current_bg}") @@ -490,17 +999,14 @@ def _on_keypress(self, obj, evt): elif key == "p": ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") - out_path = Path.cwd() / f"view_{ts}_hires.png" + scale = 2 if shift_pressed else 1 + suffix = "_2x" if scale == 2 else "" + out_path = Path.cwd() / f"view_{ts}{suffix}.png" try: - highres_size = (2000, 2000) - fury.window.record( - scene=self.scene, - out_path=str(out_path), - size=highres_size, - reset_camera=False, - ) + width, height = self._save_screenshot(out_path, scale=scale) print( - f"Saved high-resolution screenshot to {out_path} ({highres_size[0]}x{highres_size[1]})" + f"Saved screenshot to {out_path} " + f"({width}x{height}, {scale}x viewer resolution)" ) except Exception as e: print(f"Failed to save screenshot: {e}") @@ -541,11 +1047,7 @@ def _lod_on(self, obj=None, evt=None): self.actor0.SetVisibility(False) self.actor_fast.SetVisibility(True) - # Apply clipping only if it was previously enabled - if self.clipping_active: - mapper = self.actor_fast.GetMapper() - mapper.RemoveAllClippingPlanes() - mapper.AddClippingPlane(self.plane_fn) + self._apply_clipping_planes() except Exception: pass self._render_now() @@ -556,11 +1058,7 @@ def _lod_off(self, obj=None, evt=None): self.actor_fast.SetVisibility(False) # Hide the fast actor self.actor0.SetVisibility(True) # Show the full-res actor - # Apply clipping to the full-res actor if it was previously enabled - if self.clipping_active: - mapper = self.actor0.GetMapper() - mapper.RemoveAllClippingPlanes() - mapper.AddClippingPlane(self.plane_fn) + self._apply_clipping_planes() except Exception: pass self._render_now() @@ -571,6 +1069,53 @@ def _set_default_camera(self) -> None: self.scene.elevation(10) self.scene.zoom(1.1) + def _set_camera_preset(self, preset: str) -> None: + camera = self.scene.GetActiveCamera() + center = np.asarray(self.center, dtype=float) + diagonal = max( + 1.0, + float( + np.linalg.norm( + [ + self.bounds[1] - self.bounds[0], + self.bounds[3] - self.bounds[2], + self.bounds[5] - self.bounds[4], + ] + ) + ), + ) + directions = { + "front": (np.array([0.0, -1.0, 0.0]), (0.0, 0.0, 1.0)), + "side": (np.array([1.0, 0.0, 0.0]), (0.0, 0.0, 1.0)), + "top": (np.array([0.0, 0.0, 1.0]), (0.0, 1.0, 0.0)), + "isometric": (np.array([1.0, -1.0, 1.0]), (0.0, 0.0, 1.0)), + } + direction, view_up = directions[preset] + direction /= np.linalg.norm(direction) + camera.SetFocalPoint(*center) + camera.SetPosition(*(center + 1.8 * diagonal * direction)) + camera.SetViewUp(*view_up) + camera.OrthogonalizeViewUp() + self.scene.ResetCameraClippingRange() + self._render_now() + print(f"Camera preset: {preset}") + + def _toggle_projection(self) -> None: + camera = self.scene.GetActiveCamera() + use_parallel = not bool(camera.GetParallelProjection()) + camera.SetParallelProjection(use_parallel) + if use_parallel: + size = np.array( + [ + self.bounds[1] - self.bounds[0], + self.bounds[3] - self.bounds[2], + self.bounds[5] - self.bounds[4], + ] + ) + camera.SetParallelScale(0.55 * float(np.max(size))) + self._render_now() + print(f"Projection: {'parallel' if use_parallel else 'perspective'}") + def _record_orbit( self, video_path: str | Path, video_frames: int, video_fps: int ) -> None: @@ -590,7 +1135,11 @@ def _record_orbit( f"Recording FURY orbit video: {n_frames} frames at {fps} fps -> {out_path}" ) - with tempfile.TemporaryDirectory(prefix="cardiotensor_fury_orbit_") as tmpdir: + scratch_root = out_path.parent / ".cardiotensor_scratch" + scratch_root.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory( + prefix="cardiotensor_fury_orbit_", dir=scratch_root + ) as tmpdir: tmpdir_path = Path(tmpdir) frame_paths = [] try: @@ -633,6 +1182,7 @@ def run( scene=self.scene, size=self.window_size, reset_camera=False ) self.showm.initialize() + self._build_controls_panel() iren = self.showm.iren iren.SetDesiredUpdateRate(60.0) @@ -640,7 +1190,7 @@ def run( # ensure anti-alias looks good when idle try: - self.showm.renwin.SetMultiSamples(0) + self.showm.window.SetMultiSamples(0) self.scene.enable_anti_aliasing("fxaa") except Exception: pass @@ -663,13 +1213,24 @@ def run( self.plane_widget.AddObserver( vtk.vtkCommand.EndInteractionEvent, self._sync_plane_from_widget ) + self._setup_box_widget() self.showm.iren.AddObserver("KeyPressEvent", self._on_keypress) + # An ExitEvent observer replaces VTK's default TerminateApp callback, + # so our handler must save the session and explicitly end interaction. + self._closing = False + self.showm.iren.AddObserver("ExitEvent", self._close_window) self._set_default_camera() + if self.restore_session: + self._restore_session() print( - "Keys: O toggle plane, H hide gizmo, I flip side, R reset plane, +/- thickness, B background, S scale bar, P save PNG, V record orbit" + "Keys: O plane clipping, H plane gizmo, I flip plane, R reset plane; " + "C rotatable crop box, G box gizmo, X reset box; L material controls; " + "1 front, 2 side, 3 top, 4 isometric, 5 projection; +/- thickness; " + "B background, S scale bar, Ctrl+S session, P PNG, Shift+P 2x PNG, " + "V orbit" ) self.showm.start() else: @@ -718,6 +1279,13 @@ def show_streamlines( tube_sides: int = 9, color_range: tuple[float, float] | None = None, color_label: str = "Angle (deg)", + random_seed: int | None = None, + session_path: str | Path | None = None, + restore_session: bool = True, + session_settings: dict | None = None, + streamlines_file: str | Path | None = None, + opacity: float = 1.0, + quality: str = "interactive", ): print(f"Initial number of streamlines: {len(streamlines_xyz)}") full_mins, full_maxs = _compute_streamline_bounds(streamlines_xyz) @@ -764,18 +1332,17 @@ def show_streamlines( if not streamlines_xyz: raise ValueError("No streamlines left after downsampling or filtering.") + rng = random.Random(random_seed) if subsample_factor > 1: print(f"Subsampling: keeping 1 in every {subsample_factor} streamlines") total = len(streamlines_xyz) - keep_idx = sorted( - random.sample(range(total), max(1, total // subsample_factor)) - ) + keep_idx = sorted(rng.sample(range(total), max(1, total // subsample_factor))) streamlines_xyz = [streamlines_xyz[i] for i in keep_idx] color_values = [color_values[i] for i in keep_idx] if max_streamlines is not None and len(streamlines_xyz) > max_streamlines: print(f"Limiting to max {max_streamlines} streamlines") - keep_idx = sorted(random.sample(range(len(streamlines_xyz)), max_streamlines)) + keep_idx = sorted(rng.sample(range(len(streamlines_xyz)), max_streamlines)) streamlines_xyz = [streamlines_xyz[i] for i in keep_idx] color_values = [color_values[i] for i in keep_idx] @@ -806,6 +1373,14 @@ def show_streamlines( lut_range = (float(color_range[0]), float(color_range[1])) print(f"Colorbar range: min={lut_range[0]:.3f}, max={lut_range[1]:.3f}") print(f"Rendering mode: {mode}") + if int(spline_subdiv) <= 1: + print("Visual smoothing: off (original streamline points preserved)") + else: + print( + "Visual smoothing: VTK cardinal spline with " + f"{int(spline_subdiv)} segments per streamline" + ) + print(f"FURY quality preset: {quality}") if colormap is None: lut = fury.actor.colormap_lookup_table( @@ -831,6 +1406,13 @@ def show_streamlines( tube_sides=tube_sides, color_range=lut_range, color_label=color_label, + flat_values=flat_colors, + session_path=session_path, + restore_session=restore_session, + session_settings=session_settings, + streamlines_file=streamlines_file, + opacity=opacity, + quality=quality, ) viewer.run( interactive=interactive, diff --git a/src/cardiotensor/visualization/pyvista_plotting_streamlines.py b/src/cardiotensor/visualization/pyvista_plotting_streamlines.py new file mode 100644 index 0000000..3443c57 --- /dev/null +++ b/src/cardiotensor/visualization/pyvista_plotting_streamlines.py @@ -0,0 +1,1344 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import datetime +import random +from pathlib import Path + +import numpy as np + +_NAMED_COLORS = { + "white": (1.0, 1.0, 1.0), + "black": (0.0, 0.0, 0.0), + "gray": (0.5, 0.5, 0.5), + "grey": (0.5, 0.5, 0.5), + "lightgray": (0.9, 0.9, 0.9), + "lightgrey": (0.9, 0.9, 0.9), + "charcoal": (0.05, 0.06, 0.08), + "midnight": (0.02, 0.03, 0.07), +} + + +def _load_pyvista(): + try: + import pyvista as pv + except ImportError as err: + raise ImportError( + "PyVista backend requested but pyvista is not installed. " + "Install the project with the updated dependencies, or run `pip install pyvista`." + ) from err + return pv + + +def _parse_background_color(color) -> tuple[float, float, float] | str: + if color is None: + return _NAMED_COLORS["white"] + if isinstance(color, str): + return _NAMED_COLORS.get(color.lower(), color) + if isinstance(color, tuple | list) and len(color) == 3: + return tuple(float(c) for c in color) + raise TypeError("background_color must be a string or 3-tuple") + + +def _is_dark_color(color: tuple[float, float, float] | str) -> bool: + if isinstance(color, str): + color = _NAMED_COLORS.get(color.lower(), (0.0, 0.0, 0.0)) + r, g, b = color + return 0.2126 * r + 0.7152 * g + 0.0722 * b < 0.45 + + +def _downsample_streamline(streamline: np.ndarray, factor: int = 2) -> np.ndarray: + return streamline if len(streamline) < 3 or factor <= 1 else streamline[::factor] + + +def _compute_streamline_bounds( + streamlines_xyz: list[np.ndarray], +) -> tuple[np.ndarray, np.ndarray]: + if not streamlines_xyz: + raise ValueError("No streamlines available to compute bounds.") + + mins = np.full(3, np.inf) + maxs = np.full(3, -np.inf) + for streamline in streamlines_xyz: + if len(streamline) == 0: + continue + mins = np.minimum(mins, np.min(streamline, axis=0)) + maxs = np.maximum(maxs, np.max(streamline, axis=0)) + if not np.all(np.isfinite(mins)): + raise ValueError("No non-empty streamlines available to compute bounds.") + return mins, maxs + + +def _print_box_shape(label: str, mins: np.ndarray, maxs: np.ndarray) -> None: + size = maxs - mins + print( + f"{label}: \n" + f"x=[{mins[0]:.2f}, {maxs[0]:.2f}], \n" + f"y=[{mins[1]:.2f}, {maxs[1]:.2f}], \n" + f"z=[{mins[2]:.2f}, {maxs[2]:.2f}] \n" + f"-> shape=({size[0]:.2f}, {size[1]:.2f}, {size[2]:.2f})\n" + ) + + +def _split_streamline_by_bounds( + sl: np.ndarray, cl: np.ndarray, x_min, x_max, y_min, y_max, z_min, z_max +): + within = ( + (sl[:, 0] >= x_min) + & (sl[:, 0] <= x_max) + & (sl[:, 1] >= y_min) + & (sl[:, 1] <= y_max) + & (sl[:, 2] >= z_min) + & (sl[:, 2] <= z_max) + ) + if not np.any(within): + return [], [] + + w = within.astype(np.int8) + trans = np.diff(np.pad(w, (1, 1), constant_values=0)) + starts = np.where(trans == +1)[0] + ends = np.where(trans == -1)[0] + + segs, cols = [], [] + for s, e in zip(starts, ends): + seg = sl[s:e] + col = cl[s:e] + if len(seg) > 1: + segs.append(seg) + cols.append(col) + return segs, cols + + +def _prepare_streamlines( + streamlines_xyz: list[np.ndarray], + color_values: list[np.ndarray], + downsample_factor: int, + max_streamlines: int | None, + filter_min_len: int | None, + subsample_factor: int, + crop_bounds: ( + tuple[tuple[float, float], tuple[float, float], tuple[float, float]] | None + ), + random_seed: int | None = None, +) -> tuple[list[np.ndarray], list[np.ndarray], tuple[float, float]]: + if len(streamlines_xyz) != len(color_values): + raise ValueError("color_values must contain one array per streamline") + + print(f"Initial number of streamlines: {len(streamlines_xyz)}") + full_mins, full_maxs = _compute_streamline_bounds(streamlines_xyz) + _print_box_shape("Full streamline box", full_mins, full_maxs) + + if crop_bounds is not None: + (x_min, x_max), (y_min, y_max), (z_min, z_max) = crop_bounds + print(f"Cropping streamlines within bounds: {crop_bounds}") + new_streamlines, new_colors = [], [] + for sl, cl in zip(streamlines_xyz, color_values): + segs, cols = _split_streamline_by_bounds( + np.asarray(sl), np.asarray(cl), x_min, x_max, y_min, y_max, z_min, z_max + ) + if segs: + new_streamlines.extend(segs) + new_colors.extend(cols) + streamlines_xyz, color_values = new_streamlines, new_colors + if not streamlines_xyz: + raise ValueError("No streamlines intersect the crop box.") + cropped_mins, cropped_maxs = _compute_streamline_bounds(streamlines_xyz) + _print_box_shape("Cropped streamline box", cropped_mins, cropped_maxs) + print("Cropping applied.") + else: + print("No cropping applied.") + + print(f"Downsampling points by factor {downsample_factor}") + if filter_min_len is not None: + print(f"Filtering out streamlines shorter than {filter_min_len} points") + + ds_streamlines, ds_colors = [], [] + for sl, cl in zip(streamlines_xyz, color_values): + sl_arr = np.asarray(sl, dtype=np.float32) + cl_arr = np.asarray(cl, dtype=np.float32).ravel() + if len(sl_arr) != len(cl_arr): + raise ValueError("Each color array must match its streamline length") + ds_sl = _downsample_streamline(sl_arr, downsample_factor) + ds_cl = _downsample_streamline(cl_arr, downsample_factor) + if len(ds_sl) > 1 and (filter_min_len is None or len(ds_sl) >= filter_min_len): + ds_streamlines.append(ds_sl) + ds_colors.append(ds_cl) + + streamlines_xyz, color_values = ds_streamlines, ds_colors + if not streamlines_xyz: + raise ValueError("No streamlines left after downsampling or filtering.") + + rng = random.Random(random_seed) + if subsample_factor > 1: + print(f"Subsampling: keeping 1 in every {subsample_factor} streamlines") + total = len(streamlines_xyz) + keep_idx = sorted(rng.sample(range(total), max(1, total // subsample_factor))) + streamlines_xyz = [streamlines_xyz[i] for i in keep_idx] + color_values = [color_values[i] for i in keep_idx] + + if max_streamlines is not None and len(streamlines_xyz) > max_streamlines: + print(f"Limiting to max {max_streamlines} streamlines") + keep_idx = sorted(rng.sample(range(len(streamlines_xyz)), max_streamlines)) + streamlines_xyz = [streamlines_xyz[i] for i in keep_idx] + color_values = [color_values[i] for i in keep_idx] + + flat_colors = np.concatenate([np.asarray(c).ravel() for c in color_values]).astype( + np.float32 + ) + min_val = float(np.nanmin(flat_colors)) + max_val = float(np.nanmax(flat_colors)) + print(f"Final number of streamlines to render: {len(streamlines_xyz)}") + print(f"Coloring range: min={min_val:.3f}, max={max_val:.3f}") + return streamlines_xyz, color_values, (min_val, max_val) + + +def _crop_streamlines_by_bounds( + streamlines_xyz: list[np.ndarray], + color_values: list[np.ndarray], + bounds: tuple[float, float, float, float, float, float], +) -> tuple[list[np.ndarray], list[np.ndarray]]: + x_min, x_max, y_min, y_max, z_min, z_max = bounds + cropped_streamlines, cropped_colors = [], [] + for sl, cl in zip(streamlines_xyz, color_values): + segs, cols = _split_streamline_by_bounds( + np.asarray(sl), np.asarray(cl), x_min, x_max, y_min, y_max, z_min, z_max + ) + cropped_streamlines.extend(segs) + cropped_colors.extend(cols) + return cropped_streamlines, cropped_colors + + +def _format_crop_cli_bounds( + bounds: tuple[float, float, float, float, float, float], +) -> str: + return ( + f"--crop-x {bounds[0]:.2f} {bounds[1]:.2f} " + f"--crop-y {bounds[2]:.2f} {bounds[3]:.2f} " + f"--crop-z {bounds[4]:.2f} {bounds[5]:.2f}" + ) + + +def _smooth_polydata_with_vtk_spline( + poly, + pv, + subdivisions: int, +): + if subdivisions <= 1: + return poly + try: + from vtkmodules.vtkFiltersGeneral import vtkSplineFilter + except ImportError: + try: + from vtk import vtkSplineFilter + except ImportError: + print("Visual smoothing unavailable: VTK spline filter is not installed.") + return poly + + spline_filter = vtkSplineFilter() + spline_filter.SetInputData(poly) + spline_filter.SetSubdivideToSpecified() + spline_filter.SetNumberOfSubdivisions(int(subdivisions)) + spline_filter.Update() + return pv.wrap(spline_filter.GetOutput()).copy(deep=True) + + +def _print_render_summary( + streamlines_xyz: list[np.ndarray], + color_values: list[np.ndarray], + data_range: tuple[float, float], + color_range: tuple[float, float] | None, + mode: str, + line_width: float, + tube_sides: int, + downsample_factor: int, + subsample_factor: int, + max_streamlines: int | None, + filter_min_len: int | None, + crop_bounds, + spline_subdiv: int, +) -> None: + lengths = np.array([len(sl) for sl in streamlines_xyz], dtype=np.int64) + total_points = int(lengths.sum()) + total_segments = int(np.maximum(lengths - 1, 0).sum()) + mins, maxs = _compute_streamline_bounds(streamlines_xyz) + box_size = maxs - mins + diagonal = float(np.linalg.norm(box_size)) + scalar_points = int(sum(len(np.asarray(c).ravel()) for c in color_values)) + effective_range = color_range or data_range + approx_tube_polys = total_segments * max(3, int(tube_sides)) * 2 + + print("\nPyVista render summary") + print("-" * 24) + print(f"Backend/mode: pyvista / {mode}") + print(f"Streamlines: {len(streamlines_xyz):,}") + print( + "Points per streamline: " + f"mean={float(lengths.mean()):.1f}, " + f"median={float(np.median(lengths)):.1f}, " + f"min={int(lengths.min())}, max={int(lengths.max())}" + ) + print(f"Total points: {total_points:,}") + print(f"Total segments: {total_segments:,}") + print(f"Scalar samples: {scalar_points:,}") + print( + "Bounds: " + f"x=[{mins[0]:.2f}, {maxs[0]:.2f}], " + f"y=[{mins[1]:.2f}, {maxs[1]:.2f}], " + f"z=[{mins[2]:.2f}, {maxs[2]:.2f}]" + ) + print( + "Box size: " + f"dx={box_size[0]:.2f}, dy={box_size[1]:.2f}, dz={box_size[2]:.2f}, " + f"diagonal={diagonal:.2f}" + ) + print(f"Data color range: min={data_range[0]:.3f}, max={data_range[1]:.3f}") + print( + f"Colorbar range: min={effective_range[0]:.3f}, " + f"max={effective_range[1]:.3f}" + ) + print( + f"Requested crop bounds: {crop_bounds if crop_bounds is not None else 'none'}" + ) + print( + "Filtering: " + f"downsample_factor={downsample_factor}, " + f"subsample_factor={subsample_factor}, " + f"max_streamlines={max_streamlines}, " + f"min_length={filter_min_len}, " + f"spline_subdiv={spline_subdiv}" + ) + print( + "Geometry: " + f"line_width/radius={line_width}, tube_sides={tube_sides}, " + f"approx_tube_triangles={approx_tube_polys:,}" + ) + if mode == "tube" and approx_tube_polys > 1_000_000: + print( + "Performance hint: many streamlines. Press M to switch to line mode for faster navigation." + ) + elif mode == "line": + print( + "Performance hint: line mode is best for navigation; tube mode is best for final figures." + ) + print("-" * 24) + + +def _streamlines_to_polydata( + streamlines_xyz: list[np.ndarray], + color_values: list[np.ndarray], + scalar_name: str, +): + pv = _load_pyvista() + lengths = np.fromiter( + (len(streamline) for streamline in streamlines_xyz), + dtype=np.int64, + count=len(streamlines_xyz), + ) + n_points = int(lengths.sum()) + points = np.concatenate(streamlines_xyz, axis=0).astype(np.float32, copy=False) + scalars = np.concatenate(color_values).astype(np.float32, copy=False) + + # VTK stores lines as [length, point ids..., length, point ids...]. Build + # that array in NumPy instead of appending millions of Python integers. + line_starts = np.empty(len(lengths), dtype=np.int64) + line_starts[0] = 0 + if len(lengths) > 1: + np.cumsum(lengths[:-1] + 1, out=line_starts[1:]) + lines = np.empty(n_points + len(lengths), dtype=np.int64) + lines[line_starts] = lengths + point_slots = np.ones(len(lines), dtype=bool) + point_slots[line_starts] = False + lines[point_slots] = np.arange(n_points, dtype=np.int64) + + poly = pv.PolyData(points) + poly.lines = lines + poly.point_data[scalar_name] = scalars + poly.set_active_scalars(scalar_name) + return poly + + +def show_streamlines_pyvista( + streamlines_xyz: list[np.ndarray], + color_values: list[np.ndarray], + mode: str = "tube", + line_width: float = 2.0, + interactive: bool = True, + screenshot_path: str | None = None, + video_path: str | None = None, + video_fps: int = 30, + video_frames: int = 120, + window_size: tuple[int, int] = (1000, 900), + downsample_factor: int = 2, + max_streamlines: int | None = None, + filter_min_len: int | None = None, + subsample_factor: int = 1, + crop_bounds: ( + tuple[tuple[float, float], tuple[float, float], tuple[float, float]] | None + ) = None, + colormap=None, + background_color: str | tuple[float, float, float] | None = None, + tube_sides: int = 12, + color_range: tuple[float, float] | None = None, + color_label: str = "Angle (deg)", + opacity: float = 1.0, + show_axes: bool = True, + show_bounds: bool = False, + shadows: bool = False, + spline_subdiv: int = 1, + random_seed: int | None = None, +) -> None: + pv = _load_pyvista() + mode = mode.lower().strip() + if mode not in {"tube", "line"}: + raise ValueError("mode must be 'tube' or 'line'") + + streamlines_xyz, color_values, data_range = _prepare_streamlines( + streamlines_xyz=streamlines_xyz, + color_values=color_values, + downsample_factor=downsample_factor, + max_streamlines=max_streamlines, + filter_min_len=filter_min_len, + subsample_factor=subsample_factor, + crop_bounds=crop_bounds, + random_seed=random_seed, + ) + spline_subdiv = max(1, int(spline_subdiv)) + if spline_subdiv <= 1: + print("Visual smoothing: off") + clim = color_range or data_range + print(f"Rendering backend: pyvista, mode: {mode}") + print(f"Colorbar range: min={clim[0]:.3f}, max={clim[1]:.3f}") + _print_render_summary( + streamlines_xyz=streamlines_xyz, + color_values=color_values, + data_range=data_range, + color_range=color_range, + mode=mode, + line_width=line_width, + tube_sides=tube_sides, + downsample_factor=downsample_factor, + subsample_factor=subsample_factor, + max_streamlines=max_streamlines, + filter_min_len=filter_min_len, + crop_bounds=crop_bounds, + spline_subdiv=spline_subdiv, + ) + + base_streamlines_xyz = [np.asarray(sl, dtype=np.float32) for sl in streamlines_xyz] + base_color_values = [ + np.asarray(cl, dtype=np.float32).ravel() for cl in color_values + ] + + scalar_name = color_label or "Scalar" + + def _make_display_polydata( + display_streamlines: list[np.ndarray], display_colors: list[np.ndarray] + ): + display_poly = _streamlines_to_polydata( + display_streamlines, display_colors, scalar_name + ) + if spline_subdiv <= 1: + return display_poly + + before_points = int(display_poly.n_points) + display_poly = _smooth_polydata_with_vtk_spline( + display_poly, + pv, + subdivisions=spline_subdiv, + ) + display_poly.set_active_scalars(scalar_name) + after_points = int(display_poly.n_points) + display_segments = max(0, after_points - len(display_streamlines)) + approx_tube_triangles = display_segments * max(3, int(tube_sides)) * 2 + print( + "Visual smoothing: " + "VTK vtkSplineFilter/cardinal spline, " + f"{spline_subdiv} segments per streamline " + f"({before_points:,} -> {after_points:,} points)" + ) + if mode == "tube": + print( + "Display geometry after smoothing: " + f"approx_tube_triangles={approx_tube_triangles:,}" + ) + return display_poly + + poly = _make_display_polydata(streamlines_xyz, color_values) + + bg = _parse_background_color(background_color) + is_dark = _is_dark_color(bg) + text_color = "white" if is_dark else "black" + bounds_color = (0.75, 0.78, 0.82) if is_dark else (0.25, 0.28, 0.32) + + plotter = pv.Plotter(window_size=window_size, off_screen=not interactive) + if is_dark: + plotter.set_background(bg, top=(0.12, 0.14, 0.18)) + else: + plotter.set_background(bg, top=(0.96, 0.97, 0.99)) + + # Headlight: attached to the camera, always illuminates from the viewer's + # direction — identical behaviour to FURY's default lighting. Even brightness + # across the whole sample regardless of viewing angle. + try: + plotter.clear_lights() + plotter.add_light(pv.Light(light_type="headlight", intensity=0.85)) + except Exception: + pass + + scalar_bar_args = { + "title": color_label, + "vertical": True, + "title_font_size": 16, + "label_font_size": 12, + "color": text_color, + "position_x": 0.88, + "position_y": 0.22, + "width": 0.08, + "height": 0.55, + } + + state = { + "axes": False, + "bounds": False, + "box": False, + "box_actor": None, + "background_dark": is_dark, + "shadows": False, + "streamline_actor": None, + "actors": {"tube": None, "line": None}, + "render_mode": mode, + "tube_thickness": max(float(line_width), 0.5), + "line_thickness": max(float(line_width) * 0.4, 0.5), # thinner LOD lines + "crop_widget": None, + "crop_widget_active": False, + "crop_bounds": None, + "auto_lod": interactive and mode == "tube", + "lod_active": False, + # lighting params — kept in state so sliders can read/write them + "ambient": 0.60, + "diffuse": 0.45, + "specular": 0.08, + "light_intensity": 0.85, + "light_sliders_visible": False, + } + tube_sides = max(3, int(tube_sides)) + + bounds = poly.bounds + full_bounds = tuple(float(v) for v in bounds) + + def _enable_shadows() -> None: + try: + plotter.enable_shadows() + state["shadows"] = True + print("Shadows ON") + except Exception as err: + print(f"Shadows unavailable in this PyVista/VTK setup: {err}") + + def _disable_shadows() -> None: + try: + plotter.disable_shadows() + except Exception: + pass + state["shadows"] = False + print("Shadows OFF") + + def _toggle_shadows() -> None: + if state["shadows"]: + _disable_shadows() + else: + _enable_shadows() + _render() + + if shadows: + _enable_shadows() + + def _add_streamline_actor(render_mode: str, show_scalar_bar: bool) -> None: + # Both modes use the same line PolyData with render_lines_as_tubes — no + # poly.tube() call, so startup and thickness changes are instant. + # Tube mode enables lighting so the GPU shader computes per-pixel tube normals: + # silhouette edges are dark (→ natural black outline), lit faces are bright. + lw = ( + state["tube_thickness"] + if render_mode == "tube" + else state["line_thickness"] + ) + mesh_kwargs = { + "scalars": scalar_name, + "cmap": colormap, + "clim": clim, + "opacity": opacity, + "name": f"streamlines_{render_mode}", + "line_width": lw, + "render_lines_as_tubes": True, + } + if render_mode == "tube": + # Read current lighting values from state so interactive sliders take effect. + mesh_kwargs.update( + { + "lighting": True, + "ambient": state["ambient"], + "diffuse": state["diffuse"], + "specular": state["specular"], + "specular_power": 10, + } + ) + else: + # Line mode: flat colours, no shading, used for LOD during camera motion. + mesh_kwargs["lighting"] = False + + if show_scalar_bar: + mesh_kwargs["scalar_bar_args"] = scalar_bar_args + else: + mesh_kwargs["show_scalar_bar"] = False + + actor = plotter.add_mesh(poly, reset_camera=False, **mesh_kwargs) + try: + actor.GetProperty().SetLineWidth(lw) + except Exception: + pass + actor.SetVisibility(render_mode == state["render_mode"]) + state["actors"][render_mode] = actor + if render_mode == state["render_mode"]: + state["streamline_actor"] = actor + + def _ensure_streamline_actor( + render_mode: str, show_scalar_bar: bool = False + ) -> None: + if state["actors"][render_mode] is None: + _add_streamline_actor(render_mode, show_scalar_bar=show_scalar_bar) + else: + lw = ( + state["tube_thickness"] + if render_mode == "tube" + else state["line_thickness"] + ) + try: + state["actors"][render_mode].GetProperty().SetLineWidth(lw) + except Exception: + pass + + def _show_streamline_mode(render_mode: str) -> None: + _ensure_streamline_actor(render_mode) + for mode_name, actor in state["actors"].items(): + if actor is not None: + actor.SetVisibility(mode_name == render_mode) + state["streamline_actor"] = state["actors"][render_mode] + state["lod_active"] = False + + def _show_interaction_lod(*args) -> None: + if not state["auto_lod"] or state["render_mode"] != "tube": + return + _ensure_streamline_actor("line") + tube_actor = state["actors"].get("tube") + line_actor = state["actors"].get("line") + if tube_actor is None or line_actor is None: + return + tube_actor.SetVisibility(False) + line_actor.SetVisibility(True) + state["streamline_actor"] = line_actor + state["lod_active"] = True + _render() + + def _hide_interaction_lod(*args) -> None: + if not state["lod_active"]: + return + tube_actor = state["actors"].get("tube") + line_actor = state["actors"].get("line") + if tube_actor is not None: + tube_actor.SetVisibility(True) + state["streamline_actor"] = tube_actor + if line_actor is not None: + line_actor.SetVisibility(False) + state["lod_active"] = False + _render() + + def _add_interaction_observer(event_name: str, callback) -> bool: + iren = getattr(plotter, "iren", None) + candidates = [iren, getattr(iren, "interactor", None)] + for candidate in candidates: + if candidate is None: + continue + add_observer = getattr(candidate, "add_observer", None) + if add_observer is not None: + add_observer(event_name, callback) + return True + add_observer = getattr(candidate, "AddObserver", None) + if add_observer is not None: + add_observer(event_name, callback) + return True + return False + + _add_streamline_actor(state["render_mode"], show_scalar_bar=True) + if state["render_mode"] == "tube": + _add_streamline_actor("line", show_scalar_bar=False) + + def _make_outline_mesh(): + xmin, xmax, ymin, ymax, zmin, zmax = poly.bounds + points = np.array( + [ + [xmin, ymin, zmin], + [xmax, ymin, zmin], + [xmax, ymax, zmin], + [xmin, ymax, zmin], + [xmin, ymin, zmax], + [xmax, ymin, zmax], + [xmax, ymax, zmax], + [xmin, ymax, zmax], + ], + dtype=np.float32, + ) + edges = np.array( + [ + [0, 1], + [1, 2], + [2, 3], + [3, 0], + [4, 5], + [5, 6], + [6, 7], + [7, 4], + [0, 4], + [1, 5], + [2, 6], + [3, 7], + ], + dtype=np.int64, + ) + lines = np.column_stack([np.full(len(edges), 2), edges]).ravel() + outline = pv.PolyData(points) + outline.lines = lines + return outline + + outline_mesh = _make_outline_mesh() + + def _show_bounds() -> None: + try: + plotter.show_bounds( + grid="front", + location="outer", + all_edges=False, + color=bounds_color, + font_size=10, + xtitle="X", + ytitle="Y", + ztitle="Z", + ) + except TypeError: + plotter.show_bounds( + grid="front", + location="outer", + all_edges=False, + color=bounds_color, + font_size=10, + xlabel="X", + ylabel="Y", + zlabel="Z", + ) + state["bounds"] = True + + def _hide_bounds() -> None: + try: + plotter.remove_bounds_axes() + except Exception: + try: + plotter.remove_bounds_axis() + except Exception: + pass + state["bounds"] = False + + def _show_axes() -> None: + try: + plotter.add_axes(line_width=2, color=text_color) + except TypeError: + plotter.add_axes() + state["axes"] = True + + def _hide_axes() -> None: + try: + plotter.hide_axes() + except Exception: + pass + state["axes"] = False + + def _render() -> None: + try: + plotter.render() + except Exception: + pass + + def _toggle_background() -> None: + if state["background_dark"]: + plotter.set_background("white", top=(0.96, 0.97, 0.99)) + state["background_dark"] = False + print("Background set to white") + else: + plotter.set_background("black", top=(0.12, 0.14, 0.18)) + state["background_dark"] = True + print("Background set to black") + _render() + + def _toggle_bounds() -> None: + if state["bounds"]: + _hide_bounds() + print("Bounds grid OFF") + else: + _show_bounds() + print("Bounds grid ON") + _render() + + def _toggle_axes() -> None: + if state["axes"]: + _hide_axes() + print("Axes OFF") + else: + _show_axes() + print("Axes ON") + _render() + + def _add_box_actor(): + return plotter.add_mesh( + outline_mesh, + color=bounds_color, + line_width=3, + render_lines_as_tubes=True, + name="streamline_bounding_box", + reset_camera=False, + ) + + def _toggle_box() -> None: + if state["box"]: + try: + plotter.remove_actor(state["box_actor"]) + except Exception: + pass + state["box_actor"] = None + state["box"] = False + print("Bounding box OFF") + else: + state["box_actor"] = _add_box_actor() + state["box"] = True + print("Bounding box ON") + _render() + + def _capture_camera_state(): + camera = plotter.camera + try: + return { + "position": camera.GetPosition(), + "focal_point": camera.GetFocalPoint(), + "view_up": camera.GetViewUp(), + "parallel_scale": camera.GetParallelScale(), + "view_angle": camera.GetViewAngle(), + "parallel_projection": camera.GetParallelProjection(), + } + except Exception: + return None + + def _restore_camera_state(camera_state) -> None: + if camera_state is None: + return + camera = plotter.camera + try: + camera.SetPosition(camera_state["position"]) + camera.SetFocalPoint(camera_state["focal_point"]) + camera.SetViewUp(camera_state["view_up"]) + camera.SetParallelScale(camera_state["parallel_scale"]) + camera.SetViewAngle(camera_state["view_angle"]) + if camera_state["parallel_projection"]: + camera.ParallelProjectionOn() + else: + camera.ParallelProjectionOff() + plotter.reset_camera_clipping_range() + except Exception: + pass + + def _set_displayed_streamlines( + new_streamlines: list[np.ndarray], new_colors: list[np.ndarray] + ) -> None: + nonlocal poly, outline_mesh + camera_state = _capture_camera_state() + + for actor in state["actors"].values(): + if actor is not None: + try: + plotter.remove_actor(actor) + except Exception: + pass + state["actors"] = {"tube": None, "line": None} + state["streamline_actor"] = None + + poly = _make_display_polydata(new_streamlines, new_colors) + outline_mesh = _make_outline_mesh() + + if state["box"]: + try: + plotter.remove_actor(state["box_actor"]) + except Exception: + pass + state["box_actor"] = _add_box_actor() + + if state["bounds"]: + _hide_bounds() + _show_bounds() + + _ensure_streamline_actor(state["render_mode"], show_scalar_bar=True) + if state["auto_lod"]: + _ensure_streamline_actor("line") + _show_streamline_mode(state["render_mode"]) + _restore_camera_state(camera_state) + _render() + + def _print_interactive_crop_summary( + bounds: tuple[float, float, float, float, float, float], + cropped_streamlines: list[np.ndarray], + ) -> None: + lengths = np.array([len(sl) for sl in cropped_streamlines], dtype=np.int64) + mins, maxs = _compute_streamline_bounds(cropped_streamlines) + box_size = maxs - mins + print("\nInteractive crop applied") + print("-" * 24) + print(f"Streamlines: {len(cropped_streamlines):,}") + print( + "Points per streamline: " + f"mean={float(lengths.mean()):.1f}, " + f"median={float(np.median(lengths)):.1f}, " + f"min={int(lengths.min())}, max={int(lengths.max())}" + ) + print( + "Crop box: " + f"x=[{bounds[0]:.2f}, {bounds[1]:.2f}], " + f"y=[{bounds[2]:.2f}, {bounds[3]:.2f}], " + f"z=[{bounds[4]:.2f}, {bounds[5]:.2f}]" + ) + print( + "Displayed bounds: " + f"x=[{mins[0]:.2f}, {maxs[0]:.2f}], " + f"y=[{mins[1]:.2f}, {maxs[1]:.2f}], " + f"z=[{mins[2]:.2f}, {maxs[2]:.2f}]" + ) + print( + "Displayed size: " + f"dx={box_size[0]:.2f}, dy={box_size[1]:.2f}, dz={box_size[2]:.2f}" + ) + print(f"Reproduce from CLI: {_format_crop_cli_bounds(bounds)}") + print("-" * 24) + + def _apply_interactive_crop(box) -> None: + try: + selected_bounds = tuple(float(v) for v in box.bounds) + except Exception as err: + print(f"Interactive crop failed to read box bounds: {err}") + return + + cropped_streamlines, cropped_colors = _crop_streamlines_by_bounds( + base_streamlines_xyz, base_color_values, selected_bounds + ) + if not cropped_streamlines: + print("Interactive crop ignored: no streamlines inside the selected box.") + return + + state["crop_bounds"] = selected_bounds + _set_displayed_streamlines(cropped_streamlines, cropped_colors) + _print_interactive_crop_summary(selected_bounds, cropped_streamlines) + + def _disable_crop_widget(quiet: bool = False) -> None: + try: + plotter.clear_box_widgets() + except Exception: + widget = state.get("crop_widget") + if widget is not None: + try: + widget.Off() + except Exception: + pass + state["crop_widget"] = None + state["crop_widget_active"] = False + if not quiet: + print( + "Interactive crop box OFF; current crop is kept. Press X to clear it." + ) + _render() + + def _enable_crop_widget() -> None: + if state["crop_widget_active"]: + return + widget_bounds = state["crop_bounds"] or full_bounds + try: + state["crop_widget"] = plotter.add_box_widget( + callback=_apply_interactive_crop, + bounds=widget_bounds, + factor=1.0, + rotation_enabled=False, + color=bounds_color, + use_planes=False, + outline_translation=True, + interaction_event="end", + ) + except TypeError: + state["crop_widget"] = plotter.add_box_widget( + callback=_apply_interactive_crop, + bounds=widget_bounds, + factor=1.0, + rotation_enabled=False, + color=bounds_color, + use_planes=False, + outline_translation=True, + ) + except Exception as err: + print(f"Interactive crop box unavailable: {err}") + return + state["crop_widget_active"] = True + print( + "Interactive crop box ON: drag the box faces, crop updates when you release. " + "Press X to clear crop." + ) + _render() + + def _toggle_interactive_crop() -> None: + if state["crop_widget_active"]: + _disable_crop_widget() + else: + _enable_crop_widget() + + def _clear_interactive_crop() -> None: + was_active = state["crop_widget_active"] + if was_active: + _disable_crop_widget(quiet=True) + state["crop_bounds"] = None + _set_displayed_streamlines(base_streamlines_xyz, base_color_values) + print( + "Interactive crop cleared; showing all streamlines after command-line filters." + ) + if was_active: + _enable_crop_widget() + + def _redraw_current_streamline_actor() -> None: + render_mode = state["render_mode"] + camera_state = _capture_camera_state() + actor = state["actors"].get(render_mode) + if actor is not None: + try: + plotter.remove_actor(actor) + except Exception: + pass + state["actors"][render_mode] = None + _ensure_streamline_actor(render_mode) + _show_streamline_mode(render_mode) + _restore_camera_state(camera_state) + + def _change_thickness(factor: float) -> None: + # Fast path: just update line_width on the existing actor — no geometry rebuild. + if state["render_mode"] == "tube": + lower, upper = 0.5, 200.0 + state["tube_thickness"] = min( + upper, max(lower, state["tube_thickness"] * factor) + ) + lw = state["tube_thickness"] + print(f"Thickness: {lw:.1f} px") + else: + lower, upper = 0.5, 100.0 + state["line_thickness"] = min( + upper, max(lower, state["line_thickness"] * factor) + ) + lw = state["line_thickness"] + print(f"Thickness: {lw:.1f} px") + + # Update all relevant actors in-place (no rebuild). + for mode_name, act in state["actors"].items(): + if act is None: + continue + expected_lw = ( + state["tube_thickness"] + if mode_name == "tube" + else state["line_thickness"] + ) + try: + act.GetProperty().SetLineWidth(expected_lw) + except Exception: + pass + _render() + + def _increase_thickness() -> None: + _change_thickness(1.25) + + def _decrease_thickness() -> None: + _change_thickness(0.8) + + def _toggle_fast_mode() -> None: + # Toggle between tube mode (thick, Phong-shaded, dark outlines) and + # line mode (thin, flat colours, fastest for navigation). + camera_state = _capture_camera_state() + state["render_mode"] = "line" if state["render_mode"] == "tube" else "tube" + state["auto_lod"] = interactive and state["render_mode"] == "tube" + _show_streamline_mode(state["render_mode"]) + _restore_camera_state(camera_state) + _render() + if state["render_mode"] == "line": + print( + f"Line mode ON (flat colours, width={state['line_thickness']:.1f} px)" + ) + else: + print( + f"Tube mode ON (Phong shading, width={state['tube_thickness']:.1f} px)" + ) + + def _reset_view() -> None: + plotter.view_isometric() + plotter.camera.zoom(1.15) + _render() + print("View reset") + + def _save_screenshot() -> None: + ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + out_path = Path.cwd() / f"pyvista_view_{ts}.png" + plotter.screenshot(str(out_path)) + print(f"Saved screenshot to {out_path}") + + def _record_orbit(out_path: str | Path, n_frames: int, fps: int) -> None: + """Orbit the camera once and write every frame. + + Uses plotter.open_movie() / open_gif() + manual write_frame() loop, + following the PyVista MP4 Movie example pattern: + open_movie → (show if off-screen) → write_frame × N → close_movie + This avoids the orbit_on_path(write_frames=True) internal bug. + .mp4 requires imageio-ffmpeg: pip install imageio-ffmpeg + """ + out_path = Path(out_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + camera_state = _capture_camera_state() + + viewup = [0, 0, 1] + focal_point = ( + 0.5 * (full_bounds[0] + full_bounds[1]), + 0.5 * (full_bounds[2] + full_bounds[3]), + 0.5 * (full_bounds[4] + full_bounds[5]), + ) + + try: + orbit_positions = plotter.generate_orbital_path( + n_points=n_frames, viewup=viewup, factor=2.0 + ).points + except Exception as err: + print(f"Could not generate orbital path: {err}") + return + + suffix = out_path.suffix.lower() + try: + if suffix == ".gif": + plotter.open_gif(str(out_path)) + else: + plotter.open_movie(str(out_path), framerate=fps) + except Exception as err: + print( + f"Could not open writer for {out_path}: {err}\n" + "Tip: for .mp4 install imageio-ffmpeg: pip install imageio-ffmpeg" + ) + _restore_camera_state(camera_state) + return + + print(f"Recording {n_frames} frames at {fps} fps → {out_path} …") + try: + from tqdm import tqdm + + frame_iter = tqdm( + orbit_positions, total=n_frames, unit="frame", desc="Encoding" + ) + except ImportError: + frame_iter = orbit_positions # plain iterator, no progress bar + + try: + for i, pos in enumerate(frame_iter): + plotter.camera.position = tuple(float(v) for v in pos) + plotter.camera.focal_point = focal_point + plotter.camera.up = viewup + plotter.reset_camera_clipping_range() + plotter.render() + plotter.write_frame() + except Exception as err: + print(f"Error while recording frame {i}: {err}") + finally: + try: + plotter.close_movie() + except Exception: + pass + + _restore_camera_state(camera_state) + _render() + print(f"Saved orbit video to {out_path}") + + def _record_orbit_interactive() -> None: + ts = datetime.datetime.now().strftime("%Y%m%d_%H%M%S") + out_path = Path.cwd() / f"pyvista_orbit_{ts}.mp4" + _record_orbit(out_path, n_frames=video_frames, fps=video_fps) + + # ------------------------------------------------------------------ lighting + def _apply_lighting_to_actors() -> None: + """Push current state ambient/diffuse/specular onto every tube actor.""" + for act in state["actors"].values(): + if act is None: + continue + try: + prop = act.GetProperty() + prop.SetAmbient(state["ambient"]) + prop.SetDiffuse(state["diffuse"]) + prop.SetSpecular(state["specular"]) + except Exception: + pass + + def _apply_light_intensity(value: float) -> None: + """Set intensity on all lights in the renderer.""" + try: + lights = plotter.renderer.GetLights() + lights.InitTraversal() + light = lights.GetNextItem() + while light: + light.SetIntensity(value) + light = lights.GetNextItem() + except Exception: + pass + + def _toggle_light_sliders() -> None: + if state["light_sliders_visible"]: + # Hide: remove all slider widgets (crop widget is not a slider, safe to clear) + try: + plotter.clear_slider_widgets() + except Exception: + pass + state["light_sliders_visible"] = False + print("Lighting sliders OFF") + _render() + return + + # Show three sliders on the left edge of the window. + slider_style = { + "slider_width": 0.03, + "tube_width": 0.005, + } + + def _cb_ambient(value: float) -> None: + state["ambient"] = round(value, 3) + _apply_lighting_to_actors() + _render() + + def _cb_diffuse(value: float) -> None: + state["diffuse"] = round(value, 3) + _apply_lighting_to_actors() + _render() + + def _cb_light(value: float) -> None: + state["light_intensity"] = round(value, 3) + _apply_light_intensity(value) + _render() + + try: + plotter.add_slider_widget( + _cb_ambient, + rng=[0.0, 1.0], + value=state["ambient"], + title="Ambient", + pointa=(0.02, 0.74), + pointb=(0.22, 0.74), + color=text_color, + **slider_style, + ) + plotter.add_slider_widget( + _cb_diffuse, + rng=[0.0, 1.0], + value=state["diffuse"], + title="Diffuse", + pointa=(0.02, 0.55), + pointb=(0.22, 0.55), + color=text_color, + **slider_style, + ) + plotter.add_slider_widget( + _cb_light, + rng=[0.0, 1.5], + value=state["light_intensity"], + title="Light", + pointa=(0.02, 0.36), + pointb=(0.22, 0.36), + color=text_color, + **slider_style, + ) + state["light_sliders_visible"] = True + print( + "Lighting sliders ON (Ambient / Diffuse / Light intensity) — press I to hide" + ) + except Exception as err: + print(f"Could not create lighting sliders: {err}") + _render() + + # ------------------------------------------------------------------ + + def _print_shortcuts() -> None: + print( + "PyVista keys: B background, G bounds grid, S bounding box, " + "C crop box, X clear crop, M tube/line toggle, L shadows, I lighting sliders, " + "+/- thickness, A axes, R reset view, P save PNG, V record orbit, H show keys" + ) + + if show_bounds: + _show_bounds() + + if show_axes: + _show_axes() + + try: + plotter.enable_anti_aliasing("fxaa") + except Exception: + pass + + plotter.add_key_event("b", _toggle_background) + plotter.add_key_event("g", _toggle_bounds) + plotter.add_key_event("s", _toggle_box) + plotter.add_key_event("c", _toggle_interactive_crop) + plotter.add_key_event("x", _clear_interactive_crop) + plotter.add_key_event("l", _toggle_shadows) + plotter.add_key_event("i", _toggle_light_sliders) + plotter.add_key_event("m", _toggle_fast_mode) + plotter.add_key_event("plus", _increase_thickness) + plotter.add_key_event("equal", _increase_thickness) + plotter.add_key_event("minus", _decrease_thickness) + plotter.add_key_event("a", _toggle_axes) + plotter.add_key_event("r", _reset_view) + plotter.add_key_event("p", _save_screenshot) + plotter.add_key_event("v", _record_orbit_interactive) + plotter.add_key_event("h", _print_shortcuts) + + if state["auto_lod"]: + start_ok = _add_interaction_observer( + "StartInteractionEvent", _show_interaction_lod + ) + end_ok = _add_interaction_observer("EndInteractionEvent", _hide_interaction_lod) + if start_ok and end_ok: + print( + "PyVista auto-fast interaction ON: using lines while moving, tubes when idle." + ) + else: + print( + "PyVista auto-fast interaction unavailable in this PyVista/VTK setup." + ) + + plotter.view_isometric() + plotter.camera.zoom(1.15) + + if interactive: + _print_shortcuts() + + if not interactive and not screenshot_path and not video_path: + raise ValueError( + "Must specify screenshot_path or video_path when interactive=False." + ) + + if screenshot_path: + Path(screenshot_path).parent.mkdir(parents=True, exist_ok=True) + + if video_path and not interactive: + # PyVista MP4 pattern: open_movie → show(auto_close=False) → write_frame × N → close + # show() must be called first to initialise the renderer before any write_frame(). + plotter.show(auto_close=False) + _record_orbit(video_path, n_frames=video_frames, fps=video_fps) + plotter.close() + else: + plotter.show(screenshot=screenshot_path, auto_close=True) diff --git a/src/cardiotensor/visualization/streamlines.py b/src/cardiotensor/visualization/streamlines.py index ca2356a..88560fe 100644 --- a/src/cardiotensor/visualization/streamlines.py +++ b/src/cardiotensor/visualization/streamlines.py @@ -57,6 +57,7 @@ def visualize_streamlines( filter_min_len: int | None = None, downsample_factor: int = 1, max_streamlines: int | None = None, + top_clusters: int | None = None, crop_bounds: tuple | None = None, # ((xmin,xmax),(ymin,ymax),(zmin,zmax)) interactive: bool = True, screenshot_path: str | None = None, @@ -74,6 +75,11 @@ def visualize_streamlines( pyvista_show_axes: bool = True, pyvista_show_bounds: bool = False, pyvista_shadows: bool = False, + random_seed: int | None = None, + session_path: str | Path | None = None, + restore_session: bool = True, + session_settings: dict | None = None, + fury_quality: str = "interactive", ): """ Visualize .trk streamlines with per-point angle-based coloring. @@ -93,29 +99,90 @@ def visualize_streamlines( raise ValueError("Only .trk input is supported here") print(f"Loading .trk streamlines: {p}") - streamlines_xyz, attrs = load_trk_streamlines( - p - ) # attrs is dict[str, List[np.ndarray]] - attrs = { - name.upper(): [np.asarray(arr, dtype=np.float32) for arr in seq] - for name, seq in attrs.items() + streamlines_xyz, attrs, streamline_attrs = load_trk_streamlines( + p, include_per_streamline=True + ) + attrs = {name.upper(): seq for name, seq in attrs.items()} + streamline_attrs = { + name.upper(): values for name, values in streamline_attrs.items() } - # ---- Inform the user of available angle fields ---- + if top_clusters is not None: + if top_clusters < 1: + raise ValueError("top_clusters must be at least 1") + if "CLUSTER_ID" not in attrs: + raise ValueError( + "--top-clusters requires a clustered TRK with a per-point " + "cluster_id field" + ) + + cluster_ids = [] + for values in attrs["CLUSTER_ID"]: + values = np.asarray(values, dtype=np.float64).reshape(-1) + if ( + not values.size + or not np.all(np.isfinite(values)) + or not np.allclose(values, values[0]) + or not np.isclose(values[0], round(float(values[0]))) + ): + raise ValueError( + "Each streamline must have one constant integer cluster_id" + ) + cluster_ids.append(int(round(float(values[0])))) + + cluster_ids = np.asarray(cluster_ids, dtype=np.int64) + unique_ids, counts = np.unique(cluster_ids, return_counts=True) + cluster_sizes = { + int(cluster_id): int(count) + for cluster_id, count in zip(unique_ids, counts) + } + if "CLUSTER_SIZE" in streamline_attrs: + stored_sizes = np.asarray( + streamline_attrs["CLUSTER_SIZE"], dtype=np.float64 + ).reshape(len(streamlines_xyz), -1)[:, 0] + for cluster_id, size in zip(cluster_ids, stored_sizes): + if np.isfinite(size) and size > 0: + cluster_sizes[int(cluster_id)] = max( + cluster_sizes[int(cluster_id)], int(round(float(size))) + ) + + ranked_ids = sorted( + (int(cluster_id) for cluster_id in unique_ids), + key=lambda cluster_id: (-cluster_sizes[cluster_id], cluster_id), + ) + selected_ids = set(ranked_ids[:top_clusters]) + keep = [ + index + for index, cluster_id in enumerate(cluster_ids) + if int(cluster_id) in selected_ids + ] + print( + f"Top-cluster filter: keeping {len(selected_ids):,}/" + f"{len(unique_ids):,} largest clusters and {len(keep):,}/" + f"{len(streamlines_xyz):,} streamlines" + ) + streamlines_xyz = [streamlines_xyz[index] for index in keep] + attrs = { + name: [values[index] for index in keep] + for name, values in attrs.items() + } + + # ---- Inform the user of available stored scalar fields ---- available = list(attrs.keys()) - # Also say that AZ and EL can be computed even if missing print( - "\n🎨 Available angle fields in this .trk:", + "\n🎨 Available per-point fields in this .trk:", available if available else "None stored", ) print( "💡 Note: 'az' and 'el' can still be computed on-the-fly from streamline geometry." ) - print("🧭 You can use: color_by = ha, ia, az, el, elevation, azimuth\n") + options = ["elevation", "azimuth", "az", "el", *available] + print(f"🧭 You can use: color_by = {', '.join(options)}\n") # Decide the color scalar color_mode = color_by.lower().strip() + stored_fields = {name.lower(): name for name in available} color_values: list[np.ndarray] | None = None color_range: tuple[float, float] | None = None color_label = "Angle (deg)" @@ -125,7 +192,9 @@ def visualize_streamlines( color_range = ANGLE_RANGES[key] color_label = f"{key} (deg)" if key in attrs: - color_values = attrs[key] + color_values = [ + np.asarray(arr, dtype=np.float32).reshape(-1) for arr in attrs[key] + ] else: if key in {"AZ", "EL"}: az_list, el_list = _compute_az_el_from_streamlines(streamlines_xyz) @@ -143,8 +212,30 @@ def visualize_streamlines( color_label = ( "Elevation (deg)" if color_mode == "elevation" else "Azimuth (deg)" ) + elif color_mode in stored_fields: + key = stored_fields[color_mode] + color_values = [] + field_min = np.inf + field_max = -np.inf + for values in attrs[key]: + values = np.asarray(values, dtype=np.float32) + if values.ndim != 1: + raise ValueError(f"Per-point field '{key}' must contain one scalar") + color_values.append(values) + finite = values[np.isfinite(values)] + if finite.size: + field_min = min(field_min, float(finite.min())) + field_max = max(field_max, float(finite.max())) + if not np.isfinite(field_min): + raise ValueError(f"Per-point field '{key}' contains no finite values") + if field_min == field_max: + field_max = field_min + 1.0 + color_range = (field_min, field_max) + color_label = key else: - raise ValueError("color_by must be one of: ha, ia, az, el, elevation, azimuth") + raise ValueError( + f"Unknown color_by '{color_by}'. Available: {', '.join(options)}" + ) # Default colormap selection if colormap is None: @@ -152,8 +243,10 @@ def visualize_streamlines( colormap = cm.viridis elif color_mode in {"ha", "ia", "elevation"}: colormap = helix_angle_cmap - else: + elif color_mode in {"az", "azimuth"}: colormap = cm.hsv + else: + colormap = cm.viridis backend = backend.lower().strip() render_mode = mode.lower().strip() @@ -187,6 +280,13 @@ def visualize_streamlines( tube_sides=tube_sides, color_range=color_range, color_label=color_label, + random_seed=random_seed, + session_path=session_path, + restore_session=restore_session, + session_settings=session_settings, + streamlines_file=p, + opacity=pyvista_opacity, + quality=fury_quality, ) elif backend == "pyvista": from cardiotensor.visualization.pyvista_plotting_streamlines import ( @@ -219,6 +319,7 @@ def visualize_streamlines( show_bounds=pyvista_show_bounds, shadows=pyvista_shadows, spline_subdiv=spline_subdiv, + random_seed=random_seed, ) else: raise ValueError("backend must be one of: fury, pyvista") diff --git a/src/cardiotensor/visualization/vector_field.py b/src/cardiotensor/visualization/vector_field.py index d16183b..abdb394 100644 --- a/src/cardiotensor/visualization/vector_field.py +++ b/src/cardiotensor/visualization/vector_field.py @@ -6,7 +6,6 @@ """ import math -import sys from pathlib import Path import numpy as np @@ -99,8 +98,9 @@ def visualize_vector_field( vector_field_path = Path(vector_field_path) if not vector_field_path.exists(): - print(f"❌ Vector field path does not exist: {vector_field_path}") - sys.exit(1) + raise FileNotFoundError( + f"Vector field path does not exist: {vector_field_path}" + ) # Load input volume just for shape data_reader_vol = DataReader(vector_field_path) @@ -112,12 +112,9 @@ def visualize_vector_field( # Downsample if needed if bin_factor > 1: - downsample_vector_volume( + vec_load_dir = downsample_vector_volume( vector_field_path, bin_factor, vector_field_path.parent ) - vec_load_dir = ( - vector_field_path.parent / f"bin{bin_factor}" / vector_field_path.name - ) start_binned = start_idx // bin_factor end_binned = math.ceil(end_idx / bin_factor) if end_idx else None else: diff --git a/tests/test_DataReader.py b/tests/test_DataReader.py index 7079b47..7caa5d1 100644 --- a/tests/test_DataReader.py +++ b/tests/test_DataReader.py @@ -6,6 +6,7 @@ import tifffile from cardiotensor.utils.DataReader import DataReader +from cardiotensor.utils.image_io import initialize_zarr_vector_field @pytest.fixture @@ -43,3 +44,103 @@ def test_read_single_tiff_stack(tmp_path: Path): assert vol.shape == (5, 10, 10) assert np.array_equal(vol[2], stack[2]) + + +def test_read_single_tiff_stack_loads_only_requested_pages(tmp_path: Path, monkeypatch): + stack = np.arange(6 * 8 * 9, dtype=np.uint16).reshape(6, 8, 9) + stack_path = tmp_path / "stack.tif" + tifffile.imwrite(stack_path, stack, photometric="minisblack") + + requested_keys = [] + original_asarray = tifffile.TiffFile.asarray + + def record_key(self, key=None, **kwargs): + requested_keys.append(key) + return original_asarray(self, key=key, **kwargs) + + monkeypatch.setattr(tifffile.TiffFile, "asarray", record_key) + volume = DataReader(stack_path).load_volume(2, 5) + + np.testing.assert_array_equal(volume, stack[2:5]) + assert requested_keys == [slice(2, 5)] + + +def test_single_tiff_fallback_memmap_uses_data_scratch(tmp_path: Path, monkeypatch): + stack = np.arange(3 * 4 * 5, dtype=np.uint16).reshape(3, 4, 5) + stack_path = tmp_path / "stack.tif" + tifffile.imwrite(stack_path, stack, photometric="minisblack") + reader = DataReader(stack_path) + captured = {} + + class FakeSeries: + pages = [object()] + + def asarray(self, *, out): + captured["out"] = Path(out) + mapped = np.memmap( + out, mode="w+", dtype=stack.dtype, shape=stack.shape + ) + mapped[:] = stack + return mapped + + class FakeTiff: + series = [FakeSeries()] + + def __enter__(self): + return self + + def __exit__(self, exc_type, exc_value, traceback): + return False + + monkeypatch.setattr(tifffile, "TiffFile", lambda path: FakeTiff()) + + volume = reader.load_volume(1, 3) + + np.testing.assert_array_equal(volume, stack[1:3]) + assert captured["out"].parent.parent == tmp_path / ".cardiotensor_scratch" + assert not captured["out"].parent.exists() + + +def test_mhd_reads_requested_z_range_as_copy_on_write_memmap(tmp_path: Path): + volume = np.arange(5 * 6 * 7, dtype=np.int16).reshape(5, 6, 7) + raw_path = tmp_path / "volume.raw" + raw_path.write_bytes(volume.astype(" 0 + + +def test_trk_analysis_requests_lazy_geometry_only(monkeypatch, tmp_path): + calls = [] + tractogram = SimpleNamespace( + streamlines=iter([np.array([[0, 0, 0], [1, 0, 0]], dtype=float)]) + ) + + def fake_load(path, **kwargs): + calls.append((path, kwargs)) + return SimpleNamespace(tractogram=tractogram) + + monkeypatch.setattr("nibabel.streamlines.load", fake_load) + path = tmp_path / "streamlines.trk" + + streamlines = list(analysis.iter_streamlines(path)) + + assert len(streamlines) == 1 + assert calls == [(str(path), {"lazy_load": True})] + + +def test_fixed_memory_kde_and_normalized_histogram(): + rng = np.random.default_rng(123) + samples = rng.normal(size=100_000) + grid = np.linspace(-5, 5, 512) + + density = analysis.gaussian_kde_1d(samples, grid) + + assert density.shape == grid.shape + assert np.all(np.isfinite(density)) + assert 0.98 < np.trapezoid(density, grid) < 1.02 + + fig, ax = plt.subplots() + analysis.plot_hist( + ax, + samples, + bins=np.linspace(-5, 5, 51), + xlabel="value", + normalize=True, + label="sample", + ) + heights = [patch.get_height() for patch in ax.patches] + assert np.isclose(max(heights), 1.0) + plt.close(fig) + + +def test_quickbundlesx_saves_clusters_centroids_and_membership(tmp_path): + streamlines = np.empty(5, dtype=object) + streamlines[0] = np.array([[0, 0, 0], [5, 0, 0], [10, 0, 0]]) + streamlines[1] = np.array([[0, 0.2, 0], [5, 0.2, 0], [10, 0.2, 0]]) + streamlines[2] = np.array([[10, 0.1, 0], [5, 0.1, 0], [0, 0.1, 0]]) + streamlines[3] = np.array([[0, 10, 0], [5, 10, 0], [10, 10, 0]]) + streamlines[4] = np.array([[0, 10.2, 0], [5, 10.2, 0], [10, 10.2, 0]]) + path = tmp_path / "streamlines.npy" + np.save(path, streamlines, allow_pickle=True) + + result = analysis.run_quickbundles( + path=path, + key=None, + label="heart", + outdir=tmp_path, + voxel_size=1.0, + min_points=2, + thresholds=[5.0, 1.0], + cluster_points=8, + max_streamlines=0, + seed=0, + cluster_min_size=3, + ) + + assert result["method"] == "QuickBundlesX" + assert result["eligible_streamlines"] == 5 + assert result["sampled_streamlines"] == 5 + assert result["clusters"] == 2 + assert result["saved_clusters"] == 1 + assert result["csv"].exists() + with result["csv"].open(newline="") as input_file: + rows = list(csv.DictReader(input_file)) + assert len(rows) == 2 + assert [row["saved_centroid"] for row in rows] == ["True", "False"] + + membership = np.load(result["membership"]) + np.testing.assert_array_equal( + membership["source_streamline_index"], np.arange(5) + ) + np.testing.assert_array_equal( + np.sort(np.bincount(membership["cluster_id"])), [2, 3] + ) + + import nibabel as nib + + centroids = nib.streamlines.load(str(result["centroids"])).tractogram + assert len(centroids.streamlines) == 1 + np.testing.assert_array_equal( + centroids.data_per_streamline["cluster_size"].ravel(), [3] + ) + for cluster_id, point_values in enumerate( + centroids.data_per_point["cluster_id"] + ): + assert np.all(point_values == cluster_id) + + members = nib.streamlines.load(str(result["members"])).tractogram + assert len(members.streamlines) == 3 + assert [len(streamline) for streamline in members.streamlines] == [3, 3, 3] + np.testing.assert_array_equal( + members.data_per_streamline["cluster_size"].ravel(), [3, 3, 3] + ) + for point_values in members.data_per_point["cluster_id"]: + assert np.all(point_values == 0) + + quickbundles = analysis.run_quickbundles( + path=path, + key=None, + label="heart_qb", + outdir=tmp_path, + voxel_size=1.0, + min_points=2, + thresholds=[1.0], + cluster_points=8, + max_streamlines=0, + seed=0, + max_clusters=1, + ) + + assert quickbundles["method"] == "QuickBundles" + assert quickbundles["clusters"] == 2 + assert quickbundles["saved_clusters"] == 1 + assert len( + nib.streamlines.load(str(quickbundles["members"])).tractogram.streamlines + ) == 3 + with quickbundles["csv"].open(newline="") as input_file: + rows = list(csv.DictReader(input_file)) + assert [row["saved_centroid"] for row in rows] == ["True", "False"] + + +def test_generate_streamlines_writes_configured_spacing_in_mm(monkeypatch, tmp_path): + captured = {} + monkeypatch.setattr( + generator_cli, + "read_conf_file", + lambda path: { + "OUTPUT_PATH": str(tmp_path), + "MASK_PATH": None, + "VECTOR_FORMAT": "zarr", + "ANGLE_MODE": "ha_ia", + "VOXEL_SIZE": 16.495, + }, + ) + monkeypatch.setattr( + generator_cli, "vector_field_path", lambda output, storage_format: output + ) + monkeypatch.setattr( + generator_cli, + "generate_streamlines_from_params", + lambda **kwargs: captured.update(kwargs), + ) + monkeypatch.setattr( + sys, + "argv", + ["cardio-generate-streamlines", str(tmp_path / "parameters.conf")], + ) + + generator_cli.script() + + np.testing.assert_allclose(captured["voxel_sizes_zyx"], [0.016495] * 3) diff --git a/tests/test_convert_trk_to_am.py b/tests/test_convert_trk_to_am.py new file mode 100644 index 0000000..5320ad4 --- /dev/null +++ b/tests/test_convert_trk_to_am.py @@ -0,0 +1,47 @@ +import sys + +import numpy as np + +import cardiotensor.scripts.convert_trk_to_am as converter + + +def test_converter_skips_one_point_streamlines(tmp_path, monkeypatch): + input_path = tmp_path / "streamlines.trk" + input_path.touch() + streamlines = [ + np.array([[0.0, 0.0, 0.0]]), + np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), + ] + attrs = { + "HA": [np.array([5.0]), np.array([10.0, 20.0])], + "IA": [np.array([6.0]), np.array([30.0, 40.0])], + } + captured = {} + + monkeypatch.setattr( + converter, + "load_trk_streamlines", + lambda path: (streamlines, attrs), + ) + + def capture_write(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(converter, "write_spatialgraph_am", capture_write) + monkeypatch.setattr( + sys, + "argv", + [ + "cardio-trk_2_am", + str(input_path), + "--edge-scalar-sources", + "ha,ia", + ], + ) + + converter.script() + + assert len(captured["streamlines_xyz"]) == 1 + assert len(captured["streamlines_xyz"][0]) == 2 + np.testing.assert_allclose(captured["edge_scalar"]["HA"], [15.0]) + np.testing.assert_allclose(captured["edge_scalar"]["IA"], [35.0]) diff --git a/tests/test_downsampling.py b/tests/test_downsampling.py index f25eb70..562d1bc 100644 --- a/tests/test_downsampling.py +++ b/tests/test_downsampling.py @@ -1,11 +1,20 @@ from __future__ import annotations import math +import importlib from pathlib import Path import numpy as np +import pytest -from cardiotensor.utils.downsampling import process_vector_block +from cardiotensor.utils.downsampling import ( + _downsample_vector_axes, + _process_vector_range, + downsample_vector_volume, + downsample_volume, + process_vector_block, +) +from cardiotensor.utils.image_io import initialize_zarr_vector_field # --------------------------------------------------------------------------- # Helpers @@ -76,3 +85,184 @@ def test_process_vector_block_output_dtype(tmp_path: Path): result = np.load(out_file) assert result.dtype == np.float32 + +def test_downsample_vector_axes_does_not_cancel_opposite_signs(): + vectors = np.zeros((3, 2, 2, 2), dtype=np.float32) + vectors[0, 0] = 1.0 + vectors[0, 1] = -1.0 + + result = _downsample_vector_axes(vectors, bin_factor=2) + + assert result.shape == (3, 1, 1) + np.testing.assert_allclose(result[:, 0, 0], [1.0, 0.0, 0.0], atol=1e-6) + np.testing.assert_allclose(np.linalg.norm(result[:, 0, 0]), 1.0, atol=1e-6) + + +def test_downsample_vector_volume_reads_zarr(tmp_path: Path): + store = initialize_zarr_vector_field(tmp_path / "source", (2, 4, 4)) + vector_slice = np.zeros((3, 4, 4), dtype=np.float32) + vector_slice[0] = 1.0 + store.write_slice(vector_slice, 0) + store.write_slice(-vector_slice, 1) + + output_path = downsample_vector_volume( + store.path, bin_factor=2, output_dir=tmp_path / "downsampled" + ) + + result = np.load(output_path / "eigen_vec_000000.npy") + assert result.shape == (3, 2, 2) + np.testing.assert_allclose(result[0], 1.0, atol=1e-6) + np.testing.assert_allclose(result[1:], 0.0, atol=1e-6) + + +def test_downsample_vector_volume_applies_mask_before_binning(tmp_path: Path): + store = initialize_zarr_vector_field(tmp_path / "source", (2, 4, 4)) + vector_slice = np.zeros((3, 4, 4), dtype=np.float32) + vector_slice[1] = 1.0 + vector_slice[0, :2, :2] = 1.0 + vector_slice[1, :2, :2] = 0.0 + store.write_slice(vector_slice, 0) + store.write_slice(vector_slice, 1) + + mask_dir = tmp_path / "myocardium_mask" + mask_dir.mkdir() + mask_slice = np.zeros((2, 2), dtype=np.uint8) + mask_slice[0, 0] = 1 + np.save(mask_dir / "mask_000000.npy", mask_slice) + + output_dir = tmp_path / "downsampled" + output_path = downsample_vector_volume(store.path, 2, output_dir) + unmasked = np.load(output_path / "eigen_vec_000000.npy") + np.testing.assert_allclose(unmasked[:, 1, 1], [0.0, 1.0, 0.0], atol=1e-6) + + output_path = downsample_vector_volume( + store.path, 2, output_dir, mask_path=mask_dir + ) + masked = np.load(output_path / "eigen_vec_000000.npy") + np.testing.assert_allclose(masked[:, 0, 0], [1.0, 0.0, 0.0], atol=1e-6) + np.testing.assert_allclose(masked[:, 1, 1], 0.0, atol=1e-6) + + +def test_empty_mask_block_skips_vector_read(tmp_path: Path): + class VectorReader: + shape = (3, 2, 4, 4) + + def load_volume(self, *args, **kwargs): + raise AssertionError("empty mask block should not read vectors") + + class MaskReader: + shape = (1, 2, 2) + + def load_volume(self, *args, **kwargs): + return np.zeros((1, 2, 2), dtype=np.uint8) + + _process_vector_range( + VectorReader(), + start_index=0, + end_index=2, + bin_factor=2, + output_dir=tmp_path, + output_index=0, + mask_reader=MaskReader(), + mask_y_indices=np.array([0, 0, 1, 1]), + mask_x_indices=np.array([0, 0, 1, 1]), + ) + + result = np.load(tmp_path / "eigen_vec_000000.npy") + np.testing.assert_array_equal(result, np.zeros((3, 2, 2), dtype=np.float32)) + + +def test_streamline_seed_selection_uses_mask(tmp_path: Path, monkeypatch): + tractography = importlib.import_module( + "cardiotensor.tractography.generate_streamlines" + ) + vector_path = tmp_path / "vectors" + fa_path = tmp_path / "FA" + angle_path = tmp_path / "HA" + mask_path = tmp_path / "mask" + for path in (vector_path, fa_path, angle_path, mask_path): + path.mkdir() + + myocardium = np.zeros((2, 2, 2), dtype=np.uint8) + myocardium[:, 0, 0] = 1 + + class FakeReader: + def __init__(self, path): + self.path = Path(path) + self.shape = (3, 2, 2, 2) if self.path == vector_path else (2, 2, 2) + + def load_region(self, **kwargs): + return np.ones((3, 2, 2, 2), dtype=np.float32) + + def load_volume(self, **kwargs): + if self.path == mask_path: + return myocardium.copy() + return np.ones((2, 2, 2), dtype=np.float32) + + captured = {} + + def capture_seed_mask(seed_mask, num_seeds, random_seed): + captured["seed_mask"] = seed_mask.copy() + raise RuntimeError("seed mask captured") + + monkeypatch.setattr(tractography, "DataReader", FakeReader) + monkeypatch.setattr(tractography, "_select_seed_points", capture_seed_mask) + + with pytest.raises(RuntimeError, match="seed mask captured"): + tractography.generate_streamlines_from_params( + vector_field_dir=vector_path, + output_dir=tmp_path / "output", + fa_dir=fa_path, + angle_dir=angle_path, + mask_path=mask_path, + fa_seed_min=0.5, + ) + + np.testing.assert_array_equal(captured["seed_mask"], myocardium.astype(bool)) + + +def test_streamline_generation_rejects_incomplete_scalar_outputs( + tmp_path: Path, monkeypatch +): + tractography = importlib.import_module( + "cardiotensor.tractography.generate_streamlines" + ) + vector_path = tmp_path / "vectors" + fa_path = tmp_path / "FA" + angle_path = tmp_path / "HA" + for path in (vector_path, fa_path, angle_path): + path.mkdir() + + class FakeReader: + def __init__(self, path): + self.path = Path(path) + self.volume_info = {"type": "npy"} + self.shape = (3, 8, 4, 4) if self.path == vector_path else (7, 4, 4) + + monkeypatch.setattr(tractography, "DataReader", FakeReader) + + with pytest.raises(RuntimeError, match="Orientation outputs are incomplete"): + tractography.generate_streamlines_from_params( + vector_field_dir=vector_path, + output_dir=tmp_path / "output", + fa_dir=fa_path, + angle_dir=angle_path, + ) + + +def test_scalar_downsampling_rebuilds_unmarked_cache(tmp_path: Path): + source = tmp_path / "source" + source.mkdir() + for index in range(4): + np.save(source / f"FA_{index:06d}.npy", np.full((4, 4), 255, np.uint8)) + + stale_dir = tmp_path / "output" / "bin2" / "FA" + stale_dir.mkdir(parents=True) + stale_path = stale_dir / "FA_000000.tif" + stale_path.write_bytes(b"stale") + + downsample_volume(source, 2, tmp_path / "output", subfolder="FA", out_ext="tif") + + assert stale_path.stat().st_size > len(b"stale") + assert (stale_dir / "FA_000001.tif").exists() + assert (stale_dir / ".source_complete").exists() diff --git a/tests/test_orientation_computation_functions.py b/tests/test_orientation_computation_functions.py index e1e9bcd..36cd014 100644 --- a/tests/test_orientation_computation_functions.py +++ b/tests/test_orientation_computation_functions.py @@ -1,8 +1,14 @@ +import inspect +from concurrent.futures import ThreadPoolExecutor from pathlib import Path import matplotlib.pyplot as plt import numpy as np +import pytest +import zarr +import cardiotensor.orientation.orientation_computation_functions as orientation_functions +import cardiotensor.orientation.orientation_computation_pipeline as orientation_pipeline from cardiotensor.orientation.orientation_computation_functions import ( adjust_start_end_index, calculate_structure_tensor, @@ -12,10 +18,26 @@ interpolate_points, orient_vectors_z_positive, plot_images, + remove_padding, rotate_vectors_to_new_axis, write_images, ) -from cardiotensor.utils.image_io import write_vector_field +from cardiotensor.orientation.orientation_computation_pipeline import ( + _ensure_structure_tensor_memory, + _memory_aware_worker_count, + _safe_low_memory_chunk_size, + check_already_processed, + compute_orientation, +) +from cardiotensor.tractography.generate_streamlines import _select_seed_points +from cardiotensor.utils.DataReader import DataReader +from cardiotensor.utils.image_io import ( + ZARR_VECTOR_FORMAT_VERSION, + ZarrVectorFieldStore, + initialize_zarr_vector_field, + open_zarr_vector_field, + write_vector_field, +) def test_interpolate_points_linear(): @@ -81,6 +103,112 @@ def test_structure_tensor_and_fa(): assert np.all(fa_map >= 0) and np.all(fa_map <= 1), "FA values out of range [0,1]" +def test_structure_tensor_can_skip_eigenvalue_allocation(monkeypatch): + captured = {} + + def fake_parallel(volume, sigma, rho, **kwargs): + captured.update(kwargs) + vectors = np.zeros((3, *volume.shape), dtype=np.float32) + return None, None, vectors + + monkeypatch.setattr( + orientation_functions, "parallel_structure_tensor_analysis", fake_parallel + ) + values, vectors = calculate_structure_tensor( + np.ones((2, 3, 4), dtype=np.uint16), + sigma=1.0, + rho=1.0, + devices=["cpu"], + return_eigenvalues=False, + ) + + assert values is None + assert captured["eigenvalues"] is None + assert vectors.shape == (3, 2, 3, 4) + + +def test_structure_tensor_can_use_memory_mapped_outputs(tmp_path: Path, monkeypatch): + captured = {} + + def fake_parallel(volume, sigma, rho, **kwargs): + captured.update(kwargs) + kwargs["eigenvectors"][:] = 1 + kwargs["eigenvalues"][:] = 2 + kwargs["progress_callback_fn"](16, 16) + return None, kwargs["eigenvalues"], kwargs["eigenvectors"] + + monkeypatch.setattr( + orientation_functions, "parallel_structure_tensor_analysis", fake_parallel + ) + values, vectors = calculate_structure_tensor( + np.ones((2, 3, 4), dtype=np.uint8), + sigma=1.0, + rho=1.0, + devices=["cpu"], + memmap_dir=tmp_path, + ) + + assert isinstance(values, np.memmap) + assert isinstance(vectors, np.memmap) + np.testing.assert_array_equal(values, 2) + np.testing.assert_array_equal(vectors, 1) + + +def test_vector_only_pipeline_preserves_input_dtype_and_cleans_memmaps( + tmp_path: Path, monkeypatch +): + volume_dir = tmp_path / "volume" + volume_dir.mkdir() + for z in range(3): + np.save( + volume_dir / f"slice_{z:03d}.npy", + np.full((4, 5), z + 1, dtype=np.uint16), + ) + + captured = {} + + def fake_structure_tensor(volume, sigma, rho, **kwargs): + captured["input_dtype"] = volume.dtype + captured.update(kwargs) + scratch = Path(kwargs["memmap_dir"]) + captured["scratch"] = scratch + vectors = np.lib.format.open_memmap( + scratch / "eigenvectors.npy", + mode="w+", + dtype=np.float32, + shape=(3, *volume.shape), + ) + vectors[:] = 0 + vectors[0] = 1 + return None, vectors + + monkeypatch.setattr( + orientation_pipeline, "calculate_structure_tensor", fake_structure_tensor + ) + output_dir = tmp_path / "output" + compute_orientation( + volume_path=str(volume_dir), + output_dir=str(output_dir), + sigma=0.1, + rho=0.1, + axis_points=np.array([(2, 2, 0), (2, 2, 2)]), + write_vectors=True, + write_angles=False, + use_gpu=False, + low_memory=True, + ) + + assert captured["input_dtype"] == np.uint16 + assert captured["return_eigenvalues"] is False + assert captured["scratch"].parent == ( + output_dir / ".cardiotensor_scratch" + ).resolve() + assert not captured["scratch"].exists() + vector_store = open_zarr_vector_field(output_dir) + assert vector_store.vectors.shape == (3, 3, 4, 5) + assert np.all(vector_store.completed[:]) + + def test_vector_rotation_and_angles(): """ Test rotation of a 2D vector field slice (3, Y, X) to a new axis. @@ -99,10 +227,24 @@ def test_vector_rotation_and_angles(): ) assert isinstance(rotated_same, np.ndarray) assert rotated_same.shape == vector_slice.shape - assert np.allclose(rotated_same, vector_slice, atol=1e-6), ( - "Rotation to same axis should not change vectors" + assert np.allclose( + rotated_same, vector_slice, atol=1e-6 + ), "Rotation to same axis should not change vectors" + assert rotated_same.dtype == np.float32 + + +def test_vector_rotation_preserves_input_and_vector_magnitude(): + vector_slice = np.zeros((3, 1, 1), dtype=np.float32) + vector_slice[:, 0, 0] = [2.0, 0.0, 0.0] + original = vector_slice.copy() + + rotated = rotate_vectors_to_new_axis( + vector_slice, np.array([1.0, 0.0, 0.0], dtype=np.float32) ) + np.testing.assert_array_equal(vector_slice, original) + np.testing.assert_allclose(np.linalg.norm(rotated[:, 0, 0]), 2.0, atol=1e-6) + def test_orient_vectors_z_positive_flips_negative_z_vectors(): vector_slice = np.zeros((3, 2, 2), dtype=np.float32) @@ -152,6 +294,8 @@ def test_compute_angles_use_full_vector_without_projection(): expected_intrusion = np.rad2deg(np.arctan2(1, np.hypot(2, 3))) np.testing.assert_allclose(helical, [[expected_helical]], atol=1e-6) np.testing.assert_allclose(intrusion, [[expected_intrusion]], atol=1e-6) + assert helical.dtype == np.float32 + assert intrusion.dtype == np.float32 def test_compute_angles_can_return_projected_intrusion_values(): @@ -208,6 +352,317 @@ def test_write_images_and_vectors(tmp_path: Path): assert eigen_vec_files, "Vector field .npy file was not created" +def test_check_already_processed_prints_missing_file(tmp_path: Path, capsys): + assert not check_already_processed( + str(tmp_path), + 0, + 1, + write_vectors=True, + write_angles=False, + output_format="tif", + vector_format="npy", + ) + + output = capsys.readouterr().out + assert "Missing output file(s) for slice 000000" in output + assert "eigen_vec_000000.npy" in output + + +def test_check_already_processed_prints_corrupted_file(tmp_path: Path, capsys): + vector_dir = tmp_path / "eigen_vec" + vector_dir.mkdir() + vector_path = vector_dir / "eigen_vec_000000.npy" + vector_path.write_bytes(b"invalid") + + with pytest.warns(UserWarning, match="potentially corrupted"): + assert not check_already_processed( + str(tmp_path), + 0, + 1, + write_vectors=True, + write_angles=False, + output_format="tif", + vector_format="npy", + ) + + output = capsys.readouterr().out + assert "Corrupted output file(s) found for slice 000000" in output + assert str(vector_path) in output + + +def test_write_zarr_vector_field_roundtrip_and_restart_marker(tmp_path: Path): + store = initialize_zarr_vector_field(tmp_path, (3, 5, 7)) + vector_slice = np.ones((3, 5, 7), dtype=np.float32) + vector_slice[:, :2, :2] = 0 + + assert not check_already_processed( + str(tmp_path), + 1, + 2, + write_vectors=True, + write_angles=False, + output_format="tif", + vector_format="zarr", + zarr_store=store, + ) + + write_vector_field( + vector_slice, + start_index=1, + output_dir=str(tmp_path), + slice_idx=0, + vector_format="zarr", + zarr_store=store, + ) + + assert store.vectors.chunks == (3, 1, 5, 7) + assert store.vectors.shards == (3, 1, 5, 7) + assert "BloscCodec" in type(store.vectors.compressors[0]).__name__ + assert bool(store.completed[1]) + assert check_already_processed( + str(tmp_path), + 1, + 2, + write_vectors=True, + write_angles=False, + output_format="tif", + vector_format="zarr", + zarr_store=store, + ) + + reader = DataReader(tmp_path / "eigen_vec.zarr") + loaded = reader.load_volume(start_index=1, end_index=2) + assert reader.shape == (3, 3, 5, 7) + np.testing.assert_array_equal(loaded[:, 0], vector_slice) + + +def test_zarr_vector_slices_can_be_written_concurrently(tmp_path: Path): + initialize_zarr_vector_field(tmp_path, (4, 8, 9)) + + def write_slice(z: int) -> None: + local_store = open_zarr_vector_field(tmp_path, mode="r+") + data = np.full((3, 8, 9), z + 1, dtype=np.float32) + local_store.write_slice(data, z) + + with ThreadPoolExecutor(max_workers=4) as executor: + list(executor.map(write_slice, range(4))) + + store = open_zarr_vector_field(tmp_path) + np.testing.assert_array_equal(store.completed[:], np.ones(4, dtype=bool)) + for z in range(4): + np.testing.assert_array_equal(store.vectors[:, z], z + 1) + + +def test_interrupted_zarr_write_does_not_set_completion_marker(tmp_path: Path): + store = initialize_zarr_vector_field(tmp_path, (2, 5, 7)) + vector_slice = np.ones((3, 5, 7), dtype=np.float32) + + class InterruptedVectors: + shape = store.vectors.shape + + def __setitem__(self, key, value): + z = key[1] + store.vectors[:, z, :2, :] = value[:, :2, :] + raise OSError("simulated interrupted write") + + interrupted_store = ZarrVectorFieldStore( + store.path, InterruptedVectors(), store.completed + ) + with pytest.raises(OSError, match="interrupted"): + interrupted_store.write_slice(vector_slice, 0) + + assert not bool(store.completed[0]) + store.write_slice(vector_slice, 0) + assert bool(store.completed[0]) + np.testing.assert_array_equal(store.vectors[:, 0], vector_slice) + + +def test_existing_zarr_store_rejects_incompatible_shards(tmp_path: Path): + group = zarr.open_group( + store=str(tmp_path / "eigen_vec.zarr"), mode="w", zarr_format=3 + ) + group.create_array( + "vectors", + shape=(3, 2, 1024, 1024), + chunks=(3, 1, 512, 512), + shards=(3, 1, 512, 512), + dtype="float32", + fill_value=0.0, + compressors=zarr.codecs.BloscCodec( + cname="zstd", + clevel=3, + shuffle=zarr.codecs.BloscShuffle.bitshuffle, + ), + dimension_names=("component", "z", "y", "x"), + ) + group.create_array( + "completed", shape=(2,), chunks=(1,), dtype="bool", fill_value=False + ) + group.attrs.update( + { + "cardiotensor_format": "vector_field", + "cardiotensor_format_version": ZARR_VECTOR_FORMAT_VERSION, + "axis_order": ["component", "z", "y", "x"], + "components": ["x", "y", "z"], + "masked_fill_value": 0.0, + } + ) + + with pytest.raises(ValueError, match="shards"): + initialize_zarr_vector_field(tmp_path, (2, 1024, 1024)) + + +def test_open_zarr_store_rejects_incompatible_compressor(tmp_path: Path): + group = zarr.open_group( + store=str(tmp_path / "eigen_vec.zarr"), mode="w", zarr_format=3 + ) + group.create_array( + "vectors", + shape=(3, 2, 4, 4), + chunks=(3, 1, 4, 4), + shards=(3, 1, 4, 4), + dtype="float32", + fill_value=0.0, + compressors=zarr.codecs.ZstdCodec(level=3), + dimension_names=("component", "z", "y", "x"), + ) + group.create_array( + "completed", shape=(2,), chunks=(1,), dtype="bool", fill_value=False + ) + group.attrs.update( + { + "cardiotensor_format": "vector_field", + "cardiotensor_format_version": ZARR_VECTOR_FORMAT_VERSION, + "axis_order": ["component", "z", "y", "x"], + "components": ["x", "y", "z"], + "masked_fill_value": 0.0, + } + ) + + with pytest.raises(ValueError, match="Blosc"): + open_zarr_vector_field(tmp_path) + + +def test_memory_aware_workers_limit_massive_angle_slices(): + workers = _memory_aware_worker_count( + num_slices=10, + height=20_000, + width=20_000, + write_angles=True, + available_memory_bytes=64 * 1024**3, + cpu_count=32, + ) + assert workers == 1 + + +def test_memory_aware_workers_respect_slice_and_cpu_counts(): + workers = _memory_aware_worker_count( + num_slices=3, + height=64, + width=64, + write_angles=True, + available_memory_bytes=64 * 1024**3, + cpu_count=2, + ) + assert workers == 2 + + +def test_remove_padding_returns_views_unless_copy_is_requested(): + volume = np.arange(5 * 2 * 3, dtype=np.float32).reshape(5, 2, 3) + values = np.broadcast_to(volume, (3, *volume.shape)).copy() + vectors = values.copy() + + cropped_volume, cropped_values, cropped_vectors = remove_padding( + volume, values, vectors, padding_start=1, padding_end=1 + ) + + np.testing.assert_array_equal(cropped_volume, volume[1:4]) + np.testing.assert_array_equal(cropped_values, values[:, 1:4]) + np.testing.assert_array_equal(cropped_vectors, vectors[:, 1:4]) + assert np.shares_memory(cropped_volume, volume) + assert np.shares_memory(cropped_values, values) + assert np.shares_memory(cropped_vectors, vectors) + + copied_volume, copied_values, copied_vectors = remove_padding( + volume, values, vectors, padding_start=1, padding_end=1, copy=True + ) + assert not np.shares_memory(copied_volume, volume) + assert not np.shares_memory(copied_values, values) + assert not np.shares_memory(copied_vectors, vectors) + + +def test_structure_tensor_memory_check_fails_before_large_allocation(): + with pytest.raises(MemoryError, match="Reduce N_CHUNK"): + _ensure_structure_tensor_memory( + (10, 10, 10), + available_memory_bytes=10_000, + ) + + +def test_low_memory_chunk_size_shrinks_failed_64_gib_dataset(): + safe_chunk, requested_peak, safe_peak = _safe_low_memory_chunk_size( + 50, + 10_920, + 7_660, + 7_385, + np.uint16, + padding=9, + include_eigenvalues=True, + has_mask=True, + write_angles=True, + sigma=0.6, + rho=4.0, + truncate=2.0, + available_memory_bytes=64 * 1024**3, + cpu_count=8, + ) + + assert safe_chunk == 16 + assert requested_peak > 64 * 1024**3 + assert safe_peak <= int(64 * 1024**3 * 0.82) + + +def test_low_memory_runtime_guard_rejects_oversized_padded_chunk(): + with pytest.raises(MemoryError, match="padded low-memory chunk"): + _ensure_structure_tensor_memory( + (59, 7_660, 7_385), + include_eigenvalues=True, + low_memory=True, + available_memory_bytes=64 * 1024**3, + input_dtype=np.uint16, + has_mask=True, + write_angles=True, + sigma=0.6, + rho=4.0, + truncate=2.0, + ) + + +def test_compute_orientation_keeps_new_options_at_end_of_public_signature(): + assert list(inspect.signature(compute_orientation).parameters)[-3:] == [ + "vector_format", + "low_memory", + "low_memory_dir", + ] + + +def test_compute_orientation_rejects_no_requested_outputs(): + with pytest.raises(ValueError, match="At least one"): + compute_orientation("missing-volume", write_vectors=False, write_angles=False) + + +def test_seed_selection_is_reproducible_and_returns_valid_coordinates(): + seed_mask = np.ones((4, 5, 6), dtype=bool) + first = _select_seed_points(seed_mask, num_seeds=12, random_seed=42) + second = _select_seed_points(seed_mask, num_seeds=12, random_seed=42) + different = _select_seed_points(seed_mask, num_seeds=12, random_seed=7) + + np.testing.assert_array_equal(first, second) + assert not np.array_equal(first, different) + assert first.shape == (12, 3) + assert np.all(seed_mask[tuple(first.T)]) + + def test_plot_images_with_vector_overlay(tmp_path: Path): plt.switch_backend("Agg") diff --git a/tests/test_orientation_pipeline_helpers.py b/tests/test_orientation_pipeline_helpers.py new file mode 100644 index 0000000..17c7ea2 --- /dev/null +++ b/tests/test_orientation_pipeline_helpers.py @@ -0,0 +1,47 @@ +import numpy as np + +from cardiotensor.orientation.orientation_computation_pipeline import ( + _centerline_neighborhood, + _normalize_vectors_in_place, +) + + +def test_centerline_neighborhood_uses_global_slice_index(): + center_line = np.column_stack( + ( + np.arange(20, dtype=np.float64), + np.zeros(20), + np.arange(20, dtype=np.float64), + ) + ) + + result = _centerline_neighborhood(center_line, global_slice_idx=10, buffer=2) + + np.testing.assert_array_equal(result, center_line[8:13]) + + +def test_normalize_vectors_keeps_masked_zero_vectors_finite(): + vectors = np.zeros((3, 1, 1, 2), dtype=np.float32) + vectors[:, 0, 0, 1] = [3.0, 4.0, 0.0] + + returned = _normalize_vectors_in_place(vectors) + + assert returned is vectors + np.testing.assert_array_equal(vectors[:, 0, 0, 0], [0.0, 0.0, 0.0]) + np.testing.assert_allclose(vectors[:, 0, 0, 1], [0.6, 0.8, 0.0]) + assert np.all(np.isfinite(vectors)) + + +def test_normalize_vectors_allocates_norms_one_slice_at_a_time(monkeypatch): + vectors = np.ones((3, 3, 2, 4), dtype=np.float32) + original_norm = np.linalg.norm + shapes = [] + + def recording_norm(array, *args, **kwargs): + shapes.append(array.shape) + return original_norm(array, *args, **kwargs) + + monkeypatch.setattr(np.linalg, "norm", recording_norm) + _normalize_vectors_in_place(vectors) + + assert shapes == [(3, 2, 4)] * 3 diff --git a/tests/test_public_api.py b/tests/test_public_api.py new file mode 100644 index 0000000..94177da --- /dev/null +++ b/tests/test_public_api.py @@ -0,0 +1,96 @@ +from cardiotensor import ( + DataReader, + calculate_intensities, + calculate_structure_tensor, + compute_azimuth_and_elevation, + compute_fraction_anisotropy, + compute_helical_and_intrusion_angles, + compute_orientation, + convert_to_8bit, + export_vector_field_to_vtk, + find_end_points, + generate_streamlines_from_params, + generate_streamlines_from_vector_field, + load_npz_streamlines, + load_trk_streamlines, + plot_intensity, + read_conf_file, + save_intensity, + write_spatialgraph_am, +) +from cardiotensor.analysis import ( + calculate_intensities as analysis_calculate_intensities, +) +from cardiotensor.analysis import find_end_points as analysis_find_end_points +from cardiotensor.analysis import plot_intensity as analysis_plot_intensity +from cardiotensor.analysis import save_intensity as analysis_save_intensity +from cardiotensor.orientation import ( + calculate_structure_tensor as orientation_calculate_structure_tensor, +) +from cardiotensor.orientation import ( + compute_azimuth_and_elevation as orientation_compute_azimuth_and_elevation, +) +from cardiotensor.orientation import ( + compute_fraction_anisotropy as orientation_compute_fraction_anisotropy, +) +from cardiotensor.orientation import ( + compute_helical_and_intrusion_angles as orientation_compute_helical_and_intrusion_angles, +) +from cardiotensor.orientation import ( + compute_orientation as orientation_compute_orientation, +) +from cardiotensor.orientation import ( + rotate_vectors_to_new_axis as orientation_rotate_vectors_to_new_axis, +) +from cardiotensor.tractography import ( + generate_streamlines_from_params as tractography_generate_streamlines_from_params, +) +from cardiotensor.tractography import ( + generate_streamlines_from_vector_field as tractography_generate_streamlines_from_vector_field, +) +from cardiotensor.utils import DataReader as utils_data_reader +from cardiotensor.utils import convert_to_8bit as utils_convert_to_8bit +from cardiotensor.utils import ( + export_vector_field_to_vtk as utils_export_vector_field_to_vtk, +) +from cardiotensor.utils import load_npz_streamlines as utils_load_npz_streamlines +from cardiotensor.utils import load_trk_streamlines as utils_load_trk_streamlines +from cardiotensor.utils import read_conf_file as utils_read_conf_file +from cardiotensor.utils import write_spatialgraph_am as utils_write_spatialgraph_am +from cardiotensor.visualization import ( + visualize_streamlines, + visualize_vector_field, +) + + +def test_root_public_api_exports(): + assert DataReader is utils_data_reader + assert calculate_intensities is analysis_calculate_intensities + assert calculate_structure_tensor is orientation_calculate_structure_tensor + assert compute_azimuth_and_elevation is orientation_compute_azimuth_and_elevation + assert compute_fraction_anisotropy is orientation_compute_fraction_anisotropy + assert compute_helical_and_intrusion_angles is orientation_compute_helical_and_intrusion_angles + assert compute_orientation is orientation_compute_orientation + assert convert_to_8bit is utils_convert_to_8bit + assert export_vector_field_to_vtk is utils_export_vector_field_to_vtk + assert find_end_points is analysis_find_end_points + assert generate_streamlines_from_params is tractography_generate_streamlines_from_params + assert ( + generate_streamlines_from_vector_field + is tractography_generate_streamlines_from_vector_field + ) + assert load_npz_streamlines is utils_load_npz_streamlines + assert load_trk_streamlines is utils_load_trk_streamlines + assert plot_intensity is analysis_plot_intensity + assert read_conf_file is utils_read_conf_file + assert save_intensity is analysis_save_intensity + assert write_spatialgraph_am is utils_write_spatialgraph_am + + +def test_grouped_public_imports_expose_visualization_helpers(): + assert callable(visualize_streamlines) + assert callable(visualize_vector_field) + + +def test_grouped_public_imports_expose_orientation_helper(): + assert callable(orientation_rotate_vectors_to_new_axis) diff --git a/tests/test_slurm_launcher.py b/tests/test_slurm_launcher.py new file mode 100644 index 0000000..f8332a8 --- /dev/null +++ b/tests/test_slurm_launcher.py @@ -0,0 +1,162 @@ +from pathlib import Path + +import numpy as np + +from cardiotensor.launcher.slurm_launcher import monitor_job_output, slurm_launcher +from cardiotensor.utils.image_io import initialize_zarr_vector_field + + +def _create_angle_outputs(output_dir: Path, names: tuple[str, ...], count: int): + for name in names: + folder = output_dir / name + folder.mkdir(parents=True) + for index in range(count): + (folder / f"{name}_{index:06d}.tif").touch() + + +def test_monitor_waits_for_zarr_when_angles_are_complete( + tmp_path: Path, monkeypatch, capsys +): + _create_angle_outputs(tmp_path, ("HA", "IA", "FA"), count=2) + store = initialize_zarr_vector_field(tmp_path, (2, 4, 5)) + store.completed[0] = True + + sleep_calls = [] + + def complete_zarr_after_first_poll(seconds): + sleep_calls.append(seconds) + store.completed[1] = True + + monkeypatch.setattr( + "cardiotensor.launcher.slurm_launcher.time.sleep", + complete_zarr_after_first_poll, + ) + + monitor_job_output( + output_directory=str(tmp_path), + start_index=0, + end_index_exclusive=2, + output_format="tif", + write_angles=True, + write_vectors=True, + vector_format="zarr", + poll_interval_sec=1, + ) + + output = capsys.readouterr().out + assert "HA: 2/2 | IA: 2/2 | FA: 2/2 | Zarr vectors: 1/2" in output + assert "HA: 2/2 | IA: 2/2 | FA: 2/2 | Zarr vectors: 2/2" in output + assert sleep_calls == [1] + + +def test_monitor_supports_projected_angle_names(tmp_path: Path, capsys): + _create_angle_outputs(tmp_path, ("HA_projected", "IA_projected", "FA"), count=1) + + monitor_job_output( + output_directory=str(tmp_path), + start_index=0, + end_index_exclusive=1, + output_format="tif", + write_angles=True, + write_vectors=False, + projected=True, + ) + + output = capsys.readouterr().out + assert "HA_projected: 1/1 | IA_projected: 1/1 | FA: 1/1" in output + + +def test_slurm_low_memory_dry_run_reduces_unsafe_chunk( + tmp_path: Path, monkeypatch, capsys +): + params = { + "IMAGES_PATH": "large-volume", + "MASK_PATH": "large-mask", + "OUTPUT_PATH": str(tmp_path / "output"), + "OUTPUT_FORMAT": "jp2", + "ANGLE_MODE": "ha_ia", + "WRITE_ANGLES": True, + "WRITE_VECTORS": True, + "VECTOR_FORMAT": "npy", + "LOW_MEMORY": True, + "SIGMA": 0.6, + "RHO": 4.0, + "TRUNCATE": 2.0, + "VERTICAL_PADDING": 9, + "N_CHUNK": 50, + "TEST": False, + } + + class FakeReader: + shape = (10_920, 7_660, 7_385) + dtype = np.dtype(np.uint16) + + def __init__(self, path): + pass + + monkeypatch.setattr( + "cardiotensor.launcher.slurm_launcher.read_conf_file", lambda _: params + ) + monkeypatch.setattr("cardiotensor.launcher.slurm_launcher.DataReader", FakeReader) + + slurm_launcher( + "parameters.conf", + start_index=0, + end_index=100, + chunk_size=50, + cpus_per_task=8, + mem_gb=64, + array_parallel=4, + log_dir=str(tmp_path / "logs"), + submit_dir=str(tmp_path / "submit"), + monitor=False, + dry_run=True, + ) + + output = capsys.readouterr().out + assert "adjusted SLURM chunk_size from 50 to 16" in output + script = next((tmp_path / "submit").glob("*.slurm")).read_text() + assert "IMAGES_PER_JOB=16" in script + assert "echo TMPDIR: ${TMPDIR:-/tmp}" in script + + +def test_slurm_skips_complete_chunks(tmp_path: Path, monkeypatch, capsys): + output_dir = tmp_path / "output" + _create_angle_outputs(output_dir, ("HA", "IA", "FA"), count=2) + params = { + "IMAGES_PATH": "volume", + "OUTPUT_PATH": str(output_dir), + "OUTPUT_FORMAT": "tif", + "ANGLE_MODE": "ha_ia", + "WRITE_ANGLES": True, + "WRITE_VECTORS": False, + "N_CHUNK": 2, + "TEST": False, + } + + class FakeReader: + shape = (4, 8, 8) + dtype = np.dtype(np.uint16) + + def __init__(self, path): + pass + + monkeypatch.setattr( + "cardiotensor.launcher.slurm_launcher.read_conf_file", lambda _: params + ) + monkeypatch.setattr("cardiotensor.launcher.slurm_launcher.DataReader", FakeReader) + + slurm_launcher( + "parameters.conf", + chunk_size=2, + log_dir=str(tmp_path / "logs"), + submit_dir=str(tmp_path / "submit"), + monitor=False, + dry_run=True, + ) + + output = capsys.readouterr().out + assert "1 complete job(s) skipped, 1 job(s) needed" in output + scripts = list((tmp_path / "submit").glob("*.slurm")) + assert len(scripts) == 1 + assert "START_INDEX_BASE=2" in scripts[0].read_text() diff --git a/tests/test_streamlines_io_utils.py b/tests/test_streamlines_io_utils.py index 3998d8a..a0ae1c1 100644 --- a/tests/test_streamlines_io_utils.py +++ b/tests/test_streamlines_io_utils.py @@ -131,6 +131,18 @@ def test_normalize_attrs_already_degrees(): np.testing.assert_allclose(result["IA"][0], [-0.5, 0.0, 0.5], atol=1e-6) +def test_normalize_attrs_preserves_positive_degrees_across_field(): + attrs = { + "HA": [ + np.array([5.0], dtype=np.float32), + np.array([10.0, 20.0], dtype=np.float32), + ] + } + result = normalize_attrs_to_degrees(attrs) + np.testing.assert_allclose(result["HA"][0], [5.0]) + np.testing.assert_allclose(result["HA"][1], [10.0, 20.0]) + + # --------------------------------------------------------------------------- # compute_elevation_angles # --------------------------------------------------------------------------- diff --git a/tests/test_tractography_streamlines.py b/tests/test_tractography_streamlines.py new file mode 100644 index 0000000..72a0288 --- /dev/null +++ b/tests/test_tractography_streamlines.py @@ -0,0 +1,64 @@ +import numpy as np + +from cardiotensor.tractography.generate_streamlines import ( + _fa_to_unit_interval, + generate_streamlines_from_vector_field, + trace_streamline, + trilinear_interpolate_vector, +) + + +def test_trilinear_interpolate_vector_aligns_reversed_corners_to_reference(): + vector_field = np.zeros((3, 1, 1, 2), dtype=np.float32) + vector_field[:, 0, 0, 0] = [1.0, 0.0, 0.0] + vector_field[:, 0, 0, 1] = [-1.0, 0.0, 0.0] + + raw = trilinear_interpolate_vector(vector_field, (0.0, 0.0, 0.5)) + aligned = trilinear_interpolate_vector( + vector_field, + (0.0, 0.0, 0.5), + reference_vector=np.array([1.0, 0.0, 0.0]), + ) + + assert np.allclose(raw, [0.0, 0.0, 0.0]) + assert np.allclose(aligned, [1.0, 0.0, 0.0]) + + +def test_trace_streamline_continues_across_reversed_axis_region(): + vector_field = np.zeros((3, 3, 3, 8), dtype=np.float32) + vector_field[0, :, :, :4] = 1.0 + vector_field[0, :, :, 4:] = -1.0 + + streamline = trace_streamline( + start_pt=(1.0, 1.0, 1.0), + vector_field=vector_field, + step_length=0.5, + max_steps=10, + angle_threshold=30.0, + ) + + assert len(streamline) > 6 + assert streamline[-1][2] > 4.0 + + +def test_fa_to_unit_interval_decodes_uint8_and_preserves_float_data(): + encoded = np.array([0, 128, 255], dtype=np.uint8) + decoded = _fa_to_unit_interval(encoded) + np.testing.assert_allclose(decoded, [0.0, 128.0 / 255.0, 1.0]) + assert decoded.dtype == np.float32 + + scientific = np.array([0.0, 0.25, 1.0], dtype=np.float32) + np.testing.assert_array_equal(_fa_to_unit_interval(scientific), scientific) + + +def test_streamline_minimum_is_at_least_two_points(): + vector_field = np.zeros((3, 1, 1, 1), dtype=np.float32) + streamlines = generate_streamlines_from_vector_field( + vector_field, + seed_points=np.array([[0, 0, 0]]), + min_length_pts=0, + max_steps=1, + bidirectional=False, + ) + + assert streamlines == [] diff --git a/tests/test_utils.py b/tests/test_utils.py index b4940fa..31e7749 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,9 +1,40 @@ +from pathlib import Path +from types import SimpleNamespace from unittest import TestCase import numpy as np import pytest -from cardiotensor.utils.utils import convert_to_8bit, read_conf_file +from cardiotensor.utils import utils +from cardiotensor.utils.utils import ( + _memory_available_from_cgroup, + convert_to_8bit, + get_available_memory_bytes, + read_conf_file, +) + + +def test_memory_available_from_cgroup_v2(tmp_path: Path): + proc_cgroup = tmp_path / "cgroup" + proc_cgroup.write_text("0::/slurm/job_123\n") + job_cgroup = tmp_path / "sys_fs_cgroup" / "slurm" / "job_123" + job_cgroup.mkdir(parents=True) + (job_cgroup / "memory.max").write_text("1000") + (job_cgroup / "memory.current").write_text("250") + + assert _memory_available_from_cgroup(proc_cgroup, tmp_path / "sys_fs_cgroup") == 750 + + (job_cgroup / "memory.current").write_text("1000") + assert _memory_available_from_cgroup(proc_cgroup, tmp_path / "sys_fs_cgroup") == 0 + + +def test_available_memory_uses_cgroup_limit(monkeypatch): + monkeypatch.setattr( + utils.psutil, "virtual_memory", lambda: SimpleNamespace(available=1000) + ) + monkeypatch.setattr(utils, "_memory_available_from_cgroup", lambda: 400) + + assert get_available_memory_bytes() == 400 def test_convert_to_8bit(): @@ -13,6 +44,19 @@ def test_convert_to_8bit(): assert out.min() >= 0 and out.max() <= 255 +def test_convert_to_8bit_skips_percentiles_for_explicit_range(monkeypatch): + def fail_if_called(*args, **kwargs): + raise AssertionError("nanpercentile should not be called") + + monkeypatch.setattr(np, "nanpercentile", fail_if_called) + out = convert_to_8bit( + np.array([0.0, 0.5, 1.0], dtype=np.float32), + min_value=0.0, + max_value=1.0, + ) + np.testing.assert_array_equal(out, [0, 127, 255]) + + def test_read_conf_file(tmp_path): # Create dummy directories for IMAGES_PATH and MASK_PATH images_dir = tmp_path / "images" @@ -24,8 +68,8 @@ def test_read_conf_file(tmp_path): conf_file = tmp_path / "test.conf" conf_file.write_text( "[DATASET]\n" - f"IMAGES_PATH = {images_dir}\n" - f"MASK_PATH = {mask_file}\n" + "IMAGES_PATH = images\n" + "MASK_PATH = mask.tif\n" "VOXEL_SIZE = 0.5\n\n" "[STRUCTURE TENSOR CALCULATION]\n" "SIGMA = 2.0\n" @@ -34,6 +78,8 @@ def test_read_conf_file(tmp_path): "VERTICAL_PADDING = 10.0\n" "N_CHUNK = 50\n" "USE_GPU = True\n" + "LOW_MEMORY = True\n" + "LOW_MEMORY_DIR = scratch\n" "WRITE_VECTORS = True\n" "REVERSE = False\n\n" "[ANGLE CALCULATION]\n" @@ -63,7 +109,10 @@ def test_read_conf_file(tmp_path): assert pytest.approx(config["RHO"]) == 1.5 assert config["N_CHUNK"] == 50 assert config["USE_GPU"] is True + assert config["LOW_MEMORY"] is True + assert config["LOW_MEMORY_DIR"] == str(tmp_path / "scratch") assert config["WRITE_VECTORS"] is True + assert config["VECTOR_FORMAT"] == "zarr" assert config["REVERSE"] is False # Angles @@ -78,6 +127,23 @@ def test_read_conf_file(tmp_path): assert config["SHOW_QUIVER"] is True # Output - assert config["OUTPUT_PATH"] == "./output" + assert config["OUTPUT_PATH"] == str(tmp_path / "output") assert config["OUTPUT_FORMAT"] == "tif" assert config["OUTPUT_TYPE"] == "rgb" + + +def test_read_conf_file_rejects_no_requested_outputs(tmp_path): + images_dir = tmp_path / "images" + images_dir.mkdir() + conf_file = tmp_path / "invalid.conf" + conf_file.write_text( + "[DATASET]\n" + "IMAGES_PATH = images\n\n" + "[STRUCTURE TENSOR CALCULATION]\n" + "WRITE_VECTORS = False\n\n" + "[ANGLE CALCULATION]\n" + "WRITE_ANGLES = False\n" + ) + + with pytest.raises(ValueError, match="At least one"): + read_conf_file(conf_file) diff --git a/tests/test_visualize_streamlines_performance.py b/tests/test_visualize_streamlines_performance.py new file mode 100644 index 0000000..396e4bc --- /dev/null +++ b/tests/test_visualize_streamlines_performance.py @@ -0,0 +1,715 @@ +import json +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest + +import cardiotensor.scripts.visualize_streamlines as visualize_cli +import cardiotensor.visualization.streamlines as streamline_visualization +from cardiotensor.scripts.visualize_streamlines import ( + _apply_session_settings, + _discover_trk_color_fields, +) +from cardiotensor.visualization import fury_plotting_streamlines +from cardiotensor.visualization import pyvista_plotting_streamlines + + +def test_fury_orbit_frames_use_video_output_scratch(monkeypatch, tmp_path): + viewer = object.__new__(fury_plotting_streamlines.StreamlineViewer) + viewer.actor_fast = SimpleNamespace(SetVisibility=lambda visible: None) + viewer.actor0 = SimpleNamespace(SetVisibility=lambda visible: None) + viewer.scene = SimpleNamespace(azimuth=lambda degrees: None) + viewer.window_size = (100, 100) + viewer._set_default_camera = lambda: None + + monkeypatch.setattr( + fury_plotting_streamlines.fury.window, + "record", + lambda scene, out_path, size, reset_camera: Path(out_path).write_bytes(b"png"), + ) + + video_path = tmp_path / "analysis" / "orbit.mp4" + + def fake_video_writer(frame_paths, out_path, fps): + scratch_root = video_path.parent / ".cardiotensor_scratch" + assert all(path.parent.parent == scratch_root for path in frame_paths) + Path(out_path).write_bytes(b"video") + + monkeypatch.setattr( + fury_plotting_streamlines, + "_write_frame_sequence_to_video", + fake_video_writer, + ) + + viewer._record_orbit(video_path, video_frames=2, video_fps=30) + + assert video_path.exists() + assert not list( + (video_path.parent / ".cardiotensor_scratch").glob( + "cardiotensor_fury_orbit_*" + ) + ) + + +def test_color_field_discovery_uses_lazy_trk_loading(monkeypatch, tmp_path): + load_calls = [] + tractogram = SimpleNamespace(data_per_point={"HA": [], "IA": []}) + + def fake_load(path, **kwargs): + load_calls.append((path, kwargs)) + return SimpleNamespace(tractogram=tractogram) + + monkeypatch.setattr("nibabel.streamlines.load", fake_load) + + fields = _discover_trk_color_fields(tmp_path / "streamlines.trk") + + assert fields == ["HA", "IA"] + assert load_calls == [(str(tmp_path / "streamlines.trk"), {"lazy_load": True})] + + +def test_visualizer_accepts_custom_per_point_scalar(monkeypatch, tmp_path): + trk_path = tmp_path / "centroids.trk" + trk_path.write_bytes(b"trk") + streamlines = [ + np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32), + np.array([[0, 1, 0], [1, 1, 0]], dtype=np.float32), + ] + cluster_ids = [ + np.array([0, 0], dtype=np.float32), + np.array([1, 1], dtype=np.float32), + ] + captured = {} + + monkeypatch.setattr( + streamline_visualization, + "load_trk_streamlines", + lambda path, include_per_streamline=False: ( + streamlines, + {"cluster_id": cluster_ids}, + {}, + ), + ) + monkeypatch.setattr( + fury_plotting_streamlines, + "show_streamlines", + lambda **kwargs: captured.update(kwargs), + ) + + streamline_visualization.visualize_streamlines( + trk_path, color_by="cluster_id", interactive=False + ) + + np.testing.assert_array_equal(captured["color_values"][0], [0, 0]) + np.testing.assert_array_equal(captured["color_values"][1], [1, 1]) + assert captured["color_range"] == (0.0, 1.0) + assert captured["color_label"] == "CLUSTER_ID" + + +def test_top_clusters_keeps_every_member_of_the_largest_clusters( + monkeypatch, tmp_path +): + trk_path = tmp_path / "members.trk" + trk_path.write_bytes(b"trk") + streamlines = [ + np.array([[0, index, 0], [1, index, 0]], dtype=np.float32) + for index in range(6) + ] + cluster_ids = [0, 0, 1, 1, 1, 2] + cluster_values = [ + np.full(len(streamline), cluster_id, dtype=np.float32) + for streamline, cluster_id in zip(streamlines, cluster_ids) + ] + captured = {} + + monkeypatch.setattr( + streamline_visualization, + "load_trk_streamlines", + lambda path, include_per_streamline=False: ( + streamlines, + {"cluster_id": cluster_values}, + {"cluster_size": np.array([[2], [2], [3], [3], [3], [1]])}, + ), + ) + monkeypatch.setattr( + fury_plotting_streamlines, + "show_streamlines", + lambda **kwargs: captured.update(kwargs), + ) + + streamline_visualization.visualize_streamlines( + trk_path, + color_by="cluster_id", + top_clusters=2, + interactive=False, + ) + + assert len(captured["streamlines_xyz"]) == 5 + assert { + int(values[0]) for values in captured["color_values"] + } == {0, 1} + + +def test_top_clusters_rejects_normal_tractogram(monkeypatch, tmp_path): + trk_path = tmp_path / "normal.trk" + trk_path.write_bytes(b"trk") + streamlines = [ + np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32) + ] + monkeypatch.setattr( + streamline_visualization, + "load_trk_streamlines", + lambda path, include_per_streamline=False: ( + streamlines, + {"HA": [np.array([0, 0], dtype=np.float32)]}, + {}, + ), + ) + + with pytest.raises(ValueError, match="requires a clustered TRK"): + streamline_visualization.visualize_streamlines( + trk_path, + color_by="HA", + top_clusters=1, + interactive=False, + ) + + +def test_cli_accepts_hyphenated_custom_color_field(monkeypatch, tmp_path): + trk_path = tmp_path / "centroids.trk" + trk_path.write_bytes(b"trk") + captured = {} + + monkeypatch.setattr( + visualize_cli, "_discover_trk_color_fields", lambda path: ["cluster_id"] + ) + monkeypatch.setattr( + visualize_cli, "visualize_streamlines", lambda **kwargs: captured.update(kwargs) + ) + monkeypatch.setattr( + visualize_cli.sys, + "argv", + [ + "cardio-visualize-streamlines", + str(trk_path), + "--color-by", + "cluster-id", + "--top-clusters", + "3", + ], + ) + + visualize_cli.script() + + assert captured["color_by"] == "cluster_id" + assert captured["top_clusters"] == 3 + + +def test_cli_uses_default_session_path_when_none_is_given(monkeypatch, tmp_path): + trk_path = tmp_path / "streamlines.trk" + trk_path.write_bytes(b"trk") + captured = {} + + monkeypatch.setattr(visualize_cli, "_discover_trk_color_fields", lambda path: []) + monkeypatch.setattr( + visualize_cli, "visualize_streamlines", lambda **kwargs: captured.update(kwargs) + ) + monkeypatch.setattr( + visualize_cli.sys, "argv", ["cardio-visualize-streamlines", str(trk_path)] + ) + + visualize_cli.script() + + assert captured["session_path"] == tmp_path / "streamlines_session.json" + assert captured["restore_session"] is False + + +def test_cli_asks_before_restoring_default_session(monkeypatch, tmp_path): + trk_path = tmp_path / "streamlines.trk" + trk_path.write_bytes(b"trk") + session_path = tmp_path / "streamlines_session.json" + session_path.write_text( + json.dumps( + { + "streamlines_file": str(trk_path), + "streamlines_size": 3, + "settings": {"line_width": 7.0}, + } + ) + ) + monkeypatch.setattr(visualize_cli, "_discover_trk_color_fields", lambda path: []) + monkeypatch.setattr( + visualize_cli.sys, "stdin", SimpleNamespace(isatty=lambda: True) + ) + monkeypatch.setattr( + visualize_cli.sys, "argv", ["cardio-visualize-streamlines", str(trk_path)] + ) + + for answer, expected_restore, expected_width in ( + ("y", True, 7.0), + ("n", False, 4.0), + ): + captured = {} + monkeypatch.setattr("builtins.input", lambda prompt, value=answer: value) + monkeypatch.setattr( + visualize_cli, + "visualize_streamlines", + lambda **kwargs: captured.update(kwargs), + ) + + visualize_cli.script() + + assert captured["session_path"] == session_path + assert captured["restore_session"] is expected_restore + assert captured["line_width"] == expected_width + + +def test_fury_window_close_saves_once_and_terminates(): + calls = {"saved": 0, "exited": 0, "widgets_disabled": 0} + + class FakeWidget: + def EnabledOff(self): + calls["widgets_disabled"] += 1 + + viewer = fury_plotting_streamlines.StreamlineViewer.__new__( + fury_plotting_streamlines.StreamlineViewer + ) + viewer._closing = False + viewer.plane_widget = FakeWidget() + viewer.box_widget = FakeWidget() + viewer.showm = SimpleNamespace( + exit=lambda: calls.__setitem__("exited", calls["exited"] + 1) + ) + viewer._autosave_session = lambda: calls.__setitem__("saved", calls["saved"] + 1) + + viewer._close_window() + viewer._close_window() + + assert calls == {"saved": 1, "exited": 1, "widgets_disabled": 2} + + +def test_pyvista_connectivity_is_built_correctly(monkeypatch): + class FakePolyData: + def __init__(self, points): + self.points = points + self.lines = None + self.point_data = {} + self.active_scalars = None + + def set_active_scalars(self, name): + self.active_scalars = name + + fake_pyvista = SimpleNamespace(PolyData=FakePolyData) + monkeypatch.setattr( + pyvista_plotting_streamlines, "_load_pyvista", lambda: fake_pyvista + ) + + streamlines = [ + np.array([[0, 0, 0], [1, 0, 0]], dtype=np.float32), + np.array([[2, 0, 0], [2, 1, 0], [2, 2, 0]], dtype=np.float32), + ] + colors = [ + np.array([10, 20], dtype=np.float32), + np.array([30, 40, 50], dtype=np.float32), + ] + + poly = pyvista_plotting_streamlines._streamlines_to_polydata( + streamlines, colors, "HA" + ) + + np.testing.assert_array_equal(poly.points, np.concatenate(streamlines)) + np.testing.assert_array_equal(poly.lines, [2, 0, 1, 3, 2, 3, 4]) + np.testing.assert_array_equal(poly.point_data["HA"], [10, 20, 30, 40, 50]) + assert poly.active_scalars == "HA" + + +def test_pyvista_bounds_ignore_empty_streamlines(): + streamlines = [ + np.empty((0, 3), dtype=np.float32), + np.array([[1, 5, -2], [4, 3, 8]], dtype=np.float32), + np.array([[-1, 7, 0], [2, 6, 3]], dtype=np.float32), + ] + + mins, maxs = pyvista_plotting_streamlines._compute_streamline_bounds(streamlines) + + np.testing.assert_array_equal(mins, [-1, 3, -2]) + np.testing.assert_array_equal(maxs, [4, 7, 8]) + + +def test_fury_screenshot_uses_current_viewer_window(monkeypatch, tmp_path): + calls = {} + + class FakeRenderWindow: + def Render(self): + calls["rendered"] = True + + class FakeOutput: + def __init__(self, scale): + self.scale = scale + + def GetDimensions(self): + return (3200 * self.scale, 1800 * self.scale, 1) + + class FakeCapture: + scale = 1 + + def SetInput(self, render_window): + calls["capture_input"] = render_window + + def SetInputBufferTypeToRGB(self): + calls["rgb"] = True + + def Update(self): + calls["updated"] = True + + def SetScale(self, scale): + self.scale = scale + calls["scale"] = scale + + def GetOutputPort(self): + return self + + def GetOutput(self): + return FakeOutput(self.scale) + + class FakeWriter: + def SetFileName(self, filename): + calls["filename"] = filename + + def SetInputConnection(self, connection): + calls["writer_input"] = connection + + def Write(self): + calls["written"] = True + + capture = FakeCapture() + monkeypatch.setattr( + fury_plotting_streamlines.vtk, + "vtkWindowToImageFilter", + lambda: capture, + ) + monkeypatch.setattr(fury_plotting_streamlines.vtk, "vtkPNGWriter", FakeWriter) + + render_window = FakeRenderWindow() + viewer = fury_plotting_streamlines.StreamlineViewer.__new__( + fury_plotting_streamlines.StreamlineViewer + ) + viewer.showm = SimpleNamespace(window=render_window) + out_path = tmp_path / "view.png" + + size = viewer._save_screenshot(out_path) + + assert size == (3200, 1800) + assert calls["capture_input"] is render_window + assert calls["filename"] == str(out_path) + assert calls["rendered"] and calls["updated"] and calls["written"] + + size = viewer._save_screenshot(out_path, scale=2) + assert size == (6400, 3600) + assert calls["scale"] == 2 + + +def test_fury_spline_subdivision_and_crop_box_rotation(): + streamlines = [ + np.array([[index, index % 2, 0] for index in range(10)], dtype=np.float32) + ] + colors = [np.arange(10, dtype=np.float32)] + lut = fury_plotting_streamlines.actor.colormap_lookup_table(scale_range=(0, 9)) + + viewer = fury_plotting_streamlines.StreamlineViewer( + streamlines, + colors, + "line", + 2, + (800, 600), + lut, + spline_subdiv=2, + ) + assert viewer.actor_spline_subdiv == 2 + viewer.actor0.GetMapper().Update() + assert viewer.actor0.GetMapper().GetInput().GetNumberOfPoints() == 3 + + viewer.showm = fury_plotting_streamlines.window.ShowManager( + scene=viewer.scene, size=(800, 600), reset_camera=False + ) + viewer._setup_box_widget() + assert viewer.box_widget.GetRotationEnabled() == 1 + + viewer.box_clipping_active = True + viewer._apply_clipping_planes() + assert viewer.actor0.GetMapper().GetNumberOfClippingPlanes() == 6 + + viewer.clipping_active = True + viewer._apply_clipping_planes() + assert viewer.actor0.GetMapper().GetNumberOfClippingPlanes() == 7 + + +def test_saved_settings_are_defaults_but_cli_options_win(): + args = SimpleNamespace( + line_width=2.0, + subsample=1, + crop_x=None, + crop_y=None, + crop_z=None, + width=800, + height=800, + hide_axes=False, + show_bounds=False, + hide_bounds=False, + shadows=False, + no_shadows=False, + ) + settings = { + "line_width": 6.0, + "tube_thickness": 0.25, + "subsample_factor": 4, + "crop_bounds": [[1, 2], [3, 4], [5, 6]], + "window_size": [1600, 1000], + "show_axes": False, + "show_bounds": True, + "shadows": True, + } + + _apply_session_settings(args, settings, {"--line-width"}) + + assert args.line_width == 2.0 + assert args.subsample == 4 + assert (args.crop_x, args.crop_y, args.crop_z) == ([1, 2], [3, 4], [5, 6]) + assert (args.width, args.height) == (1600, 1000) + assert args.hide_axes is True + assert args.show_bounds is True + assert args.shadows is True + + _apply_session_settings(args, settings, set()) + assert args.line_width == 0.25 + + +def test_pyvista_subsampling_is_reproducible(): + streamlines = [ + np.array([[index, 0, 0], [index, 1, 0]], dtype=np.float32) + for index in range(20) + ] + colors = [np.array([0, 1], dtype=np.float32) for _ in streamlines] + + first, _, _ = pyvista_plotting_streamlines._prepare_streamlines( + streamlines, + colors, + downsample_factor=1, + max_streamlines=None, + filter_min_len=None, + subsample_factor=4, + crop_bounds=None, + random_seed=1234, + ) + second, _, _ = pyvista_plotting_streamlines._prepare_streamlines( + streamlines, + colors, + downsample_factor=1, + max_streamlines=None, + filter_min_len=None, + subsample_factor=4, + crop_bounds=None, + random_seed=1234, + ) + + assert [line[0, 0] for line in first] == [line[0, 0] for line in second] + + +def test_ctrl_s_saves_complete_fury_session(tmp_path): + class FakeCamera: + def GetPosition(self): + return (10, 20, 30) + + def GetFocalPoint(self): + return (1, 2, 3) + + def GetViewUp(self): + return (0, 0, 1) + + def GetClippingRange(self): + return (0.1, 5000) + + def GetViewAngle(self): + return 30 + + def GetParallelProjection(self): + return False + + def GetParallelScale(self): + return 12 + + class FakePlane: + def GetOrigin(self, output): + output[:] = [4, 5, 6] + + def GetNormal(self, output): + output[:] = [0, 1, 0] + + trk_path = tmp_path / "streamlines.trk" + trk_path.write_bytes(b"trk") + session_path = tmp_path / "my_view.json" + viewer = fury_plotting_streamlines.StreamlineViewer.__new__( + fury_plotting_streamlines.StreamlineViewer + ) + viewer.session_path = session_path + viewer.session_settings = {"subsample_factor": 4, "random_seed": 1234} + viewer.streamlines_file = trk_path + viewer.scene = SimpleNamespace(GetActiveCamera=lambda: FakeCamera()) + viewer.plane_rep = FakePlane() + viewer.plane_widget = SimpleNamespace(GetEnabled=lambda: True) + viewer.showm = SimpleNamespace(window=SimpleNamespace(GetSize=lambda: (1800, 1200))) + viewer.window_size = (800, 800) + viewer.linewidth = 2.5 + viewer.current_bg = (0.0, 0.0, 0.0) + viewer.clipping_active = True + viewer.scale_bar_on = False + viewer.material = { + "ambient": 0.4, + "diffuse": 0.7, + "specular": 0.2, + "opacity": 0.9, + } + viewer.quality = "publication" + viewer.box_rep = None + viewer.box_widget = None + viewer.box_clipping_active = False + viewer.controls_visible = False + + class ControlS: + def GetKeySym(self): + return "s" + + def GetControlKey(self): + return 1 + + viewer._on_keypress(ControlS(), None) + saved = json.loads(session_path.read_text()) + + assert saved["format"] == "cardiotensor-fury-session" + assert saved["streamlines_size"] == 3 + assert saved["settings"]["random_seed"] == 1234 + assert saved["settings"]["line_width"] == 2.5 + assert saved["settings"]["tube_thickness"] == 2.5 + assert saved["view"]["camera"]["position"] == [10, 20, 30] + assert saved["view"]["clipping_plane"] == { + "origin": [4, 5, 6], + "normal": [0, 1, 0], + "enabled": True, + "gizmo_visible": True, + } + assert saved["view"]["window_size"] == [1800, 1200] + + +def test_fury_session_restores_camera_plane_and_controls(tmp_path): + calls = {} + session_path = tmp_path / "my_view.json" + session_path.write_text( + json.dumps( + { + "view": { + "camera": { + "position": [10, 20, 30], + "focal_point": [1, 2, 3], + "view_up": [0, 0, 1], + "clipping_range": [0.1, 5000], + "view_angle": 25, + "parallel_projection": True, + "parallel_scale": 12, + }, + "clipping_plane": { + "origin": [4, 5, 6], + "normal": [0, 1, 0], + "enabled": True, + "gizmo_visible": True, + }, + "background_color": [1, 1, 1], + "scale_bar_visible": False, + } + } + ) + ) + + class FakeCamera: + def __getattr__(self, name): + if name.startswith("Set"): + return lambda *values: calls.__setitem__(name, values) + raise AttributeError(name) + + class FakePlaneRepresentation: + origin = [0, 0, 0] + normal = [1, 0, 0] + + def SetOrigin(self, *values): + self.origin = list(values) + + def SetNormal(self, *values): + self.normal = list(values) + + def GetOrigin(self, output): + output[:] = self.origin + + def GetNormal(self, output): + output[:] = self.normal + + def UpdatePlacement(self): + calls["plane_updated"] = True + + class FakePlane: + def SetOrigin(self, *values): + calls["plane_origin"] = values + + def SetNormal(self, *values): + calls["plane_normal"] = values + + class FakeMapper: + def RemoveAllClippingPlanes(self): + calls["planes_removed"] = calls.get("planes_removed", 0) + 1 + + def AddClippingPlane(self, plane): + calls["planes_added"] = calls.get("planes_added", 0) + 1 + + class FakeWidget: + def EnabledOn(self): + calls["gizmo"] = True + + def EnabledOff(self): + calls["gizmo"] = False + + mapper_a = FakeMapper() + mapper_b = FakeMapper() + camera = FakeCamera() + scene = SimpleNamespace( + GetActiveCamera=lambda: camera, + SetBackground=lambda *values: calls.__setitem__("background", values), + add=lambda actor: calls.__setitem__("scale_bar_added", actor), + rm=lambda actor: calls.__setitem__("scale_bar_removed", actor), + ) + viewer = fury_plotting_streamlines.StreamlineViewer.__new__( + fury_plotting_streamlines.StreamlineViewer + ) + viewer.session_path = session_path + viewer.scene = scene + viewer.plane_rep = FakePlaneRepresentation() + viewer.plane_fn = FakePlane() + viewer.actor0 = SimpleNamespace(GetMapper=lambda: mapper_a) + viewer.actor_fast = SimpleNamespace(GetMapper=lambda: mapper_b) + viewer.plane_widget = FakeWidget() + viewer.box_rep = None + viewer.box_widget = None + viewer.box_clipping_active = False + viewer.clipping_active = False + viewer.current_bg = (0, 0, 0) + viewer.scale_bar = object() + viewer.scale_bar_on = True + viewer.showm = SimpleNamespace(render=lambda: calls.__setitem__("rendered", True)) + + assert viewer._restore_session() is True + assert calls["SetPosition"] == (10, 20, 30) + assert calls["SetParallelProjection"] == (True,) + assert calls["plane_origin"] == ([4, 5, 6],) + assert calls["plane_normal"] == ([0, 1, 0],) + assert calls["planes_added"] == 2 + assert calls["gizmo"] is True + assert calls["background"] == (1.0, 1.0, 1.0) + assert calls["scale_bar_removed"] is viewer.scale_bar + assert calls["rendered"] is True