Quick Start

Quick Start#

In this notebook, we will guide you through the basic steps of using EvoX.

# install evox, skip it if you have already installed evox
try:
    import evox
except ImportError:
    !pip install --disable-pip-version-check --upgrade -q evox
    import evox

from evox import algorithms, problems, workflows, monitors, use_state

import jax.numpy as jnp
from jax import random

Create an algorithm and a problem#

To demostrate, we will create a PSO algorithm and an Ackley function.

For more detailed list, please refer to our API documentation. List of Algorithms and List of Problems.

pso = algorithms.PSO(
    lb=jnp.full(shape=(2,), fill_value=-32),
    ub=jnp.full(shape=(2,), fill_value=32),
    pop_size=100,
)
ackley = problems.numerical.Ackley()

Now we want to run the algorithm against the problem. To accomplish this, we need to create a workflow which represents the overall process of evolutionary computation.

workflow = workflows.StdWorkflow(pso, ackley)

Since we adopt the functional programming paradigm. We must explicitly initialize and use the state of a module. To initialize, call init with a pseudorandom number generators key (PRNGKey).

key = random.PRNGKey(42)
state = workflow.init(key)

The state represents the mutatable variables within the whole workflow, including those inside the algorithm and the problem. For example, the population in an algorithm is part of the state, because it will be changing across iterations.

Now, call step on the workflow to execute one iteration.

state = workflow.step(state)

To run multiple iterations, wrap it inside a for-loop.

# run the workflow for 100 steps
for i in range(100):
    state = workflow.step(state)

Notice that we are passing state as an argument of step and it returns a new state. This is exactly how the stateful computation works in functional programming.

And you may also notice that the step doesn’t give any feedback, like the result of the optimization. This is because we are missing another component in our workflow. Introducing monitor.

Monitor#

Monitor is a standard way to monitor the intermediate values inside a optimization process. Information like fitness or population can be easily obtained by the monitor.

Now, create a “Evalution monitor”

monitor = monitors.EvalMonitor()

The monitor can be plugged into the workflow.

workflow = workflows.StdWorkflow(
    pso,
    ackley,
    monitors=[monitor],
)

Now, re-initialize the workflow, and executed it again.

# init the workflow
state = workflow.init(key)
# run the workflow for 5 steps
for i in range(5):
    state = workflow.step(state)

To access the minimum fitness achieved, we use the monitor submodule. However, since the state variable belongs to the workflow variable, it cannot be accessed directly by its submodule, monitor. Instead, we need to use the use_state wrapper.

Note

In EvoX, whenever you want to call a method from a submodule, the use_state wrapper is required. This wrapper handles nested submodules at any depth, so you don’t need to worry about the module structure.

best_fitness, state = use_state(monitor.get_best_fitness)(state)
print(f"Best fitness so far: {best_fitness}")

Best fitness so far: 4.297701835632324

Execute another 5 iterations, and the minimum fitness will change accordingly.

for i in range(50):
    state = workflow.step(state)

# since we run another 50 iterations, the result should get better.
best_fitness, state = use_state(monitor.get_best_fitness)(state)
print(f"Best fitness so far: {best_fitness}")

Best fitness so far: 0.00029850006103515625

best_solution, state = use_state(monitor.get_best_solution)(state)
print(f"Best solution so far: {best_solution}")

Best solution so far: [-8.6565415e-05  5.9473794e-05]

It shows the best solution is (-8.6565415e-05, 5.9473794e-05), which is close to the global minimum at (0,0). 🥳

Additionally, please note that the best fitness remains the same as last time. This is because we are using the same key when initializing the workflow as before. This deterministic behavior in EvoX allows others to easily reproduce your results.

finally, we can use the built-in plot function to plot the loss curve.

monitor.plot()
Contents