Add pycurv-gpu dispatcher in curvature.run_pycurv.

README.md

@@ -16,6 +16,7 @@ Everything is driven by a single `morphometrics` command plus a `config.yml` fil
 - [Quick start](#quick-start)
 - [Example data](#example-data)
 - [The pipeline](#the-pipeline)
+- [GPU curvature (pycurv-gpu)](#gpu-curvature-pycurv-gpu)
 - [Analysis & visualization](#analysis--visualization)
 - [Reference](#reference)
 - [Troubleshooting](#troubleshooting)
@@ -32,7 +33,8 @@ The fastest, easiest starting point for most Linux boxes, and now for Mac as wel
 1. Clone the repository: `git clone https://github.com/baradlab/surface_morphometrics.git`
 2. Create the environment (this also installs the toolkit and the `morphometrics` command via `pip install -e .`): `conda env create -f environment.yml`
 3. Activate it: `conda activate morphometrics`
-4. Check the install: `morphometrics --help` should list the pipeline subcommands.
+4. *(GPU machines)* Install a CUDA-enabled PyTorch build into the same env — see [pytorch.org](https://pytorch.org/get-started/locally/). `pycurv-gpu` is already installed by `environment.yml`.
+5. Check the install: `morphometrics --help` should list the pipeline subcommands.
 
 > Older Ubuntu installs (and some other Linux distributions) have known issues with graph-tool. If the main environment file fails, try `conda env create -f environment-ubuntu.yml` instead (tested on Ubuntu 22.04 LTS).
 
@@ -106,14 +108,32 @@ Each step reads a `config.yml` and writes its outputs into the configured `work_
 | 0 | Create a config | `morphometrics new_config` |
 | – | *(optional)* [Validate the setup](#checking-your-setup-and-progress) | `morphometrics validate config.yml` |
 | 1 | Segmentations → meshes | `morphometrics make_meshes config.yml` |
-| 2 | Curvature (pycurv) | `morphometrics pycurv config.yml` |
+| 2 | Curvature (pycurv) | `morphometrics pycurv config.yml` ([GPU](#gpu-curvature-pycurv-gpu) optional) |
 | 3 | *(optional)* [Mesh refinement](#mesh-refinement-optional) | `morphometrics refine_mesh config.yml` → `accept_refinement` |
 | 4 | Distances & orientations | `morphometrics distances_orientations config.yml` |
 | 5 | [Thickness](#data-organization-for-thickness-and-refinement) (needs tomograms) | `morphometrics sample_density config.yml` → `measure_thickness` |
 | 6 | Aggregate statistics | `morphometrics stats config.yml <name>` |
 
 Most steps also accept a single input (e.g. `morphometrics pycurv config.yml TE1_OMM.surface.vtp`, `morphometrics distances_orientations config.yml TE1.mrc`) so you can parallelize per tomogram on a cluster.
 
+### GPU curvature (pycurv-gpu)
+Step 2 can use [pycurv-gpu](https://github.com/hemanthkapa/pycurv-gpu) instead of CPU pycurv. Outputs (`.vtp` / `.csv` / `.gt`) keep the same names and properties, so distances, thickness, refinement, and the GUI work unchanged.
+
+Enable it either way:
+```yaml
+# in config.yml under curvature_measurements
+use_gpu: true
+# gpu_device: null        # optional: cuda (null = auto; never falls back to CPU)
+# gpu_batch_size: 1024    # optional tuning
+```
+```bash
+morphometrics pycurv config.yml --gpu
+morphometrics pycurv config.yml TE1_OMM.surface.vtp --gpu
+```
+`--gpu` overrides `use_gpu: false`. `refine_mesh` also honors `use_gpu` from config for its final pycurv pass.
+
+Requirements: CUDA-capable GPU + a CUDA PyTorch build in the morphometrics env (`pycurv-gpu` itself is installed by `environment.yml`). Without CUDA, leave `use_gpu: false` (default) and use CPU pycurv.
+
 ### Configuration
 `morphometrics new_config` writes a fully-commented `config.yml` into the current directory (`-o NAME` for a different name); edit it for your project. For a stripped-down starting point, `morphometrics new_config --simple` writes a minimal config with just the most commonly-adjusted settings — directories, `segmentation_values`, cores, the main meshing options, and the thickness/refinement averaging radii; everything omitted falls back to documented defaults, so a partial config still runs. A few starting tips:
 - For higher-quality meshes with near-equilateral triangles, set `isotropic_remesh: true` and `simplify: false` in `surface_generation`. `target_area` of 1.0–3.0 nm² works well (smaller = finer but slower).
@@ -249,6 +269,7 @@ The toolkit is the `surface_morphometrics` Python package; the pipeline steps ar
 - `Gaussian or Mean curvature of X has a large computation error` — safe to ignore (pycurv cleans these up); they are suppressed by default.
 - MRC files from AMIRA (and some other software, e.g. Dragonfly) lack proper machine stamps; open with `mrcfile.open(filename, permissive=True)`.
 - If pycurv seems to hang indefinitely, try setting `cores: 1` in the config.
+- `GPU pycurv requires PyTorch` / `CUDA is not available` — install a CUDA torch build into the morphometrics env, or leave `use_gpu: false`.
 
 ---
 
@@ -276,7 +297,7 @@ Two config keys were also renamed: `data_dir` → `seg_dir`, and `max_triangles`
 ---
 
 ## Dependencies
-Numpy, Scipy, Pandas, mrcfile, Click, Matplotlib, starfile, Pymeshlab, and [PyCurv](https://github.com/kalemaria/pycurv) (which pulls in Pyto and graph-tool). The conda environment files install everything; see [Installation](#installation).
+Numpy, Scipy, Pandas, mrcfile, Click, Matplotlib, starfile, Pymeshlab, [PyCurv](https://github.com/kalemaria/pycurv) (which pulls in Pyto and graph-tool), and [pycurv-gpu](https://github.com/hemanthkapa/pycurv-gpu) for optional GPU curvature. The conda environment files install everything except a CUDA PyTorch build (install that separately on GPU machines); see [Installation](#installation).
 
 ---
 

environment-ubuntu.yml

@@ -29,6 +29,7 @@ dependencies:
     - pymeshlab>=2023.12
     - git+https://github.com/vladanl/Pyto.git
     - git+https://github.com/bbarad/pycurv.git
+    - git+https://github.com/hemanthkapa/pycurv-gpu.git
     - git+https://github.com/hemanthkapa/napari-vedo-bridge.git
     - napari-tomoslice
     - napari-timelapse-processor

environment.yml

@@ -26,8 +26,9 @@ dependencies:
   - pip:
     - git+https://github.com/vladanl/Pyto.git
     - git+https://github.com/bbarad/pycurv.git
+    - git+https://github.com/hemanthkapa/pycurv-gpu.git
     - git+https://github.com/hemanthkapa/napari-vedo-bridge.git
     - napari-tomoslice
     - napari-timelapse-processor
     # Install this toolkit and the `morphometrics` command (pulls remaining pip deps, e.g. starfile)
-    - -e .
+    - -e .

surface_morphometrics/config_template.yml

@@ -28,6 +28,9 @@ curvature_measurements:
   radius_hit: 9 # This corresponds to the radius of the smallest feature of interest, roughly, for neighborhood determination. 8-15 seems to be a good range for mitochondria
   min_component: 30 # The minimum number of triangles for a component to be considered for curvature measurement.
   exclude_borders: 1 # Values greater than zero exclude the border from the curvature calculation by n nm/angstroms.
+  use_gpu: false # Set true to use pycurv-gpu (CUDA) instead of CPU pycurv.
+  gpu_device: null # Optional torch device: cuda. null = auto-select.
+  gpu_batch_size: 1024 # SSSP/voting batch size for pycurv-gpu.
 distance_and_orientation_measurements:
   mindist: 3 # Minimum distance between two points for them to be considered for distance measurement.
   maxdist: 400 # Maximum distance between two points for them to be considered for distance measurement.

surface_morphometrics/config_utils.py

@@ -40,6 +40,9 @@
         "radius_hit": 9,
         "min_component": 30,
         "exclude_borders": 1,
+        "use_gpu": False,
+        "gpu_device": None,
+        "gpu_batch_size": 1024,
     },
     "distance_and_orientation_measurements": {
         "mindist": 3,

surface_morphometrics/curvature.py

@@ -14,6 +14,96 @@
 
 from .curvature_calculation import new_workflow, extract_curvatures_after_new_workflow
 
+
+def _require_gpu_backend(device=None):
+    """Fail fast if pycurv-gpu / torch / a CUDA device are unavailable."""
+    try:
+        import torch
+    except ImportError as e:
+        raise RuntimeError(
+            "GPU pycurv requires PyTorch. Install a CUDA torch build "
+            "(https://pytorch.org/get-started/locally/) into the morphometrics env."
+        ) from e
+    try:
+        from core.api import run_pipeline  # noqa: F401
+    except ImportError as e:
+        raise RuntimeError(
+            "GPU pycurv requires pycurv-gpu. Reinstall the morphometrics env "
+            "(environment.yml includes pycurv-gpu) or: "
+            "pip install git+https://github.com/hemanthkapa/pycurv-gpu.git"
+        ) from e
+
+    cuda_ok = torch.cuda.is_available()
+
+    if device is None:
+        if cuda_ok:
+            return 'cuda'
+        raise RuntimeError(
+            "use_gpu was requested but CUDA is not available. "
+            "Install a CUDA PyTorch build on a GPU machine, or set use_gpu: false "
+            "to use CPU pycurv. GPU mode will not fall back to CPU torch."
+        )
+
+    resolved = str(device).lower()
+    if resolved in ('cpu', 'cpu:0'):
+        raise RuntimeError(
+            "gpu_device='cpu' is not supported. use_gpu requires CUDA; "
+            "set use_gpu: false to run CPU pycurv instead."
+        )
+    if resolved == 'mps':
+        raise RuntimeError(
+            "gpu_device='mps' is not supported. use_gpu requires CUDA; "
+            "set use_gpu: false to run CPU pycurv instead."
+        )
+    if resolved.startswith('cuda') and not cuda_ok:
+        raise RuntimeError(
+            "use_gpu was requested with gpu_device={!r} but CUDA is not available. "
+            "Install a CUDA PyTorch build, or set use_gpu: false for CPU pycurv."
+            .format(device)
+        )
+    return resolved
+
+
+def run_pycurv_gpu(filename, folder, scale=1, radius_hit=10, min_component=30,
+                   exclude_borders=0, cores=16, remove_wrong_borders=False,
+                   device=None, batch_size=1024):
+    """Run pycurv-gpu on a vtp surface file (drop-in for run_pycurv)."""
+    assert filename.endswith(".surface.vtp"), "Surface must be a vtp file ending with .surface.vtp"
+    basename = filename[:-len(".surface.vtp")]
+    device = _require_gpu_backend(device)
+
+    from core.api import run_pipeline
+
+    t_begin = time.time()
+    print("\nCalculating curvatures (GPU) for {}".format(basename))
+    run_pipeline(
+        os.path.join(folder, filename),
+        output_dir=folder,
+        radius_hit=radius_hit,
+        pixel_size=scale,
+        min_component=min_component,
+        exclude_borders=exclude_borders,
+        remove_wrong_borders=remove_wrong_borders,
+        device=device,
+        batch_size=batch_size,
+        write_gt=True,
+        write_vtp=True,
+        cores=cores,
+    )
+
+    t_end = time.time()
+    duration = t_end - t_begin
+    minutes, seconds = divmod(duration, 60)
+    if not folder.endswith("/"):
+        folder += "/"
+    rh_str = str(int(radius_hit)) if radius_hit == int(radius_hit) else str(radius_hit)
+    print('\nTotal GPU pycurv time: {} min {} s'.format(int(minutes), seconds))
+    print("Final outputs written:")
+    print("VTP file for paraview: " + folder + basename + f'.AVV_rh{rh_str}.vtp')
+    print("CSV file for pandas based quantification: " + folder + basename + f'.AVV_rh{rh_str}.csv')
+    print("GT file for further morphometrics quantification: " + folder + basename + f'.AVV_rh{rh_str}.gt')
+
+
 @click.command()
 @click.argument('filename')
 @click.argument('folder')
@@ -34,7 +124,9 @@ def run_pycurv_cli(filename, folder, scale, radius_hit, min_component, exclude_b
                 raise SystemExit(1)
     run_pycurv(filename, folder, scale, radius_hit, min_component, exclude_borders, cores, remove_wrong_borders)
 
-def run_pycurv(filename, folder, scale=1, radius_hit=10, min_component=30, exclude_borders=0, cores=16, remove_wrong_borders=False):
+def run_pycurv(filename, folder, scale=1, radius_hit=10, min_component=30, exclude_borders=0,
+               cores=16, remove_wrong_borders=False, use_gpu=False, gpu_device=None,
+               gpu_batch_size=1024):
     """Run pycurv on a vtp surface file and extract curvatures
 
     filename (str): vtp surface file
@@ -45,7 +137,18 @@ def run_pycurv(filename, folder, scale=1, radius_hit=10, min_component=30, exclu
     exclude_borders (int): distance in surface units (angstroms or nm) to exclude from the curvature calculation. Default is 0.
     cores (int): number of cores to use for the calculation. Default is 16.
     remove_wrong_borders (bool): eat back the surface before calculations. If using screened poisson workflow leave this False.
+    use_gpu (bool): if True, dispatch to pycurv-gpu instead of CPU pycurv.
+    gpu_device (str or None): optional torch device ('cuda', 'cpu', 'mps').
+    gpu_batch_size (int): SSSP/voting batch size for pycurv-gpu.
     """
+    if use_gpu:
+        return run_pycurv_gpu(
+            filename, folder, scale=scale, radius_hit=radius_hit,
+            min_component=min_component, exclude_borders=exclude_borders, cores=cores,
+            remove_wrong_borders=remove_wrong_borders, device=gpu_device,
+            batch_size=gpu_batch_size,
+        )
+
     assert filename.endswith(".surface.vtp"), "Surface must be a vtp file of a surface, ending with .surface.vtp"
     basename = filename[:-len(".surface.vtp")]
     runtimes_file = "{}{}_runtimes.csv".format(folder, basename)