Docs / Python SDK / Decorators

Decorators

The Axon Python SDK is a small library of decorators that you apply to your Python classes and methods. They serve as a contract between your code and the FluxSim platform. None of them change the runtime behavior of your code. They are all no-ops at runtime. Their purpose is to give the FluxSim code generator the information it needs to wire your blocks together into a runnable simulation.

from axon import axon_block_type, axon_setup, axon_step_2, axon_simulation_feature

How Codegen Works

The .ax canvas is a visual graph where you wire blocks together. Each node on the canvas corresponds to a Python class in your simulator project. When you save the graph, or save any Python file in the simulator builder, FluxSim scans your project's source files, reads the decorators via AST analysis, and generates simulation.py, the actual runnable simulation that executes inside your deployed Docker container.

The decorators tell codegen things it cannot infer from the code alone: which method is the setup step, which methods run in which loop, which class is a loader, and which methods expose live outputs to other blocks. You never write or edit simulation.py by hand. It is always generated.

The generated simulation has two nested loops:

# Pseudocode — actual output is in simulation.py

for _patient in cohort_loop.patients():

    # @axon_setup methods run here — once per patient
    block.load()

    # @axon_step_1 methods run here — once per patient
    loader_result = loader.load(cohort, patient_id)
    block.load(loader_result)

    t = 0
    while t < ProjectConstants.duration:

        # @axon_step_2 methods run here — every timestep
        output = block_a.sample(block_b.get_state())

        t += ProjectConstants.dt

dt and duration are project-level constants defined in the Project Constants panel of the simulator builder. They control the timestep size and the total simulation duration for every patient in the cohort.

The decorators are imported from axon, a small library that ships with every simulator project. All decorators are no-ops at runtime. They exist purely to inform codegen.

@axon_block_type

Marks a class as an Axon block, making it visible to the simulator builder and eligible to appear as a node on the .ax canvas. This is required on every class that participates in the simulation graph.

ParameterTypeDescription
categorystrEither 'block' for a standard simulation block, or 'loader' for a data loader. See @axon_loader.
from axon import axon_block_type, axon_setup, axon_step_2

@axon_block_type('block')
class BisMeasurement:

    @axon_setup
    def load(self):
        self.ec50 = 3.08
        self.gamma = 1.47

    @axon_step_2
    def sample(self, state=(0.0, 0.0, 0.0)):
        Ce, _, _, _ = state
        return self.e0 - self.emax * (Ce ** self.gamma) / (self.ec50 ** self.gamma + Ce ** self.gamma)

@axon_setup

Marks a method as the initialization step for a block. The setup method is called once per patient, before either simulation loop begins. This is where you initialize state variables, set model parameters, and prepare anything that needs to be reset between patients.

The setup method must be named load by convention. It takes no arguments beyond self.

@axon_setup
def load(self):
    # Eleveld et al. 2018 — population BIS model parameters
    self.e0 = 93.0
    self.emax = 93.0
    self.ec50 = 3.08
    self.gamma = 1.47

@axon_step_1 / @axon_step_2

Marks a method as a simulation step. The distinction between axon_step_1 and axon_step_2 determines which loop the method is called in.

  • @axon_step_1 runs once per patient, inside the cohort loop but before the timestep loop. Use this for patient-specific setup that depends on data loaded at runtime, such as initializing a plant state from patient demographics.
  • @axon_step_2 runs every timestep, inside the inner simulation loop. Use this for methods that compute outputs from inputs at each tick, such as a PD model or a controller update.
# @axon_step_1: runs once per patient, resolves patient-specific initial conditions
@axon_step_1
def load(self, cohort: str, patient_id: int):
    state0, params = self.loader.load(cohort, patient_id)
    self.state = state0
    self.params = params

# @axon_step_2: runs every timestep dt
@axon_step_2
def sample(self, state=(0.0, 0.0, 0.0)):
    Ce, _, _, _ = state
    bis = self.e0 - self.emax * (Ce ** self.gamma) / (self.ec50 ** self.gamma + Ce ** self.gamma)
    return max(0.0, bis)

In the generated simulation.py, this maps directly to the loop structure described in How Codegen Works.

@axon_finish

Marks a method that runs once at the end of the simulation for each patient, after the timestep loop completes. Use this for any cleanup, final scoring, or artifact writing that should happen after the full simulation trace has been produced.

@axon_finish
def finish(self):
    self.capture.write()

@axon_loader

Marks a class as a data loader, bound to a specific block. A loader replaces the block's @axon_setup inputs on the canvas with its own inputs, typically things like cohort and patient_id. Visually, you see the loader's inputs on the block itself. In the generated code, codegen calls the loader, takes its outputs, and feeds them into the block's setup method.

ParameterTypeDescription
loader_classstrThe name of the block class this loader is bound to.
from axon import axon_block_type, axon_setup, axon_getter, load_patient

@axon_block_type('loader')
@axon_loader('SchniderPlant')
class SchniderPlantLoader:

    @axon_setup
    def load(self, cohort: str, patient_id: int):
        patient = load_patient(cohort, patient_id)
        self.state0 = compute_initial_state(patient)
        self.params = compute_params(patient)

    @axon_getter
    def get_state0(self):
        return self.state0

    @axon_getter
    def get_params(self):
        return self.params

Codegen resolves this automatically. In simulation.py, you will see the loader instantiated and called, with its outputs passed directly into the block's setup:

# Generated output — you do not write this
schniderPlantLoader = SchniderPlantLoader()
state0, constant_params = schniderPlantLoader.load(ProjectConstants.cohort, cohort_loop.get_patient_id())
schniderplant.load(state0, constant_params)

@axon_getter

Marks a method as a live output of a block. A getter exposes a class variable as a real-time output port on the canvas. When another block wires into a getter, codegen generates a call to that getter at each timestep and passes the result directly as an argument to the receiving block's step method.

This is different from a standard return value. A getter reads the current value of a class variable at the moment it is called, rather than returning a snapshot captured at some earlier point in the simulation.

@axon_getter
def get_state(self):
    return self.state

If a wire connects SchniderPlant.get_state to BisMeasurement.sample, the generated code looks like:

# Generated output — called every timestep
bis = bismeasurement.sample(schniderplant.get_state())

@axon_controller

Marks a class as the controller entry point. This tells the simulator which class to substitute in place of the baked-in controller when running a simulation. When you click Simulate, the simulator deploys your controller project and codegen identifies the @axon_controller class as the implementation to inject.

from axon import axon_controller, axon_setup, axon_step_2

@axon_controller
class PidController:

    @axon_setup
    def load(self):
        self.setpoint = 50
        self.kp = 2.0
        self.ki = 0.1
        self.kd = 0.5
        self.integral = 0.0
        self.prev_error = 0.0

    @axon_step_2
    def update(self, bis: float) -> float:
        error = self.setpoint - bis
        self.integral += error * ProjectConstants.dt
        derivative = (error - self.prev_error) / ProjectConstants.dt
        self.prev_error = error
        return self.kp * error + self.ki * self.integral + self.kd * derivative

Only one class per controller project should carry @axon_controller.

@axon_simulation_feature

Marks a method as a configurable simulation parameter. Methods decorated with @axon_simulation_feature appear as controls in the Simulation Config panel, accessible via the gear icon to the left of the Simulate button. This lets you vary simulation behavior at runtime without modifying source code or regenerating simulation.py.

The decorated method acts as a setter. Codegen injects a call to it before the simulation loop begins, passing the value configured in the UI. The type annotation on the method's argument determines what kind of control is rendered: a toggle for bool, a number field for float or int.

ParameterTypeDescription
categorystrGroups this feature under a named section in the config panel. Defaults to 'General'.
defaultanyThe value used if the feature has never been configured. Shown as the initial value in the config panel.
from axon import axon_block_type, axon_setup, axon_step_2, axon_simulation_feature

@axon_block_type('block')
class BisMeasurement:

    @axon_setup
    def load(self):
        self.e0 = 93.0
        self.emax = 93.0
        self.ec50 = 3.08
        self.gamma = 1.47

    @axon_step_2
    def sample(self, state=(0.0, 0.0, 0.0)):
        Ce, _, _, _ = state
        jitter = 0.0
        if self._noise_enable:
            A = self._noise_amplitude
            jitter = A * np.random.uniform() - A / 2.0
        bis = self.e0 - self.emax * (Ce ** self.gamma) / (self.ec50 ** self.gamma + Ce ** self.gamma)
        return max(0.0, bis + jitter)

    @axon_simulation_feature(category='BIS Measurement Noise', default=True)
    def noise_enable(self, enable: bool):
        self._noise_enable = enable

    @axon_simulation_feature(category='BIS Measurement Noise', default=5.0)
    def noise_amplitude(self, amplitude: float):
        self._noise_amplitude = amplitude

The feature setter is called before @axon_setup runs. Do not initialize the backing variable in load() — the decorator's default is the single source of truth and will always be applied first.

The generated injection in simulation.py looks like this:

# Simulation feature configuration
bismeasurement.noise_enable(_sim_config["noise_enable"])
bismeasurement.noise_amplitude(_sim_config["noise_amplitude"])

The _sim_config dict is populated from the values set in the Simulation Config panel and passed to the container at runtime via the SIM_CONFIG environment variable. Each simulation run snapshots the config used, so results are always traceable to the exact parameters that produced them.

@axon_optimize

Marks a setter method on a controller class as an optimization target. Methods decorated with @axon_optimize appear in the Optimization Config panel, accessible via the gear icon next to the Optimize button. Each decorated method defines one dimension of the search space, with a configurable min, max, and default.

When you click Optimize, FluxSim runs a random search over the configured bounds. For each iteration, it samples a candidate value for every @axon_optimize parameter, injects the setter calls into a copy of simulation.py before @axon_setup runs, and executes the full cohort simulation. The iteration with the lowest aggregate loss is recorded as the best candidate.

ParameterTypeDescription
minfloatLower bound of the search range. Also used as the default if default is not specified.
maxfloatUpper bound of the search range.
defaultfloatValue used for regular (non-optimization) simulate runs. This is what gets injected before load() on every normal run.
from axon import axon_controller, axon_setup, axon_step_2, axon_optimize

@axon_controller
class PidController:

    @axon_optimize(min=3500, max=4500, default=4000)
    def kp(self, kp):
        self._kp = kp

    @axon_optimize(min=500, max=700, default=600)
    def ki(self, ki):
        self._ki = ki

    @axon_optimize(min=1800, max=2200, default=2000)
    def kd(self, kd):
        self._kd = kd

    @axon_setup
    def load(self, dt):
        self._dt = dt
        self._setpoint = 50
        self._integral = 0.0
        self._prev_bis = None

    @axon_step_2
    def step(self, bis: float) -> float:
        error = bis - self._setpoint
        self._integral += error * self._dt
        d_term = 0.0 if self._prev_bis is None else (bis - self._prev_bis) / self._dt
        self._prev_bis = bis
        return self._kp * error + self._ki * self._integral + self._kd * d_term

Do not initialize optimized parameters in load(). The setter is always called before load() — either with the sampled candidate during optimization, or with the default value during a normal simulate run. The default is the single source of truth for what your controller runs with when you are not optimizing.

Each optimization run is versioned. Results are stored under versions/vN/opt1/, opt2/, etc. with a manifest capturing the candidate values, per-patient loss, and the sim config snapshot used. If save best only is enabled in the Optimization Config panel, only iterations that improve on the current best loss are persisted to disk.