Stacking distributions or changing distribution after assignment

Questions v5
1

Hi,

I have some code that looks something like this:

mu_coefs = np.random.random((3,3,2))
sigma_coefs = np.ones((3,3,2))
coefs = pm.Normal("coefs", mu=mu_coefs, sigma=sigma_coefs, dims=(3,3,2))

and sometimes i would like to do something like this:

BoundNormal = pm.Bound(pm.Normal, lower=0.0)
coefs[1,1] = BoundNormal("test", mu=1, sigma=0.5)

but it gives an error:

TypeError: 'TensorVariable' object does not support item assignment

Ideally, solution would be to “stack” distributions, for example, if I could create and array of distributions and then merge them into one, I would be over the moon.
example:

coefs = stack([[pm.Normal(),pm.Normal(),BoundNormal()],
      [pm.Normal(),pm.Normal(),BoundNormal()]])

Any tips?

Best

2

Try pm.math.stack/concatenate. However note that you want to use as few individual variables as possible and make use of shape to obtain multiple independent variables from the same distribution due to vectorization and having to track less variables.

So while you can do:

coeffs = pm.math.stack([
  pm.Normal("x1", 0, 1),
  pm.Normal("x2", 1, 1),
  pm.Normal("x3", 2, 1),
  pm.HalfNormal("x4", 1),
  pm.HalfNormal("x5", 1)
])

You are probably better of doing:

coeffs = pm.math.concatenate([
  pm.Normal("x1-x3", [0, 1, 2], 1, shape=3),
  pm.HalfNormal("x4-x5", 1, shape=2),
], axis=0)  # I never remember which axis it should be

More details here: Distribution Dimensionality — PyMC 4.1.7 documentation

3

Is it possible to retain the same dimensionality?
Im trying this:

count = 0
coefs = []
for m,s in zip(mu_coefs.flatten(), sigma_coefs.flatten()):
    temp = pm.Normal(str(count), mu=m, sigma=s)
    count += 1
    coefs.append(temp)
lag_coefs = pm.math.stack(coefs,dims=(2,3,3))

which works but the rest of the code expects dimensions to be (2,3,3).

I checked and this:

from functools import partial
draw = partial(pm.draw, random_seed=2137)
draw(lag_coefs)

prints flat array - 18x1

4

You can use reshape, but what is the meaning of the dimensions? You don’t want to mix them wrongly.

lag_coefs = lagcoefs.reshape(2, 3, 3)
5

basically I want to have this:

but with additional restrictions on priors, sometimes you may want to have a variable to be set to 0 or maybe bounded below or above 0.

Sadly, applying the above + reshape does not work (dimensions seems to be fine) - it crashes with:

SamplingError: Initial evaluation of model at starting point failed!
Starting values:
{'intercept': array([-0.58209429, -0.926367  , -0.9039506 ]), 'noise': array([ 0.22777612, -0.78495135, -0.69520265]), '0': array(-0.2404607), '1': array(-0.15676309), '2': array(0.36900768), '3': array(0.68119166), '4': array(0.68974787), '5': array(1.12546836), '6': array(1.58719748), '7': array(-0.37563589), '8': array(-0.43870746), '9': array(0.19989166), '10': array(-0.17859524), '11': array(1.60743432), '12': array(0.0366893), '13': array(1.40441257), '14': array(0.28183592), '15': array(0.04049973), '16': array(0.90593601), '17': array(1.1111694)}

Initial evaluation results:
{'intercept': -3.76, 'noise': -3.33, '0': -3.61, '1': -0.86, '2': 0.25, '3': -0.26, '4': -0.72, '5': -0.82, '6': -1.09, '7': -5.58, '8': -15.93, '9': -0.98, '10': -1.49, '11': -1.32, '12': -0.76, '13': -0.97, '14': 0.48, '15': -0.56, '16': -0.69, '17': -0.77, 'obs': -inf}

And this is the code I’m trying to run:

lags = 2
mu_coefs = np.array([[[ 0.06232432,  0.46381767,  0.37192259],
        [ 0.6867632 ,  1.4492859 ,  0.26156156],
        [ 1.01963577,  1.19915437,  0.28733603]],

       [[ 0.03062758,  0.18815452,  0.64496428],
        [ 0.36147662,  0.02588367,  0.69575532],
        [0.09893386,  0.23010376,  0.57175569]]])
sigma_coefs = mu_coefs

coords={
    "lags": reversed(range(-lags, 0)),
    "vars": ("fedfunds", "inflation","ogap"),
    "cross_vars": ("fedfunds", "inflation","ogap"),
    "time": range(len(data) - lags),
}

with pm.Model(coords=coords) as BVAR_model:
    intercept = pm.Normal("intercept", mu=0, sigma=1, dims=("vars",))
    
    #as in tutorial
    #lag_coefs = pm.Normal("lag_coefs", mu=mu_coefs, sigma=sigma_coefs, dims=("lags", "vars", "cross_vars"))
    noise = pm.Normal("noise",mu = np.zeros(3), sigma = np.ones(3), dims=("vars",)) #SUM_OLS
    
    ##########my modification
    count = 0
    coefs = []
    for m,s in zip(mu_coefs.flatten(), sigma_coefs.flatten()):
        temp = pm.Normal(str(count), mu=m, sigma=s)
        count += 1
        coefs.append(temp)
    lag_coefs = pm.math.stack(coefs).reshape((2, 3, 3))
  ###########
    ar_fedfunds = pm.math.sum(
        [pm.math.sum(lag_coefs[i, 0] * data.values[lags-(i+1): -(i+1)], axis=-1)
        for i in range(lags)], axis=0) 

    ar_inflation = pm.math.sum([
        pm.math.sum(lag_coefs[i, 1] * data.values[lags-(i+1): -(i+1)], axis=-1)
        for i in range(lags)
    ], axis=0)   

    ar_ogap = pm.math.sum([
        pm.math.sum(lag_coefs[i, 2] * data.values[lags-(i+1): -(i+1)], axis=-1)
        for i in range(lags)
    ], axis=0)  

    mean = intercept + pm.math.stack([ar_fedfunds, ar_inflation,ar_ogap], axis=-1)

    obs = pm.Normal("obs", mu=mean, sigma=noise, observed=data[lags:], dims=("time", "vars"))

without the modification, sampling error does not appear.

6

I didn’t look careful at your code to see why the logp becomes -inf.

For the more general question, you can use set_subtensor: Basic Tensor Functionality — Aesara 0+untagged.50.g1d899ee.dirty documentation

7

Thanks a lot.

lag_coefs = aesara.tensor.subtensor.set_subtensor(lag_coefs[1,1,1], pm.Normal("test",mu=1,sigma=1))