Refactor to resample template per task

Templates for different takss _could_ be different resolutions. This refactors the resampling such that a template is generated per task for the longitudinal workflow.

Author: kaitj

src/rbc/bids/longitudinal/functional.py

@@ -50,8 +50,9 @@ def resolve_longitudinal_func(
             without=["space"],
         )
 
+    task = func_df["task"].unique()[0]
     return {
-        "template": tpl_q.expect(tpl_df, res="bold", suffix="T1w"),
+        "template": tpl_q.expect(tpl_df, suffix="T1w", res=task),
         "anat_to_template_xfm": tpl_q.expect(
             tpl_df,
             suffix="xfm",

src/rbc/bids/longitudinal/template.py

@@ -21,13 +21,13 @@ class TemplateInputs(NamedTuple):
         sub: Subject label.
         sessions: Per-input session labels (parallel to ``files``).
         files: Per-session preprocessed T1w brain volumes.
-        bold_ref: First BOLD volume for grid reference.
+        bold_ref: Per-session preprocessed BOLD volumes (all tasks).
     """
 
     sub: str
     sessions: list[str]
     files: list[Path]
-    bold_ref: Path | None
+    bold_files: dict[str, Path]
 
 
 def discover_template_inputs(
@@ -58,17 +58,13 @@ def discover_template_inputs(
         # the mri_robust_template invocation.
         pl.col("space").is_null(),
     )
-    bold_ref_rows = df.filter(
+    bold_rows = df.filter(
         pl.col("ses") != "longitudinal",
         pl.col("datatype") == "func",
+        pl.col("desc") == "preproc",
         pl.col("suffix") == "bold",
         pl.col("space").is_null(),
     )
-    bold_ref = (
-        None
-        if bold_ref_rows.is_empty()
-        else Path(bold_ref_rows["root"][0]) / bold_ref_rows["path"][0]
-    )
 
     inputs: list[TemplateInputs] = []
     skipped: list[str] = []
@@ -81,8 +77,18 @@ def discover_template_inputs(
         files = [
             Path(row["root"]) / row["path"] for row in sub_group.iter_rows(named=True)
         ]
+        # Filter for first found session; only single reference per task is needed
+        sub_bold = bold_rows.filter(
+            (pl.col("sub") == sub) & (pl.col("ses") == sessions[0])
+        )
+        bold_files = {
+            row["task"]: Path(row["root"]) / row["path"]
+            for row in sub_bold.iter_rows(named=True)
+        }
         inputs.append(
-            TemplateInputs(sub=sub, sessions=sessions, files=files, bold_ref=bold_ref)
+            TemplateInputs(
+                sub=sub, sessions=sessions, files=files, bold_files=bold_files
+            )
         )
     return inputs, skipped
 
@@ -96,8 +102,8 @@ def export_template(tpl: Bids, outputs: LongitudinalTemplateOutputs) -> None:
         outputs: Results from the longitudinal template workflow.
     """
     tpl.save(outputs.template, suffix=Suffix.T1W)
-    if outputs.bold_template is not None:
-        tpl.save(outputs.bold_template, res="bold", suffix=Suffix.T1W)
+    for btask, bold_template in outputs.bold_templates.items():
+        tpl.save(bold_template, res=btask, suffix=Suffix.T1W)
     for ses, xfm in zip(outputs.sessions, outputs.transforms, strict=True):
         ses_label = bids_safe_label(ses)
         tpl.save(

src/rbc/orchestration/longitudinal/template.py

@@ -49,7 +49,7 @@ def process_subject(
         sub=inputs.sub,
         sessions=inputs.sessions,
         in_files=inputs.files,
-        bold_ref=inputs.bold_ref,
+        bold_files=inputs.bold_files,
     )
     tpl = pipe_ctx.bids(datatype=Datatype.ANAT).derive(ses="longitudinal")
     export_template(tpl, outputs)

src/rbc/workflows/longitudinal/template.py

@@ -17,7 +17,7 @@
 from rbc.core.longitudinal.resampling import resample_img_to_bold_grid
 
 if TYPE_CHECKING:
-    from collections.abc import Sequence
+    from collections.abc import Mapping, Sequence
     from pathlib import Path
 
 _logger = logging.getLogger("rbc")
@@ -34,7 +34,7 @@ class LongitudinalTemplateOutputs(NamedTuple):
     """
 
     template: Path
-    bold_template: Path | None
+    bold_templates: dict[str, Path]
     sessions: list[str]
     transforms: list[Path]
 
@@ -43,15 +43,15 @@ def generate_subject_template(
     sub: str,
     sessions: Sequence[str],
     in_files: Sequence[Path],
-    bold_ref: Path | None = None,
+    bold_files: Mapping[str, Path],
 ) -> LongitudinalTemplateOutputs:
     """Build a robust template and ITK transforms for one subject.
 
     Args:
         sub: Subject label (without the ``sub-`` prefix).
         sessions: Session labels parallel to ``in_files``.
         in_files: Per-session preprocessed T1w volumes (e.g. brain-extracted).
-        bold_ref: Reference bold volume to resample template for functional data.
+        bold_files: Reference bold volumes to resample template for functional data.
 
     Returns:
         :class:`LongitudinalTemplateOutputs` ready for BIDS export.
@@ -70,13 +70,15 @@ def generate_subject_template(
         in_xfms=robust.transforms,
     )
 
-    bold_template = (
-        resample_img_to_bold_grid(bold_ref, robust.template) if bold_ref else None
-    )
+    _logger.info("Creating reference volumes for each functional task")
+    bold_templates = {
+        btask: resample_img_to_bold_grid(bfile, robust.template)
+        for btask, bfile in bold_files.items()
+    }
 
     return LongitudinalTemplateOutputs(
         template=robust.template,
-        bold_template=bold_template,
+        bold_templates=bold_templates,
         sessions=list(sessions),
         transforms=itk_xfms,
     )

tests/unit/bids/test_longitudinal_template.py

@@ -155,7 +155,7 @@ def test_writes_template_and_xfms(self, tmp_path: Path) -> None:
 
         outputs = LongitudinalTemplateOutputs(
             template=template_src,
-            bold_template=template_bold_src,
+            bold_templates={"test": template_bold_src},
             sessions=["baseline", "vis2"],
             transforms=[xfm_baseline, xfm_vis2],
         )