Issue #91: team-02 - Mean trajectory dispersion simulation

docs/notebooks/plotting_dispersion_analysis.ipynb

@@ -0,0 +1,133 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from matplotlib import pyplot as plt\n",
+    "from matplotlib.pyplot import cm\n",
+    "from doctest import testfile\n",
+    "from rocketpy import Environment, Rocket, Flight, SolidMotor, Function\n",
+    "from rocketpy.utilities import set_parabolic_trajectory, plot_dispersion\n",
+    "import numpy as np\n",
+    "import datetime\n",
+    "\n",
+    "def test_flight():\n",
+    "    test_env = Environment(\n",
+    "        railLength=5,\n",
+    "        latitude=32.990254,\n",
+    "        longitude=-106.974998,\n",
+    "        elevation=1400,\n",
+    "        datum=\"WGS84\",\n",
+    "    )\n",
+    "    tomorrow = datetime.date.today() + datetime.timedelta(days=1)\n",
+    "    test_env.setDate(\n",
+    "        (tomorrow.year, tomorrow.month, tomorrow.day, 12)\n",
+    "    )  # Hour given in UTC time\n",
+    "\n",
+    "    test_motor = SolidMotor(\n",
+    "        thrustSource=\"data/motors/Cesaroni_M1670.eng\",\n",
+    "        burnOut=3.9,\n",
+    "        grainNumber=5,\n",
+    "        grainSeparation=5 / 1000,\n",
+    "        grainDensity=1815,\n",
+    "        grainOuterRadius=33 / 1000,\n",
+    "        grainInitialInnerRadius=15 / 1000,\n",
+    "        grainInitialHeight=120 / 1000,\n",
+    "        nozzleRadius=33 / 1000,\n",
+    "        throatRadius=11 / 1000,\n",
+    "        interpolationMethod=\"linear\",\n",
+    "    )\n",
+    "\n",
+    "    test_rocket = Rocket(\n",
+    "        motor=test_motor,\n",
+    "        radius=127 / 2000,\n",
+    "        mass=19.197 - 2.956,\n",
+    "        inertiaI=6.60,\n",
+    "        inertiaZ=0.0351,\n",
+    "        distanceRocketNozzle=-1.255,\n",
+    "        distanceRocketPropellant=-0.85704,\n",
+    "        powerOffDrag=\"data/calisto/powerOffDragCurve.csv\",\n",
+    "        powerOnDrag=\"data/calisto/powerOnDragCurve.csv\",\n",
+    "    )\n",
+    "\n",
+    "    test_rocket.setRailButtons([0.2, -0.5])\n",
+    "\n",
+    "    NoseCone = test_rocket.addNose(\n",
+    "        length=0.55829, kind=\"vonKarman\", distanceToCM=0.71971\n",
+    "    )\n",
+    "    FinSet = test_rocket.addFins(\n",
+    "        4, span=0.100, rootChord=0.120, tipChord=0.040, distanceToCM=-1.04956\n",
+    "    )\n",
+    "    Tail = test_rocket.addTail(\n",
+    "        topRadius=0.0635, bottomRadius=0.0435, length=0.060, distanceToCM=-1.194656\n",
+    "    )\n",
+    "\n",
+    "    def drogueTrigger(p, y):\n",
+    "        # p = pressure\n",
+    "        # y = [x, y, z, vx, vy, vz, e0, e1, e2, e3, w1, w2, w3]\n",
+    "        # activate drogue when vz < 0 m/s.\n",
+    "        return True if y[5] < 0 else False\n",
+    "\n",
+    "    def mainTrigger(p, y):\n",
+    "        # p = pressure\n",
+    "        # y = [x, y, z, vx, vy, vz, e0, e1, e2, e3, w1, w2, w3]\n",
+    "        # activate main when vz < 0 m/s and z < 800 m.\n",
+    "        return True if y[5] < 0 and y[2] < 800 else False\n",
+    "\n",
+    "    Main = test_rocket.addParachute(\n",
+    "        \"Main\",\n",
+    "        CdS=10.0,\n",
+    "        trigger=mainTrigger,\n",
+    "        samplingRate=105,\n",
+    "        lag=1.5,\n",
+    "        noise=(0, 8.3, 0.5),\n",
+    "    )\n",
+    "\n",
+    "    Drogue = test_rocket.addParachute(\n",
+    "        \"Drogue\",\n",
+    "        CdS=1.0,\n",
+    "        trigger=drogueTrigger,\n",
+    "        samplingRate=105,\n",
+    "        lag=1.5,\n",
+    "        noise=(0, 8.3, 0.5),\n",
+    "    )\n",
+    "\n",
+    "    test_flight = Flight(\n",
+    "        rocket=test_rocket, environment=test_env, inclination=85, heading=0\n",
+    "    )\n",
+    "\n",
+    "    return test_flight\n",
+    "\n",
+    "flight2 = test_flight()\n",
+    "flight3 = test_flight()\n",
+    "\n",
+    "set_parabolic_trajectory(flight2, t=(0, 600), x=(0, 2000), y=(0, 3000), zmax=10000, n=601)\n",
+    "set_parabolic_trajectory(flight3, t=(0, 300), x=(0, 1000), y=(0, 5000), zmax=15000, n=301)\n",
+    "\n",
+    "plot_dispersion([flight2, flight3], postProcess=False)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.10.4 64-bit",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.4"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "2a3265b3f13718e60c35107355e230194cea6bb45665fbc0d844c53e8c491ff6"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

rocketpy/utilities.py

@@ -5,6 +5,8 @@
 
 import numpy as np
 from scipy.integrate import solve_ivp
+from matplotlib import pyplot as plt
+from matplotlib.pyplot import cm
 
 from .Environment import Environment
 from .Function import Function
@@ -197,3 +199,213 @@ def du(z, u):
         velocityFunction()
 
     return altitudeFunction, velocityFunction, final_sol
+
+
+# TODO: Needs tests
+
+def traj(flight, postProcess=True, n=1000):
+    '''
+    Convert a flight into a trajectory
+
+    flight          Flight
+    postProcess     if True, make sure the flight is postProcessed
+    n               Number of points for reinterpolation
+
+    Returns
+        trajectory  ndarray([(s, x, y, z), ...])
+    '''
+    if postProcess and not flight.postProcessed:
+        flight.postProcess()
+    s = np.linspace(0, 1, n)
+    tq = flight.x.source[:, 0]
+    t0 = tq.min()
+    tf = tq.max()
+    dt = tf - t0
+    t = t0 + dt * s
+    x = flight.x(t)
+    y = flight.y(t)
+    z = flight.z(t)
+    return np.stack([s, x, y, z], 1)
+
+def flight_stats(flights, postProcess=True, n=1000):
+    '''
+    Prepare statistics for a set of flights
+
+    flights         iterable of Flight
+    postProcess     if True, make sure each Flight has been postProcessed before proceeding
+    n               number of timepoints to sample during reinterpolation
+
+    Returns
+        mean        trajectory [(s, x, y, z), ...] representing the mean non-dimensionalized flight
+        lower       trajectory [(s, x, y, z), ...] representing a flight one standard deviation lower than the mean
+        upper       trajectory [(s, x, y, z), ...] representing a flight one standard deviation higher than the mean
+    '''
+    # Convert flights into ndarrays
+    # Non-dimensionalize time by their respective flight durations
+    trajs = np.zeros([len(flights), 1000, 4])
+    for i, flight in enumerate(flights):
+        trajs[i, :, :] = traj(flight, postProcess=postProcess, n=n)
+
+    # Calculate mean trajectory
+    mean = sum(trajs) / len(trajs)
+
+    # Calculate stddev at each time point
+    stddev = np.sqrt(np.sum((trajs[:, :, 1 :] - mean[:, 1 :])**2, axis=0) / len(trajs))[:, 2]
+
+    # Calculate upper trajectory
+    upper = mean.copy()
+    upper[:, 3] += stddev
+    lower = mean.copy()
+    lower[:, 3] -= stddev
+
+    return mean, lower, upper
+
+def axes3d():
+    '''
+    Get a set of axes with 3d projection
+
+    Returns
+        ax          axis with 3d projection
+    '''
+    plt.figure()
+    return plt.subplot(projection='3d')
+
+def plot_traj(traj, color, linestyle="-", elevation=None, projections=True, ax=None):
+    '''
+    Plot trajectory in 3d
+
+    traj            ndarray([(t, x, y, z), ...])
+    color           matplotlib color
+    linestyle       style of line for plotting, default "-"
+    elevation (m)   environment elevation
+    projections     (unimplemented) if True, plot projections of trajectory onto principal planes
+    ax              axis to plot on
+
+    Returns
+        ax          axis which was plotted to
+    '''
+    # Get axis if one is not provided
+    if ax is None:
+        plt.figure()
+        ax = plt.subplot(projection='3d')
+
+    # If environment elevation is not supplied, use minimum z value
+    if elevation is None:
+        elevation = traj[:, 3].min()
+
+    # Does not currently work correctly
+    # If plane projections are enabled, plot plane projections first
+    # if projections:
+    #     ax.plot(traj[:, 1], traj[:, 2], zs=0, zdir="z", linestyle="--", color=color)
+    #     ax.plot(traj[:, 1], traj[:, 3] - elevation, zs=elevation, zdir="y", linestyle="--", color=color)
+    #     ax.plot(traj[:, 2], traj[:, 2] - elevation, zs=elevation, zdir="x", linestyle="--", color=color)
+
+    # Plot 3d trajectory curve
+    ax.plot(traj[:, 1], traj[:, 2], traj[:, 3] - elevation, linestyle=linestyle, linewidth="2", color=color)
+
+    # Plot origin
+    ax.scatter(0, 0, 0, color="k")
+
+    # Set labels
+    ax.set_xlabel("X - East (m)")
+    ax.set_ylabel("Y - North (m)")
+    ax.set_zlabel("Z - Altitude Above Ground Level (m)")
+    ax.set_title("Flight Trajectory")
+
+    # Set view orientation
+    ax.view_init(15, 45)
+
+    return ax
+
+def plot_flight(flight, color, postProcess=True, n=1000, projections=True, ax=None):
+    '''
+    Plot flight in 3d
+
+    flight          Flight
+    color           matplotlib color
+    postProcess     if True, make sure the flight is postProcessed before proceeding
+    n               number of points for reinterpolation
+    projections     (unimplemented) if True, show projections of trajectory onto principal planes
+    ax              axis onto which to plot, if None, an axis is initialized automatically
+
+    Returns
+        ax          the axis onto which the flight was plotted
+    '''
+    return plot_traj(traj(flight, postProcess, n), color, elevation=flight.env.elevation, projections=projections, ax=ax)
+
+def plot_flights(flights, colors=None, elevations=None, postProcess=True, n=1000, projections=True, ax=None):
+    '''
+    Plot flights or trajectories automatically in 3d
+
+    flights         iterable of Flights or trajectories
+    colors          iterable of matplotlib colors, if None, the rainbow is used
+    elevations      elevations for trajectories, if None, elevations are inferred
+    postProcess     if True, make sure Flights are postProcessed before proceeding
+    n               number of points for reinterpolation
+    projections     (unimplemented) if True, plot projections of trajectories onto principal planes
+    ax              axis onto which to plot, if None axis is intitialized automatically
+
+    Returns
+        ax          axis which was plotted to
+    '''
+    # If colors is not given, select colors automatically
+    if colors is None:
+        colors = cm.rainbow(np.linspace(0, 1, len(flights)))
+
+    # If elevations are not given, do not give elevations
+    if elevations is None:
+        elevations = [None for _ in flights]
+
+    # Plot each flight or trajectory
+    for flight, color, elevation in zip(flights, colors, elevations):
+        if isinstance(flight, np.ndarray):
+            ax = plot_traj(flight, color, elevation=elevation, projections=projections, ax=ax)
+        else:
+            ax = plot_flight(flight, color, postProcess=postProcess, n=n, projections=projections, ax=ax)
+
+    return ax
+
+def plot_dispersion(flights, postProcess=True, n=1000, projections=True, ax=None):
+    '''
+    Plot the results of a dispersion analysis with mean trajectory
+
+    flights         Flights for which trajectories will be plotted, len(flights) >= 1
+    postProcess     if True, make sure flights have been postProcessed before proceeding
+    n               number of reinterpolation points
+    projections     (unimplemented) if True, plot projections of trajectories onto principal axes
+    ax              axis onto which to plot
+
+    Returns
+        ax          axis which was plotted to
+    '''
+    mean, lower, upper = flight_stats(flights, postProcess=postProcess, n=n)
+    ax = plot_flights(flights, ['k' for _ in range(len(flights))], postProcess=postProcess, n=n, projections=projections, ax=ax)
+    ax = plot_traj(mean, 'r', linestyle="-", elevation=flights[0].env.elevation, projections=projections, ax=ax)
+    ax = plot_traj(lower, 'r', linestyle="--", elevation=flights[0].env.elevation, projections=projections, ax=ax)
+    ax = plot_traj(upper, 'r', linestyle="--", elevation=flights[0].env.elevation, projections=projections, ax=ax)
+    return ax
+
+def set_parabolic_trajectory(flight, t, x, y, zmax, n):
+    '''
+    Create a parabolic trajectory embedded in 3d space
+
+    flight      Flight object to modify
+    t           (t0, tf)
+    x           (x0, xf)
+    y           (y0, yf)
+    zmax        apogee
+    n           number of interpolation points
+
+    Returns
+        trajectory  [(t, x, y, z), ...]
+    '''
+    ts = np.linspace(*t, n)
+    xs = np.linspace(*x, n)
+    ys = np.linspace(*y, n)
+    zs = -4 * zmax * (ts - t[0]) * (ts - t[1]) / (t[1] - t[0])**2
+
+    flight.x = Function(np.stack([ts, xs], 1))
+    flight.y = Function(np.stack([ts, ys], 1))
+    flight.z = Function(np.stack([ts, zs], 1))
+
+    return flight