Add `ZeroSumNormal` distribution #6121

AlexAndorra · 2022-09-12T14:51:23Z

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

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 · 2022-09-12T15:01:51Z

pymc/distributions/continuous.py

+
+        return super().dist([sigma], zerosum_axes=zerosum_axes, **kwargs)
+
+    # TODO: This is if we want ZeroSum constraint on other dists than Normal


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

pymc/distributions/continuous.py

pymc/distributions/transforms.py

codecov · 2022-09-12T15:10:46Z

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 · 2022-09-12T15:34:45Z

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

pymc/tests/distributions/test_continuous.py

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

ricardoV94

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

pymc/distributions/continuous.py

AlexAndorra · 2022-09-15T13:37:41Z

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

pymc/distributions/multivariate.py

pymc/distributions/transforms.py

lucianopaz · 2022-09-16T19:17:10Z

pymc/distributions/transforms.py

+
+
+def extend_axis(array, axis):
+    n = array.shape[axis] + 1


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

Did you find your notes @lucianopaz ?

pymc/tests/distributions/test_multivariate.py

pymc/distributions/multivariate.py

pymc/tests/distributions/test_multivariate.py

pymc/distributions/multivariate.py

pymc/tests/distributions/test_multivariate.py

pymc/distributions/__init__.py

pymc/distributions/transforms.py

pymc/distributions/multivariate.py

pymc/distributions/transforms.py

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

pymc/distributions/transforms.py

Stale

Use None as default value for zerosum_axes

6260c84

AlexAndorra added enhancements major Include in major changes release notes section labels Sep 12, 2022

AlexAndorra requested a review from ricardoV94 September 12, 2022 14:51

AlexAndorra self-assigned this Sep 12, 2022

AlexAndorra requested review from aseyboldt and lucianopaz September 12, 2022 14:51

drbenvincent mentioned this pull request Sep 12, 2022

WIP ZeroSumNormal example notebook pymc-devs/pymc-examples#210

Draft

6 tasks

ricardoV94 reviewed Sep 12, 2022

View reviewed changes

pymc/distributions/continuous.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Sep 12, 2022

View reviewed changes

pymc/distributions/transforms.py Outdated Show resolved Hide resolved

Add tests for ZSN

af96016

ricardoV94 reviewed Sep 12, 2022

View reviewed changes

AlexAndorra and others added 5 commits September 12, 2022 20:51

Reorder dispatched functions

71e5651

Test pylint

3cadb26

Ignore type check on normalize_axis_tuple

a66c586

Disable mypy on import of normalize_axis_tuple

e3be495

Remove base class in tests

759de36

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

ricardoV94 reviewed Sep 14, 2022

View reviewed changes

pymc/distributions/continuous.py Outdated Show resolved Hide resolved

pymc/distributions/continuous.py Outdated Show resolved Hide resolved

AlexAndorra added 5 commits September 15, 2022 15:28

Use pytest parametrize

a5a1e45

Add pm.draw to tests

c9eea6e

Test moment

0582d7c

Add change size test

0bdcdd7

Move ZSN to multivariate.py

854ef4c

Move ZSN tests to test_multivariate.py

fd3aefa

lucianopaz previously requested changes Sep 16, 2022

View reviewed changes

Add check if zerosum_axes is iterable in dist method

e94e4f1

AlexAndorra requested a review from ricardoV94 September 30, 2022 14:31

ricardoV94 reviewed Sep 30, 2022

View reviewed changes

pymc/tests/distributions/test_multivariate.py Outdated Show resolved Hide resolved

ricardoV94 changed the title ~~Add ZeroSumNormal distribution 🔥~~ Add ZeroSumNormal distribution Oct 5, 2022

AlexAndorra added 7 commits October 5, 2022 18:25

Add ZSN logp test

f363118

Fix test_inconsistent_steps_and_shape

64eca5c

Integrate review comments

c5e76c9

Solve freaking pre-commit issues

08c9df0

Put assert_zerosum_axes at top of test class

c120f7e

Improve error message of get_support_shape

ba5f3a1

Nicer format for ZSN logp test

48dafe9

AlexAndorra requested a review from ricardoV94 October 5, 2022 17:40

ricardoV94 reviewed Oct 5, 2022

View reviewed changes

pymc/distributions/multivariate.py Outdated Show resolved Hide resolved

pymc/tests/distributions/test_multivariate.py Outdated Show resolved Hide resolved

pymc/distributions/__init__.py Show resolved Hide resolved

AlexAndorra added 4 commits October 6, 2022 10:13

Increase tolerance for test_zsn_variance

6612a24

Add ZSN to docs

6b07a2a

Refactor ZSN docs

135ed47

Better latex in ZSN docs

cba0187

ricardoV94 requested changes Oct 6, 2022

View reviewed changes

pymc/distributions/transforms.py Show resolved Hide resolved

pymc/distributions/multivariate.py Outdated Show resolved Hide resolved

Add ZeroSumTransform to docs

566f308

ricardoV94 requested changes Oct 7, 2022

View reviewed changes

pymc/distributions/transforms.py Outdated Show resolved Hide resolved

Remove mention of default value in ZS transform docs

5954e65

Co-authored-by: Ricardo Vieira <28983449+ricardoV94@users.noreply.github.com>

AlexAndorra requested a review from ricardoV94 October 7, 2022 14:08

ricardoV94 reviewed Oct 7, 2022

View reviewed changes

pymc/distributions/transforms.py Outdated Show resolved Hide resolved

Update pymc/distributions/transforms.py

3e72922

ricardoV94 approved these changes Oct 7, 2022

View reviewed changes

ricardoV94 removed the major Include in major changes release notes section label Oct 7, 2022

ricardoV94 force-pushed the add-zerosumnormal branch from 9ef5240 to 3e72922 Compare October 7, 2022 14:34

ricardoV94 merged commit 9aeb6b5 into main Oct 7, 2022

ricardoV94 deleted the add-zerosumnormal branch November 2, 2022 12:25

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `ZeroSumNormal` distribution #6121

Add `ZeroSumNormal` distribution #6121

AlexAndorra commented Sep 12, 2022 •

edited by ricardoV94

Loading

ricardoV94 Sep 12, 2022

codecov bot commented Sep 12, 2022 •

edited

Loading

AlexAndorra commented Sep 12, 2022

ricardoV94 left a comment

AlexAndorra commented Sep 15, 2022

lucianopaz left a comment

lucianopaz Sep 16, 2022

AlexAndorra Sep 18, 2022


		return super().dist([sigma], zerosum_axes=zerosum_axes, **kwargs)

		# TODO: This is if we want ZeroSum constraint on other dists than Normal

Add ZeroSumNormal distribution #6121

Add ZeroSumNormal distribution #6121

Conversation

AlexAndorra commented Sep 12, 2022 • edited by ricardoV94 Loading

Major / Breaking Changes

Bugfixes / New features

Docs / Maintenance

ricardoV94 Sep 12, 2022

Choose a reason for hiding this comment

codecov bot commented Sep 12, 2022 • edited Loading

Codecov Report

AlexAndorra commented Sep 12, 2022

ricardoV94 left a comment

Choose a reason for hiding this comment

AlexAndorra commented Sep 15, 2022

lucianopaz left a comment

Choose a reason for hiding this comment

lucianopaz Sep 16, 2022

Choose a reason for hiding this comment

AlexAndorra Sep 18, 2022

Choose a reason for hiding this comment

Add `ZeroSumNormal` distribution #6121

Add `ZeroSumNormal` distribution #6121

AlexAndorra commented Sep 12, 2022 •

edited by ricardoV94

Loading

codecov bot commented Sep 12, 2022 •

edited

Loading