Add keyword args to allow for stripping run and monitor types from GenericNeXusWorkflow

src/ess/reduce/nexus/workflow.py

@@ -3,8 +3,10 @@
 
 """Workflow and workflow components for interacting with NeXus files."""
 
+from collections.abc import Sequence
 from typing import Any
 
+import networkx as nx
 import sciline
 import scipp as sc
 import scippnexus as snx
@@ -518,11 +520,48 @@ def LoadDetectorWorkflow() -> sciline.Pipeline:
     return wf
 
 
-def GenericNeXusWorkflow() -> sciline.Pipeline:
-    """Generic workflow for loading detector and monitor data from a NeXus file."""
+def GenericNeXusWorkflow(
+    *,
+    run_types: Sequence[sciline.typing.Key] | None = None,
+    monitor_types: Sequence[sciline.typing.Key] | None = None,
+) -> sciline.Pipeline:
+    """
+    Generic workflow for loading detector and monitor data from a NeXus file.
+
+    Parameters
+    ----------
+    run_types:
+        List of run types to include in the workflow. If not provided, all run types
+        are included. It is recommended to specify run types to avoid creating very
+        large workflows.
+    monitor_types:
+        List of monitor types to include in the workflow. If not provided, all monitor
+        types are included. It is recommended to specify monitor types to avoid creating
+        very large workflows.
+
+    Returns
+    -------
+    :
+        The workflow.
+    """
+    if monitor_types is not None and run_types is None:
+        raise ValueError("run_types must be specified if monitor_types is specified")
     wf = sciline.Pipeline(
         (*_common_providers, *_monitor_providers, *_detector_providers)
     )
     wf[DetectorBankSizes] = DetectorBankSizes({})
     wf[PreopenNeXusFile] = PreopenNeXusFile(False)
+
+    g = wf.underlying_graph
+    ancestors = set()
+    # DetectorData and MonitorData are the "final" outputs, so finding and removing all
+    # their ancestors is what we need to strip unused run and monitor types.
+    for rt in run_types or ():
+        ancestors |= nx.ancestors(g, DetectorData[rt])
+        ancestors.add(DetectorData[rt])
+        for mt in monitor_types or ():
+            ancestors |= nx.ancestors(g, MonitorData[rt, mt])
+            ancestors.add(MonitorData[rt, mt])
+    if run_types is not None:
+        g.remove_nodes_from(set(g.nodes) - ancestors)
     return wf

tests/nexus/workflow_test.py

@@ -8,10 +8,14 @@
 from ess.reduce import data
 from ess.reduce.nexus import compute_component_position, workflow
 from ess.reduce.nexus.types import (
+    BackgroundRun,
     DetectorData,
     Filename,
     Monitor1,
+    Monitor2,
+    Monitor3,
     MonitorData,
+    NeXusComponentLocationSpec,
     NeXusName,
     NeXusTransformation,
     SampleRun,
@@ -397,3 +401,36 @@ def test_generic_nexus_workflow() -> None:
     assert 'source_position' in da.coords
     assert da.bins is not None
     assert da.dims == ('event_time_zero',)
+
+
+def test_generic_nexus_workflow_raises_if_monitor_types_but_not_run_types_given() -> (
+    None
+):
+    with pytest.raises(ValueError, match='run_types'):
+        GenericNeXusWorkflow(monitor_types=[Monitor1])
+
+
+def test_generic_nexus_workflow_includes_only_given_run_and_monitor_types() -> None:
+    wf = GenericNeXusWorkflow(run_types=[SampleRun], monitor_types=[Monitor1, Monitor3])
+    graph = wf.underlying_graph
+    assert DetectorData[SampleRun] in graph
+    assert DetectorData[BackgroundRun] not in graph
+    assert MonitorData[SampleRun, Monitor1] in graph
+    assert MonitorData[SampleRun, Monitor2] not in graph
+    assert MonitorData[SampleRun, Monitor3] in graph
+    assert MonitorData[BackgroundRun, Monitor1] not in graph
+    assert MonitorData[BackgroundRun, Monitor2] not in graph
+    assert MonitorData[BackgroundRun, Monitor3] not in graph
+    # Many other keys are also removed, this is just an example
+    assert NeXusComponentLocationSpec[Monitor1, SampleRun] in graph
+    assert NeXusComponentLocationSpec[Monitor2, SampleRun] not in graph
+    assert NeXusComponentLocationSpec[Monitor3, SampleRun] in graph
+    assert NeXusComponentLocationSpec[snx.NXdetector, SampleRun] in graph
+    assert NeXusComponentLocationSpec[snx.NXsample, SampleRun] in graph
+    assert NeXusComponentLocationSpec[snx.NXsource, SampleRun] in graph
+    assert NeXusComponentLocationSpec[Monitor1, BackgroundRun] not in graph
+    assert NeXusComponentLocationSpec[Monitor2, BackgroundRun] not in graph
+    assert NeXusComponentLocationSpec[Monitor3, BackgroundRun] not in graph
+    assert NeXusComponentLocationSpec[snx.NXdetector, BackgroundRun] not in graph
+    assert NeXusComponentLocationSpec[snx.NXsample, BackgroundRun] not in graph
+    assert NeXusComponentLocationSpec[snx.NXsource, BackgroundRun] not in graph