Solve an ODE: the harmonic oscillator

from __future__ import annotations

import matplotlib.pyplot as plt
from numpy import array
from numpy import cos
from numpy import linspace
from numpy import ndarray
from numpy import sin

from gemseo import create_discipline
from gemseo.core.discipline.discipline import Discipline
from gemseo.disciplines.ode.ode_discipline import ODEDiscipline
from gemseo.problems.ode.oscillator_discipline import OscillatorDiscipline

This tutorial describes how to instantiate an ODEDiscipline to solve a first-order ordinary differential equation. A first-order ODE is a differential equation that can be written as

\[\frac{ds(t)}{dt} = f(t, s(t))\]

where the right-hand side of the equation \(f(t, s(t))\) is a function of time \(t\) and of a state \(s(t)\) that returns another state \(\frac{ds(t)}{dt}\) (see Hartman, Philip (2002) [1964], Ordinary differential equations, Classics in Applied Mathematics, vol. 38, Philadelphia: Society for Industrial and Applied Mathematics). To solve this equation, initial conditions must be set:

\[s(t_0) = s_0\]

For this example, we consider the harmonic oscillator equation as an example of a time-independent problem.

Solving the harmonic oscillator problem

The harmonic problem describes the position over time of a body oscillating with a frequency \(\omega/ (2 \pi)\) according to the following ODE:

\[\frac{d^2 x(t)}{dt^2} + \omega^2 x(t) = 0\]

where \(x(t)\) is the position coordinate at time \(t\).

To solve the harmonic oscillator problem with GEMSEO, we start by modeling this second-order equation as a first-order equation. Let us define \(y = \frac{dx}{dt}\) and the state vector \(s\) as \(s = \begin{pmatrix}x\\y\end{pmatrix}\). Then the harmonic oscillator problem can be rewritten as:

\[\begin{split}\frac{ds(t)}{dt} = f(t, s(t)) = \begin{pmatrix} y(t) \\ - \omega^2 x(t) \end{pmatrix}\end{split}\]

The analytical solution of the harmonic oscillator problem characterized by an angular velocity \(\omega\), and with initial position and velocity \(x_0\) and \(v_0\) respectively, at the initial time \(t = 0\) is:

\[x(t) = x_0 \cos(\omega t) + (v_0 / \omega) \sin(\omega t)\]

Step 1 : Definition of the dynamic

# As the first step, we introduce a discipline defining the dynamics of the problem,
# that is a discipline representing the right-hand side of the ODE.
# For a harmonic oscillator, the discipline can be an :class:`.AutoPyDiscipline`.
#
# The discipline describing the RHS must include, among its inputs, the time variable
# and the variables defining the state (in the case of the harmonic oscillator, the
# variables `position` and `velocity`). The outputs of the discipline must be the
# time derivatives of the state variables (here `position_dot` and `velocity_dot`).
# All inputs that are neither the time variable nor the state variables, are denoted as
# design variables.

_time = array([0.0])
position_1 = array([1.0])
velocity_1 = array([0.0])
omega_1 = array([2.0])


def rhs_function(
    time: ndarray = _time,
    position: ndarray = position_1,
    velocity: ndarray = velocity_1,
    omega: ndarray = omega_1,
):
    position_dot = velocity
    velocity_dot = -(omega**2) * position

    return position_dot, velocity_dot  # noqa: RET504


rhs_discipline = create_discipline(
    "AutoPyDiscipline",
    py_func=rhs_function,
    grammar_type=Discipline.GrammarType.SIMPLE,
)

Step 2: Initialization of the ODEDiscipline

Once the discipline representing the right-hand side of the ODE has been created, we can create an instance of ODEDiscipline, representing the initial-value problem to be solved.

The constructor of the ODEDiscipline must be provided with the arguments discipline (the discipline representing the RHS of the ODE), and times (an ArrayLike with at least two entries: the fist representing the initial time, and the last one the final time). The parameter state_names is a list of the name of the state variables in rhs_discipline, in order to differentiate them from the time variable (named "time" by default), and from the design variables (here omega). By default, the output of the ODEDiscipline are: the final time of the evaluation of the ODE ("termination_time"), the list of times for which the ODE has been solved ("times"), and the final states of the state variables (named by default "X_final", where "X" is the name of the corresponding state variable). If the boolean return_trajectories is set to True, additional outputs is provided, consisting in the trajectories of the state variables in the times listed in the output "times".

ode_discipline = ODEDiscipline(
    discipline=rhs_discipline,
    times=linspace(0.0, 10.0, 51),
    state_names=["position", "velocity"],
    return_trajectories=True,
)

Step 3: Execution of the ODEDiscipline

Once the ODEDiscipline has been initialized, it can be executed like all other disciplines in GEMSEO by the method execute().

ode_res_1 = ode_discipline.execute()

The default inputs of the ODEDiscipline are the default inputs of the underlying discipline defining the RHS of the ODE. Therefore, ode_res_1 contains the solution of the problem of a harmonic oscillator with angular velocity omega equal to \(2.0\), with initial position \(1.0\) and initial velocity \(0.0\).

Different values for the design variables and for the initial conditions can be specified by passing a suitable dictionary to the execute() method of the class:ODEDiscipline.

position_2 = array([2.0])
velocity_2 = array([0.5])
omega_2 = array([1.0])

ode_res_2 = ode_discipline.execute({
    "position": position_2,
    "velocity": velocity_2,
    "omega": omega_2,
})

The object ode_res_2 contains the solution of a different harmonic oscillator problem, characterized by an angular velocity omega equal to \(1.0\), with a positive initial velocity of \(0.5\), and an initial position deviated by \(2.0\) from the equilibrium.

Terminating events

In some cases, it can be useful not to pursue the solution of the ODE for the entire time interval, but only up to the realization of a certain condition. For example, one could be interested in the dynamic of a falling object up to its impact on the ground. In solvers like scipy.integrate.solve_ivp, the termination condition is encoded by a function with the same entries as the RHS of the ODE, returning a real value. The termination condition is fulfilled when the function crosses the threshold \(0\) (for further information, consult the SciPy documentation). In GEMSEO, the same result can be achieved by passing a list of disciplines to the constructor of ODEDiscipline, each representing a termination condition. All disciplines representing termination conditions must have the same inputs as the discipline for the dynamic, and one real output.

In this example, we consider a new ODEDiscipline, describing the trajectory of a harmonic oscillator from its starting point up to the instant when it crosses the equilibrium position.

def termination_function(
    time: ndarray = _time,
    position: ndarray = position_1,
    velocity: ndarray = velocity_1,
    omega: ndarray = omega_1,
):
    termination = position
    return termination  # noqa: RET504


termination_discipline = create_discipline(
    "AutoPyDiscipline",
    py_func=termination_function,
    grammar_type=Discipline.GrammarType.SIMPLE,
)

Here, termination_discipline is the AutoPyDiscipline encoding the termination condition. In order to include this condition in the solution of the ODE, we pass the tuple (termination_discipline,) as the argument termination_event_disciplines of the constructor of ODEDiscipline.

ode_discipline_termination = ODEDiscipline(
    discipline=rhs_discipline,
    times=linspace(0.0, 10.0, 51),
    state_names=["position", "velocity"],
    termination_event_disciplines=(termination_discipline,),
    return_trajectories=True,
)

ode_res_termination = ode_discipline_termination.execute()

Examining the results

The outcomes of the discipline executions can be accessed by passing the names of the corresponding outputs to ode_res_1, ode_res_2, and ode_res_termination.

plt.plot(ode_res_1["times"], ode_res_1["position_trajectory"])
plt.plot(ode_res_2["times"], ode_res_2["position_trajectory"])
plt.plot(ode_res_termination["times"], ode_res_termination["position_trajectory"])
plt.legend(["ode_res_1", "ode_res_2", "ode_res_termination"])
plt.show()

These results can also be compared with the analytical solutions.

analytic_res_1 = position_1 * cos(omega_1 * ode_res_1["times"]) + (
    velocity_1 / omega_1
) * sin(omega_1 * ode_res_1["times"])
analytic_res_2 = position_2 * cos(omega_2 * ode_res_2["times"]) + (
    velocity_2 / omega_2
) * sin(omega_2 * ode_res_2["times"])
analytic_res_termination = position_1 * cos(
    omega_1 * ode_res_termination["times"]
) + velocity_1 / omega_1 * sin(omega_1 * ode_res_termination["times"])

plt.plot(ode_res_1["times"], analytic_res_1, "r", label="Analytical solution")
plt.plot(
    ode_res_1["times"],
    ode_res_1["position_trajectory"],
    "b--",
    label="Numerical solution",
)
plt.legend(["Analytical solution", "Solution by ODEDiscipline"])
plt.show()

plt.plot(ode_res_2["times"], analytic_res_2, "r", label="Analytical solution")
plt.plot(
    ode_res_2["times"],
    ode_res_2["position_trajectory"],
    "b--",
    label="Numerical solution",
)
plt.legend(["Analytical solution", "Solution by ODEDiscipline"])
plt.show()


plt.plot(
    ode_res_termination["times"],
    analytic_res_termination,
    "r",
    label="Analytical solution",
)
plt.plot(
    ode_res_termination["times"],
    ode_res_termination["position_trajectory"],
    "b--",
    label="Numerical solution",
)
plt.legend(["Analytical solution", "Solution by ODEDiscipline"])
plt.show()

Shortcut

The class OscillatorDiscipline is available in the package gemseo.problems.ode, so it just needs to be imported to be used.

ode_oscillator_discipline = OscillatorDiscipline(
    omega=1.0, times=linspace(0.0, 10.0, 51)
)
ode_oscillator_discipline.execute()

{'time': array([0.]), 'position': array([0.]), 'velocity': array([1.]), 'position_final': array([-0.54402111]), 'velocity_final': array([-0.83907153]), 'termination_time': array([10.]), 'times': array([ 0. ,  0.2,  0.4,  0.6,  0.8,  1. ,  1.2,  1.4,  1.6,  1.8,  2. ,
        2.2,  2.4,  2.6,  2.8,  3. ,  3.2,  3.4,  3.6,  3.8,  4. ,  4.2,
        4.4,  4.6,  4.8,  5. ,  5.2,  5.4,  5.6,  5.8,  6. ,  6.2,  6.4,
        6.6,  6.8,  7. ,  7.2,  7.4,  7.6,  7.8,  8. ,  8.2,  8.4,  8.6,
        8.8,  9. ,  9.2,  9.4,  9.6,  9.8, 10. ])}