diff --git a/docs/source/Support/bskReleaseNotes.rst b/docs/source/Support/bskReleaseNotes.rst index 69473b4fa4..e10bca109e 100644 --- a/docs/source/Support/bskReleaseNotes.rst +++ b/docs/source/Support/bskReleaseNotes.rst @@ -33,6 +33,10 @@ Version |release| - Added a new github workflow job ``canary`` to routinely check the compatibility of latest python dependencies with python 3.13 on the latest mac-os. - Fixed a bug in :ref:`spiceInterface` where multiple instances of the module were not properly managing SPICE kernel references, leading to potential conflicts and data corruption. - Deprecated :ref:`SpacecraftSystem`. It was never completed and we have other ways to connect spacecraft components +- Support for stochastic dynamics (dynamics driven by Stochastic Differential Equations). Added two stochastic integrators: + :ref:`svStochIntegratorW2Ito1` and :ref:`svStochIntegratorW2Ito2`. These are 2-weak-order stochastic integrators. +- Added :ref:`scenarioStochasticDrag`, which illustrates how to use a state driven by stochastic dynamics to + model randomly evolving atmospheric density (and thus drag force). Version 2.7.0 (April 20, 2025) diff --git a/examples/_default.rst b/examples/_default.rst index 91427bf00c..bcf0ff208d 100644 --- a/examples/_default.rst +++ b/examples/_default.rst @@ -126,7 +126,7 @@ Atmospheric Drag :maxdepth: 1 Satellite Drag Deorbit about Earth - + Satellite Drag with Stochastic (Random) Density Access to Communication Locations ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/examples/mujoco/scenarioStochasticDrag.py b/examples/mujoco/scenarioStochasticDrag.py new file mode 100644 index 0000000000..40b2e2b8b8 --- /dev/null +++ b/examples/mujoco/scenarioStochasticDrag.py @@ -0,0 +1,470 @@ +# +# ISC License +# +# Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder +# +# Permission to use, copy, modify, and/or distribute this software for any +# purpose with or without fee is hereby granted, provided that the above +# copyright notice and this permission notice appear in all copies. +# +# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. +""" +This scenario uses :ref:`MJScene` as its ``DynamicObject``, although it is +not its main focus. For context on dynamics using :ref:`MJScene`, check out: + +#. ``examples/mujoco/scenarioReactionWheel.py`` + +This scripts shows how to simulate a point-mass spacecraft using :ref:`MJScene` +(mujoco-based dynamics) with cannonball drag force and a stochastic atmospheric +density model. The script illustrates how one defines and handles dynamic states +that are driven by stochastic terms. + +The spacecraft definition is given in ``CANNONBALL_SCENE_XML``; it contains a single +body for which we set its mass directly. We use the :ref:`NBodyGravity` model to compute +the gravity acting on this body due to a point-mass Earth. + +The drag force on acting on the body is computed in the ``CannonballDrag`` model +defined in this scenario. This takes as input the state of the spacecraft (so that +its velocity and orientation) and the atmospheric density. It outputs a force vector +on the body-fixed reference frame (more accurately, the body's center of mass frame/site). + +The atmospheric density used on the drag model is given by:: + + densityStochastic = densityExponential * (1 + densityCorrection) + +where ``densityExponential`` is computed by the :ref:`exponentialAtmosphere` model +while ``densityCorrection`` is a stochastic process centered at zero. This process +is modeled as an Ornstein-Uhlenbeck (OU) process, whose Stochastic Differential +Equation (SDE) is given by:: + + d(densityCorrection) = -theta*densityCorrection*dt + sigma*dW + +where ``theta`` and ``sigma`` are the terms that characterize the OU process. +In the SDE above, ``-theta*densityCorrection`` represents the 'drift' term (the +classical derivative). ``sigma``, on the other hand, represents the 'diffusion' +term, which maps the influence of the noise to the state. + +The ``densityStochastic`` and ``densityCorrection`` are computed in +``StochasticAtmosphere``. ``densityCorrection`` must be modeled as a stochastic +dynamic state (since its evolution is given by its derivative and noise diffusion). + +Illustration of Simulation Results +---------------------------------- +The following images illustrate a possible simulation result. + +The orbit is plotted in the orbital plane: + +.. image:: /_images/Scenarios/scenarioStochasticDrag_orbit.svg + :align: center + +The altitude as a function of time is plotted. + +.. image:: /_images/Scenarios/scenarioStochasticDrag_altitude.svg + :align: center + +The atmospheric density as a function of altitude is plotted in lin-log space. +This shows two lines: the deterministic, exponential density (should appear +linear); and the stochastic density. + +.. image:: /_images/Scenarios/scenarioStochasticDrag_density.svg + :align: center + +The atmospheric density correction, which should have a standard deviation +of 0.15. + +.. image:: /_images/Scenarios/scenarioStochasticDrag_densityDiff.svg + :align: center + +The magnitude of drag force over time is plotted in lin-log space. + +.. image:: /_images/Scenarios/scenarioStochasticDrag_drag.svg + :align: center + +""" + +from Basilisk.architecture import sysModel +from Basilisk.architecture import messaging +from Basilisk.simulation import mujoco +from Basilisk.simulation import StatefulSysModel +from Basilisk.simulation import dynParamManager +from Basilisk.simulation import svIntegrators +from Basilisk.simulation import pointMassGravityModel +from Basilisk.simulation import NBodyGravity +from Basilisk.simulation import exponentialAtmosphere +from Basilisk.utilities import SimulationBaseClass +from Basilisk.utilities import macros +from Basilisk.utilities import orbitalMotion +from Basilisk.utilities import simIncludeGravBody +from Basilisk.utilities import simSetPlanetEnvironment +from Basilisk.utilities import RigidBodyKinematics as rbk + +import numpy as np +import matplotlib.pyplot as plt + +"""An cannon-ball spacecraft. + +We set its inertial properties directly, instead of defining them +through their geometry. The spacecraft has no sub-bodies or any +moving parts. +""" +CANNONBALL_SCENE_XML = r""" + + + + + + + + +""" + +def run(showPlots: bool = False): + """Main function, see scenario description. + + Args: + showPlots (bool, optional): If True, simulation results are plotted and show. + Defaults to False. + """ + initialAlt = 250 # km + planet = simIncludeGravBody.BODY_DATA["earth"] + + # Set up a circular orbit using classical orbit elements + oe = orbitalMotion.ClassicElements() + oe.a = planet.radEquator + initialAlt * 1000 # meters + oe.e = 0 + oe.i = 33.3 * macros.D2R + oe.Omega = 48.2 * macros.D2R + oe.omega = 347.8 * macros.D2R + oe.f = 85.3 * macros.D2R + rN, vN = orbitalMotion.elem2rv(planet.mu, oe) + oe = orbitalMotion.rv2elem(planet.mu, rN, vN) + orbitPeriod = 2. * np.pi / np.sqrt(planet.mu / oe.a**3) + + dt = 10 # s + tf = 7.1 * orbitPeriod # s + + # Create a simulation, process, and task as usual + scSim = SimulationBaseClass.SimBaseClass() + process = scSim.CreateNewProcess("test") + process.addTask(scSim.CreateNewTask("test", macros.sec2nano(dt))) + + # Create the Mujoco scene (our MuJoCo DynamicObject) + # Load the XML file that defines the system from a file + scene = mujoco.MJScene(CANNONBALL_SCENE_XML) + scSim.AddModelToTask("test", scene) + + # Set the integrator of the DynamicObject to W2Ito2 + # This is a stochastic integrator, necessary since we have + # states (related to density) that are driven by stochastic + # dynamics + integ = svIntegrators.svStochIntegratorW2Ito2(scene) + scene.setIntegrator(integ) + + ### Get mujoco body and site (point, frame) of interest + body: mujoco.MJBody = scene.getBody("ball") + dragPoint: mujoco.MJSite = body.getCenterOfMass() + + ### Create the NBodyGravity model + # add model to the dynamics task of the scene + gravity = NBodyGravity.NBodyGravity() + scene.AddModelToDynamicsTask(gravity) + + # Create a point-mass gravity source + gravityModel = pointMassGravityModel.PointMassGravityModel() + gravityModel.muBody = planet.mu + gravity.addGravitySource("earth", gravityModel, True) + + # Create a gravity target from the mujoco body + gravity.addGravityTarget("ball", body) + + ### Create the density model + atmo = exponentialAtmosphere.ExponentialAtmosphere() + atmo.ModelTag = "ExpAtmo" + simSetPlanetEnvironment.exponentialAtmosphere(atmo, "earth") + atmo.addSpacecraftToModel(dragPoint.stateOutMsg) + + # Will be updated with the task period + scSim.AddModelToTask("test", atmo) + + stochasticAtmo = StochasticAtmosphere( + densityCorrectionStandardDeviation=0.15, + densityCorrectionTheta=.01 + ) + stochasticAtmo.ModelTag = "StochasticExpAtmo" + + stochasticAtmo.atmoDensInMsg.subscribeTo( atmo.envOutMsgs[0] ) + + # The StochasticAtmosphere is a StatefulSysModel with a state + # with both drift (regular derivative) and diffusion (stochastic) + # so we need to add to both DynamicsTask and DiffusionDynamicsTask + scene.AddModelToDynamicsTask(stochasticAtmo) + scene.AddModelToDiffusionDynamicsTask(stochasticAtmo) + + ### Create drag force model + drag = CannonballDrag( + projectedArea=10, + dragCoeff=2.2 + ) + drag.ModelTag = "DragEff" + + drag.frameInMsg.subscribeTo( dragPoint.stateOutMsg ) + drag.atmoDensInMsg.subscribeTo( stochasticAtmo.atmoDensOutMsg ) + + # The CannonballDrag model will produce a force vector message. + # We must connect it to an actuator so that the force is applied + dragActuator: mujoco.MJForceActuator = scene.addForceActuator( "dragForce", dragPoint ) + dragActuator.forceInMsg.subscribeTo( drag.forceOutMsg ) + + # Must be added to the dynamics task because it impacts the + # dynamics of the scene (sets a force on a body) + scene.AddModelToDynamicsTask(drag) + + ### Add recorders + # Record the state of the 'ball' body through + # the ``stateOutMsg`` of its 'origin' site (i.e. frame). + bodyStateRecorder = body.getOrigin().stateOutMsg.recorder() + scSim.AddModelToTask("test", bodyStateRecorder) + + deterministicDensityRecorder = stochasticAtmo.atmoDensInMsg.recorder() + scSim.AddModelToTask("test", deterministicDensityRecorder) + + densityRecorder = stochasticAtmo.atmoDensOutMsg.recorder() + scSim.AddModelToTask("test", densityRecorder) + + dragRecorder = drag.forceOutMsg.recorder() + scSim.AddModelToTask("test", dragRecorder) + + # Initialize the simulation + scSim.InitializeSimulation() + + # Set initial position and velocity + body.setPosition(rN) + body.setVelocity(vN) + + # Run the simulation + scSim.ConfigureStopTime(macros.sec2nano(tf)) + scSim.ExecuteSimulation() + + figures = plotOrbits( + timeAxis=bodyStateRecorder.times(), + posData=bodyStateRecorder.r_BN_N, + velData=bodyStateRecorder.v_BN_N, + dragForce=dragRecorder.force_S, + deterministicDenseData=deterministicDensityRecorder.neutralDensity, + denseData=densityRecorder.neutralDensity, + oe=oe, + mu=planet.mu, + planetRadius=planet.radEquator + ) + + if showPlots: + plt.show() + + return figures + +def plotOrbits(timeAxis, posData, velData, dragForce, deterministicDenseData, denseData, oe, mu, planetRadius): + """ + Plot the results of the stochastic drag simulation, including orbit, altitude, + density, density difference, and drag force over time. + + Args: + timeAxis: Array of time values. + posData: Position data array. + velData: Velocity data array. + dragForce: Drag force data array. + deterministicDenseData: Deterministic atmospheric density data. + denseData: Stochastic atmospheric density data. + oe: Classical orbital elements object. + mu: Gravitational parameter. + planetRadius: Radius of the planet. + Returns: + Dictionary of matplotlib figure objects. + """ + # draw the inertial position vector components + figureList = {} + + figureList["orbit"], ax = plt.subplots() + ax.axis('equal') + planetColor = '#008800' + ax.add_artist(plt.Circle((0, 0), planetRadius / 1000, color=planetColor)) + # draw the actual orbit + rData = [] + fData = [] + for idx in range(0, len(posData)): + oeData = orbitalMotion.rv2elem(mu, posData[idx], velData[idx]) + rData.append(oeData.rmag) + fData.append(oeData.f + oeData.omega - oe.omega) + ax.plot(rData * np.cos(fData) / 1000, rData * np.sin(fData) / 1000, color='#aa0000', linewidth=1.0 + ) + ax.set_xlabel('$i_e$ Cord. [km]') + ax.set_ylabel('$i_p$ Cord. [km]') + + # draw altitude as a function of time + figureList["altitude"], ax = plt.subplots() + ax.ticklabel_format(useOffset=False, style='plain') + alt = (np.array(rData) - planetRadius) / 1000 + ax.plot(timeAxis * macros.NANO2HOUR, alt) + ax.set_xlabel('$t$ [h]') + ax.set_ylabel('Altitude [km]') + + # draw density as a function of altitude + figureList["density"], ax = plt.subplots() + ax.semilogy(alt, denseData, label="Stochastic") + ax.semilogy(alt, deterministicDenseData, label="Exponential") + ax.legend(loc="upper right") + ax.set_xlabel('Altitude [km]') + ax.set_ylabel('$\\rho$ [kg/m$^3$]') + + # draw density as a function of altitude + figureList["densityDiff"], ax = plt.subplots() + ax.plot(timeAxis * macros.NANO2HOUR, (denseData / deterministicDenseData) - 1) + ax.set_xlabel('Time [hr]') + ax.set_ylabel(r'$(\rho_{stoch} / \rho_{exp}) - 1$ [-]') + + # draw drag as a function of time + figureList["drag"], ax = plt.subplots() + ax.semilogy(timeAxis * macros.NANO2HOUR, np.linalg.norm(dragForce, 2, 1)) + ax.set_xlabel('$t$ [hr]') + ax.set_ylabel('$|F_drag|$ [N]') + + return figureList + +class StochasticAtmosphere(StatefulSysModel.StatefulSysModel): + """This model takes in an ``AtmoPropsMsg`` and outputs another + ``AtmoPropsMsg`` with the density perturbed by some stochastic correction + term. + + The output density is given by: + + densityOut = densityIn * (1 + densityCorrection) + + where ``densityCorrection`` is thus the relative density correction. This is + modeled as an Ornstein-Uhlenbeck (OU) process: + + d(densityCorrection) = -theta*densityCorrection*dt + sigma*dW + + where ``theta`` and ``sigma`` are the terms that characterize the OU + process. In the constructor of this class, ``theta`` is set through + ``densityCorrectionTheta``, while ``sigma`` is computed from + the constructor input ``densityCorrectionStandardDeviation``. + """ + + def __init__(self, densityCorrectionStandardDeviation: float, densityCorrectionTheta: float): + """ + Initialize the StochasticAtmosphere model. + + Args: + densityCorrectionStandardDeviation (float): Stationary standard deviation + of the Ornstein-Uhlenbeck process that drives the relative atmospheric density + correction. + densityCorrectionTheta (float): ``theta`` term of the + Ornstein-Uhlenbeck process that drives the relative atmospheric density correction. + """ + super().__init__() + self.densityCorrectionStandardDeviation = densityCorrectionStandardDeviation + self.densityCorrectionTheta = densityCorrectionTheta + + self.atmoDensInMsg = messaging.AtmoPropsMsgReader() + self.atmoDensOutMsg = messaging.AtmoPropsMsg() + + @property + def densityCorrectionSigma(self): + """ + Compute the ``sigma`` term of the Ornstein-Uhlenbeck process that + drives the atmospheric density correction. + + Returns: + float: The sigma value for the OU process. + """ + return self.densityCorrectionStandardDeviation * np.sqrt(2*self.densityCorrectionTheta) + + def registerStates(self, registerer: StatefulSysModel.DynParamRegisterer): + """ + Register the stochastic state with the simulation's state manager. + + Args: + registerer: The DynParamRegisterer used to register states. + """ + self.densityCorrectionState: dynParamManager.StateData = registerer.registerState(1, 1, "densityCorrection") + self.densityCorrectionState.setNumNoiseSources(1) + + def UpdateState(self, CurrentSimNanos: int): + """ + Update the state at each integrator step, computing both + the drift and diffusion terms of the density correction. + Also reads the input atmospheric density and writes the + perturbed output density. + + Args: + CurrentSimNanos (int): Current simulation time in nanoseconds. + """ + densityCorrection = self.densityCorrectionState.getState()[0][0] + + # See class docstring for the dynamic equations driving + # this correction term + self.densityCorrectionState.setDerivative([[ -self.densityCorrectionTheta*densityCorrection ]]) + self.densityCorrectionState.setDiffusion([[ self.densityCorrectionSigma ]], index=0) + + atmoDensIn: messaging.AtmoPropsMsgPayload = self.atmoDensInMsg() + + atmoDensOut = messaging.AtmoPropsMsgPayload() + atmoDensOut.neutralDensity = atmoDensIn.neutralDensity * (1 + densityCorrection) + atmoDensOut.localTemp = atmoDensIn.localTemp + self.atmoDensOutMsg.write(atmoDensOut, CurrentSimNanos, self.moduleID) + + +class CannonballDrag(sysModel.SysModel): + """ + Implements a cannonball drag force model for a spacecraft. + Computes the drag force based on atmospheric density and spacecraft velocity. + """ + + def __init__(self, projectedArea: float, dragCoeff: float): + """ + Initialize the CannonballDrag model. + + Args: + projectedArea (float): Projected area of the spacecraft [m^2]. + dragCoeff (float): Drag coefficient (dimensionless). + """ + super().__init__() + self.projectedArea = projectedArea + self.dragCoeff = dragCoeff + + self.atmoDensInMsg = messaging.AtmoPropsMsgReader() + self.frameInMsg = messaging.SCStatesMsgReader() + + self.forceOutMsg = messaging.ForceAtSiteMsg() + + def UpdateState(self, CurrentSimNanos: int): + """ + Update the drag force at each integrator step. + + Args: + CurrentSimNanos (int): Current simulation time in nanoseconds. + """ + density = self.atmoDensInMsg().neutralDensity + + # N frame: inertial frame + # B frame: body-fixed frame + frame: messaging.SCStatesMsgPayload = self.frameInMsg() + dcm_BN = rbk.MRP2C(frame.sigma_BN) + v_BN_B = dcm_BN @ frame.v_BN_N + vNorm = np.linalg.norm(v_BN_B) + v_hat_BN_B = v_BN_B / vNorm + + dragForce_B = 0.5 * self.dragCoeff * vNorm**2 * self.projectedArea * density * (-v_hat_BN_B) + + payload = messaging.ForceAtSiteMsgPayload() + payload.force_S = dragForce_B + self.forceOutMsg.write(payload, CurrentSimNanos, self.moduleID) + + +if __name__ == "__main__": + run(True) diff --git a/src/simulation/dynamics/Integrators/_UnitTest/test_stochasticIntegrators.py b/src/simulation/dynamics/Integrators/_UnitTest/test_stochasticIntegrators.py new file mode 100644 index 0000000000..8b00248e2f --- /dev/null +++ b/src/simulation/dynamics/Integrators/_UnitTest/test_stochasticIntegrators.py @@ -0,0 +1,890 @@ +# +# ISC License +# +# Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder +# +# Permission to use, copy, modify, and/or distribute this software for any +# purpose with or without fee is hereby granted, provided that the above +# copyright notice and this permission notice appear in all copies. +# +# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. +""" +Tests the Euler-Mayurama and SRK integrator described in: + +Tang, X., Xiao, A. Efficient weak second-order stochastic Runge-Kutta methods +for Itô stochastic differential equations. Bit Numer Math 57, 241-260 (2017). +https://doi.org/10.1007/s10543-016-0618-9 +""" +from __future__ import annotations + +from typing import Callable, List, Literal, get_args, Any +import tqdm +import itertools + +import pytest +import numpy as np +import numpy.typing as npt +import numpy.testing +import matplotlib.pyplot as plt + +from Basilisk.utilities import macros +from Basilisk.utilities import SimulationBaseClass +from Basilisk.utilities import pythonVariableLogger +from Basilisk.simulation import svIntegrators +from Basilisk.simulation import dynParamManager + +try: + from Basilisk.simulation import mujoco + from Basilisk.simulation import StatefulSysModel + couldImportMujoco = True +except: + couldImportMujoco = False + +SRKMethod = Literal["W2Ito1", "W2Ito2"] +Method = Literal["W2Ito1", "W2Ito2", "EulerMayurama"] +METHODS = get_args(Method) + +# Function of the form y = f(t,x) where x and y are vectors of the same size +DynFunc = Callable[[float, npt.NDArray[np.float64]], npt.NDArray[np.float64]] + +# Coefficient sets (from Tang & Xiao Table 2, W2-Ito1 & W2-Ito2) +W2Ito1Coefficients = { + 'alpha': np.array([1/6, 2/3, 1/6]), + 'beta0': np.array([-1, 1, 1]), + 'beta1': np.array([2, 0, -2]), + + 'A0': np.array([[0.0, 0.0, 0.0], + [1/2, 0.0, 0.0], + [-1, 2, 0]]), + + 'B0': np.array([[0.0, 0.0, 0.0], + [(6-np.sqrt(6))/10, 0.0, 0.0], + [(3+2*np.sqrt(6))/5, 0.0, 0.0]]), + + 'A1': np.array([[0.0, 0.0, 0.0], + [1/4, 0.0, 0.0], + [1/4, 0.0, 0.0]]), + + 'B1': np.array([[0.0, 0.0, 0.0], + [1/2, 0.0, 0.0], + [-1/2, 0.0, 0.0]]), + + 'B2': np.array([[0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 0.0, 0.0]]), +} + +W2Ito2Coefficients = { + 'alpha': np.array([1/6, 1/3, 1/3, 1/6]), + 'beta0': np.array([0, -1, 1, 1]), + 'beta1': np.array([0, 2, 0, -2]), + + 'A0': np.array([[0.0, 0.0, 0.0, 0.0], + [1/2, 0.0, 0.0, 0.0], + [0.0, 1/2, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0]]), + + 'B0': np.array([[0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0]]), + + 'A1': np.array([[0.0, 0.0, 0.0, 0.0], + [1/2, 0.0, 0.0, 0.0], + [1/2, 0.0, 0.0, 0.0], + [1/2, 0.0, 0.0, 0.0]]), + + 'B1': np.array([[0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 1/2, 0.0, 0.0], + [0.0, -1/2, 0.0, 0.0]]), + + 'B2': np.array([[0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0]]), +} + +SRK_COEFFICIENTS: dict[SRKMethod, dict[str, npt.NDArray[Any]]] = { + "W2Ito1": W2Ito1Coefficients, + "W2Ito2": W2Ito2Coefficients +} + +def srk2Integrate( + f: DynFunc, + g_list: List[DynFunc], + x0: npt.NDArray[np.float64], + dt: float, + tf: float, + rng_seed: int, + alpha: npt.NDArray[np.float64], + beta0: npt.NDArray[np.float64], + beta1: npt.NDArray[np.float64], + A0: npt.NDArray[np.float64], + B0: npt.NDArray[np.float64], + A1: npt.NDArray[np.float64], + B1: npt.NDArray[np.float64], + B2: npt.NDArray[np.float64], + ): + """ + Generic s-stage SRK integrator for vector SDE: + dX = f(t,X) dt + sum_k g_list[k](t,X) dW_k + + Method described in Tang & Xiao. + Args: + f: Drift function. + g_list: List of diffusion functions. + x0: Initial state. + dt: Time step. + tf: Final time. + rng_seed: Random seed. + alpha, beta0, beta1, A0, B0, A1, B1, B2: SRK coefficients. + Returns: + Array of state trajectories, including time as the first column. + """ + + def wrapped_f(full_state: npt.NDArray[Any]): + t = full_state[0] + x = full_state[1:] + return np.concatenate([[1], f(t, x)]) + + wrapped_g_list = [] + for g in g_list: + def wrapped_g(full_state: npt.NDArray[Any], g: DynFunc = g): + t = full_state[0] + x = full_state[1:] + return np.concatenate([[0], g(t, x)]) + wrapped_g_list.append(wrapped_g) + + s = alpha.size + n = x0.size + m = len(g_list) + + nsteps = int(np.floor(tf / dt)) + + x = np.zeros([nsteps+1, n+1], dtype=float) + x[0,0] = 0 + x[0,1:] = x0 + + rng = svIntegrators.SRKRandomVariableGenerator() + rng.setSeed(rng_seed) + + # throw the first away, similar to how Basilisk + # calls the equationsOfMotion once at the beginning + rng.generate(m, 0) + + for step in range(nsteps): + X = x[step,:].copy() + + # generate random variables + rvs: svIntegrators.SRKRandomVariables = rng.generate(m, dt) + Ik: List[List[float]] = rvs.Ik + Ikl: List[List[float]] = rvs.Ikl + xi: float = rvs.xi + + # allocate stage arrays + H0 = [X.copy() for _ in range(s)] + Hk = [[X.copy() for _ in range(s)] for _ in range(m)] + + for i in range(s): + + # compute H0 stages + sumA = np.zeros(n+1) + sumB = np.zeros(n+1) + for j in range(s): + sumA += A0[i, j] * wrapped_f(H0[j]) * dt + for k in range(m): + sumB += B0[i, j] * wrapped_g_list[k](Hk[k][j]) * Ik[k] + H0[i] = X + sumA + sumB + + # compute Hk stages + for k in range(m): + sumA = np.zeros(n+1) + sumB1 = np.zeros(n+1) + sumB2 = np.zeros(n+1) + for j in range(s): + sumA += A1[i, j] * wrapped_f(H0[j]) * dt + sumB1 += B1[i, j] * wrapped_g_list[k](Hk[k][j]) * xi + for l in range(m): + if l != k: + sumB2 += B2[i, j] * wrapped_g_list[l](Hk[k][j]) * Ikl[k][l] + Hk[k][i] = X + sumA + sumB1 + sumB2 + + # combine increments + drift = np.zeros(n+1) + for i in range(s): + drift += alpha[i] * wrapped_f(H0[i]) * dt + + diffusion = np.zeros(n+1) + for k in range(m): + for i in range(s): + diffusion += beta0[i] * wrapped_g_list[k](Hk[k][i]) * Ik[k] + diffusion += beta1[i] * wrapped_g_list[k](Hk[k][i]) * Ikl[k][k] + + x[step+1,:] = X + drift + diffusion + + return x + +def eulerMayuramaIntegrate( + f: DynFunc, + g_list: List[DynFunc], + x0: npt.NDArray[np.float64], + dt: float, + tf: float, + rng_seed: int + ): + """ + Euler-Mayurama integrator for the vector SDE: + dX = f(t,X) dt + sum_k g_list[k](t,X) dW_k + + Args: + f: Drift function. + g_list: List of diffusion functions. + x0: Initial state. + dt: Time step. + tf: Final time. + rng_seed: Random seed. + Returns: + Array of state trajectories, including time as the first column. + """ + n = x0.size + m = len(g_list) + + nsteps = int(np.floor(tf / dt)) + + t = dt * np.arange(nsteps+1) + x = np.zeros([nsteps+1, n], dtype=float) + x[0,:] = x0 + + rng = svIntegrators.EulerMayuramaRandomVariableGenerator() + rng.setSeed(rng_seed) + + # throw the first away, similar to how Basilisk + # calls the equationsOfMotion once at the beginning + rng.generate(m, 0) + + for step in range(nsteps): + # generate random variables + pseudoTimeSteps = rng.generate(m, dt) + + x[step+1,:] = x[step,:] + f(t[step], x[step,:])*dt + for k, g in enumerate(g_list): + x[step+1,:] += g(t[step], x[step,:])*pseudoTimeSteps[k] + + return np.column_stack([t, x]) + +class ExponentialSystem: + """A simple deterministic system with one state: dx/dt = x*t.""" + + def __init__(self, x0: float = 1): + """Initialize the system with initial state x0.""" + self.x0 = np.array([x0]) + + def f(self, t: float, x: npt.NDArray[np.float64]): + """Drift function for the exponential system.""" + return np.array([t*x[0]]) + + g = [] + +class OrnsteinUhlenbeckSystem: + """A process defined by + + dx = theta*(mu - x)*dt + sigma*dW + """ + def __init__(self, x0: float = .1, mu: float = 0, theta: float = .1, sigma: float = .01): + """Initialize the OU process parameters.""" + self.x0 = np.array([x0]) + self.mu = mu + self.theta = theta + self.sigma = sigma + + def f(self, t: float, x: npt.NDArray[np.float64]): + """Drift function for the OU process.""" + return np.array([self.theta*(self.mu - x[0])]) + + def g1(self, t: float, x: npt.NDArray[np.float64]): + """Diffusion function for the OU process.""" + return np.array([self.sigma]) + + @property + def g(self): + """Return list of diffusion functions.""" + return [self.g1] + + def mean(self, t: float): + """E[x(t)]""" + return np.array([ + self.x0[0]*np.exp(-self.theta*t) + self.mu*(1 - np.exp(-self.theta*t)) + ]) + + def var(self, t: float): + """Var(x(t))""" + return np.array([ + self.sigma**2/(2*self.theta) * (1- np.exp(-self.theta*t)) + ]) + +class ComplexOrnsteinUhlenbeckSystem: + """A process defined by + + dx = -theta*x*dt + sigma_x1*dW_1 + sigma_x2*dW_2 + dy = -theta*y*dt + sigma_y1*dW_1 + sigma_y2*dW_2 + """ + def __init__( + self, + x0: float = .1, y0: float = -.1, + theta_x: float = .1, theta_y: float = .073, + sigma_x1: float = .015, sigma_x2: float = .011, + sigma_y1: float = 0, sigma_y2: float = .029 + ): + """Initialize the complex OU process parameters.""" + self.x0 = np.array([x0, y0]) + self.theta_x = theta_x + self.theta_y = theta_y + self.sigma_x1 = sigma_x1 + self.sigma_x2 = sigma_x2 + self.sigma_y1 = sigma_y1 + self.sigma_y2 = sigma_y2 + + def f(self, t: float, x: npt.NDArray[np.float64]): + """Drift function for the complex OU process.""" + return np.array([-self.theta_x*x[0], -self.theta_y*x[1]]) + + def g1(self, t: float, x: npt.NDArray[np.float64]): + """First diffusion function for the complex OU process.""" + return np.array([self.sigma_x1, self.sigma_y1]) + + def g2(self, t: float, x: npt.NDArray[np.float64]): + """Second diffusion function for the complex OU process.""" + return np.array([self.sigma_x2, self.sigma_y2]) + + @property + def g(self): + """Return list of diffusion functions.""" + return [self.g1, self.g2] + +class Example1System: + """Example 1 dynamical system in: + + Tang, X., Xiao, A. Efficient weak second-order stochastic Runge-Kutta methods + for Itô stochastic differential equations. Bit Numer Math 57, 241-260 (2017). + https://doi.org/10.1007/s10543-016-0618-9 + """ + def __init__(self, y1_0: float = 1, y2_0: float = 1): + """Initialize the Example 1 system.""" + self.x0 = np.array([y1_0, y2_0]) + + def f(self, t: float, x: npt.NDArray[np.float64]): + """Drift function for Example 1 system.""" + y1, y2 = x + return np.array([ + -273/512*y1, + -1/160*y1 - (785/512 - np.sqrt(2)/8)*y2 + ]) + + def g1(self, t: float, x: npt.NDArray[np.float64]): + """First diffusion function for Example 1 system.""" + y1, y2 = x + return np.array([ + 1/4*y1, + (1 - 2*np.sqrt(2)/8)/4*y2 + ]) + + def g2(self, t: float, x: npt.NDArray[np.float64]): + """Second diffusion function for Example 1 system.""" + y1, y2 = x + return np.array([ + 1/16*y1, + 1/10*y1 + 1/16*y2 + ]) + + @property + def g(self): + """Return list of diffusion functions.""" + return [self.g1, self.g2] + +def getBasiliskSim(method: Method, dt: float, x0: npt.NDArray[np.float64], f: DynFunc, g: List[DynFunc], seed: int | None): + """ + Set up and return a Basilisk simulation for a given SDE and integrator method. + + Args: + method: Integration method ("W2Ito1", "W2Ito2", or "EulerMayurama"). + dt: Time step. + x0: Initial state. + f: Drift function. + g: List of diffusion functions. + seed: RNG seed (or None for random). + + Returns: + Tuple of (scSim, stateModel, integratorObject, stateLogger). + """ + + # Declared inside, since StatefulSysModel may be undefined if not running with mujoco + class GenericStochasticStateModel(StatefulSysModel.StatefulSysModel): + + def __init__(self, x0: npt.NDArray[np.float64], f: DynFunc, g: List[DynFunc]): + super().__init__() + self.x0 = x0 + self.f = f + self.g = g + + def registerStates(self, registerer: StatefulSysModel.DynParamRegisterer): + """Called once during InitializeSimulation""" + # We get one noise source per diffusion function g: + m = len(self.g) + n = self.x0.size + + self.states: List[dynParamManager.StateData] = [] + for i in range(n): + self.states.append( registerer.registerState(1, 1, f"y{i+1}") ) + self.states[-1].setNumNoiseSources(m) + self.states[-1].setState([[self.x0[i]]]) + + # We want every noise source to be shared between states + for k in range(m): + registerer.registerSharedNoiseSource([ + (state, k) + for state in self.states + ]) + + def UpdateState(self, CurrentSimNanos: int): + """Called at every integrator step""" + m = len(self.g) + n = self.x0.size + + t = macros.NANO2SEC * CurrentSimNanos + + # Collect current state into a numpy array + x = np.array([state.getState()[0][0] for state in self.states]) + + # Compute f and store in the derivative of the states + deriv = self.f(t, x) + for i in range(n): + self.states[i].setDerivative( [[deriv[i]]] ) + + # Compute g_k and store in the diffusion of the states + for k in range(m): + diff = self.g[k](t, x) + for i in range(n): + self.states[i].setDiffusion( [[diff[i]]], index=k ) + + # Create sim, process, and task + scSim = SimulationBaseClass.SimBaseClass() + dynProcess = scSim.CreateNewProcess("test") + dynProcess.addTask(scSim.CreateNewTask("test", macros.sec2nano(dt))) + + scene = mujoco.MJScene("") # empty scene, no multi-body dynamics + scSim.AddModelToTask("test", scene) + + if method == "W2Ito1": + integratorClass = svIntegrators.svStochIntegratorW2Ito1 + elif method == "W2Ito2": + integratorClass = svIntegrators.svStochIntegratorW2Ito2 + elif method == "EulerMayurama": + integratorClass = svIntegrators.svStochIntegratorMayurama + else: + raise NotImplementedError(method) + + integratorObject = integratorClass(scene) + if seed is not None: + integratorObject.setRNGSeed(seed) + scene.setIntegrator(integratorObject) + + stateModel = GenericStochasticStateModel(x0, f, g) + stateModel.ModelTag = "testModel" + + # The same model computes both the drift and diffusion of the state + # so they must be added to both tasks + scene.AddModelToDynamicsTask(stateModel) + scene.AddModelToDiffusionDynamicsTask(stateModel) + + stateLogger = pythonVariableLogger.PythonVariableLogger( + {"x": lambda _: np.array([state.getState()[0][0] for state in stateModel.states])} + ) + scSim.AddModelToTask("test", stateLogger) + + scSim.InitializeSimulation() + + return scSim, stateModel, integratorObject, stateLogger + +def estimateErrorAndEmpiricalVariance( + computeTrajectory: Callable[[], npt.NDArray[np.float64]], + G: Callable[[npt.NDArray[np.float64]], float], + estimateGOnTrajectory: float, + M1: int, + M2: int, + ): + r"""Computes the error and empirical variance according to the + equations described in Section 4 of Tang & Xiao. + + Args: + computeTrajectory (Callable[[], npt.NDArray[np.float64]]): A function that, + when called, realizes one simulation by forward propagating the system from + t0. The result is the state at the end point. In the paper: ``y_N``. + G (Callable[[npt.NDArray[np.float64]], float]): A differentiable function + that takes in a state vector and outputs a single number. + estimateGOnTrajectory (float): The estimate of the application of the function + G on the random variable that is the last state of the propagated system. + In the paper: ``E[G(y(t_N))]`` + M1 (int): Number of batches. + M2 (int): Number of trajectories per batch. + + Returns: + tuple[float, float]: The integrator error (``\hat{\mu}`` in the paper), and + the empirical variance (``\hat{\sigma}_\mu^2`` in the paper). + """ + + muHat = [0. for _ in range(M1)] + for j, _ in tqdm.tqdm(itertools.product(range(M1), range(M2)), total=M1*M2): + muHat[j] += (G(computeTrajectory()) - estimateGOnTrajectory)/M2 + + muHatAvg = sum(muHat) / M1 + + sigmaMuSquared = 0 + for j in range(M1): + sigmaMuSquared += (muHat[j] - muHatAvg)**2/(M1-1) + + return muHatAvg, sigmaMuSquared + + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +@pytest.mark.parametrize("method", METHODS) +def test_deterministic(method: Method, plot: bool = False): + """ + Test deterministic integration (no diffusion) for all integrator methods. + Compares Basilisk and Python implementations against the analytical solution + for the exponential system dx/dt = x*t. + + Args: + method: Integration method ("W2Ito1", "W2Ito2", or "EulerMayurama"). + plot: If True, plot the relative error. + """ + + dt = .01 + tf = 1 + seed = 10 + + system = ExponentialSystem() + + scSim, stateModel, integratorObject, stateLogger = getBasiliskSim(method, dt, system.x0, system.f, system.g, seed) + scSim.ConfigureStopTime( macros.sec2nano(tf) ) + scSim.ExecuteSimulation() + + tBasilisk = macros.NANO2SEC* stateLogger.times() + xBasilisk = stateLogger.x + + if method == "EulerMayurama": + txPython = eulerMayuramaIntegrate(system.f, system.g, system.x0, dt, tf, seed) + else: + txPython = srk2Integrate(system.f, system.g, system.x0, dt, tf, seed, **SRK_COEFFICIENTS[method]) + tPython = txPython[:,0] + xPython = txPython[:,1] + + xExpected = np.exp( tPython **2 / 2 ) + + if plot: + fig, ax = plt.subplots() + ax.plot(tBasilisk, 100*(xBasilisk - xExpected)/xExpected, label="Basilisk") + ax.plot(tPython, 100*(xPython - xExpected)/xExpected, ls="--", label="Python") + ax.legend() + ax.set_ylabel("Relative Error [%]") + plt.show() + + # The Python and Basilisk implementation should behave identially + numpy.testing.assert_allclose( + xBasilisk[-1], + xPython[-1], + atol=1e-12, + rtol=0 + ) + + if method == "EulerMayurama": + expectedIntegrationError = .05 + elif method == "W2Ito1": + expectedIntegrationError = 1e-6 + elif method == "W2Ito2": + expectedIntegrationError = 1e-8 + + # The Basilisk should have some integration error w.r.t analytical solution + numpy.testing.assert_allclose( + xBasilisk[-1], + xExpected[-1], + atol=expectedIntegrationError, + rtol=0 + ) + + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +@pytest.mark.parametrize("method", METHODS) +def test_ou(method: Method, plot: bool = False): + """ + Test Ornstein-Uhlenbeck process integration for all integrator methods. + Compares Basilisk and Python implementations for a single realization. + + Args: + method: Integration method. + plot: If True, plot the state trajectories. + """ + + dt = 1 + tf = 10 + seed = 10 + + system = OrnsteinUhlenbeckSystem() + + scSim, stateModel, integratorObject, stateLogger = getBasiliskSim(method, dt, system.x0, system.f, system.g, seed) + scSim.ConfigureStopTime( macros.sec2nano(tf) ) + scSim.ExecuteSimulation() + + tBasilisk = macros.NANO2SEC* stateLogger.times() + xBasilisk = stateLogger.x + + if method == "EulerMayurama": + txPython = eulerMayuramaIntegrate(system.f, system.g, system.x0, dt, tf, seed) + else: + txPython = srk2Integrate(system.f, system.g, system.x0, dt, tf, seed, **SRK_COEFFICIENTS[method]) + tPython = txPython[:,0] + xPython = txPython[:,1] + + if plot: + fig, ax = plt.subplots() + ax.plot(tBasilisk, xBasilisk, label="Basilisk") + ax.plot(tPython, xPython, ls="--", label="Python") + ax.legend() + ax.set_ylabel("x") + ax.set_xlabel("t") + plt.show() + + # The Python and Basilisk implementation should behave identially + numpy.testing.assert_allclose( + xBasilisk[-1], + xPython[-1], + atol=1e-12, + rtol=0 + ) + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +@pytest.mark.parametrize("method", METHODS) +def test_ouComplex(method: Method, plot: bool = False): + """ + Test integration of a two-dimensional coupled Ornstein-Uhlenbeck process. + Compares Basilisk and Python implementations for a single realization. + + Args: + method: Integration method. + plot: If True, plot the state trajectories. + """ + + dt = 1 + tf = 10 + seed = 10 + + system = ComplexOrnsteinUhlenbeckSystem() + + scSim, stateModel, integratorObject, stateLogger = getBasiliskSim(method, dt, system.x0, system.f, system.g, seed) + scSim.ConfigureStopTime( macros.sec2nano(tf) ) + scSim.ExecuteSimulation() + + tBasilisk = macros.NANO2SEC* stateLogger.times() + xBasilisk = stateLogger.x + + if method == "EulerMayurama": + txPython = eulerMayuramaIntegrate(system.f, system.g, system.x0, dt, tf, seed) + else: + txPython = srk2Integrate(system.f, system.g, system.x0, dt, tf, seed, **SRK_COEFFICIENTS[method]) + tPython = txPython[:,0] + xPython = txPython[:,1:] + + if plot: + fig, axs = plt.subplots(2, sharex=True) + for i, ax in enumerate(axs): + ax.plot(tBasilisk, xBasilisk[:,i], label="Basilisk") + ax.plot(tPython, xPython[:,i], ls="--", label="Python") + ax.legend() + ax.set_ylabel("xy"[i]) + ax.set_xlabel("t") + plt.show() + + # The Python and Basilisk implementation should behave identially + numpy.testing.assert_allclose( + xBasilisk[-1,:], + xPython[-1,:], + atol=1e-12, + rtol=0 + ) + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +@pytest.mark.parametrize("method", METHODS) +def test_example1(method: Method, plot: bool = False): + """ + Test Example 1 system from Tang & Xiao (2017) for all integrator methods. + Compares Basilisk and Python implementations for a single realization. + + Args: + method: Integration method. + plot: If True, plot the state trajectories. + """ + + dt = 1 + tf = 10 + seed = 10 + + system = Example1System() + + scSim, stateModel, integratorObject, stateLogger = getBasiliskSim(method, dt, system.x0, system.f, system.g, seed) + scSim.ConfigureStopTime( macros.sec2nano(tf) ) + scSim.ExecuteSimulation() + + tBasilisk = macros.NANO2SEC* stateLogger.times() + xBasilisk = stateLogger.x + + if method == "EulerMayurama": + txPython = eulerMayuramaIntegrate(system.f, system.g, system.x0, dt, tf, seed) + else: + txPython = srk2Integrate(system.f, system.g, system.x0, dt, tf, seed, **SRK_COEFFICIENTS[method]) + tPython = txPython[:,0] + xPython = txPython[:,1:] + + if plot: + fig, axs = plt.subplots(2, sharex=True) + for i, ax in enumerate(axs): + ax.plot(tBasilisk, xBasilisk[:,i], label="Basilisk") + ax.plot(tPython, xPython[:,i], ls="--", label="Python") + ax.legend() + ax.set_ylabel(f"y{i+1}") + ax.set_xlabel("t") + plt.show() + + # The Python and Basilisk implementation should behave identially + numpy.testing.assert_allclose( + xBasilisk[-1,:], + xPython[-1,:], + atol=1e-12, + rtol=0 + ) + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +@pytest.mark.parametrize("method", METHODS) +@pytest.mark.parametrize("figureOfMerit", ["mean", "variance"]) +def test_validateOu(method: Method, figureOfMerit: Literal["mean", "variance"]): + """ + Validate the weak accuracy of the integrators for the Ornstein-Uhlenbeck process. + Compares the empirical mean or variance of the final state to the analytical value + using multiple Monte Carlo batches. + + Args: + method: Integration method. + figureOfMerit: "mean" or "variance" to validate. + """ + + dt = .05 + tf = 5 + + system = OrnsteinUhlenbeckSystem() + + def basiliskTrajectory(): + scSim, stateModel, integratorObject, stateLogger = getBasiliskSim(method, dt, system.x0, system.f, system.g, None) + scSim.ConfigureStopTime( macros.sec2nano(tf) ) + scSim.ExecuteSimulation() + + xBasilisk = stateLogger.x + + return np.array([xBasilisk[-1]]) + + # G is some differentiable function on the final state after + # propagation + # estimateGOnTrajectory should be the correct E[G(x(t))] + + if figureOfMerit == "mean": + def G(arr: npt.NDArray[np.float64]) -> float: + return arr[0] + + # E[G(x(t))] = E[x(t)] + estimateGOnTrajectory = system.mean(tf)[0] + + elif figureOfMerit == "variance": + def G(arr: npt.NDArray[np.float64]) -> float: + return arr[0]**2 + + # E[G(x(t))] = E[x(t)**2] = Var(x(t)) + E[x(t)]**2 + estimateGOnTrajectory = system.var(tf)[0] + system.mean(tf)[0]**2 + + err, varErr = estimateErrorAndEmpiricalVariance( + basiliskTrajectory, + G, + estimateGOnTrajectory, + M1 = 10, + M2 = 100 + ) + fourSigma = 4*np.sqrt(varErr) + + print(method, figureOfMerit, "error", err, "+-", fourSigma) + + # We expect the error to be zero, but we allow some tolerance + # given that the error is estimated with a certain variance + np.testing.assert_allclose( + err, + 0, + atol = fourSigma, + rtol = 0 + ) + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +@pytest.mark.parametrize("method", METHODS) +def test_validateExample1(method: Method): + """ + Validate the weak accuracy of the integrators for Example 1 from Tang & Xiao (2017). + Compares the empirical variance of the final state to the analytical value using + multiple Monte Carlo batches. + + Args: + method: Integration method. + """ + + dt = 2.**-3 + tf = 4 + + system = Example1System() + + def basiliskTrajectory(): + scSim, stateModel, integratorObject, stateLogger = getBasiliskSim(method, dt, system.x0, system.f, system.g, None) + scSim.ConfigureStopTime( macros.sec2nano(tf) ) + scSim.ExecuteSimulation() + + xBasilisk = stateLogger.x + + return xBasilisk[-1,:] + + # Use the same G and E[G(x(t))] as in Example 1 in the paper + def G(arr: npt.NDArray[np.float64]) -> float: + return arr[1]**2 + + estimateGOnTrajectory = 149/150*np.exp(-5/2*tf) +1/150*np.exp(-tf) + + err, varErr = estimateErrorAndEmpiricalVariance( + basiliskTrajectory, + G, + estimateGOnTrajectory, + M1 = 10, + M2 = 100 + ) + fourSigma = 4*np.sqrt(varErr) + + print(method, "variance error", err, "+-", fourSigma) + + # We expect the error to be zero, but we allow some tolerance + # given that the error is estimated with a certain variance + np.testing.assert_allclose( + err, + 0, + atol = fourSigma, + rtol = 0 + ) + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/src/simulation/dynamics/Integrators/svIntegrators.i b/src/simulation/dynamics/Integrators/svIntegrators.i index eb775bfc04..dda02a76ad 100755 --- a/src/simulation/dynamics/Integrators/svIntegrators.i +++ b/src/simulation/dynamics/Integrators/svIntegrators.i @@ -21,12 +21,16 @@ %{ #include #include "../_GeneralModuleFiles/stateVecIntegrator.h" + #include "../_GeneralModuleFiles/stateVecStochasticIntegrator.h" #include "../_GeneralModuleFiles/svIntegratorRungeKutta.h" #include "../_GeneralModuleFiles/svIntegratorRK4.h" #include "svIntegratorEuler.h" #include "svIntegratorRK2.h" #include "svIntegratorRKF45.h" #include "svIntegratorRKF78.h" + #include "svStochIntegratorW2Ito1.h" + #include "svStochIntegratorW2Ito2.h" + #include "svStochIntegratorMayurama.h" #include "architecture/_GeneralModuleFiles/sys_model.h" #include "../_GeneralModuleFiles/dynamicObject.h" %} @@ -37,6 +41,7 @@ from Basilisk.architecture.swig_common_model import * # The following maps store the RK base classes w.r.t their stage number _rk_base_classes = {} _rk_adaptive_base_classes = {} +_rk_stochastic_base_classes = {} %} %include @@ -52,11 +57,15 @@ _rk_adaptive_base_classes = {} } } +%include "swig_eigen.i" + %include "sys_model.i" %include "../_GeneralModuleFiles/stateVecIntegrator.h" +%include "../_GeneralModuleFiles/stateVecStochasticIntegrator.h" %include "../_GeneralModuleFiles/svIntegratorRungeKutta.h" %include "../_GeneralModuleFiles/svIntegratorAdaptiveRungeKutta.h" +%include "../_GeneralModuleFiles/svIntegratorWeakStochasticRungeKutta.h" // We add a constructor for svIntegratorRungeKutta and svIntegratorAdaptiveRungeKutta // These are useful for us to build these classes on the Python side without having @@ -67,7 +76,7 @@ _rk_adaptive_base_classes = {} std::vector> aMatrix, std::vector bArray, std::vector cArray - ) + ) { RKCoefficients coefficients; for (size_t i = 0; i < numberStages; i++) @@ -89,7 +98,7 @@ _rk_adaptive_base_classes = {} std::vector bStarArray, std::vector cArray, double largestOrder - ) + ) { RKAdaptiveCoefficients coefficients; for (size_t i = 0; i < numberStages; i++) @@ -104,13 +113,50 @@ _rk_adaptive_base_classes = {} } } +%extend svIntegratorWeakStochasticRungeKutta { + svIntegratorWeakStochasticRungeKutta( + DynamicObject* dynIn, + std::vector alpha, + std::vector beta0, + std::vector beta1, + std::vector> A0, + std::vector> B0, + std::vector> A1, + std::vector> B1, + std::vector> B2, + std::vector c0, + std::vector c1 + ) + { + SRKCoefficients coefficients; + + std::copy_n(alpha.begin(), numberStages, coefficients.alpha.begin()); + std::copy_n(beta0.begin(), numberStages, coefficients.beta0.begin()); + std::copy_n(beta1.begin(), numberStages, coefficients.beta1.begin()); + std::copy_n(c0.begin(), numberStages, coefficients.c0.begin()); + std::copy_n(c1.begin(), numberStages, coefficients.c1.begin()); + + for (size_t i = 0; i < numberStages; i++) { + std::copy_n(A0.at(i).begin(), numberStages, coefficients.A0.at(i).begin()); + std::copy_n(B0.at(i).begin(), numberStages, coefficients.B0.at(i).begin()); + std::copy_n(A1.at(i).begin(), numberStages, coefficients.A1.at(i).begin()); + std::copy_n(B1.at(i).begin(), numberStages, coefficients.B1.at(i).begin()); + std::copy_n(B2.at(i).begin(), numberStages, coefficients.B2.at(i).begin()); + } + + return new svIntegratorWeakStochasticRungeKutta(dynIn, coefficients); + } +} + %define TEMPLATE_HELPER(stageNumber) %template(svIntegratorRungeKutta ## stageNumber) svIntegratorRungeKutta< ## stageNumber>; %template(svIntegratorAdaptiveRungeKutta ## stageNumber) svIntegratorAdaptiveRungeKutta< ## stageNumber>; +%template(svIntegratorWeakStochasticRungeKutta ## stageNumber) svIntegratorWeakStochasticRungeKutta< ## stageNumber>; %pythoncode %{ _rk_base_classes[## stageNumber] = svIntegratorRungeKutta ## stageNumber _rk_adaptive_base_classes[## stageNumber] = svIntegratorAdaptiveRungeKutta ## stageNumber +_rk_stochastic_base_classes[## stageNumber] = svIntegratorWeakStochasticRungeKutta ## stageNumber %} %enddef @@ -129,6 +175,9 @@ TEMPLATE_HELPER(13) %include "svIntegratorRK2.h" %include "svIntegratorRKF45.h" %include "svIntegratorRKF78.h" +%include "svStochIntegratorW2Ito1.h" +%include "svStochIntegratorW2Ito2.h" +%include "svStochIntegratorMayurama.h" // The following methods allow users to create new Runge-Kutta // methods simply by providing their coefficients on the Python side @@ -142,7 +191,7 @@ def _validate_coefficients(a_coefficients, **array_coefficients): assert a_coefficients.ndim == 2 except: raise ValueError("a_coefficients must be a square matrix or a non-ragged sequence of sequences") - + if a_coefficients.shape[0] != a_coefficients.shape[1]: raise ValueError("a_coefficients must be a square matrix") @@ -170,11 +219,11 @@ def svIntegratorRungeKutta( left side. Args: - a_coefficients (Sequence[Sequence[float]] | np.ndarray): + a_coefficients (Sequence[Sequence[float]] | np.ndarray): "a" matrix in the Butcher table. - b_coefficients (Sequence[float] | np.ndarray): + b_coefficients (Sequence[float] | np.ndarray): "b" array in the Butcher table - c_coefficients (Sequence[float] | np.ndarray): + c_coefficients (Sequence[float] | np.ndarray): "c" array in the Butcher table Returns: @@ -208,25 +257,79 @@ def svIntegratorAdaptiveRungeKutta( largest_order (float): The order of the higher-order RK method used in this adaptive RK method. For example, for RKF45, largest_order should be 5. - a_coefficients (Sequence[Sequence[float]] | np.ndarray): + a_coefficients (Sequence[Sequence[float]] | np.ndarray): "a" matrix in the Butcher table. - b_coefficients (Sequence[float] | np.ndarray): + b_coefficients (Sequence[float] | np.ndarray): "b" array in the Butcher table - b_star_coefficients (Sequence[float] | np.ndarray): + b_star_coefficients (Sequence[float] | np.ndarray): "b" array for the lower order method in the Butcher table - c_coefficients (Sequence[float] | np.ndarray): + c_coefficients (Sequence[float] | np.ndarray): "c" array in the Butcher table Returns: StateVecIntegrator: A Runge-Kutta integrator object """ _validate_coefficients( - a_coefficients, b_coefficients=b_coefficients, + a_coefficients, b_coefficients=b_coefficients, b_star_coefficients=b_star_coefficients, c_coefficients=c_coefficients) stages = len(b_coefficients) return _rk_adaptive_base_classes[stages](dynamic_object, a_coefficients, b_coefficients, b_star_coefficients, c_coefficients, largest_order) +def _validate_srk_coefficients(array_coefficients, matrix_coefficients): + stages = len(next(iter(array_coefficients.values()))) + + for input_name, array_coefficient in array_coefficients.items(): + if stages != len(array_coefficient): + raise ValueError( + f"All arrays must have length {stages}, but {input_name} has length {len(array_coefficient)}.") + + for matrix_name, matrix in matrix_coefficients.items(): + if len(matrix) != stages or any(len(row) != stages for row in matrix): + raise ValueError( + f"Matrix {matrix_name} must have dimensions {stages}x{stages}, but has dimensions {len(matrix)}x{len(matrix[0]) if matrix else 0}.") + +def svIntegratorWeakStochasticRungeKutta( + dynamic_object, + alpha: Union[Sequence[float], np.ndarray], + beta0: Union[Sequence[float], np.ndarray], + beta1: Union[Sequence[float], np.ndarray], + A0: Union[Sequence[Sequence[float]], np.ndarray], + B0: Union[Sequence[Sequence[float]], np.ndarray], + A1: Union[Sequence[Sequence[float]], np.ndarray], + B1: Union[Sequence[Sequence[float]], np.ndarray], + B2: Union[Sequence[Sequence[float]], np.ndarray], + c0: Union[Sequence[float], np.ndarray], + c1: Union[Sequence[float], np.ndarray], +) -> StateVecStochasticIntegrator: + """Generates an explicit, weak stochastic Runge-Kutta integrator from the given coefficients. + + Args: + alpha, beta0, beta1, c0, c1: Arrays with length equal to the number of stages. + A0, B0, A1, B1, B2: Matrices with dimensions stages x stages. + + Returns: + StateVecStochasticIntegrator: A weak stochastic Runge-Kutta integrator object. + """ + _validate_srk_coefficients( + array_coefficients={'alpha': alpha, 'beta0': beta0, 'beta1': beta1, 'c0': c0, 'c1': c1}, + matrix_coefficients={'A0': A0, 'B0': B0, 'A1': A1, 'B1': B1, 'B2': B2} + ) + stages = len(alpha) + return _rk_stochastic_base_classes[stages]( + dynamic_object, + alpha, + beta0, + beta1, + A0, + B0, + A1, + B1, + B2, + c0, + c1 + ) + import sys protectAllClasses(sys.modules[__name__]) %} diff --git a/src/simulation/dynamics/Integrators/svStochIntegratorMayurama.cpp b/src/simulation/dynamics/Integrators/svStochIntegratorMayurama.cpp new file mode 100644 index 0000000000..fef4350af6 --- /dev/null +++ b/src/simulation/dynamics/Integrators/svStochIntegratorMayurama.cpp @@ -0,0 +1,50 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ +#include "svStochIntegratorMayurama.h" + +Eigen::VectorXd EulerMayuramaRandomVariableGenerator::generate(size_t m, double h) +{ + // purge any hidden states on the rv + this->normal_rv.reset(); + + // Generate m standard normal rv samples scaled by sqrt(h) + Eigen::VectorXd pseudoTimeSteps = std::sqrt(h)*Eigen::VectorXd::Ones(m); + for (size_t i = 0; i < m; i++) + { + pseudoTimeSteps[i] *= this->normal_rv(this->rng); + } + return pseudoTimeSteps; +} + +void svStochIntegratorMayurama::integrate(double currentTime, double timeStep) +{ + // Compute the derivatives and diffusions + for (auto dynPtr : this->dynPtrs) { + dynPtr->equationsOfMotion(currentTime, timeStep); + dynPtr->equationsOfMotionDiffusion(currentTime, timeStep); + } + + // this is a map (ExtendedStateId -> noise index) for each of the noise sources + // (so the length of this vector should be m) + std::vector stateIdToNoiseIndexMaps = getStateIdToNoiseIndexMaps(); + + auto pseudoTimeSteps = this->rvGenerator.generate(stateIdToNoiseIndexMaps.size(), timeStep); + + this->propagateState(timeStep, pseudoTimeSteps, stateIdToNoiseIndexMaps); +} diff --git a/src/simulation/dynamics/Integrators/svStochIntegratorMayurama.h b/src/simulation/dynamics/Integrators/svStochIntegratorMayurama.h new file mode 100644 index 0000000000..0c751540da --- /dev/null +++ b/src/simulation/dynamics/Integrators/svStochIntegratorMayurama.h @@ -0,0 +1,88 @@ +/* + ISC License + +Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + +Permission to use, copy, modify, and/or distribute this software for any +purpose with or without fee is hereby granted, provided that the above +copyright notice and this permission notice appear in all copies. + +THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + +*/ + +#ifndef svStochIntegratorMayurama_h +#define svStochIntegratorMayurama_h + +#include "../_GeneralModuleFiles/dynamicObject.h" +#include "../_GeneralModuleFiles/dynParamManager.h" +#include "../_GeneralModuleFiles/stateVecStochasticIntegrator.h" +#include "../_GeneralModuleFiles/extendedStateVector.h" + +#include + + +/** Random Number Generator for the Euler-Mayurama integrator */ +class EulerMayuramaRandomVariableGenerator +{ +public: + /** Sets the seed for the RNG */ + inline void setSeed(size_t seed) {rng.seed(seed);} + + /** Generates m normally distributed RV samples, with mean + * zero and standard deviation equal to sqrt(h) + */ + Eigen::VectorXd generate(size_t m, double h); + +protected: + /** Random Number Generator */ + std::mt19937 rng{std::random_device{}()}; + + /** Standard normally distributed random variable */ + std::normal_distribution normal_rv{0., 1.}; +}; + +/** The 1-weak 1-strong order Euler-Mayurama stochastic integrator. + * + * For an SDE of the form: + * + * dx = f(t,x)*dt + sum_i g_i(t,x)*dW + * + * The integrator, with timestep h, computes: + * + * x_{n+1} = x_n + f(t,x)*h + sum_i g_i(t,x)*sqrt(h)*N_i + * + * where N_i are independent and identically distributed standard normal + * random variables. + */ +class svStochIntegratorMayurama : public StateVecStochasticIntegrator { +public: + + /** Uses same constructor as StateVecStochasticIntegrator */ + using StateVecStochasticIntegrator::StateVecStochasticIntegrator; + + /** Sets the seed for the Random Number Generator used by this integrator. + * + * As a stochastic integrator, random numbers are drawn during each time step. + * By default, a randomly generated seed is used each time. + * + * If the seed is set, the integrator will always draw the same numbers + * during time-stepping. + */ + inline void setRNGSeed(size_t seed) {rvGenerator.setSeed(seed);}; + + /** Performs the integration of the associated dynamic objects up to time currentTime+timeStep */ + virtual void integrate(double currentTime, double timeStep) override; + +public: + /** Random Number Generator for the integrator */ + EulerMayuramaRandomVariableGenerator rvGenerator; +}; + +#endif // svStochIntegratorMayurama_h diff --git a/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito1.cpp b/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito1.cpp new file mode 100644 index 0000000000..70ca2822b3 --- /dev/null +++ b/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito1.cpp @@ -0,0 +1,55 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ +#include "svStochIntegratorW2Ito1.h" + +svStochIntegratorW2Ito1::svStochIntegratorW2Ito1(DynamicObject* dyn) + : svIntegratorWeakStochasticRungeKutta(dyn, svStochIntegratorW2Ito1::getCoefficients()) +{} + +SRKCoefficients<3> svStochIntegratorW2Ito1::getCoefficients() +{ + SRKCoefficients<3> coefficients; + + coefficients.A0[1][0] = .5; + coefficients.A0[2][0] = -1; + coefficients.A0[2][1] = 2; + + coefficients.A1[1][0] = .25; + coefficients.A1[2][0] = .25; + + coefficients.B0[1][0] = (6 - std::sqrt(6))/10.; + coefficients.B0[2][0] = (3 + 2*std::sqrt(6))/5.; + + coefficients.B1[1][0] = .5; + coefficients.B1[2][0] = -.5; + + coefficients.B2[1][0] = 1; + + coefficients.alpha = {1./6., 2./3., 1./6.}; + coefficients.beta0 = {-1, 1, 1}; + coefficients.beta1 = { 2, 0, -2}; + + // c0[j] = sum_{p=1..s}( a0[j][p] ) + coefficients.c0 = {0, .5, 1}; + + // c1[j] = sum_{p=1..s}( a1[j][p] ) + coefficients.c1 = {0, .25, .25}; + + return coefficients; +} diff --git a/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito1.h b/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito1.h new file mode 100644 index 0000000000..e5209fcae3 --- /dev/null +++ b/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito1.h @@ -0,0 +1,41 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ + +#ifndef svStochIntegratorW2Ito1_h +#define svStochIntegratorW2Ito1_h + +#include "../_GeneralModuleFiles/svIntegratorWeakStochasticRungeKutta.h" + +/** @brief Weak-order 2, ODE-order 3 stochastic integrator + * + * Method is described in the following paper, coefficients according to + * Table 2: + * + * Tang, X., Xiao, A. Efficient weak second-order stochastic Runge–Kutta methods + * for Itô stochastic differential equations. Bit Numer Math 57, 241–260 (2017). + * https://doi.org/10.1007/s10543-016-0618-9 + */ +class svStochIntegratorW2Ito1 : public svIntegratorWeakStochasticRungeKutta<3> { +public: + svStochIntegratorW2Ito1(DynamicObject* dyn); //!< Constructor +private: + static SRKCoefficients<3> getCoefficients(); +}; + +#endif /* svStochIntegratorW2Ito1_h */ diff --git a/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito2.cpp b/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito2.cpp new file mode 100644 index 0000000000..d0803a0a82 --- /dev/null +++ b/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito2.cpp @@ -0,0 +1,56 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ +#include "svStochIntegratorW2Ito2.h" + +svStochIntegratorW2Ito2::svStochIntegratorW2Ito2(DynamicObject* dyn) + : svIntegratorWeakStochasticRungeKutta(dyn, svStochIntegratorW2Ito2::getCoefficients()) +{} + +SRKCoefficients<4> svStochIntegratorW2Ito2::getCoefficients() +{ + SRKCoefficients<4> coefficients; + + coefficients.A0[1][0] = .5; + coefficients.A0[2][1] = .5; + coefficients.A0[3][2] = 1; + + coefficients.A1[1][0] = .5; + coefficients.A1[2][0] = .5; + coefficients.A1[3][0] = .5; + + coefficients.B0[2][1] = 1; + coefficients.B0[3][1] = 1; + + coefficients.B1[2][1] = .5; + coefficients.B1[3][1] = -.5; + + coefficients.B2[2][1] = 1; + + coefficients.alpha = {1./6., 1./3., 1./3., 1./6.}; + coefficients.beta0 = {0, -1, 1, 1}; + coefficients.beta1 = {0, 2, 0, -2}; + + // c0[j] = sum_{p=1..s}( a0[j][p] ) + coefficients.c0 = {0, .5, .5, 1}; + + // c1[j] = sum_{p=1..s}( a1[j][p] ) + coefficients.c1 = {0, .5, .5, .5}; + + return coefficients; +} diff --git a/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito2.h b/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito2.h new file mode 100644 index 0000000000..fac12fd0a1 --- /dev/null +++ b/src/simulation/dynamics/Integrators/svStochIntegratorW2Ito2.h @@ -0,0 +1,41 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ + +#ifndef svStochIntegratorW2Ito2_h +#define svStochIntegratorW2Ito2_h + +#include "../_GeneralModuleFiles/svIntegratorWeakStochasticRungeKutta.h" + +/** @brief Weak-order 2, ODE-order 4 stochastic integrator + * + * Method is described in the following paper, coefficients according to + * Table 3: + * + * Tang, X., Xiao, A. Efficient weak second-order stochastic Runge–Kutta methods + * for Itô stochastic differential equations. Bit Numer Math 57, 241–260 (2017). + * https://doi.org/10.1007/s10543-016-0618-9 + */ +class svStochIntegratorW2Ito2 : public svIntegratorWeakStochasticRungeKutta<4> { +public: + svStochIntegratorW2Ito2(DynamicObject* dyn); //!< Constructor +private: + static SRKCoefficients<4> getCoefficients(); +}; + +#endif /* svStochIntegratorW2Ito2_h */ diff --git a/src/simulation/dynamics/_GeneralModuleFiles/dynParamManager.cpp b/src/simulation/dynamics/_GeneralModuleFiles/dynParamManager.cpp index b1982a2791..d758124d7f 100755 --- a/src/simulation/dynamics/_GeneralModuleFiles/dynParamManager.cpp +++ b/src/simulation/dynamics/_GeneralModuleFiles/dynParamManager.cpp @@ -61,10 +61,17 @@ void StateVector::scaleStates(double scaleFactor) } } -void StateVector::propagateStates(double dt) +void StateVector::propagateStates(double dt, const std::unordered_map>& pseudoTimeSteps) { for (const auto& [key, value] : stateMap) { - value->propagateState(dt); + if (pseudoTimeSteps.count(key) > 0) + { + value->propagateState(dt, pseudoTimeSteps.at(key)); + } + else + { + value->propagateState(dt); + } } } @@ -92,7 +99,10 @@ void DynParamManager::updateStateVector(const StateVector& newState) this->stateContainer.setStates(newState); } -void DynParamManager::propagateStateVector(double dt) { this->stateContainer.propagateStates(dt); } +void DynParamManager::propagateStateVector(double dt, const std::unordered_map>& pseudoTimeSteps) +{ + this->stateContainer.propagateStates(dt, pseudoTimeSteps); +} Eigen::MatrixXd* DynParamManager::createProperty(std::string propName, const Eigen::MatrixXd& propValue) @@ -141,3 +151,40 @@ void DynParamManager::setPropertyValue(const std::string propName, const Eigen:: it->second = propValue; } } + +void +DynParamManager::registerSharedNoiseSource(std::vector> sharedNoises) +{ + std::vector> noiseIds; + + for (auto&& [stateData, noiseIndex] : sharedNoises) + { + if ( + this->stateContainer.stateMap.count(stateData.getName()) == 0 + || this->stateContainer.stateMap.at(stateData.getName()).get() != &stateData + ) + { + throw std::runtime_error("Given StateData '" + + stateData.getName() + "' does not belong to this dynParamManager."); + } + + if (noiseIndex >= stateData.getNumNoiseSources()) + { + throw std::runtime_error("Cannot share noise index '" + + std::to_string(noiseIndex) + "' of StateData '" + stateData.getName() + + "' because that StateData only has " + std::to_string(stateData.getNumNoiseSources()) + + " noise sources."); + } + + noiseIds.emplace_back(stateData.getName(), noiseIndex); + } + + // Save all noiseId in the sharedNoiseMap with the + // same ID (a simple counter to guarantee uniqueness) + size_t sharedNoiseId = this->sharedNoiseMapIdCounter; + this->sharedNoiseMapIdCounter++; + for (auto&& noiseId : noiseIds) + { + sharedNoiseMap[noiseId] = sharedNoiseId; + } +} diff --git a/src/simulation/dynamics/_GeneralModuleFiles/dynParamManager.h b/src/simulation/dynamics/_GeneralModuleFiles/dynParamManager.h index fb85e9b7ab..6b645aef87 100755 --- a/src/simulation/dynamics/_GeneralModuleFiles/dynParamManager.h +++ b/src/simulation/dynamics/_GeneralModuleFiles/dynParamManager.h @@ -27,6 +27,33 @@ #include #include #include +#include + +/// @cond DOXYGEN_IGNORE +template +void hash_combine(std::size_t& seed, const T& v, const Rest&... rest) +{ + seed ^= std::hash{}(v) + 0x9e3779b9 + (seed << 6) + (seed >> 2); + (hash_combine(seed, rest), ...); +} + +namespace std // Inject hash for std::pair into std:: +{ + /** Hash implementation for ``std::pair``, + * allows using it as the key in maps. + */ + template<> struct hash> + { + /** Produce hash from ``std::pair`` */ + std::size_t operator()(const std::pair& input) const noexcept + { + std::size_t h = 0; + hash_combine(h, input.first, input.second); + return h; + } + }; +} +/// @endcond /** StateVector represents an ordered collection of StateData, * with each state having a unique name. @@ -71,7 +98,7 @@ class StateVector { void scaleStates(double scaleFactor); /** Scales the states of this StateVector by the given delta time */ - void propagateStates(double dt); + void propagateStates(double dt, const std::unordered_map>& pseudoTimeSteps = {}); }; /** A class that manages a set of states and properties. */ @@ -141,7 +168,44 @@ class DynParamManager { void updateStateVector(const StateVector& newState); /** Propagates the states managed by this class a given delta time. */ - void propagateStateVector(double dt); + void propagateStateVector(double dt, const std::unordered_map>& pseudoTimeSteps = {}); + + /** Used when more than one state have dynamics perturbed + * by the same noise process. + * + * For example, consider the following SDE: + * + * dx_0 = f_0(t,x)dt + g_00(t,x)dW_0 + g_01(t,x)dW_1 + * dx_1 = f_1(t,x)dt + g_11(t,x)dW_1 + * + * In this case, state 'x_0' is affected by 2 sources of noise + * and 'x_1' by 1 source of noise. However, the source 'W_1' is + * shared between 'x_0' and 'x_1'. + * + * This function is called like: + * + * \code + * dynParamManager.registerSharedNoiseSource({ + * {myStateX0, 1}, + * {myStateX1, 0} + * }); + * \endcode + * + * which means that the 2nd noise source of the ``StateData`` 'myStateX0' + * and the first noise source of the ``StateData`` 'myStateX1' actually + * correspond to the same noise process. + */ + void registerSharedNoiseSource(std::vector>); + + /** When multiple states share a noise source, then this map + * maps their (name, noiseIndex) to the same number + */ + std::unordered_map, size_t> sharedNoiseMap; +protected: + /** A counter used to ensure that the values in ``sharedNoiseMap`` + * are unique. + */ + size_t sharedNoiseMapIdCounter = 0; }; template , we need to declare these manually +%traits_swigtype(StateData); +%fragment(SWIG_Traits_frag(StateData)); + +%template() std::pair; +%template() std::vector>; + +%extend DynParamManager { + // SWIG doesnt like const StateData& so we have to use const StateData* and convert + void registerSharedNoiseSource(std::vector> list) { + // Convert from pointer pairs to reference pairs + std::vector> refList; + refList.reserve(list.size()); + for (const auto& p : list) { + refList.emplace_back(*p.first, p.second); + } + $self->registerSharedNoiseSource(refList); + } +} + +%ignore DynParamManager::registerSharedNoiseSource(std::vector>); + %include "simulation/dynamics/_GeneralModuleFiles/dynParamManager.h" // Used in testing in Python, we use the default template arguments diff --git a/src/simulation/dynamics/_GeneralModuleFiles/dynamicObject.h b/src/simulation/dynamics/_GeneralModuleFiles/dynamicObject.h index 1c408a1972..94acf26359 100755 --- a/src/simulation/dynamics/_GeneralModuleFiles/dynamicObject.h +++ b/src/simulation/dynamics/_GeneralModuleFiles/dynamicObject.h @@ -47,9 +47,25 @@ class DynamicObject : public SysModel { /** Hooks the dyn-object into Basilisk architecture */ virtual void UpdateState(uint64_t callTime) = 0; - /** Computes F = Xdot(X,t) */ + /** Computes the time derivative of the states: + * + * dx = f(t,x) dt + * + * ``equationsOfMotion`` is f(t,x) is the equation above. + */ virtual void equationsOfMotion(double t, double timeStep) = 0; + /** Computes the diffusion of the states: + * + * dx = f(t,x) dt + g_0(t,x)dW_0 + g_1(t,x)dW_1 + ... + g_{n-1}(t,x)dW_{n-1} + * + * ``equationsOfMotionDiffusion`` is equivalent to evaluating + * g_0(t,x), g_1(t,x), ..., g_{n-1}(t,x) in the equation above. + * + * Note that not all ``DynamicObjects`` may support this functionality. + */ + virtual void equationsOfMotionDiffusion(double t, double timeStep) {}; + /** Performs pre-integration steps */ virtual void preIntegration(double callTime) = 0; diff --git a/src/simulation/dynamics/_GeneralModuleFiles/extendedStateVector.cpp b/src/simulation/dynamics/_GeneralModuleFiles/extendedStateVector.cpp index 587c9a674a..8b60706b56 100644 --- a/src/simulation/dynamics/_GeneralModuleFiles/extendedStateVector.cpp +++ b/src/simulation/dynamics/_GeneralModuleFiles/extendedStateVector.cpp @@ -29,9 +29,31 @@ ExtendedStateVector ExtendedStateVector::fromStateDerivs(const std::vector - functor) const +ExtendedStateVector +ExtendedStateVector::fromStateDiffusions(const std::vector& dynPtrs, + const StateIdToIndexMap& stateIdToNoiseIndexMap) +{ + ExtendedStateVector result; + for (auto&& [stateId, noiseIndex] : stateIdToNoiseIndexMap) + { + result.emplace(stateId, dynPtrs.at(stateId.first)->dynManager.getStateObject(stateId.second)->getStateDiffusion(noiseIndex)); + } + return result; +} + +std::vector +ExtendedStateVector::fromStateDiffusions(const std::vector& dynPtrs, + const std::vector& stateIdToNoiseIndexMaps) +{ + std::vector result; + for (auto&& stateIdToNoiseIndexMap : stateIdToNoiseIndexMaps) + result.push_back( fromStateDiffusions(dynPtrs, stateIdToNoiseIndexMap) ); + return result; +} + +ExtendedStateVector +ExtendedStateVector::map( + std::function functor) const { ExtendedStateVector result; result.reserve(this->size()); @@ -115,6 +137,23 @@ void ExtendedStateVector::setDerivatives(std::vector& dynPtrs) c }); } +void ExtendedStateVector::setDiffusions( + std::vector& dynPtrs, + const StateIdToIndexMap& stateIdToNoiseIndexMap +) const +{ + this->apply([&dynPtrs, &stateIdToNoiseIndexMap](const size_t& dynObjIndex, + const std::string& stateName, + const Eigen::MatrixXd& thisDerivative) { + dynPtrs.at(dynObjIndex) + ->dynManager.stateContainer.stateMap.at(stateName) + ->setDiffusion( + thisDerivative, + stateIdToNoiseIndexMap.at({dynObjIndex, stateName}) + ); + }); +} + ExtendedStateVector ExtendedStateVector::fromStateData(const std::vector& dynPtrs, std::function functor) @@ -130,3 +169,32 @@ ExtendedStateVector::fromStateData(const std::vector& dynPtrs, return result; } + +// ostream<< overload, useful for easily printing the ExtendedStateVector +std::ostream& operator<<(std::ostream& os, const ExtendedStateVector& myMap) { + std::map, Eigen::MatrixXd> ordered(myMap.begin(), myMap.end()); + + // Aligned column formatting with 6 decimal precision and consistent indentation + Eigen::IOFormat alignedFmt( + Eigen::StreamPrecision, // use full stream precision + Eigen::DontAlignCols, // prevents broken alignment with indentation + " ", // coefficient separator + "\n", // row separator + " ", // row prefix (indent) + "" // row postfix + ); + + for (const auto& [key, matrix] : ordered) { + const auto& [id, name] = key; + os << "(" << id << ", \"" << name << "\"):\n"; + + if (matrix.rows() == 1 && matrix.cols() == 1) { + os << " " << matrix(0, 0) << "\n"; + } else if (matrix.cols() == 1) { + os << matrix.transpose().format(alignedFmt) << "\n"; + } else { + os << matrix.format(alignedFmt) << "\n"; + } + } + return os; +} diff --git a/src/simulation/dynamics/_GeneralModuleFiles/extendedStateVector.h b/src/simulation/dynamics/_GeneralModuleFiles/extendedStateVector.h index 2671f4ce61..018f8e217a 100644 --- a/src/simulation/dynamics/_GeneralModuleFiles/extendedStateVector.h +++ b/src/simulation/dynamics/_GeneralModuleFiles/extendedStateVector.h @@ -59,17 +59,26 @@ makes performing state-wise operations easier. */ using ExtendedStateId = std::pair; -/** ExtendedStateIdHash is required to make ExtendedStateId hashable (usable as a key in a map) */ -struct ExtendedStateIdHash { - /** Generates a hash value (integer) from an ExtendedStateId object */ - std::size_t operator()(const ExtendedStateId& p) const +/// @cond DOXYGEN_IGNORE +namespace std // Inject hash for ExtendedStateId into std:: +{ + /** Hash implementation for ``std::pair``, + * allows using it as the key in maps. + */ + template<> struct hash { - auto seed = std::hash{}(p.first); - // Algorithm from boost::hash_combine - seed ^= std::hash{}(p.second) + 0x9e3779b9 + (seed << 6) + (seed >> 2); - return seed; - } -}; + /** Produce hash from ``std::pair`` */ + std::size_t operator()(const ExtendedStateId& input) const noexcept + { + std::size_t h = 0; + hash_combine(h, input.first, input.second); + return h; + } + }; +} +/// @endcond + +using StateIdToIndexMap = std::unordered_map; /** * Conceptually similar to StateVector, this class allows us to handle @@ -78,7 +87,7 @@ struct ExtendedStateIdHash { * It also supports several utility functions. */ class ExtendedStateVector - : public std::unordered_map { + : public std::unordered_map { public: /** * Builds a ExtendedStateVector from all states in the given @@ -92,6 +101,19 @@ class ExtendedStateVector */ static ExtendedStateVector fromStateDerivs(const std::vector& dynPtrs); + /** + * Extracts the diffusion at the specified noise index for the states + * present in ``stateIdToNoiseIndexMap``. + */ + static ExtendedStateVector + fromStateDiffusions(const std::vector& dynPtrs, const StateIdToIndexMap& stateIdToNoiseIndexMap); + + /** + * Calls and returns the result of ``fromStateDiffusions`` for each map in ``stateIdToNoiseIndexMaps``. + */ + static std::vector + fromStateDiffusions(const std::vector& dynPtrs, const std::vector& stateIdToNoiseIndexMaps); + /** * This method will call the given std::function for every * state in the ExtendedStateVector. The arguments to the functor @@ -140,9 +162,24 @@ class ExtendedStateVector /** Calls StateData::setDerivative for every entry in this */ void setDerivatives(std::vector& dynPtrs) const; + /** Calls StateData::setDiffusion for every entry in this + * + * Note that this extendedStateVector may contain the diffusion + * corresponding to different noise sources for each state. Thus, + * the input ``stateIdToNoiseIndexMap`` is necessary to set for + * which noise source this diffusion is given. + */ + void setDiffusions( + std::vector& dynPtrs, + const StateIdToIndexMap& stateIdToNoiseIndexMap + ) const; + private: static ExtendedStateVector fromStateData(const std::vector& dynPtrs, std::function); }; +// ostream<< overload, useful for easily printing the ExtendedStateVector +std::ostream& operator<<(std::ostream& os, const ExtendedStateVector& myMap); + #endif /* extendedStateVector_h */ diff --git a/src/simulation/dynamics/_GeneralModuleFiles/stateData.cpp b/src/simulation/dynamics/_GeneralModuleFiles/stateData.cpp index 0b886304a9..6ca50042e0 100755 --- a/src/simulation/dynamics/_GeneralModuleFiles/stateData.cpp +++ b/src/simulation/dynamics/_GeneralModuleFiles/stateData.cpp @@ -24,6 +24,11 @@ std::unique_ptr StateData::clone() const { auto result = std::make_unique(this->stateName, this->state); result->stateDeriv = this->stateDeriv; + result->stateDiffusion.resize(this->stateDiffusion.size()); + for (size_t i = 0; i < this->stateDiffusion.size(); i++) + { + result->stateDiffusion[i] = this->stateDiffusion[i]; + } return result; } @@ -35,14 +40,43 @@ StateData::StateData(std::string inName, const Eigen::MatrixXd & newState) setDerivative( Eigen::MatrixXd::Zero(state.innerSize(), state.outerSize()) ); } +void StateData::setNumNoiseSources(size_t numSources) +{ + stateDiffusion.resize(numSources); + for (size_t i = 0; i < numSources; i++) + { + stateDiffusion[i].resizeLike(state); + + } +} + +size_t StateData::getNumNoiseSources() const +{ + return stateDiffusion.size(); +} + void StateData::setState(const Eigen::MatrixXd & newState) { state = newState; } -void StateData::propagateState(double dt) +void StateData::propagateState(double dt, std::vector pseudoStep) { state += stateDeriv * dt; + + if (getNumNoiseSources() > 0 && pseudoStep.size() == 0) + { + auto errorMsg = "State " + this->getName() + " has stochastic dynamics, but " + + "the integrator tried to propagate it without pseudoSteps. Are you sure " + + "you are using a stochastic integrator?"; + bskLogger.bskLog(BSK_ERROR, errorMsg.c_str()); + throw std::invalid_argument(errorMsg); + } + + for (size_t i = 0; i < getNumNoiseSources(); i++) + { + state += stateDiffusion.at(i) * pseudoStep.at(i); + } } @@ -51,6 +85,21 @@ void StateData::setDerivative(const Eigen::MatrixXd & newDeriv) stateDeriv = newDeriv; } +void StateData::setDiffusion(const Eigen::MatrixXd & newDiffusion, size_t index) +{ + if (index < stateDiffusion.size()) + { + stateDiffusion[index] = newDiffusion; + } + else + { + auto errorMsg = "Tried to set diffusion index greater than number of noise sources configured: " + + std::to_string(index) + " >= " + std::to_string(stateDiffusion.size()); + bskLogger.bskLog(BSK_ERROR, errorMsg.c_str()); + throw std::out_of_range(errorMsg); + } +} + void StateData::scaleState(double scaleFactor) { state *= scaleFactor; diff --git a/src/simulation/dynamics/_GeneralModuleFiles/stateData.h b/src/simulation/dynamics/_GeneralModuleFiles/stateData.h index ab99701414..32f6d35ed8 100755 --- a/src/simulation/dynamics/_GeneralModuleFiles/stateData.h +++ b/src/simulation/dynamics/_GeneralModuleFiles/stateData.h @@ -30,6 +30,7 @@ class StateData public: Eigen::MatrixXd state; //!< [-] State value storage Eigen::MatrixXd stateDeriv; //!< [-] State derivative value storage + std::vector stateDiffusion; //!< [-] State diffusion value storage const std::string stateName; //!< [-] Name of the state BSKLogger bskLogger; //!< -- BSK Logging @@ -46,18 +47,66 @@ class StateData /** Destructor */ virtual ~StateData() = default; + /** Sets the number of noise sources for this state. + * + * This is used for stochastic dynamics, where the + * evolutions of the state are driven by a set of independent + * noise sources: + * + * dx = f(t,x)dt + g_0(t,x)dW_0 + g_1(t,x)dW_1 + ... + g_{n-1}(t,x)dW_{n-1} + * + * where dW_i are independent Wiener processes. The number of + * noise sources is equal to the number of diffusion matrices + * that are used to drive the stochastic dynamics (n above). + * + * @param numSources The number of noise sources + */ + void setNumNoiseSources(size_t numSources); + + /** Get how many independent sources of noise drive the dynamics + * of this state. + * + * Any number greater than zero indicates that this state + * is driven by a stochastic differential equation. + */ + size_t getNumNoiseSources() const; + /** Updates the value of the state */ void setState(const Eigen::MatrixXd& newState); /** Updates the derivative of the value of the state */ void setDerivative(const Eigen::MatrixXd& newDeriv); + /** Updates the diffusion of the value of the state. + * + * This is used for stochastic dynamics, where the + * evolutions of the state are driven by a set of independent + * noise sources: + * + * dx = f(t,x)dt + g_0(t,x)dW_0 + g_1(t,x)dW_1 + ... + g_{n-1}(t,x)dW_{n-1} + * + * where dW_i are independent Wiener processes. The diffusion + * matrices are used to drive the stochastic dynamics (g_i above). + * + * @param newDiffusion The new diffusion matrix + * @param index The index of the diffusion matrix to update. + * This must be less than the number of noise sources. + */ + void setDiffusion(const Eigen::MatrixXd& newDiffusion, size_t index); + /** Retrieves a copy of the current state */ Eigen::MatrixXd getState() const { return state; } /** Retrieves a copy of the current state derivative */ Eigen::MatrixXd getStateDeriv() const { return stateDeriv; } + /** Retrieves a copy of the current state diffusion + * + * @param index The index of the diffusion matrix to retrieve. + * This must be less than the number of noise sources. + */ + Eigen::MatrixXd getStateDiffusion(size_t index) const { return stateDiffusion.at(index); } + /** Returns the name of the state */ std::string getName() const { return stateName; } @@ -79,8 +128,20 @@ class StateData /** Adds the values of the other state to this state */ void addState(const StateData& other); - /** Propagates the state in time with the stored derivative */ - virtual void propagateState(double dt); + /** + * @brief Propagates the state over a time step. + * + * This method integrates the position state using the state derivative + * over the given time step:: + * + * x += f(t,x)*h + g_0(t,x)*pseudoStep[0] + g_1(t,x)*pseudoStep[1] ... + * + * @param h The time step for propagation. + * @param pseudoStep For states driven by stochastic dynamics, this + * represents the random pseudotimestep. The length of this input must + * match the number of noise sources of this state (``getNumNoiseSources()``) + */ + virtual void propagateState(double h, std::vector pseudoStep = {}); }; #endif /* STATE_DATA_H */ diff --git a/src/simulation/dynamics/_GeneralModuleFiles/stateVecStochasticIntegrator.cpp b/src/simulation/dynamics/_GeneralModuleFiles/stateVecStochasticIntegrator.cpp new file mode 100644 index 0000000000..64d05a5951 --- /dev/null +++ b/src/simulation/dynamics/_GeneralModuleFiles/stateVecStochasticIntegrator.cpp @@ -0,0 +1,93 @@ +#include + +#include "stateVecStochasticIntegrator.h" +#include "dynamicObject.h" + +std::vector +StateVecStochasticIntegrator::getStateIdToNoiseIndexMaps() +{ + std::vector stateIdToNoiseIndexMaps; + + std::unordered_map sharedNoiseIndexToVectorIndex; + + auto getMapForNoise = [&](size_t dynIndex, const std::string& stateName, size_t noiseIndex) + -> StateIdToIndexMap& + { + std::pair extendedNoiseId = {stateName, noiseIndex}; + + // This noise source is not shared with any other state + if (dynPtrs.at(dynIndex)->dynManager.sharedNoiseMap.count(extendedNoiseId) == 0) + { + return stateIdToNoiseIndexMaps.emplace_back(); + } + + // States that share a noise source have the same value + // at the sharedNoiseMap + size_t sharedNoiseUUID = dynPtrs.at(dynIndex)->dynManager.sharedNoiseMap.at(extendedNoiseId); + + // This noise source is shared, but this is the first time we + // generate the map for this noise source + if (sharedNoiseIndexToVectorIndex.count(sharedNoiseUUID) == 0) + { + // register that this sharedNoiseUUID corresponds to the + // map at this position in the diffusions vector + // and the same for the stateIdToNoiseIndexMaps + sharedNoiseIndexToVectorIndex[sharedNoiseUUID] = stateIdToNoiseIndexMaps.size(); + return stateIdToNoiseIndexMaps.emplace_back(); + } + + // This noise source is shared, and the diffusion should be + // placed at an map that was already created + size_t vecIndex = sharedNoiseIndexToVectorIndex.at(sharedNoiseUUID); + return stateIdToNoiseIndexMaps.at(vecIndex); + }; + + for (size_t dynIndex = 0; dynIndex < dynPtrs.size(); dynIndex++) { + for (auto&& [stateName, stateData] : + dynPtrs.at(dynIndex)->dynManager.stateContainer.stateMap) { + + for (size_t noiseIndex = 0; noiseIndex < stateData->getNumNoiseSources(); noiseIndex++) + { + getMapForNoise(dynIndex, stateName, noiseIndex)[{dynIndex, stateName}] = noiseIndex; + } + } + } + + return stateIdToNoiseIndexMaps; +} + +void +StateVecStochasticIntegrator::propagateState( + double timeStep, + const Eigen::VectorXd& pseudoTimeSteps, + const std::vector& stateIdToNoiseIndexMaps +) +{ + // This map will hold the vector of pseudoTimeSteps to use + // for each state. If a state has 2 sources of noise, this + // map will hold a vector of size 2 with the value of the pseudo + // time step used for each time step + std::vector< std::unordered_map> > localPseudoTimeSteps{dynPtrs.size()}; + + auto getLocalPseudoTForState = [&](size_t dynPtrIndex, const std::string& state) -> std::vector& + { + if (localPseudoTimeSteps.at(dynPtrIndex).count(state) > 0) return localPseudoTimeSteps.at(dynPtrIndex).at(state); + auto nNoise = dynPtrs.at(dynPtrIndex)->dynManager.getStateObject(state)->getNumNoiseSources(); + return localPseudoTimeSteps.at(dynPtrIndex).emplace(state, nNoise).first->second; + }; + + // stateIdToNoiseIndexMaps is used to map between generated + // sources of noise and the ones used locally for each state data + for (size_t k = 0; k < stateIdToNoiseIndexMaps.size(); k++) + { + for (auto&& [stateId, noiseIndex] : stateIdToNoiseIndexMaps.at(k)) + { + getLocalPseudoTForState(stateId.first, stateId.second).at(noiseIndex) = pseudoTimeSteps[k]; + } + } + + for (size_t dynPtrIndex = 0; dynPtrIndex < dynPtrs.size(); dynPtrIndex++) + { + this->dynPtrs.at(dynPtrIndex)->dynManager.propagateStateVector(timeStep, localPseudoTimeSteps.at(dynPtrIndex)); + } +} diff --git a/src/simulation/dynamics/_GeneralModuleFiles/stateVecStochasticIntegrator.h b/src/simulation/dynamics/_GeneralModuleFiles/stateVecStochasticIntegrator.h new file mode 100644 index 0000000000..5f8ef161ce --- /dev/null +++ b/src/simulation/dynamics/_GeneralModuleFiles/stateVecStochasticIntegrator.h @@ -0,0 +1,107 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ + + +#ifndef stateVecStochasticIntegrator_h +#define stateVecStochasticIntegrator_h + +#include +#include +#include + +#include "../_GeneralModuleFiles/stateVecIntegrator.h" +#include "../_GeneralModuleFiles/extendedStateVector.h" + +class StateData; + +/*! @brief state vector integrator class for stochastic dynamics */ +class StateVecStochasticIntegrator : public StateVecIntegrator +{ + +public: + using StateVecIntegrator::StateVecIntegrator; + + /** Returns a vector with length equal to the total number of noise + * sources in the system. Each index contains a map that indicates + * how that noise source maps to individual noise indices for each + * relevant ``StateData``. + * + * For example, consider the following SDE: + * + * dx_0 = f_0(t,x)dt + g_00(t,x)dW_0 + g_01(t,x)dW_1 + * dx_1 = f_1(t,x)dt + g_11(t,x)dW_1 + * + * In this case, state 'x_0' is affected by 2 sources of noise + * and 'x_1' by 1 source of noise. The source 'W_1' is + * shared between 'x_0' and 'x_1'. + * + * This function would return (pseudo-code): + * + * [ + * {(0, "x_0") -> 0}, + * {(0, "x_0") -> 1, (0, "x_1") -> 0}, + * ] + * + * because there are 2 unique sources of noise, the first of + * which affects the first noise index of 'x_0' and the second + * of which affects the second noise index of 'x_0' and the + * first noise index of "x_1". + */ + std::vector getStateIdToNoiseIndexMaps(); + + /** This method calls ``propagateStateVector`` on every state of + * the dynamic objects handled by this integrator. + * + * The length of ``pseudoTimeSteps`` must match the size of + * ``stateIdToNoiseIndexMaps``. The ``stateIdToNoiseIndexMaps`` + * are used to map each ``pseudoTimeStep`` to a noise index + * for each state. + * + * Consider the following SDE: + * + * dx_0 = f_0(t,x)dt + g_00(t,x)dW_0 + g_01(t,x)dW_1 + * dx_1 = f_1(t,x)dt + g_11(t,x)dW_1 + * + * Assume that the values of the drift and diffusion have + * been computed and are stored in the state objects. If this + * method is called with: + * + * - ``timeStep``: h + * - ``pseudoTimeSteps``: [W_0, W_1] + * - ``stateIdToNoiseIndexMaps``: + * [ + * {(0, "x_0") -> 0}, + * {(0, "x_0") -> 1, (0, "x_1") -> 0}, + * ] + * + * Then the states will be changed by: + * + * x_0 += f_0(t,x)*h + g_00(t,x)*W_0 + g_01(t,x)*W_1 + * x_1 += f_1(t,x)*h + g_11(t,x)*W_1 + * + */ + void propagateState( + double timeStep, + const Eigen::VectorXd& pseudoTimeSteps, + const std::vector& stateIdToNoiseIndexMaps + ); +}; + + +#endif /* StateVecIntegrator_h */ diff --git a/src/simulation/dynamics/_GeneralModuleFiles/svIntegratorAdaptiveRungeKutta.h b/src/simulation/dynamics/_GeneralModuleFiles/svIntegratorAdaptiveRungeKutta.h index 7c06d26194..b205aaa618 100644 --- a/src/simulation/dynamics/_GeneralModuleFiles/svIntegratorAdaptiveRungeKutta.h +++ b/src/simulation/dynamics/_GeneralModuleFiles/svIntegratorAdaptiveRungeKutta.h @@ -244,11 +244,11 @@ class svIntegratorAdaptiveRungeKutta : public svIntegratorRungeKutta dynObjectStateSpecificRelTol; + std::unordered_map dynObjectStateSpecificRelTol; /** Holds the maximum absolute truncation error allowed for specific states of specific dynamic * objects*/ - std::unordered_map dynObjectStateSpecificAbsTol; + std::unordered_map dynObjectStateSpecificAbsTol; }; template diff --git a/src/simulation/dynamics/_GeneralModuleFiles/svIntegratorWeakStochasticRungeKutta.h b/src/simulation/dynamics/_GeneralModuleFiles/svIntegratorWeakStochasticRungeKutta.h new file mode 100644 index 0000000000..91bb5a42cd --- /dev/null +++ b/src/simulation/dynamics/_GeneralModuleFiles/svIntegratorWeakStochasticRungeKutta.h @@ -0,0 +1,537 @@ +/* + ISC License + +Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + +Permission to use, copy, modify, and/or distribute this software for any +purpose with or without fee is hereby granted, provided that the above +copyright notice and this permission notice appear in all copies. + +THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES +WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF +MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR +ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES +WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN +ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF +OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + +*/ + +#ifndef svIntegratorWeakStochasticRungeKutta_h +#define svIntegratorWeakStochasticRungeKutta_h + +#include "../_GeneralModuleFiles/dynamicObject.h" +#include "../_GeneralModuleFiles/dynParamManager.h" +#include "../_GeneralModuleFiles/stateVecStochasticIntegrator.h" +#include "extendedStateVector.h" + +#include +#include +#include +#include +#include +#include +#include + +/** + * Stores the coefficients necessary to use the Stochastic Runge-Kutta method. + * + * The Butcher table looks like: + * + * c0 | A0 | B0 | + * c1 | A1 | B1 | B2 + * ----|--------------------- + * | alpha | beta0 | beta1 + * + * See more in the description of svIntegratorWeakStochasticRungeKutta. + */ +template struct SRKCoefficients { + + /** Array with size = numberStages */ + using StageSizedArray = std::array; + + /** Square matrix with size = numberStages */ + using StageSizedMatrix = std::array; + + // = {} performs zero-initialization + StageSizedArray alpha = {}; /**< "alpha" coefficient array of the SRK method */ + StageSizedArray beta0 = {}; /**< "beta0" coefficient array of the SRK method */ + StageSizedArray beta1 = {}; /**< "beta1" coefficient array of the SRK method */ + StageSizedMatrix A0 = {}; /**< "A0" coefficient matrix of the SRK method */ + StageSizedMatrix B0 = {}; /**< "B0" coefficient matrix of the SRK method */ + StageSizedMatrix A1 = {}; /**< "A1" coefficient matrix of the SRK method */ + StageSizedMatrix B1 = {}; /**< "B1" coefficient matrix of the SRK method */ + StageSizedMatrix B2 = {}; /**< "B2" coefficient matrix of the SRK method */ + + /** "c0" coefficient array of the SRK method + * + * c0 is the row sum of the A(0) matrix: + * + * c0 = sum_j ​a(0)_ij + */ + StageSizedArray c0 = {}; + + /** "c1" coefficient array of the SRK method + * + * c1 is the row sum of the A(1) matrix: + * + * c1 = sum_j ​a(1)_ij + */ + StageSizedArray c1 = {}; +}; + +/** Random variables used in the Weak Stochastic Integrator */ +struct SRKRandomVariables +{ + Eigen::VectorXd Ik; /**< See Eq 3.2 in Tang & Xiao */ + Eigen::MatrixXd Ikl; /**< See Eq 3.3 in Tang & Xiao */ + Eigen::VectorXd Ikk; /**< Same as Ikl(k,k) */ + double xi; /**< See Eq 3.2 in Tang & Xiao */ +}; + +/** Random Generator for the integrator described in: + * + * Tang, X., Xiao, A. Efficient weak second-order stochastic Runge–Kutta methods + * for Itô stochastic differential equations. Bit Numer Math 57, 241–260 (2017). + * https://doi.org/10.1007/s10543-016-0618-9 + */ +class SRKRandomVariableGenerator +{ +public: + + /** Sets the seed to the RNG */ + inline void setSeed(size_t seed) {rng.seed(seed);} + + /** Generates tha random values necessary for one step + * of the integrator.abs_y_is_huge + * + * @param m number of noise sources + * @param h time step, in seconds + */ + inline SRKRandomVariables generate(size_t m, double h) + { + // ensure there's no hidden state on the RVs so setting the seed + // is always consistent + uniform_rv.reset(); + bernoulli_rv.reset(); + + SRKRandomVariables result; + result.Ik.resize(m); + result.Ikl.resize(m, m); + result.Ikk.resize(m); + + // 1) sample I_k per (3.2) + for(size_t k=0; kl) result.Ikl(k,l) = 0.5*(result.Ik(l) + eta2 * result.Ik(k)); + else result.Ikl(k,l) = 0.5*(result.Ik(k)*result.Ik(k) - h); + } + + result.Ikk(k) = result.Ikl(k,k); + } + + return result; + } + +protected: + /** Random Number Generator */ + std::mt19937 rng{std::random_device{}()}; + + /** Uniform random variable, from 0 to 1 */ + std::uniform_real_distribution uniform_rv; + + /** Bernoulli random variable (either 0 or 1, with equal probability) */ + std::bernoulli_distribution bernoulli_rv; +}; + +/** + * The svIntegratorWeakStochasticRungeKutta class implements a state integrator + * that provides weak solutions to problems with stochastic dynamics (SDEs). + * + * The method is described in: + * + * Tang, X., Xiao, A. Efficient weak second-order stochastic Runge–Kutta methods + * for Itô stochastic differential equations. Bit Numer Math 57, 241–260 (2017). + * https://doi.org/10.1007/s10543-016-0618-9 + * + * The method described in said paper is designed for autonomous systems (where + * f and g do not depend on time), so we need to modify it for non-autonomous systems. + * This is simple: we just need to treat the time as a state (with f=1 and g=0). + * + * The resulting pseudocode for the (explicit) method is (note that the order of the + * sums and some elements has moved around to reflect the code): + * + * --- (3.1*) stage definitions --- + * for i = 1..s: + * // drift‐stage + * H0[i] = y_n + * + h * sum_{j=1..i-1}( a0[i][j] * f( t_n + c0[j]*h, H0[j] ) ) + * + sum_{k=1..m}( Ihat[k] * sum_{j=1..i-1}( b0[i][j] * g[k]( t_n + c1[j]*h, Hk[j] ) ) ) + * + * // diffusion‐stages, one per noise source k + * for k = 1..m: + * Hk[i] = y_n + * + h * sum_{j=1..i-1}( a1[i][j] * f( t_n + c0[j]*h, H0[j] ) ) + * + xi * sum_{j=1..i-1}( b1[i][j] * g[k]( t_n + c1[j]*h, Hk[j] ) ) + * + sum_{l=1..m, l!=k}( Ihat_kl[k][l] * sum_{j=1..i-1}( b2[i][j] * g[l]( t_n + c1[i]*h, Hl[i] ) ) ); + * + * where + * c0[j] = sum_{p=1..s}( a0[j][p] ) + * c1[j] = sum_{p=1..s}( a1[j][p] ) + * + * --- (3.2) driving‐variable distributions --- + * For each k: + * P( Ihat[k] = +sqrt(3*h) ) = 1/6 + * P( Ihat[k] = -sqrt(3*h) ) = 1/6 + * P( Ihat[k] = 0 ) = 2/3 + * + * xi = eta1 * sqrt(h) // eta1 is {+1,−1} with equal prob. + * + * --- (3.3) mixed‐integral approximations --- + * for each pair (k,l): + * if k < l: + * Ihat_kl[k][l] = 0.5 * ( Ihat[l] - eta2 * Ihat[l] ) + * else if k > l: + * Ihat_kl[k][l] = 0.5 * ( Ihat[l] + eta2 * Ihat[l] ) + * else // k == l: + * Ihat_kl[k][k] = 0.5 * ( Ihat[k]*Ihat[k] * xi - xi ) + * + * --- (3.1*) state evolution --- + * y_n+1 = y_n + * + h * sum_{i=1..s}( alpha[i] * f( t_n + c0[i]*h, H0[i] ) ) + * + sum_{k=1..m}( Ihat[k] * sum_{i=1..s}( beta0[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) ) + * + sum_{k=1..m}( Ihat_kl[k][k] * sum_{i=1..s}( beta1[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) ) + */ +template class svIntegratorWeakStochasticRungeKutta : public StateVecStochasticIntegrator { +public: + static_assert(numberStages > 0, "One cannot declare Runge Kutta integrators of stage 0"); + + /** Creates an explicit RK integrator for the given DynamicObject using the passed + * coefficients.*/ + svIntegratorWeakStochasticRungeKutta(DynamicObject* dynIn, const SRKCoefficients& coefficients); + + /** Sets the seed for the Random Number Generator used by this integrator. + * + * As a stochastic integrator, random numbers are drawn during each time step. + * By default, a randomly generated seed is used each time. + * + * If the seed is set, the integrator will always draw the same numbers + * during time-stepping. + */ + inline void setRNGSeed(size_t seed) {rvGenerator.setSeed(seed);} + + /** Performs the integration of the associated dynamic objects up to time currentTime+timeStep */ + virtual void integrate(double currentTime, double timeStep) override; + +public: + /** Random Number Generator for the integrator */ + SRKRandomVariableGenerator rvGenerator; + +protected: + /** + * Can be used by subclasses to support passing coefficients + * that are subclasses of SRKCoefficients + */ + svIntegratorWeakStochasticRungeKutta(DynamicObject* dynIn, + std::unique_ptr>&& coefficients); + + /** + * Computes the derivatives of every state given a time and current states. + * + * Internally, this sets the states on the dynamic objects and + * calls the equationsOfMotion methods. + */ + ExtendedStateVector + computeDerivatives(double time, double timeStep); + + /** + * Computes the diffusions of every state given a time and current states. + * + * Internally, this sets the states on the dynamic objects and + * calls the equationsOfMotionDiffusion methods. + */ + std::vector + computeDiffusions(double time, double timeStep, const std::vector& stateIdToNoiseIndexMaps); + + /** + * Computes the diffusions for the states and noise index + * in ``stateIdToNoiseIndexMap`` given a time and current states. + * + * Internally, this sets the states on the dynamic objects and + * calls the equationsOfMotionDiffusion methods. + */ + ExtendedStateVector computeDiffusions(double time, + double timeStep, + const StateIdToIndexMap& stateIdToNoiseIndexMap); + + /** Utility function, computes: + * + * result = sum_{i=0..length-1} factors[i]*vectors[i] + */ + ExtendedStateVector scaledSum(const std::array& factors, const std::array& vectors, size_t length); + +protected: + // coefficients is stored as a pointer to support polymorphism + /** Coefficients to be used in the method */ + const std::unique_ptr> coefficients; +}; + +template +svIntegratorWeakStochasticRungeKutta::svIntegratorWeakStochasticRungeKutta( + DynamicObject* dynIn, + const SRKCoefficients& coefficients) + : StateVecStochasticIntegrator(dynIn), + coefficients(std::make_unique>(coefficients)) +{ +} + +template +svIntegratorWeakStochasticRungeKutta::svIntegratorWeakStochasticRungeKutta( + DynamicObject* dynIn, + std::unique_ptr>&& coefficients) + : StateVecStochasticIntegrator(dynIn), coefficients(std::move(coefficients)) +{ +} + +template +void svIntegratorWeakStochasticRungeKutta::integrate(double currentTime, double timeStep) +{ + ExtendedStateVector currentState = ExtendedStateVector::fromStates(dynPtrs); + + // this is a map (ExtendedStateId -> noise index) for each of the noise sources + // (so the length of this vector should be m) + std::vector stateIdToNoiseIndexMaps = getStateIdToNoiseIndexMaps(); + + // Sample the random variables + auto rv = rvGenerator.generate(stateIdToNoiseIndexMaps.size(), timeStep); + + std::array f_H0; // f(H_i^(0)) for i=0..s-1 + std::vector> g_Hk; // g^k(H_i^(k)) for i=0..s-1; k=0..m-1 + + // i = 0 (first stage) we do here becase we can optimize it + // Note that current state == H_0^(0) == H_0^(k) + f_H0.at(0) = computeDerivatives( + currentTime, + timeStep + ); + + std::vector diffs = computeDiffusions( + currentTime, + timeStep, + stateIdToNoiseIndexMaps + ); + for (auto&& diff : std::move(diffs)) + { + g_Hk.emplace_back().at(0) = std::move(diff); + } + + for (size_t i = 1; i < numberStages; i++) + { + /* // drift‐stage + * H0[i] = y_n + * + h * sum_{j=1..i-1}( a0[i][j] * f( t_n + c0[j]*h, H0[j] ) ) + * + sum_{k=1..m}( Ihat[k] * sum_{j=1..i-1}( b0[i][j] * g[k]( t_n + c1[j]*h, Hk[j] ) ) ) + */ + + // y_n + // and store the state in the dynPtrs + currentState.setStates(dynPtrs); + + // sum_{j=1..i-1}( a0[i][j] * f( t_n + c0[j]*h, H0[j] ) ) + // and store the derivative in the dynPtrs + scaledSum(coefficients->A0.at(i), f_H0, i).setDerivatives(dynPtrs); + + // sum_{k=1..m}( Ihat[k] * sum_{j=1..i-1}( b0[i][j] * g[k]( t_n + c1[j]*h, Hk[j] ) ) ) + for (size_t k = 0; k < stateIdToNoiseIndexMaps.size(); k++) + { + // sum_{j=1..i-1}( b0[i][j] * g[k]( t_n + c1[j]*h, Hk[j] ) ) + // and store the diffusion for the noise source k in the dynPtrs + scaledSum(coefficients->B0.at(i), g_Hk.at(k), i).setDiffusions(dynPtrs, stateIdToNoiseIndexMaps.at(k)); + } + + // This sets the current state to H_i^(0) + propagateState(timeStep, rv.Ik, stateIdToNoiseIndexMaps); + + // This computes f(t_n + c0[i]*h, H0[i]) for the next steps + f_H0.at(i) = computeDerivatives(currentTime + coefficients->c0.at(i)*timeStep, timeStep); + + + /* // diffusion‐stages, one per noise source k + * for k = 1..m: + * Hk[i] = y_n + * + h * sum_{j=1..i-1}( a1[i][j] * f( t_n + c0[j]*h, H0[j] ) ) + * + xi * sum_{j=1..i-1}( b1[i][j] * g[k]( t_n + c1[j]*h, Hk[j] ) ) + * + sum_{l=1..m, l!=k}( Ihat_kl[k][l] * sum_{j=1..i-1}( b2[i][j] * g[l]( t_n + c1[i]*h, Hl[i] ) ) ); + */ + for (size_t k = 0; k < stateIdToNoiseIndexMaps.size(); k++) + { + // y_n + // and store the state in the dynPtrs + currentState.setStates(dynPtrs); + + // sum_{j=1..i-1}( a0[i][j] * f( t_n + c0[j]*h, H0[j] ) ) + // and store the derivative in the dynPtrs + scaledSum(coefficients->A1.at(i), f_H0, i).setDerivatives(dynPtrs); + + // sum_{j=1..i-1}( b1[i][j] * g[k]( t_n + c1[j]*h, Hk[j] ) ) + // and store the diffusion for the noise source k in the dynPtrs + scaledSum(coefficients->B1.at(i), g_Hk.at(k), i).setDiffusions(dynPtrs, stateIdToNoiseIndexMaps.at(k)); + + for (size_t l = 0; l < stateIdToNoiseIndexMaps.size(); l++) + { + if (l==k) continue; + + // sum_{j=1..i-1}( b2[i][j] * g[l]( t_n + c1[i]*h, Hl[i] ) ) + // and store the diffusion for the noise source k in the dynPtrs + scaledSum(coefficients->B2.at(i), g_Hk.at(l), i).setDiffusions(dynPtrs, stateIdToNoiseIndexMaps.at(l)); + } + + // The kth element uses pseudo timestep xi, the rest Ikl[k,l] + Eigen::VectorXd pseudoTimeStep(stateIdToNoiseIndexMaps.size()); + for (size_t l = 0; l < stateIdToNoiseIndexMaps.size(); l++) + { + if (l==k) pseudoTimeStep(l) = rv.xi; + else pseudoTimeStep(l) = rv.Ikl(k,l); + } + + // This sets the current state to H_i^(k) + propagateState(timeStep, pseudoTimeStep, stateIdToNoiseIndexMaps); + + // This computes g^(k)(t_n + c1[i]*h, Hk[i]) for the next steps + g_Hk.at(k).at(i) = computeDiffusions(currentTime + coefficients->c1.at(i)*timeStep, timeStep, stateIdToNoiseIndexMaps.at(k)); + } + } + + + + /* y_n+1 = y_n + * + h * sum_{i=1..s}( alpha[i] * f( t_n + c0[i]*h, H0[i] ) ) + * + sum_{k=1..m}( Ihat[k] * sum_{i=1..s}( beta0[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) ) + * + sum_{k=1..m}( Ihat_kl[k][k] * sum_{i=1..s}( beta1[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) ) + */ + + // y_n + // and store the state in the dynPtrs + currentState.setStates(dynPtrs); + + // sum_{i=1..s}( alpha[i] * f( t_n + c0[i]*h, H0[i] ) ) + // and store the derivative in the dynPtrs + scaledSum(coefficients->alpha, f_H0, numberStages).setDerivatives(dynPtrs); + + // sum_{k=1..m}( Ihat[k] * sum_{i=1..s}( beta0[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) ) + for (size_t k = 0; k < stateIdToNoiseIndexMaps.size(); k++) + { + // sum_{i=1..s}( beta0[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) + // and store the diffusion for the noise source k in the dynPtrs + scaledSum(coefficients->beta0, g_Hk.at(k), numberStages).setDiffusions(dynPtrs, stateIdToNoiseIndexMaps.at(k)); + } + + /* + * This computes: + * + * y_n + * + h * sum_{i=1..s}( alpha[i] * f( t_n + c0[i]*h, H0[i] ) ) + * + sum_{k=1..m}( Ihat[k] * sum_{i=1..s}( beta0[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) ) + * + * BUT! We're missing the 'last line' of the full y_n+1 + * because we need to set the diffusions to different values + */ + propagateState(timeStep, rv.Ik, stateIdToNoiseIndexMaps); + + // sum_{k=1..m}( Ihat_kl[k] * sum_{i=1..s}( beta1[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) ) + for (size_t k = 0; k < stateIdToNoiseIndexMaps.size(); k++) + { + // sum_{i=1..s}( beta1[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) + // and store the diffusion for the noise source k in the dynPtrs + scaledSum(coefficients->beta1, g_Hk.at(k), numberStages).setDiffusions(dynPtrs, stateIdToNoiseIndexMaps.at(k)); + } + + /* + * This computes: + * + * y_n+1 = y_n + * + h * sum_{i=1..s}( alpha[i] * f( t_n + c0[i]*h, H0[i] ) ) + * + sum_{k=1..m}( Ihat[k] * sum_{i=1..s}( beta0[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) ) + * + sum_{k=1..m}( Ihat_kl[k][k] * sum_{i=1..s}( beta1[i] * g[k]( t_n + c1[i]*h, Hk[i] ) ) ) + * + * By propagating the previously computed state (see previous call + * to propagateState) with only the last row. To do so, we replaced the + * diffusions with those of the last row and set the timeStep to zero + * so that the contribution of the derivative would not be double counted. + */ + propagateState(0, rv.Ikk, stateIdToNoiseIndexMaps); + + // now the dynPtrs have y_n+1 in them!! +} + +template +ExtendedStateVector +svIntegratorWeakStochasticRungeKutta::computeDerivatives(double time, double timeStep) +{ + for (auto dynPtr : this->dynPtrs) { + dynPtr->equationsOfMotion(time, timeStep); + } + + return ExtendedStateVector::fromStateDerivs(this->dynPtrs); +} + +template +std::vector +svIntegratorWeakStochasticRungeKutta::computeDiffusions(double time, + double timeStep, + const std::vector& stateIdToNoiseIndexMaps + ) +{ + for (auto dynPtr : this->dynPtrs) { + dynPtr->equationsOfMotionDiffusion(time, timeStep); + } + + return ExtendedStateVector::fromStateDiffusions(this->dynPtrs, stateIdToNoiseIndexMaps); +} + +template +ExtendedStateVector +svIntegratorWeakStochasticRungeKutta::computeDiffusions(double time, + double timeStep, + const StateIdToIndexMap& stateIdToNoiseIndexMap + ) +{ + for (auto dynPtr : this->dynPtrs) { + dynPtr->equationsOfMotionDiffusion(time, timeStep); + } + + return ExtendedStateVector::fromStateDiffusions(this->dynPtrs, stateIdToNoiseIndexMap); +} + +template +inline ExtendedStateVector +svIntegratorWeakStochasticRungeKutta::scaledSum(const std::array& factors, + const std::array& vectors, + size_t length) +{ + ExtendedStateVector result = vectors.at(0) * factors.at(0); + for (size_t i = 1; i < length; i++) { + if (factors.at(i) == 0) continue; + result += vectors.at(i) * factors.at(i); + } + return result; +} + +#endif /* svIntegratorWeakStochasticRungeKutta_h */ diff --git a/src/simulation/dynamics/spacecraft/spacecraft.cpp b/src/simulation/dynamics/spacecraft/spacecraft.cpp index f13382528e..8a565ce867 100755 --- a/src/simulation/dynamics/spacecraft/spacecraft.cpp +++ b/src/simulation/dynamics/spacecraft/spacecraft.cpp @@ -449,6 +449,11 @@ void Spacecraft::equationsOfMotion(double integTimeSeconds, double timeStep) } } +void Spacecraft::equationsOfMotionDiffusion(double integTimeSeconds, double timeStep) +{ + this->equationsOfMotion(integTimeSeconds, timeStep); +} + /*! Prepare for integration process @param integrateToThisTime Time to integrate to */ diff --git a/src/simulation/dynamics/spacecraft/spacecraft.h b/src/simulation/dynamics/spacecraft/spacecraft.h index c167e43bbe..376ac2c38b 100755 --- a/src/simulation/dynamics/spacecraft/spacecraft.h +++ b/src/simulation/dynamics/spacecraft/spacecraft.h @@ -87,6 +87,7 @@ class Spacecraft : public DynamicObject{ void UpdateState(uint64_t CurrentSimNanos); //!< Runtime hook back into Basilisk arch void linkInStates(DynParamManager& statesIn); //!< Method to get access to the hub's states void equationsOfMotion(double integTimeSeconds, double timeStep); //!< This method computes the equations of motion for the whole system + void equationsOfMotionDiffusion(double integTimeSeconds, double timeStep); //!< This method computes the diffusion equations of motion for the whole system void addStateEffector(StateEffector *newSateEffector); //!< Attaches a stateEffector to the system void addDynamicEffector(DynamicEffector *newDynamicEffector); //!< Attaches a dynamicEffector void preIntegration(double callTime) final; //!< method to perform pre-integration steps diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.cpp index d6f88a9479..37130661ef 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.cpp @@ -31,7 +31,21 @@ void MJQPosStateData::configure(mjModel* mujocoModel) this->stateDeriv.resize(mujocoModel->nv, 1); } -void MJQPosStateData::propagateState(double dt) +void MJQPosStateData::propagateState(double dt, std::vector pseudoStep) { mj_integratePos(this->mujocoModel, this->state.data(), this->stateDeriv.data(), dt); + + if (getNumNoiseSources() > 0 && pseudoStep.size() == 0) + { + auto errorMsg = "State " + this->getName() + " has stochastic dynamics, but " + + "the integrator tried to propagate it without pseudoSteps. Are you sure " + + "you are using a stochastic integrator?"; + bskLogger.bskLog(BSK_ERROR, errorMsg.c_str()); + throw std::invalid_argument(errorMsg); + } + + for (size_t i = 0; i < getNumNoiseSources(); i++) + { + mj_integratePos(this->mujocoModel, this->state.data(), this->stateDiffusion.at(i).data(), pseudoStep.at(i)); + } } diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.h index 8c391e2e57..ce810ee0db 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.h @@ -66,11 +66,16 @@ class MJQPosStateData : public StateData * @brief Propagates the state over a time step. * * This method integrates the position state using the state derivative - * over the given time step. + * over the given time step:: * - * @param dt The time step for propagation. + * x += f(t,x)*h + g_0(t,x)*pseudoStep[0] + g_1(t,x)*pseudoStep[1] ... + * + * @param h The time step for propagation. + * @param pseudoStep For states driven by stochastic dynamics, this + * represents the random pseudotimestep. The length of this input must + * match the number of noise sources of this state (``getNumNoiseSources()``) */ - void propagateState(double dt) override; + void propagateState(double h, std::vector pseudoStep = {}) override; protected: mjModel* mujocoModel; ///< Pointer to the MuJoCo model associated with the state. diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp index 69ffaf74a3..03ecb077b6 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp @@ -32,6 +32,7 @@ #include #include #include +#include MJScene::MJScene(std::string xml, const std::vector& files) : spec(*this, xml, files) @@ -59,13 +60,31 @@ void MJScene::AddFwdKinematicsToDynamicsTask(int32_t priority) this->AddModelToDynamicsTask(this->ownedSysModel.back().get(), priority); } -void MJScene::SelfInit() { this->dynamicsTask.SelfInitTaskList(); } +void MJScene::AddModelToDiffusionDynamicsTask(SysModel* model, int32_t priority) +{ + this->dynamicsDiffusionTask.AddNewObject(model, priority); +} + +void MJScene::AddFwdKinematicsToDiffusionDynamicsTask(int32_t priority) +{ + this->ownedSysModel.emplace_back(std::make_unique(*this)); + this->ownedSysModel.back()->ModelTag = "FwdKinematics" + std::to_string(this->ownedSysModel.size()-1); + this->AddModelToDiffusionDynamicsTask(this->ownedSysModel.back().get(), priority); +} + +void MJScene::SelfInit() +{ + this->dynamicsTask.SelfInitTaskList(); + this->dynamicsDiffusionTask.SelfInitTaskList(); +} void MJScene::Reset(uint64_t CurrentSimNanos) { this->timeBefore = CurrentSimNanos * NANO2SEC; this->dynamicsTask.TaskName = "Dynamics:" + this->ModelTag; this->dynamicsTask.ResetTaskList(CurrentSimNanos); + this->dynamicsDiffusionTask.TaskName = "DiffusionDynamics:" + this->ModelTag; + this->dynamicsDiffusionTask.ResetTaskList(CurrentSimNanos); this->initializeDynamics(); this->writeOutputStateMessages(CurrentSimNanos); } @@ -95,16 +114,31 @@ void MJScene::initializeDynamics() } // Register the states of the models in the dynamics task - for (auto[_, sysModelPtr] : this->dynamicsTask.TaskModels) + std::unordered_set alreadyRegisteredModels; + auto registerStatesOnSysModel = [this, &alreadyRegisteredModels](SysModel* sysModelPtr) { if (auto statefulSysModelPtr = dynamic_cast(sysModelPtr)) { + // Don't registerStates in a model twice! + if (alreadyRegisteredModels.count(statefulSysModelPtr) > 0) return; + statefulSysModelPtr->registerStates(DynParamRegisterer( this->dynManager, sysModelPtr->ModelTag.empty() ? std::string("model") : sysModelPtr->ModelTag + "_" + std::to_string(sysModelPtr->moduleID) + "_" )); + + alreadyRegisteredModels.emplace(statefulSysModelPtr); } + }; + + for (auto[_, sysModelPtr] : this->dynamicsTask.TaskModels) + { + registerStatesOnSysModel(sysModelPtr); + } + for (auto[_, sysModelPtr] : this->dynamicsDiffusionTask.TaskModels) + { + registerStatesOnSysModel(sysModelPtr); } } @@ -206,6 +240,12 @@ void MJScene::equationsOfMotion(double t, double timeStep) } } +void MJScene::equationsOfMotionDiffusion(double t, double timeStep) +{ + auto nanos = static_cast(t * SEC2NANO); + this->dynamicsDiffusionTask.ExecuteTaskList(nanos); +} + void MJScene::preIntegration(double callTime) { this->timeStep = callTime - this->timeBefore; } void MJScene::postIntegration(double callTime) @@ -221,6 +261,7 @@ void MJScene::postIntegration(double callTime) // time with the final/integrated state. This also calls // MJFwdKinematics::fwdKinematics equationsOfMotion(callTime, 0); + equationsOfMotionDiffusion(callTime, 0); } else { diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h index 26dfcd6101..4fd0d41020 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h @@ -290,6 +290,12 @@ class MJScene : public DynamicObject * task typically compute and apply forces and torques on the system * that depend on the simulation time or the state of the multi-body. * + * Models in this task should contribute to computing ``f`` in: + * + * dx = f(t,x) dt + * + * where ``x`` represents the state of this dynamic object. + * * @warning `SysModel` added to the dynamics task should be memoryless. * That is, any output message computed should depend strictly on * input messages and the current time/states. It should save values @@ -303,6 +309,23 @@ class MJScene : public DynamicObject */ void AddModelToDynamicsTask(SysModel* model, int32_t priority = -1); + /** @brief Adds a model to the diffusion dynamics task. + * + * This task and function are only relevant when the dynamics of + * the system are stochastic. + * + * Similar to ``AddModelToDynamicsTask``, except that models in this + * task contribute to computing ``g`` in: + * + * dx = f(t,x)dt + g(t,x)dW + * + * where ``x`` is the state of the system, ``f`` represents the + * drift term of the dynamics (classic time derivative), and + * ``g`` is the diffusion term of the dynamics (which evaluates + * the impact of the random 'noise' ``W`` on the dynamics). + */ + void AddModelToDiffusionDynamicsTask(SysModel* model, int32_t priority = -1); + /** * @brief Adds forward kinematics to the dynamics task. * @@ -334,6 +357,17 @@ class MJScene : public DynamicObject */ void AddFwdKinematicsToDynamicsTask(int32_t priority); + /** + * @brief Adds forward kinematics to the diffusion dynamics task. + * + * See ``AddModelToDiffusionDynamicsTask`` and ``AddFwdKinematicsToDynamicsTask``. + * + * By default, the diffusion dynamics task has no forward kinematics + * model, so one must be added if the diffusion term depends on the + * forward kinematics of the multibody. + */ + void AddFwdKinematicsToDiffusionDynamicsTask(int32_t priority); + /** * @brief Calls `SelfInit` on all system models in the dynamics task. */ @@ -367,11 +401,30 @@ class MJScene : public DynamicObject * joints... These information is translated, through MuJoCo, into * first derivatives of the joint states and the mass of bodies. * + * Computes ``f`` in: + * + * dx = f(t,x) dt + * + * where ``x`` represents the current state of the dynamics. + * * @param t The current simulation time in seconds. * @param timeStep The integration time step. */ void equationsOfMotion(double t, double timeStep) override; + /** + * @brief Computes the diffusion of the dynamics of the system. + * + * Only relevant for systems with stochastic dynamics. + * + * Computes ``g`` in: + * + * dx = f(t,x) dt + g(t,x)dW + * + * where ``x`` represents the current state of the dynamics. + */ + void equationsOfMotionDiffusion(double t, double timeStep) override; + /** * @brief Prepares the system before actual integration. * @@ -558,6 +611,7 @@ class MJScene : public DynamicObject bool forwardKinematicsStale = true; ///< Flag indicating stale forward kinematics. SysModelTask dynamicsTask; ///< Task managing models involved in the dynamics of this scene. + SysModelTask dynamicsDiffusionTask; ///< Task managing models involved in the diffusion stochastic dynamics of this scene. std::vector> ownedSysModel; ///< System models that should be cleared on this scene destruction. MJQPosStateData* qposState; ///< Position state data. diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.h index 8ec93423d1..2eb758f5db 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.h @@ -64,6 +64,39 @@ class DynParamRegisterer ); } + /** Used when more than one state have dynamics perturbed + * by the same noise process. + * + * For example, consider the following SDE: + * + * dx_0 = f_0(t,x)dt + g_00(t,x)dW_0 + g_01(t,x)dW_1 + * dx_1 = f_1(t,x)dt + g_11(t,x)dW_1 + * + * In this case, state 'x_0' is affected by 2 sources of noise + * and 'x_1' by 1 source of noise. However, the source 'W_1' is + * shared between 'x_0' and 'x_1'. + * + * This function is called like: + * + * \code + * dynParamManager.registerSharedNoiseSource({ + * {myStateX0, 1}, + * {myStateX1, 0} + * }); + * \endcode + * + * which means that the 2nd noise source of the ``StateData`` 'myStateX0' + * and the first noise source of the ``StateData`` 'myStateX1' actually + * correspond to the same noise process. + * + * NOTE: Some stochastic integrators do not support sharing noise sources. + * In this case, this method should raise ``std::logic_error``. + */ + inline void registerSharedNoiseSource(std::vector> in) + { + manager.registerSharedNoiseSource(std::move(in)); + } + protected: DynParamManager& manager; ///< wrapped manager std::string stateNamePrefix; ///< prefix added to all registered state names diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.i b/src/simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.i new file mode 100644 index 0000000000..e60a298bbd --- /dev/null +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.i @@ -0,0 +1,51 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ +%module cStatefulSysModel +%{ + #include "StatefulSysModel.h" +%} + +%include "architecture/utilities/bskLogging.h" +%import "architecture/_GeneralModuleFiles/sys_model.i" +%import "simulation/dynamics/_GeneralModuleFiles/dynParamManager.i" + +// We don't need to construct the DynParamRegisterer on the Python side +%ignore DynParamRegisterer::DynParamRegisterer; + +// Current limitation of SWIG for complex templated types like +// std::pair, we need to declare these manually +%traits_swigtype(StateData); +%fragment(SWIG_Traits_frag(StateData)); + +%extend DynParamRegisterer { + // SWIG doesnt like const StateData& so we have to use const StateData* and convert + void registerSharedNoiseSource(std::vector> list) { + // Convert from pointer pairs to reference pairs + std::vector> refList; + refList.reserve(list.size()); + for (const auto& p : list) { + refList.emplace_back(*p.first, p.second); + } + $self->registerSharedNoiseSource(refList); + } +} + +%ignore DynParamRegisterer::registerSharedNoiseSource(std::vector>); + +%include "StatefulSysModel.h" diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/pyStatefulSysModel.i b/src/simulation/mujocoDynamics/_GeneralModuleFiles/pyStatefulSysModel.i index f893eade38..23d03467b8 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/pyStatefulSysModel.i +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/pyStatefulSysModel.i @@ -16,7 +16,7 @@ OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. */ -%module(directors="1",threads="1") StatefulSysModel +%module(directors="1",threads="1",package="Basilisk.simulation") StatefulSysModel %{ #include "StatefulSysModel.h" %} @@ -37,7 +37,7 @@ from Basilisk.architecture.swig_common_model import * %feature("pythonappend") StatefulSysModel::StatefulSysModel %{ self.__super_init_called__ = True%} %rename("_StatefulSysModel") StatefulSysModel; -%include "StatefulSysModel.h" +%include "StatefulSysModel.i" %template(registerState) DynParamRegisterer::registerState; diff --git a/src/tests/test_scenarioMujoco.py b/src/tests/test_scenarioMujoco.py index d627e7833c..585919d799 100644 --- a/src/tests/test_scenarioMujoco.py +++ b/src/tests/test_scenarioMujoco.py @@ -24,6 +24,7 @@ import importlib import pytest +from Basilisk.utilities import unitTestSupport try: from Basilisk.simulation import mujoco @@ -32,8 +33,9 @@ except: couldImportMujoco = False +THIS_FOLDER = os.path.dirname(__file__) SCENARIO_FOLDER = os.path.join( - os.path.dirname(__file__), "..", "..", "examples", "mujoco" + THIS_FOLDER, "..", "..", "examples", "mujoco" ) SCENARIO_FILES = [ filename[:-3] @@ -49,8 +51,11 @@ @pytest.mark.scenarioTest def test_scenarios(scenario: str): module = importlib.import_module(scenario) - module.run() # Every mujoco scenario should have a run function + figures = module.run() # Every mujoco scenario should have a run function + if figures is not None: + for pltName, plt in figures.items(): + unitTestSupport.saveScenarioFigure(f"{scenario}_{pltName}", plt, THIS_FOLDER) if __name__ == "__main__": pytest.main([__file__])