pymc3: Truncated Normal mixture

I want to make a mixture of two TruncatedNormal distributions in pymc3.

I am trying to modify this piece of documentation. See example #2 for Poisson.

import pymc3 as pm
import numpy as np

x_obs = [x for x in np.random.randn(250) if 1>x>0] + [x for x in np.random.randn(250) if 1.2>x>0.2]
x = np.array(x_obs)
with pm.Model() as model:
    w = pm.Dirichlet("w", [1,1])
    mu1 = pm.TruncatedNormal("mu1", mu=x.mean(), sigma=1, lower=0, upper=1)
    tau1 = pm.Gamma("tau1", mu=1/x.std()**2, sigma=1.0)
    mu2 = pm.TruncatedNormal("mu2", mu=x.mean(), sigma=1, lower=0, upper=1)
    tau2 = pm.Gamma("tau2", mu=1/x.std()**2, sigma=1.0)
    tn1 = pm.TruncatedNormal.dist("tn1",
                                  mu=mu1,
                                  tau=tau1,
                                  lower=0,
                                  upper=1)
    tn2 = pm.TruncatedNormal.dist("tn2",
                                  mu=mu2,
                                  tau=tau2,
                                  lower=0,
                                  upper=1)
    like = pm.Mixture(name="like",
                      w=w,
                      comp_dists=[tn1, tn2],
                      observed = x)

The error I am getting is:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-7-0287db35d04c> in <module>
      5 #     mu2 = pm.TruncatedNormal("mu2", mu=x.mean(), sigma=1, lower=0, upper=1)
      6 #     tau2 = pm.Gamma("tau2", mu=1/x.std()**2, sigma=1.0)
----> 7     tn1 = pm.TruncatedNormal.dist("tn1",
      8                                   mu=mu1,
      9                                   tau=tau1,

~/.conda/envs/user_env/lib/python3.9/site-packages/pymc3/distributions/distribution.py in dist(cls, *args, **kwargs)
    128     def dist(cls, *args, **kwargs):
    129         dist = object.__new__(cls)
--> 130         dist.__init__(*args, **kwargs)
    131         return dist
    132 

TypeError: __init__() got multiple values for argument 'mu'

What is the correct way to apply mixture models for truncated normals?

Answer

You don’t need to give the distribution a name. The string tn1 passed as first the argument to pm.TruncatedNormal.dist is interpreted as mu (as you pass a named mu also, you get the exception). Try

tn1 = pm.TruncatedNormal.dist(mu=mu1, tau=tau1, lower=0, upper=1)

(and tn2 accordingly).