Extract row entities to find refs + file naming

Author: kaitj

src/rbc/bids/longitudinal/template.py

@@ -7,7 +7,7 @@
 
 import polars as pl
 
-from rbc.bids import Suffix, bids_safe_label
+from rbc.bids import FUNC_GROUP_ENTITIES, Suffix, bids_safe_label
 
 if TYPE_CHECKING:
     from rbc.bids import Bids
@@ -77,23 +77,19 @@ 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])
-        ).unique(subset=("task", "root", "path"))
+        ).unique(subset=(*FUNC_GROUP_ENTITIES, "root", "path"))
         # Check each task is unique, otherwise raise assertion error with details
-        if sub_bold.height != sub_bold["task"].n_unique():
+        if sub_bold.height != sub_bold.unique().height:
             conflicts = (
-                sub_bold.filter(pl.col("task").is_duplicated())
-                .group_by("task")
+                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(
-                    f"Task '{row['task']}': {row['paths']}"
-                    for row in conflicts.iter_rows(named=True)
-                )
+                + "\n".join(str(dict(row)) for row in conflicts.iter_rows(named=True))
             )
         bold_files = {
             row["task"]: Path(row["root"]) / row["path"]

src/rbc/orchestration/longitudinal/all.py

@@ -15,10 +15,17 @@
 
 from tqdm import tqdm
 
-from rbc.bids import FUNC_GROUP_ENTITIES, Datatype, Suffix, extract_entities, load_table
+from rbc.bids import (
+    ANAT_GROUP_ENTITIES,
+    FUNC_GROUP_ENTITIES,
+    Datatype,
+    Suffix,
+    extract_entities,
+    load_table,
+)
 from rbc.bids.longitudinal.template import discover_template_inputs
 from rbc.bids.metrics import export_metrics
-from rbc.bids.session import iter_session_files
+from rbc.bids.session import _FUNC_ENTITY_KEYS, iter_session_files
 from rbc.context import RunContext
 from rbc.orchestration import Filters, RunnerConfig, init_runner
 from rbc.orchestration.longitudinal._iter import iter_sessions_with_template
@@ -109,7 +116,9 @@ def run(
         input_dirs, output_dir, filters=filters, verbose=verbose
     ):
         # Anatomical
-        for _, anat_df in session.anat.group_by(("run", "acq"), maintain_order=True):
+        for _, anat_df in session.anat.group_by(
+            ANAT_GROUP_ENTITIES, maintain_order=True
+        ):
             anat_outputs = process_anat(
                 pipe_ctx=pipe_ctx, anat_df=anat_df, tpl_df=tpl_df
             )
@@ -124,7 +133,7 @@ def run(
             )
 
             row = func_df.filter(suffix=Suffix.BOLD).row(0, named=True)
-            ents = extract_entities(row, ["task", "run"])
+            ents = extract_entities(row, list(_FUNC_ENTITY_KEYS))
             func_q = pipe_ctx.bids(datatype=Datatype.FUNC, entities=ents)
             func_long = func_q.derive(space="longitudinal")
 

src/rbc/orchestration/longitudinal/anatomical.py

@@ -51,7 +51,7 @@ def process_anat(
         Workflow outputs for in-memory handoff to downstream stages.
     """
     anat_df = anat_df.filter(pl.col("space").is_null())
-    ents = extract_entities(anat_df.row(0, named=True), ["run"])
+    ents = extract_entities(anat_df.row(0, named=True), ["run", "acq", "rec", "echo"])
 
     anat_q = pipe_ctx.bids(datatype=Datatype.ANAT)
     tpl_q = anat_q.derive(ses="longitudinal")

src/rbc/orchestration/longitudinal/functional.py

@@ -10,7 +10,7 @@
     export_longitudinal_func,
     resolve_longitudinal_func,
 )
-from rbc.bids.session import iter_session_files
+from rbc.bids.session import _FUNC_ENTITY_KEYS, iter_session_files
 from rbc.orchestration import Filters, RunnerConfig, init_runner
 from rbc.orchestration.longitudinal._iter import iter_sessions_with_template
 from rbc.workflows.longitudinal.functional import (
@@ -53,7 +53,7 @@ def process_func(
         Workflow outputs for in-memory handoff to downstream stages.
     """
     row = func_df.filter(suffix=Suffix.BOLD).row(0, named=True)
-    ents = extract_entities(row, ["task", "run"])
+    ents = extract_entities(row, list(_FUNC_ENTITY_KEYS))
 
     func_q = pipe_ctx.bids(datatype=Datatype.FUNC, entities=ents)
     tpl_q = pipe_ctx.bids(datatype=Datatype.ANAT).derive(ses="longitudinal")

tests/unit/bids/test_longitudinal_template.py

@@ -23,6 +23,11 @@
     "sub",
     "ses",
     "space",
+    "acq",
+    "dir",
+    "echo",
+    "part",
+    "rec",
     "task",
     "run",
     "desc",
@@ -50,6 +55,11 @@ def _anat_row(sub: str, ses: str, space: str | None = None) -> tuple:
         space,
         None,
         None,
+        None,
+        None,
+        None,
+        None,
+        None,
         "brain",
         "/data",
         path,
@@ -58,7 +68,24 @@ def _anat_row(sub: str, ses: str, space: str | None = None) -> tuple:
 
 def _func_row(sub: str, ses: str, task: str = "rest") -> tuple:
     path = f"sub-{sub}/ses-{ses}/func/sub-{sub}_ses-{ses}_task-{task}_bold.nii.gz"
-    return ("func", "bold", ".nii.gz", sub, ses, None, task, None, None, "/data", path)
+    return (
+        "func",
+        "bold",
+        ".nii.gz",
+        sub,
+        ses,
+        None,
+        None,
+        None,
+        None,
+        None,
+        None,
+        task,
+        None,
+        None,
+        "/data",
+        path,
+    )
 
 
 class TestDiscoverTemplateInputs:

tests/unit/orchestration/test_longitudinal_template.py

@@ -23,6 +23,11 @@
     "sub",
     "ses",
     "space",
+    "acq",
+    "dir",
+    "echo",
+    "part",
+    "rec",
     "task",
     "run",
     "desc",
@@ -42,6 +47,11 @@ def _brain_row(sub: str, ses: str) -> tuple:
         None,
         None,
         None,
+        None,
+        None,
+        None,
+        None,
+        None,
         "brain",
         "/data",
         path,