# Modified version of Poisson NMF modeling via NUTS sampler is very slow

**URL:** https://discourse.julialang.org/t/modified-version-of-poisson-nmf-modeling-via-nuts-sampler-is-very-slow/83460
**Category:** New to Julia
**Tags:** question
**Created:** [June 28, 2022, 1:44pm UTC](https://discourse.julialang.org/t/modified-version-of-poisson-nmf-modeling-via-nuts-sampler-is-very-slow/83460 "2022-06-28T13:44:07Z")
**Posts on this page:** 6
**Page:** 1

<div class="post-metadata">

### Author: ![hahawalk](https://avatars.discourse-cdn.com/v4/letter/h/b9bd4f/32.png) [@hahawalk](https://discourse.julialang.org/u/hahawalk)
#### Post date: [June 28, 2022, 1:44pm UTC](https://discourse.julialang.org/t/modified-version-of-poisson-nmf-modeling-via-nuts-sampler-is-very-slow/83460/1 "2022-06-28T13:44:07Z")

</div>

Hi! I am new to Turing.jl and trying to using NUTS to infer the parameters from a modified version of Poisson Non-negative Matrix Factorization(NMF) model. The demo codes are listed as below. It worked for a small number of draws but when I tried 500 draws, it took me like 12 hours to run the sampler and got killed. I wonder if any part of my codes is incorrect or there are any ways to accelerate the sampling procedure? I really appreciate your help!

```julia
eps_k = 200.0
a0_k = 6.0
b0_k = eps_k*(a0_k - 1)

alpha = 1.0
a1 = 1.0
b1 = 1.0

@model function PoissonNMF(x, r_k, ::Type{T} = Float64) where {T}

  mu = Vector{T}(undef, K)
  a = Vector{T}(undef, J)
  t = Matrix{T}(undef, J, K)
  r_inf = Matrix{T}(undef, I, K_u) 

  for j = 1:J
     a[j] ~ InverseGamma(a1, b1)
  end 

  for i = 1:K
       mu[i] ~ InverseGamma(a0_k,b0_k)
  end

  for j = 1:J, k = 1:K
    t[j,k] ~ Gamma(a[j],mu[k]/a[j])
  end

  for kk = 1:K_k
    r_k[:,kk] ~ Dirichlet(I, alpha)
  end

  for ku = 1:K_u
    r_inf[:,ku] ~ Dirichlet(I, alpha)
  end

  r = hcat(r_k, r_inf)
  rate = r*t'

  for i = 1:I, j = 1:J
    x[i,j] ~ Poisson(rate[i,j])
  end
end

model = PoissonNMF(x, r_k)
draws = 500
acp_prob = 0.65
lf_steps = 5
ts = time()
println("The sampling starts")
flush(stdout)
chn = sample(model, NUTS(lf_steps, acp_prob), draws, init_ϵ = 0.000389862060546875)

```

---

<div class="post-metadata">

### Author: ![Christopher\_Fisher](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/christopher_fisher/32/26132_2.png) [@Christopher\_Fisher](https://discourse.julialang.org/u/Christopher_Fisher)
#### Post date: [June 28, 2022, 9:42pm UTC](https://discourse.julialang.org/t/modified-version-of-poisson-nmf-modeling-via-nuts-sampler-is-very-slow/83460/2 "2022-06-28T21:42:43Z")

</div>

Hi,

In general, eliminating global variables will help improve speed (e.g., J, K a1, etc.) . In this particular case, I’m not sure that is the primary problem.

Another potential issue is the auto-diff backend. In Turing, the default forward mode works well for a low number of parameters, such as 5-10. In your case, it looks like you might have much more. If that is true, I recommend using [ReverseDiff](https://turing.ml/dev/docs/using-turing/autodiff). Unfortunately, it requires you to vectorize your code in order to get good performance. That might provide an improvement of several orders of magnitude.

If you can provide a fully runnable example, someone might be able to provide more guidance.

---

<div class="post-metadata">

### Author: ![hahawalk](https://avatars.discourse-cdn.com/v4/letter/h/b9bd4f/32.png) [@hahawalk](https://discourse.julialang.org/u/hahawalk)
#### Post date: [July 5, 2022, 3:47pm UTC](https://discourse.julialang.org/t/modified-version-of-poisson-nmf-modeling-via-nuts-sampler-is-very-slow/83460/3 "2022-07-05T15:47:21Z")

</div>

Hi! Thanks for the comments! I have checked how to vectorize my codes and still am confused about vectorizing the parameters like this part

```julia
for j = 1:J, k = 1:K
    t[j,k] ~ Gamma(a[j],mu[k]/a[j])
  end

```

Yeah, here is a demo code with synthetic data.

```julia
using Turing
using Optim
using Printf

x = [67 15 79 114 33; 21 16 43 41 4; 7 11 14 17 6; 1 3 5 3 0; 119 22 89 193 80; 28 5 27 48 31; 19 4 14 30 16; 1 4 13 4 2; 13 11 18 16 14; 14 18 49 47 16; 106 17 68 131 86; 4 7 17 11 4]
println(size(x))

r_k = [0.0955774895744073 0.18949879345012058; 7.189496934450597e-9 0.0009537832013093777; 0.01183664879213095 0.0038035674109844556; 0.01906680209476258 0.0012755013375099333; 0.0865186633906054 0.3979584390193475; 0.09657592077201838 0.07261594631055265; 0.022973198407369454 0.07093309190047935; 0.00041223552282597023 4.410682039031856e-5; 0.09570716261434926 0.010371759404443884; 0.010831561823427063 0.001405200855652656; 0.5065116350310219 0.2402186872546876; 0.053988674787584665 0.010921123034521816]
println(size(r_k))

#####hyperparameter set-up######
eps_k = 200.0
a0_k = 6.0
b0_k = eps_k*(a0_k - 1)

eps_u = 200.0
a0_u = 6.0
b0_u = eps_u*(a0_u - 1)

alpha = 1.0
a1 = 1.0
b1 = 1.0

@model function PoissonNMF(x, r_k, ::Type{T} = Float64) where {T}

  I,J = size(x)
  K = 4
  _,K_k = size(r_k)
  K_u = K - K_k
  mu = Vector{T}(undef, K)
  a = Vector{T}(undef, J)
  t = Matrix{T}(undef, J, K)
  r_inf = Matrix{T}(undef, I, K_u)

  for j = 1:J
     a[j] ~ InverseGamma(a1, b1)
  end

  for i = 1:K
       mu[i] ~ InverseGamma(a0_k,b0_k)
  end

  for j = 1:J, k = 1:K
    t[j,k] ~ Gamma(a[j],mu[k]/a[j])
  end

  for kk = 1:K_k
    r_k[:,kk] ~ Dirichlet(I, alpha)
  end

  for ku = 1:K_u
    r_inf[:,ku] ~ Dirichlet(I, alpha)
  end

  r = hcat(r_k, r_inf)
  rate = r*t'

  for i = 1:I, j = 1:J
    x[i,j] ~ Poisson(rate[i,j])
  end
end

model = PoissonNMF(x, r_k)
draws = 500
acp_prob = 0.65
lf_steps = 5
ts = time()
println("The sampling starts")
flush(stdout)
chn = sample(model, NUTS(lf_steps, acp_prob), draws, init_ϵ = 0.000389862060546875)

```

---

<div class="post-metadata">

### Author: ![ElOceanografo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/eloceanografo/32/624_2.png) [@ElOceanografo](https://discourse.julialang.org/u/ElOceanografo)
#### Post date: [July 5, 2022, 4:35pm UTC](https://discourse.julialang.org/t/modified-version-of-poisson-nmf-modeling-via-nuts-sampler-is-very-slow/83460/4 "2022-07-05T16:35:27Z")

</div>

Check out this section of the Turing performance tips:  
[https://turing.ml/dev/docs/using-turing/performancetips#special-care-for-codetrackercode-and-codezygotecode](https://turing.ml/dev/docs/using-turing/performancetips#special-care-for-codetrackercode-and-codezygotecode)  
Instead of preallocating your arrays and then filling them using a loop, you probably want to use `filldist` (for your priors) and `arraydist` (for the Poisson observations). For instance, you’d initialize the prior for `a` like

```julia
a ~ filldist(InverseGamma(a1, b1), J)

```

and you’d write the observation likelihood as

```julia
x ~ arraydist(Poisson.(rate))

```

---

<div class="post-metadata">

### Author: ![hahawalk](https://avatars.discourse-cdn.com/v4/letter/h/b9bd4f/32.png) [@hahawalk](https://discourse.julialang.org/u/hahawalk)
#### Post date: [July 6, 2022, 1:41am UTC](https://discourse.julialang.org/t/modified-version-of-poisson-nmf-modeling-via-nuts-sampler-is-very-slow/83460/5 "2022-07-06T01:41:21Z")

</div>

Thanks! I vectorized priors a, mu, r\_k and r\_inf. It works pretty well!  
May I ask if there are any ways to vectorize t?

---

<div class="post-metadata">

### Author: ![ElOceanografo](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/eloceanografo/32/624_2.png) [@ElOceanografo](https://discourse.julialang.org/u/ElOceanografo)
#### Post date: [July 6, 2022, 3:21pm UTC](https://discourse.julialang.org/t/modified-version-of-poisson-nmf-modeling-via-nuts-sampler-is-very-slow/83460/6 "2022-07-06T15:21:36Z")

</div>

This should work:

```julia
t ~ arraydist(Gamma.(a, mu' ./ a))

```

Doing `mu' ./ a` will expand to a J x K matrix, and then the call to `Gamma.()` knows to broadcast `a` over each column.