diff --git a/docs/source/Support/bskReleaseNotes.rst b/docs/source/Support/bskReleaseNotes.rst index 88a00cd2ad..4ba0200765 100644 --- a/docs/source/Support/bskReleaseNotes.rst +++ b/docs/source/Support/bskReleaseNotes.rst @@ -74,6 +74,10 @@ Version |release| The new option is useful when the user wants to ensure that the simulation runs for at least the specified time, instead of at most the specified time. - Fixed documentation quote typos that caused documentation build errors with doxygen version 1.15 and newer. +- 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.8.0 (August 30, 2025) diff --git a/examples/_default.rst b/examples/_default.rst index c6826af614..5507172f55 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/scenarioFormationFlyingStochasticDrag.py b/examples/mujoco/scenarioFormationFlyingStochasticDrag.py new file mode 100644 index 0000000000..130e5fd3c3 --- /dev/null +++ b/examples/mujoco/scenarioFormationFlyingStochasticDrag.py @@ -0,0 +1,1053 @@ +# +# 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. + +from dataclasses import dataclass +from collections import namedtuple +from typing import Optional, Dict, Any +import os + +import numpy as np +import matplotlib.pyplot as plt +from matplotlib.collections import LineCollection + +from Basilisk.architecture import messaging +from Basilisk.simulation import mujoco +from Basilisk.simulation import svIntegrators +from Basilisk.simulation import pointMassGravityModel +from Basilisk.simulation import NBodyGravity +from Basilisk.simulation import MJAerodynamicDrag +from Basilisk.simulation import MJLinearTimeInvariantSystem +from Basilisk.simulation import MJMeanRevertingNoise +from Basilisk.simulation import exponentialAtmosphere +from Basilisk.utilities import SimulationBaseClass +from Basilisk.simulation import MJCmdForceInertialToForceAtSite +from Basilisk.utilities import macros +from Basilisk.utilities import orbitalMotion +from Basilisk.utilities import simIncludeGravBody +from Basilisk.utilities import simSetPlanetEnvironment +from Basilisk.fswAlgorithms import orbElemOffset +from Basilisk.simulation import orbElemConvert + +"""Two cannon-ball spacecraft with optional drag, feedback control, and feedforward drag cancellation. + +The chief and deputy spacecraft are modeled as simple point-mass bodies with sliding joints. +Their inertial properties are set directly, with no sub-bodies or moving parts. +""" + +CHIEF_DEPUTY_SCENE_XML = r""" + + + + + + + + + + + + + + + + + +""" + + +@dataclass +class ScenarioConfig: + """Scenario configuration toggles for the chief-deputy simulation.""" + drag: bool = True + control: bool = True + feedforward: bool = True + integralControl: bool = True + stochastic: bool = False + + +@dataclass +class PlotConfig: + """Plot configuration and output control.""" + showPlots: bool = False + saveFolder: Optional[str] = None + + plotChiefInertial: bool = True + plotChiefDrag: bool = True + plotChiefDensity: bool = True + plotChiefCd: bool = True + plotControl: bool = True + plotOeChief: bool = True + plotOeDeputy: bool = True + plotOeDiff: bool = True + plotHillXY: bool = True + + +@dataclass +class SimulationModels: + """Handles to all models and messages that must remain in scope.""" + scSim: SimulationBaseClass.SimBaseClass + integrator: svIntegrators.StateVecIntegrator + scene: mujoco.MJScene + chiefBody: mujoco.MJBody + deputyBody: mujoco.MJBody + gravity: NBodyGravity.NBodyGravity + atmo: exponentialAtmosphere.ExponentialAtmosphere + + # Drag models + dragChief: Optional[MJAerodynamicDrag.AerodynamicDrag] = None + dragDeputy: Optional[MJAerodynamicDrag.AerodynamicDrag] = None + stochDragGeomChief: Optional[MJMeanRevertingNoise.StochasticDragCoeff] = None + + # Control models + orbElemConvertChief: Optional[orbElemConvert.OrbElemConvert] = None + orbElemConvertDeputy: Optional[orbElemConvert.OrbElemConvert] = None + orbElemTarget: Optional[orbElemOffset.OrbElemOffset] = None + oeFeedback: Optional[MJLinearTimeInvariantSystem.OrbitalElementControl] = None + inertialToSiteForce: Optional[MJCmdForceInertialToForceAtSite.CmdForceInertialToForceAtSite] = None + + # Feedforward drag models + orbElemConvertTarget: Optional[orbElemConvert.OrbElemConvert] = None + feedforwardDragTarget: Optional[MJAerodynamicDrag.AerodynamicDrag] = None + feedforwardDragDeputy: Optional[MJLinearTimeInvariantSystem.ForceAtSiteLTI] = None + + # Force actuators kept alive explicitly + controlActuatorDeputy: Optional[mujoco.MJForceActuator] = None + feedforwardDragActuatorTarget: Optional[mujoco.MJForceActuator] = None + feedforwardDragActuatorDeputy: Optional[mujoco.MJForceActuator] = None + + # Messages that must not be garbage collected + dragGeomChiefMsg: Optional[messaging.DragGeometryMsg] = None + dragGeomDeputyMsg: Optional[messaging.DragGeometryMsg] = None + stochDragGeomChiefMsg: Optional[messaging.DragGeometryMsg] = None + offsetElementsMsg: Optional[messaging.ClassicElementsMsg] = None + + +@dataclass +class RecorderHandles: + """Handles to all recorders used in the simulation.""" + chiefState: Any + deputyState: Any + dragChief: Optional[Any] = None + densityChief: Optional[Any] = None + cdChief: Optional[Any] = None + oeFeedbackForce: Optional[Any] = None + feedforwardDragTarget: Optional[Any] = None + feedforwardDragDeputy: Optional[Any] = None + + +@dataclass +class InitialOrbitInfo: + """Initial orbit geometry for the chief-deputy pair.""" + semiMajorAxis: float + orbitPeriod: float + + +@dataclass +class SimulationTimeSeries: + """All time history data required for plotting and analysis.""" + times: np.ndarray + posChief: np.ndarray + velChief: np.ndarray + posDeputy: np.ndarray + velDeputy: np.ndarray + oeChief: np.ndarray + oeDeputy: np.ndarray + oeDiff: np.ndarray + hillPosition: np.ndarray + hillVelocity: np.ndarray + + +def run(scenarioConfig: ScenarioConfig = ScenarioConfig(), plotConfig: PlotConfig = PlotConfig(showPlots=False)) -> Dict[str, plt.Figure]: + """ + Run the chief-deputy relative motion scenario with optional drag, feedback control, + and drag feedforward cancellation. + + This function builds the simulation, sets initial conditions, runs the dynamics, + collects time histories, generates figures, and optionally shows and/or saves them. + + Parameters + ---------- + scenarioConfig + High level configuration toggles for drag, control, feedforward, etc. + plotConfig + Plot visibility and output configuration. + + Returns + ------- + dict + Mapping from figure name to matplotlib Figure instance. + """ + planet = simIncludeGravBody.BODY_DATA["earth"] + + # Proportional and integral gains on classical elements (with mean anomaly) + kProp = np.array([1e5, 1e5, 0, 0, 10, 10]) + kInt = np.array([100, 0, 0, 0, 0, 0]) + + # Controller target delta elements + targetOeDiffPayload = messaging.ClassicElementsMsgPayload( + a=0, + e=0.0005, + i=0, + Omega=0, + omega=0, + f=0, + ) + + # Time step based on whether stochastic integration is used + dt = 0.1 if scenarioConfig.stochastic else 10.0 + + # Build all simulation models and their connections + models = createSimulationModels( + planet=planet, + scenarioConfig=scenarioConfig, + targetOeDiffPayload=targetOeDiffPayload, + kProp=kProp, + kInt=kInt, + dt=dt, + ) + + # Attach recorders that sample the main quantities of interest + recorders = addRecorders(models) + + # Set initial chief and deputy orbits and propagate to Cartesian states + orbitInfo = setInitialConditions(models, planet) + + # Run the simulation for a fixed number of orbits + runSimulation(models, orbitInfo) + + # Extract time series from recorders and postprocess into useful coordinates + timeSeries = extractTimeSeries(planet, recorders) + + # Generate figures, optionally save to disk, and optionally show interactively + figures = createFigures( + scenarioConfig=scenarioConfig, + plotConfig=plotConfig, + models=models, + recorders=recorders, + timeSeries=timeSeries, + targetOeDiffPayload=targetOeDiffPayload, + kProp=kProp, + ) + + if plotConfig.showPlots: + plt.show() + + return figures + + +def createSimulationModels( + planet: Any, + scenarioConfig: ScenarioConfig, + targetOeDiffPayload: messaging.ClassicElementsMsgPayload, + kProp: np.ndarray, + kInt: np.ndarray, + dt: float, +) -> SimulationModels: + """Create and connect all simulation models, returning a container with live references.""" + # Create simulation, process, and task + scSim = SimulationBaseClass.SimBaseClass() + process = scSim.CreateNewProcess("test") + task = scSim.CreateNewTask("test", macros.sec2nano(dt)) + process.addTask(task) + + # Create the MuJoCo scene and add it as a dynamic object + scene = mujoco.MJScene(CHIEF_DEPUTY_SCENE_XML) + scene.extraEoMCall = True + scSim.AddModelToTask("test", scene, 1) + + # Select integrator + if scenarioConfig.stochastic: + integrator = svIntegrators.svStochIntegratorW2Ito1(scene) + else: + integrator = svIntegrators.svIntegratorRKF45(scene) + integrator.relTol = 1e-10 + integrator.absTol = 1e-10 + scene.setIntegrator(integrator) + + # Get MuJoCo bodies + chiefBody: mujoco.MJBody = scene.getBody("chief") + deputyBody: mujoco.MJBody = scene.getBody("deputy") + + # Central body gravity + gravity = NBodyGravity.NBodyGravity() + gravity.ModelTag = "gravity" + scene.AddModelToDynamicsTask(gravity) + + gravityModel = pointMassGravityModel.PointMassGravityModel() + gravityModel.muBody = planet.mu + gravity.addGravitySource("earth", gravityModel, True) + gravity.addGravityTarget("chief", chiefBody) + gravity.addGravityTarget("deputy", deputyBody) + + # Exponential atmosphere model + atmo = exponentialAtmosphere.ExponentialAtmosphere() + atmo.ModelTag = "ExpAtmo" + simSetPlanetEnvironment.exponentialAtmosphere(atmo, "earth") + atmo.addSpacecraftToModel(chiefBody.getCenterOfMass().stateOutMsg) + atmo.addSpacecraftToModel(deputyBody.getCenterOfMass().stateOutMsg) + scene.AddModelToDynamicsTask(atmo) + + # Initialize container for all models/messages that must stay in scope + models = SimulationModels( + scSim=scSim, + integrator=integrator, + scene=scene, + chiefBody=chiefBody, + deputyBody=deputyBody, + gravity=gravity, + atmo=atmo, + ) + + # Optional drag on chief and deputy + if scenarioConfig.drag: + configureDragModels(models, scenarioConfig) + + # Optional feedback control based on orbital elements + if scenarioConfig.control: + configureControlModels(models, planet, targetOeDiffPayload, kProp, kInt, scenarioConfig) + + # Optional drag feedforward cancellation + if scenarioConfig.feedforward: + configureFeedforwardModels(models, planet) + + return models + + +def configureDragModels(models: SimulationModels, scenarioConfig: ScenarioConfig) -> None: + """Attach aerodynamic drag models to chief and deputy, with optional stochastic Cd.""" + scene = models.scene + atmo = models.atmo + chiefBody = models.chiefBody + deputyBody = models.deputyBody + + # Chief drag model + dragChief = MJAerodynamicDrag.AerodynamicDrag() + dragChief.ModelTag = "dragChief" + scene.AddModelToDynamicsTask(dragChief) + dragChief.applyTo(chiefBody) + + dragGeomChiefMsg = messaging.DragGeometryMsg() + dragGeomChiefMsg.write( + messaging.DragGeometryMsgPayload( + projectedArea=1, + dragCoeff=30, + r_CP_S=[0, 0, 0], + ) + ) + + stochDragGeomChief = None + stochDragGeomChiefMsg = None + + if scenarioConfig.stochastic: + # Stochastic mean reverting perturbation on Cd for the chief + stochDragGeomChief = MJMeanRevertingNoise.StochasticDragCoeff() + stochDragGeomChief.ModelTag = "stochDragGeomChief" + scene.AddModelToDynamicsTask(stochDragGeomChief) + + stochDragGeomChief.setStationaryStd(0.15) + stochDragGeomChief.setTimeConstant(1800.0) + + stochDragGeomChief.dragGeomInMsg.subscribeTo(dragGeomChiefMsg) + dragChief.dragGeometryInMsg.subscribeTo(stochDragGeomChief.dragGeomOutMsg) + + stochDragGeomChiefMsg = dragGeomChiefMsg + else: + dragChief.dragGeometryInMsg.subscribeTo(dragGeomChiefMsg) + + dragChief.atmoDensInMsg.subscribeTo(atmo.envOutMsgs[0]) + + # Deputy drag model + dragDeputy = MJAerodynamicDrag.AerodynamicDrag() + dragDeputy.ModelTag = "dragDeputy" + scene.AddModelToDynamicsTask(dragDeputy) + dragDeputy.applyTo(deputyBody) + + dragGeomDeputyMsg = messaging.DragGeometryMsg() + dragGeomDeputyMsg.write( + messaging.DragGeometryMsgPayload( + projectedArea=1, + dragCoeff=100, + r_CP_S=[0, 0, 0], + ) + ) + dragDeputy.atmoDensInMsg.subscribeTo(atmo.envOutMsgs[1]) + dragDeputy.dragGeometryInMsg.subscribeTo(dragGeomDeputyMsg) + + # Store references in models to prevent garbage collection + models.dragChief = dragChief + models.dragDeputy = dragDeputy + models.stochDragGeomChief = stochDragGeomChief + models.dragGeomChiefMsg = dragGeomChiefMsg + models.dragGeomDeputyMsg = dragGeomDeputyMsg + models.stochDragGeomChiefMsg = stochDragGeomChiefMsg + + +def configureControlModels( + models: SimulationModels, + planet: Any, + targetOeDiffPayload: messaging.ClassicElementsMsgPayload, + kProp: np.ndarray, + kInt: np.ndarray, + scenarioConfig: ScenarioConfig, +) -> None: + """Configure orbital element feedback control and force mapping.""" + scene = models.scene + chiefBody = models.chiefBody + deputyBody = models.deputyBody + + # Convert SC states to classical orbital elements + orbElemConvertChief = orbElemConvert.OrbElemConvert() + orbElemConvertChief.mu = planet.mu + orbElemConvertChief.scStateInMsg.subscribeTo(chiefBody.getCenterOfMass().stateOutMsg) + scene.AddModelToDynamicsTask(orbElemConvertChief) + + orbElemConvertDeputy = orbElemConvert.OrbElemConvert() + orbElemConvertDeputy.mu = planet.mu + orbElemConvertDeputy.scStateInMsg.subscribeTo(deputyBody.getCenterOfMass().stateOutMsg) + scene.AddModelToDynamicsTask(orbElemConvertDeputy) + + # Offset elements target (chief plus target delta) + offsetElementsMsg = messaging.ClassicElementsMsg() + offsetElementsMsg.write(targetOeDiffPayload) + + orbElemTarget = orbElemOffset.OrbElemOffset() + orbElemTarget.useMeanAnomalyOffset = True + orbElemTarget.mainElementsInMsg.subscribeTo(orbElemConvertChief.elemOutMsg) + orbElemTarget.offsetElementsInMsg.subscribeTo(offsetElementsMsg) + scene.AddModelToDynamicsTask(orbElemTarget) + + # Orbital element feedback controller + oeFeedback = MJLinearTimeInvariantSystem.OrbitalElementControl() + oeFeedback.ModelTag = "oeFeedback" + oeFeedback.mu = planet.mu + oeFeedback.setProportionalGain(np.diag(kProp)) + + if scenarioConfig.integralControl: + oeFeedback.setIntegralGain(np.diag(kInt)) + + oeFeedback.targetOEInMsg.subscribeTo(orbElemTarget.elementsOutMsg) + oeFeedback.currentOEInMsg.subscribeTo(orbElemConvertDeputy.elemOutMsg) + scene.AddModelToDynamicsTask(oeFeedback) + + # Map commanded inertial force to site-fixed force at deputy center of mass + controlActuatorDeputy: mujoco.MJForceActuator = scene.addForceActuator( + "deputyControl", deputyBody.getCenterOfMass() + ) + + inertialToSiteForce = MJCmdForceInertialToForceAtSite.CmdForceInertialToForceAtSite() + inertialToSiteForce.ModelTag = "inertialToSiteForce" + inertialToSiteForce.siteFrameStateInMsg.subscribeTo(deputyBody.getCenterOfMass().stateOutMsg) + inertialToSiteForce.cmdForceInertialInMsg.subscribeTo(oeFeedback.forceOutMsg) + controlActuatorDeputy.forceInMsg.subscribeTo(inertialToSiteForce.forceOutMsg) + scene.AddModelToDynamicsTask(inertialToSiteForce) + + # Keep references alive + models.orbElemConvertChief = orbElemConvertChief + models.orbElemConvertDeputy = orbElemConvertDeputy + models.orbElemTarget = orbElemTarget + models.offsetElementsMsg = offsetElementsMsg + models.oeFeedback = oeFeedback + models.inertialToSiteForce = inertialToSiteForce + models.controlActuatorDeputy = controlActuatorDeputy + + +def configureFeedforwardModels(models: SimulationModels, planet: Any) -> None: + """Configure drag feedforward models using the deputy drag as a disturbance estimate.""" + scene = models.scene + atmo = models.atmo + deputyBody = models.deputyBody + orbElemTarget = models.orbElemTarget + dragGeomChiefMsg = models.dragGeomChiefMsg + dragDeputy = models.dragDeputy + + if orbElemTarget is None: + raise RuntimeError("Feedforward drag requested but control is not configured.") + + if dragDeputy is None or dragGeomChiefMsg is None: + raise RuntimeError("Feedforward drag requested but deputy drag model is not configured.") + + + # Feedforward drag on desired position, computed at that location and applied to deputy + orbElemConvertTarget = orbElemConvert.OrbElemConvert() + orbElemConvertTarget.mu = planet.mu + orbElemConvertTarget.ModelTag = "orbElemConvertTarget" + scene.AddModelToDynamicsTask(orbElemConvertTarget) + + orbElemConvertTarget.elemInMsg.subscribeTo( orbElemTarget.elementsOutMsg ) + + feedforwardDragTarget = MJAerodynamicDrag.AerodynamicDrag() + feedforwardDragTarget.ModelTag = "feedforwardDragTarget" + scene.AddModelToDynamicsTask(feedforwardDragTarget) + + feedforwardDragTarget.dragGeometryInMsg.subscribeTo( dragGeomChiefMsg ) + feedforwardDragTarget.atmoDensInMsg.subscribeTo(atmo.envOutMsgs[0]) + feedforwardDragTarget.referenceFrameStateInMsg.subscribeTo(orbElemConvertTarget.scStateOutMsg) + + feedforwardDragActuatorTarget: mujoco.MJForceActuator = scene.addForceActuator( + "feedforwardDragActuatorTarget", deputyBody.getCenterOfMass() + ) + feedforwardDragActuatorTarget.forceInMsg.subscribeTo(feedforwardDragTarget.forceOutMsg) + + # Map deputy drag into an equal and opposite force on deputy + feedforwardDragDeputy = MJLinearTimeInvariantSystem.ForceAtSiteLTI() + feedforwardDragDeputy.ModelTag = "feedforwardDragDeputy" + scene.AddModelToDynamicsTask(feedforwardDragDeputy) + + feedforwardDragDeputy.setD(-np.eye(3)) + feedforwardDragDeputy.inMsg.subscribeTo(dragDeputy.forceOutMsg) + + feedforwardDragActuatorDeputy: mujoco.MJForceActuator = scene.addForceActuator( + "feedforwardDragActuatorDeputy", deputyBody.getCenterOfMass() + ) + feedforwardDragActuatorDeputy.forceInMsg.subscribeTo(feedforwardDragDeputy.outMsg) + + # Store references + models.orbElemConvertTarget = orbElemConvertTarget + models.feedforwardDragTarget = feedforwardDragTarget + models.feedforwardDragDeputy = feedforwardDragDeputy + models.feedforwardDragActuatorTarget = feedforwardDragActuatorTarget + models.feedforwardDragActuatorDeputy = feedforwardDragActuatorDeputy + + +def addRecorders(models: SimulationModels) -> RecorderHandles: + """Attach all recorders used for postprocessing and keep them in a container.""" + scSim = models.scSim + chiefBody = models.chiefBody + deputyBody = models.deputyBody + + recordDt = macros.min2nano(1.0) + + # Inertial states of chief and deputy + chiefStateRecorder = chiefBody.getCenterOfMass().stateOutMsg.recorder(recordDt) + deputyStateRecorder = deputyBody.getCenterOfMass().stateOutMsg.recorder(recordDt) + + scSim.AddModelToTask("test", chiefStateRecorder, 0) + scSim.AddModelToTask("test", deputyStateRecorder, 0) + + dragChiefRecorder = None + densityChiefRecorder = None + cdChiefRecorder = None + oeFeedbackForceRecorder = None + feedforwardDragTargetRecorder = None + feedforwardDragDeputyRecorder = None + + # Chief drag and density recorders + if models.dragChief is not None: + dragChiefRecorder = models.dragChief.forceOutMsg.recorder(recordDt) + densityChiefRecorder = models.dragChief.atmoDensInMsg.recorder(recordDt) + scSim.AddModelToTask("test", dragChiefRecorder, 0) + scSim.AddModelToTask("test", densityChiefRecorder, 0) + + # Stochastic Cd recorder + if models.stochDragGeomChief is not None: + cdChiefRecorder = models.stochDragGeomChief.dragGeomOutMsg.recorder(recordDt) + scSim.AddModelToTask("test", cdChiefRecorder, 0) + + # Orbital element feedback force + if models.inertialToSiteForce is not None: + oeFeedbackForceRecorder = models.inertialToSiteForce.forceOutMsg.recorder(recordDt) + scSim.AddModelToTask("test", oeFeedbackForceRecorder, 0) + + # Feedforward drag recorders + if models.feedforwardDragTarget is not None: + feedforwardDragTargetRecorder = models.feedforwardDragTarget.forceOutMsg.recorder(recordDt) + scSim.AddModelToTask("test", feedforwardDragTargetRecorder, 0) + + if models.feedforwardDragDeputy is not None: + feedforwardDragDeputyRecorder = models.feedforwardDragDeputy.outMsg.recorder(recordDt) + scSim.AddModelToTask("test", feedforwardDragDeputyRecorder, 0) + + return RecorderHandles( + chiefState=chiefStateRecorder, + deputyState=deputyStateRecorder, + dragChief=dragChiefRecorder, + densityChief=densityChiefRecorder, + cdChief=cdChiefRecorder, + oeFeedbackForce=oeFeedbackForceRecorder, + feedforwardDragTarget=feedforwardDragTargetRecorder, + feedforwardDragDeputy=feedforwardDragDeputyRecorder, + ) + + +def setInitialConditions(models: SimulationModels, planet: Any) -> InitialOrbitInfo: + """Set initial orbital elements for chief and deputy and propagate to Cartesian states.""" + models.scSim.InitializeSimulation() + + chiefBody = models.chiefBody + deputyBody = models.deputyBody + + # Chief orbit between 250 and 300 km altitude + initialMinAlt = 250.0 # km + initialMaxAlt = 300.0 # km + initialRadiusPeriapsis = planet.radEquator + initialMinAlt * 1000.0 + initialRadiusApoapsis = planet.radEquator + initialMaxAlt * 1000.0 + initialSemiMajorAxis = (initialRadiusPeriapsis + initialRadiusApoapsis) / 2.0 + + initialOeChief = orbitalMotion.ClassicElements() + initialOeChief.a = initialSemiMajorAxis + initialOeChief.e = 1.0 - initialRadiusPeriapsis / initialSemiMajorAxis + initialOeChief.i = 33.3 * macros.D2R + initialOeChief.Omega = 48.2 * macros.D2R + initialOeChief.omega = 347.8 * macros.D2R + initialOeChief.f = 85.3 * macros.D2R + + rNChief, vNChief = orbitalMotion.elem2rv(planet.mu, initialOeChief) + chiefBody.setPosition(rNChief) + chiefBody.setVelocity(vNChief) + + # Apply small differential elements to deputy + initialDeltaOe = (0.0001, 0.0002, -0.0003, 0.0005, 0.0005, 0.0006) + initialOeDeputy = orbitalMotion.ClassicElements() + initialOeDeputy.a = initialOeChief.a * (1.0 + initialDeltaOe[0]) + initialOeDeputy.e = initialOeChief.e + initialDeltaOe[1] + initialOeDeputy.i = initialOeChief.i + initialDeltaOe[2] + initialOeDeputy.Omega = initialOeChief.Omega + initialDeltaOe[3] + initialOeDeputy.omega = initialOeChief.omega + initialDeltaOe[4] + + initialMeanAnomalyChief = orbitalMotion.E2M( + orbitalMotion.f2E(initialOeChief.f, initialOeChief.e), + initialOeChief.e, + ) + initialMeanAnomalyDeputy = initialMeanAnomalyChief + initialDeltaOe[5] + initialOeDeputy.f = orbitalMotion.E2f( + orbitalMotion.M2E(initialMeanAnomalyDeputy, initialOeDeputy.e), + initialOeDeputy.e, + ) + + rNDeputy, vNDeputy = orbitalMotion.elem2rv(planet.mu, initialOeDeputy) + deputyBody.setPosition(rNDeputy) + deputyBody.setVelocity(vNDeputy) + + # Compute orbit period for the initial semi major axis + initialOrbitPeriod = 2.0 * np.pi / np.sqrt(planet.mu / initialSemiMajorAxis ** 3) + + return InitialOrbitInfo(semiMajorAxis=initialSemiMajorAxis, orbitPeriod=initialOrbitPeriod) + + +def runSimulation(models: SimulationModels, orbitInfo: InitialOrbitInfo) -> None: + """Initialize and execute the simulation for a fixed multiple of the orbit period.""" + scSim = models.scSim + + finalTime = 30.0 * orbitInfo.orbitPeriod + scSim.ConfigureStopTime(macros.sec2nano(finalTime)) + scSim.ExecuteSimulation() + + +def extractTimeSeries( + planet: Any, + recorders: RecorderHandles, +) -> SimulationTimeSeries: + """Convert raw recorder data into orbital elements, relative coordinates, and time arrays.""" + chiefState = recorders.chiefState + deputyState = recorders.deputyState + + times = chiefState.times() * macros.NANO2HOUR + posChief = chiefState.r_BN_N + velChief = chiefState.v_BN_N + posDeputy = deputyState.r_BN_N + velDeputy = deputyState.v_BN_N + + oeChief = rv2elemVector(planet.mu, posChief, velChief) + oeDeputy = rv2elemVector(planet.mu, posDeputy, velDeputy) + + # Difference in semimajor axis and eccentricity is simple subtraction + oeDiff = np.copy(oeDeputy) + oeDiff[:, :2] -= oeChief[:, :2] + + # Angular elements must be wrapped to avoid jumps near 360 deg + oeDiff[:, 2:] = ((oeDeputy[:, 2:] - oeChief[:, 2:] + 180.0) % 360.0) - 180.0 + + hillPosition, hillVelocity = deputyToHillFrame(posChief, velChief, posDeputy, velDeputy) + + return SimulationTimeSeries( + times=times, + posChief=posChief, + velChief=velChief, + posDeputy=posDeputy, + velDeputy=velDeputy, + oeChief=oeChief, + oeDeputy=oeDeputy, + oeDiff=oeDiff, + hillPosition=hillPosition, + hillVelocity=hillVelocity, + ) + +COLORS = namedtuple('Hcset', 'blue yellow red black')('#004488', '#DDAA33', '#BB5566', '#000000') + +def createFigures( + scenarioConfig: ScenarioConfig, + plotConfig: PlotConfig, + models: SimulationModels, + recorders: RecorderHandles, + timeSeries: SimulationTimeSeries, + targetOeDiffPayload: messaging.ClassicElementsMsgPayload, + kProp: np.ndarray, +) -> Dict[str, plt.Figure]: + """Generate all requested figures, optionally saving them as PDF files.""" + figures: Dict[str, plt.Figure] = {} + + times = timeSeries.times + posChief = timeSeries.posChief + oeChief = timeSeries.oeChief + oeDeputy = timeSeries.oeDeputy + oeDiff = timeSeries.oeDiff + hillPosition = timeSeries.hillPosition + + # Ensure save folder exists if requested + if plotConfig.saveFolder is not None: + os.makedirs(plotConfig.saveFolder, exist_ok=True) + + def maybeSave(figName: str, fig: plt.Figure) -> None: + if plotConfig.saveFolder is not None: + filePath = os.path.join(plotConfig.saveFolder, f"{figName}.pdf") + fig.savefig(filePath, bbox_inches="tight") + + # Chief inertial trajectory + if plotConfig.plotChiefInertial: + fig, ax = plt.subplots(figsize=[5, 5]) + plotGradient(ax, posChief[:, 0] / 1000.0, posChief[:, 1] / 1000.0) + ax.set_xlabel(r"${}^\mathcal{N} {x}$ [km]") + ax.set_ylabel(r"${}^\mathcal{N} {y}$ [km]") + ax.axis("equal") + figures["chiefInertial"] = fig + maybeSave("chiefInertial", fig) + + # Chief drag + if plotConfig.plotChiefDrag and recorders.dragChief is not None: + fig, ax = plt.subplots(figsize=[8,4]) + chiefDragMagnitude = np.linalg.norm(recorders.dragChief.force_S, axis=1) + ax.plot(times, chiefDragMagnitude, color=COLORS.blue) + ax.set_ylabel("Chief Drag [m/s$^2$]") + ax.set_xlabel("Time [hr]") + figures["chiefDrag"] = fig + maybeSave("chiefDrag", fig) + + # Chief density + if plotConfig.plotChiefDensity and recorders.densityChief is not None: + fig, ax = plt.subplots(figsize=[8,4]) + ax.plot(times, recorders.densityChief.neutralDensity, color=COLORS.blue) + ax.set_ylabel("Chief Density [kg/m$^3$]") + ax.set_xlabel("Time [hr]") + figures["chiefDensity"] = fig + maybeSave("chiefDensity", fig) + + # Chief Cd + if plotConfig.plotChiefCd and recorders.cdChief is not None: + fig, ax = plt.subplots(figsize=[8,4]) + ax.plot(times, recorders.cdChief.dragCoeff, color=COLORS.blue) + ax.set_ylabel("$C_D$ [-]") + ax.set_xlabel("Time [hr]") + figures["chiefCd"] = fig + maybeSave("chiefCd", fig) + + # Control forces + if plotConfig.plotControl and scenarioConfig.control: + fig, ax = plt.subplots(figsize=[8,4]) + + totalControl = np.zeros([len(times), 3]) + plotTotalControl = True + + if recorders.feedforwardDragTarget is not None and recorders.feedforwardDragDeputy is not None: + feedforwardForceVec = recorders.feedforwardDragTarget.force_S + recorders.feedforwardDragDeputy.force_S + feedforwardForce = np.linalg.norm(feedforwardForceVec, axis=1) + ax.plot(times, feedforwardForce, label="Feedforward", color=COLORS.blue) + + totalControl += feedforwardForceVec + else: + plotTotalControl = False + + if recorders.oeFeedbackForce is not None: + feedforwardForceVec = recorders.oeFeedbackForce.force_S + feedbackForce = np.linalg.norm(feedforwardForceVec, axis=1) + ax.plot(times, feedbackForce, label="Feedback", color=COLORS.red) + + totalControl += feedforwardForceVec + else: + plotTotalControl = False + + if plotTotalControl: + ax.plot(times, np.linalg.norm(totalControl, axis=1), label="Total", ls="--", color=COLORS.yellow) + + ax.legend(loc="upper right") + ax.set_xlabel("Time [hr]") + ax.set_ylabel("Control [m/s$^2$]") # Deputy mass is 1, so force equals acceleration + figures["control"] = fig + maybeSave("control", fig) + + # Orbital element plots + labels = ("a [m]", "e [-]", "i [deg]", r"$\Omega$ [deg]", r"$\omega$ [deg]", r"$M$ [deg]") + targetDiffVector = [ + targetOeDiffPayload.a, + targetOeDiffPayload.e, + targetOeDiffPayload.i, + targetOeDiffPayload.Omega, + targetOeDiffPayload.omega, + targetOeDiffPayload.f, + ] + indicesToPlot = list(range(6)) + + def plotOeSet(name: str, oe: np.ndarray, isDiff: bool, enabled: bool) -> None: + if not enabled: + return + + n = len(indicesToPlot) + layout = {1: (1, 1), 2: (1, 2), 3: (1, 3), 4: (2, 2), 5: (2, 3), 6: (2, 3)} + nrows, ncols = layout[n] + fig, axs = plt.subplots(nrows=nrows, ncols=ncols, sharex=True, figsize=[10, 6]) + axsArray = np.atleast_1d(axs).flatten() + + for i, idx in enumerate(indicesToPlot): + ax = axsArray[i] + ax.plot(times, oe[:, idx], color=COLORS.blue) + ylabel = (r"$\Delta$" if isDiff else "") + labels[idx] + ax.set_ylabel(ylabel) + ax.ticklabel_format(useOffset=False, axis="y") + + # Draw controller target for states that are regulated + if isDiff and models.oeFeedback is not None and kProp[idx] != 0: + targetValue = targetDiffVector[idx] + if idx > 1: + targetValue = np.rad2deg(targetValue) + ax.axhline(targetValue, ls="--", color=COLORS.red) + + # Label shared x axis on bottom row + for ax in axsArray[-ncols:]: + ax.set_xlabel("Time [hr]") + + figures[f"oe{name.capitalize()}"] = fig + maybeSave(f"oe{name.capitalize()}", fig) + + plotOeSet("chief", oeChief, isDiff=False, enabled=plotConfig.plotOeChief) + plotOeSet("deputy", oeDeputy, isDiff=False, enabled=plotConfig.plotOeDeputy) + plotOeSet("diff", oeDiff, isDiff=True, enabled=plotConfig.plotOeDiff) + + # Hill frame deputy relative motion + if plotConfig.plotHillXY: + fig, ax = plt.subplots(figsize=[4, 4.8]) + plotGradient(ax, hillPosition[:, 0] / 1000.0, hillPosition[:, 1] / 1000.0) + ax.set_xlabel(r"${}^\mathcal{H} {x}$ [km]") + ax.set_ylabel(r"${}^\mathcal{H} {y}$ [km]") + ax.axis("equal") + figures["hillXY"] = fig + maybeSave("hillXY", fig) + + return figures + + +def rv2elemVector(mu: float, rBNN: np.ndarray, vBNN: np.ndarray) -> np.ndarray: + """Convert arrays of Cartesian states to classical orbital elements with mean anomaly.""" + oe = np.empty((rBNN.shape[0], 6)) + for i in range(rBNN.shape[0]): + oeI = orbitalMotion.rv2elem(mu, rBNN[i, :], vBNN[i, :]) + meanAnomaly = orbitalMotion.E2M(orbitalMotion.f2E(oeI.f, oeI.e), oeI.e) + oe[i, :] = ( + oeI.a, + oeI.e, + np.rad2deg(oeI.i), + np.rad2deg(oeI.Omega), + np.rad2deg(oeI.omega), + np.rad2deg(meanAnomaly), + ) + return oe + + +def deputyToHillFrame( + nRC: np.ndarray, + nVC: np.ndarray, + nRD: np.ndarray, + nVD: np.ndarray, +) -> (np.ndarray, np.ndarray): + """ + Transform deputy inertial state to Hill frame relative position and velocity. + + Parameters + ---------- + nRC + Chief position in inertial frame, shape (N, 3). + nVC + Chief velocity in inertial frame, shape (N, 3). + nRD + Deputy position in inertial frame, shape (N, 3). + nVD + Deputy velocity in inertial frame, shape (N, 3). + + Returns + ------- + hillRho + Deputy relative position in Hill frame, shape (N, 3). + hillRhoPrime + Deputy relative velocity in Hill frame, shape (N, 3). + """ + def rowNorm(x: np.ndarray) -> np.ndarray: + return np.linalg.norm(x, axis=1, keepdims=True) + + rHat = nRC / rowNorm(nRC) + hVec = np.cross(nRC, nVC, axis=1) + hHat = hVec / rowNorm(hVec) + yHat = np.cross(hHat, rHat, axis=1) + + hillDcmN = np.stack([rHat, yHat, hHat], axis=1) + + hMagnitude = rowNorm(np.cross(nRC, nVC, axis=1))[:, 0] + radiusMagnitude = rowNorm(nRC)[:, 0] + trueAnomalyDot = hMagnitude / (radiusMagnitude ** 2) + omega = np.stack( + [np.zeros_like(trueAnomalyDot), np.zeros_like(trueAnomalyDot), trueAnomalyDot], + axis=1, + ) + + hillRho = np.einsum("nij,nj->ni", hillDcmN, (nRD - nRC)) + hillRhoPrime = np.einsum("nij,nj->ni", hillDcmN, (nVD - nVC)) - np.cross(omega, hillRho, axis=1) + + return hillRho, hillRhoPrime + + +def plotGradient( + ax: plt.Axes, + x: np.ndarray, + y: np.ndarray, + *, + cmap: str = "viridis", + linewidth: float = 2.0, + arrows: bool = True, + arrowEvery: int = 40, + showEndpoints: bool = False, + **kwargs: Any, +) -> LineCollection: + """ + Plot a line whose color varies along its length, optionally with arrows and endpoints. + + Parameters + ---------- + ax + Matplotlib axes on which to draw. + x, y + Coordinate arrays. + cmap + Name of the colormap to use. + linewidth + Width of the plotted line. + arrows + If True, draw arrows along the trajectory. + arrowEvery + Spacing between arrows in number of points. + showEndpoints + If True, mark start and end points. + + Returns + ------- + LineCollection + The line collection added to the axes. + """ + x = np.asarray(x) + y = np.asarray(y) + + points = np.column_stack([x, y]).reshape(-1, 1, 2) + segments = np.concatenate([points[:-1], points[1:]], axis=1) + + lineCollection = LineCollection(segments, cmap=cmap, linewidth=linewidth, **kwargs) + lineCollection.set_array(np.linspace(0.0, 1.0, len(x))) + ax.add_collection(lineCollection) + + if arrows: + for i in range(0, len(x) - arrowEvery, arrowEvery): + ax.annotate( + "", + xy=(x[i + 1], y[i + 1]), + xytext=(x[i], y[i]), + arrowprops=dict( + arrowstyle="->", + color="k", + lw=1, + shrinkA=0, + shrinkB=0, + ), + ) + + if showEndpoints: + ax.scatter(x[0], y[0], color="red", zorder=3) + ax.scatter(x[-1], y[-1], color="blue", zorder=3) + + ax.autoscale() + ax.set_aspect("equal") + return lineCollection + +def runMultipleCases(): + cases = { + "noDragSimpleControl": ScenarioConfig( + drag=False, + control=True, + feedforward=False, + integralControl=False, + stochastic=False, + ), + "simpleControl": ScenarioConfig( + drag=True, + control=True, + feedforward=False, + integralControl=False, + stochastic=False, + ), + "withFeedforward": ScenarioConfig( + drag=True, + control=True, + feedforward=True, + integralControl=False, + stochastic=False, + ), + "withIntegral": ScenarioConfig( + drag=True, + control=True, + feedforward=False, + integralControl=True, + stochastic=False, + ), + "withFeedforwardAndIntegral": ScenarioConfig( + drag=True, + control=True, + feedforward=True, + integralControl=True, + stochastic=False, + ), + "stochasticWithFeedforwardAndIntegral": ScenarioConfig( + drag=True, + control=True, + feedforward=True, + integralControl=True, + stochastic=True, + ), + } + for folder, config in cases.items(): + run(config, PlotConfig(saveFolder=f"scenarioFormationFlyingStochasticDrag/{folder}")) + print("Finished case:", folder) + plt.close("all") + +if __name__ == "__main__": + runMultipleCases() + scenarioConfig = ScenarioConfig( + drag=True, + control=True, + feedforward=True, + integralControl=False, + stochastic=False, + ) + plotConfig = PlotConfig( + showPlots=True, + saveFolder=None, + ) + # run(scenarioConfig, plotConfig) diff --git a/examples/mujoco/scenarioStochasticDrag.py b/examples/mujoco/scenarioStochasticDrag.py new file mode 100644 index 0000000000..d243da621b --- /dev/null +++ b/examples/mujoco/scenarioStochasticDrag.py @@ -0,0 +1,485 @@ +# +# 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, + densityCorrectionTimeConstant=1.8 * 1 * 60 # s + ) + 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. Alternatively, one can express the process as: + + d(densityCorrection) = -densityCorrection/tau*dt + sqrt(2/tau)*sigma_st*dW + + where ``tau`` represents the time constant of the process, and ``sigma_st`` + is its stationary standard deviation. These parameters are set in the constructor + of this class. + """ + + def __init__(self, densityCorrectionStandardDeviation: float, densityCorrectionTimeConstant: float): + """ + Initialize the StochasticAtmosphere model. + + Args: + densityCorrectionStandardDeviation (float): Stationary standard deviation + of the Ornstein-Uhlenbeck process that drives the relative atmospheric density + correction. + densityCorrectionTimeConstant (float): time constant of the + Ornstein-Uhlenbeck process that drives the relative atmospheric density correction. + """ + super().__init__() + self.densityCorrectionStandardDeviation = densityCorrectionStandardDeviation + self.densityCorrectionTimeConstant = densityCorrectionTimeConstant + + self.atmoDensInMsg = messaging.AtmoPropsMsgReader() + self.atmoDensOutMsg = messaging.AtmoPropsMsg() + + @property + def densityCorrectionTheta(self): + """ + Compute the ``theta`` term of the Ornstein-Uhlenbeck process that + drives the atmospheric density correction. + + Returns: + float: The theta value for the OU process. + """ + return 1/self.densityCorrectionTimeConstant + + @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.densityCorrectionTimeConstant) + + 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/architecture/_GeneralModuleFiles/py_sys_model.i b/src/architecture/_GeneralModuleFiles/py_sys_model.i index ae43f1229f..854c1b094e 100644 --- a/src/architecture/_GeneralModuleFiles/py_sys_model.i +++ b/src/architecture/_GeneralModuleFiles/py_sys_model.i @@ -1,3 +1,20 @@ +/* + 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(directors="1",threads="1",package="Basilisk.architecture") sysModel @@ -13,49 +30,158 @@ import sys import traceback from Basilisk.architecture.swig_common_model import * %} + %include "std_string.i" %include "swig_conly_data.i" %include "architecture/utilities/bskLogging.h" +/* Director support for the C++ SysModel base */ %feature("director") SysModel; -%feature("pythonappend") SysModel::SysModel %{ - self.__super_init_called__ = True%} + +/* Expose the raw SWIG proxy as _SysModel */ %rename("_SysModel") SysModel; %include "sys_model.i" %pythoncode %{ -class SuperInitChecker(type): - - def __call__(cls, *a, **kw): - rv = super(SuperInitChecker, cls).__call__(*a, **kw) - if not getattr(rv, "__super_init_called__", False): - error_msg = ( - "Need to call parent __init__ in SysModel subclasses:\n" - f"class {cls.__name__}(sysModel.SysModel):\n" - " def __init__(...):\n" - " super().__init__()" - ) - raise SyntaxError(error_msg) - return rv - def logError(func): - """Decorator that prints any exceptions that happen when - the original function is called, and then raises them again.""" - def inner(*arg, **kwargs): + """Decorator that prints any exceptions raised by func, then re raises them.""" + def inner(*args, **kwargs): try: - return func(*arg, **kwargs) + return func(*args, **kwargs) except Exception: traceback.print_exc() raise + # Mark so we do not wrap twice + inner._sysmodelLogged = True return inner -class SysModel(_SysModel, metaclass=SuperInitChecker): - bskLogger: BSKLogger = None + +class SysModelMixin: + """Pure Python mixin for Basilisk system models. + + Provides: + * Automatic logError wrapping for overridden director methods + from any SWIG C++ proxy bases in the MRO. + * Enforcement that subclasses call super().__init__(). + """ + + def __init__(self, *args, **kwargs): + # Cooperative init; super() will eventually call the SWIG proxy __init__ + super().__init__(*args, **kwargs) + # Mark that the mixin init has been called on this instance. + self._sysmodelBaseInitSeen = True + + @classmethod + def _wrapOverriddenDirectorMethods(cls): + """Wrap all overridden director methods with logError. + + A method is considered a director method if: + * it exists on at least one SWIG C++ proxy base + * and it is overridden in cls with a different callable + """ + cppBases = [] + for base in cls.mro()[1:]: + if base is object: + continue + # Heuristic: SWIG proxy types have __swig_destroy__ or similar + if hasattr(base, "__swig_destroy__"): + cppBases.append(base) + + if not cppBases: + return + + # Collect candidate method names from all C++ bases + candidateNames = set() + for base in cppBases: + for name in dir(base): + if name.startswith("_"): + continue + if name in ("__init__",): + continue + candidateNames.add(name) + + for name in candidateNames: + subAttr = getattr(cls, name, None) + if not callable(subAttr): + continue + + # Find the first base that defines this name + baseAttr = None + for base in cppBases: + if hasattr(base, name): + baseAttr = getattr(base, name) + break + + if baseAttr is None: + continue + + # Only wrap if the subclass actually overrides the method + if subAttr is baseAttr: + continue + + # Avoid double wrapping + if getattr(subAttr, "_sysmodelLogged", False): + continue + + setattr(cls, name, logError(subAttr)) def __init_subclass__(cls): - # Make it so any exceptions in UpdateState and Reset - # print any exceptions before returning control to - # C++ (at which point exceptions will crash the program) - cls.UpdateState = logError(cls.UpdateState) - cls.Reset = logError(cls.Reset) + """Hook that runs for every Python subclass of SysModelMixin. + + It: + * wraps all overridden director methods from SWIG C++ proxy bases + in the MRO with logError + * enforces super().__init__() on subclasses + """ + super().__init_subclass__() + + # Wrap overridden director methods + cls._wrapOverriddenDirectorMethods() + + origInit = cls.__init__ + + # Avoid double wrapping if someone subclasses a subclass. + if getattr(origInit, "_sysmodelWrappedInit", False): + return + + def wrappedInit(self, *args, **kwargs): + # Clear flag before running user __init__. + object.__setattr__(self, "_sysmodelBaseInitSeen", False) + origInit(self, *args, **kwargs) + if not getattr(self, "_sysmodelBaseInitSeen", False): + errorMsg = ( + "Need to call parent __init__ in SysModel based subclasses:\n" + f"class {cls.__name__}(...):\n" + " def __init__(...):\n" + " super().__init__()" + ) + raise SyntaxError(errorMsg) + + wrappedInit._sysmodelWrappedInit = True + cls.__init__ = wrappedInit + + +class SysModel(SysModelMixin, _SysModel): + """Python friendly base class for Basilisk system models. + + Inherits: + * SysModelMixin (pure Python) + * _SysModel (SWIG proxy for C++ SysModel) + + Typical usage for new models: + from Basilisk.architecture import sysModel + + class MyModel(sysModel.SysModel): + def __init__(self): + super().__init__() + ... + + def UpdateState(self, CurrentSimNanos): + ... + + def Reset(self, CurrentSimNanos): + ... + """ + + bskLogger: BSKLogger = None %} diff --git a/src/architecture/_GeneralModuleFiles/swig_eigen.i b/src/architecture/_GeneralModuleFiles/swig_eigen.i index c5eff26d7d..a1d730b1f6 100644 --- a/src/architecture/_GeneralModuleFiles/swig_eigen.i +++ b/src/architecture/_GeneralModuleFiles/swig_eigen.i @@ -413,44 +413,51 @@ void fillPyObjList(PyObject *input, const Eigen::Quaterniond %define EIGEN_MAT_WRAP(type, typeCheckPrecedence) +/* Python -> C++: by value */ %typemap(in, fragment="pyObjToEigenMatrix") type { $1 = pyObjToEigenMatrix($input); if (PyErr_Occurred()) SWIG_fail; } +/* For member assignment (e.g. setting a struct field) */ %typemap(memberin) type { $1 = std::move($input); } +/* Python -> C++: non-const reference (copy-in, no write-back) */ %typemap(in, fragment="pyObjToEigenMatrix") type & { $1 = new type; *$1 = pyObjToEigenMatrix($input); if (PyErr_Occurred()) SWIG_fail; } +/* Free temporary created for non-const reference in typemap(in) */ %typemap(freearg) type & { delete $1; } -%typemap(typecheck, fragment="checkPyObjectIsMatrixLike", precedence= ## typeCheckPrecedence) type { - // PyErr_Fetch and PyErr_Restore preserve/restore the error status before this function - // We use the error flag to check whether conversion is valid, but we do not want to - // alter the previous error status of the program +/* Overload resolution; check whether input can be converted to type */ +%typemap(typecheck, + fragment="checkPyObjectIsMatrixLike", + precedence= ## typeCheckPrecedence) type { + // Preserve prior error state PyObject *ty, *value, *traceback; PyErr_Fetch(&ty, &value, &traceback); pyObjToEigenMatrix($input); - $1 = ! PyErr_Occurred(); + $1 = !PyErr_Occurred(); PyErr_Restore(ty, value, traceback); } +/* C++ -> Python: by value */ %typemap(out, optimal="1", fragment="fillPyObjList") type { $result = PyList_New(0); fillPyObjList($result, $1); if (PyErr_Occurred()) SWIG_fail; } +/* C++ -> Python: pointer return */ %typemap(out, fragment="fillPyObjList") type * { if(!($1)) { @@ -464,32 +471,102 @@ void fillPyObjList(PyObject *input, const Eigen::Quaterniond if (PyErr_Occurred()) SWIG_fail; } -%typemap(typecheck) type & = type; -%typemap(typecheck) type && = type; +/* Director: Python override return -> C++ Eigen (by value) */ +%typemap(directorout, fragment="pyObjToEigenMatrix") type { + $result = pyObjToEigenMatrix($input); + if (PyErr_Occurred()) { + Swig::DirectorTypeMismatchException::raise( + PyExc_TypeError, + "Could not convert Python return value to Eigen matrix " + "for director method" + ); + } +} + +/* Director: Python override return -> C++ Eigen pointer (allocate new) */ +%typemap(directorout, fragment="pyObjToEigenMatrix") type * { + type *tmp = new type; + *tmp = pyObjToEigenMatrix($input); + if (PyErr_Occurred()) { + delete tmp; + Swig::DirectorTypeMismatchException::raise( + PyExc_TypeError, + "Could not convert Python return value to Eigen matrix pointer " + "for director method" + ); + } + $result = tmp; +} + +/* Director: C++ Eigen argument -> Python list (by value) */ +%typemap(directorin, fragment="fillPyObjList") type { + $input = PyList_New(0); + fillPyObjList($input, $1); + if (PyErr_Occurred()) { + Swig::DirectorTypeMismatchException::raise( + PyExc_TypeError, + "Could not convert Eigen matrix argument for director method" + ); + } +} -%typemap(in) type && = type &; -%typemap(freearg) type && = type &; +/* Director: C++ Eigen pointer argument -> Python object */ +%typemap(directorin, fragment="fillPyObjList") type * { + if (!$1) { + Py_INCREF(Py_None); + $input = Py_None; + } else { + $input = PyList_New(0); + fillPyObjList($input, *$1); + if (PyErr_Occurred()) { + Swig::DirectorTypeMismatchException::raise( + PyExc_TypeError, + "Could not convert Eigen matrix pointer argument for director method" + ); + } + } +} + +/* Aliases for references in type system and director typemaps */ +%typemap(typecheck) type & = type; +%typemap(typecheck) type && = type; + +%typemap(in) type && = type &; +%typemap(freearg) type && = type &; + +%typemap(directorin) type & = type; +%typemap(directorout) type & = type; %enddef +/* ============================================================ + Rotation wrappers (MRPd, Quaterniond) + ============================================================ */ + %define EIGEN_ROT_WRAP(type, typeCheckPrecedence) +/* Python -> C++: by value */ %typemap(in, fragment="pyObjToRotation") type { $1 = pyObjToRotation($input); if (PyErr_Occurred()) SWIG_fail; } +/* Python -> C++: non-const reference (copy-in, no write-back) */ %typemap(in, fragment="pyObjToRotation") type & { $1 = new type; *$1 = pyObjToRotation($input); if (PyErr_Occurred()) SWIG_fail; } +/* Free temporary created for non-const reference in typemap(in) */ %typemap(freearg) type & { delete $1; } -%typemap(typecheck, fragment="getInputSize", precedence= ## typeCheckPrecedence) type { +/* Overload resolution for rotation-like inputs */ +%typemap(typecheck, + fragment="getInputSize", + precedence= ## typeCheckPrecedence) type { auto maybeSize = getInputSize($input); if (!maybeSize) @@ -507,14 +584,16 @@ void fillPyObjList(PyObject *input, const Eigen::Quaterniond } } +/* C++ -> Python: by value */ %typemap(out, optimal="1", fragment="fillPyObjList") type { $result = PyList_New(0); fillPyObjList($result, $1); if (PyErr_Occurred()) SWIG_fail; } +/* C++ -> Python: pointer return */ %typemap(out, fragment="fillPyObjList") type * { - if(!($1)) + if (!($1)) { $result = SWIG_Py_Void(); } @@ -526,26 +605,90 @@ void fillPyObjList(PyObject *input, const Eigen::Quaterniond if (PyErr_Occurred()) SWIG_fail; } -%typemap(typecheck) type & = type; -%typemap(typecheck) type && = type; +/* Director: Python override return -> C++ rotation (by value) */ +%typemap(directorout, fragment="pyObjToRotation") type { + $result = pyObjToRotation($input); + if (PyErr_Occurred()) { + Swig::DirectorTypeMismatchException::raise( + PyExc_TypeError, + "Could not convert Python return value to Eigen rotation " + "for director method" + ); + } +} + +/* Director: Python override return -> C++ rotation pointer (allocate new) */ +%typemap(directorout, fragment="pyObjToRotation") type * { + type *tmp = new type; + *tmp = pyObjToRotation($input); + if (PyErr_Occurred()) { + delete tmp; + Swig::DirectorTypeMismatchException::raise( + PyExc_TypeError, + "Could not convert Python return value to Eigen rotation pointer " + "for director method" + ); + } + $result = tmp; +} + +/* Director: C++ rotation argument -> Python list (by value) */ +%typemap(directorin, fragment="fillPyObjList") type { + $input = PyList_New(0); + fillPyObjList($input, $1); + if (PyErr_Occurred()) { + Swig::DirectorTypeMismatchException::raise( + PyExc_TypeError, + "Could not convert Eigen rotation argument for director method" + ); + } +} -%typemap(in) type && = type &; -%typemap(freearg) type && = type &; +/* Director: C++ rotation pointer argument -> Python object */ +%typemap(directorin, fragment="fillPyObjList") type * { + if (!$1) { + Py_INCREF(Py_None); + $input = Py_None; + } else { + $input = PyList_New(0); + fillPyObjList($input, *$1); + if (PyErr_Occurred()) { + Swig::DirectorTypeMismatchException::raise( + PyExc_TypeError, + "Could not convert Eigen rotation pointer argument for director method" + ); + } + } +} + +/* Aliases for references */ +%typemap(typecheck) type & = type; +%typemap(typecheck) type && = type; + +%typemap(in) type && = type &; +%typemap(freearg) type && = type &; + +%typemap(directorin) type & = type; +%typemap(directorout) type & = type; %enddef +/* ============================================================ + Concrete bindings + ============================================================ */ + // Second value is the precedence for overloads. If two methods are overloaded // with different Eigen inputs, then the precedence determines which ones are // attempted first. // -// Fully fixed-size types should have precendence 155 < precedence < 160 +// Fully fixed-size types should have precedence 155 < precedence < 160 // Types partially fixed-size should have precedence 160 < precedence < 165 -// Fully dynamically sized should have precedence 165 < precendence < 170 +// Fully dynamically sized should have precedence 165 < precedence < 170 // -// Boolean matrices should have lower precendence than integer matrices, -// which should have lower precendence than double matrices. +// Boolean matrices should have lower precedence than integer matrices, +// which should have lower precedence than double matrices. // -// For reference, std::array has precendence 155, and std::vector has precendence 160 +// For reference, std::array has precedence 155, and std::vector has precedence 160. // // If we were to pass a list of three integers from python to an overloaded method 'foo', // this would be the order in which they would be chosen (from first to last): @@ -568,12 +711,11 @@ EIGEN_MAT_WRAP(Eigen::MatrixX3d,164) EIGEN_MAT_WRAP(Eigen::MatrixXi, 169) EIGEN_MAT_WRAP(Eigen::MatrixXd, 169) -// Rotation wrappers work so that they can be set from Python on any -// of the three ways we have of representing rotations: MRPs (3x1 vector), -// quaternions (4x1 vector), or rotation matrix (3x3 vector). -// -// Their return type, however, is always an MRPd. +// Rotation wrappers work so that they can be set from Python as any of: +// * MRPs (3x1 vector), +// * quaternions (4x1 vector), +// * rotation matrix (3x3). // -// The precedence should be that of fixed-size double eigen matrices (159). +// Here we bind MRPd and Quaterniond as concrete rotation types. EIGEN_ROT_WRAP(Eigen::MRPd , 159) EIGEN_ROT_WRAP(Eigen::Quaterniond, 159) diff --git a/src/architecture/msgPayloadDefC/DragGeometryMsgPayload.h b/src/architecture/msgPayloadDefC/DragGeometryMsgPayload.h new file mode 100644 index 0000000000..737d3976f0 --- /dev/null +++ b/src/architecture/msgPayloadDefC/DragGeometryMsgPayload.h @@ -0,0 +1,31 @@ +/* + 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 DRAG_GEOMETRY_MSG_PAYLPAD_H +#define DRAG_GEOMETRY_MSG_PAYLPAD_H + + +//! @brief Container for basic drag parameters. +typedef struct { + double projectedArea; //!< [m^2] Area of spacecraft projected in velocity direction + double dragCoeff; //!< [-] Nondimensional drag coefficient + double r_CP_S[3]; //!< [m] Position of center of pressure relative to the center of frame S, given in frame S [m] +}DragGeometryMsgPayload; + +#endif diff --git a/src/fswAlgorithms/orbitControl/orbElemOffset/_UnitTest/test_orbElemOffset.py b/src/fswAlgorithms/orbitControl/orbElemOffset/_UnitTest/test_orbElemOffset.py new file mode 100644 index 0000000000..5210b5d2c7 --- /dev/null +++ b/src/fswAlgorithms/orbitControl/orbElemOffset/_UnitTest/test_orbElemOffset.py @@ -0,0 +1,132 @@ +# +# ISC License +# +# Copyright (c) 2025, Autonomous Vehicle Systems +# Laboratory, University of Colorado 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. +# + +import pytest + +from Basilisk.utilities import SimulationBaseClass +from Basilisk.utilities import macros +from Basilisk.utilities import orbitalMotion + +from Basilisk.architecture import messaging +from Basilisk.fswAlgorithms import orbElemOffset + + +@pytest.mark.parametrize("useMeanAnomalyOffset", [False, True]) +def test_orbElemOffset(useMeanAnomalyOffset): + """ + Functional test for OrbElemOffset: + + - When useMeanAnomalyOffset is False: + All ClassicElements fields are summed element wise: + out = main + offset + + - When useMeanAnomalyOffset is True: + All fields except f are still summed element wise. The f field + is computed by interpreting offset.f as a mean anomaly increment + DeltaM (rad) and performing + + E_main = f2E(main.f, main.e) + M_main = E2M(E_main, main.e) + e_out = main.e + offset.e + M_out = M_main + offset.f + f_out = E2f(M2E(M_out, e_out), e_out) + """ + + # 1. Create test simulation + scSim = SimulationBaseClass.SimBaseClass() + proc = scSim.CreateNewProcess("TestProc") + task = scSim.CreateNewTask("unitTask", macros.sec2nano(0.1)) + proc.addTask(task) + + # 2. Construct the module + module = orbElemOffset.OrbElemOffset() + module.ModelTag = "OrbElemOffset" + module.useMeanAnomalyOffset = useMeanAnomalyOffset + scSim.AddModelToTask("unitTask", module) + + # 3. Create two ClassicElements messages + main = messaging.ClassicElementsMsgPayload() + offset = messaging.ClassicElementsMsgPayload() + + # Main elements (choose a moderately eccentric orbit) + main.a = 1.0 + main.e = 0.1 + main.i = 0.2 + main.Omega = 0.3 + main.omega = 0.4 + main.f = 0.5 + + # Offsets: small increments + offset.a = 2.0 + offset.e = 0.02 + offset.i = 0.03 + offset.Omega = 0.04 + offset.omega = 0.05 + # For the angle offset, this will be interpreted either as df or DeltaM + offset.f = 0.06 + + msgMain = messaging.ClassicElementsMsg().write(main) + msgOffset = messaging.ClassicElementsMsg().write(offset) + + # 4. Subscribe module input messages + module.mainElementsInMsg.subscribeTo(msgMain) + module.offsetElementsInMsg.subscribeTo(msgOffset) + + # 5. Record output + recorder = module.elementsOutMsg.recorder() + scSim.AddModelToTask("unitTask", recorder) + + # Initialize and run one step + scSim.InitializeSimulation() + scSim.ConfigureStopTime(macros.sec2nano(0.1)) + scSim.ExecuteSimulation() + + # 6. Extract result - only one write should have occurred + assert recorder.a.shape[0] > 0 + idx = 0 + + # Common expected fields: always simple sums + expected_a = main.a + offset.a + expected_e = main.e + offset.e + expected_i = main.i + offset.i + expected_Omega = main.Omega + offset.Omega + expected_omega = main.omega + offset.omega + + # Check non angle fields + assert recorder.a[idx] == pytest.approx(expected_a) + assert recorder.e[idx] == pytest.approx(expected_e) + assert recorder.i[idx] == pytest.approx(expected_i) + assert recorder.Omega[idx] == pytest.approx(expected_Omega) + assert recorder.omega[idx] == pytest.approx(expected_omega) + + # 7. Check the anomaly handling + if not useMeanAnomalyOffset: + # Plain true anomaly increment + expected_f = main.f + offset.f + assert recorder.f[idx] == pytest.approx(expected_f) + else: + recorderM = orbitalMotion.E2M(orbitalMotion.f2E(recorder.f[idx], recorder.e[idx]), recorder.e[idx]) + mainM = orbitalMotion.E2M(orbitalMotion.f2E(main.f, main.e), main.e) + assert (recorderM - mainM) == pytest.approx(offset.f, rel=1e-12, abs=1e-12) + + +if __name__ == "__main__": + # Run the test standalone for quick checking + test_orbElemOffset(useMeanAnomalyOffset=False) + test_orbElemOffset(useMeanAnomalyOffset=True) diff --git a/src/fswAlgorithms/orbitControl/orbElemOffset/orbElemOffset.cpp b/src/fswAlgorithms/orbitControl/orbElemOffset/orbElemOffset.cpp new file mode 100644 index 0000000000..0dd210250a --- /dev/null +++ b/src/fswAlgorithms/orbitControl/orbElemOffset/orbElemOffset.cpp @@ -0,0 +1,73 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab + + 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 "orbElemOffset.h" + +#include "architecture/utilities/orbitalMotion.h" + +void OrbElemOffset::Reset(uint64_t /*CurrentSimNanos*/) +{ + if (!this->mainElementsInMsg.isLinked()) { + this->bskLogger.bskLog( + BSK_ERROR, + "OrbElemOffset::Reset: mainElementsInMsg is not linked." + ); + } + + if (!this->offsetElementsInMsg.isLinked()) { + this->bskLogger.bskLog( + BSK_ERROR, + "OrbElemOffset::Reset: offsetElementsInMsg is not linked." + ); + } +} + +void OrbElemOffset::UpdateState(uint64_t CurrentSimNanos) +{ + ClassicElementsMsgPayload outPayload{}; + const ClassicElementsMsgPayload main = this->mainElementsInMsg(); + const ClassicElementsMsgPayload offset = this->offsetElementsInMsg(); + + // Element wise addition for all fields except f which is handled separately + outPayload.a = main.a + offset.a; + outPayload.e = main.e + offset.e; + outPayload.i = main.i + offset.i; + outPayload.Omega = main.Omega + offset.Omega; + outPayload.omega = main.omega + offset.omega; + + if (useMeanAnomalyOffset) { + // Interpret offset.f as a mean anomaly increment DeltaM + // Step 1: mean anomaly of the main orbit + const double e_main = main.e; + const double e_out = outPayload.e; + + double E_main = f2E(main.f, e_main); + double M_main = E2M(E_main, e_main); + double M_offset = offset.f; + + // Step 2: final mean anomaly and conversion back to true anomaly + double M_out = M_main + M_offset; + double E_out = M2E(M_out, e_out); + outPayload.f = E2f(E_out, e_out); + } else { + // Plain true anomaly increment + outPayload.f = main.f + offset.f; + } + + this->elementsOutMsg.write(&outPayload, this->moduleID, CurrentSimNanos); +} diff --git a/src/fswAlgorithms/orbitControl/orbElemOffset/orbElemOffset.h b/src/fswAlgorithms/orbitControl/orbElemOffset/orbElemOffset.h new file mode 100644 index 0000000000..0b60910aef --- /dev/null +++ b/src/fswAlgorithms/orbitControl/orbElemOffset/orbElemOffset.h @@ -0,0 +1,119 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab + + 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 ORB_ELEM_OFFSET_H +#define ORB_ELEM_OFFSET_H + +#include + +#include "architecture/_GeneralModuleFiles/sys_model.h" +#include "architecture/messaging/messaging.h" +#include "architecture/msgPayloadDefC/ClassicElementsMsgPayload.h" +#include "architecture/utilities/bskLogging.h" + +/*! + * @brief Classic orbital elements combiner with optional mean anomaly offset handling. + * + * This module reads two ClassicElementsMsgPayload inputs: + * + * - mainElementsInMsg: nominal or "base" classical orbital elements + * - offsetElementsInMsg: element offsets to apply to the base elements + * + * For most fields the module performs a simple element wise addition + * + * \code + * out.a = main.a + offset.a + * out.e = main.e + offset.e + * out.i = main.i + offset.i + * out.Omega = main.Omega + offset.Omega + * out.omega = main.omega + offset.omega + * \endcode + * + * The treatment of the true anomaly field f depends on the flag useMeanAnomalyOffset + * + * - If useMeanAnomalyOffset is false (default) then the offsetElementsInMsg.f + * field is interpreted as a true anomaly increment and the output is + * + * \code + * out.f = main.f + offset.f + * \endcode + * + * - If useMeanAnomalyOffset is true then offsetElementsInMsg.f is interpreted + * as a mean anomaly offset \f$\Delta M\f$ in radians The module then + * + * 1 Converts the main elements (main.f main.e) to mean anomaly M_main + * 2 Forms the final eccentricity as e_out = main.e + offset.e + * 3 Forms the final mean anomaly M_out = M_main + offset.f + * 4 Converts (M_out e_out) back to a true anomaly f_out + * + * and sets out.f = f_out + * + * All other elements in ClassicElementsMsgPayload are left to zero. + * + * The output payload is written to elementsOutMsg. + */ +class OrbElemOffset : public SysModel +{ +public: + /** Default constructor */ + OrbElemOffset() = default; + + /*! + * @brief Reset the internal state and validate input message connections. + * + * This method checks that both mainElementsInMsg and offsetElementsInMsg + * are connected If either is not linked an error is logged through + * the module logger + * + * @param CurrentSimNanos Current simulation time in nanoseconds (unused) + */ + void Reset(uint64_t CurrentSimNanos) override; + + /*! + * @brief Compute the combined classical orbital elements and write them out. + * + * The module reads both input messages, constructs the output classical + * elements according to the rules described in the class documentation, + * and writes the result to elementsOutMsg + * + * @param CurrentSimNanos Current simulation time in nanoseconds + */ + void UpdateState(uint64_t CurrentSimNanos) override; + +public: + /*! + * @brief Toggle interpretation of offsetElementsInMsg.f as a mean anomaly offset. + * + * - false: offsetElementsInMsg.f is treated as a true anomaly increment + * - true: offsetElementsInMsg.f is treated as a mean anomaly offset \f$\Delta M\f$ + */ + bool useMeanAnomalyOffset = false; + + //! Nominal or "main" classical elements input message + ReadFunctor mainElementsInMsg; + + //! Classical elements offset input message + ReadFunctor offsetElementsInMsg; + + //! Output classical elements message for the combined elements + Message elementsOutMsg; + + BSKLogger bskLogger; //!< BSK Logging +}; + +#endif /* ORB_ELEM_OFFSET_H */ diff --git a/src/fswAlgorithms/orbitControl/orbElemOffset/orbElemOffset.i b/src/fswAlgorithms/orbitControl/orbElemOffset/orbElemOffset.i new file mode 100644 index 0000000000..0fc10e79af --- /dev/null +++ b/src/fswAlgorithms/orbitControl/orbElemOffset/orbElemOffset.i @@ -0,0 +1,48 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab + + 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 orbElemOffset + +%include "architecture/utilities/bskException.swg" +%default_bsk_exception(); + +%{ + #include "orbElemOffset.h" +%} + +%pythoncode %{ +from Basilisk.architecture.swig_common_model import * +%} + +%include "std_string.i" +%include "swig_conly_data.i" + +/* Base class and SysModel helpers */ +%include "sys_model.i" + +/* OrbElemOffset class */ +%include "orbElemOffset.h" + +/* Message payloads used by OrbElemOffset */ +%include "architecture/msgPayloadDefC/ClassicElementsMsgPayload.h" +struct ClassicElementsMsg_C; + +%pythoncode %{ +import sys +protectAllClasses(sys.modules[__name__]) +%} 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 c429a069e2..1479d7b6a5 100755 --- a/src/simulation/dynamics/Integrators/svIntegrators.i +++ b/src/simulation/dynamics/Integrators/svIntegrators.i @@ -25,12 +25,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" %} @@ -41,6 +45,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 @@ -56,11 +61,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 @@ -108,13 +117,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 @@ -133,6 +179,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 @@ -231,6 +280,60 @@ def svIntegratorAdaptiveRungeKutta( 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 059418df0d..96a84ff035 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(uint64_t callTimeNanos) = 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/dragEffector/_UnitTest/test_atmoDrag.py b/src/simulation/dynamics/dragEffector/_UnitTest/test_atmoDrag.py index 038383512d..a994ad3a75 100644 --- a/src/simulation/dynamics/dragEffector/_UnitTest/test_atmoDrag.py +++ b/src/simulation/dynamics/dragEffector/_UnitTest/test_atmoDrag.py @@ -342,7 +342,7 @@ def run(show_plots, orbitCase, planetCase): for idx in range(0, len(posData)): oeData = orbitalMotion.rv2elem(mu, posData[idx, 0:3], velData[idx, 0:3]) smaData.append(oeData.a/1000.) - plt.plot(posData[:, 0]*macros.NANO2SEC/P, smaData + plt.plot(dataLog.times()*macros.NANO2SEC/P, smaData ,color='#aa0000', ) plt.xlabel('Time [orbits]') diff --git a/src/simulation/dynamics/spacecraft/spacecraft.cpp b/src/simulation/dynamics/spacecraft/spacecraft.cpp index 363420d949..10bb92cd8c 100755 --- a/src/simulation/dynamics/spacecraft/spacecraft.cpp +++ b/src/simulation/dynamics/spacecraft/spacecraft.cpp @@ -446,6 +446,11 @@ void Spacecraft::equationsOfMotion(double integTimeSeconds, double timeStep) } } +void Spacecraft::equationsOfMotionDiffusion(double integTimeSeconds, double timeStep) +{ + this->equationsOfMotion(integTimeSeconds, timeStep); +} + /*! Prepare for integration process @param integrateToThisTimeNanos Time to integrate to */ diff --git a/src/simulation/dynamics/spacecraft/spacecraft.h b/src/simulation/dynamics/spacecraft/spacecraft.h index cec0819b0f..315a7be009 100755 --- a/src/simulation/dynamics/spacecraft/spacecraft.h +++ b/src/simulation/dynamics/spacecraft/spacecraft.h @@ -84,6 +84,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(uint64_t callTimeNanos) final; //!< method to perform pre-integration steps diff --git a/src/simulation/mujocoDynamics/PIDControllers/PIDControllers.i b/src/simulation/mujocoDynamics/PIDControllers/PIDControllers.i index a5e6c8577c..13c3dc5c14 100644 --- a/src/simulation/mujocoDynamics/PIDControllers/PIDControllers.i +++ b/src/simulation/mujocoDynamics/PIDControllers/PIDControllers.i @@ -19,6 +19,10 @@ %module MJPIDControllers + +%include "architecture/utilities/bskException.swg" +%default_bsk_exception(); + %{ #include "JointPIDController.h" %} @@ -31,7 +35,7 @@ from Basilisk.architecture.swig_common_model import * %include "exception.i" %include "sys_model.i" -%include "simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.h" +%include "simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.i" %include "simulation/mujocoDynamics/_GeneralModuleFiles/PIDController.h" // JointPIDController diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp index ca200fd705..0e9657018a 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp @@ -27,6 +27,28 @@ #include +namespace +{ + /** Returns true if the first three scalar joints are translational + * and along the main axis ([1,0,0], [0,1,0], [0,0,1]). + */ + bool areJoints3DTranslation(std::list& joints) + { + if (joints.size() < 3) return false; + + size_t idx = 0; + for (auto&& joint : joints) + { + if (idx == 3) break; + if (joint.isHinge()) return false; + if (std::fabs(joint.getAxis()[idx] - 1) > 1e-10) return false; + ++idx; + } + + return true; + } +} + MJBody::MJBody(mjsBody* body, MJSpec& spec) : MJObject(body), spec(spec) { @@ -152,20 +174,40 @@ MJFreeJoint & MJBody::getFreeJoint() void MJBody::setPosition(const Eigen::Vector3d& position) { - if (!this->freeJoint.has_value()) { - this->getSpec().getScene().logAndThrow("Tried to set position in non-free body " + + if (this->freeJoint.has_value()) { + this->freeJoint.value().setPosition(position); + } else if (areJoints3DTranslation(scalarJoints)) + { + size_t idx = 0; + for (auto&& joint : scalarJoints) + { + if (idx == 3) break; + joint.setPosition(position[idx]); + ++idx; + } + } else { + this->getSpec().getScene().logAndThrow("Tried to set position in a body with no 'free' joint or no 3D translational joints " + name); } - this->freeJoint.value().setPosition(position); } void MJBody::setVelocity(const Eigen::Vector3d& velocity) { - if (!this->freeJoint.has_value()) { - this->getSpec().getScene().logAndThrow("Tried to set velocity in non-free body " + + if (this->freeJoint.has_value()) { + this->freeJoint.value().setPosition(velocity); + } else if (areJoints3DTranslation(scalarJoints)) + { + size_t idx = 0; + for (auto&& joint : scalarJoints) + { + if (idx == 3) break; + joint.setVelocity(velocity[idx]); + ++idx; + } + } else { + this->getSpec().getScene().logAndThrow("Tried to set velocity in a body with no 'free' joint or no 3D translational joints " + name); } - this->freeJoint.value().setVelocity(velocity); } void MJBody::setAttitude(const Eigen::MRPd& attitude) diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp index 0324ebb029..d2f8e096e3 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp @@ -53,7 +53,7 @@ void MJJoint::configure(const mjModel* m) this->qvelAdr = m->jnt_dofadr[this->getId()]; } -void MJJoint::checkInitialized() +void MJJoint::checkInitialized() const { if (!this->qposAdr.has_value()) { return body.getSpec().getScene().logAndThrow( @@ -69,6 +69,20 @@ MJScalarJoint::MJScalarJoint(mjsJoint* joint, MJBody& body) ) {} +Eigen::Vector3d MJScalarJoint::getAxis() const +{ + checkInitialized(); + const auto m = this->body.getSpec().getMujocoModel(); + return Eigen::Vector3d(m->jnt_axis + (this->qposAdr.value() * 3)); +} + +bool MJScalarJoint::isHinge() const +{ + checkInitialized(); + const auto m = this->body.getSpec().getMujocoModel(); + return m->jnt_type[this->qposAdr.value()] == mjJNT_HINGE; +} + void MJScalarJoint::configure(const mjModel* m) { MJJoint::configure(m); diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.h index 6c672d2485..32a67025aa 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.h @@ -89,7 +89,7 @@ class MJJoint : public MJObject * * Throws an exception if initialization has not been completed. */ - void checkInitialized(); + void checkInitialized() const; protected: MJBody& body; ///< Reference to the body the joint is attached to. @@ -122,6 +122,20 @@ class MJScalarJoint : public MJJoint */ MJScalarJoint(mjsJoint* joint, MJBody& body); + /** + * @brief Returns the axis of rotation for hinge joints and the direction + * of translation for slide joints. + * + * Returned vector is normalized. + */ + Eigen::Vector3d getAxis() const; + + /** + * @brief Returns true if this is a rotational joint, false if it's a + * translational slide joint. + */ + bool isHinge() const; + /** * @brief Sets the position of the joint. * 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 143ae0b3ea..4c95fcfcbb 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) @@ -63,13 +64,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); } @@ -99,16 +118,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); } } @@ -209,7 +243,13 @@ void MJScene::equationsOfMotion(double t, double timeStep) } } -void MJScene::preIntegration(uint64_t callTimeNanos) { this->timeStep = diffNanoToSec(callTimeNanos, this->timeBeforeNanos); } +void MJScene::equationsOfMotionDiffusion(double t, double timeStep) +{ + auto nanos = static_cast(t * SEC2NANO); + this->dynamicsDiffusionTask.ExecuteTaskList(nanos); +} + +void MJScene::preIntegration(uint64_t callTime) { this->timeStep = diffNanoToSec(callTime, this->timeBeforeNanos); } void MJScene::postIntegration(uint64_t callTimeNanos) { @@ -226,6 +266,7 @@ void MJScene::postIntegration(uint64_t callTimeNanos) // 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 3038a2eea8..dbf21a9fc1 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h @@ -314,6 +314,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 @@ -327,6 +333,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. * @@ -358,6 +381,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. */ @@ -391,11 +425,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. * @@ -582,6 +635,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..1a2789272c --- /dev/null +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.i @@ -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. + + */ +%module cStatefulSysModel + +%include "architecture/utilities/bskException.swg" +%default_bsk_exception(); + +%{ + #include "simulation/mujocoDynamics/_GeneralModuleFiles/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 8cdae82764..ebd5d2ba40 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/pyStatefulSysModel.i +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/pyStatefulSysModel.i @@ -16,14 +16,17 @@ 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 "architecture/utilities/bskException.swg" +%default_bsk_exception(); + %{ #include "StatefulSysModel.h" %} %pythoncode %{ import sys -import traceback from Basilisk.architecture.swig_common_model import * %} @@ -35,22 +38,14 @@ from Basilisk.architecture.swig_common_model import * %ignore DynParamRegisterer::DynParamRegisterer; %feature("director") StatefulSysModel; -%feature("pythonappend") StatefulSysModel::StatefulSysModel %{ - self.__super_init_called__ = True%} %rename("_StatefulSysModel") StatefulSysModel; -%include "StatefulSysModel.h" +%include "StatefulSysModel.i" %template(registerState) DynParamRegisterer::registerState; %pythoncode %{ -class StatefulSysModel(_StatefulSysModel, metaclass=Basilisk.architecture.sysModel.SuperInitChecker): - bskLogger: BSKLogger = None - - def __init_subclass__(cls): - # Make it so any exceptions in UpdateState and Reset - # print any exceptions before returning control to - # C++ (at which point exceptions will crash the program) - cls.UpdateState = Basilisk.architecture.sysModel.logError(cls.UpdateState) - cls.Reset = Basilisk.architecture.sysModel.logError(cls.Reset) - cls.registerStates = Basilisk.architecture.sysModel.logError(cls.registerStates) +from Basilisk.architecture.sysModel import SysModelMixin + +class StatefulSysModel(SysModelMixin, _StatefulSysModel): + """Python wrapper for the C++ StatefulSysModel.""" %} diff --git a/src/simulation/mujocoDynamics/aerodynamicDrag/_UnitTest/test_aerodynamicDrag.py b/src/simulation/mujocoDynamics/aerodynamicDrag/_UnitTest/test_aerodynamicDrag.py new file mode 100644 index 0000000000..9e56efa02b --- /dev/null +++ b/src/simulation/mujocoDynamics/aerodynamicDrag/_UnitTest/test_aerodynamicDrag.py @@ -0,0 +1,366 @@ +from typing import Literal + +import pytest +import numpy as np +import numpy.testing as npt + +from Basilisk.utilities import SimulationBaseClass +from Basilisk.architecture import messaging +from Basilisk.utilities import macros +from Basilisk.utilities import RigidBodyKinematics as rbk +from Basilisk.utilities import simIncludeGravBody +from Basilisk.utilities import orbitalMotion +from Basilisk.simulation import pointMassGravityModel +from Basilisk.simulation import exponentialAtmosphere + +try: + from Basilisk.simulation import mujoco + from Basilisk.simulation import MJAerodynamicDrag + from Basilisk.simulation import NBodyGravity + + couldImportMujoco = True +except: + couldImportMujoco = False + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +def test_aerodynamicDrag(): + """ + Unit test for classes that inherit from MeanRevertingNoise. + """ + dt = 1 # s + + scSim = SimulationBaseClass.SimBaseClass() + dynProcess = scSim.CreateNewProcess("test") + dynProcess.addTask(scSim.CreateNewTask("test", macros.sec2nano(dt))) + + # Create the MJScene from a simple cannonball body + scene = mujoco.MJScene("") # empty scene, no multi-body dynamics + scSim.AddModelToTask("test", scene) + + density = 2 # kg/m^3 + cd = 1.5 + area = 2.75 # m^2 + r_CP_S = [1, 0, 0] # m + + v_SN_N = [0, 0.5, 0] # m/s + sigma_SN = [0.1, 0.2, 0.3] + + atmoMsg = messaging.AtmoPropsMsg() + atmoMsg.write(messaging.AtmoPropsMsgPayload(neutralDensity=density)) + + dragGeometryMsg = messaging.DragGeometryMsg() + dragGeometryMsg.write(messaging.DragGeometryMsgPayload( + projectedArea = area, + dragCoeff = cd, + r_CP_S = r_CP_S + )) + + siteFrameMsg = messaging.SCStatesMsg() + siteFrameMsg.write(messaging.SCStatesMsgPayload( + sigma_BN = sigma_SN, + v_BN_N = v_SN_N + )) + + drag = MJAerodynamicDrag.AerodynamicDrag() + drag.dragGeometryInMsg.subscribeTo(dragGeometryMsg) + drag.atmoDensInMsg.subscribeTo(atmoMsg) + drag.referenceFrameStateInMsg.subscribeTo(siteFrameMsg) + scene.AddModelToDynamicsTask(drag) + + scSim.InitializeSimulation() + scSim.ConfigureStopTime(macros.sec2nano(dt)) + scSim.ExecuteSimulation() + + force_S = drag.forceOutMsg.read().force_S + torque_S = drag.torqueOutMsg.read().torque_S + + # Expected + # Rotate velocity from N to S using MRPs + C_BN = np.array(rbk.MRP2C(sigma_SN)) # maps N components to B=S components + v_S = C_BN.dot(v_SN_N) + v_mag = np.linalg.norm(v_SN_N) # drag magnitude depends on speed, frame invariant + if v_mag <= 1e-12: + F_exp = np.zeros(3) + T_exp = np.zeros(3) + else: + F_mag = 0.5 * density * (v_mag ** 2) * cd * area + v_hat_S = v_S / np.linalg.norm(v_S) + F_exp = -F_mag * v_hat_S + T_exp = np.cross(r_CP_S, F_exp) + + npt.assert_allclose(force_S, F_exp, rtol=0.0, atol=1e-12) + npt.assert_allclose(torque_S, T_exp, rtol=0.0, atol=1e-12) + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +def test_orbit(orbitCase: Literal["LPO", "LTO"], planetCase: Literal["earth", "mars"], showPlots = False): + def cannonballDragBFrame(dragCoeff, density, area, velInertial, sigmaBN): + speed = np.linalg.norm(velInertial) + if speed <= 0.0: + return np.zeros(3) + + dragDirInertial = -velInertial / speed + dcmBN = rbk.MRP2C(sigmaBN) + dragDirBody = dcmBN.dot(dragDirInertial) + + return 0.5 * dragCoeff * density * area * speed**2 * dragDirBody + + + scSim = SimulationBaseClass.SimBaseClass() + processName = "dragProcess" + taskName = "dragTask" + dynamicsProcess = scSim.CreateNewProcess(processName) + stepNanos = macros.sec2nano(1.0) + dynamicsProcess.addTask(scSim.CreateNewTask(taskName, stepNanos)) + + bodyName = "satellite" + dragSceneXml = fr""" + + + + + + + + + + + """ + scene = mujoco.MJScene(dragSceneXml) + scene.extraEoMCall = True + scSim.AddModelToTask(taskName, scene) + + mjBody = scene.getBody(bodyName) + + planetData = simIncludeGravBody.BODY_DATA[planetCase] + gravitationalParameter = planetData.mu + planetRadius = planetData.radEquator + + gravityModel = NBodyGravity.NBodyGravity() + gravityModel.ModelTag = "gravity" + scene.AddModelToDynamicsTask(gravityModel) + + pointMassModel = pointMassGravityModel.PointMassGravityModel() + pointMassModel.muBody = gravitationalParameter + gravityModel.addGravitySource(planetCase, pointMassModel, True) + gravityModel.addGravityTarget(bodyName, mjBody) + + if planetCase == "earth": + baseDensity = 1.217 + scaleHeight = 8500.0 + else: + baseDensity = 0.020 + scaleHeight = 11100.0 + + atmoModel = exponentialAtmosphere.ExponentialAtmosphere() + atmoModel.ModelTag = "ExpAtmo" + atmoModel.planetRadius = planetRadius + atmoModel.baseDensity = baseDensity + atmoModel.scaleHeight = scaleHeight + scene.AddModelToDynamicsTask(atmoModel) + + comStateMsg = mjBody.getCenterOfMass().stateOutMsg + atmoModel.addSpacecraftToModel(comStateMsg) + + projectedArea = 10.0 + dragCoeff = 2.2 + + dragGeometryMsg = messaging.DragGeometryMsg() + dragGeometryMsg.write( + messaging.DragGeometryMsgPayload( + projectedArea=projectedArea, + dragCoeff=dragCoeff, + r_CP_S=[1.0, 0.0, 0.0], + ) + ) + + dragModel = MJAerodynamicDrag.AerodynamicDrag() + dragModel.ModelTag = "drag" + scene.AddModelToDynamicsTask(dragModel) + forceTorqueDrag: mujoco.MJForceTorqueActuator = dragModel.applyTo(mjBody) + dragModel.atmoDensInMsg.subscribeTo(atmoModel.envOutMsgs[0]) + dragModel.dragGeometryInMsg.subscribeTo(dragGeometryMsg) + + if orbitCase == "LPO": + orbAltMin = 300e3 # m + orbAltMax = orbAltMin + elif orbitCase == "LTO": + orbAltMin = 300e3 # m + orbAltMax = 800e3 # m + + rMin = planetRadius + orbAltMin + rMax = planetRadius + orbAltMax + + elements = orbitalMotion.ClassicElements() + elements.a = (rMin+rMax)/2.0 + elements.e = 1.0 - rMin/elements.a + elements.i = 0.0 + elements.Omega = 0.0 + elements.omega = 0.0 + elements.f = 0.0 + + meanMotion = np.sqrt(gravitationalParameter / elements.a**3) + orbitPeriod = 2.0 * np.pi / meanMotion + finalTimeNanos = macros.sec2nano(orbitPeriod) + + stateRecorder = comStateMsg.recorder() + atmoRecorder = atmoModel.envOutMsgs[0].recorder() + dragForceRecorder = forceTorqueDrag.forceInMsg.recorder() + + scSim.AddModelToTask(taskName, stateRecorder) + scSim.AddModelToTask(taskName, atmoRecorder) + scSim.AddModelToTask(taskName, dragForceRecorder) + + scSim.InitializeSimulation() + + positionInertial, velocityInertial = orbitalMotion.elem2rv(gravitationalParameter, elements) + mjBody.setPosition(positionInertial) + mjBody.setVelocity(velocityInertial) + + scSim.ConfigureStopTime(finalTimeNanos) + scSim.ExecuteSimulation() + + positionData = stateRecorder.r_BN_N + velocityData = stateRecorder.v_BN_N + attitudeData = stateRecorder.sigma_BN + densityData = atmoRecorder.neutralDensity + dragForceData = dragForceRecorder.force_S + + numSteps = dragForceData.shape[0] + referenceDrag = np.zeros_like(dragForceData) + + for timeIndex in range(numSteps): + referenceDrag[timeIndex, :] = cannonballDragBFrame( + dragCoeff=dragCoeff, + density=densityData[timeIndex], + area=projectedArea, + velInertial=velocityData[timeIndex, :], + sigmaBN=attitudeData[timeIndex, :], + ) + + npt.assert_allclose( + dragForceData, + referenceDrag, + rtol=0, + atol=1e-12, + ) + + if showPlots: + import matplotlib.pyplot as plt + + timesSec = stateRecorder.times() * macros.NANO2SEC + timesOrbit = timesSec / orbitPeriod + + # Inertial position components + plt.figure() + fig = plt.gcf() + ax = fig.gca() + ax.ticklabel_format(useOffset=False, style="plain") + for idx in range(3): + plt.plot( + timesOrbit, + positionData[:, idx] / 1e3, + label=f"$r_{{BN,{idx}}}$", + ) + plt.legend(loc="lower right") + plt.xlabel("Time [orbits]") + plt.ylabel("Inertial position [km]") + + # Orbit in perifocal frame + elementsInit = orbitalMotion.rv2elem(gravitationalParameter, positionInertial, velocityInertial) + bParam = elementsInit.a * np.sqrt(1.0 - elementsInit.e**2) + pParam = elementsInit.a * (1.0 - elementsInit.e**2) + + plt.figure(figsize=tuple(np.array((1.0, bParam / elementsInit.a)) * 4.75), dpi=100) + axisLimits = np.array( + [-elementsInit.rApoap, elementsInit.rPeriap, -bParam, bParam] + ) / 1e3 * 1.25 + plt.axis(axisLimits) + + fig = plt.gcf() + ax = fig.gca() + planetColor = "#008800" + planetRadiusKm = planetRadius / 1e3 + ax.add_artist(plt.Circle((0.0, 0.0), planetRadiusKm, color=planetColor)) + + rDataList = [] + fDataList = [] + for idx in range(positionData.shape[0]): + elementsStep = orbitalMotion.rv2elem( + gravitationalParameter, + positionData[idx, 0:3], + velocityData[idx, 0:3], + ) + rDataList.append(elementsStep.rmag) + fDataList.append(elementsStep.f + elementsStep.omega - elementsInit.omega) + + rDataArray = np.array(rDataList) + fDataArray = np.array(fDataList) + + plt.plot( + rDataArray * np.cos(fDataArray) / 1e3, + rDataArray * np.sin(fDataArray) / 1e3, + color="#aa0000", + linewidth=3.0, + ) + + fGrid = np.linspace(0.0, 2.0 * np.pi, 200) + rGrid = pParam / (1.0 + elementsInit.e * np.cos(fGrid)) + plt.plot( + rGrid * np.cos(fGrid) / 1e3, + rGrid * np.sin(fGrid) / 1e3, + "--", + color="#555555", + ) + plt.xlabel("$i_e$ coord [km]") + plt.ylabel("$i_p$ coord [km]") + plt.grid() + + # Semi major axis vs time + plt.figure() + fig = plt.gcf() + ax = fig.gca() + ax.ticklabel_format(useOffset=False, style="plain") + smaData = [] + for idx in range(positionData.shape[0]): + elementsStep = orbitalMotion.rv2elem( + gravitationalParameter, + positionData[idx, 0:3], + velocityData[idx, 0:3], + ) + smaData.append(elementsStep.a / 1e3) + smaData = np.array(smaData) + plt.plot(timesOrbit, smaData, color="#aa0000") + plt.xlabel("Time [orbits]") + plt.ylabel("SMA [km]") + + # Density vs time + plt.figure() + fig = plt.gcf() + ax = fig.gca() + ax.ticklabel_format(useOffset=False, style="sci") + timesAtmoSec = atmoRecorder.times() * macros.NANO2SEC + plt.plot(timesAtmoSec, densityData) + plt.title("Density vs time") + plt.xlabel("Time [s]") + plt.ylabel("Density [kg/m^3]") + + # Drag vs reference + plt.figure() + fig = plt.gcf() + ax = fig.gca() + ax.ticklabel_format(useOffset=False, style="plain") + for idx in range(3): + plt.plot(timesOrbit, dragForceData[:, idx], label=f"F_mod_{idx}") + plt.plot(timesOrbit, referenceDrag[:, idx], "--", label=f"F_ref_{idx}") + plt.xlabel("Time [orbits]") + plt.ylabel("Drag force [N]") + plt.legend() + plt.title("Drag force vs reference") + + plt.show() + + +if __name__ == "__main__": + assert couldImportMujoco + # test_aerodynamicDrag() + test_orbit("LTO","earth",True) diff --git a/src/simulation/mujocoDynamics/aerodynamicDrag/aerodynamicDrag.cpp b/src/simulation/mujocoDynamics/aerodynamicDrag/aerodynamicDrag.cpp new file mode 100644 index 0000000000..bbbb1296c9 --- /dev/null +++ b/src/simulation/mujocoDynamics/aerodynamicDrag/aerodynamicDrag.cpp @@ -0,0 +1,92 @@ +#include "aerodynamicDrag.h" + +#include + +void +AerodynamicDrag::UpdateState(uint64_t CurrentSimNanos) +{ + const double density = this->atmoDensInMsg().neutralDensity; + + const auto geomPayload = this->dragGeometryInMsg(); + const double cd = geomPayload.dragCoeff; + const double area = geomPayload.projectedArea; + const Eigen::Map r_CP_s(geomPayload.r_CP_S); + + // 'S' denotes the 'site' reference frame + const auto siteStatePayload = this->referenceFrameStateInMsg(); + Eigen::MRPd sigma_SN; + sigma_SN = Eigen::Map(siteStatePayload.sigma_BN); + const Eigen::Matrix3d dcm_SN = sigma_SN.toRotationMatrix().transpose(); + + // inertial velocity of the center of the S frame, expressed in the S frame + const auto v_S = dcm_SN*Eigen::Map(siteStatePayload.v_BN_N); + const auto v = v_S.norm(); + + if (v < 1e-10) + { + this->forceOutMsg.write(&this->forceOutMsg.zeroMsgPayload, this->moduleID, CurrentSimNanos); + this->torqueOutMsg.write(&this->torqueOutMsg.zeroMsgPayload, this->moduleID, CurrentSimNanos); + return; + } + + const auto v_hat_S = v_S / v; + + const Eigen::Vector3d force_S = -0.5 * cd * area * (v*v) * density * v_hat_S; + const Eigen::Vector3d torque_S = r_CP_s.cross(force_S); + + ForceAtSiteMsgPayload forcePayload = this->forceOutMsg.zeroMsgPayload; + TorqueAtSiteMsgPayload torquePayload = this->torqueOutMsg.zeroMsgPayload; + std::copy_n(force_S.data(), 3, forcePayload.force_S); + std::copy_n(torque_S.data(), 3, torquePayload.torque_S); + + this->forceOutMsg.write(&forcePayload, this->moduleID, CurrentSimNanos); + this->torqueOutMsg.write(&torquePayload, this->moduleID, CurrentSimNanos);} + +void AerodynamicDrag::Reset(uint64_t /*CurrentSimNanos*/) +{ + if (!this->atmoDensInMsg.isLinked()) { + bskLogger.bskLog(BSK_ERROR, + "AerodynamicDrag: atmoDensInMsg is not linked. " + "Call atmoDensInMsg.subscribeTo() before simulation."); + MJBasilisk::detail::logAndThrow( + "AerodynamicDrag missing input: atmoDensInMsg", &bskLogger); + } + + if (!this->dragGeometryInMsg.isLinked()) { + bskLogger.bskLog(BSK_ERROR, + "AerodynamicDrag: dragGeometryInMsg is not linked. " + "Call dragGeometryInMsg.subscribeTo() before simulation."); + MJBasilisk::detail::logAndThrow( + "AerodynamicDrag missing input: dragGeometryInMsg", &bskLogger); + } + + if (!this->referenceFrameStateInMsg.isLinked()) { + bskLogger.bskLog(BSK_ERROR, + "AerodynamicDrag: referenceFrameStateInMsg is not linked. " + "Call referenceFrameStateInMsg.subscribeTo() before simulation."); + MJBasilisk::detail::logAndThrow( + "AerodynamicDrag missing input: referenceFrameStateInMsg", &bskLogger); + } +} + + +MJForceTorqueActuator& +AerodynamicDrag::applyTo(MJBody& body) +{ + return applyTo(body.getCenterOfMass()); +} + +MJForceTorqueActuator& +AerodynamicDrag::applyTo(MJSite& site) +{ + this->referenceFrameStateInMsg.subscribeTo( &site.stateOutMsg ); + return applyTo(site.getBody().getSpec().getScene().addForceTorqueActuator(ModelTag + "_aerodynamicDrag_" + site.getName(), site)); +} + +MJForceTorqueActuator& +AerodynamicDrag::applyTo(MJForceTorqueActuator& actuator) +{ + actuator.forceInMsg.subscribeTo( &this->forceOutMsg ); + actuator.torqueInMsg.subscribeTo( &this->torqueOutMsg ); + return actuator; +} diff --git a/src/simulation/mujocoDynamics/aerodynamicDrag/aerodynamicDrag.h b/src/simulation/mujocoDynamics/aerodynamicDrag/aerodynamicDrag.h new file mode 100644 index 0000000000..f9c57b7e73 --- /dev/null +++ b/src/simulation/mujocoDynamics/aerodynamicDrag/aerodynamicDrag.h @@ -0,0 +1,131 @@ +/* + 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 AERODYNAMIC_DRAG_H +#define AERODYNAMIC_DRAG_H + +#include "architecture/utilities/bskLogging.h" +#include "architecture/messaging/messaging.h" +#include "architecture/_GeneralModuleFiles/sys_model.h" +#include "simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h" + +#include "architecture/msgPayloadDefC/DragGeometryMsgPayload.h" +#include "architecture/msgPayloadDefC/AtmoPropsMsgPayload.h" +#include "architecture/msgPayloadDefC/ForceAtSiteMsgPayload.h" +#include "architecture/msgPayloadDefC/TorqueAtSiteMsgPayload.h" + +/** + * @brief Aerodynamic drag model that computes force and torque at a MuJoCo site. + * + * Updates a force and a torque to represent quasi-steady aerodynamic drag acting + * at a site fixed on a body. The model reads: + * - Geometry: drag coefficient and projected reference area. + * - Atmosphere: neutral density. + * - Reference frame state: position, velocity, and attitude of the site frame. + * + * Force direction is opposite the relative flow. Magnitude follows + * \f$ F = \tfrac{1}{2}\rho\,v^2 C_D A \f$. + * Torque is \f$ \tau = r_{CP/S} \times F \f$ in the site frame. + * + * The force and torque output in `forceOutMsg` and `torqueOutMsg` are given + * in the site reference frame whose state is given in `referenceFrameStateInMsg`. + * + * Frames and notation: + * - N: inertial. + * - S: site-fixed reference frame where outputs are expressed. + */ +class AerodynamicDrag : public SysModel +{ +public: + /** + * @brief Bind this model to an MJBody target. + * + * If this method is used, the `dragGeometryInMsg` and + * `atmoDensInMsg` must be linked. + * + * @param body MuJoCo body that will receive the drag force and torque at it's center of mass. + * @return Reference to the actuator on which the forceand torque is applied. + */ + MJForceTorqueActuator& applyTo(MJBody& body); + + /** + * @brief Bind this model to an MJSite target. + * + * If this method is used, only the `dragGeometryInMsg` and + * `atmoDensInMsg` must be linked. + * + * @param site MuJoCo site where the force application point is defined. + * @return Reference to the actuator on which the forceand torque is applied. + */ + MJForceTorqueActuator& applyTo(MJSite& site); + + /** + * @brief Bind this model to an existing force/torque actuator. + * + * If this method is used, the `dragGeometryInMsg`, + * `atmoDensInMsg`, and `referenceFrameStateInMsg` must be linked. + * + * @param actuator Precreated actuator that will receive the computed wrench. + * @return Reference to the same actuator. + * + * @note You + */ + MJForceTorqueActuator& applyTo(MJForceTorqueActuator& actuator); + + /** @name Framework interface @{ */ + + /** + * @brief Advance the model and publish outputs. + * @param CurrentSimNanos Current simulation time in nanoseconds. + */ + void UpdateState(uint64_t CurrentSimNanos) override; + + /** + * @brief Validate all input messages are linked. + * @param CurrentSimNanos Current simulation time in nanoseconds. + */ + void Reset(uint64_t CurrentSimNanos) override; + + /** @} */ + +public: + /** @name I/O Messages @{ */ + + /** @brief Geometry input. */ + ReadFunctor dragGeometryInMsg; + + /** @brief Atmospheric properties input. */ + ReadFunctor atmoDensInMsg; + + /** @brief State of the reference frame . */ + ReadFunctor referenceFrameStateInMsg; + + /** @brief Output force at the site in the site reference frame. */ + Message forceOutMsg; + + /** @brief Output torque in the site reference frame.*/ + Message torqueOutMsg; + + /** @} */ + + /** @brief Logger */ + BSKLogger bskLogger; +}; + + +#endif diff --git a/src/simulation/mujocoDynamics/aerodynamicDrag/aerodynamicDrag.i b/src/simulation/mujocoDynamics/aerodynamicDrag/aerodynamicDrag.i new file mode 100644 index 0000000000..60677b456c --- /dev/null +++ b/src/simulation/mujocoDynamics/aerodynamicDrag/aerodynamicDrag.i @@ -0,0 +1,52 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado 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 MJAerodynamicDrag + +%include "architecture/utilities/bskException.swg" +%default_bsk_exception(); + +%{ + #include "aerodynamicDrag.h" +%} + +%pythoncode %{ + from Basilisk.architecture.swig_common_model import * +%} +%include "std_string.i" +%include "swig_conly_data.i" + +%include "sys_model.i" +%include "aerodynamicDrag.h" + +%include "architecture/msgPayloadDefC/DragGeometryMsgPayload.h" +struct DragGeometryMsg_C; +%include "architecture/msgPayloadDefC/AtmoPropsMsgPayload.h" +struct AtmoPropsMsg_C; +%include "architecture/msgPayloadDefC/SCStatesMsgPayload.h" +struct SCStatesMsg_C; +%include "architecture/msgPayloadDefC/ForceAtSiteMsgPayload.h" +struct ForceAtSiteMsg_C; +%include "architecture/msgPayloadDefC/TorqueAtSiteMsgPayload.h" +struct TorqueAtSiteMsg_C; + +%pythoncode %{ +import sys +protectAllClasses(sys.modules[__name__]) +%} diff --git a/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/_UnitTest/test_cmdForceInertialToForceAtSite.py b/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/_UnitTest/test_cmdForceInertialToForceAtSite.py new file mode 100644 index 0000000000..26beb97867 --- /dev/null +++ b/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/_UnitTest/test_cmdForceInertialToForceAtSite.py @@ -0,0 +1,102 @@ +import pytest +import numpy as np +import numpy.testing as npt + +from Basilisk.utilities import SimulationBaseClass +from Basilisk.architecture import messaging +from Basilisk.utilities import macros +from Basilisk.utilities import RigidBodyKinematics as rbk + +try: + from Basilisk.simulation import mujoco + from Basilisk.simulation import MJCmdForceInertialToForceAtSite + couldImportMujoco = True +except Exception: + couldImportMujoco = False + + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +def test_cmdForceInertialToForceAtSiteSCStates(): + dt = 1.0 + + scSim = SimulationBaseClass.SimBaseClass() + proc = scSim.CreateNewProcess("p") + proc.addTask(scSim.CreateNewTask("t", macros.sec2nano(dt))) + + scene = mujoco.MJScene("") + scSim.AddModelToTask("t", scene) + + # Inputs + F_N = np.array([10.0, -2.0, 5.0]) # N in inertial frame + sigma_BN = np.array([0.1, 0.2, 0.3]) # MRPs for S=B wrt N + + cmdMsg = messaging.CmdForceInertialMsg() + cmdMsg.write(messaging.CmdForceInertialMsgPayload(forceRequestInertial=F_N.tolist())) + + scMsg = messaging.SCStatesMsg() + scMsg.write(messaging.SCStatesMsgPayload(sigma_BN=sigma_BN.tolist())) + + mod = MJCmdForceInertialToForceAtSite.CmdForceInertialToForceAtSite() + mod.cmdForceInertialInMsg.subscribeTo(cmdMsg) + mod.siteFrameStateInMsg.subscribeTo(scMsg) + scene.AddModelToDynamicsTask(mod) + + scSim.InitializeSimulation() + scSim.ConfigureStopTime(macros.sec2nano(dt)) + scSim.ExecuteSimulation() + + F_S = np.array(mod.forceOutMsg.read().force_S, dtype=float) + + # Expected: F_S = C_SN * F_N, and C_SN == C_BN from MRPs + C_BN = np.array(rbk.MRP2C(sigma_BN)) + F_S_exp = C_BN.dot(F_N) + + npt.assert_allclose(F_S, F_S_exp, rtol=0.0, atol=1e-12) + + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +def test_cmdForceInertialToForceAtSiteNavAtt(): + dt = 1.0 + + scSim = SimulationBaseClass.SimBaseClass() + proc = scSim.CreateNewProcess("p") + proc.addTask(scSim.CreateNewTask("t", macros.sec2nano(dt))) + + scene = mujoco.MJScene("") + scSim.AddModelToTask("t", scene) + + F_N = np.array([3.0, 4.0, 0.0]) + sigmaAtt_BN = np.array([0.05, -0.02, 0.01]) # will be used + sigmaState_BN = np.array([0.3, 0.0, 0.0]) # should be ignored if NavAtt is linked + + cmdMsg = messaging.CmdForceInertialMsg() + cmdMsg.write(messaging.CmdForceInertialMsgPayload(forceRequestInertial=F_N.tolist())) + + attMsg = messaging.NavAttMsg() + attMsg.write(messaging.NavAttMsgPayload(sigma_BN=sigmaAtt_BN.tolist())) + + scMsg = messaging.SCStatesMsg() + scMsg.write(messaging.SCStatesMsgPayload(sigma_BN=sigmaState_BN.tolist())) + + mod = MJCmdForceInertialToForceAtSite.CmdForceInertialToForceAtSite() + mod.cmdForceInertialInMsg.subscribeTo(cmdMsg) + mod.siteAttInMsg.subscribeTo(attMsg) # preferred + mod.siteFrameStateInMsg.subscribeTo(scMsg) # fallback + scene.AddModelToDynamicsTask(mod) + + scSim.InitializeSimulation() + scSim.ConfigureStopTime(macros.sec2nano(dt)) + scSim.ExecuteSimulation() + + F_S = np.array(mod.forceOutMsg.read().force_S, dtype=float) + + C_BN_att = np.array(rbk.MRP2C(sigmaAtt_BN)) + F_S_exp = C_BN_att.dot(F_N) + + npt.assert_allclose(F_S, F_S_exp, rtol=0.0, atol=1e-12) + + +if __name__ == "__main__": + assert couldImportMujoco + test_cmdForceInertialToForceAtSiteSCStates() + test_cmdForceInertialToForceAtSiteNavAtt() diff --git a/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/cmdForceInertialToForceAtSite.cpp b/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/cmdForceInertialToForceAtSite.cpp new file mode 100644 index 0000000000..0b7f113e03 --- /dev/null +++ b/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/cmdForceInertialToForceAtSite.cpp @@ -0,0 +1,43 @@ +#include "cmdForceInertialToForceAtSite.h" + +#include +#include "architecture/utilities/avsEigenMRP.h" + +void +CmdForceInertialToForceAtSite::UpdateState(uint64_t CurrentSimNanos) +{ + // 'S' denotes the 'site' reference frame + Eigen::MRPd sigma_SN; + if (this->siteAttInMsg.isLinked()) + { + sigma_SN = Eigen::Map(this->siteAttInMsg().sigma_BN); + } + else + { + sigma_SN = Eigen::Map(this->siteFrameStateInMsg().sigma_BN); + } + const Eigen::Matrix3d dcm_SN = sigma_SN.toRotationMatrix().transpose(); + + Eigen::Map force_N{this->cmdForceInertialInMsg().forceRequestInertial}; + const Eigen::Vector3d force_S = dcm_SN*force_N; + + ForceAtSiteMsgPayload forcePayload = this->forceOutMsg.zeroMsgPayload; + std::copy_n(force_S.data(), 3, forcePayload.force_S); + this->forceOutMsg.write(&forcePayload, this->moduleID, CurrentSimNanos); +} + +void CmdForceInertialToForceAtSite::Reset(uint64_t /*CurrentSimNanos*/) +{ + if (!this->cmdForceInertialInMsg.isLinked()) { + bskLogger.bskLog(BSK_ERROR, + "CmdForceInertialToForceAtSite: cmdForceInertialInMsg is not linked. " + "Call cmdForceInertialInMsg.subscribeTo() before simulation."); + } + + if (!this->siteFrameStateInMsg.isLinked() && !this->siteAttInMsg.isLinked()) { + bskLogger.bskLog(BSK_ERROR, + "CmdForceInertialToForceAtSite: siteFrameStateInMsg and siteAttInMsg are not linked. At least one must be linked." + "Call siteFrameStateInMsg.subscribeTo() or siteAttInMsg.subscribeTo() before simulation."); + } + +} diff --git a/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/cmdForceInertialToForceAtSite.h b/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/cmdForceInertialToForceAtSite.h new file mode 100644 index 0000000000..0ad685735a --- /dev/null +++ b/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/cmdForceInertialToForceAtSite.h @@ -0,0 +1,80 @@ +/* + 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 CMD_FORCE_INERTIAL_TO_FORCE_AT_SITE_H +#define CMD_FORCE_INERTIAL_TO_FORCE_AT_SITE_H + +#include "architecture/utilities/bskLogging.h" +#include "architecture/messaging/messaging.h" +#include "architecture/_GeneralModuleFiles/sys_model.h" + +#include "architecture/msgPayloadDefC/CmdForceInertialMsgPayload.h" +#include "architecture/msgPayloadDefC/ForceAtSiteMsgPayload.h" +#include "architecture/msgPayloadDefC/SCStatesMsgPayload.h" +#include "architecture/msgPayloadDefC/NavAttMsgPayload.h" + +/** + * @brief Convert a commanded force in the inertial frame into a force at a site frame. + * + * Computes a force vector expressed in the site frame S from an input force expressed in the + * inertial frame N. Uses either the site attitude message (preferred) or the site state message + * to obtain orientation S relative to N. The output message contains the force expressed in S. + */ +class CmdForceInertialToForceAtSite : public SysModel +{ +public: + /** @name Framework interface @{ */ + + /** + * @brief Advance the model and publish outputs. + * Reads forceRequestInertial in N, builds C_SN from sigma_BN, computes F_S, and writes forceOutMsg. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + */ + void UpdateState(uint64_t CurrentSimNanos) override; + + /** + * @brief Validate input links. + * @param CurrentSimNanos Current simulation time in nanoseconds. + */ + void Reset(uint64_t CurrentSimNanos) override; + + /** @} */ + +public: + /** @name I/O Messages @{ */ + + /// Commanded force in N + ReadFunctor cmdForceInertialInMsg; + + /// Site state providing sigma_BN as a fallback orientation source + ReadFunctor siteFrameStateInMsg; + + /// Site attitude providing sigma_BN as the primary orientation source + ReadFunctor siteAttInMsg; + + /// Output force expressed in S + Message forceOutMsg; + + /** @} */ + + /// Logger + BSKLogger bskLogger; +}; + +#endif diff --git a/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/cmdForceInertialToForceAtSite.i b/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/cmdForceInertialToForceAtSite.i new file mode 100644 index 0000000000..a34dcb3bbd --- /dev/null +++ b/src/simulation/mujocoDynamics/cmdForceInertialToForceAtSite/cmdForceInertialToForceAtSite.i @@ -0,0 +1,50 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado 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 MJCmdForceInertialToForceAtSite + +%include "architecture/utilities/bskException.swg" +%default_bsk_exception(); + +%{ + #include "cmdForceInertialToForceAtSite.h" +%} + +%pythoncode %{ + from Basilisk.architecture.swig_common_model import * +%} +%include "std_string.i" +%include "swig_conly_data.i" + +%include "sys_model.i" +%include "cmdForceInertialToForceAtSite.h" + +%include "architecture/msgPayloadDefC/CmdForceInertialMsgPayload.h" +struct CmdForceInertialMsg_C; +%include "architecture/msgPayloadDefC/NavAttMsgPayload.h" +struct NavAttMsg_C; +%include "architecture/msgPayloadDefC/SCStatesMsgPayload.h" +struct SCStatesMsg_C; +%include "architecture/msgPayloadDefC/ForceAtSiteMsgPayload.h" +struct ForceAtSiteMsg_C; + +%pythoncode %{ +import sys +protectAllClasses(sys.modules[__name__]) +%} diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/_UnitTest/test_linearTimeInvariantSystem.py b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/_UnitTest/test_linearTimeInvariantSystem.py new file mode 100644 index 0000000000..5c4997997f --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/_UnitTest/test_linearTimeInvariantSystem.py @@ -0,0 +1,260 @@ +# +# 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. +import pytest +import numpy as np +import numpy.testing as npt +import matplotlib.pyplot as plt + +from Basilisk.utilities import SimulationBaseClass +from Basilisk.architecture import messaging +from Basilisk.utilities import macros + +try: + from Basilisk.simulation import mujoco + from Basilisk.simulation import MJLinearTimeInvariantSystem + couldImportMujoco = True +except Exception: + couldImportMujoco = False + + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +@pytest.mark.parametrize("usePythonSubclass", [False, True]) +def test_linearTimeInvariantSystemFirstOrder(usePythonSubclass: bool, + showPlots: bool = False): + """ + Unit test for LinearTimeInvariantSystem: + + - When usePythonSubclass == False: + Uses the C++ SingleActuatorLTI subclass. + + - When usePythonSubclass == True: + Uses a Python subclass of LinearTimeInvariantSystem that overrides + readInput and writeOutput via directors. + + Both implement the same first order system: + + x_dot = -a x + a u, y = x, u = 1 (constant step) + + so that in continuous time: + + x(t) = 1 - exp(-a t) + """ + # Simulation setup + dt = 0.01 # s + tf = 2.0 # s + + # First order dynamics parameter + a = 1.0 # 1/s; time constant tau = 1/a = 1 s + + scSim = SimulationBaseClass.SimBaseClass() + dynProcess = scSim.CreateNewProcess("test") + dynProcess.addTask(scSim.CreateNewTask("test", macros.sec2nano(dt))) + + # Empty MuJoCo scene: just a container for system models + scene = mujoco.MJScene("") + scSim.AddModelToTask("test", scene) + + # Constant input message u = 1.0 + cmdMsg = messaging.SingleActuatorMsg() + cmdPayload = messaging.SingleActuatorMsgPayload() + cmdPayload.input = 1.0 + cmdMsg.write(cmdPayload) + + if usePythonSubclass: + # Python subclass of LinearTimeInvariantSystem + class PyFirstOrderLTI(MJLinearTimeInvariantSystem.LinearTimeInvariantSystem): + def __init__(self): + super().__init__() + # Input/output messaging + self.inMsg = messaging.SingleActuatorMsgReader() + self.outMsg = messaging.SingleActuatorMsg() + # First order system: x_dot = -a x + a u, y = x + A = np.array([[-a]]) + B = np.array([[a]]) + C = np.array([[1.0]]) + D = np.array([[0.0]]) + self.setA(A) + self.setB(B) + self.setC(C) + self.setD(D) + + def getInputSize(self) -> int: + return 1 + + def getOutputSize(self) -> int: + return 1 + + def readInput(self, CurrentSimNanos): + # Return 1x1 vector from SingleActuatorMsg + payload = self.inMsg() + return np.array([[payload.input]]) + + def writeOutput(self, CurrentSimNanos, y): + # Write y[0] to SingleActuatorMsg + payload = messaging.SingleActuatorMsgPayload() + payload.input = y[0][0] + self.outMsg.write(payload, self.moduleID, CurrentSimNanos) + + system = PyFirstOrderLTI() + system.inMsg.subscribeTo(cmdMsg) + system.readInput(0) + + else: + # C++ subclass: SingleActuatorLTI + system = MJLinearTimeInvariantSystem.SingleActuatorLTI() + # First order system: x_dot = -a x + a u, y = x + A = np.array([[-a]]) + B = np.array([[a]]) + C = np.array([[1.0]]) + D = np.array([[0.0]]) + system.setA(A) + system.setB(B) + system.setC(C) + system.setD(D) + system.inMsg.subscribeTo(cmdMsg) + + # Add system to the scene dynamics + scene.AddModelToDynamicsTask(system) + + # Recorder for the system output + outRecorder = system.outMsg.recorder() + scSim.AddModelToTask("test", outRecorder) + + # Initialize and run + scSim.InitializeSimulation() + + # Basic size checks use the LinearTimeInvariantSystem API + assert system.getInputSize() == 1 + assert system.getOutputSize() == 1 + assert system.getStateSize() == 1 + + scSim.ConfigureStopTime(macros.sec2nano(tf)) + scSim.ExecuteSimulation() + + # Extract data + tNanos = outRecorder.times() + yVals = np.asarray(outRecorder.input) # SingleActuatorMsgPayload.input + + # Convert times to seconds for plotting / diagnostics + t = tNanos * 1.0e-9 + + if showPlots: + fig, ax = plt.subplots() + ax.plot(t, yVals, label="y(t)") + ax.set_xlabel("Time [s]") + ax.set_ylabel("Output y") + ax.grid(True) + ax.legend() + plt.show() + + # Continuous time target: x(t) = 1 - exp(-a t); y = x + yTargetFinal = 1.0 - np.exp(-a * tf) + + # Use the last sample + yFinal = float(yVals[-1]) + + # Assert that final value is reasonably close to the continuous solution + # Tolerance is relaxed a bit to allow for integration and discretization error + npt.assert_allclose(yFinal, yTargetFinal, rtol=0.02, atol=1e-2) + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +def test_linearTimeInvariantSystemSecondOrder(showPlots: bool = False): + # Simulation setup + dt = 0.01 # s + tf = 2.0 # s + + scSim = SimulationBaseClass.SimBaseClass() + dynProcess = scSim.CreateNewProcess("test") + dynProcess.addTask(scSim.CreateNewTask("test", macros.sec2nano(dt))) + + # Empty MuJoCo scene: just a container for system models + scene = mujoco.MJScene("") + scSim.AddModelToTask("test", scene) + + # Constant input message u = 1.0 + cmdMsg = messaging.SingleActuatorMsg() + cmdPayload = messaging.SingleActuatorMsgPayload() + cmdPayload.input = 1.0 + cmdMsg.write(cmdPayload) + + # Second order critically damped LTI: wn = 10 rad/s, zeta = 1, k = 1 + wn = 10.0 + zeta = 1.0 + + system = MJLinearTimeInvariantSystem.SingleActuatorLTI() + system.configureSecondOrder(wn, zeta) # default k = 1 + system.inMsg.subscribeTo(cmdMsg) + + # Add system to the scene dynamics + scene.AddModelToDynamicsTask(system) + + # Recorder for the system output + outRecorder = system.outMsg.recorder() + scSim.AddModelToTask("test", outRecorder) + + # Initialize and run + scSim.InitializeSimulation() + + # Basic size checks use the LinearTimeInvariantSystem API + assert system.getInputSize() == 1 + assert system.getOutputSize() == 1 + assert system.getStateSize() == 2 + + scSim.ConfigureStopTime(macros.sec2nano(tf)) + scSim.ExecuteSimulation() + + # Extract data + tNanos = outRecorder.times() + yVals = np.asarray(outRecorder.input) # SingleActuatorMsgPayload.input + + # Convert times to seconds for plotting / diagnostics + t = tNanos * 1.0e-9 + + if showPlots: + fig, ax = plt.subplots() + ax.plot(t, yVals, label="y(t)") + ax.set_xlabel("Time [s]") + ax.set_ylabel("Output y") + ax.grid(True) + ax.legend() + plt.show() + + # Analytic critically damped step response for: + # G(s) = wn^2 / (s^2 + 2*wn*s + wn^2), unit step input + # y(t) = 1 - exp(-wn t) * (1 + wn t) + yAnalytic = 1.0 - np.exp(-wn * t) * (1.0 + wn * t) + + # 1) Final value near 1.0 + npt.assert_allclose(yVals[-1], 1.0, atol=0.005) + + # 2) No overshoot beyond a small numerical margin + assert np.max(yVals) <= 1.01 + + # 3) Shape close to analytic critically damped response + npt.assert_allclose( + yVals, + yAnalytic, + rtol=0.1, + atol=0.05 + ) + + +if __name__ == "__main__": + assert couldImportMujoco + test_linearTimeInvariantSystemFirstOrder(usePythonSubclass=False, showPlots=True) + # test_linearTimeInvariantSystemFirstOrder(usePythonSubclass=True, showPlots=True) + test_linearTimeInvariantSystemSecondOrder(showPlots=True) diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/_UnitTest/test_orbitalElementControl.py b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/_UnitTest/test_orbitalElementControl.py new file mode 100644 index 0000000000..1d4ed846ad --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/_UnitTest/test_orbitalElementControl.py @@ -0,0 +1,182 @@ +# +# 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. + +import pytest +import numpy as np + +from Basilisk.architecture import messaging +from Basilisk.fswAlgorithms import meanOEFeedback +from Basilisk.utilities import SimulationBaseClass +from Basilisk.utilities import macros +from Basilisk.utilities import orbitalMotion + +try: + from Basilisk.simulation import mujoco + from Basilisk.simulation import MJLinearTimeInvariantSystem + couldImportMujoco = True +except Exception: + couldImportMujoco = False + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +@pytest.mark.parametrize("integral", (True, False)) +def test_orbitalElementControl(integral: bool): + """Checks that the MJLinearTimeInvariantSystem.OrbitalElementControl behaves + the same way as meanOEFeedback with J2 = 0. + + Both the proportional and integral gains are tested. + """ + unitTaskName = "unitTask" + unitProcessName = "TestProcess" + + sim = SimulationBaseClass.SimBaseClass() + proc = sim.CreateNewProcess(unitProcessName) + dt = macros.sec2nano(0.1) + proc.addTask(sim.CreateNewTask(unitTaskName, dt)) + + # MuJoCo scene just to host OrbitalElementControl + scene = mujoco.MJScene("") + sim.AddModelToTask(unitTaskName, scene) + + # Gravitational parameter and J2 + mu = orbitalMotion.MU_EARTH * 1e9 # [m^3/s^2] + req = orbitalMotion.REQ_EARTH * 1e3 # [m] + J2_small = 1e-12 # effectively zero J2 for "truth" + + # Chief and deputy classic elements (same as meanOEFeedback test) + oe_c = orbitalMotion.ClassicElements() + oe_c.a = 20000e3 + oe_c.e = 0.1 + oe_c.i = 0.2 + oe_c.Omega = 0.3 + oe_c.omega = 0.4 + oe_c.f = 0.5 + r_c, v_c = orbitalMotion.elem2rv(mu, oe_c) + + oe_d = orbitalMotion.ClassicElements() + oe_d.a = (1 + 0.0006) * 7000e3 + oe_d.e = 0.2 + 0.0005 + oe_d.i = 0.0 + 0.0004 + oe_d.Omega = 0.0 + 0.0003 + oe_d.omega = 0.0 + 0.0002 + oe_d.f = 0.0001 + r_d, v_d = orbitalMotion.elem2rv(mu, oe_d) + + # meanOEFeedback module (to generate truth values) + mean_module = meanOEFeedback.meanOEFeedback() + mean_module.ModelTag = "meanOEFeedback_orbitElemControlTruth" + mean_module.targetDiffOeMean = [0.0] * 6 + mean_module.mu = mu + mean_module.req = req + mean_module.J2 = J2_small + mean_module.K = [ + 1e7, 0.0, 0.0, 0.0, 0.0, 0.0, + 0.0, 1e7, 0.0, 0.0, 0.0, 0.0, + 0.0, 0.0, 1e7, 0.0, 0.0, 0.0, + 0.0, 0.0, 0.0, 1e7, 0.0, 0.0, + 0.0, 0.0, 0.0, 0.0, 1e7, 0.0, + 0.0, 0.0, 0.0, 0.0, 0.0, 1e7 + ] + + sim.AddModelToTask(unitTaskName, mean_module) + + # Chief NavTrans message + chiefNavPayload = messaging.NavTransMsgPayload() + chiefNavPayload.timeTag = 0 + chiefNavPayload.r_BN_N = r_c + chiefNavPayload.v_BN_N = v_c + chiefNavPayload.vehAccumDV = [0.0, 0.0, 0.0] + chiefNavMsg = messaging.NavTransMsg().write(chiefNavPayload) + + # Deputy NavTrans message + deputyNavPayload = messaging.NavTransMsgPayload() + deputyNavPayload.timeTag = 0 + deputyNavPayload.r_BN_N = r_d + deputyNavPayload.v_BN_N = v_d + deputyNavPayload.vehAccumDV = [0.0, 0.0, 0.0] + deputyNavMsg = messaging.NavTransMsg().write(deputyNavPayload) + + mean_module.chiefTransInMsg.subscribeTo(chiefNavMsg) + mean_module.deputyTransInMsg.subscribeTo(deputyNavMsg) + + # Truth recorder + meanLog = mean_module.forceOutMsg.recorder() + sim.AddModelToTask(unitTaskName, meanLog) + + # OrbitalElementControl module under test + ctrl = MJLinearTimeInvariantSystem.OrbitalElementControl() + ctrl.ModelTag = "OrbitalElementControl" + ctrl.mu = mu + + # the proportional gain is applied to the instantaenous difference + # in orbital elements, the integral gain is applied to a continuous + # state whose derivative is the difference in orbital elements. + # The difference in orbital elements is constant, and we run this + # scenario for 10 seconds, so the integral gain / 10 should be equivalent + # to the proportional gain. + K_mat = np.diag([1e7] * 6) + if integral: + ctrl.setIntegralGain(K_mat / 10) + else: + ctrl.setProportionalGain(K_mat) + + # ClassicElements messages for target and current. + # Use the same osculating elements as the meanOEFeedback test, since J2 is tiny. + targetPayload = messaging.ClassicElementsMsgPayload() + targetPayload.a = oe_c.a + targetPayload.e = oe_c.e + targetPayload.i = oe_c.i + targetPayload.Omega = oe_c.Omega + targetPayload.omega = oe_c.omega + targetPayload.f = oe_c.f + + currentPayload = messaging.ClassicElementsMsgPayload() + currentPayload.a = oe_d.a + currentPayload.e = oe_d.e + currentPayload.i = oe_d.i + currentPayload.Omega = oe_d.Omega + currentPayload.omega = oe_d.omega + currentPayload.f = oe_d.f + + targetMsg = messaging.ClassicElementsMsg().write(targetPayload) + currentMsg = messaging.ClassicElementsMsg().write(currentPayload) + + ctrl.targetOEInMsg.subscribeTo(targetMsg) + ctrl.currentOEInMsg.subscribeTo(currentMsg) + + # Add controller under MuJoCo scene + scene.AddModelToDynamicsTask(ctrl) + + # Controller recorder + ctrlLog = ctrl.forceOutMsg.recorder() + sim.AddModelToTask(unitTaskName, ctrlLog) + + # Initialize and run a single step + sim.InitializeSimulation() + sim.ConfigureStopTime(macros.sec2nano(10)) + sim.ExecuteSimulation() + + # Extract outputs + mean_force = meanLog.forceRequestInertial[-1] + ctrl_force = ctrlLog.forceRequestInertial[-1] + + # Compare forces + np.testing.assert_allclose(ctrl_force, mean_force, rtol=1e-6, atol=1e-6) + + +if __name__ == "__main__": + assert couldImportMujoco + test_orbitalElementControl(True) diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/forceAtSiteLTI.cpp b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/forceAtSiteLTI.cpp new file mode 100644 index 0000000000..90ab5a5afa --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/forceAtSiteLTI.cpp @@ -0,0 +1,44 @@ +/* + 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 "forceAtSiteLTI.h" + +Eigen::VectorXd ForceAtSiteLTI::readInput(uint64_t /*CurrentSimNanos*/) +{ + // Read the current payload from the input message + const ForceAtSiteMsgPayload payload = this->inMsg(); + + // Construct a 3 by 1 input vector from the force components + Eigen::VectorXd u(3); + u(0) = payload.force_S[0]; + u(1) = payload.force_S[1]; + u(2) = payload.force_S[2]; + + return u; +} + +void ForceAtSiteLTI::writeOutput(uint64_t CurrentSimNanos, + const Eigen::VectorXd &y) +{ + ForceAtSiteMsgPayload payload{}; + payload.force_S[0] = y(0); + payload.force_S[1] = y(1); + payload.force_S[2] = y(2); + + this->outMsg.write(&payload, this->moduleID, CurrentSimNanos); +} diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/forceAtSiteLTI.h b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/forceAtSiteLTI.h new file mode 100644 index 0000000000..6a37f75490 --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/forceAtSiteLTI.h @@ -0,0 +1,114 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab + + 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 FORCE_AT_SITE_LTI_H +#define FORCE_AT_SITE_LTI_H + +#include "linearTimeInvariantSystem.h" + +#include "architecture/messaging/messaging.h" +#include "architecture/msgPayloadDefC/ForceAtSiteMsgPayload.h" + +/** + * @brief Linear 3D force model based on LinearTimeInvariantSystem. + * + * This class implements a three input, three output LTI system where the + * input and output are carried through ForceAtSiteMsgPayload messages. + * + * Input: + * u = [Fx, Fy, Fz]^T in the site S frame, read from inMsg().force_S. + * + * Output: + * y = [Fx, Fy, Fz]^T in the site S frame, written to outMsg.force_S. + * + * The internal dynamics are defined by the A, B, C, D matrices stored in the + * base class LinearTimeInvariantSystem. + */ +class ForceAtSiteLTI : public LinearTimeInvariantSystem +{ +public: + /** + * @brief Default constructor. + * + * The system matrices must be configured by the user before running + * the simulation, either by direct calls to setA, setB, setC, setD + * or through convenience configuration helpers. + */ + ForceAtSiteLTI() = default; + + /** + * @brief Get the dimension of the input vector. + * + * This model is strictly 3D, so the input dimension is always 3. + * + * @return Number of inputs (always 3). + */ + size_t getInputSize() const override { return 3; } + + /** + * @brief Get the dimension of the output vector. + * + * This model is strictly 3D, so the output dimension is always 3. + * + * @return Number of outputs (always 3). + */ + size_t getOutputSize() const override { return 3; } + + /** + * @brief Read the current input vector from the subscribed input message. + * + * This method constructs a 3 by 1 Eigen::VectorXd whose elements are + * taken from the force_S array of the ForceAtSiteMsgPayload read via + * inMsg(). + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + * (Unused in this implementation.) + * @return Input vector u of size 3. + */ + Eigen::VectorXd readInput(uint64_t CurrentSimNanos) override; + + /** + * @brief Write the current output vector to the output message. + * + * The first three elements of the output vector y are written to the + * force_S array of the ForceAtSiteMsgPayload and sent on outMsg. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + * @param y Output vector of size at least 3. + */ + void writeOutput(uint64_t CurrentSimNanos, + const Eigen::VectorXd &y) override; + +public: + /** + * @brief Output message carrying the 3D force command at the site. + * + * The force_S array is populated from the first three elements of + * the output vector y computed by LinearTimeInvariantSystem. + */ + Message outMsg; + + /** + * @brief Input message read functor providing the 3D force at the site. + * + * The force_S array of the payload is mapped to the input vector u. + */ + ReadFunctor inMsg; +}; + +#endif /* FORCE_AT_SITE_LTI_H */ diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/linearTimeInvariantSystem.cpp b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/linearTimeInvariantSystem.cpp new file mode 100644 index 0000000000..324a4deda5 --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/linearTimeInvariantSystem.cpp @@ -0,0 +1,411 @@ +/* + 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 "linearTimeInvariantSystem.h" +#include + +void LinearTimeInvariantSystem::registerStates(DynParamRegisterer registerer) +{ + const size_t stateSize = this->getStateSize(); + if (stateSize > 0) { + xState = registerer.registerState(stateSize, 1, "x"); + xState->setState(Eigen::VectorXd::Constant(stateSize, 0.0)); // default to x0 = 0 + } +} + +void LinearTimeInvariantSystem::UpdateState(uint64_t CurrentSimNanos) +{ + const size_t stateSize = this->getStateSize(); + const size_t inputSize = this->getInputSize(); + const size_t outputSize = this->getOutputSize(); + + Eigen::VectorXd u; + if (inputSize > 0) { + u = Eigen::VectorXd::Zero(static_cast(inputSize)); + u = readInput(CurrentSimNanos); + + if (static_cast(u.size()) != inputSize) { + bskLogger.bskLog( + BSK_ERROR, + "LinearTimeInvariantSystem::UpdateState: readInput returned vector with incorrect size."); + } + } + + Eigen::VectorXd y = Eigen::VectorXd::Zero(static_cast(outputSize)); + + if (xState && stateSize > 0) { + Eigen::VectorXd x = xState->getState(); + + if (static_cast(x.size()) != stateSize) { + bskLogger.bskLog( + BSK_ERROR, + "LinearTimeInvariantSystem::UpdateState: stored state has incorrect size."); + } + + Eigen::VectorXd dx = Eigen::VectorXd::Zero(static_cast(stateSize)); + if (A.cols() > 0) { + dx += A * x; + } + if (B.cols() > 0 && inputSize > 0) { + dx += B * u; + } + xState->setDerivative(dx); + + if (C.cols() > 0) { + y += C * x; + } + } + + if (D.size() > 0 && inputSize > 0) { + y += D * u; + } + + writeOutput(CurrentSimNanos, y); +} + +void LinearTimeInvariantSystem::Reset(uint64_t /*CurrentSimNanos*/) +{ + const size_t stateSize = this->getStateSize(); + const size_t inputSize = this->getInputSize(); + const size_t outputSize = this->getOutputSize(); + + // ======================= + // Matrix A (state x state) + // ======================= + if (A.size() > 0) { + const size_t A_rows = static_cast(A.rows()); + const size_t A_cols = static_cast(A.cols()); + + if (A_rows != stateSize || A_cols != stateSize) { + bskLogger.bskLog( + BSK_ERROR, + ( + std::string("LinearTimeInvariantSystem::Reset: matrix A has inconsistent dimensions. ") + + "Expected size [" + + std::to_string(stateSize) + " x " + std::to_string(stateSize) + + "] based on getStateSize(), but received [" + + std::to_string(A_rows) + " x " + std::to_string(A_cols) + + "]." + ).c_str() + ); + } + } + + // ======================= + // Matrix B (state x input) + // ======================= + if (B.size() > 0) { + const size_t B_rows = static_cast(B.rows()); + const size_t B_cols = static_cast(B.cols()); + + if (B_rows != stateSize || B_cols != inputSize) { + bskLogger.bskLog( + BSK_ERROR, + ( + std::string("LinearTimeInvariantSystem::Reset: matrix B has inconsistent dimensions. ") + + "Expected size [" + + std::to_string(stateSize) + " x " + std::to_string(inputSize) + + "] based on getStateSize() and getInputSize(), but received [" + + std::to_string(B_rows) + " x " + std::to_string(B_cols) + + "]." + ).c_str() + ); + } + } + + // ========================== + // Matrix C (output x state) + // ========================== + if (C.size() > 0) { + const size_t C_rows = static_cast(C.rows()); + const size_t C_cols = static_cast(C.cols()); + + if (C_rows != outputSize || C_cols != stateSize) { + bskLogger.bskLog( + BSK_ERROR, + ( + std::string("LinearTimeInvariantSystem::Reset: matrix C has inconsistent dimensions. ") + + "Expected size [" + + std::to_string(outputSize) + " x " + std::to_string(stateSize) + + "] based on getOutputSize() and getStateSize(), but received [" + + std::to_string(C_rows) + " x " + std::to_string(C_cols) + + "]." + ).c_str() + ); + } + } + + // ========================= + // Matrix D (output x input) + // ========================= + if (D.size() > 0) { + const size_t D_rows = static_cast(D.rows()); + const size_t D_cols = static_cast(D.cols()); + + if (D_rows != outputSize || D_cols != inputSize) { + bskLogger.bskLog( + BSK_ERROR, + ( + std::string("LinearTimeInvariantSystem::Reset: matrix D has inconsistent dimensions. ") + + "Expected size [" + + std::to_string(outputSize) + " x " + std::to_string(inputSize) + + "] based on getOutputSize() and getInputSize(), but received [" + + std::to_string(D_rows) + " x " + std::to_string(D_cols) + + "]." + ).c_str() + ); + } + } +} + +const Eigen::MatrixXd &LinearTimeInvariantSystem::getA() const +{ + return A; +} + +void LinearTimeInvariantSystem::setA(const Eigen::MatrixXd &Ain) +{ + A = Ain; +} + +const Eigen::MatrixXd &LinearTimeInvariantSystem::getB() const +{ + return B; +} + +void LinearTimeInvariantSystem::setB(const Eigen::MatrixXd &Bin) +{ + B = Bin; +} + +const Eigen::MatrixXd &LinearTimeInvariantSystem::getC() const +{ + return C; +} + +void LinearTimeInvariantSystem::setC(const Eigen::MatrixXd &Cin) +{ + C = Cin; +} + +const Eigen::MatrixXd &LinearTimeInvariantSystem::getD() const +{ + return D; +} + +void LinearTimeInvariantSystem::setD(const Eigen::MatrixXd &Din) +{ + D = Din; +} + +void LinearTimeInvariantSystem::configureSecondOrder(double wn, double zeta, double k) +{ + if (wn <= 0.0) { + bskLogger.bskLog( + BSK_ERROR, + "configureSecondOrder: natural frequency wn must be positive."); + return; + } + if (zeta < 0.0) { + bskLogger.bskLog( + BSK_ERROR, + "configureSecondOrder: damping ratio zeta must be non negative."); + return; + } + + A.resize(2, 2); + B.resize(2, 1); + C.resize(1, 2); + D.resize(1, 1); + + const double wn2 = wn * wn; + + A << 0.0 , 1.0, + -wn2, -2.0 * zeta * wn; + + B << 0.0, + k * wn2; + + C << 1.0, 0.0; + + D << 0.0; +} + +void LinearTimeInvariantSystem::configureSecondOrder(const Eigen::VectorXd &wn, + const Eigen::VectorXd &zeta, + const Eigen::VectorXd &k) +{ + const Eigen::Index n = wn.size(); + + if (n <= 0) { + bskLogger.bskLog( + BSK_ERROR, + "LinearTimeInvariantSystem::configureSecondOrder(MIMO): input vectors must be non empty." + ); + return; + } + + if (zeta.size() != n || k.size() != n) { + bskLogger.bskLog( + BSK_ERROR, + ( + std::string("LinearTimeInvariantSystem::configureSecondOrder(MIMO): size mismatch. ") + + "Expected wn, zeta, k to all have length " + std::to_string(static_cast(n)) + ". " + "Received lengths wn=" + std::to_string(static_cast(wn.size())) + + ", zeta=" + std::to_string(static_cast(zeta.size())) + + ", k=" + std::to_string(static_cast(k.size())) + "." + ).c_str() + ); + return; + } + + // Validate entries + for (Eigen::Index i = 0; i < n; ++i) { + const double wn_i = wn(i); + const double zeta_i = zeta(i); + + if (wn_i <= 0.0) { + bskLogger.bskLog( + BSK_ERROR, + ( + std::string("LinearTimeInvariantSystem::configureSecondOrder(MIMO): wn(") + + std::to_string(static_cast(i)) + + ") must be positive, but is " + std::to_string(wn_i) + "." + ).c_str() + ); + return; + } + + if (zeta_i < 0.0) { + bskLogger.bskLog( + BSK_ERROR, + ( + std::string("LinearTimeInvariantSystem::configureSecondOrder(MIMO): zeta(") + + std::to_string(static_cast(i)) + + ") must be non negative, but is " + std::to_string(zeta_i) + "." + ).c_str() + ); + return; + } + } + + // Allocate matrices for MIMO second-order system: n channels, each 2 states + const Eigen::Index stateSize = 2 * n; + const Eigen::Index inputSize = n; + const Eigen::Index outputSize = n; + + A.setZero(stateSize, stateSize); + B.setZero(stateSize, inputSize); + C.setZero(outputSize, stateSize); + D.setZero(outputSize, inputSize); + + for (Eigen::Index i = 0; i < n; ++i) { + const double wn_i = wn(i); + const double zeta_i = zeta(i); + const double k_i = k(i); + + const double wn2 = wn_i * wn_i; + + const Eigen::Index rowPos = 2 * i; // position state index + const Eigen::Index rowVel = 2 * i + 1; // velocity state index + + // A block for channel i: + // [ 0 1 ] + // [ -wn_i^2 -2*zeta_i*wn_i ] + A(rowPos, rowVel) = 1.0; + A(rowVel, rowPos) = -wn2; + A(rowVel, rowVel) = -2.0 * zeta_i * wn_i; + + // B column i: [0, k_i * wn_i^2]^T in rows (rowPos, rowVel) + B(rowVel, i) = k_i * wn2; + + // C row i: [1 0] on the position state of channel i + C(i, rowPos) = 1.0; + + // D(i, i) remains zero (direct feedthrough = 0) + } +} + +Eigen::VectorXd +LinearTimeInvariantSystem::getX() +{ + if (!xState) { + bskLogger.bskLog( + BSK_ERROR, + "LinearTimeInvariantSystem::getX: state has not been registered."); + return Eigen::VectorXd(); + } + return xState->getState(); +} + +void LinearTimeInvariantSystem::setX(const Eigen::VectorXd &xin) +{ + if (!xState) { + bskLogger.bskLog( + BSK_ERROR, + "LinearTimeInvariantSystem::setX: state has not been registered."); + return; + } + + const size_t stateSize = this->getStateSize(); + if (static_cast(xin.size()) != stateSize) { + bskLogger.bskLog( + BSK_ERROR, + "LinearTimeInvariantSystem::setX: input state vector has incorrect size."); + return; + } + + xState->setState(xin); +} + +size_t LinearTimeInvariantSystem::getInputSize() const +{ + if (B.cols() > 0) { + return static_cast(B.cols()); + } + if (D.cols() > 0) { + return static_cast(D.cols()); + } + return 0; +} + +size_t LinearTimeInvariantSystem::getStateSize() const +{ + if (A.cols() > 0) { + return static_cast(A.cols()); + } + if (B.rows() > 0) { + return static_cast(B.rows()); + } + if (C.cols() > 0) { + return static_cast(C.cols()); + } + return 0; +} + +size_t LinearTimeInvariantSystem::getOutputSize() const +{ + if (C.rows() > 0) { + return static_cast(C.rows()); + } + if (D.rows() > 0) { + return static_cast(D.rows()); + } + return 0; +} diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/linearTimeInvariantSystem.h b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/linearTimeInvariantSystem.h new file mode 100644 index 0000000000..97f5d33746 --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/linearTimeInvariantSystem.h @@ -0,0 +1,316 @@ +/* + 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 LINEAR_TIME_INVARIANT_SYSTEM_H +#define LINEAR_TIME_INVARIANT_SYSTEM_H + +#include +#include + +#include + +#include "simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.h" +#include "architecture/utilities/bskLogging.h" + +/** + * @brief Linear time invariant system of the form + * x_dot = A x + B u, y = C x + D u. + * + * Subclasses must implement the input and output mapping methods. + */ +class LinearTimeInvariantSystem : public StatefulSysModel { +public: + /** @brief Default constructor. */ + LinearTimeInvariantSystem() = default; + + /** @name Framework interface @{ */ + + /** + * @brief Register the state vector with the dynamics framework. + * + * If the state dimension is greater than zero, this function allocates + * a state block of size getStateSize() and initializes it to zero. + * + * @param registerer State registration helper provided by the framework. + */ + void registerStates(DynParamRegisterer registerer) override; + + /** + * @brief Compute the state derivative and system output at the current time. + * + * This method reads the current input vector using readInput, computes + * the state derivative + * x_dot = A x + B u and the output y = C x + D u, and writes the output + * using writeOutput. If no state is registered, only the direct feedthrough + * term y = D u is evaluated if applicable. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + */ + void UpdateState(uint64_t CurrentSimNanos) override; + + /** + * @brief Perform consistency checks and one time initialization. + * + * This method validates that the dimensions of the system matrices A, B, + * C, and D are consistent with getStateSize, getInputSize, and + * getOutputSize. Any inconsistency is reported through the logger. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + */ + virtual void Reset(uint64_t CurrentSimNanos) override; + + /** @} */ + + /** @name Parameter getters and setters + * Accessors for the system matrices A, B, C, and D. + * @{ */ + + /** + * @brief Get the state matrix A. + * + * @return Constant reference to the A matrix. + */ + const Eigen::MatrixXd &getA() const; + + /** + * @brief Set the state matrix A. + * + * This matrix defines the homogeneous term in the state equation + * x_dot = A x + B u. + * + * @param Ain New A matrix. + */ + void setA(const Eigen::MatrixXd &Ain); + + /** + * @brief Get the input matrix B. + * + * @return Constant reference to the B matrix. + */ + const Eigen::MatrixXd &getB() const; + + /** + * @brief Set the input matrix B. + * + * This matrix defines the contribution of the input vector u to the + * state equation x_dot = A x + B u. + * + * @param Bin New B matrix. + */ + void setB(const Eigen::MatrixXd &Bin); + + /** + * @brief Get the output matrix C. + * + * @return Constant reference to the C matrix. + */ + const Eigen::MatrixXd &getC() const; + + /** + * @brief Set the output matrix C. + * + * This matrix maps the state vector x to the output equation + * y = C x + D u. + * + * @param Cin New C matrix. + */ + void setC(const Eigen::MatrixXd &Cin); + + /** + * @brief Get the feedthrough matrix D. + * + * @return Constant reference to the D matrix. + */ + const Eigen::MatrixXd &getD() const; + + /** + * @brief Set the feedthrough matrix D. + * + * This matrix maps the input vector u directly to the output equation + * y = C x + D u. + * + * @param Din New D matrix. + */ + void setD(const Eigen::MatrixXd &Din); + + /** + * @brief Configure the A, B, C, D matrices as a standard SISO second order model. + * + * The resulting transfer function is: + * G(s) = k * wn^2 / (s^2 + 2*zeta*wn*s + wn^2) + * + * State space realization: + * x = [position, velocity]^T + * x_dot = [ 0 1 ] x + [0 ] u + * [ -wn^2 -2*zeta*wn ] [k * wn^2] + * y = [1 0] x + * + * @param wn Natural frequency in rad/s. Must be positive. + * @param zeta Damping ratio (dimensionless). Must be non negative. + * @param k Static gain. Default is unity. + */ + void configureSecondOrder(double wn, double zeta, double k = 1.0); + + /** + * @brief Configure A, B, C, D as a decoupled MIMO second order model. + * + * Each element of the input vectors defines one independent SISO channel: + * + * G_i(s) = k_i * wn_i^2 / (s^2 + 2*zeta_i*wn_i*s + wn_i^2) + * + * The full state vector is + * x = [pos_0, vel_0, pos_1, vel_1, ..., pos_{n-1}, vel_{n-1}]^T + * + * For n channels, the resulting dimensions are: + * A: (2n x 2n), B: (2n x n), C: (n x 2n), D: (n x n). + * + * @param wn Vector of natural frequencies in rad/s. Each wn(i) must be positive. + * @param zeta Vector of damping ratios (dimensionless). Each zeta(i) must be non negative. + * @param k Vector of static gains. Must have the same length as wn and zeta. + */ + void configureSecondOrder(const Eigen::VectorXd &wn, + const Eigen::VectorXd &zeta, + const Eigen::VectorXd &k); + + /** @} */ + + /** @name State access + * Accessors for the internal state vector. + * @{ */ + + /** + * @brief Get the current state vector x. + * + * If the state has not been registered yet, an empty vector is returned + * and an error is logged. + * + * @return Copy of the current state vector. + */ + Eigen::VectorXd getX(); + + /** + * @brief Set the current state vector x. + * + * If the state has been registered and the supplied vector has the same + * dimension as the stored state, the state is updated. Otherwise an + * error is logged and the call is ignored. + * + * @param xin New state vector. + */ + void setX(const Eigen::VectorXd &xin); + + /** @} */ + + /** + * @brief Get the dimension of the input vector u. + * + * The size is inferred from the B or D matrices when they are set. + * + * @return Size of the input vector. + */ + virtual size_t getInputSize() const; + + /** + * @brief Get the dimension of the state vector x. + * + * The size is inferred from the A, B, or C matrices when they are set. + * + * @return Size of the state vector. + */ + virtual size_t getStateSize() const; + + /** + * @brief Get the dimension of the output vector y. + * + * The size is inferred from the C or D matrices when they are set. + * + * @return Size of the output vector. + */ + virtual size_t getOutputSize() const; + + /** + * @brief Logger used to report errors and informational messages. + */ + BSKLogger bskLogger; + +protected: + /** + * @brief Read the current input vector u from the simulation. + * + * Subclasses must implement this method to construct the input vector + * of size getInputSize() from the relevant input messages at the given + * simulation time. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + * + * @return Input vector u. + */ + virtual Eigen::VectorXd readInput(uint64_t CurrentSimNanos) = 0; + + /** + * @brief Write the current output vector y to the simulation. + * + * Subclasses must implement this method to map the output vector + * of size getOutputSize() to the appropriate output messages at the + * given simulation time. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + * @param y Output vector to write. + */ + virtual void writeOutput(uint64_t CurrentSimNanos, + const Eigen::VectorXd &y) = 0; + +private: + /** + * @brief Pointer to the registered state data for the state vector x. + * + * This pointer is set during registerStates and is used to read and + * update the state and its derivative. + */ + StateData *xState = nullptr; + + /** + * @brief State matrix A. + * + * Defines the homogeneous dynamics x_dot = A x + B u. + */ + Eigen::MatrixXd A; + + /** + * @brief Input matrix B. + * + * Defines the input contribution to the dynamics x_dot = A x + B u. + */ + Eigen::MatrixXd B; + + /** + * @brief Output matrix C. + * + * Maps the state vector to the output y = C x + D u. + */ + Eigen::MatrixXd C; + + /** + * @brief Feedthrough matrix D. + * + * Maps the input vector directly to the output y = C x + D u. + */ + Eigen::MatrixXd D; +}; + +#endif /* LINEAR_TIME_INVARIANT_SYSTEM_H */ diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/linearTimeInvariantSystem.i b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/linearTimeInvariantSystem.i new file mode 100644 index 0000000000..84defe2d3f --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/linearTimeInvariantSystem.i @@ -0,0 +1,98 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab + + 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(directors="1",threads="1") MJLinearTimeInvariantSystem + +%include "architecture/utilities/bskException.swg" +%default_bsk_exception(); + +%{ + // C++ includes visible to the generated wrapper + #include "linearTimeInvariantSystem.h" + #include "forceAtSiteLTI.h" + #include "singleActuatorLTI.h" + #include "orbitalElementControl.h" +%} + +/* Python helpers consistent with your other modules */ +%pythoncode %{ +from Basilisk.architecture.swig_common_model import * +%} + +/* Common SWIG helpers */ +%include "swig_eigen.i" +%include "std_string.i" +%include "exception.i" + +/* Basilisk system-model base and helpers */ +%import "simulation/mujocoDynamics/_GeneralModuleFiles/pyStatefulSysModel.i" + +/* ============================ + Base class: LinearTimeInvariantSystem + ============================ */ + +/* Enable directors so Python subclasses can override virtual methods + (readInput, writeOutput, etc.). */ +%feature("director") LinearTimeInvariantSystem; + +%rename("_LinearTimeInvariantSystem") LinearTimeInvariantSystem; +%include "linearTimeInvariantSystem.h" + +/* ============================ + SingleActuatorLTI + ============================ */ + +%ignore SingleActuatorLTI::readInput; +%ignore SingleActuatorLTI::writeOutput; +%include "singleActuatorLTI.h" + +%include "architecture/msgPayloadDefC/SingleActuatorMsgPayload.h" +struct SingleActuatorMsgPayload_C; + +/* ============================ + ForceAtSiteLTI + ============================ */ + +%ignore ForceAtSiteLTI::readInput; +%ignore ForceAtSiteLTI::writeOutput; +%include "forceAtSiteLTI.h" + +%include "architecture/msgPayloadDefC/ForceAtSiteMsgPayload.h" +struct ForceAtSiteMsgPayload_C; + +/* ============================ + OrbitalElementControl + ============================ */ + +%ignore OrbitalElementControl::readInput; +%ignore OrbitalElementControl::writeOutput; +%include "orbitalElementControl.h" + +%include "architecture/msgPayloadDefC/ClassicElementsMsgPayload.h" +struct ClassicElementsMsgPayload_C; + +%include "architecture/msgPayloadDefC/CmdForceInertialMsgPayload.h" +struct CmdForceInertialMsgPayload_C; + +/* Final Python-side wrapper for the base LTI class */ +%pythoncode %{ +from Basilisk.architecture.sysModel import SysModelMixin + +class LinearTimeInvariantSystem(SysModelMixin, _LinearTimeInvariantSystem): + """Python wrapper for the C++ LinearTimeInvariantSystem.""" +%} diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/orbitalElementControl.cpp b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/orbitalElementControl.cpp new file mode 100644 index 0000000000..3162adcbc1 --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/orbitalElementControl.cpp @@ -0,0 +1,240 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab + + 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 "orbitalElementControl.h" + +#include + +#include "architecture/utilities/orbitalMotion.h" + +namespace { + +double angleWarp(double angle) +{ + const double width = 2.0 * M_PI; + double adjusted = angle; + while (adjusted > M_PI) { + adjusted -= width; + } + while (adjusted < -M_PI) { + adjusted += width; + } + return adjusted; +} + +/** + * @brief Compute classical Gauss control matrix B (6x3) for given elements. + * + * See: Schaub and Junkins, Analytical Mechanics of Space Systems. + * + * Rows: da, de, di, dOmega, domega, dM + * Cols: f_R, f_T, f_N in Hill or LVLH frame. + */ +Eigen::Matrix calc_B_cl(double mu, const ClassicElements &oe_cl) +{ + Eigen::Matrix B; + B.setZero(); + + const double a = oe_cl.a; + const double e = oe_cl.e; + const double i = oe_cl.i; + const double omega = oe_cl.omega; + const double f = oe_cl.f; + + const double theta = omega + f; + const double eta = std::sqrt(1.0 - e * e); + const double b = a * eta; + const double n = std::sqrt(mu / std::pow(a, 3)); + const double h = n * a * b; + const double p = a * (1.0 - e * e); + const double r = p / (1.0 + e * std::cos(f)); + + // B matrix (rows: da, de, di, dOmega, domega, dM; cols: f_R, f_T, f_N) + B(0, 0) = 2.0 * std::pow(a, 2) * e * std::sin(f) / h / a; + B(0, 1) = 2.0 * std::pow(a, 2) * p / (h * r) / a; + B(0, 2) = 0.0; + + B(1, 0) = p * std::sin(f) / h; + B(1, 1) = ((p + r) * std::cos(f) + r * e) / h; + B(1, 2) = 0.0; + + B(2, 0) = 0.0; + B(2, 1) = 0.0; + B(2, 2) = r * std::cos(theta) / h; + + B(3, 0) = 0.0; + B(3, 1) = 0.0; + B(3, 2) = r * std::sin(theta) / (h * std::sin(i)); + + B(4, 0) = -p * std::cos(f) / (h * e); + B(4, 1) = (p + r) * std::sin(f) / (h * e); + B(4, 2) = -r * std::sin(theta) * std::cos(i) / (h * std::sin(i)); + + B(5, 0) = eta * (p * std::cos(f) - 2.0 * r * e) / (h * e); + B(5, 1) = -eta * (p + r) * std::sin(f) / (h * e); + B(5, 2) = 0.0; + + return B; +} + +} // anonymous namespace + +void +OrbitalElementControl::Reset(uint64_t CurrentSimNanos) +{ + // Base class checks A/B/C/D consistency. + LinearTimeInvariantSystem::Reset(CurrentSimNanos); + + if (this->mu <= 0.0) { + bskLogger.bskLog( + BSK_ERROR, + "OrbitalElementControl::Reset: mu must be set to a positive value."); + } + + if (!this->targetOEInMsg.isLinked()) { + bskLogger.bskLog( + BSK_ERROR, + "OrbitalElementControl::Reset: targetOEInMsg was not connected."); + } + if (!this->currentOEInMsg.isLinked()) { + bskLogger.bskLog( + BSK_ERROR, + "OrbitalElementControl::Reset: currentOEInMsg was not connected."); + } +} + +void +OrbitalElementControl::setProportionalGain(const Eigen::MatrixXd& K) +{ + const size_t K_rows = static_cast(K.rows()); + const size_t K_cols = static_cast(K.cols()); + + if (K_rows != 6 || K_cols != 6) { + bskLogger.bskLog( + BSK_ERROR, + ( + std::string("OrbitalElementControl::setProportionalGain: gain K has inconsistent dimensions. ") + + "Expected size [6x6], but received [" + + std::to_string(K_rows) + " x " + std::to_string(K_cols) + + "]." + ).c_str() + ); + } + + this->setD(K); +} + +void +OrbitalElementControl::setIntegralGain(const Eigen::MatrixXd& K) +{ + const size_t K_rows = static_cast(K.rows()); + const size_t K_cols = static_cast(K.cols()); + + if (K_rows != 6 || K_cols != 6) { + bskLogger.bskLog( + BSK_ERROR, + ( + std::string("OrbitalElementControl::setIntegralGain: gain K has inconsistent dimensions. ") + + "Expected size [6x6], but received [" + + std::to_string(K_rows) + " x " + std::to_string(K_cols) + + "]." + ).c_str() + ); + } + + this->setB(Eigen::MatrixXd::Identity(6, 6)); + this->setC(K); +} + +Eigen::VectorXd +OrbitalElementControl::readInput(uint64_t /*CurrentSimNanos*/) +{ + const size_t n = this->getInputSize(); + Eigen::VectorXd result = Eigen::VectorXd::Zero(static_cast(n)); + + const ClassicElementsMsgPayload& targetOE = this->targetOEInMsg(); + const ClassicElementsMsgPayload& currentOE = this->currentOEInMsg(); + + // Relative OE error: [da/a, de, di, dOmega, domega, dM] + if (targetOE.a != 0.0) { + result(0) = (currentOE.a - targetOE.a) / targetOE.a; + } else { + result(0) = currentOE.a - targetOE.a; + } + + result(1) = currentOE.e - targetOE.e; + result(2) = angleWarp(currentOE.i - targetOE.i); + result(3) = angleWarp(currentOE.Omega - targetOE.Omega); + result(4) = angleWarp(currentOE.omega - targetOE.omega); + + const double currentE = f2E(currentOE.f, currentOE.e); + const double currentM = E2M(currentE, currentOE.e); + const double targetE = f2E(targetOE.f, targetOE.e); + const double targetM = E2M(targetE, targetOE.e); + + result(5) = angleWarp(currentM - targetM); + + return result; +} + +void +OrbitalElementControl::writeOutput(uint64_t CurrentSimNanos, const Eigen::VectorXd &y) +{ + // Copy current classical elements into ClassicElements struct. + const ClassicElementsMsgPayload& inputElement = this->currentOEInMsg(); + ClassicElements elements; + elements.a = inputElement.a; + elements.e = inputElement.e; + elements.i = inputElement.i; + elements.Omega = inputElement.Omega; + elements.omega = inputElement.omega; + elements.f = inputElement.f; + + // Compute Gauss control matrix B (6x3) that maps force_LVLH to d(oe). + Eigen::Matrix B = calc_B_cl(this->mu, elements); + Eigen::Vector3d force_LVLH = -B.transpose() * y; + + // Current inertial position and velocity from elements. + double r_N_arr[3]; + double v_N_arr[3]; + elem2rv(this->mu, &elements, r_N_arr, v_N_arr); + + auto r_N = Eigen::Map(r_N_arr); + auto v_N = Eigen::Map(v_N_arr); + + // Build LVLH DCM H_N with rows R, T, N. + Eigen::Vector3d ir = r_N.normalized(); + Eigen::Vector3d hVec = r_N.cross(v_N); + Eigen::Vector3d ih = hVec.normalized(); + Eigen::Vector3d it = ih.cross(ir); + + Eigen::Matrix3d H_N; + H_N.row(0) = ir.transpose(); // radial + H_N.row(1) = it.transpose(); // along track + H_N.row(2) = ih.transpose(); // cross track + + // Transform LVLH force to inertial frame. + Eigen::Vector3d force_N = H_N.transpose() * force_LVLH; + + CmdForceInertialMsgPayload payload; + payload.forceRequestInertial[0] = force_N(0); + payload.forceRequestInertial[1] = force_N(1); + payload.forceRequestInertial[2] = force_N(2); + + this->forceOutMsg.write(&payload, moduleID, CurrentSimNanos); +} diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/orbitalElementControl.h b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/orbitalElementControl.h new file mode 100644 index 0000000000..66496662f9 --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/orbitalElementControl.h @@ -0,0 +1,224 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab + + 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 ORBITAL_ELEMENT_CONTROL_H +#define ORBITAL_ELEMENT_CONTROL_H + +#include "linearTimeInvariantSystem.h" + +#include "architecture/messaging/messaging.h" +#include "architecture/msgPayloadDefC/ClassicElementsMsgPayload.h" +#include "architecture/msgPayloadDefC/CmdForceInertialMsgPayload.h" + +/** + * @class OrbitalElementControl + * @brief LTI based orbital element feedback controller that outputs an + * inertial force command. + * + * This class derives from LinearTimeInvariantSystem and implements a + * controller in classical orbital element coordinates. The LTI input + * is a 6 element vector of orbital element errors + * + * u = [da/a, de, di, dOmega, domega, dM]^T + * + * constructed from target and current classical elements. The LTI + * output y is also a 6 element vector in the same orbital element + * space. + * + * The module then uses the Gauss planetary control matrix B(oe) to map + * y to a Hill frame force, and finally transforms this force into the + * inertial frame using the current orbit geometry. The inertial force + * is written to a CmdForceInertialMsgPayload. + * + * The internal state dimension and A, B, C, D matrices are configured + * through the LinearTimeInvariantSystem interface and are not fixed by + * this class. + */ +class OrbitalElementControl : public LinearTimeInvariantSystem +{ +public: + /** + * @brief Default constructor. + * + * Constructs an OrbitalElementControl instance with mu set to zero + * and unconnected input and output messages. The user must set mu, + * configure the LTI matrices, and connect the input and output + * messages before the module is used in a simulation. + */ + OrbitalElementControl() = default; + + /** + * @brief Set the proportional gain matrix for the orbital element controller. + * + * This method configures the proportional feedback gain used in the + * internal LinearTimeInvariantSystem realization of the control law. + * + * The proportional term maps the instantaneous 6x1 orbital-element error + * vector into a 6x1 control output before it is transformed by the Gauss + * planetary-equations B-matrix in writeOutput(). + * + * Internally, this method assigns the gain matrix K to the D matrix of the + * underlying LTI system: + * + * y = C x + D u + * + * where: + * - u is the six-element orbital-element error vector + * - D = K implements the proportional action + * + * @param K 6x6 proportional gain matrix. + */ + void setProportionalGain(const Eigen::MatrixXd& K); + + + /** + * @brief Set the integral gain matrix for the orbital element controller. + * + * This method configures the integral feedback gain for the controller, + * implemented using the state dynamics of the underlying + * LinearTimeInvariantSystem. The integral term accumulates the + * orbital-element error over time and contributes to the output force. + * + * The structure set here is: + * + * ẋ = B u with B = I [6x6] + * y = C x + D u with C = K_I + * + * This realizes: + * + * x(t) = ∫ u(t) dt (integrator state) + * y = K_I x (integral feedback) + * + * @param K 6x6 integral gain matrix. + */ + void setIntegralGain(const Eigen::MatrixXd& K); + + + /** + * @brief Reset the module and validate configuration. + * + * This method first calls LinearTimeInvariantSystem::Reset to + * validate A, B, C, and D against the sizes returned by + * getStateSize, getInputSize, and getOutputSize. It then checks + * that the gravitational parameter mu is positive and that both + * targetOEInMsg and currentOEInMsg are linked. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + */ + void Reset(uint64_t CurrentSimNanos) override; + + /** + * @brief Size of the LTI input vector. + * + * The input to the LTI system is the 6 element orbital element + * error vector + * + * u = [da/a, de, di, dOmega, domega, dM]^T. + * + * @return Number of input elements, always 6. + */ + size_t getInputSize() const override { return 6; } + + /** + * @brief Size of the LTI output vector. + * + * The LTI output y is a 6 element vector in orbital element space. + * It is not directly the applied force. Instead, writeOutput uses y + * together with the Gauss control matrix B to compute a Hill frame + * force, then transforms that force into the inertial frame and + * writes it to the forceOutMsg. + * + * @return Number of output elements, always 6. + */ + size_t getOutputSize() const override { return 6; } + + /** + * @brief Read the current input vector for the LTI system. + * + * This method reads the target and current classical orbital element + * messages and constructs the 6 element orbital element error vector + * + * u = [da/a, de, di, dOmega, domega, dM]^T, + * + * where: + * - da/a is the relative semi major axis error + * - de is the eccentricity error + * - di is the inclination error + * - dOmega is the RAAN error + * - domega is the argument of periapsis error + * - dM is the mean anomaly error + * + * Angular differences are wrapped into [-pi, pi]. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + * @return Input vector u as an Eigen::VectorXd of length + * getInputSize(). + */ + Eigen::VectorXd readInput(uint64_t CurrentSimNanos) override; + + /** + * @brief Map the LTI output vector to an inertial force command and + * write the output message. + * + * The argument y is the output of the LTI system in orbital element + * space, with dimension getOutputSize() = 6. This method performs + * the following steps: + * + * 1. Read the current classical orbital elements from + * currentOEInMsg and construct a ClassicElements struct. + * 2. Compute the 6 by 3 Gauss control matrix B(oe) that maps + * Hill frame force to orbital element rates. + * 3. Interpret y as an orbital elements feedback vector and + * compute a Hill frame force, for example via + * f_H = B^T * y. + * 4. Convert the current elements to inertial position and + * velocity and compute the Hill to inertial direction cosine + * matrix. + * 5. Rotate the Hill frame force into the inertial frame and + * write it into forceOutMsg as a CmdForceInertialMsgPayload. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + * @param y LTI output vector in orbital element space with length + * getOutputSize(). + */ + void writeOutput(uint64_t CurrentSimNanos, + const Eigen::VectorXd &y) override; + +public: + /** + * @brief Gravitational parameter mu [m^3/s^2]. + * + * This parameter must be set to the central body's gravitational + * parameter before use. The Reset method checks that mu is positive + * and reports an error if it is not. + */ + double mu = 0.0; + + /** @brief Target classical orbital elements input. */ + ReadFunctor targetOEInMsg; + + /** + * @brief Current classical orbital elements input. */ + ReadFunctor currentOEInMsg; + + /** + * @brief Inertial force command output message. */ + Message forceOutMsg; +}; + +#endif /* ORBITAL_ELEMENT_CONTROL_H */ diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/singleActuatorLTI.cpp b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/singleActuatorLTI.cpp new file mode 100644 index 0000000000..963134bf0a --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/singleActuatorLTI.cpp @@ -0,0 +1,12 @@ +#include "singleActuatorLTI.h" + +Eigen::VectorXd SingleActuatorLTI::readInput(uint64_t CurrentSimNanos) +{ + return Eigen::VectorXd::Constant(1, this->inMsg().input); +} + +void SingleActuatorLTI::writeOutput(uint64_t CurrentSimNanos, const Eigen::VectorXd& y) +{ + auto payload = SingleActuatorMsgPayload{y(0)}; + this->outMsg.write(&payload, moduleID, CurrentSimNanos); +} diff --git a/src/simulation/mujocoDynamics/linearTimeInvariantSystem/singleActuatorLTI.h b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/singleActuatorLTI.h new file mode 100644 index 0000000000..89942b67ab --- /dev/null +++ b/src/simulation/mujocoDynamics/linearTimeInvariantSystem/singleActuatorLTI.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 SINGLE_ACTUATOR_LTI_H +#define SINGLE_ACTUATOR_LTI_H + +#include "linearTimeInvariantSystem.h" + +#include "architecture/messaging/messaging.h" +#include "architecture/msgPayloadDefC/SingleActuatorMsgPayload.h" + +/** + * @brief Linear SISO actuator model built on top of LinearTimeInvariantSystem. + * + * This class implements a single-input, single-output LTI system where + * both the input and output are carried through a SingleActuatorMsgPayload. + * + * The internal dynamics are defined by the A, B, C, D matrices stored in the + * base class LinearTimeInvariantSystem. + */ +class SingleActuatorLTI : public LinearTimeInvariantSystem +{ +public: + /** + * @brief Default constructor. + * + * The underlying LinearTimeInvariantSystem matrices must be configured + * (for example, via configureSecondOrder or direct setA/setB/setC/setD + * calls) before the simulation is run. + */ + SingleActuatorLTI() = default; + + /** + * @brief Get the dimension of the input vector. + * + * This model is strictly SISO, so the input dimension is always 1. + * + * @return Number of inputs (always 1). + */ + size_t getInputSize() const override { return 1; } + + /** + * @brief Get the dimension of the output vector. + * + * This model is strictly SISO, so the output dimension is always 1. + * + * @return Number of outputs (always 1). + */ + size_t getOutputSize() const override { return 1; } + + /** + * @brief Read the current input vector from the subscribed input message. + * + * This method constructs a 1x1 Eigen::VectorXd whose single element is + * taken from inMsg().input. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + * @return Input vector u of size 1. + */ + Eigen::VectorXd readInput(uint64_t CurrentSimNanos) override; + + /** + * @brief Write the current output vector to the output message. + * + * The first element of the output vector y is written to outMsg as the + * 'input' field of SingleActuatorMsgPayload. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + * @param y Output vector of size 1. + */ + void writeOutput(uint64_t CurrentSimNanos, + const Eigen::VectorXd &y) override; + +public: + /** + * @brief Output message carrying the actuator command. + * + * The 'input' field is populated from the first element of the output + * vector y computed by LinearTimeInvariantSystem. + */ + Message outMsg; + + /** + * @brief Input message read functor providing the actuator command. + * + * The 'input' field of the payload is mapped to the input vector u(0). + */ + ReadFunctor inMsg; +}; + +#endif /* SINGLE_ACTUATOR_LTI_H */ diff --git a/src/simulation/mujocoDynamics/meanRevertingNoise/_UnitTest/test_meanRevertingNoise.py b/src/simulation/mujocoDynamics/meanRevertingNoise/_UnitTest/test_meanRevertingNoise.py new file mode 100644 index 0000000000..ad305c94fa --- /dev/null +++ b/src/simulation/mujocoDynamics/meanRevertingNoise/_UnitTest/test_meanRevertingNoise.py @@ -0,0 +1,140 @@ +# +# 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. +import pytest +import numpy as np +import numpy.testing as npt +import matplotlib.pyplot as plt + +from Basilisk.utilities import SimulationBaseClass +from Basilisk.simulation import svIntegrators +from Basilisk.architecture import messaging +from Basilisk.utilities import macros + +try: + from Basilisk.simulation import mujoco + from Basilisk.simulation import MJMeanRevertingNoise + + couldImportMujoco = True +except: + raise + couldImportMujoco = False + +@pytest.mark.skipif(not couldImportMujoco, reason="Compiled Basilisk without --mujoco") +@pytest.mark.parametrize("usePython", [False, True]) +def test_meanRevertingNoise(usePython: bool = False, showPlots: bool = False): + """ + Unit test for classes that inherit from MeanRevertingNoise. + """ + dt = .002 # s + tf = 100 # s + + sigmaStationary = 0.8 # kg/m^3 + timeConstant = 0.1 # s + + scSim = SimulationBaseClass.SimBaseClass() + dynProcess = scSim.CreateNewProcess("test") + dynProcess.addTask(scSim.CreateNewTask("test", macros.sec2nano(dt))) + + # Create the MJScene from a simple cannonball body + scene = mujoco.MJScene("") # empty scene, no multi-body dynamics + scSim.AddModelToTask("test", scene) + + # must use strong integrator, since weak integrators do not produce + # the right sample-wise dynamic shape + integratorObject = svIntegrators.svStochIntegratorMayurama(scene) + integratorObject.setRNGSeed(0) + scene.setIntegrator(integratorObject) + + nominal = messaging.AtmoPropsMsg() + nominal.write(messaging.AtmoPropsMsgPayload(neutralDensity=1)) + + if usePython: + class PyStochasticAtmDensity(MJMeanRevertingNoise.MeanRevertingNoise): + + def __init__(self): + super().__init__() + self.atmoDensInMsg = messaging.AtmoPropsMsgReader() + self.atmoDensOutMsg = messaging.AtmoPropsMsg() + + def writeOutput(self, CurrentSimNanos, x): + payload = messaging.AtmoPropsMsgPayload( + neutralDensity = self.atmoDensInMsg().neutralDensity * (1 + x) + ) + self.atmoDensOutMsg.write(payload, self.moduleID, CurrentSimNanos) + + stochasticAtm = PyStochasticAtmDensity() + + else: + stochasticAtm = MJMeanRevertingNoise.StochasticAtmDensity() + stochasticAtm.setStationaryStd(sigmaStationary) + stochasticAtm.setTimeConstant(timeConstant) + stochasticAtm.atmoDensInMsg.subscribeTo( nominal ) + scene.AddModelToDynamicsTask(stochasticAtm) + + atmoDensRecorder = stochasticAtm.atmoDensOutMsg.recorder() + + scSim.AddModelToTask("test", atmoDensRecorder) + + scSim.InitializeSimulation() + + assert stochasticAtm.getTimeConstant() == timeConstant + assert stochasticAtm.getStationaryStd() == sigmaStationary + + scSim.ConfigureStopTime(macros.sec2nano(tf)) + scSim.ExecuteSimulation() + + t = atmoDensRecorder.times() + dens = atmoDensRecorder.neutralDensity + + # Plot results + if showPlots: + fig, ax = plt.subplots() + ax.plot(t, dens) + ax.set_xlabel("Time [s]") + ax.set_ylabel("Stochastic Density [kg/m^3]") + + plt.show() + + y = np.asarray(dens) + + # targets + meanTarget = 1.0 + varTarget = sigmaStationary**2 + phiTarget = np.exp(-dt / timeConstant) + + # estimates + meanEst = float(np.mean(y)) + varEst = float(np.var(y, ddof=1)) + yc0 = y[:-1] - meanEst + yc1 = y[1:] - meanEst + phiEst = float(np.dot(yc0, yc1) / np.dot(yc0, yc0)) + + phiphiEstClamped = min(max(phiEst, 1e-6), 0.999999) + timeConstantEst = -dt / np.log(phiphiEstClamped) + + print(f"mean={meanEst:.3f} var={varEst:.3f} tau_hat={timeConstantEst:.3f}") + + # assertions + npt.assert_allclose(meanEst, meanTarget, atol=0.2) + npt.assert_allclose(varEst, varTarget, rtol=0.3) + npt.assert_allclose(phiEst, phiTarget, rtol=0.03, atol=0.003) + npt.assert_allclose(timeConstantEst, timeConstant, rtol=0.05, atol=0.1) + + +if __name__ == "__main__": + assert couldImportMujoco + test_meanRevertingNoise(True, True) diff --git a/src/simulation/mujocoDynamics/meanRevertingNoise/meanRevertingNoise.cpp b/src/simulation/mujocoDynamics/meanRevertingNoise/meanRevertingNoise.cpp new file mode 100644 index 0000000000..7fe6ab91ac --- /dev/null +++ b/src/simulation/mujocoDynamics/meanRevertingNoise/meanRevertingNoise.cpp @@ -0,0 +1,93 @@ +/* + 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 "meanRevertingNoise.h" + +void MeanRevertingNoise::registerStates(DynParamRegisterer registerer) { + xState = registerer.registerState(1, 1, "meanRevertingState"); + xState->setNumNoiseSources(1); + // Default initial condition + xState->setState(Eigen::Matrix::Constant(0.0)); +} + +void MeanRevertingNoise::UpdateState(uint64_t CurrentSimNanos) { + // Safety + if (!xState) { + MJBasilisk::detail::logAndThrow( + "MeanRevertingNoise: state not registered", + &bskLogger + ); + } + + // Current x + const Eigen::MatrixXd &x = xState->getState(); + const double x0 = x(0, 0); + + // Drift: dx = -theta * x + Eigen::Matrix drift; + drift(0, 0) = -(1.0 / timeConstant) * x0; + xState->setDerivative(drift); + + // Diffusion: sigma * dW on noise source 0 + Eigen::Matrix diff; + diff(0, 0) = sigmaStationary * std::sqrt(2.0 / timeConstant); + xState->setDiffusion(diff, 0); + + // Delegate output + writeOutput(CurrentSimNanos, x0); +} + +double MeanRevertingNoise::getStateValue() const { + if (!xState) { + MJBasilisk::detail::logAndThrow( + "MeanRevertingNoise: getStateValue before initialization", + &const_cast(bskLogger) + ); + } + return xState->getState()(0, 0); +} + +void MeanRevertingNoise::setStateValue(double val) { + if (!xState) { + MJBasilisk::detail::logAndThrow( + "MeanRevertingNoise: setStateValue before initialization", + &bskLogger + ); + } + xState->setState(Eigen::Matrix::Constant(val)); +} + +void MeanRevertingNoise::setTimeConstant(double t) { + if (t <= 0.0) { + MJBasilisk::detail::logAndThrow( + "MeanRevertingNoise::setTimeConstant: tau must be > 0", + &bskLogger + ); + } + timeConstant = t; +} + +void MeanRevertingNoise::setStationaryStd(double s) { + if (s < 0.0) { + MJBasilisk::detail::logAndThrow( + "MeanRevertingNoise::setStationaryStd: sigma_st must be >= 0", + &bskLogger + ); + } + sigmaStationary = s; +} diff --git a/src/simulation/mujocoDynamics/meanRevertingNoise/meanRevertingNoise.h b/src/simulation/mujocoDynamics/meanRevertingNoise/meanRevertingNoise.h new file mode 100644 index 0000000000..3cd7622e10 --- /dev/null +++ b/src/simulation/mujocoDynamics/meanRevertingNoise/meanRevertingNoise.h @@ -0,0 +1,146 @@ +/* + 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 MEAN_REVERTING_NOISE_H +#define MEAN_REVERTING_NOISE_H + +#include +#include +#include +#include + +#include "architecture/utilities/bskLogging.h" +#include "architecture/messaging/messaging.h" +#include "simulation/mujocoDynamics/_GeneralModuleFiles/StatefulSysModel.h" +#include "simulation/mujocoDynamics/_GeneralModuleFiles/MJUtils.h" + +/** + * @class MeanRevertingNoise + * @brief Base class that provides a scalar Ornstein–Uhlenbeck state x and its + * stochastic evolution. Derived classes implement output handling. + * + * The state x follows + * \f[ + * dx = -\frac{1}{\tau}\,x\,dt + \sqrt{\frac{2}{\tau}}\,\sigma_{\text{st}}\,dW + * \f] + * The base class: + * 1) registers a 1-by-1 state with one noise source + * 2) computes drift and diffusion each step + * 3) calls the user hook writeOutput(CurrentSimNanos, x) + */ +class MeanRevertingNoise : public StatefulSysModel { +public: + /** + * @brief Default constructor. Parameters remain at defaults. + */ + MeanRevertingNoise() = default; + + /** @name Framework interface @{ */ + + /** + * @brief Register the scalar OU state with one noise source. + * @param registerer Simulation state registerer. + */ + void registerStates(DynParamRegisterer registerer) override; + + /** + * @brief Update the OU state drift and diffusion, then call writeOutput. + * @param CurrentSimNanos Current simulation time in nanoseconds. + */ + void UpdateState(uint64_t CurrentSimNanos) override; + + /** @} */ + + /** @name Parameter getters and setters in stationary form @{ */ + + /** + * @brief Get stationary standard deviation \f$\sigma_{\text{st}}\f$. + */ + double getStationaryStd() const { return sigmaStationary; } + /** + * @brief Set stationary standard deviation \f$\sigma_{\text{st}}\f$. + * @param s Nonnegative expected. + */ + void setStationaryStd(double s); + + /** + * @brief Get time constant \f$\tau\f$ [s]. + */ + double getTimeConstant() const { return timeConstant; } + + /** + * @brief Set time constant \f$\tau\f$ [s]. + * @param t Positive expected. + */ + void setTimeConstant(double t); + + /** @} */ + + /** @name Canonical OU parameter accessors @{ */ + + /** + * @brief Get \f$\theta = 1/\tau\f$ [s^{-1}]. + */ + double getTheta() const { return 1.0 / timeConstant; } + + /** + * @brief Get diffusion amplitude \f$\sigma = \sqrt{2/\tau}\,\sigma_{\text{st}}\f$. + */ + double getSigma() const { return sigmaStationary * std::sqrt(2.0 / timeConstant); } + + /** @} */ + + /** @name State access @{ */ + + /** + * @brief Get the current scalar state x. + * @throws std::runtime_error if called before registerStates. + */ + double getStateValue() const; + + /** + * @brief Set the scalar state x. + * @param val New value. + * @throws std::runtime_error if called before registerStates. + */ + void setStateValue(double val); + + /** @} */ + + /// Logger + BSKLogger bskLogger; + +protected: + /** + * @brief Pure virtual output hook. Called every step after drift and diffusion + * are set. Use this to read inputs if needed and write outputs. + * @param CurrentSimNanos Current simulation time in nanoseconds. + * @param x Current scalar state value. + */ + virtual void writeOutput(uint64_t CurrentSimNanos, double x) = 0; + +private: + // State storage + StateData *xState = nullptr; + + // Parameters in stationary form + double sigmaStationary = 0.0; ///< Stationary std dev of x. + double timeConstant = 1.0; ///< OU time constant in seconds. +}; + +#endif diff --git a/src/simulation/mujocoDynamics/meanRevertingNoise/meanRevertingNoise.i b/src/simulation/mujocoDynamics/meanRevertingNoise/meanRevertingNoise.i new file mode 100644 index 0000000000..8e9e0344fe --- /dev/null +++ b/src/simulation/mujocoDynamics/meanRevertingNoise/meanRevertingNoise.i @@ -0,0 +1,85 @@ +/* + ISC License + + Copyright (c) 2025, Autonomous Vehicle Systems Lab + + 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(directors="1",threads="1") MJMeanRevertingNoise + +%include "architecture/utilities/bskException.swg" +%default_bsk_exception(); + +%{ + #include "meanRevertingNoise.h" + #include "stochasticAtmDensity.h" + #include "stochasticDragCoeff.h" +%} + +%pythoncode %{ +from Basilisk.architecture.swig_common_model import * +%} + +/* Common SWIG helpers */ +%include "swig_eigen.i" +%include "std_string.i" +%include "exception.i" + +/* Basilisk system-model base and helpers */ +%import "simulation/mujocoDynamics/_GeneralModuleFiles/pyStatefulSysModel.i" + +/* ============================ + Base class: MeanRevertingNoise + ============================ */ + +/* Enable directors so Python subclasses can override virtual methods */ +%feature("director") MeanRevertingNoise; + +/* Rename the raw SWIG proxy to _MeanRevertingNoise so we can define + * a Python-side wrapper class named MeanRevertingNoise that also + * inherits from the Python StatefulSysModel wrapper. + */ +%rename("_MeanRevertingNoise") MeanRevertingNoise; +%include "meanRevertingNoise.h" + +/* ============================ + StochasticAtmDensity + ============================ */ + +/* Keep these concrete C++ classes as they are. We do not want them + * to use the Python StatefulSysModel wrapper, and we still do not + * expose their writeOutput directly. + */ +%ignore StochasticAtmDensity::writeOutput; +%include "stochasticAtmDensity.h" + +%include "architecture/msgPayloadDefC/AtmoPropsMsgPayload.h" +struct AtmoPropsMsgPayload_C; + +/* ============================ + StochasticDragCoeff + ============================ */ + +%ignore StochasticDragCoeff::writeOutput; +%include "stochasticDragCoeff.h" + +%include "architecture/msgPayloadDefC/DragGeometryMsgPayload.h" +struct DragGeometryMsgPayload_C; + +%pythoncode %{ +from Basilisk.architecture.sysModel import SysModelMixin + +class MeanRevertingNoise(SysModelMixin, _MeanRevertingNoise): + """Python wrapper for the C++ MeanRevertingNoise.""" +%} diff --git a/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticAtmDensity.cpp b/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticAtmDensity.cpp new file mode 100644 index 0000000000..0a25f4ded6 --- /dev/null +++ b/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticAtmDensity.cpp @@ -0,0 +1,26 @@ +/* + 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 "stochasticAtmDensity.h" + +void StochasticAtmDensity::writeOutput(uint64_t CurrentSimNanos, double x) +{ + AtmoPropsMsgPayload out = atmoDensInMsg(); + out.neutralDensity *= 1.0 + x; + atmoDensOutMsg.write(&out, CurrentSimNanos, this->moduleID); +} diff --git a/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticAtmDensity.h b/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticAtmDensity.h new file mode 100644 index 0000000000..45f28a7ef4 --- /dev/null +++ b/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticAtmDensity.h @@ -0,0 +1,80 @@ +/* + 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 STOCHASTIC_ATM_DENSITY_H +#define STOCHASTIC_ATM_DENSITY_H + +#include "meanRevertingNoise.h" + +#include "architecture/msgPayloadDefC/AtmoPropsMsgPayload.h" + +/** + * @class StochasticAtmDensity + * @brief Applies a mean-reverting (Ornstein–Uhlenbeck) correction to atmospheric density. + * + * This class derives from MeanRevertingNoise, which provides the scalar + * OU state \f$x\f$ evolving as: + * \f[ + * dx = -\frac{1}{\tau}\, x \, dt + \sqrt{\frac{2}{\tau}}\, \sigma_{\text{st}}\, dW + * \f] + * + * The corrected density is: + * \f[ + * \rho_\text{out} = \rho_\text{in}\,(1 + x) + * \f] + * + * The state \f$x\f$ is stored and propagated by the base class. + * This class only specifies how the state modifies the atmosphere message. + */ +class StochasticAtmDensity : public MeanRevertingNoise { +public: + + /** @name I/O Messages @{ */ + + /** + * @brief Input atmospheric properties message. + * + * This message is read each step to obtain the unperturbed atmospheric density + * and temperature before applying the stochastic correction. + */ + ReadFunctor atmoDensInMsg; + + /** + * @brief Output atmospheric properties message. + * + * This message is written each step with the corrected neutral density value. + */ + Message atmoDensOutMsg; + + /** @} */ + +protected: + /** + * @brief Apply the OU correction factor and write output. + * + * Called automatically by MeanRevertingNoise::UpdateState(), after + * the stochastic state has been updated and before the simulation advances. + * + * @param CurrentSimNanos Current simulation time in nanoseconds. + * @param x Current value of the scalar mean-reverting correction factor. + */ + void writeOutput(uint64_t CurrentSimNanos, double x) override; +}; + + +#endif diff --git a/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticDragCoeff.cpp b/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticDragCoeff.cpp new file mode 100644 index 0000000000..7ad881c71d --- /dev/null +++ b/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticDragCoeff.cpp @@ -0,0 +1,26 @@ +/* + 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 "stochasticDragCoeff.h" + +void StochasticDragCoeff::writeOutput(uint64_t CurrentSimNanos, double x) +{ + DragGeometryMsgPayload out = dragGeomInMsg(); + out.dragCoeff *= (1.0 + x); + dragGeomOutMsg.write(&out, CurrentSimNanos, this->moduleID); +} diff --git a/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticDragCoeff.h b/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticDragCoeff.h new file mode 100644 index 0000000000..9e6510f0ea --- /dev/null +++ b/src/simulation/mujocoDynamics/meanRevertingNoise/stochasticDragCoeff.h @@ -0,0 +1,66 @@ +/* + 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 STOCHASTIC_DRAG_COEFF_H +#define STOCHASTIC_DRAG_COEFF_H + +#include "meanRevertingNoise.h" +#include "architecture/msgPayloadDefC/DragGeometryMsgPayload.h" + +/** + * @class StochasticDragCoeff + * @brief Applies a mean-reverting Ornstein–Uhlenbeck correction to the drag coefficient. + * + * The scalar state \( x \) evolves as + * \( dx = -(1/\tau)\,x\,dt + \sqrt{2/\tau}\,\sigma_{\text{st}}\,dW \). + * The outgoing coefficient is modified as + * \( C_{D,\text{out}} = C_{D,\text{in}}(1 + x) \). + * + * The projected area and \( r_{CP,S} \) are passed through unchanged. + */ +class StochasticDragCoeff : public MeanRevertingNoise { +public: + /** @name I/O Messages @{ */ + + /** + * @brief Input drag geometry message. + * + * Read each step to obtain the nominal geometry and coefficient. + */ + ReadFunctor dragGeomInMsg; + + /** + * @brief Output drag geometry message with corrected drag coefficient. + */ + Message dragGeomOutMsg; + + /** @} */ + +protected: + /** + * @brief Apply the OU correction to \( C_D \) and write the result. + * + * Called by MeanRevertingNoise::UpdateState after drift and diffusion are set. + * + * @param CurrentSimNanos Simulation time in nanoseconds. + * @param x Current scalar correction factor. + */ + void writeOutput(uint64_t CurrentSimNanos, double x) override; +}; + +#endif 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__])