Configuring custom MDAEngine with YAML files

[JAMoreno-Larios]

Hi,

Thank you for working in making Micro-Manager scripting better for everyone.
I'm building a custom hyperspectral microscope using pymmcore-plus and useq-schema to control a Liquid Crystal Tunable Filter and a monochromator via serial port. For that, I subclassed the MDAEngine to set up these devices through the metadata field.

Currently, I am defining my experiments with the following functions, depending on what I need:

# Define configuration function
def configure_experiment_illumination(exposure, ext_wavelengths, em_wavelengths,
                                      loops, z_positions, channel_names,
                                      lctf_device_name):
    """
    Generates a list of MDAEvents following the next proceedure.
    We define exposure, the emission LCTF wavelength, the dichroic
    mirror to use and the number of times we want to repeat the measurement.
    We should be able to get a time series for a given configuration.
    """
    # Since we just want to generate lists, no additional validation is needed.
    experiment = []
    for position in range(0, len(z_positions)):
        for channel in channel_names:
            for excitation in ext_wavelengths:
                for emission in em_wavelengths:
                    for timepoint in range(0, loops):
                        event = useq.MDAEvent(exposure=exposure,
                                              z_pos=z_positions[position],
                                              channel=channel,
                                              keep_shutter_open=True,
                                              metadata = {
                                                  'timepoint' : timepoint,
                                                  'monochromator' : excitation,
                                                  'lctf' : (lctf_device_name,
                                                            'Wavelength', str(emission)),
                                                  'z_pos_index' : position,
                                                  'z_pos' : z_positions[position],})
                        experiment.append(event)
    return experiment

def configure_experiment_sets(exposure,
                              mirror_ext_em_set,
                              loops, z_positions,
                              lctf_device_name):
    """
    Similar to configure_experiment_illumination, but we expect the channel
    names, an excitation wavelength, and emission wavelengths to be input
    as a list of tuples,
    i. e. sets = [(mirror1, ext_wav1, list_of_em_wavelengths1), 
                  (mirror2, ext_wav2, list_of_em_wavelengths2),
                  (mirror3, ext_wav3, list_of_em_wavelengths3),...]
    so on and so forth.
    """
    # Since we just want to generate lists, no additional validation is needed.
    experiment = []
    for position in range(0, len(z_positions)):
        for channel, excitation, em_wavelengths in mirror_ext_em_set:
            for emission in em_wavelengths:
                for timepoint in range(0, loops):
                    event = useq.MDAEvent(exposure=exposure,
                                          z_pos=z_positions[position],
                                          channel=channel,
                                          keep_shutter_open=True,
                                          metadata = {
                                              'timepoint' : timepoint,
                                              'monochromator' : excitation,
                                              'lctf' : (lctf_device_name,
                                                        'Wavelength', str(emission)),
                                              'z_pos_index' : position,
                                              'z_pos' : z_positions[position],})
                    experiment.append(event)
    return experiment

With my approach, I cannot define tiling, and if someone different than me wants to use my scripts, I fear they will not be able to set up experiments easily. I wish to be able to define the experiments using a YAML file, what would be a good starting point?

Thanks,

[tlambert03]

I'm building a custom hyperspectral microscope using pymmcore-plus and useq-schema to control a Liquid Crystal Tunable Filter and a monochromator via serial port

fun! :)

With my approach, I cannot define tiling, and if someone different than me wants to use my scripts, I fear they will not be able to set up experiments easily. I wish to be able to define the experiments using a YAML file, what would be a good starting point?

yeah I agree... let's see whether we can do better using MDASequence instead of manually building the list of MDAEvents.

I think what I would do in this scenario is try to put as much info into the "channel" object as possible, and then deal with the idiosyncrasies of your spectral setup in your custom engine. So for example your configure_experiment_illumination definition is quite close to this MDASequence:

channel_names = ["A", "B",]
seq = useq.MDASequence(
    channels=[useq.Channel(config=ch, exposure=50) for ch in channel_names],
    z_plan=useq.ZRangeAround(range=3, step=1),
    time_plan=useq.TIntervalLoops(interval=0, loops=10),
    keep_shutter_open_across=("z",),
    metadata={"lctf": {"device": "YourLCTFDeviceName"}},
)

(note that that expands to something very similar to your list of MDAEvents):

>>> list(seq)
[
    MDAEvent(index={'t': 0, 'c': 0, 'z': 0}, channel=Channel(config='A'), exposure=50.0, min_start_time=0.0, z_pos=-1.5, keep_shutter_open=True, reset_event_timer=True),
    MDAEvent(index={'t': 0, 'c': 0, 'z': 1}, channel=Channel(config='A'), exposure=50.0, min_start_time=0.0, z_pos=-0.5, keep_shutter_open=True),
    MDAEvent(index={'t': 0, 'c': 0, 'z': 2}, channel=Channel(config='A'), exposure=50.0, min_start_time=0.0, z_pos=0.5, keep_shutter_open=True),
    MDAEvent(index={'t': 0, 'c': 0, 'z': 3}, channel=Channel(config='A'), exposure=50.0, min_start_time=0.0, z_pos=1.5),
    MDAEvent(index={'t': 0, 'c': 1, 'z': 0}, channel=Channel(config='B'), exposure=50.0, min_start_time=0.0, z_pos=-1.5, keep_shutter_open=True),
    MDAEvent(index={'t': 0, 'c': 1, 'z': 1}, channel=Channel(config='B'), exposure=50.0, min_start_time=0.0, z_pos=-0.5, keep_shutter_open=True),
    MDAEvent(index={'t': 0, 'c': 1, 'z': 2}, channel=Channel(config='B'), exposure=50.0, min_start_time=0.0, z_pos=0.5, keep_shutter_open=True),
    MDAEvent(index={'t': 0, 'c': 1, 'z': 3}, channel=Channel(config='B'), exposure=50.0, min_start_time=0.0, z_pos=1.5),
    ...

and that MDASequence object dumps quite nicely to YAML:

>>> print(seq.yaml())
channels:
- config: A
  exposure: 50.0
- config: B
  exposure: 50.0
keep_shutter_open_across:
- z
metadata:
  lctf:
    device: YourLCTFDeviceName
time_plan:
  interval: 0.0
  loops: 10
z_plan:
  range: 3.0
  step: 1.0

now, to decide how to handle spectral... I need to understand your hardware setup a bit more... From your code it looks like you're scanning both excitation and emission? (i.e. collect each em_wavelength at every excitation wavelength?). I think there are two options:

1. if you can possibly express this using standard micromanager channel configurations, such that each combination of ex/em can be expressed as a micromanager configuration group, and activated using core.setConfig(groupName, configName)... then that's great, and it means we don't have to do as much on the custom engine side.
2. On the other hand, if the LCTF device can't be controlled easily by micro-manager, and you must set up your hardware using python-side logic... then I would try to do as much of it as possible in your custom MDAEngine, probably by overriding the _set_event_channel method (which we could make public if you want. That would then do all the work of converting your event Channel object into hardware setup. As a slightly hacky solution... you could even use the Channel.config (channel_name) to store something like ex450_m550 and then parse that string in your Engine and set up what you need.

---

I expect you'll have follow up questions... but that would be my initial advice: Try to express it directly as an MDASequence, rather than a list of MDAEvents, combine that with your custom MDAEngine (which we can also discuss), and then let's dig into ways in which expressing it as an MDASequence feels insufficient. If we can't figure out a nice solution with the current useq model, we can discuss ways in which it could be improved.

[JAMoreno-Larios]

Thank you for answering my question so quickly!

now, to decide how to handle spectral... I need to understand your hardware setup a bit more... From your code it looks like you're scanning both excitation and emission?

I do need to scan both excitation and emission. Our monochromator ranges from 250 to 650nm and our LCTF ranges from 420 to 730nm, that alone would balloon the number of channels that need to be defined in a normal micromanager configuration group. Currently, I'm defining the intervals with np.linspace(...)

I overloaded MDAEngine's contructor and setup_single_event(...) as follows:

class HyperspectralEngine(MDAEngine):
    """
    Custon MultiDimensional Acquisiton Engine for pymmcore-plus.
    We add logic to control the attached LCTF and monochromator.
    """
    def __init__(self, mmc: CMMCorePlus, 
                 monochromator : Monochromator,
                 use_hardware_sequencing: bool = False) -> None:
        """
        Inheriting from MDAEngine, we extend to include an instance
        of our monochromator control class
        """
        super().__init__(mmc, use_hardware_sequencing) #MDAEngine init
        self._monochromator = monochromator
        
    def setup_single_event(self, event: useq.MDAEvent) -> None:
        """
        Overriding to include LCFT control
        
        """
        if event.keep_shutter_open:
            ...

        if event.x_pos is not None or event.y_pos is not None:
            self._set_event_position(event)
        if event.z_pos is not None:
            self._set_event_z(event)

        if event.channel is not None:
            try:
                self._mmc.setConfig(event.channel.group, event.channel.config)
            except Exception as e:
                logger.warning("Failed to set channel. %s", e)
        if event.exposure is not None:
            try:
                self._mmc.setExposure(event.exposure)
            except Exception as e:
                logger.warning("Failed to set exposure. %s", e)

        if (
            # (if autoshutter wasn't set at the beginning of the sequence
            # then it never matters...)
            self._autoshutter_was_set
            # if we want to leave the shutter open after this event, and autoshutter
            # is currently enabled...
            and event.keep_shutter_open
            and self._mmc.getAutoShutter()
        ):
            # we have to disable autoshutter and open the shutter
            self._mmc.setAutoShutter(False)
            self._mmc.setShutterOpen(True)
        
        # We set the LCTF
        if 'lctf' in event.metadata:
            try:
                self._mmc.setProperty(*event.metadata['lctf'])
            except Exception as e:
                logger.warning("Failed to set LCTF. %s", e)
                
        if 'monochromator' in event.metadata:
            try:
                self._monochromator.set_wavelength(event.metadata['monochromator'])
            except Exception as e:
                logger.warning("Failed to set monochromator. %s", e)
                

    def exec_single_event(self, event: useq.MDAEvent) -> Iterator[PImagePayload]:
        """Execute a single (non-triggered) event and return the image data.

        This method is not part of the PMDAEngine protocol (it is called by
        `exec_event`, which *is* part of the protocol), but it is made public
        in case a user wants to subclass this engine and override this method.
        """
        try:
            self._mmc.snapImage()
        except Exception as e:
            logger.warning("Failed to snap image. %s", e)
            return ()
        if not event.keep_shutter_open:
            self._mmc.setShutterOpen(False)
        yield ImagePayload(self._mmc.getImage(), event, self._mmc.getTags())

I agree that the first step would be to be able to generate my experiments as an MDASequence. For defining the excitation and emission ranges, I could look into how you define the time or z plans for inspiration.

Thanks,

Edit:
What do you think about modifying the Axis class defined in _util.py? Essentially we are treating the excitation and emission wavelengths as additional axis in our workflow. That would also entail that I need to modify the iterator to handle custom axis, right?

[JAMoreno-Larios]

A little update on what I've been working:
I decided to modify how MDASequence is generated by declaring my own Hyperspectral Channel and Iterator Subclasses, as well as modifying _iter_sequence, too.

Hyperspectral Channel subclass and iterator:

class HyperspectralChannel(Channel):
...
    # Define new variables
    lctf_val: Optional[int] = Field(None, ge=420, le=730)
    monochromator_val: Optional[int] = Field(None, ge=250, le=650)


class HyperChannelPlan(Channel):
    """
    Subclass of Channel, that defines an acquisition channel while specifying
    how to use an iterable component. In our case, it is the LCTF and
    the monochromator.

  ...

    # Define new variables
    # lctf_device_name: str
    lctf_start: int
    lctf_stop: Optional[int] = Field(0)
    lctf_step: Optional[int] = Field(1)
    monochromator_start: int
    monochromator_stop: Optional[int] = Field(0)
    monochromator_step: Optional[int] = Field(1)

    # Define the iterator
    def __iter__(self) -> Iterator[HyperspectralChannel]:  # type: ignore
        hyperChannels = self.generate_spectral_plan()
        yield from hyperChannels

    def generate_range(self, start, stop, step):
        """
        Aux method to generate a sequence of values
        """
        if stop == 0:
            return [start]
        stop = stop + step  # make sure we include the last point
        return [int(x) for x in range(start, stop, step)]

    def generate_spectral_plan(self):
        """
        Helper function to generate the LCTF values to use.
        """
        # Define ranges - can add logic to provide a list of points
        lctf_range = self.generate_range(self.lctf_start,
                                         self.lctf_stop,
                                         self.lctf_step)
        monochromator_range = self.generate_range(self.monochromator_start,
                                                  self.monochromator_stop,
                                                  self.monochromator_step)
        # Set up our generator
        hyperChannels = [
            HyperspectralChannel(config=self.config,
                                 group=self.group,
                                 exposure=self.exposure,
                                 do_stack=self.do_stack,
                                 z_offset=self.z_offset,
                                 acquire_every=self.acquire_every,
                                 camera=self.camera,
                                 lctf_val=int(x),
                                 monochromator_val=int(y)) for y in
            monochromator_range for x in lctf_range
        ]

        return hyperChannels

Custom _iter_sequence.py

def _iter_sequence(
    sequence: MDASequence,
    *,
    base_event_kwargs: MDAEventDict | None = None,
    event_kwarg_overrides: MDAEventDict | None = None,
    position_offsets: PositionDict | None = None,
    _last_t_idx: int = -1,
) -> Iterator[MDAEvent]:
    order = _used_axes(sequence)
...
        if channel:
            event_kwargs["channel"] = EventChannel.model_construct(
                config=channel.config, group=channel.group
            )

            if channel.exposure is not None:
                event_kwargs["exposure"] = channel.exposure
            # We add our custom data here
            # We verify if the metadata dict is already set, if not, create it
            if "metadata" not in event_kwargs:
                event_kwargs.update(metadata=dict())
            if channel.lctf_val is not None:
                event_kwargs["metadata"].update(lctf=channel.lctf_val)
            if channel.monochromator_val is not None:
                event_kwargs["metadata"].update(
                    monochromator=channel.monochromator_val)
...

Custom MDASequence

class HyperMDASequence(MDASequence):

    metadata: dict[str, Any] = Field(default_factory=dict)
    axis_order: tuple[str, ...] = AXES
    # note that these are BOTH just `Sequence[Position]` but we retain the distinction
    # here so that WellPlatePlans are preserved in the model instance.
    stage_positions: Union[WellPlatePlan, tuple[Position, ...]] = Field(  # type: ignore
        default_factory=tuple, union_mode="left_to_right"
    )
    grid_plan: Optional[MultiPointPlan] = Field(
        default=None, union_mode="left_to_right"
    )
    channels: Union[HyperChannelPlan,
                    tuple[Channel, ...]] = Field(default_factory=tuple)
...
    @field_validator("channels", mode="before")
    def _validate_channels(cls, value: Any) -> Union[HyperChannelPlan, tuple[Channel, ...]]:
        if isinstance(value, HyperChannelPlan):
            value = list(value)  # Expand our iterator
        if isinstance(value, str) or not isinstance(
                value, Sequence) or not isinstance:  # pragma: no cover
            raise ValueError(f"channels must be a sequence, got {type(value)}")
        channels = []
        for v in value:
            if isinstance(v, Channel):
                channels.append(v)
            elif isinstance(v, str):
                channels.append(Channel.model_construct(config=v))
            elif isinstance(v, dict):
                channels.append(Channel(**v))
            else:  # pragma: no cover
                raise ValueError(f"Invalid Channel argument: {value!r}")
        return tuple(channels)
...

With this I can generate all the events using a MDASequence subclass, however, my script generates lots of channels and I haven't work on their serialization. Maybe instead of using a HyperspectralChannel generator, I can work on the iterators so that they read the appropiate LCTF and monochromator start, stop, and step parameters from a YAML file and generate the MDAEvents.

[tlambert03]

hey @JAMoreno-Larios

Sorry, for the delay. I'd say you're roughly thinking in the right direction. At this point, creating customized runtime variants of the schema models that are also themselves easily serializable is still a challenge. I do think that modify _iter_sequence itself is probably something you would ideally avoid, as it's very likely to result in something that diverges or breaks in the future.

I will briefly mention that I have what is effectively a rewrite of the core model that would probably make this use case a little bit easier (i.e. it would be easier for you to simply define a new pydantic model, and define how it should be iterated over), but it's still a work in progress. #192 I'll let you know if i make progress on that.

Until then, i do think you should try hard not to define ex/em as new axes (since they aren't really new axes...). It is worth pointing out that you could indeed probably express this with the current system, it's just that your YAML would be very long. You're looking for a more concise way to express many channels with just a few parameters, which definitely makes sense, it's just not something that the schema can currently express with brevity. So, it does look like you'll currently need to decide where to make the tradeoff: either on the brevity of the your YAML file, or the complexity of the runtime handling. If possible, I'd go with the ugly yaml file that sticks with the same runtime logic... but i get the desire to have it more concise :)

[JAMoreno-Larios]

Hi @tlambert03

I went and further modifyed _iter_sequence to be able to expand the given HyperChannelPlan class I defined.
Here's a brief snippet at the beginning of _iter_sequence:

def _iter_sequence(
    sequence: MDASequence,
    *,
    base_event_kwargs: MDAEventDict | None = None,
    event_kwarg_overrides: MDAEventDict | None = None,
    position_offsets: PositionDict | None = None,
    _last_t_idx: int = -1,
) -> Iterator[MDAEvent]:
 
    # Let's expand the channels first
    expanded_channels = []
    for channelPlan in sequence.channels:
        expanded_list = list(channelPlan)
        for chan in expanded_list:
            expanded_channels.append(chan)

    # Now let's update our sequence
    # sequence.channels = tuple(expanded_channels)
    sequence = sequence.model_copy(
        update={'channels': tuple(expanded_channels)}
    )

    # We continue with the iterator as usual
    order = _used_axes(sequence)
...

At this moment, my control scripts still need some user input, but I think I managed to streamline how to configure new experiments:

# User and experiment ID
user_name = "jaml"  # User initials
sample_name = "rings40x"  # Sample name
metadata_filename = "metadata.json"

# Set acquisition axis order
axis_order = "tpcz"  # Rightmost position changes faster

# Camera settings
cam_exposure = 1000  # Exposure in miliseconds

# Set positions here
# VERIFY THAT THE POSITIONS ARE VALID UNDER uMANAGER
# IF CARELESS, YOU CAN DESTROY THE OBJECTIVE
x_pos = None  # If None, use what we have in uManager
y_pos = None  # If None, use what we have in uManager
z_pos = 7000  # If None, use what we have in uManager. Needs to be set if
# using a z-plan

# Configure z-plan
do_zstack = True
nyquist_step = 0.81  # Depends on experiment
z_range = 10 * nyquist_step
z_step = nyquist_step

# Time plan
time_interval = 0.0  # Time between time points in seconds
time_loops = 5  # Number of repetitions

# Define channels, lctf and monochromator start, stop and step
channel_names = [
    'mirror_458FF',
    'mirror_517Longpass',
    'mirror_543UltraLongpass',
    'mirror_562FF',
    'mirror_642UltraLongpass'
]
monochromator_starts = [
    420, 480, 520, 540, 630
]

lctf_starts = [460, 520, 540, 560, 650]
lctf_stops = [730, 730, 730, 730, 730]
lctf_steps = [10, 10, 10, 10, 10]

# Initialize channel objects with HyperChannelPlan
channels = [
    frainbow_microscope.HyperChannelPlan(
        config=channel_name,
        group='Channel',
        exposure=cam_exposure,
        do_stack=do_zstack, z_offset=0,
        lctf_start=lctf_start,
        lctf_stop=lctf_stop,
        lctf_step=lctf_step,
        monochromator_start=monochromator_start,
    )
    for channel_name, lctf_start, lctf_stop, lctf_step,
    monochromator_start, in zip(channel_names, lctf_starts, lctf_stops,
                                lctf_steps, monochromator_starts)]

# Create experimental HyperMDASequence
experiment = frainbow_microscope.HyperMDASequence(
    axis_order=axis_order,  # Rightmost position changes faster
    # keep_shutter_open_across=('t', 'p', 'c', 'z'),
    # axis_order="tgcz",
    time_plan={"interval": time_interval, "loops": time_loops},
    stage_positions=[(x_pos, y_pos, z_pos)],  # Required, otherwise it breaks.
    # grid_plan={"rows": 2, "columns": 2},
    z_plan={"range": z_range, "step": z_step},
    channels=channels)
```

It certainly beats having to write a MDAEvent iterator from scratch!

[tlambert03]

well, I'm encouraged that you seem to be well-equipped to come up with some workarounds on your own! :) sorry it's not a bit easier at the moment. I'll definitely keep this pattern in mind as I continue to work on #iterations

[JAMoreno-Larios]

Now it's matter of people doing more ambitious experiments with our custom setup. I hope I can clean up the code and eventually release these scripts to the public. :D

I appreciate your work and assistance in this matter, I think I'll close this issue for the time being.