Filter anat suffix during load

Author: kaitj

src/rbc/cli/all.py

@@ -106,7 +106,7 @@ def main(args: AllArgs) -> int:  # noqa: C901
         session = load_session(sub_ses_group, pipe_ctx.sub, pipe_ctx.ses)
 
         # --- Anatomical (once per session, first T1w) ---
-        for _, anat_df in session.anat.filter(pl.col("suffix") == "T1w").group_by(
+        for _, anat_df in session.anat.group_by(
             _ANAT_GROUP_ENTITIES, maintain_order=True
         ):
             anat_row = anat_df.filter(suffix="T1w").row(0, named=True)

src/rbc/cli/anatomical.py

@@ -67,7 +67,7 @@ def main(args: AnatomicalArgs) -> int:
         )
         session = load_session(sub_ses_group, pipe_ctx.sub, pipe_ctx.ses)
 
-        for _, anat_df in session.anat.filter(pl.col("suffix") == "T1w").group_by(
+        for _, anat_df in session.anat.group_by(
             _ANAT_GROUP_ENTITIES, maintain_order=True
         ):
             row = anat_df.row(0, named=True)

src/rbc/cli/longitudinal.py

@@ -193,7 +193,7 @@ def main(args: LongitudinalArgs) -> int:
             raise ValueError("No longitudinal template found")
 
         if args.anatomical:
-            for _, anat_df in session.anat.filter(pl.col("suffix") == "T1w").group_by(
+            for _, anat_df in session.anat.group_by(
                 ("run", "acq"), maintain_order=True
             ):
                 _process_anat(pipe_ctx=pipe_ctx, anat_df=anat_df, tpl_df=tpl_df)

src/rbc/cli/query.py

@@ -6,6 +6,8 @@
 
 import polars as pl
 
+from rbc.core.bids import Suffix
+
 if TYPE_CHECKING:
     from collections.abc import Iterator, Sequence
 
@@ -32,7 +34,11 @@ def load_session(df: pl.DataFrame, subject: str, session: str | None) -> Session
     base: list[pl.Expr] = [pl.col("sub") == subject]
     if session is not None:
         base.append(pl.col("ses") == session)
-    anat_df = df.filter(pl.all_horizontal([*base, pl.col("datatype") == "anat"]))
+    anat_df = df.filter(
+        pl.all_horizontal(
+            [*base, pl.col("datatype") == "anat", pl.col("suffix") == Suffix.T1W]
+        )
+    )
     func_df = df.filter(pl.all_horizontal([*base, pl.col("datatype") == "func"]))
 
     return SessionTables(anat=anat_df, func=func_df if not func_df.is_empty() else None)