Plotting capability added (still to be improved)

AeroPython · Aug 9, 2016 · 07abcde · 07abcde
1 parent 5fe928e
commit 07abcde
Show file tree

Hide file tree

Showing 3 changed files with 153 additions and 7 deletions.
diff --git a/examples/oop_example.py b/examples/oop_example.py
@@ -8,14 +8,13 @@
 
 """
 import numpy as np
+import matplotlib.pyplot as plt
 
 from pyfme.aircrafts import Cessna310
-
 from pyfme.environment.environment import Environment
 from pyfme.environment.atmosphere import ISA1976
 from pyfme.environment.gravity import VerticalConstant
 from pyfme.environment.wind import NoWind
-
 from pyfme.models.systems import EulerFlatEarth
 from pyfme.simulator import BatchSimulation
 
@@ -35,17 +34,42 @@
 controls2trim = ['delta_elevator', 'delta_aileron', 'delta_rudder', 'delta_t']
 
 trimmed_controls, trimmed_system, outputs = aircraft.steady_state_flight_trim(
-    system, environment, not_trimmed_controls, TAS=120, gamma=+np.pi/180,
-    turn_rate=0.1, controls2trim=controls2trim)
+    system=system,
+    env=environment,
+    controls=not_trimmed_controls,
+    TAS=120, gamma=+0*np.pi/180, turn_rate=0.0,
+    controls2trim=controls2trim)
 
 my_simulation = BatchSimulation(aircraft, trimmed_system, environment)
 
-N = 100
-time = np.linspace(0, 10, N)
+tfin = 20  # seconds
+N = tfin * 100 + 1
+time = np.linspace(0, tfin, N)
 controls = {c_name: np.full((N,), trimmed_controls[c_name]) for
             c_name in trimmed_controls}
 
 my_simulation.set_controls(time, controls)
+
+par_list = ['x_earth', 'y_earth', 'height',
+            'psi', 'theta', 'phi',
+            'u', 'v', 'w',
+            'v_north', 'v_east', 'v_down',
+            'p', 'q', 'r',
+            'alpha', 'beta', 'TAS',
+            'T', 'p', 'rho']
+
+my_simulation.set_par_dict(par_list)
 my_simulation.run_simulation()
 
+print(my_simulation.par_dict)
+
+plt.style.use('ggplot')
+
+for ii in range(len(par_list) // 3):
+    three_params = par_list[3*ii:3*ii+3]
+    fig, ax = plt.subplots(3, 1)
+    for jj, par in enumerate(three_params):
+        ax[jj].plot(time, my_simulation.par_dict[par])
+        ax[jj].set_ylabel(par)
 
+plt.show()
diff --git a/src/pyfme/aircrafts/aircraft.py b/src/pyfme/aircrafts/aircraft.py
@@ -202,7 +202,7 @@ def steady_state_flight_trim(self, system, env, controls, TAS, gamma=0.,
         outputs = {'alpha': self.alpha, 'beta': self.beta,
                    'u': trim_system.u, 'v': trim_system.v, 'w': trim_system.w,
                    'p': trim_system.p, 'q': trim_system.q, 'r': trim_system.r,
-                   'psi': trim_system.psi, 'theta':trim_system.theta,
+                   'psi': trim_system.psi, 'theta': trim_system.theta,
                    'phi': trim_system.phi}
 
         trimmed_controls = controls

diff --git a/src/pyfme/simulator.py b/src/pyfme/simulator.py
@@ -8,6 +8,7 @@
 
 """
 from abc import abstractmethod
+import numpy as np
 
 
 class Simulation(object):
@@ -17,9 +18,37 @@ def __init__(self, aircraft, system, environment):
         self.system = system
         self.environment = environment
         self._time_step = 0
+        self.PAR_KEYS = {'T': self.environment.T,  # env
+                         'pressure': self.environment.p,
+                         'rho': self.environment.rho,
+                         'a': self.environment.a,
+                         'TAS': self.aircraft.TAS,  # aircraft
+                         'Mach': self.aircraft.Mach,
+                         'q_inf': self.aircraft.q_inf,
+                         'alpha': self.aircraft.alpha,
+                         'beta': self.aircraft.beta,
+                         'x_earth': self.system.x_earth,  # system
+                         'y_earth': self.system.y_earth,
+                         'z_earth': self.system.z_earth,
+                         'psi': self.system.psi,
+                         'theta': self.system.theta,
+                         'phi': self.system.phi,
+                         'u': self.system.u,
+                         'v': self.system.v,
+                         'w': self.system.w,
+                         'v_north': self.system.v_north,
+                         'v_east': self.system.v_east,
+                         'v_down': self.system.v_down,
+                         'p': self.system.p,
+                         'q': self.system.q,
+                         'r': self.system.r,
+                         'height': self.system.height
+                         }
+        self.par_dict = {}
 
     def time_step(self, dt):
 
+        self.save_current_par_dict()
         self._time_step += 1
         self.system.propagate(self.aircraft, dt)
         self.environment.update(self.system)
@@ -31,6 +60,14 @@ def time_step(self, dt):
     def _get_current_controls(self, ii):
         return
 
+    @abstractmethod
+    def set_par_dict(self, par_list):
+        return
+
+    @abstractmethod
+    def save_current_par_dict(self):
+        return
+
 
 class BatchSimulation(Simulation):
 
@@ -84,6 +121,52 @@ def run_simulation(self):
         for ii, t in enumerate(self.time[1:]):
             dt = t - self.time[ii]
             self.time_step(dt)
+        # Save last time step
+        self.save_current_par_dict()
+
+    def set_par_dict(self, par_list):
+
+        if self.time is None:
+            msg = "Set controls with BatchSimulation.set_controls before " \
+                  "setting the par_dict"
+            raise RuntimeError(msg)
+
+        for par_name in par_list:
+            if par_name in self.PAR_KEYS:
+                self.par_dict[par_name] = np.empty_like(self.time)
+            else:
+                msg = "{} not found in PAR_KEYS".format(par_name)
+                raise RuntimeWarning(msg)
+
+    def save_current_par_dict(self):
+        self.PAR_KEYS = {'T': self.environment.T,  # env
+                         'pressure': self.environment.p,
+                         'rho': self.environment.rho,
+                         'a': self.environment.a,
+                         'TAS': self.aircraft.TAS,  # aircraft
+                         'Mach': self.aircraft.Mach,
+                         'q_inf': self.aircraft.q_inf,
+                         'alpha': self.aircraft.alpha,
+                         'beta': self.aircraft.beta,
+                         'x_earth': self.system.x_earth,  # system
+                         'y_earth': self.system.y_earth,
+                         'z_earth': self.system.z_earth,
+                         'psi': self.system.psi,
+                         'theta': self.system.theta,
+                         'phi': self.system.phi,
+                         'u': self.system.u,
+                         'v': self.system.v,
+                         'w': self.system.w,
+                         'v_north': self.system.v_north,
+                         'v_east': self.system.v_east,
+                         'v_down': self.system.v_down,
+                         'p': self.system.p,
+                         'q': self.system.q,
+                         'r': self.system.r,
+                         'height': self.system.height
+                         }
+        for par_name, par_values in self.par_dict.items():
+            par_values[self._time_step] = self.PAR_KEYS[par_name]
 
 
 
@@ -96,3 +179,42 @@ def __init__(self, aircraft, system, environment):
     def _get_current_controls(self, ii):
         # Joystick reading
         raise NotImplementedError
+
+    def set_par_dict(self, par_list):
+
+        for par_name in par_list:
+            if par_name in self.PAR_KEYS:
+                self.par_dict[par_name] = []
+            else:
+                msg = "{} not found in PAR_KEYS".format(par_name)
+                raise RuntimeWarning(msg)
+
+    def save_current_par_dict(self):
+        self.PAR_KEYS = {'T': self.environment.T,  # env
+                         'pressure': self.environment.p,
+                         'rho': self.environment.rho,
+                         'a': self.environment.a,
+                         'TAS': self.aircraft.TAS,  # aircraft
+                         'Mach': self.aircraft.Mach,
+                         'q_inf': self.aircraft.q_inf,
+                         'alpha': self.aircraft.alpha,
+                         'beta': self.aircraft.beta,
+                         'x_earth': self.system.x_earth,  # system
+                         'y_earth': self.system.y_earth,
+                         'z_earth': self.system.z_earth,
+                         'psi': self.system.psi,
+                         'theta': self.system.theta,
+                         'phi': self.system.phi,
+                         'u': self.system.u,
+                         'v': self.system.v,
+                         'w': self.system.w,
+                         'v_north': self.system.v_north,
+                         'v_east': self.system.v_east,
+                         'v_down': self.system.v_down,
+                         'p': self.system.p,
+                         'q': self.system.q,
+                         'r': self.system.r,
+                         'height': self.system.height
+                         }
+        for par_name, par_values in self.par_dict:
+            par_values.append(self.PAR_KEYS[par_name])