Wrapper functions for Mamba models

[nsgrantham]

What is the preferred method, if any, for writing wrapper functions for Mamba models?

Suppose I am developing a package and would like to include a function to perform Bayesian linear regression modeled after Mamba's Bayesian linear regression tutorial. Ideally, the function would not require the user to mess with the Mamba internals too much, but it should allow for specifying which parameters to monitor, say, or setting hyperparameters.

For example:

using Mamba

### The following function is defined in my package src

function blinreg(;monitor::Vector{Symbol}=[:beta, :s2],
                  hyperparams::Dict{Symbol, Any}=Dict{Symbol, Any}())

  values = Dict{Symbol, Any}(
    :beta_var => 1000,
    :a_s2 => 0.001,
    :b_s2 => 0.001
  )
  merge!(values, hyperparams)

  model = Model(
    y = Stochastic(1,
      (mu, s2) ->  MvNormal(mu, sqrt(s2)),
      :y in monitor
    ),

    mu = Logical(1,
      (xmat, beta) -> xmat * beta,
      :mu in monitor
    ),

    beta = Stochastic(1,
      () -> MvNormal(2, sqrt(values[:beta_var])),
      :beta in monitor
    ),

    s2 = Stochastic(
      () -> InverseGamma(values[:a_s2], values[:b_s2]),
      :s2 in monitor,
    ),
  )

  setsamplers!(model, [NUTS([:beta, :s2])])

  return model
end

### This is where the user gets involved

## Data
line = Dict{Symbol, Any}(
  :x => [1, 2, 3, 4, 5],
  :y => [1, 3, 3, 3, 5]
)
line[:xmat] = [ones(5) line[:x]]

## Initial Values
inits = [
  Dict{Symbol, Any}(
    :y => line[:y],
    :beta => rand(Normal(0, 1), 2),
    :s2 => rand(Gamma(1, 1))
  )
for i in 1:3
]

## Hyperparams
hyperparams = Dict{Symbol, Any}(
  :beta_var => 100,
  :a_s2 => 0.1,
  :b_s2 => 0.1
)

model = blinreg(monitor = [:y, :mu], hyperparams = hyperparams)
sim = mcmc(model, line, inits, 10000, burnin=250, thin=2, chains=3)

I suppose more advanced cases may require a metaprogramming approach, e.g., placing hyperpriors on the hyperparameters, but that seems overkill for this (admittedly bare-bones) example.

[brian-j-smith]

Lots of different methods could be used for a wrapper. Your example code looks like a valid approach. I would say the "preferred" method is one that meets your design goals. Making a wrapper user-friendly so as to avoid interactions with Mamba internals is an understandable goal. In that case, you might consider moving the mcmc() function call inside blinreg, and have the user specify the inputs, inits, and iterations directly to your function. Then, the user would only need to call blinreg to get posterior draws for the model.

[bdeonovic]

I would agree with Brian. For a function like blinreg I would try to follow the structure of the other linear regression functions in julia in regards to what they return etc. So I would imagine blinreg would simply take in the y and X variables in a regression, fit the model using Mamba, and return the appropriate posterior summary statistics.

[nsgrantham]

Good suggestions. Thank you!

[grero]

Does this still work? I wanted to implement something like this, i.e. a simple wrapper that takes in some variables, constructs a model, and then returns the fitted model, but I keep getting errors like this:

ERROR: MethodError: no method matching (::##79#80)(::Mamba.Model)
The applicable method may be too new: running in world age 25452, while current world is 25456.