diff --git a/causallearn/search/FCMBased/lingam/var_lingam.py b/causallearn/search/FCMBased/lingam/var_lingam.py index ecc7439a..0cb48c82 100644 --- a/causallearn/search/FCMBased/lingam/var_lingam.py +++ b/causallearn/search/FCMBased/lingam/var_lingam.py @@ -319,10 +319,14 @@ def _pruning(self, X, B_taus, causal_order): obj = np.zeros((len(blocks))) exp = np.zeros( (len(blocks), causal_order_no + n_features * self._lags)) - for j, block in enumerate(blocks): - obj[j] = block[0][i] - exp[j:] = np.concatenate( - [block[0][ancestor_indexes].flatten(), block[1:][:].flatten()], axis=0) + + # Create 3D array to hold flattened ancestor indices for each block + ancestor_indexes_flat = blocks[:, 0, ancestor_indexes].reshape(len(blocks), -1) + # Fill obj using advanced indexing + obj[:] = blocks[:, 0, i] + # Fill exp using advanced indexing + exp[:, :causal_order_no] = ancestor_indexes_flat + exp[:, causal_order_no:] = blocks[:, 1:].reshape(len(blocks), -1) # adaptive lasso gamma = 1.0