graph_manager.py 文件源码

def __getitem__(self, thing: Any) -> sparse.coo_matrix:
        if type(thing) is slice or type(thing) is np.ndarray or type(thing) is int:
            gm = GraphManager(None, axis=self.axis)
            for key, g in self.items():
                # Slice the graph matrix properly without making it dense
                (a, b, w) = (g.row, g.col, g.data)
                indices = np.arange(g.shape[0])[thing]
                mask = np.logical_and(np.in1d(a, indices), np.in1d(b, indices))
                a = a[mask]
                b = b[mask]
                w = w[mask]
                d = dict(zip(np.sort(indices), np.arange(indices.shape[0])))
                a = np.array([d[x] for x in a])
                b = np.array([d[x] for x in b])
                gm[key] = sparse.coo_matrix((w, (a, b)), shape=(len(indices), len(indices)))
            return gm
        else:
            return self.__getattr__(thing)