Add SymbolicDistribution and Censored distributions #5169

ricardoV94 · 2021-11-10T11:31:44Z

With aesara-devs/aeppl#22 we can now compute the logprob of censored Distributions.

This PR adds those to PyMC:

import pymc as pm

with pm.Model(rng_seeder=2021) as m:
    x = pm.Normal.dist()
    y = pm.Censored('y', x, -.1, .1)

print([float(y.eval()) for _ in range(5)])
# [0.1, 0.01990278315435156, 0.1, -0.1, -0.1]

y_val = m.rvs_to_values[y]
for y_test_val in [-.2, -.1, -.09, 0, 0.09, .1, .2]:
    print(y_test_val, m.logp({y_val: y_test_val}))
# -0.2 -inf
# -0.1 -0.7761545942065803
# -0.09 -0.9229885332046727
# 0 -0.9189385332046727
# 0.09 -0.9229885332046727
# 0.1 -0.7761545942065803
# 0.2 -inf

A large chunk of this PR concerns how to accommodate aeppl derived distributions with PyMC API seamlessly. To do this I started by duplicating pymc.Distribution into pymc.DerivedDistribution and changing things as necessary. Hopefully we will be able to use the same class in the end or at least refactor large chunks that are shared...

This may not be the best way... Suggestions are welcome

What works:

TODO:

codecov · 2021-11-17T15:59:37Z

Codecov Report

Merging #5169 (4fc7dc9) into main (600fe90) will increase coverage by 0.01%.
The diff coverage is 86.39%.

@@            Coverage Diff             @@
##             main    #5169      +/-   ##
==========================================
+ Coverage   79.08%   79.09%   +0.01%     
==========================================
  Files          87       88       +1     
  Lines       14394    14516     +122     
==========================================
+ Hits        11383    11482      +99     
- Misses       3011     3034      +23

Impacted Files	Coverage Δ
pymc/distributions/distribution.py	`90.07% <77.14%> (-5.16%)`	⬇️
pymc/distributions/censored.py	`92.45% <92.45%> (ø)`
pymc/distributions/__init__.py	`100.00% <100.00%> (ø)`
pymc/distributions/bound.py	`100.00% <100.00%> (ø)`
pymc/distributions/shape_utils.py	`98.87% <100.00%> (+<0.01%)`	⬆️
pymc/distributions/simulator.py	`86.76% <100.00%> (ø)`
pymc/util.py	`76.10% <100.00%> (+1.26%)`	⬆️
pymc/parallel_sampling.py	`86.33% <0.00%> (-1.00%)`	⬇️

pymc/distributions/censored.py

pymc/distributions/distribution.py

twiecki · 2022-01-03T12:17:59Z

@ricardoV94 This is becoming a blocker for time series, how much is left to do here?

ricardoV94 · 2022-01-03T12:19:50Z

I am getting back to this one this week!

ricardoV94 · 2022-01-03T12:37:50Z

In the meantime, suggestions for the base class name? DerivedDistribution was my first suggestion because Aeppl derives the logprob terms for these graphs, but there is not much "derivation" going on in PyMC. These are just symbolic graphs that somewhere have a RandomVariable as a node / input. Other ideas:

CompositeDistribution
SymbolicDistribution
... ?

twiecki · 2022-01-03T12:41:23Z

CustomDistribution? Otherwise I like SymbolicDistribution.

brandonwillard · 2022-01-03T17:40:06Z

What purpose are these *Distribution classes serving? Remember, the original Distribution class in v4 was just an all static method-based stub/placeholder for (temporary) backwards compatibility.

ricardoV94 · 2022-01-03T18:29:28Z

They serve the same goal the "standard" Distribution class is doing right now: basically manage all non distribution kwargs: shape/dims, transformed, observed; and behind the scenes magic things like value variables and rngs

brandonwillard · 2022-01-03T18:57:02Z

They serve the same goal the "standard" Distribution class is doing right now: basically manage all non distribution kwargs: shape/dims, transformed, observed; and behind the scenes magic things like value variables and rngs

The base Distribution class was only ever supposed to make the hand-off to the Model class, which would then handle all the generic tasks (e.g. value variables, RNGs, shape/dims, etc.), and the rest was just for backward compatibility (e.g. the now completely unnecessary Disribution.dist interface). Notice how instances of these classes are never made, and that only static/class/type-level functions are used.

By extending Distribution and adding more logic to those classes, it could become considerably more difficult to unwind those temporary backward compatibility-only choices, and they'll eventually become a permanent and confusing part of v4's design.

This isn't the time or place to start addressing all that, but these are the kinds of design considerations that need to go alongside changes/additions to the relevant areas of code (i.e. Distribution-related code).

pymc/distributions/distribution.py

* Fix invalid code example in docstrings * Rename distribution parameter to dist * Use `check_dist_not_registered` helper

ricardoV94 · 2022-01-08T07:31:00Z

@brandonwillard do you have other pointers for this PR? It is currently blocked by your old requested changes concerning docstrings

Request is stale

brandonwillard

@brandonwillard do you have other pointers for this PR? It is currently blocked by your old requested changes concerning docstrings

Sorry, I dismissed the old review; otherwise, I don't have any more input for this PR right now.

ghost · 2022-01-08T21:00:53Z

Thanks @ricardoV94! Can't wait to try these out.

twiecki · 2022-01-10T04:20:26Z

pymc/distributions/distribution.py

@@ -354,11 +355,268 @@ def dist(
        return rv_out


+class SymbolicDistribution:


Will those show up in the docs?

We also didn't mention it in the release notes.

They show in the docs (you can check the RTD job in this PR). I didn't see the SymbolicDist as being user facing, that's why it's not mentioned in the release notes.

I agree, but it's such an important new addition that I think it deserves some more attention.

ricardoV94 requested a review from brandonwillard November 10, 2021 11:31

ricardoV94 marked this pull request as draft November 10, 2021 11:32

ricardoV94 added enhancements v4 request discussion labels Nov 10, 2021

ricardoV94 force-pushed the censored_distributions branch 2 times, most recently from 7639c7f to 729c221 Compare November 17, 2021 15:44

ricardoV94 force-pushed the censored_distributions branch 4 times, most recently from ec92fed to ed7e2b0 Compare November 17, 2021 17:32

ricardoV94 commented Nov 17, 2021

View reviewed changes

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

fonnesbeck mentioned this pull request Nov 18, 2021

Design for bound, truncated and censored distributions #1864

Closed

ricardoV94 force-pushed the censored_distributions branch 6 times, most recently from a0ccd82 to 34e2fab Compare November 19, 2021 10:24

brandonwillard previously requested changes Nov 26, 2021

View reviewed changes

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

ricardoV94 mentioned this pull request Dec 13, 2021

Partially observed multivariate distributions not implemented in V4 #5260

Closed

ricardoV94 mentioned this pull request Dec 31, 2021

Start adding GRW RV Op and distribution #5298

Merged

ricardoV94 mentioned this pull request Jan 3, 2022

Refactor old Distribution base class #5308

Open

ricardoV94 force-pushed the censored_distributions branch 2 times, most recently from 2329885 to 785ee15 Compare January 4, 2022 12:28

ricardoV94 marked this pull request as ready for review January 4, 2022 12:28

ricardoV94 force-pushed the censored_distributions branch from 785ee15 to 81b3dc4 Compare January 4, 2022 12:34

twiecki reviewed Jan 4, 2022

View reviewed changes

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

twiecki reviewed Jan 4, 2022

View reviewed changes

pymc/distributions/distribution.py Show resolved Hide resolved

ricardoV94 force-pushed the censored_distributions branch 3 times, most recently from d9e7db0 to 3fae5c9 Compare January 4, 2022 13:26

ricardoV94 mentioned this pull request Jan 4, 2022

Make use of new pm.Censored in notebooks about censored data pymc-devs/pymc-examples#261

Closed

twiecki approved these changes Jan 4, 2022

View reviewed changes

ricardoV94 force-pushed the censored_distributions branch 2 times, most recently from 728d816 to 766ac8d Compare January 4, 2022 13:41

ricardoV94 added 4 commits January 4, 2022 14:41

Make get_moment Op agnostic

e30568d

Create base SymbolicDistribution

46a25c0

Add Censored distributions

2cdf282

Small fixes to pm.Bound

4fc7dc9

* Fix invalid code example in docstrings * Rename distribution parameter to dist * Use `check_dist_not_registered` helper

ricardoV94 force-pushed the censored_distributions branch from 766ac8d to 4fc7dc9 Compare January 4, 2022 13:42

ricardoV94 added the logprob label Jan 4, 2022

ricardoV94 requested a review from brandonwillard January 4, 2022 16:01

brandonwillard reviewed Jan 8, 2022

View reviewed changes

ricardoV94 merged commit 664a447 into pymc-devs:main Jan 8, 2022

twiecki changed the title ~~Add Censored distributions~~ Add SymbolicDistribution and Censored distributions Jan 9, 2022

twiecki reviewed Jan 10, 2022

View reviewed changes

sethaxen mentioned this pull request Jan 11, 2022

Constructing a censored distribution JuliaStats/Distributions.jl#1468

Closed

ricardoV94 deleted the censored_distributions branch January 13, 2022 09:58

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add SymbolicDistribution and Censored distributions #5169

Add SymbolicDistribution and Censored distributions #5169

ricardoV94 commented Nov 10, 2021 •

edited

Loading

codecov bot commented Nov 17, 2021 •

edited

Loading

twiecki commented Jan 3, 2022

ricardoV94 commented Jan 3, 2022

ricardoV94 commented Jan 3, 2022

twiecki commented Jan 3, 2022

brandonwillard commented Jan 3, 2022 •

edited

Loading

ricardoV94 commented Jan 3, 2022 •

edited

Loading

brandonwillard commented Jan 3, 2022 •

edited

Loading

ricardoV94 commented Jan 8, 2022

brandonwillard left a comment

ghost commented Jan 8, 2022

twiecki Jan 10, 2022

twiecki Jan 10, 2022

ricardoV94 Jan 10, 2022

twiecki Jan 10, 2022

		@@ -354,11 +355,268 @@ def dist(
		return rv_out


		class SymbolicDistribution:

Add SymbolicDistribution and Censored distributions #5169

Add SymbolicDistribution and Censored distributions #5169

Conversation

ricardoV94 commented Nov 10, 2021 • edited Loading

codecov bot commented Nov 17, 2021 • edited Loading

Codecov Report

twiecki commented Jan 3, 2022

ricardoV94 commented Jan 3, 2022

ricardoV94 commented Jan 3, 2022

twiecki commented Jan 3, 2022

brandonwillard commented Jan 3, 2022 • edited Loading

ricardoV94 commented Jan 3, 2022 • edited Loading

brandonwillard commented Jan 3, 2022 • edited Loading

ricardoV94 commented Jan 8, 2022

brandonwillard left a comment

Choose a reason for hiding this comment

ghost commented Jan 8, 2022

twiecki Jan 10, 2022

Choose a reason for hiding this comment

twiecki Jan 10, 2022

Choose a reason for hiding this comment

ricardoV94 Jan 10, 2022

Choose a reason for hiding this comment

twiecki Jan 10, 2022

Choose a reason for hiding this comment

ricardoV94 commented Nov 10, 2021 •

edited

Loading

codecov bot commented Nov 17, 2021 •

edited

Loading

brandonwillard commented Jan 3, 2022 •

edited

Loading

ricardoV94 commented Jan 3, 2022 •

edited

Loading

brandonwillard commented Jan 3, 2022 •

edited

Loading