LUME Model Actions Framework¶

The LUME Model interface defines a common way of exposing a calculation or simulation with state through readable / writable variables. A common use case is to expose an existing simulation tool wrapped by LUME or other python libraries to control systems. To make this application easeier, lume-base includes the "actions framework" for models.

In this framework, the Variable objects are extended with an additional .get(...) and .set(...) method that acts on a simulator tool. These extended variables also include any parameters needed to talk to the simulator such as the names of elements and attributes, or scaling factors. We call these "actions". The actions may be inserted into a special LUME model called ActionModel which handles the model's get, set, and supported variables methods. Defining parameterized actions which define broad classes of what you "can do" to the simulator object encourages maintainable, readable code and ensures a model with consistency between the supported variables list and the things the model is able to do.

In [1]:

from lume.actions import Action, ActionModel
from lume.variables import ScalarVariable
from typing import Any

from lume.actions import Action, ActionModel from lume.variables import ScalarVariable from typing import Any

Mock Simulator Object¶

First, we need to define a simulator object. In practice, this would be a LUME object such as Impact or Genesis. Here, we show a simulator tool with a dict of elements which are each a dict of attributes and their values.

In [2]:

class Simulator:
    """A mockup of a simulation tool (ie Impact, GPT, Genesis)"""

    def __init__(self):
        self.ele = {
            "SOL1": {"length": 1.0, "current": 0.5},
            "CAV1": {"phase0": 0.0, "voltage": 10.0},
        }

class Simulator: """A mockup of a simulation tool (ie Impact, GPT, Genesis)""" def __init__(self): self.ele = { "SOL1": {"length": 1.0, "current": 0.5}, "CAV1": {"phase0": 0.0, "voltage": 10.0}, }

Example Action¶

Here, we make our first action. This action sets or gets any of the float attributes in the element dict. It is parameterized by the element name and attribute name allowing the same code to be reused for any of the elements in the simulator. Which one it talks to is determined by the value of ele_name and attribute_name set when the variable is created.

In [3]:

class EleActionVariable(ScalarVariable, Action[Simulator]):
    """
    This variable sets and gets the value in the elements dict of `simulator` with the given element and attribute name.

    This could, for example, be used to set and get paramters of beamline elements in an accelerator simulation code.
    """

    # Action variables are still pydantic objects since they are variables
    # We can store data in them to parameterize the set and get function
    ele_name: str
    attribute_name: str

    # Action variable stores both method to "talk to" simulator and the data needed to parameterize it
    def _get(self, simulator: Simulator) -> float:
        return simulator.ele[self.ele_name][self.attribute_name]

    def _set(self, simulator: Simulator, value: Any) -> float:
        simulator.ele[self.ele_name][self.attribute_name] = value

class EleActionVariable(ScalarVariable, Action[Simulator]): """ This variable sets and gets the value in the elements dict of `simulator` with the given element and attribute name. This could, for example, be used to set and get paramters of beamline elements in an accelerator simulation code. """ # Action variables are still pydantic objects since they are variables # We can store data in them to parameterize the set and get function ele_name: str attribute_name: str # Action variable stores both method to "talk to" simulator and the data needed to parameterize it def _get(self, simulator: Simulator) -> float: return simulator.ele[self.ele_name][self.attribute_name] def _set(self, simulator: Simulator, value: Any) -> float: simulator.ele[self.ele_name][self.attribute_name] = value

In [4]:

# Create a variable action which sets SOL1's current. Remember, it is a variable and we must define all of the variable's
# fields (such as variable name) as well as the ones that parameterize the action
sol_action = EleActionVariable(
    name="INJECTOR:SOL1:CURRENT",  # The variable's name
    value_range=(0.0, 1.0),  # [optional] The variable's range
    unit="kA",  # [optional] Units associated with the variable
    ele_name="SOL1",  # The action's parameter. The name of the element in the simulator
    attribute_name="current",  # Another action parameter. The name of the attribute internally
)

# Create a variable action which sets SOL1's current. Remember, it is a variable and we must define all of the variable's # fields (such as variable name) as well as the ones that parameterize the action sol_action = EleActionVariable( name="INJECTOR:SOL1:CURRENT", # The variable's name value_range=(0.0, 1.0), # [optional] The variable's range unit="kA", # [optional] Units associated with the variable ele_name="SOL1", # The action's parameter. The name of the element in the simulator attribute_name="current", # Another action parameter. The name of the attribute internally )

In [5]:

# Use the action to change settings in the simulator
sim = Simulator()

# Get the current
cur = sol_action._get(sim)
print("Solenoid Current: {:.2} kA".format(cur))

# Set it lower and print the current afterwards
sol_action._set(sim, 0.25)
cur = sol_action._get(sim)
print("Solenoid Current: {:.2} kA".format(cur))

# Use the action to change settings in the simulator sim = Simulator() # Get the current cur = sol_action._get(sim) print("Solenoid Current: {:.2} kA".format(cur)) # Set it lower and print the current afterwards sol_action._set(sim, 0.25) cur = sol_action._get(sim) print("Solenoid Current: {:.2} kA".format(cur))

Solenoid Current: 0.5 kA
Solenoid Current: 0.25 kA

Actions Model¶

Now, we will create a model from several actions. We will create actions that expose all of the elements and attributes inside of the simulator.

In [6]:

actions = [
    EleActionVariable(
        name="INJECTOR:SOL1:CURRENT",
        unit="kA",
        ele_name="SOL1",
        attribute_name="current",
    ),
    EleActionVariable(
        name="INJECTOR:SOL1:LENGTH",
        unit="kA",
        ele_name="SOL1",
        attribute_name="length",
    ),
    EleActionVariable(
        name="INJECTOR:CAV1:PHASE",
        unit="1/2/pi",
        ele_name="CAV1",
        attribute_name="phase0",
    ),
    EleActionVariable(
        name="INJECTOR:CAV1:VOLTAGE",
        unit="kV",
        ele_name="CAV1",
        attribute_name="voltage",
    ),
]

actions = [ EleActionVariable( name="INJECTOR:SOL1:CURRENT", unit="kA", ele_name="SOL1", attribute_name="current", ), EleActionVariable( name="INJECTOR:SOL1:LENGTH", unit="kA", ele_name="SOL1", attribute_name="length", ), EleActionVariable( name="INJECTOR:CAV1:PHASE", unit="1/2/pi", ele_name="CAV1", attribute_name="phase0", ), EleActionVariable( name="INJECTOR:CAV1:VOLTAGE", unit="kV", ele_name="CAV1", attribute_name="voltage", ), ]

In [7]:

# Make a fresh simulator object for the model
sim = Simulator()

# Construct a model from the actions
# It contains the simulator object and the action variables
model = ActionModel(simulator=sim, action_variables=actions)

# Make a fresh simulator object for the model sim = Simulator() # Construct a model from the actions # It contains the simulator object and the action variables model = ActionModel(simulator=sim, action_variables=actions)

In [8]:

# Get all values from the simulator
model.get(
    [
        "INJECTOR:SOL1:CURRENT",
        "INJECTOR:SOL1:LENGTH",
        "INJECTOR:CAV1:PHASE",
        "INJECTOR:CAV1:VOLTAGE",
    ]
)

# Get all values from the simulator model.get( [ "INJECTOR:SOL1:CURRENT", "INJECTOR:SOL1:LENGTH", "INJECTOR:CAV1:PHASE", "INJECTOR:CAV1:VOLTAGE", ] )

Out[8]:

{'INJECTOR:SOL1:CURRENT': 0.5,
 'INJECTOR:SOL1:LENGTH': 1.0,
 'INJECTOR:CAV1:PHASE': 0.0,
 'INJECTOR:CAV1:VOLTAGE': 10.0}

In [9]:

# Turn off the cavity
model.set({"INJECTOR:CAV1:VOLTAGE": 0.0})
model.get(
    [
        "INJECTOR:SOL1:CURRENT",
        "INJECTOR:SOL1:LENGTH",
        "INJECTOR:CAV1:PHASE",
        "INJECTOR:CAV1:VOLTAGE",
    ]
)

# Turn off the cavity model.set({"INJECTOR:CAV1:VOLTAGE": 0.0}) model.get( [ "INJECTOR:SOL1:CURRENT", "INJECTOR:SOL1:LENGTH", "INJECTOR:CAV1:PHASE", "INJECTOR:CAV1:VOLTAGE", ] )

Out[9]:

{'INJECTOR:SOL1:CURRENT': 0.5,
 'INJECTOR:SOL1:LENGTH': 1.0,
 'INJECTOR:CAV1:PHASE': 0.0,
 'INJECTOR:CAV1:VOLTAGE': 0.0}

Move Advanced Action Example¶

Actions contain arbitrary python code and can do more than simply setting and getting attributes. Applications include, for example, scaling units, unit calibration, composite variables (consisting of setting/getting more than one attribute in the tool).

In [10]:

class FieldIntegralVariable(ScalarVariable, Action[Simulator]):
    """
    This action demonstrates a variable which involves more than one attribute. In this case, an integrated field.
    """

    # The name of the element we act on
    ele_name: str
    cur_to_field: float = (
        1.0  # A conversion factor from current * length to integrated field
    )

    # Action variable stores both method to "talk to" simulator and the data needed to parameterize it
    def _get(self, simulator: Simulator) -> float:
        magnet_length = simulator.ele[self.ele_name]["length"]
        magnet_current = simulator.ele[self.ele_name]["current"]
        return magnet_length * magnet_current * self.cur_to_field

    def _set(self, simulator: Simulator, value: Any) -> float:
        magnet_length = simulator.ele[self.ele_name]["length"]
        magnet_current = value / magnet_length / self.cur_to_field
        simulator.ele[self.ele_name]["current"] = magnet_current

class FieldIntegralVariable(ScalarVariable, Action[Simulator]): """ This action demonstrates a variable which involves more than one attribute. In this case, an integrated field. """ # The name of the element we act on ele_name: str cur_to_field: float = ( 1.0 # A conversion factor from current * length to integrated field ) # Action variable stores both method to "talk to" simulator and the data needed to parameterize it def _get(self, simulator: Simulator) -> float: magnet_length = simulator.ele[self.ele_name]["length"] magnet_current = simulator.ele[self.ele_name]["current"] return magnet_length * magnet_current * self.cur_to_field def _set(self, simulator: Simulator, value: Any) -> float: magnet_length = simulator.ele[self.ele_name]["length"] magnet_current = value / magnet_length / self.cur_to_field simulator.ele[self.ele_name]["current"] = magnet_current

In [11]:

class ScaledEleAction(ScalarVariable, Action[Simulator]):
    """
    This action controls an attribute for an element with a linear transformation which can be used for model
    calibration or for unit conversion.
    """

    # The element name and attribute the control variable acts on
    ele_name: str
    attribute_name: str

    # Define a linear transformation mapping simulation units to control system units
    scale: float
    offset: float

    def _get(self, simulator):
        return (
            self.scale * simulator.ele[self.ele_name][self.attribute_name] + self.offset
        )

    def _set(self, simulator, value):
        simulator.ele[self.ele_name][self.attribute_name] = (
            value - self.offset
        ) / self.scale

class ScaledEleAction(ScalarVariable, Action[Simulator]): """ This action controls an attribute for an element with a linear transformation which can be used for model calibration or for unit conversion. """ # The element name and attribute the control variable acts on ele_name: str attribute_name: str # Define a linear transformation mapping simulation units to control system units scale: float offset: float def _get(self, simulator): return ( self.scale * simulator.ele[self.ele_name][self.attribute_name] + self.offset ) def _set(self, simulator, value): simulator.ele[self.ele_name][self.attribute_name] = ( value - self.offset ) / self.scale

Adding Actions to a Model¶

Actions can be added to the model after it is created allowing you to write or use other's code that generates a basic model from your simulation tool and then add additional actions customized for your situation. We demonstrate this by adding our new actions to the existing model.

In [12]:

# The composite field integral variable on the solenoid element
var = FieldIntegralVariable(name="INJECTOR:SOL1:INTEGRATED_FIELD", ele_name="SOL1")
model.register_action_variable(var)

# A scaled version of the phase
var = ScaledEleAction(
    name="INJECTOR:CAV1:PHASE_DEG",
    ele_name="CAV1",
    attribute_name="phase0",
    scale=360.0,
    offset=0.0,
)
model.register_action_variable(var)

# The composite field integral variable on the solenoid element var = FieldIntegralVariable(name="INJECTOR:SOL1:INTEGRATED_FIELD", ele_name="SOL1") model.register_action_variable(var) # A scaled version of the phase var = ScaledEleAction( name="INJECTOR:CAV1:PHASE_DEG", ele_name="CAV1", attribute_name="phase0", scale=360.0, offset=0.0, ) model.register_action_variable(var)

In [13]:

# Get the values
model.get(
    [
        "INJECTOR:SOL1:INTEGRATED_FIELD",
        "INJECTOR:SOL1:CURRENT",
        "INJECTOR:CAV1:PHASE_DEG",
        "INJECTOR:CAV1:PHASE",
    ]
)

# Get the values model.get( [ "INJECTOR:SOL1:INTEGRATED_FIELD", "INJECTOR:SOL1:CURRENT", "INJECTOR:CAV1:PHASE_DEG", "INJECTOR:CAV1:PHASE", ] )

Out[13]:

{'INJECTOR:SOL1:INTEGRATED_FIELD': 0.5,
 'INJECTOR:SOL1:CURRENT': 0.5,
 'INJECTOR:CAV1:PHASE_DEG': 0.0,
 'INJECTOR:CAV1:PHASE': 0.0}

In [14]:

# Demonstrate setting them
model.set({"INJECTOR:SOL1:INTEGRATED_FIELD": 0.75, "INJECTOR:CAV1:PHASE_DEG": 180})
model.get(
    [
        "INJECTOR:SOL1:INTEGRATED_FIELD",
        "INJECTOR:SOL1:CURRENT",
        "INJECTOR:CAV1:PHASE_DEG",
        "INJECTOR:CAV1:PHASE",
    ]
)

# Demonstrate setting them model.set({"INJECTOR:SOL1:INTEGRATED_FIELD": 0.75, "INJECTOR:CAV1:PHASE_DEG": 180}) model.get( [ "INJECTOR:SOL1:INTEGRATED_FIELD", "INJECTOR:SOL1:CURRENT", "INJECTOR:CAV1:PHASE_DEG", "INJECTOR:CAV1:PHASE", ] )

Out[14]:

{'INJECTOR:SOL1:INTEGRATED_FIELD': 0.75,
 'INJECTOR:SOL1:CURRENT': 0.75,
 'INJECTOR:CAV1:PHASE_DEG': 180.0,
 'INJECTOR:CAV1:PHASE': 0.5}

Removal of Actions¶

You may also remove actions by variable name. This can be useful in modifying existing models.

In [15]:

# Remove the action
model.unregister_action_variable("INJECTOR:SOL1:INTEGRATED_FIELD")

# Attempt to access
try:
    model.get(["INJECTOR:SOL1:INTEGRATED_FIELD"])
except ValueError as ex:
    print(ex)

# Remove the action model.unregister_action_variable("INJECTOR:SOL1:INTEGRATED_FIELD") # Attempt to access try: model.get(["INJECTOR:SOL1:INTEGRATED_FIELD"]) except ValueError as ex: print(ex)

Variable 'INJECTOR:SOL1:INTEGRATED_FIELD' is not supported by the model.

In [ ]: