PSO example

Mix.install([
  {:bia, github: "matiascr/bia"},
  {:kino, "~> 0.10.0"},
  {:kino_vega_lite, "~> 0.1.10"},
  {:vega_lite, "~> 0.1.8"},
  {:explorer, "~> 0.7.1"}
])

Setup

We start by defining aliases for the libraries we’ll use for visualizing the PSO

alias VegaLite, as: Vl
require Kino.VegaLite, as: KVl

Now, we define a Visualizer module that contains a callback

defmodule Visualizer do
  def callback(args) do
    widget = args[:opts][:widget]

    widget
    |> KVl.clear()

    widget
    |> KVl.push_many(
      Enum.map(args[:particles], &GenServer.call(&1, :get_position))
      |> Nx.stack()
      |> Nx.to_list()
      |> Enum.map(fn [x, y] -> %{x: x, y: y} end)
    )

    Process.sleep(250)
  end
end

and create it with the parameters we want for the graph

bound_up = 10.0
bound_down = -10.0

widget = fn ->
  Vl.new(width: 400, height: 400)
  |> Vl.mark(:circle)
  |> Vl.encode_field(:x, "x", type: :quantitative, scale: [domain: [bound_down, bound_up]])
  |> Vl.encode_field(:y, "y", type: :quantitative, scale: [domain: [bound_down, bound_up]])
  |> KVl.new()
  |> Kino.render()
end

We add some functions for the PSO to optimize

defmodule OptimizationFunctions do
  import Nx.Defn

  defn unimodal(t) do
    0.26 * (t[0] ** 2 + t[1] ** 2) - 0.48 * t[0] * t[1]
  end

  defn inverted_pyramid(t) do
    Nx.abs(t[0]) + Nx.abs(t[1])
  end
end

Running the PSO

Finally, we run the PSO algorithm.

In the first step, we add the parameters of the heuristic, along with other optional parameters for in the callback and widget parameters. These will allow us to inject actions that will be performed after each iteration of the PSO. In this case, we have created a graph, and it will be updated each iteration with the position of the particles.

PSO.new(
  population_size: 50,
  num_iterations: 20,
  bound_up: bound_up,
  bound_down: bound_down,
  inertia: 0.1,
  coeff_p: 2.0,
  coeff_g: 1.0,
  callback: &Visualizer.callback/1,
  widget: widget.(),
  fun: &OptimizationFunctions.inverted_pyramid/1
)
|> PSO.run()