APS 4-ID POLAR bluesky-widgets example

aps_polar_example.py

from bluesky_widgets.models.auto_plot_builders import AutoPlotter
from bluesky_widgets.models.plot_builders import Lines
from bluesky_widgets.models.plot_specs import AxesSpec, FigureSpec
from bluesky_widgets.qt.figures import QtFigures
import databroker
import numpy as np


def xanes(monitor, detector):
    absorption = np.log(np.array(monitor) / np.array(detector)).reshape(-1, 4)
    return absorption.mean(axis=1)


def xmcd(monitor, detector):
    absorption = np.log(np.array(monitor) / np.array(detector)).reshape(-1, 4)
    return absorption[:, [0, 3]].mean(axis=1) - absorption[:, [1, 2]].mean(axis=1)


def downsampled(x):
    return np.array(x).reshape(-1, 4).mean(axis=1)


class AutoXanesPlot(AutoPlotter):
    def __init__(self, monitor="Ion Ch 4", detector="Ion Ch 5"):
        super().__init__()
        self._x_to_lines = {}  # map x variable to (xanes_lines, xmcd_lines)
        self._monitor = monitor
        self._detector = detector

    @property
    def monitor(self):
        return self._monitor

    @property
    def detector(self):
        return self._detector

    def handle_new_stream(self, run, stream_name):
        if stream_name == "primary":
            # Nothing to do for this stream.
            return
        # Detect x variable from hints in metadata.
        first_scanning_dimension = run.metadata["start"]["hints"]["dimensions"][0]
        scanning_fields, _ = first_scanning_dimension
        x = scanning_fields[0]
        # If we already have a figure for this x, reuse it (over-plot).
        try:
            (xanes_lines, xmcd_lines) = self._x_to_lines[x]
        except KeyError:
            # We don't have a figure for this x. Make one.
            xanes_axes = AxesSpec(x_label=x, title="XANES")
            xmcd_axes = AxesSpec(x_label=x, title="XMCD")
            figure = FigureSpec((xanes_axes, xmcd_axes), title="XANES and XMCD")
            # Set up objects that will select the approriate data and do the
            # desired transformation for plotting.
            xanes_lines = Lines(
                x=lambda primary: downsampled(primary[x]),
                ys=[lambda primary: xanes(primary[self._monitor], primary[self._detector])],
                axes=xanes_axes,
            )
            xmcd_lines = Lines(
                x=lambda primary: downsampled(primary[x]),
                ys=[lambda primary: xmcd(primary[self._monitor], primary[self._detector])],
                axes=xmcd_axes,
            )
            self._x_to_lines[x] = (xanes_lines, xmcd_lines)
            # Keep track of these plot builders to enable *removing* runs from them.
            self.plot_builders.append(xanes_lines)
            self.plot_builders.append(xmcd_lines)
            # Appending this figures list will trigger to view to show our new figure.
            self.figures.append(figure)
        # Add this Run to the figure.
        xanes_lines.add_run(run)
        xmcd_lines.add_run(run)


catalog = databroker.catalog["POLAR-example"]
model = AutoXanesPlot()
model.add_run(catalog[-3])

view = QtFigures(model.figures)
view.show()