Add ZeroSumNormal distribution

[AlexAndorra]

This PR introduces the world famous ZeroSumNormal distribution, i.e a Normal distribution where one or several axes are constrained to sum to zero. By default, the last axis is constrained to sum to zero.

The zerosum_axes are always in the rightmost position, i.e zerosum_axes=2 means the two rightmost axes will be constrained to sum to zero (see examples below for more details).

⚠️ sigma has to be a scalar, to ensure the zero-sum constraint. The ability to specifiy a vector of sigma may be added in future versions.

Checklist

• Explain important implementation details 👆
• Make sure that the pre-commit linting/style checks pass.
• Docstrings
• Tests
• Test ZSN variance
• Test ZSN logp
• Test get_support_shape with 2D

---

Examples:

COORDS = {
    "regions": ["a", "b", "c"],
     "answers": ["yes", "no", "whatever", "don't understand question"],
}
with pm.Model(coords=COORDS) as m:
    # the zero sum axis will be 'answers'
    v = pm.ZeroSumNormal("v", dims=("regions", "answers"))

with pm.Model(coords=COORDS) as m:
    # the zero sum axes will be 'answers' and 'regions'
    v = pm.ZeroSumNormal("v", dims=("regions", "answers"), zerosum_axes=2)

with pm.Model(coords=COORDS) as m:
    # the zero sum axes will be the last two
    v = pm.ZeroSumNormal("v", shape=(3, 4, 5), zerosum_axes=2)

Major / Breaking Changes

• ...

Bugfixes / New features

• Implemented ZeroSumNormal distribution

Docs / Maintenance

• ...

[ricardoV94]

@lucianopaz @aseyboldt does any of the math require the summed distribution to be a Normal or could it be something else (e.g., StudentT)?

[codecov]

Codecov Report

Merging #6121 (5954e65) into main (e419d53) will increase coverage by 0.35%.
The diff coverage is 97.34%.

❗ Current head 5954e65 differs from pull request most recent head 3e72922. Consider uploading reports for the commit 3e72922 to get more accurate results

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #6121      +/-   ##
==========================================
+ Coverage   93.05%   93.40%   +0.35%     
==========================================
  Files          91      100       +9     
  Lines       20804    22138    +1334     
==========================================
+ Hits        19360    20679    +1319     
- Misses       1444     1459      +15     

Impacted Files	Coverage Δ
pymc/distributions/multivariate.py	92.24% <91.54%> (-0.07%)	⬇️
pymc/distributions/shape_utils.py	98.66% <97.14%> (-0.29%)	⬇️
pymc/tests/distributions/test_multivariate.py	99.44% <98.97%> (-0.06%)	⬇️
pymc/distributions/timeseries.py	82.56% <100.00%> (-0.96%)	⬇️
pymc/distributions/transforms.py	100.00% <100.00%> (ø)
pymc/tests/distributions/test_shape_utils.py	99.76% <100.00%> (+0.03%)	⬆️
pymc/tests/distributions/test_timeseries.py	95.45% <100.00%> (-0.34%)	⬇️
pymc/parallel_sampling.py	85.80% <0.00%> (-1.00%)	⬇️
pymc/tests/distributions/test_truncated.py	99.48% <0.00%> (-0.52%)	⬇️
pymc/data.py	80.08% <0.00%> (ø)
... and 13 more

[AlexAndorra]

Thanks for the first review @ricardoV94 !
I added some tests. How do they look like?

[ricardoV94]

This would make more sense in multivariate even though we are treating it as a scalar as a hack for the time being.

[AlexAndorra]

This would make more sense in multivariate even though we are treating it as a scalar as a hack for the time being.

Good point. I moved it to multivariate.py.

I also added all the tests mentioned and pushed everything. They pass locally. Let's see if they pass here 🤞

[lucianopaz]

@AlexAndorra i left a bunch of comments all around the place. I think that I’m missing the part where the RV asserts that the zero sum axis are all negative and non repeating. I also think that you need to add some tests for the logp of the distribution and for the variance of the draws

[lucianopaz]

We could maybe add a comment here saying that this is using a householder reflection plus a projection operator to move forward from the constrained space onto the zero sum manifold. I’ll look up our notes and write something here

[AlexAndorra]

Did you find your notes @lucianopaz ?

Stale