[Python] Allow SolutionArray to carry 'extra' variables

interfaces/cython/cantera/composite.py

@@ -272,7 +272,7 @@ class SolutionArray(object):
         slicing support.
     """
 
-    def __init__(self, phase, shape=(0,), states=None):
+    def __init__(self, phase, shape=(0,), states=None, extra=None):
         self._phase = phase
 
         if isinstance(shape, int):
@@ -282,7 +282,7 @@ def __init__(self, phase, shape=(0,), states=None):
             self._shape = states.shape[:-1]
             self._states = states
         else:
-            self._shape = shape
+            self._shape = tuple(shape)
             if len(shape) == 1:
                 S = [self._phase.state for _ in range(shape[0])]
             else:
@@ -297,11 +297,46 @@ def __init__(self, phase, shape=(0,), states=None):
             self._indices = list(np.ndindex(self._shape))
             self._output_dummy = self._states[..., 0]
 
+        self._extra_lists = {}
+        self._extra_arrays = {}
+        if isinstance(extra, dict):
+            for name, v in extra.items():
+                if not np.shape(v):
+                    self._extra_lists[name] = [v]*self._shape[0]
+                    self._extra_arrays[name] = np.array(self._extra_lists[name])
+                elif len(v) == self._shape[0]:
+                    self._extra_lists[name] = list(v)
+                else:
+                    raise ValueError("Unable to map extra SolutionArray"
+                                     "input for named {!r}".format(name))
+                self._extra_arrays[name] = np.array(self._extra_lists[name])
+
+        elif extra and self._shape == (0,):
+            for name in extra:
+                self._extra_lists[name] = []
+                self._extra_arrays[name] = np.array(())
+
+        elif extra:
+            raise ValueError("Initial values for extra properties must be"
+                " supplied in a dict if the SolutionArray is not initially"
+                " empty")
+
     def __getitem__(self, index):
         states = self._states[index]
         shape = states.shape[:-1]
         return SolutionArray(self._phase, shape, states)
 
+    def __getattr__(self, name):
+        if name not in self._extra_lists:
+            raise AttributeError("'{}' object has no attribute '{}'".format(
+                self.__class__.__name__, name))
+        L = self._extra_lists[name]
+        A = self._extra_arrays[name]
+        if len(L) != len(A):
+            A = np.array(L)
+            self._extra_arrays[name] = A
+        return A
+
     def append(self, state=None, **kwargs):
         """
         Append an element to the array with the specified state. Elements can
@@ -327,6 +362,9 @@ def append(self, state=None, **kwargs):
         if len(self._shape) != 1:
             raise IndexError("Can only append to 1D SolutionArray")
 
+        for name, value in self._extra_lists.items():
+            value.append(kwargs.pop(name))
+
         if state is not None:
             self._phase.state = state
 

interfaces/cython/cantera/examples/reactors/reactor1.py

@@ -13,15 +13,13 @@
 
 sim = ct.ReactorNet([r])
 time = 0.0
-times = np.zeros(100)
-states = ct.SolutionArray(gri3, 100)
+states = ct.SolutionArray(gri3, extra=['t'])
 
 print('%10s %10s %10s %14s' % ('t [s]','T [K]','P [Pa]','u [J/kg]'))
 for n in range(100):
     time += 1.e-5
     sim.advance(time)
-    times[n] = time * 1e3  # time in ms
-    states[n].TDY = r.thermo.TDY
+    states.append(r.thermo.state, t=time*1e3)
     print('%10.3e %10.3f %10.3f %14.6e' % (sim.time, r.T,
                                            r.thermo.P, r.thermo.u))
 
@@ -31,19 +29,19 @@
     import matplotlib.pyplot as plt
     plt.clf()
     plt.subplot(2, 2, 1)
-    plt.plot(times, states.T)
+    plt.plot(states.t, states.T)
     plt.xlabel('Time (ms)')
     plt.ylabel('Temperature (K)')
     plt.subplot(2, 2, 2)
-    plt.plot(times, states.X[:,gri3.species_index('OH')])
+    plt.plot(states.t, states.X[:,gri3.species_index('OH')])
     plt.xlabel('Time (ms)')
     plt.ylabel('OH Mole Fraction')
     plt.subplot(2, 2, 3)
-    plt.plot(times, states.X[:,gri3.species_index('H')])
+    plt.plot(states.t, states.X[:,gri3.species_index('H')])
     plt.xlabel('Time (ms)')
     plt.ylabel('H Mole Fraction')
     plt.subplot(2, 2, 4)
-    plt.plot(times, states.X[:,gri3.species_index('H2')])
+    plt.plot(states.t, states.X[:,gri3.species_index('H2')])
     plt.xlabel('Time (ms)')
     plt.ylabel('H2 Mole Fraction')
     plt.tight_layout()

interfaces/cython/cantera/examples/reactors/reactor2.py

@@ -67,21 +67,17 @@
 csvfile = csv.writer(outfile)
 csvfile.writerow(['time (s)','T1 (K)','P1 (Bar)','V1 (m3)',
                   'T2 (K)','P2 (Bar)','V2 (m3)'])
-states1 = ct.SolutionArray(ar)
-states2 = ct.SolutionArray(gri3)
-vol = np.zeros((n_steps, 2))
-tm = np.zeros(n_steps)
+states1 = ct.SolutionArray(ar, extra=['t', 'V'])
+states2 = ct.SolutionArray(gri3, extra=['t', 'V'])
 
 for n in range(n_steps):
     time += 4.e-4
     print(n, time, r2.T)
     sim.advance(time)
-    tm[n] = time
-    states1.append(r1.thermo.state)
-    states2.append(r2.thermo.state)
-    vol[n,:] = r1.volume, r2.volume
-    csvfile.writerow([tm[n], r1.thermo.T, r1.thermo.P, vol[n,0],
-                      r2.thermo.T, r2.thermo.P, vol[n,1]])
+    states1.append(r1.thermo.state, t=time, V=r1.volume)
+    states2.append(r2.thermo.state, t=time, V=r2.volume)
+    csvfile.writerow([time, r1.thermo.T, r1.thermo.P, r1.volume,
+                      r2.thermo.T, r2.thermo.P, r2.volume])
 outfile.close()
 print('Output written to file piston.csv')
 print('Directory: '+os.getcwd())
@@ -90,19 +86,19 @@
     import matplotlib.pyplot as plt
     plt.clf()
     plt.subplot(2,2,1)
-    h = plt.plot(tm, states1.T, 'g-', tm, states2.T, 'b-')
+    h = plt.plot(states1.t, states1.T, 'g-', states2.t, states2.T, 'b-')
     #plt.legend(['Reactor 1','Reactor 2'],2)
     plt.xlabel('Time (s)')
     plt.ylabel('Temperature (K)')
 
     plt.subplot(2,2,2)
-    plt.plot(tm, states1.P / 1e5, 'g-', tm, states2.P / 1e5, 'b-')
+    plt.plot(states1.t, states1.P / 1e5, 'g-', states2.t, states2.P / 1e5, 'b-')
     #plt.legend(['Reactor 1','Reactor 2'],2)
     plt.xlabel('Time (s)')
     plt.ylabel('Pressure (Bar)')
 
     plt.subplot(2,2,3)
-    plt.plot(tm, vol[:,0],'g-',tm, vol[:,1],'b-')
+    plt.plot(states1.t, states1.V, 'g-', states2.t, states2.V,'b-')
     #plt.legend(['Reactor 1','Reactor 2'],2)
     plt.xlabel('Time (s)')
     plt.ylabel('Volume (m$^3$)')