[Python] Add species slicing support to SolutionArray

interfaces/cython/cantera/composite.py

@@ -254,6 +254,12 @@ class SolutionArray(object):
         >>> for i,j in np.ndindex(mu.shape):
         ...     # do something with mu[i,j]
 
+    Information about a subset of species may also be accessed, using
+    parentheses to specify the species:
+
+        >>> states('CH4').Y # -> mass fraction of CH4 in each state
+        >>> states('H2','O2').partial_molar_cp # -> cp for H2 and O2
+
     Properties and functions which are not dependent on the thermodynamic state
     act as pass-throughs to the underlying `Solution` object, and are not
     converted to arrays::
@@ -279,7 +285,7 @@ class SolutionArray(object):
             shape = (shape,)
 
         if states is not None:
-            self._shape = states.shape[:-1]
+            self._shape = np.shape(states)[:-1]
             self._states = states
         else:
             self._shape = tuple(shape)
@@ -337,6 +343,10 @@ class SolutionArray(object):
             self._extra_arrays[name] = A
         return A
 
+    def __call__(self, *species):
+        return SolutionArray(self._phase[species], states=self._states,
+                             extra=self._extra_lists)
+
     def append(self, state=None, **kwargs):
         """
         Append an element to the array with the specified state. Elements can
@@ -495,7 +505,7 @@ def _make_functions():
         def getter(self):
             a = np.empty(self._shape)
             b = np.empty(self._shape)
-            c = np.empty(self._shape + (self._phase.n_species,))
+            c = np.empty(self._shape + (self._phase.n_selected_species,))
             for index in self._indices:
                 self._phase.state = self._states[index]
                 a[index], b[index], c[index] = getattr(self._phase, name)
@@ -513,7 +523,7 @@ def _make_functions():
                     self._states[index][:] = self._phase.state
             else:
                 # composition is an array with trailing dimension n_species
-                C = np.empty(self._shape + (self._phase.n_species,))
+                C = np.empty(self._shape + (self._phase.n_selected_species,))
                 C[:] = XY
                 for index in self._indices:
                     self._phase.state = self._states[index]
@@ -531,7 +541,7 @@ def _make_functions():
         return np.empty(self._shape)
 
     def empty_species(self):
-        return np.empty(self._shape + (self._phase.n_species,))
+        return np.empty(self._shape + (self._phase.n_selected_species,))
 
     def empty_total_species(self):
         return np.empty(self._shape + (self._phase.n_total_species,))

interfaces/cython/cantera/test/test_thermo.py

@@ -1296,3 +1296,14 @@ class TestSolutionArray(utilities.CanteraTest):
 
         states.TP = np.linspace(400, 500, 5), 101325
         self.assertArrayNear(states.X.squeeze(), np.ones(5))
+
+    def test_species_slicing(self):
+        states = ct.SolutionArray(self.gas, (2,5))
+        states.TPX = np.linspace(500, 1000, 5), 2e5, 'H2:0.5, O2:0.4'
+        states.equilibrate('HP')
+        self.assertArrayNear(states('H2').X.squeeze(),
+                             states.X[...,self.gas.species_index('H2')])
+
+        kk = (self.gas.species_index('OH'), self.gas.species_index('O'))
+        self.assertArrayNear(states('OH','O').partial_molar_cp,
+                             states.partial_molar_cp[...,kk])

interfaces/cython/cantera/thermo.pyx

@@ -387,6 +387,13 @@ cdef class ThermoPhase(_SolutionBase):
         def __get__(self):
             return self.thermo.nSpecies()
 
+    property n_selected_species:
+        """
+        Number of species selected for output (by slicing of Solution object)
+        """
+        def __get__(self):
+            return self._selected_species.size or self.n_species
+
     def species_name(self, k):
         """Name of the species with index *k*."""
         return pystr(self.thermo.speciesName(k))