Output Functions

Mix.install([
  {:integrator, "~> 0.1.2"},
  {:kino_vega_lite, "~> 0.1.7"}
])

Usage

An output function lets you plot the results of an integration or simulation while it’s occurring, or send the data to an animation tool. Let’s see how that works.

First, we’ll need to set up an empty chart to receive the data:

alias VegaLite, as: VL

chart =
  VL.new(
    width: 600,
    height: 400,
    title: "Solution of van der Pol Equation (μ = 1) with ode45"
  )
  |> VL.mark(:line, point: true, tooltip: true)
  |> VL.encode_field(:x, "t", type: :quantitative)
  |> VL.encode_field(:y, "x", type: :quantitative)
  |> VL.encode_field(:color, "x_value", type: :nominal)
  |> Kino.VegaLite.new()
  |> Kino.render()

Now, let’s connect an output function to the simulation, and we’ll inject a Process.sleep(50) on purpose so that the simulation takes a while. Watch the chart above while you run the simulation below.

alias Integrator.SampleEqns

t_initial = Nx.tensor(0.0, type: :f64)
t_final = Nx.tensor(20.0, type: :f64)
x_initial = Nx.tensor([2.0, 0.0], type: :f64)

output_fn = fn t, x ->
  # t and x are lists of Nx tensors

  Enum.zip(t, x)
  |> Enum.map(fn {t, x} ->
    [
      %{t: Nx.to_number(t), x: Nx.to_number(x[0]), x_value: "x[0]"},
      %{t: Nx.to_number(t), x: Nx.to_number(x[1]), x_value: "x[1]"}
    ]
  end)
  |> List.flatten()
  |> Enum.map(fn point ->
    Kino.VegaLite.push(chart, point)

    # Sleep on purpose to slow down the simulation:
    Process.sleep(50)
  end)
end

opts = [output_fn: output_fn, type: :f64]

solution =
  Integrator.integrate(&SampleEqns.van_der_pol_fn/2, [t_initial, t_final], x_initial, opts)