simple-squiggle

A restricted subset of Squiggle
Log | Files | Refs | README

rocket_trajectory_optimization.html (11158B)


      1 <!DOCTYPE html>
      2 <html lang="en">
      3 
      4 <head>
      5   <meta charset="utf-8">
      6   <title>math.js | rocket trajectory optimization</title>
      7 
      8   <script src="../../lib/browser/math.js"></script>
      9   <script src="https://cdnjs.cloudflare.com/ajax/libs/Chart.js/2.5.0/Chart.min.js"></script>
     10 
     11   <style>
     12     body {
     13       font-family: sans-serif;
     14     }
     15 
     16     #canvas-grid {
     17       display: grid;
     18       grid-template-columns: repeat(2, 1fr);
     19       gap: 5%;
     20       margin-top: 5%;
     21     }
     22 
     23     #canvas-grid>div {
     24       overflow: hidden;
     25     }
     26   </style>
     27 </head>
     28 
     29 <body>
     30   <h1>Rocket trajectory optimization</h1>
     31   <p>
     32     This example simulates the launch of a SpaceX Falcon 9 modeled using a system of ordinary differential equations.
     33   </p>
     34 
     35   <canvas id="canvas" width="1600" height="600"></canvas>
     36   <div id="canvas-grid"></div>
     37 
     38   <script>
     39     // Solve ODE `dx/dt = f(x,t), x(0) = x0` numerically. 
     40     function ndsolve(f, x0, dt, tmax) {
     41       let x = x0.clone()  // Current values of variables
     42       const result = [x]  // Contains entire solution
     43       const nsteps = math.divide(tmax, dt)   // Number of time steps
     44       for (let i = 0; i < nsteps; i++) {
     45         // Compute derivatives
     46         const dxdt = f.map(func => func(...x.toArray()))
     47         // Euler method to compute next time step
     48         const dx = math.multiply(dxdt, dt)
     49         x = math.add(x, dx)
     50         result.push(x)
     51       }
     52       return math.matrix(result)
     53     }
     54 
     55     // Import the numerical ODE solver
     56     math.import({ ndsolve })
     57 
     58     // Create a math.js context for our simulation. Everything else occurs in the context of the expression parser!
     59     const sim = math.parser()
     60 
     61     sim.evaluate("G = 6.67408e-11 m^3 kg^-1 s^-2")  // Gravitational constant
     62     sim.evaluate("mbody = 5.9724e24 kg")            // Mass of Earth
     63     sim.evaluate("mu = G * mbody")                  // Standard gravitational parameter
     64     sim.evaluate("g0 = 9.80665 m/s^2")              // Standard gravity: used for calculating prop consumption (dmdt)
     65     sim.evaluate("r0 = 6371 km")                    // Mean radius of Earth
     66     sim.evaluate("t0 = 0 s")                        // Simulation start
     67     sim.evaluate("dt = 0.5 s")                      // Simulation timestep
     68     sim.evaluate("tfinal = 149.5 s")                // Simulation duration
     69     sim.evaluate("isp_sea = 282 s")                 // Specific impulse (at sea level)
     70     sim.evaluate("isp_vac = 311 s")                 // Specific impulse (in vacuum)
     71     sim.evaluate("gamma0 = 89.99970 deg")           // Initial pitch angle (90 deg is vertical)
     72     sim.evaluate("v0 = 1 m/s")                      // Initial velocity (must be non-zero because ODE is ill-conditioned)
     73     sim.evaluate("phi0 = 0 deg")                    // Initial orbital reference angle
     74     sim.evaluate("m1 = 433100 kg")                  // First stage mass
     75     sim.evaluate("m2 = 111500 kg")                  // Second stage mass
     76     sim.evaluate("m3 = 1700 kg")                    // Third stage / fairing mass
     77     sim.evaluate("mp = 5000 kg")                    // Payload mass
     78     sim.evaluate("m0 = m1+m2+m3+mp")                // Initial mass of rocket
     79     sim.evaluate("dm = 2750 kg/s")                  // Mass flow rate
     80     sim.evaluate("A = (3.66 m)^2 * pi")             // Area of the rocket
     81     sim.evaluate("dragCoef = 0.2")                  // Drag coefficient
     82 
     83     // Define the equations of motion. We just thrust into current direction of motion, e.g. making a gravity turn.
     84     sim.evaluate("gravity(r) = mu / r^2")
     85     sim.evaluate("angVel(r, v, gamma) = v/r * cos(gamma) * rad")   // Angular velocity of rocket around moon
     86     sim.evaluate("density(r) = 1.2250 kg/m^3 * exp(-g0 * (r - r0) / (83246.8 m^2/s^2))") // Assume constant temperature
     87     sim.evaluate("drag(r, v) = 1/2 * density(r) .* v.^2 * A * dragCoef")
     88     sim.evaluate("isp(r) = isp_vac + (isp_sea - isp_vac) * density(r)/density(r0)") // pressure ~ density for constant temperature
     89     sim.evaluate("thrust(isp) = g0 * isp * dm")
     90     // It is important to maintain the same argument order for each of these functions.
     91     sim.evaluate("drdt(r, v, m, phi, gamma, t) = v sin(gamma)")
     92     sim.evaluate("dvdt(r, v, m, phi, gamma, t) = - gravity(r) * sin(gamma) + (thrust(isp(r)) - drag(r, v)) / m")
     93     sim.evaluate("dmdt(r, v, m, phi, gamma, t) = - dm")
     94     sim.evaluate("dphidt(r, v, m, phi, gamma, t) = angVel(r, v, gamma)")
     95     sim.evaluate("dgammadt(r, v, m, phi, gamma, t) = angVel(r, v, gamma) - gravity(r) * cos(gamma) / v * rad")
     96     sim.evaluate("dtdt(r, v, m, phi, gamma, t) = 1")
     97 
     98     // Remember to maintain the same variable order in the call to ndsolve.
     99     sim.evaluate("result_stage1 = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], [r0, v0, m0, phi0, gamma0, t0], dt, tfinal)")
    100 
    101     // Reset initial conditions for interstage flight
    102     sim.evaluate("dm = 0 kg/s")
    103     sim.evaluate("tfinal = 10 s")
    104     sim.evaluate("x = flatten(result_stage1[end,:])")
    105     sim.evaluate("x[3] = m2+m3+mp") // New mass after stage seperation
    106     sim.evaluate("result_interstage = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], x, dt, tfinal)")
    107 
    108     // Reset initial conditions for stage 2 flight
    109     sim.evaluate("dm = 270.8 kg/s")
    110     sim.evaluate("isp_vac = 348 s")
    111     sim.evaluate("tfinal = 350 s")
    112     sim.evaluate("x = flatten(result_interstage[end,:])")
    113     sim.evaluate("result_stage2 = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], x, dt, tfinal)")
    114 
    115     // Reset initial conditions for unpowered flight
    116     sim.evaluate("dm = 0 kg/s")
    117     sim.evaluate("tfinal = 900 s")
    118     sim.evaluate("dt = 10 s")
    119     sim.evaluate("x = flatten(result_stage2[end,:])")
    120     sim.evaluate("result_unpowered1 = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], x, dt, tfinal)")
    121 
    122     // Reset initial conditions for final orbit insertion
    123     sim.evaluate("dm = 270.8 kg/s")
    124     sim.evaluate("tfinal = 39 s")
    125     sim.evaluate("dt = 0.5 s")
    126     sim.evaluate("x = flatten(result_unpowered1[end,:])")
    127     sim.evaluate("result_insertion = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], x, dt, tfinal)")
    128 
    129     // Reset initial conditions for unpowered flight
    130     sim.evaluate("dm = 0 kg/s")
    131     sim.evaluate("tfinal = 250 s")
    132     sim.evaluate("dt = 10 s")
    133     sim.evaluate("x = flatten(result_insertion[end,:])")
    134     sim.evaluate("result_unpowered2 = ndsolve([drdt, dvdt, dmdt, dphidt, dgammadt, dtdt], x, dt, tfinal)")
    135 
    136     // Now it's time to prepare results for plotting
    137     const resultNames = ['stage1', 'interstage', 'stage2', 'unpowered1', 'insertion', 'unpowered2']
    138       .map(stageName => `result_${stageName}`)
    139 
    140     // Concat result matrices
    141     sim.set('result',
    142       math.concat(
    143         ...resultNames.map(resultName =>
    144           sim.evaluate(`${resultName}[:end-1, :]`)  // Avoid overlap
    145         ),
    146         0 // Concat in row-dimension
    147       )
    148     )
    149 
    150     const mainDatasets = resultNames.map((resultName, i) => ({
    151       label: resultName.slice(7),
    152       data: sim.evaluate(
    153         'concat('
    154         + `(${resultName}[:,4] - phi0) * r0 / rad / km,`  // Surface distance from start (in km)
    155         + `(${resultName}[:,1] - r0) / km`                // Height above surface (in km)
    156         + ')'
    157       ).toArray().map(([x, y]) => ({ x, y })),
    158       borderColor: i % 2 ? '#999' : '#dc3912',
    159       fill: false,
    160       pointRadius: 0,
    161     }))
    162     new Chart(document.getElementById('canvas'), {
    163       type: 'line',
    164       data: { datasets: mainDatasets },
    165       options: getMainChartOptions()
    166     })
    167 
    168     createChart([{
    169       label: 'velocity (in m/s)',
    170       data: sim.evaluate("result[:,[2,6]]")
    171         .toArray()
    172         .map(([v, t]) => ({ x: t.toNumber('s'), y: v.toNumber('m/s') }))
    173     }])
    174     createChart([{
    175       label: 'height (in km)',
    176       data: sim.evaluate("concat((result[:, 1] - r0), result[:, 6])")
    177         .toArray()
    178         .map(([r, t]) => ({ x: t.toNumber('s'), y: r.toNumber('km') })),
    179     }])
    180     createChart([{
    181       label: 'gamma (in deg)',
    182       data: sim.evaluate("result[:, [5,6]]")
    183         .toArray()
    184         .map(([gamma, t]) => ({ x: t.toNumber('s'), y: gamma.toNumber('deg') })),
    185     }])
    186     createChart([{
    187       label: 'acceleration (in m/s^2)',
    188       data: sim.evaluate("concat(diff(result[:, 2]) ./ diff(result[:, 6]), result[:end-1, 6])")
    189         .toArray()
    190         .map(([acc, t]) => ({ x: t.toNumber('s'), y: acc.toNumber('m/s^2') })),
    191     }])
    192     createChart([{
    193       label: 'drag acceleration (in m/s^2)',
    194       data: sim.evaluate("concat(drag(result[:, 1], result[:, 2]) ./ result[:, 3], result[:, 6])")
    195         .toArray()
    196         .map(([dragAcc, t]) => ({ x: t.toNumber('s'), y: dragAcc.toNumber('m/s^2') })),
    197     }])
    198     createChart(
    199       [
    200         {
    201           data: sim.evaluate("result[:, [1,4]]")
    202             .toArray()
    203             .map(([r, phi]) => math.rotate([r.toNumber('km'), 0], phi))
    204             .map(([x, y]) => ({ x, y })),
    205         },
    206         {
    207           data: sim.evaluate("map(0:0.25:360, function(angle) = rotate([r0/km, 0], angle))")
    208             .toArray()
    209             .map(([x, y]) => ({ x, y })),
    210           borderColor: "#999",
    211           fill: true
    212         }
    213       ],
    214       getEarthChartOptions()
    215     )
    216 
    217     // Helper functions for plotting data (nothing to learn about math.js from here on)
    218     function createChart(datasets, options = {}) {
    219       const container = document.createElement("div")
    220       document.querySelector("#canvas-grid").appendChild(container)
    221       const canvas = document.createElement("canvas")
    222       container.appendChild(canvas)
    223       new Chart(canvas, {
    224         type: 'line',
    225         data: {
    226           datasets: datasets.map(dataset => ({
    227             borderColor: "#dc3912",
    228             fill: false,
    229             pointRadius: 0,
    230             ...dataset
    231           }))
    232         },
    233         options: getChartOptions(options)
    234       })
    235     }
    236 
    237     function getMainChartOptions() {
    238       return {
    239         scales: {
    240           xAxes: [{
    241             type: 'linear',
    242             position: 'bottom',
    243             scaleLabel: {
    244               display: true,
    245               labelString: 'surface distance travelled (in km)'
    246             }
    247           }],
    248           yAxes: [{
    249             type: 'linear',
    250             scaleLabel: {
    251               display: true,
    252               labelString: 'height above surface (in km)'
    253             }
    254           }]
    255         },
    256         animation: false
    257       }
    258     }
    259 
    260     function getChartOptions(options) {
    261       return {
    262         scales: {
    263           xAxes: [{
    264             type: 'linear',
    265             position: 'bottom',
    266             scaleLabel: {
    267               display: true,
    268               labelString: 'time (in s)'
    269             }
    270           }]
    271         },
    272         animation: false,
    273         ...options
    274       }
    275     }
    276 
    277     function getEarthChartOptions() {
    278       return {
    279         aspectRatio: 1,
    280         scales: {
    281           xAxes: [{
    282             type: 'linear',
    283             position: 'bottom',
    284             min: -8000,
    285             max: 8000,
    286             display: false
    287           }],
    288           yAxes: [{
    289             type: 'linear',
    290             min: -8000,
    291             max: 8000,
    292             display: false
    293           }]
    294         },
    295         legend: { display: false }
    296       }
    297     }
    298   </script>
    299 </body>
    300 
    301 </html>