From 97f201b132988a5a842a860844e2df8ae297ab01 Mon Sep 17 00:00:00 2001
From: Hung Pham <hungpham2511@gmail.com>
Date: Thu, 22 Nov 2018 10:55:14 +0800
Subject: [PATCH] implement indexing for transfer function

---
 control/tests/xferfcn_test.py | 23 ++++++++++++++++++--
 control/xferfcn.py            | 40 +++++++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/control/tests/xferfcn_test.py b/control/tests/xferfcn_test.py
index 204c6dfd8..0d208c6db 100644
--- a/control/tests/xferfcn_test.py
+++ b/control/tests/xferfcn_test.py
@@ -5,8 +5,8 @@
 
 import unittest
 import numpy as np
-from control.statesp import StateSpace, _convertToStateSpace
-from control.xferfcn import TransferFunction, _convertToTransferFunction
+from control.statesp import StateSpace, _convertToStateSpace, rss
+from control.xferfcn import TransferFunction, _convertToTransferFunction, ss2tf
 from control.lti import evalfr
 from control.exception import slycot_check
 # from control.lti import isdtime
@@ -536,6 +536,25 @@ def testMIMO(self):
         np.testing.assert_array_almost_equal(H.num[1][0], H2.num[1][0])
         np.testing.assert_array_almost_equal(H.den[1][0], H2.den[1][0])
 
+    def testIndexing(self):
+        tm = ss2tf(rss(5, 3, 3))
+
+        # scalar indexing
+        sys01 = tm[0, 1]
+        np.testing.assert_array_almost_equal(sys01.num[0][0], tm.num[0][1])
+        np.testing.assert_array_almost_equal(sys01.den[0][0], tm.den[0][1])
+
+        # slice indexing
+        sys = tm[:2, 1:3]
+        np.testing.assert_array_almost_equal(sys.num[0][0], tm.num[0][1])
+        np.testing.assert_array_almost_equal(sys.den[0][0], tm.den[0][1])
+        np.testing.assert_array_almost_equal(sys.num[0][1], tm.num[0][2])
+        np.testing.assert_array_almost_equal(sys.den[0][1], tm.den[0][2])
+        np.testing.assert_array_almost_equal(sys.num[1][0], tm.num[1][1])
+        np.testing.assert_array_almost_equal(sys.den[1][0], tm.den[1][1])
+        np.testing.assert_array_almost_equal(sys.num[1][1], tm.num[1][2])
+        np.testing.assert_array_almost_equal(sys.den[1][1], tm.den[1][2])
+
     def testMatrixMult(self):
         """MIMO transfer functions should be multiplyable by constant
         matrices"""
diff --git a/control/xferfcn.py b/control/xferfcn.py
index e8fe084f8..b8a48c0a7 100644
--- a/control/xferfcn.py
+++ b/control/xferfcn.py
@@ -491,6 +491,46 @@ def __pow__(self, other):
         if other < 0:
             return (TransferFunction([1], [1]) / self) * (self**(other+1))
 
+    def __getitem__(self, key):
+        key1, key2 = key
+
+        # pre-process
+        if isinstance(key1, int):
+            key1 = slice(key1, key1 + 1, 1)
+        if isinstance(key2, int):
+            key2 = slice(key2, key2 + 1, 1)
+        # dim1
+        start1, stop1, step1 = key1.start, key1.stop, key1.step
+        if step1 is None:
+            step1 = 1
+        if start1 is None:
+            start1 = 0
+        if stop1 is None:
+            stop1 = len(self.num)
+        # dim1
+        start2, stop2, step2 = key2.start, key2.stop, key2.step
+        if step2 is None:
+            step2 = 1
+        if start2 is None:
+            start2 = 0
+        if stop2 is None:
+            stop2 = len(self.num[0])
+
+        num = []
+        den = []
+        for i in range(start1, stop1, step1):
+            num_i = []
+            den_i = []
+            for j in range(start2, stop2, step2):
+                num_i.append(self.num[i][j])
+                den_i.append(self.den[i][j])
+            num.append(num_i)
+            den.append(den_i)
+        if self.isctime():
+            return TransferFunction(num, den)
+        else:
+            return TransferFunction(num, den, self.dt)
+
     def evalfr(self, omega):
         """Evaluate a transfer function at a single angular frequency.