[Repo Assist] perf(gcm): vectorise inner loops in marginal_expectation

[github-actions]

🤖 This is an automated improvement from Repo Assist.

Summary

marginal_expectation() in dowhy/gcm/stats.py is the hot path for Shapley-based feature relevance, arrow strength, and intrinsic causal influence computations — it is evaluated once per coalition subset. This PR eliminates two O(batch_size) Python loops inside the function by replacing them with vectorised NumPy operations.

Changes

1. Input-filling loop → np.repeat

The old code assigned baseline values one baseline-sample at a time in a Python loop:

for index in range(adjusted_batch_size):
    inputs[
        index * n_f : (index + 1) * n_f, baseline_feature_indices
    ] = baseline_samples[offset + index, baseline_feature_indices]

The new code broadcasts all rows in one shot:

if len(baseline_feature_indices) > 0:
    inputs[:, baseline_feature_indices] = np.repeat(
        baseline_samples[offset : offset + adjusted_batch_size, :][:, baseline_feature_indices],
        n_f,
        axis=0,
    )

2. Result-aggregation loop → reshape + mean(axis=1)

The old code called np.mean() and sliced predictions inside a Python loop:

for index in range(adjusted_batch_size):
    if return_averaged_results:
        result[offset + index] = np.mean(predictions[...], axis=0)
    else:
        result[offset + index] = predictions[...]

The new code reshapes into (batch, n_f, ...) and reduces in one NumPy call:

pred_batched = predictions.reshape((adjusted_batch_size, n_f) + extra_dims)
if return_averaged_results:
    result[offset : offset + adjusted_batch_size] = list(pred_batched.mean(axis=1))
else:
    result[offset : offset + adjusted_batch_size] = list(pred_batched)

Impact

With the typical defaults (num_samples_baseline=500, max_batch_size=100), each call to marginal_expectation runs 5 batches of 100. The input-filling loop alone goes from 100 Python iterations × 5 batches = 500 iterations per call to 5 NumPy ops. Since this is evaluated once per coalition (O(2^n) exact or O(permutations × n) approximate), the saving compounds across many Shapley calls.

Correctness

Both changes are semantically equivalent to the original. The reshape handles scalar (1-D) and multi-output (2-D) prediction functions identically to the original per-sample slice/mean.

Test Status

• ✅ All 19 tests/gcm/test_stats.py tests pass
• ✅ All 21 tests/gcm/test_feature_relevance.py + test_intrinsic_influence.py + test_arrow_strength.py tests pass
• ✅ black --check and isort --check pass
• ⚠️ flake8 reports 3 E203 (whitespace before : in slices) and 1 E501 (long docstring line) — all 4 are pre-existing issues present on main before this change (verified with git stash)

Generated by 🌈 Repo Assist, see workflow run. Learn more.

To install this agentic workflow, run

gh aw add githubnext/agentics/workflows/repo-assist.md@11c9a2c442e519ff2b427bf58679f5a525353f76

[Copilot]

Pull request overview

This PR optimizes dowhy.gcm.stats.marginal_expectation()—a performance-critical routine used in Shapley-based computations in the GCM module—by removing per-sample Python loops and replacing them with NumPy vectorized operations.

Changes:

• Vectorizes filling of baseline feature columns within each batch (replacing a per-baseline loop).
• Vectorizes per-baseline aggregation of predictions via reshape(...).mean(axis=1) (replacing a per-baseline slice/mean loop).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[github-actions]

Commit pushed: b64f187

Generated by 🌈 Repo Assist, see workflow run. Learn more.

[bloebp]

Looks good to me!