fluxion

defmodule Fluxion.ODE.Rk4 do
  @moduledoc """
  A fourth-order Runge-Kutta solver for initial value problems of Ordinary
  Differential Equations of the form @y'(t) = f(t, y(t))@.

  RK4 generates a sequence of approximate @y(t)@ values defined:

  @q\\_1 = f(t\\_k, y\\_k)@

  @q\\_2 = f(t\\_k + \\frac{h}{2}, y\\_k + \\frac{h}{2}q\\_1)@

  @q\\_3 = f(t\\_k + \\frac{h}{2}, y\\_k + \\frac{h}{2}q\\_2)@

  @q\\_4 = f(t\\_k + h, y\\_k + hq\\_3)@

  @y\\_{k+1} = y\\_k + \\frac{h}{6}[q\\_1 + 2q\\_2 + 2q\\_3 + q\\_4]@
  """

  @doc """
  Compute N iterations of the RK4 method and return the computed value for
  @y(t + stride * iterations)@.

  ## Example

  For the purpose of example, start from a known function for @y@ like
  @y(t) = sin(t)@.
  This gives @y'(t) = cos(t)@.
  Given this, it's possible to approximate the real solution with the RK4
  method.

      iex> defmodule TSq do
      ...>   @behaviour Fluxion.ODE
      ...>   @impl true
      ...>   def f(x, _y_k), do: :math.cos(x)
      ...> end
      iex> steps = 10
      iex> stride = 0.01
      iex> # Compute the real value of y(t) after 10 steps
      iex> real_y_1 = :math.sin(stride*steps)
      iex> # Approximate the value of y(t) after 10 steps
      iex> approx_y_1 = Fluxion.ODE.Rk4.step(TSq, 0 , stride, steps)
      iex> # Compare the results
      iex> Fluxion.Math.is_equal(real_y_1, approx_y_1, 1.0e-10)
      :true
  """
  @spec step(
          ode :: Fluxion.ODE.t(),
          initial :: number,
          stride :: number,
          iterations :: number
        ) :: number
  def step(ode, initial, stride, iterations \\ 1) do
    iterate(ode, initial, 0, stride, iterations)
  end

  # Handle the end of the iteration loop
  defp iterate(_ode, y_k, _t_k, _h, _remaining_iterations = 0) do
    y_k
  end

  # Compute the next value in the sequence y_k
  defp iterate(ode, y_k, t_k, h, n) do
    {q1, q2, q3, q4} = middle_steps(ode, t_k, y_k, h)
    y_next = y_k + h / 6.0 * (q1 + 2.0 * q2 + 2.0 * q3 + q4)
    iterate(ode, y_next, t_k + h, h, n - 1)
  end

  # Compute the q values for the intermediate steps taken by the RK4 method
  defp middle_steps(ode, t_k, y_k, h) do
    q1 = Fluxion.ODE.f!(ode, t_k, y_k)
    q2 = Fluxion.ODE.f!(ode, t_k + h / 2.0, y_k + h / 2.0 * q1)
    q3 = Fluxion.ODE.f!(ode, t_k + h / 2.0, y_k + h / 2.0 * q2)
    q4 = Fluxion.ODE.f!(ode, t_k + h, y_k + h * q3)
    {q1, q2, q3, q4}
  end
end