updating pruning func for varlingam for faster processing

Signed-off-by: Sven <[email protected]>

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