Merge pull request #277 from ISISComputingGroup/Ticket152

Make run_number a metadata key in period-per-point mode & document/make it easier to add scan-specific metadata

Author: danielmaclaren

doc/tutorial/overview.md

@@ -190,6 +190,50 @@ As this is fairly common functionality for most plans, we have created a "standa
 For more information on callbacks, see
 [bluesky callbacks documentation](https://blueskyproject.io/bluesky/main/callbacks.html).
 
+## Metadata
+
+Bluesky provides a number of mechanisms for inserting metadata into a scan. In general, metadata may be any JSON-serialisable object including numbers, strings, lists, dictionaries, and nested combinations of those. The bluesky documentation has an {external+bluesky:doc}`extensive description of metadata mechanisms <metadata>`, but the main options are summarised below.
+
+**Persistently (for this python session)**
+```python
+RE.md["user"] = "Tom"
+RE.md["sample"] = "unobtainium"
+```
+
+**For one `RE` call**:
+```python
+RE(some_plan(), sample="unobtainium", user="Tom")
+```
+
+**Dynamically, within a plan (using {external+bluesky:py:obj}`bluesky.preprocessors.inject_md_wrapper`)**
+```python
+import bluesky.plan_stubs as bps
+from bluesky.preprocessors import inject_md_wrapper
+
+
+def some_plan(dae):
+    run_number = yield from bps.rd(dae.current_or_next_run_number_str)
+    return (yield from inject_md_wrapper(subplan(), {"run_number": run_number}))
+```
+
+**Dynamically, within a plan (using {external+bluesky:py:obj}`bluesky.preprocessors.inject_md_decorator`)**
+```python
+import bluesky.plan_stubs as bps
+from bluesky.preprocessors import inject_md_decorator
+
+
+def some_plan(dae):
+    run_number = yield from bps.rd(dae.current_or_next_run_number_str)
+    
+    @inject_md_decorator({"run_number": run_number})
+    def _inner():
+        yield from subplan()
+
+    yield from _inner()
+```
+
+In addition to the above mechanisms, many built-in bluesky plans (such as {external+bluesky:py:obj}`bluesky.plans.scan` and {py:obj}`ibex_bluesky_core.plans.scan`) take an `md` keyword argument, which can also be used to insert additional metadata for one scan.
+
 ## See also
 
 **Plans & plan-stubs**

src/ibex_bluesky_core/plans/__init__.py

@@ -31,6 +31,17 @@
 ]
 
 
+def _get_additional_md(
+    dae: "SimpleDae", *, periods: bool, save_run: bool
+) -> Generator[Msg, None, dict[str, Any]]:
+    if periods and save_run:
+        run_number = yield from bps.rd(dae.current_or_next_run_number_str)
+        return {"run_number": run_number}
+    else:
+        yield from bps.null()
+        return {}
+
+
 def scan(  # noqa: PLR0913
     dae: "SimpleDae",
     block: NamedMovable[float],
@@ -42,6 +53,7 @@ def scan(  # noqa: PLR0913
     periods: bool = True,
     save_run: bool = False,
     rel: bool = False,
+    md: dict[Any, Any] | None = None,
 ) -> Generator[Msg, None, ISISCallbacks]:
     """Scan the DAE against a Movable.
 
@@ -55,6 +67,7 @@ def scan(  # noqa: PLR0913
         periods: whether or not to use software periods.
         save_run: whether or not to save run.
         rel: whether or not to scan around the current position or use absolute positions.
+        md: Arbitrary metadata to include in this scan.
 
     """
     yield from ensure_connected(dae, block)  # type: ignore
@@ -66,13 +79,15 @@ def scan(  # noqa: PLR0913
 
     icc = _set_up_fields_and_icc(block, dae, model, periods, save_run, ax)
 
+    additional_md = yield from _get_additional_md(dae, periods=periods, save_run=save_run)
+
     @icc
     def _inner() -> Generator[Msg, None, None]:
         if rel:
             plan = bp.rel_scan
         else:
             plan = bp.scan
-        yield from plan([dae], block, start, stop, num=num)
+        yield from plan([dae], block, start, stop, num=num, md=additional_md | (md or {}))
 
     yield from _inner()
 
@@ -116,6 +131,7 @@ def adaptive_scan(  # noqa: PLR0913, PLR0917
     periods: bool = True,
     save_run: bool = False,
     rel: bool = False,
+    md: dict[Any, Any] | None = None,
 ) -> Generator[Msg, None, ISISCallbacks]:
     """Scan the DAE against a block using an adaptive scan.
 
@@ -133,6 +149,7 @@ def adaptive_scan(  # noqa: PLR0913, PLR0917
         periods: whether or not to use software periods.
         save_run: whether or not to save run.
         rel: whether or not to scan around the current position or use absolute positions.
+        md: Arbitrary metadata to include in this scan.
 
     Returns:
         an :obj:`ibex_bluesky_core.callbacks.ISISCallbacks` instance.
@@ -148,6 +165,8 @@ def adaptive_scan(  # noqa: PLR0913, PLR0917
 
     icc = _set_up_fields_and_icc(block, dae, model, periods, save_run, ax)
 
+    additional_md = yield from _get_additional_md(dae, periods=periods, save_run=save_run)
+
     @icc
     def _inner() -> Generator[Msg, None, None]:
         if rel:
@@ -164,6 +183,7 @@ def _inner() -> Generator[Msg, None, None]:
             max_step=max_step,
             target_delta=target_delta,
             backstep=True,
+            md=additional_md | (md or {}),
         )  # type: ignore
 
     yield from _inner()
@@ -185,6 +205,7 @@ def motor_scan(  # noqa: PLR0913
     periods: bool = True,
     save_run: bool = False,
     rel: bool = False,
+    md: dict[Any, Any] | None = None,
 ) -> Generator[Msg, None, ISISCallbacks]:
     """Wrap our scan() plan and create a block_rw and a DAE object.
 
@@ -205,6 +226,7 @@ def motor_scan(  # noqa: PLR0913
         periods: whether or not to use software periods.
         save_run: whether or not to save run.
         rel: whether or not to scan around the current position or use absolute positions.
+        md: Arbitrary metadata to include in this scan.
 
     Returns:
         an :obj:`ibex_bluesky_core.callbacks.ISISCallbacks` instance.
@@ -231,6 +253,7 @@ def motor_scan(  # noqa: PLR0913
             save_run=save_run,
             periods=periods,
             rel=rel,
+            md=md,
         )
     )
 
@@ -251,6 +274,7 @@ def motor_adaptive_scan(  # noqa: PLR0913
     periods: bool = True,
     save_run: bool = False,
     rel: bool = False,
+    md: dict[Any, Any] | None = None,
 ) -> Generator[Msg, None, ISISCallbacks]:
     """Wrap adaptive_scan() plan and create a block_rw and a DAE object.
 
@@ -273,6 +297,7 @@ def motor_adaptive_scan(  # noqa: PLR0913
         periods: whether or not to use software periods.
         save_run: whether or not to save run.
         rel: whether or not to scan around the current position or use absolute positions.
+        md: Arbitrary metadata to include in this scan.
 
     Returns:
         an :obj:`ibex_bluesky_core.callbacks.ISISCallbacks` instance.
@@ -300,5 +325,6 @@ def motor_adaptive_scan(  # noqa: PLR0913
             model=model,
             save_run=save_run,
             rel=rel,
+            md=md,
         )
     )

src/ibex_bluesky_core/plans/reflectometry/__init__.py

@@ -1,6 +1,7 @@
 """Plans specific to Reflectometry beamlines."""
 
 from collections.abc import Generator
+from typing import Any
 
 from bluesky import Msg
 
@@ -44,6 +45,7 @@ def refl_scan(  # noqa: PLR0913
     periods: bool = True,
     save_run: bool = False,
     rel: bool = False,
+    md: dict[Any, Any] | None = None,
 ) -> Generator[Msg, None, ISISCallbacks]:
     """Scan over a reflectometry parameter.
 
@@ -62,6 +64,7 @@ def refl_scan(  # noqa: PLR0913
         periods: whether to use periods.
         save_run: whether to save the run of the scan.
         rel: whether to use a relative scan around the current position.
+        md: Arbitrary metadata to include in this scan.
 
     Returns:
         an :obj:`ibex_bluesky_core.callbacks.ISISCallbacks` instance.
@@ -84,6 +87,7 @@ def refl_scan(  # noqa: PLR0913
             save_run=save_run,
             periods=periods,
             rel=rel,
+            md=md,
         )
     )
 
@@ -104,6 +108,7 @@ def refl_adaptive_scan(  # noqa: PLR0913
     periods: bool = True,
     save_run: bool = False,
     rel: bool = False,
+    md: dict[Any, Any] | None = None,
 ) -> Generator[Msg, None, ISISCallbacks]:
     """Perform an adaptive scan over a reflectometry parameter.
 
@@ -124,6 +129,7 @@ def refl_adaptive_scan(  # noqa: PLR0913
         periods: whether to use periods.
         save_run: whether to save the run of the scan.
         rel: whether to use a relative scan around the current position.
+        md: Arbitrary metadata to include in this scan.
 
     Returns:
         an :obj:`ibex_bluesky_core.callbacks.ISISCallbacks` instance.
@@ -147,5 +153,6 @@ def refl_adaptive_scan(  # noqa: PLR0913
             model=model,
             save_run=save_run,
             rel=rel,
+            md=md,
         )
     )

src/ibex_bluesky_core/plans/reflectometry/_det_map_align.py

@@ -172,6 +172,7 @@ def angle_scan_plan(
     *,
     angle_map: npt.NDArray[np.float64],
     flood: sc.Variable | None = None,
+    md: dict[Any, Any] | None = None,
 ) -> Generator[Msg, None, ModelResult | None]:
     """Reflectometry detector-mapping angle alignment plan.
 
@@ -187,6 +188,7 @@ def angle_scan_plan(
             This array should be aligned along a "spectrum" dimension; counts are
             divided by this array before being used in fits. This is used to
             normalise the intensities detected by each detector pixel.
+        md: Arbitrary metadata to include in this scan.
 
     """
     logger.info("Starting angle scan")
@@ -212,7 +214,7 @@ def angle_scan_plan(
         ]
     )
     def _inner() -> Generator[Msg, None, None]:
-        yield from bp.count([dae])
+        yield from bp.count([dae], md=md)
 
     yield from _inner()
 
@@ -247,6 +249,7 @@ def height_and_angle_scan_plan(  # noqa PLR0913
     angle_map: npt.NDArray[np.float64],
     rel: bool = False,
     flood: sc.Variable | None = None,
+    md: dict[Any, Any] | None = None,
 ) -> Generator[Msg, None, DetMapAlignResult]:
     """Reflectometry detector-mapping simultaneous height & angle alignment plan.
 
@@ -266,6 +269,7 @@ def height_and_angle_scan_plan(  # noqa PLR0913
             This array should be aligned along a "spectrum" dimension; counts are
             divided by this array before being used in fits. This is used to
             normalise the intensities detected by each detector pixel.
+        md: Arbitrary metadata to include in this scan.
 
     Returns:
         A dictionary containing the fit results from gaussian height and angle fits.
@@ -309,7 +313,7 @@ def _inner() -> Generator[Msg, None, None]:
         nonlocal start, stop, num
         yield from bps.mv(dae.number_of_periods, num)
         plan = bp.rel_scan if rel else bp.scan
-        yield from plan([dae], height, start, stop, num=num)
+        yield from plan([dae], height, start, stop, num=num, md=md)
 
     yield from _inner()
 

tests/plans/test_init.py

@@ -1,6 +1,9 @@
 # pyright: reportMissingParameterType=false
+import functools
+from typing import Any
 from unittest.mock import patch
 
+import bluesky.utils
 import pytest
 from bluesky.preprocessors import run_decorator
 from ophyd_async.plan_stubs import ensure_connected
@@ -121,6 +124,36 @@ def test_scan_does_normal_scan_when_relative_false(RE, dae, block):
     assert count == bp_scan.call_args[1]["num"]
 
 
+@pytest.mark.parametrize(
+    "scan_func",
+    [
+        functools.partial(scan, start=1, stop=2, num=2),
+        functools.partial(adaptive_scan, start=1, stop=2, min_step=1, max_step=2, target_delta=1),
+    ],
+)
+def test_if_in_periods_mode_and_run_saved_then_scan_start_doc_contains_run_number(
+    RE, dae, block, scan_func
+):
+    set_mock_value(dae.current_or_next_run_number_str, "12345678")
+
+    start_doc: dict[str, Any] | None = None
+
+    def _cb(typ, doc):
+        if typ == "start":
+            nonlocal start_doc
+            start_doc = doc
+
+    with (
+        patch("ibex_bluesky_core.plans.ensure_connected"),
+    ):
+        # Scan fails because DAE isn't set up right... but it still emits a start doc so that's fine
+        with pytest.raises(bluesky.utils.FailedStatus):
+            RE(scan_func(dae, block, rel=False, periods=True, save_run=True), _cb)
+
+    assert start_doc is not None
+    assert start_doc.get("run_number") == "12345678"
+
+
 def test_scan_does_relative_scan_when_relative_true(RE, dae, block):
     start = 0
     stop = 2