Numpy array style access to state-space subsystems.

jgoppert · cwrowley · commit bdb3417d1413 · 2014-08-10T20:12:16.000-04:00
Signed-off-by: Clancy Rowley &lt;cwrowley@princeton.edu&gt;
diff --git a/src/statesp.py b/src/statesp.py
@@ -28,6 +28,7 @@
 StateSpace.feedback
 StateSpace.returnScipySignalLti
 StateSpace.append
+StateSpace.__getitem__
 _convertToStateSpace
 _rss_generate
 
@@ -560,6 +561,17 @@ def append(self, other):
         D[self.outputs:,self.inputs:] = other.D
         return StateSpace(A, B, C, D, self.dt)
 
+    def __getitem__(self, indices):
+        """Array style acces"""
+        if len(indices) != 2:
+            raise IOError('must provide indices of length 2 for state space')
+        i = indices[0]
+        j = indices[1]
+        return StateSpace(self.A,
+                self.B[:,j],
+                self.C[i,:],
+                self.D[i,j], self.dt)
+
 # TODO: add discrete time check
 def _convertToStateSpace(sys, **kw):
     """Convert a system to state space form (if needed).
diff --git a/tests/statesp_test.py b/tests/statesp_test.py
@@ -204,6 +204,25 @@ def testAppendTF(self):
         np.testing.assert_array_almost_equal(sys3c.A[3:,:3], np.zeros( (2, 3)) )
 
 
+    def testArrayAccessSS(self):
+
+        sys1 = StateSpace([[1., 2.], [3., 4.]],
+                [[5., 6.], [6., 8.]],
+                [[9., 10.], [11., 12.]],
+                [[13., 14.], [15., 16.]], 1)
+
+        sys1_11 = sys1[0,1]
+        np.testing.assert_array_almost_equal(sys1_11.A,
+                sys1.A)
+        np.testing.assert_array_almost_equal(sys1_11.B,
+                sys1.B[:,1])
+        np.testing.assert_array_almost_equal(sys1_11.C,
+                sys1.C[0,:])
+        np.testing.assert_array_almost_equal(sys1_11.D,
+                sys1.D[0,1])
+
+        assert sys1.dt == sys1_11.dt
+
 class TestRss(unittest.TestCase):
     """These are tests for the proper functionality of statesp.rss."""
     
