RF/FIX: Prioritize sbref and shortest echo for SyN-SDC

sdcflows/utils/tests/test_wrangler.py

@@ -1,4 +1,5 @@
 import pytest
+from shutil import rmtree
 
 from niworkflows.utils.testing import generate_bids_skeleton
 
@@ -270,3 +271,118 @@ def test_single_reverse_pedir(tmp_path):
         # IntendedFor is a list of strings
         # REGRESSION: The result was a PyBIDS BIDSFile (fmriprep#3020)
         assert epi.metadata['IntendedFor'] == [str(bold.path.relative_to(subject_root))]
+
+
+def test_fieldmapless(tmp_path):
+    bids_dir = tmp_path / "bids"
+
+    T1w = {"suffix": "T1w"}
+    bold = {
+        "task": "rest",
+        "suffix": "bold",
+        "metadata": {
+            "RepetitionTime": 0.8,
+            "TotalReadoutTime": 0.5,
+            "PhaseEncodingDirection": "j",
+        },
+    }
+    me_metadata = [
+        {"EchoTime": 0.01 * i, **bold["metadata"]}
+        for i in range(1, 4)
+    ]
+    sbref = {**bold, **{"suffix": "sbref"}}
+    spec = {
+        "01": {
+            "anat": [T1w],
+            "func": [bold],
+        },
+    }
+    generate_bids_skeleton(bids_dir, spec)
+    layout = gen_layout(bids_dir)
+    est = find_estimators(layout=layout, subject="01", fmapless=True)
+    assert len(est) == 1
+    assert len(est[0].sources) == 2
+    clear_registry()
+    rmtree(bids_dir)
+
+    # Multi-run generates one estimator per-run
+    spec = {
+        "01": {
+            "anat": [T1w],
+            "func": [{"run": i, **bold} for i in range(1, 3)],
+        },
+    }
+    generate_bids_skeleton(bids_dir, spec)
+    layout = gen_layout(bids_dir)
+    est = find_estimators(layout=layout, subject="01", fmapless=True)
+    assert len(est) == 2
+    assert len(est[0].sources) == 2
+    clear_registry()
+    rmtree(bids_dir)
+
+    # Multi-echo should only generate one estimator
+    spec = {
+        "01": {
+            "anat": [T1w],
+            "func": [{"echo": i + 1, **bold, **{"metadata": me_metadata[i]}} for i in range(3)],
+        },
+    }
+    generate_bids_skeleton(bids_dir, spec)
+    layout = gen_layout(bids_dir)
+    est = find_estimators(layout=layout, subject="01", fmapless=True)
+    assert len(est) == 1
+    assert len(est[0].sources) == 2
+    clear_registry()
+    rmtree(bids_dir)
+
+    # Matching bold+sbref should generate only one estimator
+    spec = {
+        "01": {
+            "anat": [T1w],
+            "func": [bold, sbref],
+        },
+    }
+    generate_bids_skeleton(bids_dir, spec)
+    layout = gen_layout(bids_dir)
+    est = find_estimators(layout=layout, subject="01", fmapless=True)
+    assert len(est) == 1
+    assert len(est[0].sources) == 2
+    clear_registry()
+    rmtree(bids_dir)
+
+    # Mismatching bold+sbref should generate two sbrefs
+    spec = {
+        "01": {
+            "anat": [T1w],
+            "func": [{"acq": "A", **bold}, {"acq": "B", **sbref}],
+        },
+    }
+    generate_bids_skeleton(bids_dir, spec)
+    layout = gen_layout(bids_dir)
+    est = find_estimators(layout=layout, subject="01", fmapless=True)
+    assert len(est) == 2
+    assert len(est[0].sources) == 2
+    clear_registry()
+    rmtree(bids_dir)
+
+    # Multiecho bold+sbref should generate only one estimator
+    spec = {
+        "01": {
+            "anat": [T1w],
+            "func": [
+                {"echo": i + 1, **bold, **{"metadata": me_metadata[i]}}
+                for i in range(3)
+            ]
+            + [
+                {"echo": i + 1, **sbref, **{"metadata": me_metadata[i]}}
+                for i in range(3)
+            ],
+        },
+    }
+    generate_bids_skeleton(bids_dir, spec)
+    layout = gen_layout(bids_dir)
+    est = find_estimators(layout=layout, subject="01", fmapless=True)
+    assert len(est) == 1
+    assert len(est[0].sources) == 2
+    clear_registry()
+    rmtree(bids_dir)

sdcflows/utils/wrangler.py

@@ -26,7 +26,7 @@
 from itertools import product
 from contextlib import suppress
 from pathlib import Path
-from typing import Optional, Union, List
+from typing import Optional, Union, List, Dict, Any
 from bids.layout import BIDSLayout, BIDSFile
 from bids.utils import listify
 
@@ -493,48 +493,111 @@ def find_estimators(
         return estimators
 
     logger.debug("Attempting fmap-less estimation")
+    estimator_specs = find_anatomical_estimators(
+        anat_file=anat_file[0],
+        layout=layout,
+        subject=subject,
+        sessions=sessions,
+        base_entities=base_entities,
+        suffixes=fmapless,
+    )
+    for spec in estimator_specs:
+        try:
+            estimator = fm.FieldmapEstimation(spec)
+        except (ValueError, TypeError) as err:
+            _log_debug_estimator_fail(logger, "ANAT", spec, layout.root, str(err))
+        else:
+            _log_debug_estimation(logger, estimator, layout.root)
+            estimators.append(estimator)
+    return estimators
+
+
+def find_anatomical_estimators(
+    *,
+    anat_file: BIDSFile,
+    layout: BIDSLayout,
+    subject: str,
+    sessions: List[str],
+    base_entities: Dict[str, Any],
+    suffixes: List[str],
+) -> List[List[fm.FieldmapFile]]:
+    r"""Find anatomical estimators
+
+    Given an anatomical reference image, create lists of files for estimating
+    susceptibility distortion for the EPI images in a dataset.
+
+    Parameters
+    ----------
+    anat_file : :class:`bids.layout.BIDSFile`
+        Anatomical reference image to use in estimators.
+    layout : :class:`bids.layout.BIDSLayout`
+        An initialized PyBIDS layout.
+    subject : :class:`str`
+        Participant label for this single-subject workflow.
+    sessions : :class:`list`
+        One of more session identifiers. To use all, pass ``[None]``.
+    base_entities : :class:`dict`
+        Entities to use to query for images. These should include any filters.
+    suffixes : :class:`list`
+        EPI suffixes, for example ``["bold", "dwi"]``. Associated ``"sbref"``\s
+        will be found and used in place of BOLD/diffusion EPIs.
+    """
+
     from .epimanip import get_trt
 
-    for ses, suffix in sorted(product(sessions, fmapless)):
-        candidates = layout.get(**{**base_entities, **{'suffix': suffix, 'session': ses}})
+    subject_root = Path(layout.root) / f"sub-{subject}"
+
+    hits = set()  # Avoid duplicates
+    estimators = []
+    for ses, suffix in sorted(product(sessions, suffixes)):
+        suffixes = ["sbref", suffix]  # Order indicates preference; prefer sbref
+        datatype = {
+            "bold": "func",
+            "dwi": "dwi",
+        }[suffix]
+        candidates = layout.get(
+            **{
+                **base_entities,
+                **{"suffix": suffixes, "session": ses, "datatype": datatype},
+            }
+        )
 
         # Filter out candidates without defined PE direction
         epi_targets = []
-        pe_dirs = []
-        ro_totals = []
 
         for candidate in candidates:
             meta = candidate.get_metadata()
-            pe_dir = meta.get("PhaseEncodingDirection")
 
-            if not pe_dir:
+            if not meta.get("PhaseEncodingDirection"):
                 continue
 
-            pe_dirs.append(pe_dir)
-            ro = 1.0
+            trt = 1.0
             with suppress(ValueError):
-                ro = get_trt(meta, candidate.path)
-            ro_totals.append(ro)
-            meta.update({"TotalReadoutTime": ro})
-            epi_targets.append(fm.FieldmapFile(candidate.path, metadata=meta))
-
-        trivial_estimators = [
-            [
-                fm.FieldmapFile(
-                    anat_file[0],
-                    metadata={"IntendedFor": str(Path(epi.path).relative_to(subject_root))},
-                ),
-                epi,
-            ] for epi in epi_targets
-        ]
-
-        # TODO: Grouping could be done here; previously we grouped by (pe_dir, ro_time) pairs
-        syn_estimators = [fm.FieldmapEstimation(e) for e in trivial_estimators]
-
-        for e in syn_estimators:
-            _log_debug_estimation(logger, e, layout.root)
-
-        estimators.extend(syn_estimators)
+                trt = get_trt(meta, candidate.path)
+            meta.update({"TotalReadoutTime": trt})
+            epi_targets.append(fm.FieldmapFile(candidate, metadata=meta))
+
+        def sort_key(fmap):
+            # Return sbref before DWI/BOLD and shortest echo first
+            return suffixes.index(fmap.suffix), fmap.metadata.get("EchoTime", 1)
+
+        for target in sorted(epi_targets, key=sort_key):
+            if target.path in hits:
+                continue
+            query = {**base_entities, **target.entities}
+
+            # Find all echos, so strip from query, if present
+            query.pop("echo", None)
+
+            # Include sbref and EPI images in IntendedFor
+            # No harm in including sbrefs that won't be corrected,
+            # and ensures the hits set prevents doubling up
+            intent = [Path(epi) for epi in layout.get(suffix=suffixes, **query)]
+            metadata = {
+                "IntendedFor": [str(epi.relative_to(subject_root)) for epi in intent]
+            }
+            estimators.append([fm.FieldmapFile(anat_file, metadata=metadata), target])
+            hits.update(intent)
 
     return estimators