Coverage for src/pytribeam/workflow.py: 0%
186 statements
« prev ^ index » next coverage.py v7.6.1, created at 2026-06-16 18:30 +0000
1#!/usr/bin/python3
2"""
3Workflow Module
4===============
6This module contains functions for managing and executing the workflow of an experiment, including performing operations, setting up the experiment, and running the main experiment loop.
8Functions
9---------
10perform_operation(step_settings, step: tbt.Step, general_settings: tbt.GeneralSettings, slice_number: int) -> bool
11 Perform the operation for the specified step settings.
13perform_operation(step_settings: tbt.ImageSettings, step: tbt.Step, general_settings: tbt.GeneralSettings, slice_number: int) -> bool
14 Perform the image operation for the specified step settings.
16perform_operation(step_settings: tbt.FIBSettings, step: tbt.Step, general_settings: tbt.GeneralSettings, slice_number: int) -> bool
17 Perform the FIB operation for the specified step settings.
19perform_operation(step_settings: tbt.CustomSettings, step: tbt.Step, general_settings: tbt.GeneralSettings, slice_number: int) -> bool
20 Perform the custom operation for the specified step settings.
22perform_operation(step_settings: tbt.EBSDSettings, step: tbt.Step, general_settings: tbt.GeneralSettings, slice_number: int) -> bool
23 Perform the EBSD operation for the specified step settings.
25perform_operation(step_settings: tbt.EDSSettings, step: tbt.Step, general_settings: tbt.GeneralSettings, slice_number: int) -> bool
26 Perform the EDS operation for the specified step settings.
28perform_operation(step_settings: tbt.LaserSettings, step: tbt.Step, general_settings: tbt.GeneralSettings, slice_number: int) -> bool
29 Perform the laser operation for the specified step settings.
31ebsd_eds_conflict_free(step_sequence: List[tbt.Step]) -> bool
32 Check if the step sequence is free of EBSD and EDS conflicts.
34pre_flight_check(yml_path: Path) -> tbt.ExperimentSettings
35 Perform a pre-flight check for the experiment.
37setup_experiment(yml_path: Path) -> tbt.ExperimentSettings
38 Set up the experiment based on the YAML configuration.
40perform_step(slice_number: int, step_number: int, experiment_settings: tbt.ExperimentSettings) -> bool
41 Perform a step in the experiment.
43run_experiment_cli(start_slice: int, start_step: int, yml_path: Path)
44 Main loop for the experiment, accessed through the command line.
45"""
47# Default python modules
48# from functools import singledispatch
49from pathlib import Path
50from typing import List
51from functools import singledispatch
52import subprocess
54# 3rd party module
56# Local scripts
57import pytribeam.constants as cs
58import pytribeam.insertable_devices as devices
59import pytribeam.factory as factory
60import pytribeam.types as tbt
61import pytribeam.utilities as ut
62import pytribeam.stage as stage
63import pytribeam.log as log
64import pytribeam.laser as laser
65import pytribeam.image as img
66import pytribeam.fib as fib
69@singledispatch
70def perform_operation(
71 step_settings,
72 step: tbt.Step,
73 general_settings: tbt.GeneralSettings,
74 slice_number: int,
75) -> bool:
76 """
77 Perform the operation for the specified step settings.
79 This function performs the operation for the specified step settings, including validation.
81 Parameters
82 ----------
83 step_settings : Any
84 The step settings for the operation.
85 step : tbt.Step
86 The step object containing the operation settings.
87 general_settings : tbt.GeneralSettings
88 The general settings object.
89 slice_number : int
90 The slice number for the operation.
92 Returns
93 -------
94 bool
95 True if the operation is performed successfully.
97 Raises
98 ------
99 NotImplementedError
100 If no handler is available for the provided step settings type.
101 """
102 _ = step_settings
103 __ = step
104 ___ = general_settings
105 ____ = slice_number
106 raise NotImplementedError(f"No handler for type {type(step_settings)}")
109@perform_operation.register
110def _(
111 step_settings: tbt.ImageSettings,
112 step: tbt.Step,
113 general_settings: tbt.GeneralSettings,
114 slice_number: int,
115) -> bool:
116 """
117 Perform the image operation for the specified step settings.
119 Parameters
120 ----------
121 step_settings : tbt.ImageSettings
122 The image settings for the operation.
123 step : tbt.Step
124 The step object containing the operation settings.
125 general_settings : tbt.GeneralSettings
126 The general settings object.
127 slice_number : int
128 The slice number for the operation.
130 Returns
131 -------
132 bool
133 True if the image operation is performed successfully.
134 """
135 return img.image_operation(
136 step=step,
137 image_settings=step.operation_settings,
138 general_settings=general_settings,
139 slice_number=slice_number,
140 )
143@perform_operation.register
144def _(
145 step_settings: tbt.FIBSettings,
146 step: tbt.Step,
147 general_settings: tbt.GeneralSettings,
148 slice_number: int,
149) -> bool:
150 """
151 Perform the FIB operation for the specified step settings.
153 Parameters
154 ----------
155 step_settings : tbt.FIBSettings
156 The FIB settings for the operation.
157 step : tbt.Step
158 The step object containing the operation settings.
159 general_settings : tbt.GeneralSettings
160 The general settings object.
161 slice_number : int
162 The slice number for the operation.
164 Returns
165 -------
166 bool
167 True if the FIB operation is performed successfully.
168 """
169 # collect image
170 image_step = tbt.Step(
171 type=tbt.StepType.IMAGE,
172 name=step.name,
173 number=step.number,
174 frequency=step.frequency,
175 stage=step.stage,
176 operation_settings=step_settings.image,
177 )
178 type(image_step.operation_settings)
179 perform_operation(
180 image_step.operation_settings,
181 step=image_step,
182 general_settings=general_settings,
183 slice_number=slice_number,
184 )
185 # mill pattern
186 fib.mill_operation(
187 step=step,
188 fib_settings=step_settings,
189 general_settings=general_settings,
190 slice_number=slice_number,
191 )
193 return True
196@perform_operation.register
197def _(
198 step_settings: tbt.CustomSettings,
199 step: tbt.Step,
200 general_settings: tbt.GeneralSettings,
201 slice_number: int,
202) -> bool:
203 """
204 Perform the custom operation for the specified step settings.
206 Parameters
207 ----------
208 step_settings : tbt.CustomSettings
209 The custom settings for the operation.
210 step : tbt.Step
211 The step object containing the operation settings.
212 general_settings : tbt.GeneralSettings
213 The general settings object.
214 slice_number : int
215 The slice number for the operation.
217 Returns
218 -------
219 bool
220 True if the custom operation is performed successfully.
221 """
222 # dump out .yml with experiment info
223 slice_info_path = Path.joinpath(general_settings.exp_dir, "slice_info.yml")
224 db = {"exp_dir": str(general_settings.exp_dir), "slice_number": slice_number}
225 ut.dict_to_yml(db=db, file_path=slice_info_path)
227 output = subprocess.run(
228 [step_settings.executable_path, step_settings.script_path],
229 capture_output=True,
230 )
231 stdout, stderr = output.stdout.decode("utf-8"), output.stderr.decode("utf-8")
232 if stdout:
233 print(f"\nCustom script output: {stdout}\n")
235 if output.returncode != 0:
236 if stderr:
237 print(f"\nCustom script errors: {stderr}\n")
238 raise ValueError(
239 f"Subprocess call for script {step_settings.script_path} using executable {step_settings.executable_path} did not execute correctly."
240 )
242 slice_info_path.unlink()
243 return True
246@perform_operation.register
247def _(
248 step_settings: tbt.EBSDSettings,
249 step: tbt.Step,
250 general_settings: tbt.GeneralSettings,
251 slice_number: int,
252) -> bool:
253 """
254 Perform the EBSD operation for the specified step settings.
256 Parameters
257 ----------
258 step_settings : tbt.EBSDSettings
259 The EBSD settings for the operation.
260 step : tbt.Step
261 The step object containing the operation settings.
262 general_settings : tbt.GeneralSettings
263 The general settings object.
264 slice_number : int
265 The slice number for the operation.
267 Returns
268 -------
269 bool
270 True if the EBSD operation is performed successfully.
271 """
272 image_settings = step_settings.image
273 microscope = image_settings.microscope
275 # insert detector
276 devices.insert_EBSD(microscope=microscope)
277 if step_settings.enable_eds:
278 devices.insert_EDS(microscope=microscope)
280 # measure and log specimen current
281 found_current_na = devices.specimen_current(microscope=microscope)
282 log.specimen_current(
283 step_number=step.number,
284 step_name=step.name,
285 slice_number=slice_number,
286 log_filepath=general_settings.log_filepath,
287 dataset_name=cs.Constants.specimen_current_dataset_name,
288 specimen_current_na=found_current_na,
289 )
291 # take image
292 img.image_operation(
293 step=step,
294 image_settings=image_settings,
295 general_settings=general_settings,
296 slice_number=slice_number,
297 )
299 # set dynamic focus/tilt correction
300 dynamic_focus = image_settings.beam.settings.dynamic_focus
301 tilt_correction = image_settings.beam.settings.tilt_correction
302 img.beam_angular_correction(
303 microscope=microscope,
304 dynamic_focus=dynamic_focus,
305 tilt_correction=tilt_correction,
306 )
308 # take map
309 laser.map_ebsd()
311 # retract detector(s)
312 devices.retract_EBSD(microscope=microscope)
313 if step_settings.enable_eds:
314 devices.retract_EDS(microscope=microscope)
316 return True
319@perform_operation.register
320def _(
321 step_settings: tbt.EDSSettings,
322 step: tbt.Step,
323 general_settings: tbt.GeneralSettings,
324 slice_number: int,
325) -> bool:
326 """
327 Perform the EDS operation for the specified step settings.
329 Parameters
330 ----------
331 step_settings : tbt.EDSSettings
332 The EDS settings for the operation.
333 step : tbt.Step
334 The step object containing the operation settings.
335 general_settings : tbt.GeneralSettings
336 The general settings object.
337 slice_number : int
338 The slice number for the operation.
340 Returns
341 -------
342 bool
343 True if the EDS operation is performed successfully.
344 """
345 image_settings = step_settings.image
346 microscope = image_settings.microscope
348 # insert detector
349 devices.insert_EDS(microscope=microscope)
351 # measure and log specimen current
352 found_current_na = devices.specimen_current(microscope=microscope)
353 log.specimen_current(
354 step_number=step.number,
355 step_name=step.name,
356 slice_number=slice_number,
357 log_filepath=general_settings.log_filepath,
358 dataset_name=cs.Constants.specimen_current_dataset_name,
359 specimen_current_na=found_current_na,
360 )
362 # take image
363 img.image_operation(
364 step=step,
365 image_settings=image_settings,
366 general_settings=general_settings,
367 slice_number=slice_number,
368 )
370 # set dynamic focus/tilt correction
371 dynamic_focus = image_settings.beam.settings.dynamic_focus
372 tilt_correction = image_settings.beam.settings.tilt_correction
373 img.beam_angular_correction(
374 microscope=microscope,
375 dynamic_focus=dynamic_focus,
376 tilt_correction=tilt_correction,
377 )
379 # take map
380 laser.map_eds()
382 # retract detector
383 devices.retract_EDS(microscope=microscope)
385 return True
388@perform_operation.register
389def _(
390 step_settings: tbt.LaserSettings,
391 step: tbt.Step,
392 general_settings: tbt.GeneralSettings,
393 slice_number: int,
394) -> bool:
395 """
396 Perform the laser operation for the specified step settings.
398 Parameters
399 ----------
400 step_settings : tbt.LaserSettings
401 The laser settings for the operation.
402 step : tbt.Step
403 The step object containing the operation settings.
404 general_settings : tbt.GeneralSettings
405 The general settings object.
406 slice_number : int
407 The slice number for the operation.
409 Returns
410 -------
411 bool
412 True if the laser operation is performed successfully.
413 """
414 return laser.laser_operation(
415 step=step,
416 general_settings=general_settings,
417 slice_number=slice_number,
418 )
421# @perform_operation(tbt.StageSettings)
422# def _(
423# step_settings,
424# step: tbt.Step,
425# log_filepath: Path,
426# ) -> bool:
427# pass
430def ebsd_eds_conflict_free(step_sequence: List[tbt.Step]) -> bool:
431 """
432 Check if the step sequence is free of EBSD and EDS conflicts.
434 This function checks if the step sequence is free of EBSD and EDS conflicts.
436 Parameters
437 ----------
438 step_sequence : List[tbt.Step]
439 The step sequence to check.
441 Returns
442 -------
443 bool
444 True if the step sequence is free of EBSD and EDS conflicts.
446 Raises
447 ------
448 ValueError
449 If an EBSD or EDS conflict is found in the step sequence.
450 """
451 EBSD_EDS_conflict_msg = "Due to current limitations in 3rd party EBSD/EDS integration with the TriBeam, only one of these step types is allowed as only one map can be configured for an experiment, but EDS can be configured to be included with an EBSD type step. See User Guide for more details."
453 found_EBSD = False
454 found_EDS = False
456 for step in step_sequence:
457 if step.type == tbt.StepType.EBSD:
458 if found_EDS == True:
459 raise ValueError(
460 f"EBSD step found in sequence after EDS step was already defined. {EBSD_EDS_conflict_msg}"
461 )
462 found_EBSD = True
464 if step.type == tbt.StepType.EDS:
465 if found_EBSD == True:
466 raise ValueError(
467 f"EDS step found in sequence after EBSD step was already defined. {EBSD_EDS_conflict_msg}"
468 )
469 found_EDS = True
471 return True
474def pre_flight_check(yml_path: Path) -> tbt.ExperimentSettings:
475 """
476 Perform a pre-flight check for the experiment.
478 This function performs a pre-flight check for the experiment by validating the YAML configuration, connecting to the microscope, and validating the step sequence.
480 Parameters
481 ----------
482 yml_path : Path
483 The path to the YAML configuration file.
485 Returns
486 -------
487 tbt.ExperimentSettings
488 The validated experiment settings.
490 Raises
491 ------
492 SystemError
493 If there are issues with the EBSD or EDS camera, or if the laser control is not enabled.
494 ValueError
495 If the step sequence is not parsed correctly or if there are EBSD/EDS conflicts.
496 """
497 # get configuration from yml
498 yml_version = ut.yml_version(yml_path)
499 experiment_settings = ut.yml_to_dict(
500 yml_path_file=yml_path,
501 version=yml_version,
502 required_keys=(
503 "general",
504 "config_file_version",
505 ),
506 )
507 yml_format = ut.yml_format(version=yml_version)
509 # get general settings and validate them
510 general_db = ut.general_settings(
511 exp_settings=experiment_settings, yml_format=yml_format
512 )
513 general_settings = factory.general(
514 general_db=general_db,
515 yml_format=yml_format,
516 )
518 # whether to enable EBSD and EDS control
519 enable_EBSD = ut.enable_external_device(general_settings.EBSD_OEM)
520 enable_EDS = ut.enable_external_device(general_settings.EDS_OEM)
521 if enable_EBSD:
522 status = devices.connect_EBSD()
523 if status == tbt.RetractableDeviceState.ERROR:
524 raise SystemError("EBSD camera is connected but in error state.")
525 if enable_EDS:
526 status = devices.connect_EDS()
527 if status == tbt.RetractableDeviceState.ERROR:
528 raise SystemError("EDS camera is connected but in error state.")
530 # connect to microscope:
531 connection = general_settings.connection
532 microscope = tbt.Microscope()
533 ut.connect_microscope(
534 microscope=microscope,
535 quiet_output=True,
536 connection_host=connection.host,
537 connection_port=connection.port,
538 )
540 # get step_count and validate settings
541 num_steps = ut.step_count(exp_settings=experiment_settings, yml_format=yml_format)
542 step_sequence = [] # empty list of tbt.Step type objects
543 for step in range(1, num_steps + 1):
544 step_name, step_settings = ut.step_settings(
545 exp_settings=experiment_settings,
546 step_number_key=yml_format.step_number_key,
547 step_number_val=step,
548 yml_format=yml_format,
549 )
550 if not step_name:
551 raise KeyError(
552 f"Step name for step {step} of {num_steps} is empty. Please provide a unique name for each step in your configuration."
553 )
554 step_type = ut.step_type(
555 settings=step_settings,
556 yml_format=yml_format,
557 )
559 # validate connections for specific step types
560 if step_type == tbt.StepType.LASER:
561 laser_enabled = laser.laser_connected()
562 if not laser_enabled:
563 raise SystemError(
564 f"Step name '{step_name}' is a Laser step type but Laser control is not currently enabled. Ensure TFS laser API is installed, Laser Control application is open."
565 )
566 if (step_type == tbt.StepType.EDS) and (not enable_EDS):
567 raise SystemError(
568 f"Step name '{step_name}' is an EDS step type but EDS control is not currently enabled."
569 )
570 if (step_type == tbt.StepType.EBSD) and (not enable_EBSD):
571 raise SystemError(
572 f"Step name '{step_name}' is an EDS step type but EDS control is not currently enabled."
573 )
574 # if (step_type == tbt.StepType.EBSD_EDS) and (
575 # (not enable_EBSD) or (not enable_EDS)
576 # ):
577 # raise SystemError(
578 # f"Step name '{step_name}' is an EBSD_EDS step type but EBSD and EDS control are not both currently enabled."
579 # )
580 # create the step settings
581 step = factory.step(
582 microscope=microscope,
583 step_name=step_name,
584 step_settings=step_settings,
585 general_settings=general_settings,
586 yml_format=yml_format,
587 )
589 step_sequence.append(step)
591 if len(step_sequence) != num_steps:
592 raise ValueError(
593 f"Settings not parsed correctly, expected {num_steps} but only {len(step_sequence)} have been parsed."
594 )
596 # ensure only EBSD or EDS step type exists
597 ebsd_eds_conflict_free(step_sequence=step_sequence)
599 experiment_settings = tbt.ExperimentSettings(
600 microscope=microscope,
601 general_settings=general_settings,
602 step_sequence=step_sequence,
603 enable_EBSD=enable_EBSD,
604 enable_EDS=enable_EDS,
605 )
606 # print("Pre-flight check complete.")
607 return experiment_settings
610def setup_experiment(
611 yml_path: Path,
612) -> tbt.ExperimentSettings:
613 """
614 Set up the experiment based on the YAML configuration.
616 This function sets up the experiment by validating the YAML configuration, creating the log file, linking the stage, and retracting all devices.
618 Parameters
619 ----------
620 yml_path : Path
621 The path to the YAML configuration file.
623 Returns
624 -------
625 tbt.ExperimentSettings
626 The experiment settings.
627 """
628 # validate yml
629 experiment_settings = pre_flight_check(yml_path=yml_path)
631 log_filepath = experiment_settings.general_settings.log_filepath
632 log.create_file(log_filepath)
634 # link stage to free working distance
635 experiment_settings.microscope.specimen.stage.link()
637 # retract all devices
638 print("\tRetracting all devices...")
639 devices.retract_all_devices(
640 microscope=experiment_settings.microscope,
641 enable_EBSD=experiment_settings.enable_EBSD,
642 enable_EDS=experiment_settings.enable_EDS,
643 )
645 return experiment_settings
648def perform_step(
649 slice_number: int,
650 step_number: int,
651 experiment_settings: tbt.ExperimentSettings,
652):
653 """
654 Perform a step in the experiment.
656 This function performs a step in the experiment based on the slice number, step number, and experiment settings.
658 Parameters
659 ----------
660 slice_number : int
661 The slice number for the step.
662 step_number : int
663 The step number for the experiment.
664 experiment_settings : tbt.ExperimentSettings
665 The experiment settings.
667 Returns
668 -------
669 bool
670 True if the step is performed successfully.
671 """
672 # # breakout experiment settings elements
673 microscope = experiment_settings.microscope
674 general_settings = experiment_settings.general_settings
675 step_sequence = experiment_settings.step_sequence
676 enable_EBSD = experiment_settings.enable_EBSD
677 enable_EDS = experiment_settings.enable_EDS
679 # get operation settings, execute operation.
680 operation = step_sequence[step_number - 1] # list is 0-indexed
681 print(
682 f"Slice {slice_number}, Step {step_number} of {general_settings.step_count}, '{operation.name}', a {operation.type.value} type step."
683 )
684 if (
685 slice_number - 1
686 ) % operation.frequency != 0: # slices start at 1, perform all steps on slice 1.
687 print(
688 f"\tStep frequency is every {operation.frequency} slices, starting on slice 1. Skipping step on this slice.\n"
689 )
690 return
692 # log step_start position
693 log.position(
694 step_number=step_number,
695 step_name=operation.name,
696 slice_number=slice_number,
697 log_filepath=general_settings.log_filepath,
698 dataset_name=cs.Constants.pre_position_dataset_name,
699 current_position=factory.active_stage_position_settings(
700 microscope=microscope,
701 ),
702 )
704 # retract all devices
705 print("\tRetracting all devices...")
706 # with ut.nostdout():
707 devices.retract_all_devices(
708 microscope=microscope,
709 enable_EBSD=enable_EBSD,
710 enable_EDS=enable_EDS,
711 )
712 print("\tDevices retracted.")
714 # move stage to starting position for slice
715 stage.step_start_position(
716 microscope=microscope,
717 slice_number=slice_number,
718 operation=operation,
719 general_settings=general_settings,
720 )
722 # perform specific operation
723 perform_operation(
724 operation.operation_settings,
725 step=operation,
726 general_settings=general_settings,
727 slice_number=slice_number,
728 )
730 # log step end position
731 log.position(
732 step_number=step_number,
733 step_name=operation.name,
734 slice_number=slice_number,
735 log_filepath=general_settings.log_filepath,
736 dataset_name=cs.Constants.post_position_dataset_name,
737 current_position=factory.active_stage_position_settings(
738 microscope=microscope,
739 ),
740 )
742 # retract all devices
743 print("\tRetracting all devices...")
744 # with ut.nostdout():
745 devices.retract_all_devices(
746 microscope=microscope,
747 enable_EBSD=enable_EBSD,
748 enable_EDS=enable_EDS,
749 )
750 print("\tDevices retracted. Step Complete.\n")
753def run_experiment_cli(
754 start_slice: int,
755 start_step: int,
756 yml_path: Path,
757):
758 """
759 Main loop for the experiment, accessed through the command line.
761 This function runs the main loop for the experiment based on the specified start slice, start step, and YAML configuration file.
763 Parameters
764 ----------
765 start_slice : int
766 The starting slice number for the experiment.
767 start_step : int
768 The starting step number for the experiment.
769 yml_path : Path
770 The path to the YAML configuration file.
772 Returns
773 -------
774 None
775 """
777 experiment_settings = setup_experiment(yml_path=yml_path)
779 # warn user of any EBSD/EDS lack of control
780 warning_text = """is not enabled, you will not have access to safety
781 checking and these modalities during data collection. Please ensure
782 this detector is retracted before proceeding."""
783 if not experiment_settings.enable_EBSD:
784 print(f"\nWARNING: EBSD {warning_text}")
785 if not ut.yes_no("Continue?"):
786 print("\nExiting now...")
787 exit()
788 if not experiment_settings.enable_EDS:
789 print(f"\nWARNING: EDS {warning_text}")
790 if not ut.yes_no("Continue?"):
791 print("\nExiting now...")
792 exit()
794 # main loop
795 log.experiment_settings(
796 slice_number=start_slice,
797 step_number=start_step,
798 log_filepath=experiment_settings.general_settings.log_filepath,
799 yml_path=yml_path,
800 )
801 num_steps = len(experiment_settings.step_sequence)
802 print(
803 f"\n\nBeginning serial sectioning experiment on slice {start_slice}, step {start_step} of {num_steps}.\n"
804 )
806 for slice_number in range(
807 start_slice, experiment_settings.general_settings.max_slice_number + 1
808 ): # inclusive of max slice number
809 for step_number in range(start_step, num_steps + 1): # list is 1-indexed
810 perform_step(
811 slice_number=slice_number,
812 step_number=step_number,
813 experiment_settings=experiment_settings,
814 )
816 # reset start_step to 1 at end of slice
817 start_step = 1
819 ut.disconnect_microscope(
820 microscope=experiment_settings.microscope,
821 quiet_output=True,
822 )
824 print("\n\nExperiment complete.")
827if __name__ == "__main__":
828 pass