Address reviews, round 2

- Removed check for uniqueness of bold dataframe
- Update stale docstring: `bold_ref` -> `bold_files`
- Wrap task with `bids_safe_label` when calling `tpl.save`
- Compare affines using `np.allclose` due to floating point precision
- Use a BoldKey NamedTuple as key to unique bold files

Author: kaitj

src/rbc/bids/longitudinal/template.py

@@ -21,13 +21,36 @@ class TemplateInputs(NamedTuple):
         sub: Subject label.
         sessions: Per-input session labels (parallel to ``files``).
         files: Per-session preprocessed T1w brain volumes.
-        bold_ref: Per-session preprocessed BOLD volumes (all tasks).
+        bold_files: Per-session, unique preprocessed BOLD volumes (all tasks).
     """
 
     sub: str
     sessions: list[str]
     files: list[Path]
-    bold_files: dict[str, Path]
+    bold_files: dict[BoldKey, Path]
+
+
+# Use typing library instead of collections for static type hints
+class BoldKey(NamedTuple):
+    """Key with associated entities to use as dict key to represent bold filepath found.
+
+    Attributes:
+        task: Task associated with file.
+        run: Run associated with file.
+        acq: Acquisition associated with file.
+        dir: Direction associated with file.
+        echo: Echo associated with file.
+        part: Part associated with file.
+        rec: Recording associated with file.
+    """
+
+    task: str
+    run: str | None = None
+    acq: str | None = None
+    dir: str | None = None
+    echo: str | None = None
+    part: str | None = None
+    rec: str | None = None
 
 
 def discover_template_inputs(
@@ -80,19 +103,9 @@ def discover_template_inputs(
         sub_bold = bold_rows.filter(
             (pl.col("sub") == sub) & (pl.col("ses") == sessions[0])
         ).unique(subset=(*FUNC_GROUP_ENTITIES, "root", "path"))
-        # Check each task is unique, otherwise raise assertion error with details
-        if sub_bold.height != sub_bold.unique().height:
-            conflicts = (
-                sub_bold.filter(pl.struct(FUNC_GROUP_ENTITIES).is_duplicated())
-                .group_by(FUNC_GROUP_ENTITIES)
-                .agg(pl.format("{}/{}", "root", "path").alias("paths"))
-            )
-            raise AssertionError(
-                f"Found multiple non-matching grids for subject {sub}:\n"
-                + "\n".join(str(dict(row)) for row in conflicts.iter_rows(named=True))
-            )
-        bold_files = {
-            row["task"]: Path(row["root"]) / row["path"]
+        bold_files: dict[BoldKey, Path] = {
+            BoldKey(**{ent: row[ent] for ent in FUNC_GROUP_ENTITIES}): Path(row["root"])
+            / row["path"]
             for row in sub_bold.iter_rows(named=True)
         }
         inputs.append(
@@ -112,8 +125,8 @@ def export_template(tpl: Bids, outputs: LongitudinalTemplateOutputs) -> None:
         outputs: Results from the longitudinal template workflow.
     """
     tpl.save(outputs.template, suffix=Suffix.T1W)
-    for btask, bold_template in outputs.bold_templates.items():
-        tpl.save(bold_template, res=btask, suffix=Suffix.T1W)
+    for bold_key, bold_template in outputs.bold_templates.items():
+        tpl.save(bold_template, res=bids_safe_label(bold_key.task), suffix=Suffix.T1W)
     for ses, xfm in zip(outputs.sessions, outputs.transforms, strict=True):
         ses_label = bids_safe_label(ses)
         tpl.save(

src/rbc/core/longitudinal/resampling.py

@@ -5,6 +5,7 @@
 from typing import TYPE_CHECKING
 
 import nibabel as nib
+import numpy as np
 from nibabel.processing import resample_from_to
 
 if TYPE_CHECKING:
@@ -31,7 +32,9 @@ def resample_img_to_bold_grid(bold_ref: Path, img: Path, order: int = 3) -> Path
     if len(bold_ref_img.shape) > 3:
         bold_ref_img = bold_ref_img.slicer[..., 0]
     # If same shape, no need to resample
-    if bold_ref_img.shape == img_obj.shape:
+    if bold_ref_img.shape == img_obj.shape and np.allclose(
+        bold_ref_img.affine, img_obj.affine
+    ):
         return img
 
     img_resampled = resample_from_to(img_obj, bold_ref_img, order=order)

src/rbc/orchestration/longitudinal/qc.py

@@ -20,7 +20,7 @@
     resolve_longitudinal_qc,
     write_longitudinal_qc_tsv,
 )
-from rbc.bids.session import iter_session_files
+from rbc.bids.session import _FUNC_ENTITY_KEYS, iter_session_files
 from rbc.core.longitudinal.resampling import resample_img_to_bold_grid
 from rbc.core.niwrap import generate_exec_folder
 from rbc.core.qc.registration import registration_qc_metrics
@@ -160,7 +160,7 @@ def run(
 
         for func_df, _ in iter_session_files(session, groupby=FUNC_GROUP_ENTITIES):
             row = func_df.filter(suffix=Suffix.BOLD).row(0, named=True)
-            ents = extract_entities(row, ["task", "run"])
+            ents = extract_entities(row, _FUNC_ENTITY_KEYS)
 
             func_q = pipe_ctx.bids(datatype=Datatype.FUNC, entities=ents)
             func_long_q = func_q.derive(space="longitudinal")

src/rbc/workflows/longitudinal/template.py

@@ -20,6 +20,8 @@
     from collections.abc import Mapping, Sequence
     from pathlib import Path
 
+    from rbc.bids.longitudinal.template import BoldKey
+
 _logger = logging.getLogger("rbc")
 
 
@@ -35,7 +37,7 @@ class LongitudinalTemplateOutputs(NamedTuple):
     """
 
     template: Path
-    bold_templates: dict[str, Path]
+    bold_templates: dict[BoldKey, Path]
     sessions: list[str]
     transforms: list[Path]
 
@@ -44,7 +46,7 @@ def generate_subject_template(
     sub: str,
     sessions: Sequence[str],
     in_files: Sequence[Path],
-    bold_files: Mapping[str, Path],
+    bold_files: Mapping[BoldKey, Path],
 ) -> LongitudinalTemplateOutputs:
     """Build a robust template and ITK transforms for one subject.
 

tests/unit/bids/test_longitudinal_template.py

@@ -7,6 +7,7 @@
 import polars as pl
 
 from rbc.bids.longitudinal.template import (
+    BoldKey,
     discover_template_inputs,
     export_template,
 )
@@ -182,7 +183,7 @@ def test_writes_template_and_xfms(self, tmp_path: Path) -> None:
 
         outputs = LongitudinalTemplateOutputs(
             template=template_src,
-            bold_templates={"test": template_bold_src},
+            bold_templates={BoldKey(task="test"): template_bold_src},
             sessions=["baseline", "vis2"],
             transforms=[xfm_baseline, xfm_vis2],
         )