allow OrderedProbit and OrderedLogit to take vector inputs

[drbenvincent]

At the moment, as far as I can tell, OrderedProbit and OrderedLogit break if you try to provide vector inputs for eta and sigma parameters.

For example, if you have data like this

where your data is encoded as a vector of y values alongside a vector of grp_idx, then it looks like you cannot do this...

Approach 1 - better but doesn't work

with pm.Model() as model1:
    theta = constrainedUniform(K)
    mu = pm.Normal("mu", mu=K/2, sigma=K, shape=2)
    sigma = pm.HalfNormal("sigma", 1, shape=2)
    pm.OrderedProbit("y_obs", 
                     cutpoints=theta, 
                     eta=mu[grp_idx], 
                     sigma=sigma[grp_idx], 
                     observed=y)

We have a common set of cutpoints theta, but each of the groups has different eta and sigma. We try to pass in a vector of eta and sigma values using advanced indexing. But this does not work 🙁

Approach 2 - split into separate likelihood terms

If we cannot pass in vectors to eta or sigma, then the solution that works is to encode the data differently so that you now have observations from one group as y1 and observations from another group as y2

with pm.Model() as model2:
    theta = constrainedUniform(K)
    mu = pm.Normal("mu", mu=K/2, sigma=K, shape=2)
    sigma = pm.HalfNormal("sigma", 1, shape=2)
    pm.OrderedProbit("y1_obs", 
                     cutpoints=theta, 
                     eta=mu[0], 
                     sigma=sigma[0], 
                     observed=y1)
    pm.OrderedProbit("y2_obs", 
                     cutpoints=theta, 
                     eta=mu[1], 
                     sigma=sigma[1], 
                     observed=y2)

This approach does work, but makes model specification cumbersome, and doesn't scale well with more groups.

Not crucial for this issue, but the constrainedUniform is a proposed new distribution, outlined in pymc-devs/pymc-extras#32.

[danhphan]

Hi @drbenvincent and @ricardoV94, happy to work on this task. I will aim for the Approach 1 that allows vector inputs for eta and sigma parameters. Let's me know if you have any suggestion. Many thanks.

[twiecki]

Great stuff @danhphan, assigned to you.

[ricardoV94]

@danhphan Let us know if you need help with this one

[danhphan]

Hi @ricardoV94, sorry it is still WIP.

I have tested this code:

df = pd.read_csv('data/OrdinalProbitData1.csv', dtype={'Y': 'category'})
df.Y = df.Y.astype(int)
grp_idx = pd.Categorical(df.X).codes
K = df.Y.nunique()

with pm.Model() as model21:
    
    cutpoints = pm.Normal("cutpoints", 0.0, 1.5, shape=K-1,
                  transform=pm.distributions.transforms.ordered,
                  testval=np.arange(K-1))
    
    mu = pm.Normal("mu", mu=K/2, sigma=K, shape=2)
    sigma = pm.HalfNormal("sigma", 1, shape=2)
    
    yobs1 = pm.OrderedLogistic("y_obs1", 
                     cutpoints=cutpoints, 
                     eta=mu[grp_idx], 
                     observed=df.Y-1)
    
    yobs2 = pm.OrderedProbit("y_obs2", 
                     cutpoints=cutpoints, 
                     eta=mu[grp_idx], 
                     sigma=sigma[grp_idx], 
                     observed=df.Y-1)
    
    trace21 = pm.sample()
    
az.summary(trace21, var_names= ["cutpoints", "mu", "sigma"])

And it work for OrderedLogistic but not for OrderedProbit. The good news is that I have found where to fix: it is in dist method of _OrderedProbit class: We need to reshape the sigma to : sigma.reshape(shape=(sigma.eval().shape[0],1)).

For example, if sigma's shape is (88,) => need to convert it to shape (88,1).

However, I am not sure this is a good solution? as I am still learning on aesara TensorVariable.

### discrete.py # Line 1927
class _OrderedProbit(Categorical):
    r"""
    Underlying class for ordered probit distributions.
    See docs for the OrderedProbit wrapper class for more details on how to use it in models.
    """
    rv_op = categorical

    @classmethod
    def dist(cls, eta, cutpoints, sigma=1, *args, **kwargs):
        eta = at.as_tensor_variable(floatX(eta))
        cutpoints = at.as_tensor_variable(cutpoints)

        probits = at.shape_padright(eta) - cutpoints
        _log_p = at.concatenate(
            [
                at.shape_padright(normal_lccdf(0, sigma, probits[..., 0])),
                log_diff_normal_cdf(0, sigma, probits[..., :-1], probits[..., 1:]), # NEED TO FIX sigma SHAPE
                at.shape_padright(normal_lcdf(0, sigma, probits[..., -1])),
            ],
            axis=-1,
        )
        _log_p = at.as_tensor_variable(floatX(_log_p))
        p = at.exp(_log_p)
 I am still learning on aesara TensorVariable.
        return super().dist(p, *args, **kwargs)

[ricardoV94]

@danhphan the reshape change you are referring to can easly be done with x = x[..., np.newaxis] (or the shape_padright that the other lines are doing), which adds a dimension of 1 in the last position (in both numpy and aesara arrays).

[danhphan]

Hi @ricardoV94, that's good to know, so it is similar to numpy arrays. 😄

I will update the code in _OrderedProbit class and write some tests as well.

Cheers 👯

[drbenvincent]

This is all looking great. FYI I'm planning on writing up an example for the pymc-examples repo when #5418 is merged. Great work.

[danhphan]

Hi, thank you @drbenvincent, glad that it helps 😄

[ricardoV94]

Closed via #5418