feat: Add workflow arguments for metadata estimates and fallback TRT

effigies · effigies · commit 0595ff05c8ba · 2025-01-29T15:50:43.000-05:00
diff --git a/sdcflows/workflows/apply/correction.py b/sdcflows/workflows/apply/correction.py
@@ -29,6 +29,8 @@
 def init_unwarp_wf(
     *,
     jacobian=True,
+    use_metadata_estimates=False,
+    fallback_total_readout_time=None,
     free_mem=None,
     omp_nthreads=1,
     debug=False,
@@ -117,8 +119,16 @@ def init_unwarp_wf(
         name="outputnode",
     )
 
-    rotime = pe.Node(GetReadoutTime(), name="rotime")
+    rotime = pe.Node(
+        GetReadoutTime(
+            use_estimate=use_metadata_estimates,
+        ),
+        name="rotime",
+        run_without_submitting=True,
+    )
     rotime.interface._always_run = debug
+    if fallback_total_readout_time is not None:
+        rotime.inputs.fallback = fallback_total_readout_time
 
     # resample is memory-hungry; choose a smaller number of threads
     # if we know how much memory we have to work with
diff --git a/sdcflows/workflows/base.py b/sdcflows/workflows/base.py
@@ -36,6 +36,8 @@ def init_fmap_preproc_wf(
     omp_nthreads,
     output_dir,
     subject,
+    use_metadata_estimates=False,
+    fallback_total_readout_time=None,
     sloppy=False,
     debug=False,
     name="fmap_preproc_wf",
@@ -107,6 +109,8 @@ def init_fmap_preproc_wf(
 
     for n, estimator in enumerate(estimators, 1):
         est_wf = estimator.get_workflow(
+            use_metadata_estimates=use_metadata_estimates,
+            fallback_total_readout_time=fallback_total_readout_time,
             omp_nthreads=omp_nthreads,
             debug=debug,
             sloppy=sloppy,
diff --git a/sdcflows/workflows/fit/fieldmap.py b/sdcflows/workflows/fit/fieldmap.py
@@ -29,7 +29,14 @@
 INPUT_FIELDS = ("magnitude", "fieldmap")
 
 
-def init_fmap_wf(omp_nthreads=1, sloppy=False, debug=False, mode="phasediff", name="fmap_wf"):
+def init_fmap_wf(
+    omp_nthreads=1,
+    sloppy=False,
+    debug=False,
+    mode="phasediff",
+    name="fmap_wf",
+    **kwargs,  # Soak up unused keyword args
+):
     """
     Estimate the fieldmap based on a field-mapping MRI acquisition.
 
diff --git a/sdcflows/workflows/fit/pepolar.py b/sdcflows/workflows/fit/pepolar.py
@@ -36,6 +36,8 @@
 
 def init_topup_wf(
     grid_reference=0,
+    use_metadata_estimates=False,
+    fallback_total_readout_time=None,
     omp_nthreads=1,
     sloppy=False,
     debug=False,
@@ -122,11 +124,15 @@ def init_topup_wf(
 
     # Calculate the total readout time of each run
     readout_time = pe.MapNode(
-        GetReadoutTime(),
+        GetReadoutTime(
+            use_estimate=use_metadata_estimates,
+        ),
         name="readout_time",
         iterfield=["metadata", "in_file"],
         run_without_submitting=True,
     )
+    if fallback_total_readout_time is not None:
+        readout_time.inputs.fallback = fallback_total_readout_time
     # Average each run so that topup is not overwhelmed (see #279)
     runwise_avg = pe.MapNode(
         RobustAverage(num_threads=omp_nthreads),
diff --git a/sdcflows/workflows/fit/syn.py b/sdcflows/workflows/fit/syn.py
@@ -42,6 +42,8 @@
 def init_syn_sdc_wf(
     *,
     atlas_threshold=3,
+    use_metadata_estimates=False,
+    fallback_total_readout_time=None,
     sloppy=False,
     debug=False,
     name="syn_sdc_wf",
@@ -160,10 +162,14 @@ def init_syn_sdc_wf(
     outputnode.inputs.method = 'FLB ("fieldmap-less", SyN-based)'
 
     readout_time = pe.Node(
-        GetReadoutTime(),
+        GetReadoutTime(
+            use_estimate=use_metadata_estimates,
+        ),
         name="readout_time",
         run_without_submitting=True,
     )
+    if fallback_total_readout_time is not None:
+        readout_time.inputs.fallback = fallback_total_readout_time
 
     warp_dir = pe.Node(
         niu.Function(function=_warp_dir),
