ENH: Utility that returns the B0FieldSource of a given potential EPI target

This code adds a ``get_identifier()`` interface that allows the user find the
a list of ``B0FieldIdentifier``(s) that are associated with a given file.

ATM, it pulls the information from the ``IntendedFor`` metadata of fieldmaps,
but in the future, it should first read the ``B0FieldSource``.

Originated-by: nipreps/dmriprep#140.
Resolves: #148.

Author: oesteban

sdcflows/fieldmaps.py

@@ -1,6 +1,7 @@
 """Utilities for fieldmap estimation."""
 from pathlib import Path
 from enum import Enum, auto
+from collections import defaultdict
 import re
 import attr
 from json import loads
@@ -11,6 +12,7 @@
 
 
 _estimators = EstimatorRegistry()
+_intents = defaultdict(set)
 
 
 class MetadataError(ValueError):
@@ -367,6 +369,12 @@ def __attrs_post_init__(self):
         if self.method == EstimatorType.UNKNOWN:
             raise ValueError("Insufficient sources to estimate a fieldmap.")
 
+        intents_meta = set(
+            el
+            for f in self.sources
+            for el in listify(f.metadata.get("IntendedFor") or [])
+        )
+
         # Register this estimation method
         if not self.bids_id:
             # If not manually set, try to get it from BIDS metadata
@@ -384,10 +392,15 @@ def __attrs_post_init__(self):
             elif bids_ids:
                 object.__setattr__(self, "bids_id", bids_ids.pop())
             else:
-                object.__setattr__(self, "bids_id", _estimators.add(self.paths()))
+                bids_id = _estimators.add(self.paths())
+                object.__setattr__(self, "bids_id", bids_id)
+                for intent_file in intents_meta:
+                    _intents[intent_file].add(bids_id)
                 return
 
         _estimators[self.bids_id] = self.paths()
+        for intent_file in intents_meta:
+            _intents[intent_file].add(self.bids_id)
 
     def paths(self):
         """Return a tuple of paths that are sorted."""
@@ -424,3 +437,30 @@ def get_workflow(self, **kwargs):
             self._wf = init_syn_sdc_wf(**kwargs)
 
         return self._wf
+
+
+def get_identifier(filename, by="IntendedFor"):
+    """
+    Return the fieldmap identifier for a given filename.
+
+    Parameters
+    ----------
+    filename : :obj:`str`
+        The file we want to search in indexes.
+    by : :obj:`str` (``IntendedFor`` or ``sources``)
+        Index table in which the search should be performed.
+
+    """
+    if by == "IntendedFor":
+        return tuple(sorted(_intents.get(filename, ())))
+
+    if by == "sources":
+        return _estimators.get_key(filename)
+
+    raise KeyError(r"'{by}'")
+
+
+def clear_registry():
+    """Empty registries."""
+    _estimators.clear()
+    _intents.clear()

sdcflows/tests/test_fieldmaps.py

@@ -69,7 +69,7 @@ def test_FieldmapEstimation(testdata_dir, inputfiles, method, nsources, raises):
             fm.FieldmapEstimation(sources)
 
         # Clean up so this parameter set can be tested.
-        fm._estimators.clear()
+        fm.clear_registry()
 
     fe = fm.FieldmapEstimation(sources)
     assert fe.method == method
@@ -109,17 +109,17 @@ def test_FieldmapEstimationError(testdata_dir, inputfiles, errortype):
     """Test errors."""
     sub_dir = testdata_dir / "sub-01"
 
-    fm._estimators.clear()
+    fm.clear_registry()
 
     with pytest.raises(errortype):
         fm.FieldmapEstimation([sub_dir / f for f in inputfiles])
 
-    fm._estimators.clear()
+    fm.clear_registry()
 
 
-def test_FieldmapEstimationIdentifier(testdata_dir):
+def test_FieldmapEstimationIdentifier(monkeypatch, testdata_dir):
     """Check some use cases of B0FieldIdentifier."""
-    fm._estimators.clear()
+    fm.clear_registry()
 
     with pytest.raises(ValueError):
         fm.FieldmapEstimation(
@@ -139,7 +139,11 @@ def test_FieldmapEstimationIdentifier(testdata_dir):
         [
             fm.FieldmapFile(
                 testdata_dir / "sub-01" / "fmap/sub-01_fieldmap.nii.gz",
-                metadata={"Units": "Hz", "B0FieldIdentifier": "fmap_0"},
+                metadata={
+                    "Units": "Hz",
+                    "B0FieldIdentifier": "fmap_0",
+                    "IntendedFor": "file1.nii.gz",
+                },
             ),
             fm.FieldmapFile(
                 testdata_dir / "sub-01" / "fmap/sub-01_magnitude.nii.gz",
@@ -148,6 +152,8 @@ def test_FieldmapEstimationIdentifier(testdata_dir):
         ]
     )
     assert fe.bids_id == "fmap_0"
+    assert fm.get_identifier("file1.nii.gz") == ("fmap_0",)
+    assert not fm.get_identifier("file2.nii.gz")
 
     with pytest.raises(KeyError):
         fm.FieldmapEstimation(
@@ -163,13 +169,17 @@ def test_FieldmapEstimationIdentifier(testdata_dir):
             ]
         )  # Consistent, but already exists
 
-    fm._estimators.clear()
+    fm.clear_registry()
 
     fe = fm.FieldmapEstimation(
         [
             fm.FieldmapFile(
                 testdata_dir / "sub-01" / "fmap/sub-01_fieldmap.nii.gz",
-                metadata={"Units": "Hz", "B0FieldIdentifier": "fmap_1"},
+                metadata={
+                    "Units": "Hz",
+                    "B0FieldIdentifier": "fmap_1",
+                    "IntendedFor": ["file1.nii.gz", "file2.nii.gz"],
+                },
             ),
             fm.FieldmapFile(
                 testdata_dir / "sub-01" / "fmap/sub-01_magnitude.nii.gz",
@@ -178,5 +188,15 @@ def test_FieldmapEstimationIdentifier(testdata_dir):
         ]
     )
     assert fe.bids_id == "fmap_1"
-
-    fm._estimators.clear()
+    assert fm.get_identifier("file1.nii.gz") == ("fmap_1",)
+    assert fm.get_identifier("file2.nii.gz") == ("fmap_1",)
+    assert not fm.get_identifier("file3.nii.gz")
+    assert fm.get_identifier(
+        str(testdata_dir / "sub-01" / "fmap/sub-01_magnitude.nii.gz"), by="sources"
+    ) == ("fmap_1",)
+
+    with monkeypatch.context() as m:
+        m.setattr(fm, "_intents", {"file1.nii.gz": {"fmap_0", "fmap_1"}})
+        assert fm.get_identifier("file1.nii.gz") == ("fmap_0", "fmap_1",)
+
+    fm.clear_registry()