Fix standard errors in marginal_effects

[raashish1601]

marginal_effects (used by MarginalEffectsPlot) got two things wrong when computing the SE of each level's relative effect:

1. inverse_variance_weight returns (mean, variance), but the second value was stored as overall_sem and passed to relativize as the control SEM. So the variance of the overall mean was used where its standard error is expected.
2. relativize was called with cov_means=0.0, but the level's mean is computed from a subset of the same observations as the overall mean, so the two are correlated. With inverse variance weights, the covariance between a group's weighted mean and the overall weighted mean is the variance of the overall mean.

This PR passes np.sqrt(overall_var) as the control SEM and overall_var as cov_means.

To check the result, I simulated the arm means many times and compared the spread of the relative effect with the reported SE. With means [1, 2, 3, 4] and SEMs [0.05, 0.2, 0.1, 0.3] in two levels:

level	simulated SD	before	after
a	1.706	3.275	1.707
b	6.525	6.407	6.528

The expected values in test_marginal_effects changed accordingly. The Betas move slightly too, because the bias correction in relativize now sees the real control SEM. I also added test_marginal_effects_se_matches_simulation, which compares the SEs with a simulation like the one above. Both tests fail on main.

Ran ax/utils/stats/tests/test_statstools.py and ax/analysis/plotly/tests/test_marginal_effects.py (Linux, Python 3.12), and black and flake8 on the changed files.

[saitcakmak]

Thanks for the fix!

[meta-codesync]

@saitcakmak has imported this pull request. If you are a Meta employee, you can view this in D123916702.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 96.59%. Comparing base (a9998bd) to head (e2cd791).

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #5297      +/-   ##
==========================================
- Coverage   96.60%   96.59%   -0.01%     
==========================================
  Files         621      621              
  Lines       70952    70966      +14     
==========================================
+ Hits        68541    68550       +9     
- Misses       2411     2416       +5     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[meta-codesync]

@saitcakmak merged this pull request in f4e58e0.