Skip to content

Tuning a process

Once you have built a model of your process, a common problem you might face is tuning its parameters. Plugboard includes a built-in optimisation utility based on Ray Tune and Optuna. Using this tool you can do things like:

  • Calibrate the parameters of a process model to match observed results; and
  • Optimise a process model to maximise or minimise its output.

These capabilities are particularly useful when working with digital twins: for example given a model of a production line, you could use the tuner to work out how to maximise its output.

Tip

By using Ray Tune, Plugboard allows you to run optimisations in parallel within a Ray cluster, allowing you to explore the parameter space quickly even when working with long simulations.

Define a model to optimise

As a simple example, we'll create a simple 3-component model to calculate the maximum height of a projectile launched at a given angle and velocity. 

 flowchart LR
   horizontal@{ shape: rounded, label: Iterator<br>**horizontal** } --> trajectory@{ shape: rounded, label: Trajectory<br>**trajectory** }
   trajectory@{ shape: rounded, label: Trajectory<br>**trajectory** } --> max-height@{ shape: rounded, label: MaxHeight<br>**max-height** }

Running the model with different values of the angle and velocity parameters configured on the Trajectory component will result in different heights being found on the MaxHeight component at the end of the simulation. We will use the Tuner class to explore this parameter space and maximise the projectile height.

Setting up the components

We'll need the following components to implement the model above: 

class Iterator(Component):
    """Creates a sequence of x values."""

    io = IO(outputs=["x"])

    def __init__(self, iters: int, **kwargs: _t.Unpack[ComponentArgsDict]) -> None:
        super().__init__(**kwargs)
        self._iters = iters

    async def init(self) -> None:
        self._seq = iter(range(self._iters))

    async def step(self) -> None:
        try:
            self.x = next(self._seq)
        except StopIteration:
            await self.io.close()


class Trajectory(Component):
    """Computes the height of a projectile."""

    io = IO(inputs=["x"], outputs=["y"])

    def __init__(
        self, angle: float = 30, velocity: float = 20, **kwargs: _t.Unpack[ComponentArgsDict]
    ) -> None:
        super().__init__(**kwargs)
        self._angle_radians = math.radians(angle)
        self._v0 = velocity

    async def step(self) -> None:
        self._logger.info("Calculating trajectory", x=self.x)
        self.y = self.x * math.tan(self._angle_radians) - (9.81 * self.x**2) / (
            2 * self._v0**2 * math.cos(self._angle_radians) ** 2
        )


class MaxHeight(Component):
    """Record the maximum height achieved."""

    io = IO(inputs=["y"], outputs=["max_y"])

    def __init__(self, **kwargs: _t.Unpack[ComponentArgsDict]) -> None:
        super().__init__(**kwargs)
        self.max_y: float = 0

    async def step(self) -> None:
        self.max_y = max(self.y, self.max_y)

Instead of building a Process as we would normally do to run the model directly, we'll instead define the ProcessSpec for the model. 

process_spec = ProcessSpec(
    args=ProcessArgsSpec(
        components=[
            {"type": "hello_tuner.Iterator", "args": {"name": "horizontal", "iters": 100}},
            {
                "type": "hello_tuner.Trajectory",
                "args": {"name": "trajectory", "angle": 30, "velocity": 20},
            },
            {"type": "hello_tuner.MaxHeight", "args": {"name": "max-height"}},
        ],
        connectors=[
            {"source": "horizontal.x", "target": "trajectory.x"},
            {"source": "trajectory.y", "target": "max-height.y"},
        ],
    ),
    type="plugboard.process.LocalProcess",
)
# Check that the process spec can be built
_ = ProcessBuilder.build(spec=process_spec)

Setting up the Tuner

Next, we set up a Tuner object by configuring the angle and velocity arguments as floating point parameters, along with constraints.

Info

Plugboard supports floating point, integer and categorical variables as tunable model parameters. See the definition of ParameterSpec for details.

When building the tuner, we also specify the number of optimisation samples and how many we will allow to run in parallel on Ray. 

tuner = Tuner(
    objective=ObjectiveSpec(  # (1)!
        object_type="component",
        object_name="max-height",
        field_type="field",
        field_name="max_y",
    ),
    parameters=[
        FloatParameterSpec(  # (2)!
            object_type="component",
            object_name="trajectory",
            field_type="arg",
            field_name="angle",
            lower=0,
            upper=90,
        ),
        FloatParameterSpec(
            object_type="component",
            object_name="trajectory",
            field_type="arg",
            field_name="velocity",
            lower=0,
            upper=100,
        ),
    ],
    num_samples=40,  # (3)!
    max_concurrent=4,  # (4)!
    mode="max",  # (5)!
)
result = tuner.run(spec=process_spec)
print(
    f"Best parameters: angle={result.config['trajectory.angle']}, velocity={result.config['trajectory.velocity']}"
)
print(f"Best max height: {result.metrics['max-height.max_y']}")

  1. Set the objective, i.e. what we want our optimisation to target. In this case it is a field on the max-height component. This can be a list of objectives if you need to do multi-objective optimisation.
  2. List the tunable parameters here. The field_type can be "arg" or "initial_value". This is also where you can specify constraints on the parameters.
  3. Set the number of trials to run. More trials will take longer, but may get closer to finding the true optimum.
  4. The level of concurrency to use in Ray.
  5. Whether to minimise or maximise the objective. This must be set as a list for multi-objective optimisation.

Running this code will execute an optimisation job and print out information on each trial, along with the final optimisation result.

Tip

Since Optuna is used under the hood, you can configure the optional algorithm argument on the Tuner with additional configuration defined in OptunaSpec. For example, the storage argument allows you to save the optimisation results to a database or SQLite file. You can then use a tool like Optuna Dashboard to study the optimisation output in more detail.

Tip

You can impose arbitary constraints on variables within a Process. In your step method you can raise a ConstraintError to indicate to the Tuner that a constraint has been breached. This will cause the trial to be stopped, and the optimisation will continue trying to find parameters that don't cause the constraint violation.

Using YAML config

Plugboard's YAML config supports an optional tune section, allowing you to define optimisation jobs alongside your model configuration: 

plugboard:
  process:  # (1)!
    args:
      components:
      - type: hello_tuner.Iterator
        args:
          name: horizontal
          iters: 100
      - type: hello_tuner.Trajectory
        args:
          name: trajectory
          angle: 25
          velocity: 20
      - type: hello_tuner.MaxHeight
        args:
          name: max-height
      connectors:
      - source: horizontal.x
        target: trajectory.x
      - source: trajectory.y
        target: max-height.y
  tune:  # (2)!
    args:
      objective:
        object_name: max-height
        field_type: field
        field_name: max_y
      parameters:
      - type: ray.tune.uniform  # (3)!
        object_type: component
        object_name: trajectory
        field_type: arg
        field_name: angle
        lower: 0
        upper: 90
      - type: ray.tune.uniform
        object_type: component
        object_name: trajectory
        field_type: arg
        field_name: velocity
        lower: 0
        upper: 100
      num_samples: 40
      mode: max
      max_concurrent: 4

  1. As usual, this section defines the Process. It can also be replaced by a path to another YAML file.
  2. This section is optional, and configures the Tuner.
  3. Parameters need to reference a type, so that Plugboard knows the type of parameter to build.

Now run plugboard process tune model-with-tuner.yaml to execute the optimisation job from the CLI.

Advanced usage: complex search spaces

Occasionally you may need to define more complex search spaces, which go beyond what can be defined with a simple parameter configuration. For example:

  • Conditional parameters, e.g. where parameter a must be greater than parameter b; or
  • Looping, e.g. building up a list of tunable parameters that is of variable length.

These conditional search space functions are supported by Ray Tune and can be defined as described in the Ray documentation. To use such a function you will need to:

  1. Setup the Tuner, defining your parameters as usual;
  2. Write a custom function to define the search space, where each tunable parameter has a name of the form "{component_name.field_or_arg_name}"; then
  3. Supply your custom function to the OptunaSpec algorithm configuration.

For example, the following search space makes the velocity depend on the angle: 

def custom_space(trial: Trial) -> dict[str, _t.Any] | None:
    """Defines a custom search space for Optuna."""
    angle = trial.suggest_int("trajectory.angle", 0, 90)
    # Make velocity depend on angle
    trial.suggest_int("trajectory.velocity", angle, 100)

Then use this configuration to point the tuner to the custom_space function. 

plugboard:
  process: model-with-tuner.yaml  # (1)!
  tune:
    args:
      objective:
        object_name: max-height
        field_type: field
        field_name: max_y
      parameters:
      - type: ray.tune.uniform
        object_type: component
        object_name: trajectory
        field_type: arg
        field_name: angle
        lower: 0
        upper: 90
      - type: ray.tune.uniform
        object_type: component
        object_name: trajectory
        field_type: arg
        field_name: velocity
        lower: 0
        upper: 100
      num_samples: 40
      mode: max
      max_concurrent: 4
      algorithm:  # (2)!
        type: ray.tune.search.optuna.OptunaSearch
        space: hello_tuner.custom_space

  1. We can reference a Process from another YAML file here to avoid repetition.
  2. Add the algorithm configuration here.

Then run using plugboard process tune model-with-tuner-custom.yaml.