Add --via-template option for multi-step template registration

- Add --via-template CLI option to snakebids.yml
- Add via_template variable and validation to Snakefile
- Add get_composite_warp, get_composite_warp_inv, get_affine_warped_for_qc helper functions to common.smk
- Create viatemplatereg.smk with template-to-template registration rules
- Update downstream rules in templatereg.smk, heatmaps.smk, regionprops.smk, preproc_mri.smk to use helper functions

Agent-Logs-Url: https://github.com/khanlab/SPIMquant/sessions/9792c81e-e8e2-4bbf-805c-cfc28d084f2f

Co-authored-by: akhanf <11492701+akhanf@users.noreply.github.com>

Author: Copilot

spimquant/config/snakebids.yml

@@ -63,6 +63,16 @@ parse_args:
      - MBMv3
      - turone
 
+  --via-template:
+    help: "Intermediate template for multi-step registration. Subject is first registered to this template, then this template is registered to the final --template. Transforms are concatenated to produce a single subject-to-template composite warp. (default: %(default)s)"
+    default: null
+    choices:
+      - ABAv3
+      - DSURQE
+      - gubra
+      - MBMv3
+      - turone
+
   --atlas_segs:
     help: "Atlas segmentations to use with the chosen template (default: predefined list or all)"
     default: null

spimquant/workflow/Snakefile

@@ -57,6 +57,19 @@ stain_for_reg = None
 
 template = config["templates"][config["template"]]
 
+via_template = config.get("via_template", None)
+
+if via_template is not None:
+    if via_template == config["template"]:
+        raise ValueError(
+            f"--via-template ({via_template}) must be different from --template ({config['template']})"
+        )
+    if via_template not in config["templates"]:
+        raise ValueError(
+            f"--via-template ({via_template}) not found in configured templates: "
+            f"{list(config['templates'].keys())}"
+        )
+
 
 # first, check if there are any SPIM templates defined
 # for the stains we have
@@ -923,6 +936,11 @@ include: "rules/groupstats.smk"
 include: "rules/qc.smk"
 
 
+if via_template is not None:
+
+    include: "rules/viatemplatereg.smk"
+
+
 if config["register_to_mri"]:
 
     include: "rules/preproc_mri.smk"

spimquant/workflow/rules/common.smk

@@ -86,3 +86,80 @@ def get_regionprops_parquet(wildcards):
             suffix="regionprops.parquet",
             **inputs["spim"].wildcards,
         )
+
+
+def get_composite_warp(wildcards):
+    """Get the composite warp from subject to template space.
+
+    Returns the path to either the direct composite warp (no via-template)
+    or the combined via-template composite warp if --via-template is configured.
+    """
+    spim_wildcards = {k: getattr(wildcards, k) for k in inputs["spim"].wildcards.keys()}
+    if via_template is None:
+        return bids(
+            root=root,
+            datatype="xfm",
+            from_="subject",
+            to=wildcards.template,
+            suffix="xfm.nii.gz",
+            **spim_wildcards,
+        )
+    else:
+        return bids(
+            root=root,
+            datatype="xfm",
+            from_="subject",
+            to=wildcards.template,
+            via=via_template,
+            suffix="xfm.nii.gz",
+            **spim_wildcards,
+        )
+
+
+def get_composite_warp_inv(wildcards):
+    """Get the inverse composite warp from template to subject space.
+
+    Returns the path to either the direct inverse composite warp (no via-template)
+    or the combined via-template inverse composite warp if --via-template is configured.
+    """
+    spim_wildcards = {k: getattr(wildcards, k) for k in inputs["spim"].wildcards.keys()}
+    if via_template is None:
+        return bids(
+            root=root,
+            datatype="xfm",
+            from_=wildcards.template,
+            to="subject",
+            suffix="xfm.nii.gz",
+            **spim_wildcards,
+        )
+    else:
+        return bids(
+            root=root,
+            datatype="xfm",
+            from_=wildcards.template,
+            to="subject",
+            via=via_template,
+            suffix="xfm.nii.gz",
+            **spim_wildcards,
+        )
+
+
+def get_affine_warped_for_qc(wildcards):
+    """Get the affine-warped SPIM image for QC report.
+
+    Returns the affine-warped image in the appropriate space:
+    - Without via-template: image is warped to the final template space
+    - With via-template: image is warped to the intermediate (via) template space,
+      since that is where the affine registration is performed.
+    """
+    spim_wildcards = {k: getattr(wildcards, k) for k in inputs["spim"].wildcards.keys()}
+    target_space = via_template if via_template is not None else wildcards.template
+    return bids(
+        root=root,
+        datatype="xfm",
+        space=target_space,
+        stain=wildcards.stain,
+        desc="affinewarped",
+        suffix="SPIM.nii.gz",
+        **spim_wildcards,
+    )

spimquant/workflow/rules/heatmaps.smk

@@ -90,7 +90,7 @@ rule deform_fieldfrac_nii_to_template_nii:
             **inputs["spim"].wildcards,
         ),
         ref=rules.import_template_anat.output.anat,
-        xfm_composite=rules.compose_subject_to_template_warp.output.xfm_composite,
+        xfm_composite=get_composite_warp,
     output:
         nii=bids(
             root=root,

spimquant/workflow/rules/preproc_mri.smk

@@ -709,7 +709,7 @@ rule warp_mri_to_template_via_spim:
         ),
         ref=rules.import_template_anat.output.anat,
         xfm_composite_mri_to_spim=rules.compose_mri_to_spim_warp.output.xfm_composite,
-        xfm_composite_spim_to_template=rules.compose_subject_to_template_warp.output.xfm_composite,
+        xfm_composite_spim_to_template=get_composite_warp,
     output:
         warped=bids(
             root=root,

spimquant/workflow/rules/regionprops.smk

@@ -46,7 +46,7 @@ rule transform_regionprops_to_template:
             suffix="regionprops.parquet",
             **inputs["spim"].wildcards,
         ),
-        xfm_composite_inv=rules.compose_subject_to_template_warp.output.xfm_composite_inv,
+        xfm_composite_inv=get_composite_warp_inv,
     params:
         coord_column_names=config["coord_column_names"],
     output:

spimquant/workflow/rules/templatereg.smk

@@ -471,7 +471,7 @@ rule affine_zarr_to_template_ome_zarr:
 
 rule deform_zarr_to_template_nii:
     input:
-        xfm_composite=rules.compose_subject_to_template_warp.output.xfm_composite,
+        xfm_composite=get_composite_warp,
         ref_nii=get_template_for_reg,
     params:
         ome_zarr=inputs["spim"].path,
@@ -500,7 +500,7 @@ rule deform_zarr_to_template_nii:
 rule deform_to_template_nii_zoomed:
     input:
         ome_zarr=inputs["spim"].path,
-        xfm_composite=rules.compose_subject_to_template_warp.output.xfm_composite,
+        xfm_composite=get_composite_warp,
         ref_nii=get_template_for_reg,
     params:
         flo_opts={},  #any additional flo znimg options
@@ -544,7 +544,7 @@ rule deform_spim_nii_to_template_nii:
             **inputs["spim"].wildcards,
         ),
         ref=rules.import_template_anat.output.anat,
-        xfm_composite=rules.compose_subject_to_template_warp.output.xfm_composite,
+        xfm_composite=get_composite_warp,
     output:
         spim=bids(
             root=root,
@@ -585,7 +585,7 @@ rule deform_template_dseg_to_subject_nii:
             **inputs["spim"].wildcards,
         ),
         dseg=bids(root=root, template="{template}", seg="{seg}", suffix="dseg.nii.gz"),
-        xfm_composite_inv=rules.compose_subject_to_template_warp.output.xfm_composite_inv,
+        xfm_composite_inv=get_composite_warp_inv,
     output:
         dseg=bids(
             root=root,
@@ -679,15 +679,7 @@ rule registration_qc_report:
             suffix="SPIM.nii.gz",
             **inputs["spim"].wildcards,
         ),
-        warped_affine=bids(
-            root=root,
-            datatype="xfm",
-            space="{template}",
-            stain="{stain}",
-            desc="affinewarped",
-            suffix="SPIM.nii.gz",
-            **inputs["spim"].wildcards,
-        ),
+        warped_affine=get_affine_warped_for_qc,
         warped_deform=bids(
             root=root,
             datatype="micr",
@@ -697,7 +689,7 @@ rule registration_qc_report:
             suffix="SPIM.nii.gz",
             **inputs["spim"].wildcards,
         ),
-        warp=rules.compose_subject_to_template_warp.output.xfm_composite,
+        warp=get_composite_warp,
         dseg=lambda wildcards: bids(
             root=root,
             template=wildcards.template,