asr_update

asr.py

@@ -23,6 +23,12 @@ class SubstModel:
         Q = normalise(Q)
         self.Q = Q
 
+    def __eq__(self, other):
+        return self.IRM == other.IRM
+
+    def __hash__(self):
+        return hash(tuple(self.IRM))
+
     def getTransitionMatrix(self, t):
         """ Compute the transition probability matrix for a branch length, t
          NOTE: Multiplication of matrices is matrix multiplication (not element-wise multiplication) """
@@ -57,7 +63,21 @@ class SubstModel:
         except ValueError:
             raise RuntimeError(f"Either {parent} or {child} is not in the substitution model's alphabet.")
 
-        return self.getTransitionMatrix(t)[parent_idx, child_idx]
+        # Check for cached TM for given model and t. Models are assumed equal if IRMs are identical
+        global tmCache  # Stores any transition probability matrix computed for any model
+        try:
+            if self in tmCache.keys():
+                if not t in tmCache[self].keys():
+                    tmCache[self][t] = self.getTransitionMatrix(t)
+            else:
+                tmCache[self] = {}
+                tmCache[self][t] = self.getTransitionMatrix(t)
+        except NameError:
+            tmCache = {}
+            tmCache[self] = {}
+            tmCache[self][t] = self.getTransitionMatrix(t)
+
+        return tmCache[self][t][parent_idx, child_idx]
 
 class PhyloBNet:
     """ Discrete Bayesian network. Implementation here is specifically used for representation of discrete characters