.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "examples/ode/ode_discipline/plot_springs_discipline.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_examples_ode_ode_discipline_plot_springs_discipline.py: Solve a system of coupled ODEs ============================== .. GENERATED FROM PYTHON SOURCE LINES 23-45 .. code-block:: Python from __future__ import annotations from typing import TYPE_CHECKING from matplotlib import pyplot as plt from numpy import array from numpy import linspace from scipy.interpolate import interp1d from gemseo.core.chains.chain import MDOChain from gemseo.core.discipline import Discipline from gemseo.disciplines.auto_py import AutoPyDiscipline from gemseo.disciplines.ode.ode_discipline import ODEDiscipline from gemseo.mda.gauss_seidel import MDAGaussSeidel from gemseo.problems.ode.springs.coupled_springs_generator import ( CoupledSpringsGenerator, ) if TYPE_CHECKING: from gemseo.typing import StrKeyMapping .. GENERATED FROM PYTHON SOURCE LINES 46-79 This tutorial describes how to use the :class:`.ODEDiscipline` with coupled ODEs. Problem description ------------------- Consider a set of point masses with masses :math:`m_1,\ m_2,...\ m_n` connected by springs with stiffnesses :math:`k_1,\ k_2,...\ k_{n+1}`. The springs at each end of the system are connected to fixed points. We hereby study the response of the system to the displacement of one or multiple of the point masses. .. figure:: /_images/ode/springs.png :width: 400 :alt: Illustration of the springs-masses problem. The motion of each point mass in this system is described by the following set of ordinary differential equations (ODEs): .. math:: \begin{cases} \frac{dx_i}{dt} &= v_i \\ \frac{dv_i}{dt} &= - \frac{k_i + k_{i+1}}{m_i}x_i + \frac{k_i}{m_i}x_{i-1} + \frac{k_{i+1}}{m_i}x_{i+1} \end{cases} where :math:`x_i` is the position of the :math:`i`-th point mass and :math:`v_i` is its velocity. These equations are coupled, since the forces applied to any given mass depend on the positions of its neighbors. In this tutorial, we will use the framework of the :class:`.ODEDisciplines` to solve this set of coupled equations. .. GENERATED FROM PYTHON SOURCE LINES 81-87 Using coupled instances of :class:`.ODEDiscipline` to solve the problem ----------------------------------------------------------------------- Let us consider the problem described above in the case of two masses. First we describe the right-hand side (RHS) function of the equations of motion for each point mass. .. GENERATED FROM PYTHON SOURCE LINES 87-101 .. code-block:: Python ode_solver_name = "RK45" stiffness_0 = 1 stiffness_1 = 1 stiffness_2 = 1 mass_0 = 1 mass_1 = 1 initial_position_0 = 1 initial_position_1 = 0 initial_velocity_0 = 0 initial_velocity_1 = 0 .. GENERATED FROM PYTHON SOURCE LINES 102-103 We define the times at which to discretize the trajectories. .. GENERATED FROM PYTHON SOURCE LINES 103-105 .. code-block:: Python times = linspace(0.0, 2.0, 30) .. GENERATED FROM PYTHON SOURCE LINES 106-108 We define two disciplines to compute the RHS of the ODEs describing the dynamics of each mass. .. GENERATED FROM PYTHON SOURCE LINES 108-184 .. code-block:: Python class RHSMassDisciplineLeft(Discipline): def __init__(self, **kwargs) -> None: input_names = ("time", "position_0", "velocity_0", "position_1") output_names = ("position_0_dot", "velocity_0_dot") super().__init__(**kwargs) self.io.input_grammar.update_from_names(input_names) self.io.output_grammar.update_from_names(output_names) self.default_input_data = { "time": 0.0, "position_0": array([initial_position_0]), "velocity_0": array([initial_velocity_0]), "position_1": times * 0.0, } self.add_differentiated_inputs(["position_0", "velocity_0"]) def _run(self, input_data: StrKeyMapping): time = self.io.data["time"] position_0 = self.io.data["position_0"] velocity_0 = self.io.data["velocity_0"] position_1_vec = self.io.data["position_1"] interp_function = interp1d(times, position_1_vec, assume_sorted=True) position_1 = interp_function(time) position_0_dot = velocity_0 velocity_0_dot = ( -(stiffness_0 + stiffness_1) * position_0 + stiffness_1 * position_1 ) / mass_0 self.local_data["position_0_dot"] = position_0_dot self.local_data["velocity_0_dot"] = velocity_0_dot class RHSMassDisciplineRight(Discipline): def __init__(self, **kwargs) -> None: input_names = ("time", "position_1", "velocity_1", "position_0") output_names = ("position_1_dot", "velocity_1_dot") super().__init__(**kwargs) self.io.input_grammar.update_from_names(input_names) self.io.output_grammar.update_from_names(output_names) self.default_input_data = { "time": 0.0, "position_1": array([initial_position_1]), "velocity_1": array([initial_velocity_1]), "position_0": times * 0.0, } self.add_differentiated_inputs(["position_1", "velocity_1"]) def _run(self, input_data: StrKeyMapping): time = self.io.data["time"] position_1 = self.io.data["position_1"] velocity_1 = self.io.data["velocity_1"] position_0_vec = self.io.data["position_0"] interp_function = interp1d(times, position_0_vec, assume_sorted=True) position_0 = interp_function(time) position_1_dot = velocity_1 velocity_1_dot = ( -(stiffness_0 + stiffness_1) * position_1 + stiffness_0 * position_0 ) / mass_1 self.local_data["position_1_dot"] = position_1_dot self.local_data["velocity_1_dot"] = velocity_1_dot .. GENERATED FROM PYTHON SOURCE LINES 185-187 We can then create a list of :class:`ODEDiscipline` objects .. GENERATED FROM PYTHON SOURCE LINES 187-206 .. code-block:: Python rhs_disciplines = [RHSMassDisciplineLeft(), RHSMassDisciplineRight()] ode_disciplines = [ ODEDiscipline( rhs_discipline=rhs_discipline, times=times, state_names=(f"position_{i}", f"velocity_{i}"), time_name="time", return_trajectories=True, ode_solver_name=ode_solver_name, rtol=1e-6, atol=1e-6, ) for i, rhs_discipline in enumerate(rhs_disciplines) ] for ode_discipline in ode_disciplines: ode_discipline.execute() .. GENERATED FROM PYTHON SOURCE LINES 207-208 We apply an MDA with the Gauss-Seidel algorithm: .. GENERATED FROM PYTHON SOURCE LINES 208-211 .. code-block:: Python mda = MDAGaussSeidel(ode_disciplines) local_data = mda.execute() .. GENERATED FROM PYTHON SOURCE LINES 212-218 The coupling structure between the instances of :class:`.ODEDiscipline` can be represented by the following picture. .. image:: /_images/ode/springs-disciplines.png :width: 400 :alt: Couple, then integrate. .. GENERATED FROM PYTHON SOURCE LINES 220-221 We can plot the residuals of this MDA. .. GENERATED FROM PYTHON SOURCE LINES 221-230 .. code-block:: Python mda.plot_residual_history() plt.plot(times, local_data["position_0"], label="mass 0") plt.plot(times, local_data["position_1"], label="mass 1") plt.title("Coupling between two ODEDisciplines") plt.legend() plt.show() .. image-sg:: /examples/ode/ode_discipline/images/sphx_glr_plot_springs_discipline_001.png :alt: Coupling between two ODEDisciplines :srcset: /examples/ode/ode_discipline/images/sphx_glr_plot_springs_discipline_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 231-242 Using a single :class:`.ODEDiscipline` with coupled dynamics ------------------------------------------------------------ In the previous section, we considered the integration in time within each ODE discipline, then coupled the disciplines, as illustrated in the next figure. Another possibility to tackle this problem is to define the couplings within a discipline, as illustrated in the next figure. .. image:: /_images/ode/time_integration.png :width: 400 :alt: Couple, then integrate. .. GENERATED FROM PYTHON SOURCE LINES 242-295 .. code-block:: Python def compute_mass_0_rhs( time=0, position_0=initial_position_0, velocity_0=initial_velocity_0, position_1=initial_position_1, ): """Compute the RHS of the ODE associated with the first point mass. Args: time: The time at which to evaluate the RHS. position_0: The position of the first point mass at this time. velocity_0: The velocity of the first point mass at this time. position_1: The position of the second point mass at this time. Returns: The first- and second-order derivatives of the position of the first point mass. """ position_0_dot = velocity_0 velocity_0_dot = ( -(stiffness_0 + stiffness_1) * position_0 + stiffness_1 * position_1 ) / mass_0 return position_0_dot, velocity_0_dot def compute_mass_1_rhs( time=0, position_1=initial_position_1, velocity_1=initial_velocity_1, position_0=initial_position_0, ): """Compute the RHS of the ODE associated with the second point mass. Args: time: The time at which to evaluate the RHS. position_1: The position of the second point mass at this time. velocity_1: The velocity of the second point mass at this time. position_0: The position of the first point mass at this time. Returns: The first- and second-order derivatives of the position of the second point mass. """ position_1_dot = velocity_1 velocity_1_dot = ( -(stiffness_1 + stiffness_2) * position_1 + stiffness_1 * position_0 ) / mass_1 return position_1_dot, velocity_1_dot .. GENERATED FROM PYTHON SOURCE LINES 296-298 To do so, we can use the RHS disciplines we created earlier to define an :class:`.MDOChain`. .. GENERATED FROM PYTHON SOURCE LINES 298-308 .. code-block:: Python rhs_disciplines = [ AutoPyDiscipline(py_func=compute_rhs) for compute_rhs in [compute_mass_0_rhs, compute_mass_1_rhs] ] rhs_disciplines[0].add_differentiated_inputs(["time", "position_0", "velocity_0"]) rhs_disciplines[1].add_differentiated_inputs(["time", "position_1", "velocity_1"]) mda = MDOChain(rhs_disciplines) .. GENERATED FROM PYTHON SOURCE LINES 309-310 We then define the ODE discipline that contains the couplings and execute it. .. GENERATED FROM PYTHON SOURCE LINES 310-334 .. code-block:: Python ode_discipline = ODEDiscipline( rhs_discipline=mda, state_names={ "position_0": "position_0_dot", "velocity_0": "velocity_0_dot", "position_1": "position_1_dot", "velocity_1": "velocity_1_dot", }, return_trajectories=True, times=times, ode_solver_name=ode_solver_name, rtol=1e-12, atol=1e-12, ) local_data = ode_discipline.execute() plt.plot(times, local_data["position_0"], label="mass 0") plt.plot(times, local_data["position_1"], label="mass 1") plt.title("Coupling inside the RHS Discipline") plt.legend() plt.show() .. image-sg:: /examples/ode/ode_discipline/images/sphx_glr_plot_springs_discipline_002.png :alt: Coupling inside the RHS Discipline :srcset: /examples/ode/ode_discipline/images/sphx_glr_plot_springs_discipline_002.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 335-341 Shortcut -------- The :mod:`.springs` module provides shortcuts to access this problem. The user can define a list of masses, stiffnesses and initial positions, then create all the disciplines with a single call using the class :class:`.CoupledSpringsGenerator`. .. GENERATED FROM PYTHON SOURCE LINES 341-350 .. code-block:: Python masses = [1, 2, 1] stiffnesses = [1, 1, 1, 1] positions = [1, 0, 0] springs_and_masses = CoupledSpringsGenerator( masses=masses, stiffnesses=stiffnesses, times=times ) .. GENERATED FROM PYTHON SOURCE LINES 351-355 The method 'coupled_ode_disciplines' of :class:`.CoupledSpringsGenerator` creates a list of :class:`.ODEDiscipline`. Each discipline is used to compute the movement of a single mass. The disciplines can be coupled with one another by an MDA. .. GENERATED FROM PYTHON SOURCE LINES 355-367 .. code-block:: Python disciplines = springs_and_masses.create_coupled_ode_disciplines( atol=1e-8, ode_solver_name=ode_solver_name ) mda_shortcut = MDAGaussSeidel(disciplines) mda_result = mda_shortcut.execute({ "initial_position_0": array([positions[0]]), "initial_position_1": array([positions[1]]), "initial_position_2": array([positions[2]]), }) .. GENERATED FROM PYTHON SOURCE LINES 368-372 The method :meth:`.CoupledSpringsGenerator.discipline_with_coupled_dynamic` of :class:`.CoupledSpringsGenerator` returns a single instance of :class:`.ODEDiscipline`, whose dynamic is defined by an MDA on the dynamics of all masses in the system. .. GENERATED FROM PYTHON SOURCE LINES 372-382 .. code-block:: Python ode_discipline_shortcut = springs_and_masses.create_discipline_with_coupled_dynamics( ode_solver_name=ode_solver_name ) ode_result = ode_discipline_shortcut.execute({ "initial_position_0": array([positions[0]]), "initial_position_1": array([positions[1]]), "initial_position_2": array([positions[2]]), }) .. GENERATED FROM PYTHON SOURCE LINES 383-387 We can plot and compare the results. The trajectories of three masses computed by coupling three instances of :class:`.ODEDiscipline`. .. GENERATED FROM PYTHON SOURCE LINES 387-395 .. code-block:: Python plt.plot(times, mda_result["position_0"], label="Mass 0") plt.plot(times, mda_result["position_1"], label="Mass 1") plt.plot(times, mda_result["position_2"], label="Mass 2") plt.title("Coupling between three ODEDisciplines") plt.legend() plt.show() .. image-sg:: /examples/ode/ode_discipline/images/sphx_glr_plot_springs_discipline_003.png :alt: Coupling between three ODEDisciplines :srcset: /examples/ode/ode_discipline/images/sphx_glr_plot_springs_discipline_003.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 396-398 The trajectories of three masses computed by one instance of :class:`.ODEDiscipline` with a dynamic defined by three coupled disciplines. .. GENERATED FROM PYTHON SOURCE LINES 398-406 .. code-block:: Python plt.plot(times, ode_result["position_0"], label="Mass 0") plt.plot(times, ode_result["position_1"], label="Mass 1") plt.plot(times, ode_result["position_2"], label="Mass 2") plt.title("Triple coupling inside the RHS Discipline") plt.legend() plt.show() .. image-sg:: /examples/ode/ode_discipline/images/sphx_glr_plot_springs_discipline_004.png :alt: Triple coupling inside the RHS Discipline :srcset: /examples/ode/ode_discipline/images/sphx_glr_plot_springs_discipline_004.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 407-408 Absolute value of the difference between the two methods. .. GENERATED FROM PYTHON SOURCE LINES 408-419 .. code-block:: Python error_position_0 = abs(mda_result["position_0"] - ode_result["position_0"]) error_position_1 = abs(mda_result["position_1"] - ode_result["position_1"]) error_position_2 = abs(mda_result["position_2"] - ode_result["position_2"]) plt.plot(times, error_position_0, label="Mass 0") plt.plot(times, error_position_1, label="Mass 1") plt.plot(times, error_position_2, label="Mass 2") plt.title("Difference between the two coupling strategies") plt.legend() plt.show() .. image-sg:: /examples/ode/ode_discipline/images/sphx_glr_plot_springs_discipline_005.png :alt: Difference between the two coupling strategies :srcset: /examples/ode/ode_discipline/images/sphx_glr_plot_springs_discipline_005.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 1.139 seconds) .. _sphx_glr_download_examples_ode_ode_discipline_plot_springs_discipline.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: plot_springs_discipline.ipynb ` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: plot_springs_discipline.py ` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: plot_springs_discipline.zip ` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_