Merge pull request #389 from QuantEcon/submatrix

MAINT: Use `np.ix_` to extract submatrix
QuantEcon · Feb 6, 2018 · 49ad385 · 49ad385
2 parents 53261bd + d1d48f8
commit 49ad385
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/quantecon/graph_tools.py b/quantecon/graph_tools.py
@@ -351,7 +351,7 @@ def subgraph(self, nodes):
             A DiGraph representing the subgraph.
 
         """
-        adj_matrix = self.csgraph[nodes, :][:, nodes]
+        adj_matrix = self.csgraph[np.ix_(nodes, nodes)]
 
         weighted = True  # To copy the dtype
 

diff --git a/quantecon/markov/core.py b/quantecon/markov/core.py
@@ -408,13 +408,11 @@ def _compute_stationary(self):
             rec_classes = self.recurrent_classes_indices
             stationary_dists = np.zeros((len(rec_classes), self.n))
             for i, rec_class in enumerate(rec_classes):
-                if not self.is_sparse:  # Dense
-                    stationary_dists[i, rec_class] = \
-                        gth_solve(self.P[rec_class, :][:, rec_class])
-                else:  # Sparse
-                    stationary_dists[i, rec_class] = \
-                        gth_solve(self.P[rec_class, :][:, rec_class].toarray(),
-                                  overwrite=True)
+                P_rec_class = self.P[np.ix_(rec_class, rec_class)]
+                if self.is_sparse:
+                    P_rec_class = P_rec_class.toarray()
+                stationary_dists[i, rec_class] = \
+                    gth_solve(P_rec_class, overwrite=True)
 
         self._stationary_dists = stationary_dists