Custom step method insists not free random variable

Hi all,

I’ve got some code running where I am hoping to improve the sampling by doing a custom sampling method but have run into a bit of a wall with it.

The relevant block of code is:

with pm.Model() as model_multinomial1:
        #M=2

        # Define priors
        #sd=0.05
        eta_sep=pm.TruncatedNormal("eta_sep",mu=0,sigma=1,lower=0.0,upper=1.0,shape=2) #array of
        eta=pm.Deterministic("eta",[0.5,0.5]+eta_sep*0.05)
        theta=pm.Deterministic("theta",2*pt.arccos(pt.sqrt(eta)))
        #priors for conciseness

        #sd=a_dev
        a_sep=pm.TruncatedNormal("a_sep", mu=0, sigma=1,lower=-np.pi,upper=np.pi, shape=2)  #array of priors for conciseness
        a=pm.Deterministic("a",[0,0]+a_sep*a_dev)

        sd=0.7
        b_sep=pm.Normal("b_sep", mu=0, sigma=1,shape=2) #array of priors for conciseness
        b=pm.Deterministic("b",[0.7,0.7]+b_sep*sd)

        #Have decided to open the Volt variable to allow additional freedom
        Volt=pm.Normal("Volt",mu=V_2_dist,sigma=0.1)
        #Volt=pm.Deterministic("Volt",pt.as_tensor(V_2_dist))
        
        phi=pm.Deterministic("phi",a[:,None]+b[:,None]*pm.math.sqr(Volt))
        
        #phi1 vanishes from p expression under simplification
        p1=pm.Deterministic("p1",-pm.math.sin(theta[0])*pm.math.sin(theta[1])*pm.math.cos(phi[0])/2 + pm.math.cos(theta[0])*pm.math.cos(theta[1])/2 + 1/2)
        p2=pm.Deterministic("p2",-pm.math.cos(theta[0] - theta[1])/4 - pm.math.cos(theta[0] + theta[1])/4 - pm.math.cos(-phi[0] + theta[0] + theta[1])/8 + pm.math.cos(phi[0] - theta[0] + theta[1])/8 + pm.math.cos(phi[0] + theta[0] - theta[1])/8 - pm.math.cos(phi[0] + theta[0] + theta[1])/8 + 1/2)

        P=pm.Deterministic("P",pm.math.stack([p1,p2],axis=-1))
        likelihood=pm.Multinomial("likelihood",n=1000,p=P[0],shape=(N,2),observed=data_2)

    print(model_multinomial1.free_RVs)
    
    with model_multinomial1:

        #stepmethod=[pm.NUTS([a_sep]),pm.NUTS([a_sep,b_sep]),pm.NUTS([a_sep,b_sep,eta_sep])]

        trace_multinomial_2_HMC = pm.sample(draws=int(1e3), chains=4, cores=cpucount,step=pm.Metropolis(vars=["a_sep"]), return_inferencedata=True)

But when I try to run this in my clean anaconda environment, the printing free random variables line gives me what I expect but then throws the following error:

If I don’t give any arguments (so no vars=…) in the custom step method then it runs fine but then that defeats the purpose of me trying to do some custom sampling! Could this be a genuine bug or am I misusing the step method argument of pm.sample()?

I believe you’re supposed to pass the variable itself (a_sep), rather than the string name ("a_sep")

That works! It’s strange because I did that originally and got the error so I tried putting it as a string but in retrying it just now it works as intended :man_facepalming: I’m not gonna look a gift horse in the mouth, thanks anyway :joy:

1 Like

@jessegrabowski Turns out that I think I have isolated the underlying problem that explains the experienced fluke. So in making a minimal working example i changed my code from 2 models with custom stepmethod to just one and when there is just one then custom sampling steps work as expected, but when i tried again with my code that had two models in it then the error was still there (maybe considered a bug but is perhaps moreso just a case of poor pymc practice)

So if I have two models like so based off the GLM example in the docs:

with Model() as model1:
        # Define priors
        sigma = HalfCauchy("sigma", beta=10)
        intercept = Normal("Intercept", 0, sigma=20)
        slope = Normal("slope", 0, sigma=20)

        # Define likelihood
        likelihood = Normal("y", mu=intercept + slope * x, sigma=sigma, observed=y)

with Model() as model2:
        # Define priors
        sigma = Normal("sigma", 0,sigma=1)
        intercept = Normal("Intercept", 0, sigma=20)
        slope = Normal("slope", 0, sigma=20)

        # Define likelihood
        likelihood = Normal("y", mu=intercept + slope * x, sigma=sigma, observed=y)

with model1:
        stepmethod=[pm.NUTS([sigma]),pm.NUTS([sigma,slope, intercept])]

        samples1=sample(1000,step=stepmethod)

with model2:
        stepmethod=[pm.NUTS([sigma]),pm.NUTS([sigma,slope]), pm.NUTS([sigma,slope, intercept])]

        samples2=sample(1000,step=stepmethod)

causes the error, but in changing say model 1 variable names so all cases of {mu,slope,intercept}->{mu1,slope1,intercept1} then resolves the error meaning that sharing variable names between models screws up custom sampling methods.

It’s not really anything to do with the custom step perse; you’re overwriting the python variables when you re-declare the models, so the name “sigma” is pointing to the wrong thing. You could instead access the variables by using the model itself like a dictionary, for example mode1['sigma'] will return the correct RV named “sigma” associated with model1

1 Like

Ahh that way of accessing will be a massive help then, save me having to recomb over some of the larger models and rewrite :sweat_smile:

1 Like