Improved XDI support

environment.yml

@@ -77,6 +77,7 @@ dependencies:
   - scikit-image
   - xlrd
   - peakutils
+  - pyarrow < 11.0.0  # Tempoary fix, remove once the libort missing symbol doesn't break CI
 
   # XES analysis packages
   - scikit-learn

src/conftest.py

@@ -32,6 +32,7 @@
 from haven.instrument.dxp import DxpDetector
 from haven.instrument.dxp import add_mcas as add_dxp_mcas
 from haven.instrument.ion_chamber import IonChamber
+from haven.instrument.monochromator import Monochromator
 from haven.instrument.robot import Robot
 from haven.instrument.shutter import Shutter
 from haven.instrument.slits import ApertureSlits, BladeSlits
@@ -274,6 +275,13 @@ def aerotech_flyer(aerotech):
     yield flyer
 
 
+@pytest.fixture()
+def mono(sim_registry):
+    mono = instantiate_fake_device(Monochromator, name="monochromator")
+    sim_registry.register(mono)
+    yield mono
+
+
 @pytest.fixture()
 def aps(sim_registry):
     aps = instantiate_fake_device(ApsMachine, name="APS")

src/haven/constants.py

@@ -1,8 +1,11 @@
+import re
+
 import xraydb
 
 
 def edge_energy(edge_name):
-    E0 = xraydb.xray_edge(*edge_name.split("_")).energy
+    element, shell = re.split(r"[-_]", edge_name)
+    E0 = xraydb.xray_edge(element, shell).energy
     return E0
 
 

src/haven/exceptions.py

@@ -57,8 +57,10 @@ class InvalidPV(ValueError):
 
 
 class DocumentNotFound(RuntimeError):
-    """An attempt was made to retrieve a document from the mongodb database,
-    but the requested document was not available."""
+    """An attempt was made to use a document, but the requested document
+    was not available.
+
+    """
 
     ...
 

src/haven/plans/energy_scan.py

@@ -8,12 +8,14 @@
 
 import numpy as np
 from bluesky import plans as bp
+from bluesky.preprocessors import subs_decorator
 
 from .._iconfig import load_config
 from ..constants import edge_energy
 from ..instrument import registry
 from ..preprocessors import baseline_decorator
 from ..typing import DetectorList
+from ..xdi_writer import XDIWriter
 
 __all__ = ["energy_scan"]
 
@@ -22,6 +24,9 @@
 
 
 # @shutter_suspend_decorator()
+@subs_decorator(
+    XDIWriter("{manager_path}/{year}{month}{day}-{sample_name}-{edge}-{short_uid}.xdi")
+)
 @baseline_decorator()
 def energy_scan(
     energies: Sequence[float],
@@ -125,12 +130,21 @@ def energy_scan(
     scan_args = [(motor, energies) for motor in energy_positioners]
     scan_args += [(motor, exposure) for motor in time_positioners]
     scan_args = [item for items in scan_args for item in items]
-    # Do the actual scan
+    # Add some extra metadata
     config = load_config()
+    md_ = {"edge": E0_str, "E0": E0}
+    for positioner in energy_positioners:
+        try:
+            md_["d_spacing"] = positioner.d_spacing.get()
+        except AttributeError:
+            continue
+        else:
+            break
+    # Do the actual scan
     yield from bp.list_scan(
         real_detectors,
         *scan_args,
-        md=ChainMap(md, {"edge": E0_str, "E0": E0}, config),
+        md=ChainMap(md, md_, config),
     )
 
 

src/haven/tests/test_energy_xafs_scan.py

@@ -41,7 +41,10 @@ def I0(sim_registry):
     return I0
 
 
-def test_energy_scan_basics(mono_motor, id_gap_motor, energies, RE):
+def test_energy_scan_basics(
+    beamline_manager, mono_motor, id_gap_motor, energies, RE, tmp_path
+):
+    beamline_manager.local_storage.full_path._readback = str(tmp_path)
     exposure_time = 1e-3
     # Set up fake detectors and motors
     I0_exposure = sim.SynAxis(
@@ -77,8 +80,9 @@ def test_energy_scan_basics(mono_motor, id_gap_motor, energies, RE):
         exposure=exposure_time,
         energy_positioners=[mono_motor, id_gap_motor],
         time_positioners=[I0_exposure, It_exposure],
+        md={"edge": "Ni_K"},
     )
-    result = RE(scan)
+    result = RE(scan, sample_name="xafs_sample")
     # Check that the mono and ID gap ended up in the right position
     # time.sleep(1.0)
     assert mono_motor.readback.get() == np.max(energies)
@@ -93,6 +97,29 @@ def test_raises_on_empty_positioners(RE, energies):
         RE(energy_scan(energies, energy_positioners=[]))
 
 
+def test_saves_dspacing(mono, energies, I0, It):
+    """Does the mono's d-spacing get added to metadata."""
+    # Prepare the messages from the plan
+    mono.d_spacing._readback = 1.5418
+    msgs = list(
+        energy_scan(
+            energies,
+            detectors=[It],
+            energy_positioners=[mono],
+            time_positioners=[It.exposure_time],
+        )
+    )
+    # Find the metadata written by the plan
+    for msg in msgs:
+        if msg.command == "open_run":
+            md = msg.kwargs
+            break
+    else:
+        raise RuntimeError("No open run message found")
+    # Check for the dspacing of the mono in the metadata
+    assert md["d_spacing"] == 1.5418
+
+
 def test_single_range(mono_motor, exposure_motor, I0):
     E0 = 10000
     expected_energies = np.arange(9990, 10001, step=1)

src/haven/tests/test_instrument_registry.py

@@ -5,9 +5,12 @@
 from haven.instrument import InstrumentRegistry
 
 
-def test_register_component():
-    # Prepare registry
-    reg = InstrumentRegistry(auto_register=False)
+@pytest.fixture()
+def reg():
+    yield InstrumentRegistry(auto_register=False)
+
+
+def test_register_component(reg):
     # Create an unregistered component
     cpt = sim.SynGauss(
         "I0",
@@ -33,9 +36,8 @@ def test_register_component():
     assert cpt in results
 
 
-def test_find_missing_components():
+def test_find_missing_components(reg):
     """Test that registry raises an exception if no matches are found."""
-    reg = InstrumentRegistry()
     cpt = sim.SynGauss(
         "I0",
         sim.motor,
@@ -51,27 +53,24 @@ def test_find_missing_components():
         reg.findall(label="spam")
 
 
-def test_find_allow_missing_components():
+def test_find_allow_missing_components(reg):
     """Test that registry tolerates missing components with the
     *allow_none* argument.
 
     """
-    reg = InstrumentRegistry()
     # Get some non-existent devices and check that the right nothing is returned
     assert list(reg.findall(label="spam", allow_none=True)) == []
     assert reg.find(name="eggs", allow_none=True) is None
 
 
-def test_exceptions():
-    reg = InstrumentRegistry()
+def test_exceptions(reg):
     reg.register(Device("", name="It"))
     # Test if a non-existent labels throws an exception
     with pytest.raises(ComponentNotFound):
         reg.find(label="spam")
 
 
-def test_as_class_decorator():
-    reg = InstrumentRegistry()
+def test_as_class_decorator(reg):
     # Create a dummy decorated class
     IonChamber = type("IonChamber", (Device,), {})
     IonChamber = reg.register(IonChamber)
@@ -83,9 +82,7 @@ def test_as_class_decorator():
     assert result.name == "I0"
 
 
-def test_find_component():
-    # Prepare registry
-    reg = InstrumentRegistry()
+def test_find_component(reg):
     # Create an unregistered component
     cptA = sim.SynGauss(
         "I0",
@@ -116,9 +113,7 @@ def test_find_component():
         result = reg.find(label="ion_chamber")
 
 
-def test_find_name_by_dot_notation():
-    # Prepare registry
-    reg = InstrumentRegistry()
+def test_find_name_by_dot_notation(reg):
     # Create a simulated component
     cptA = sim.SynGauss(
         "I0",
@@ -135,9 +130,7 @@ def test_find_name_by_dot_notation():
     assert result is cptA.val
 
 
-def test_find_labels_by_dot_notation():
-    # Prepare registry
-    reg = InstrumentRegistry()
+def test_find_labels_by_dot_notation(reg):
     # Create a simulated component
     cptA = sim.SynGauss(
         "I0",
@@ -154,9 +147,7 @@ def test_find_labels_by_dot_notation():
     assert result is cptA.val
 
 
-def test_find_any():
-    # Prepare registry
-    reg = InstrumentRegistry()
+def test_find_any(reg):
     # Create an unregistered component
     cptA = sim.SynGauss(
         "I0",
@@ -206,10 +197,8 @@ def test_find_any():
     assert cptD not in result
 
 
-def test_find_by_device():
+def test_find_by_device(reg):
     """The registry should just return the device itself if that's what is passed."""
-    # Prepare registry
-    reg = InstrumentRegistry()
     # Register a component
     cptD = sim.SynGauss(
         "sample motor",
@@ -226,10 +215,8 @@ def test_find_by_device():
     assert result is cptD
 
 
-def test_find_by_list_of_names():
+def test_find_by_list_of_names(reg):
     """Will the findall() method handle lists of things to look up."""
-    # Prepare registry
-    reg = InstrumentRegistry()
     # Register a component
     cptA = sim.SynGauss(
         "sample motor A",
@@ -268,13 +255,12 @@ def test_find_by_list_of_names():
     assert cptC not in result
 
 
-def test_user_readback():
+def test_user_readback(reg):
     """Edge case where EpicsMotor.user_readback is named the same as the motor itself."""
-    registry = InstrumentRegistry()
     device = EpicsMotor("", name="epics_motor")
-    registry.register(device)
+    reg.register(device)
     # See if requesting the device.user_readback returns the proper signal
-    registry.find("epics_motor_user_readback")
+    reg["epics_motor_user_readback"]
 
 
 # -----------------------------------------------------------------------------

src/haven/tests/test_mono_ID_calibration_plan.py

@@ -66,7 +66,9 @@ def test_aligns_pitch(mono_motor, id_motor, ion_chamber, pitch2_motor):
     plan = mono_ID_calibration(
         energies=[8000], energy_motor=mono_motor, fit_model=fit_model
     )
-    messages = list(plan)
+    with pytest.warns(UserWarning), pytest.deprecated_call():
+        # Raises a warning because there's no data to fit
+        messages = list(plan)
     device_names = [getattr(m.obj, "name", None) for m in messages]
     assert "monochromator_pitch2" in device_names
 
@@ -76,7 +78,9 @@ def test_aligns_mono_energy(mono_motor, id_motor, ion_chamber, pitch2_motor):
     plan = mono_ID_calibration(
         energies=[8000], energy_motor=mono_motor, fit_model=fit_model
     )
-    messages = list(plan)
+    with pytest.warns(UserWarning), pytest.deprecated_call():
+        # Raises a warning because there's no data to fit
+        messages = list(plan)
     id_messages = [
         m for m in messages if getattr(m.obj, "name", "") == "energy_id_energy"
     ]