From cd679620045e450fad0424f31b63fda35d9c5cab Mon Sep 17 00:00:00 2001
From: Ingmar Schoegl <ischoegl@lsu.edu>
Date: Fri, 9 Aug 2019 14:53:53 -0500
Subject: [PATCH] [Thermo] add ability to sort SolutionArray objects

---
 interfaces/cython/cantera/composite.py        | 26 ++++++++++++++++++-
 interfaces/cython/cantera/test/test_thermo.py | 22 ++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)

diff --git a/interfaces/cython/cantera/composite.py b/interfaces/cython/cantera/composite.py
index b5d661c67..2374ee792 100644
--- a/interfaces/cython/cantera/composite.py
+++ b/interfaces/cython/cantera/composite.py
@@ -492,7 +492,12 @@ class SolutionArray:
             shape = (shape,)
 
         if states is not None:
-            self._shape = np.shape(states)[:-1]
+            if isinstance(states, list):
+                self._shape = (len(states),)
+            elif isinstance(states, np.ndarray):
+                self._shape = np.shape(states)[:-1]
+            else:
+                raise TypeError('invalid type')
             self._states = states
         else:
             self._shape = tuple(shape)
@@ -603,6 +608,25 @@ class SolutionArray:
         self._indices.append(len(self._indices))
         self._shape = (len(self._indices),)
 
+    def sort(self, col, reverse=False):
+        """ 
+        Sort SolutionArray by column *col*.
+
+        :param col: Column that is used to sort the SolutionArray.
+        :param reverse: If True, the sorted list is reversed (descending order).
+        """
+        if len(self._shape) != 1:
+            raise TypeError("sort only works for 1D SolutionArray objects")
+
+        indices = np.argsort(getattr(self, col))
+        if reverse:
+            indices = indices[::-1]
+        self._states = [self._states[ix] for ix in indices]
+        for k, v in self._extra_arrays.items():
+            new = v[indices]
+            self._extra_arrays[k] = new
+            self._extra_lists[k] = list(new)
+
     def equilibrate(self, *args, **kwargs):
         """ See `ThermoPhase.equilibrate` """
         for index in self._indices:
diff --git a/interfaces/cython/cantera/test/test_thermo.py b/interfaces/cython/cantera/test/test_thermo.py
index 75eabf8d8..7c7ff4c1f 100644
--- a/interfaces/cython/cantera/test/test_thermo.py
+++ b/interfaces/cython/cantera/test/test_thermo.py
@@ -1637,6 +1637,28 @@ class TestSolutionArray(utilities.CanteraTest):
         states.TP = np.linspace(400, 500, 5), 101325
         self.assertArrayNear(states.X.squeeze(), np.ones(5))
 
+    def test_sort(self):
+        np.random.seed(0)
+        t = np.random.random(101)
+        T = np.linspace(300., 1000., 101)
+        P = ct.one_atm * (1. + 10.*np.random.random(101))
+
+        states = ct.SolutionArray(self.gas, 101, extra={'t': t})
+        states.TP = T, P
+
+        states.sort('t')
+        self.assertTrue((states.t[1:] - states.t[:-1] > 0).all())
+        self.assertFalse((states.T[1:] - states.T[:-1] > 0).all())
+        self.assertFalse(np.allclose(states.P, P))
+
+        states.sort('T')
+        self.assertFalse((states.t[1:] - states.t[:-1] > 0).all())
+        self.assertTrue((states.T[1:] - states.T[:-1] > 0).all())
+        self.assertTrue(np.allclose(states.P, P))
+
+        states.sort('T', reverse=True)
+        self.assertTrue((states.T[1:] - states.T[:-1] < 0).all())
+
     def test_set_equivalence_ratio(self):
         states = ct.SolutionArray(self.gas, 8)
         phi = np.linspace(.5, 2., 8)