Create Ordered Multinomial distribution

[AlexAndorra]

This PR enables the ordered logistic constraint on multinomial observed data (i.e aggregated by trial).
Currently, OrderedLogistic only accepts data in a disaggregated format (i.e Categorical observed data).

I also added the option of having the inferred probabilities in the trace (not only the cutpoints, which are hard to interpret anyway). As I'm guessing that's mostly what people care about when using this likelihood, the option is True by default (it's only useful to disable it when memory usage is an issue).

This probably needs a few tests (inspired by those of OrderedLogistic), but here is already a proof of concept:

true_cum_p = np.array([0.1, 0.15, 0.25, 0.50, 0.65, 0.90, 1.0])
true_p = np.hstack([true_cum_p[0], true_cum_p[1:] - true_cum_p[:-1]])
fake_elections = np.random.multinomial(n=1_000, pvals=true_p, size=60)

with pm.Model() as m:
    cutpoints = pm.Normal(
        "cutpoints",
        mu=np.arange(6) - 2.5,
        sigma=1.5,
        initval=np.arange(6) - 2.5,
        transform=pm.distributions.transforms.ordered,
    )

    pm.OrderedMultinomial(
        "results",
        eta=0.0,
        cutpoints=cutpoints,
        n=fake_elections.sum(1),
        observed=fake_elections,
    )

    trace = pm.sample()

Which does recover the true probabilities:

• No breaking changes
• New tests
• Run linting / style checks
• Mention in RELEASE-NOTES.md

[ricardoV94]

I also added the option of having the inferred probabilities in the trace (not only the cutpoints, which are hard to interpret anyway).

This would also be nice for the other Ordered distributions no?

[AlexAndorra]

This would also be nice for the other Ordered distributions no?

Just added that and addressed your other comments 🥳
Only the name issue and new tests remaining -- will handle this later

[ricardoV94]

Maybe it's easier to keep the two things separate. Have a helper class that creates the raw Ordered distribution and then, optionally, creates the deterministic p. You can retrieve p from the inputs that go into the variable returned from the raw Ordered distribution. Something along the lines of:

class _OrderedMultinomial(Multinomial):
  # standard implementation without the helper Deterministic

class OrderedMultinomial:
  def __new__(cls, name, *args, save_p=True, **kwargs):
    out_rv = _OrderedMultinomial(name, *args, **kwargs)
    if save_p:
      pm.Deterministic(f'{name}_p', out_rv.owner.inputs[4])  # I think this is the p vector in the multinomial
    return out_rv

  @classmethod
  def dist(cls, *args, **kwargs):
    return _OrderedMultinomial.dist(*args, **kwargs)

No need to even check if it's inside the model context, that's done by the default _OrderedMultinomial.

[ricardoV94]

Ignore the previous idea. I think you can overwrite __new__ to save name in self.name that's probably simpler. No, doesn't work either. The previous comment goes back to being my best idea :p

[AlexAndorra]

Agreed, I was thinking the same -- that's actually how I did it for LKJCholeskyCov back in the days

[AlexAndorra]

out_rv.owner.inputs[4]) does seem to work @ricardoV94 ! Small question though: how did you know that the p vector was in the 4th position? 😅

[ricardoV94]

out_rv.owner.inputs[4]) does seem to work @ricardoV94 ! Small question though: how did you know that the p vector was in the 4th position? sweat_smile

The RVs inputs are always of the form (RandomStateSharedVariable, size, dtype, param1, param2 ..., paramN)

import aesara.tensor as at
x = at.random.multinomial(5, np.ones(3)/3, size=2)
print(x.owner.inputs)

[RandomStateSharedVariable(<RandomState(MT19937) at 0x7F125671EC40>),
 TensorConstant{(1,) of 2},
 TensorConstant{4},
 TensorConstant{5},
 TensorConstant{(3,) of 0...3333333333}]

[AlexAndorra]

The RVs inputs are always of the form (RandomStateSharedVariable, size, dtype, param1, param2 ..., paramN)

Ooooh, ok, thanks! What does a dtype of TensorConstant{4} mean though? 🤔

[ricardoV94]

The RVs inputs are always of the form (RandomStateSharedVariable, size, dtype, param1, param2 ..., paramN)

Ooooh, ok, thanks! What does a dtype of TensorConstant{4} mean though? thinking

It's a numerical code for different dtypes ("float32", "float64", "int32", and so on). I don't know where they are defined though.

4 is probably for "int64" since that's the default dtype of discrete RVs

[AlexAndorra]

Ah ok, I wouldn't have guessed that 😅
So the only thing left for this PR is adding a few tests. I think I need to write the logpdf for that but don't know how it works yet for this distribution

[ricardoV94]

So the only thing left for this PR is adding a few tests. I think I need to write the logpdf for that but don't know how it works yet for this distribution

The logp shouldn't have to be tested as it's just the multinomial. The tests should just make sure the parameters that go into it are the correct ones. Have a look at the check_pymc_params_match_rv_op tests in test_random.py (guide), I think it does just what is needed.

Plus some tests for the helper Deterministic

[ricardoV94]

LGTM!

Some conflicts in the files it seems. Other than that it looks good.

[codecov]

Codecov Report

Merging #4773 (64b1ac8) into main (125256f) will increase coverage by 0.07%.
The diff coverage is 98.11%.

@@            Coverage Diff             @@
##             main    #4773      +/-   ##
==========================================
+ Coverage   71.97%   72.05%   +0.07%     
==========================================
  Files          85       85              
  Lines       13839    13872      +33     
==========================================
+ Hits         9961     9995      +34     
+ Misses       3878     3877       -1     

Impacted Files	Coverage Δ
pymc3/distributions/__init__.py	100.00% <ø> (ø)
pymc3/distributions/multivariate.py	58.31% <95.45%> (+1.43%)	⬆️
pymc3/distributions/discrete.py	99.00% <100.00%> (+0.05%)	⬆️
pymc3/step_methods/hmc/nuts.py	96.87% <0.00%> (-0.63%)	⬇️

[ricardoV94]

The failing tests in distributions_random.py are due to the test looking for the rv_op which doesn't exist in the wrapper class. Probably just need to overwrite the relevant part in the inherited test subclasses.

[AlexAndorra]

I just added back the rv_op in the wrapper class. Tests pass locally. Is that a good solution?

[ricardoV94]

Is that a good solution?

Actually I think testing the wrapped distribution might the best. Just test the _Ordered* classes instead.

This shouldn't require any changes to the wrapper classes

[AlexAndorra]

All tests are green 🥳
Just added a mention in RELEASE-NOTES.md

[ricardoV94]

Looks good! Left a minor suggestion that can be ignored.

[ricardoV94]

The failing test comes from #4771. It can be ignored.

Finally, do we need to add an entry in the docs for the new distribution or is it automatic? I added it

[AlexAndorra]

Merging 🍾 Thanks for your help and reviews @ricardoV94 !