Actions - #15
Merged
Merged
Conversation
Contributor
There was a problem hiding this comment.
Pull request overview
This PR refactors LUMEBmadModel to use LUME’s action-based API (ActionModel + ActionVariable) and introduces a new lume_bmad.actions module that encapsulates Tao/Bmad interactions as action variables, replacing the legacy transformer + control/output variable dictionaries.
Changes:
- Migrated
LUMEBmadModeltoActionModel, registering Tao lattice/comb outputs and track-mode-dependent variables as action variables. - Added new action variable implementations (element scalars, lattice stats, comb stats, beam-at-element, and screen-related variables).
- Updated tests to use the new action-variable interface and new track type enum values (
"single"/"beam"), and removed the legacy transformer.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
lume_bmad/actions.py |
Adds new action variable classes for Tao/Bmad interaction, including screen-related abstractions. |
lume_bmad/model.py |
Refactors the model to ActionModel, registers action variables, and adds dynamic variable refresh based on track_type. |
lume_bmad/utils.py |
Reworks Tao output-variable construction to return StatVariable/CombStatVariable instances. |
tests/test_basic.py |
Updates tests to construct the model with action variables and validate dynamic outputs under beam/single modes. |
lume_bmad/transformer.py |
Removes the legacy transformer abstraction. |
Comments suppressed due to low confidence (2)
lume_bmad/actions.py:221
tests/test_basic.pyimports and usesScreenVariable, butlume_bmad/actions.pycurrently does not define it, so the import will fail and the screen test cannot run. AddScreenVariable(e.g., as a backward-compatible alias/subclass ofScreenImageVariable).
def _get(self, simulator: Tao) -> Any:
_ = simulator
return self.screen_spec.shape[self.index]
lume_bmad/model.py:255
reset()callsself.set(self._initial_state), but_initial_stateis captured fromupdate_state()and therefore includes many read-only outputs (e.g.,mat6,vec0,name). IfActionModel.set()attempts to write all provided keys, this will fail for read-only action variables. Consider resetting only writable variables (or only the originally supplied control/action variables).
def reset(self):
"""Reset the model to its initial state."""
logger.info("Resetting model to initial state")
self.set(self._initial_state)
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+158
to
+165
| # turn tao eager mode off to speed up setting multiple variables | ||
| self.simulator.cmd("set global lattice_calc_on = F") | ||
|
|
||
| # remove particle group variables | ||
| for var in [f"{ele}_beam" for ele in self._dump_locations]: | ||
| if var in self._variables.keys(): | ||
| self._variables.pop(var) | ||
|
|
||
| if len(output) > 0: | ||
| logger.warning("Warning while setting track_type: %s", "".join(output)) | ||
| warnings.warn(f"Warning while setting track_type: {''.join(output)}") | ||
| values.pop("track_type") | ||
|
|
||
| # handle setting the input beam separately | ||
| if "input_beam" in values.keys(): | ||
| input_beam = values.pop("input_beam") | ||
| fname = getcwd() + "/input_beam.h5" | ||
| logger.debug("Writing input beam to %s", fname) | ||
| input_beam.write(fname) | ||
| self.tao.cmd(f"set beam_init position_file = {fname}") | ||
| # set control variables using their respective set methods | ||
| super()._set(values) | ||
|
|
||
| # reset comb length for tracking outputs | ||
| self.tao.cmd(f"set beam comb_ds_save = {self.comb_ds_save}") | ||
| tao_model_output_variables = get_tao_output_variables(self.tao) | ||
| self._variables.update(tao_model_output_variables) | ||
| # after setting all variables, turn eager mode back on | ||
| self.simulator.cmd("set global lattice_calc_on = T") |
Comment on lines
+182
to
+186
| try: | ||
| self._state[name] = self.supported_variables[name]._get(self.simulator) | ||
| except Exception as e: | ||
| logger.error("Error getting variable %s: %s", name, str(e)) | ||
| raise e |
Comment on lines
+69
to
+71
| control_variables = [ | ||
| EleScalarVariable(name="qf:B1_GRADIENT", units="1/m^2"), | ||
| EleScalarVariable(name="qd:B1_GRADIENT", units="1/m^2"), |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request introduces a major refactor of the
LUMEBmadModelinlume_bmad/model.pyto adopt the newActionModelandActionVariableabstractions, replacing the legacy control/output variable and transformer system. It also adds a comprehensive set of action variable classes inlume_bmad/actions.pyto encapsulate Bmad model interactions. The refactor streamlines variable registration, improves extensibility, and clarifies the model's dynamic output handling, especially for beam tracking modes.Key changes include:
Adoption of ActionModel and ActionVariables:
LUMEBmadModelto inherit fromActionModeland useActionVariableinstances for all variable interactions, removing the old control/output variable dictionaries and transformer logic. Variable registration and state management now use the new action-based interface. [1] [2] [3]Dynamic Variable Registration and Output Handling:
_refresh_dynamic_action_variablesto synchronize model outputs (such as comb statistics and dumped beam distributions) with the current tracking mode (track_type). This ensures that only relevant variables are registered and available depending on whether the model is in single or beam tracking mode.Introduction of Action Variable Classes:
lume_bmad/actions.pyfor scalar, ND, enum, and particle group variables, as well as specialized screen and beam variables. These classes encapsulate the logic for getting and setting values in the Tao simulator, providing a modular and extensible interface for model actions.Simplification and Cleanup:
Improved State Update and Setting Logic:
_get,_set, andupdate_statemethods to delegate variable access to the registered action variables, ensuring consistent handling and reducing duplicated logic. Also improved handling of simulator modes and eager/lazy evaluation for performance. [1] [2]These changes modernize the model interface, improve maintainability, and lay the groundwork for further extensibility.