From d1d48f8e230cb1cf835f030b374976b3f18e31c7 Mon Sep 17 00:00:00 2001 From: Daisuke Oyama Date: Sun, 4 Feb 2018 22:22:13 +0900 Subject: [PATCH] MAINT: Use `np.ix_` to extract submatrix --- quantecon/graph_tools.py | 2 +- quantecon/markov/core.py | 12 +++++------- 2 files changed, 6 insertions(+), 8 deletions(-) diff --git a/quantecon/graph_tools.py b/quantecon/graph_tools.py index 2bc1cbcda..fe5463f83 100644 --- 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 index a8815a1e9..b7d5570f6 100644 --- 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