@dilumaluthge and I are pleased to announce the release of SossMLJ.jl, which makes it easy to use models from the Soss.jl probabilistic programming language in MLJ.jl.

Here’s a linear regression example, expressed as a special case of a generalized linear model. First,we import the necessary packages:

```
using Distributions
using MLJBase
using Soss
using SossMLJ
using Statistics
```

Next we write the model in Soss:

```
m = @model X, s, t begin
p = size(X, 2) # number of features
β ~ Normal(0, s) |> iid(p) # coefficients
σ ~ HalfNormal(t) # dispersion
η = X * β # linear predictor
μ = η # `μ = g⁻¹(η) = η`
y ~ For(eachindex(μ)) do j
Normal(μ[j], σ) # `Yᵢ ~ Normal(mean=μᵢ, variance=σ²)`
end
end;
```

Now let’s generate some synthetic features. We’ll have two continuous features and two binary categorical features.

```
num_rows = 1_000
x1 = randn(num_rows)
x2 = randn(num_rows)
x3 = Int.(rand(num_rows) .> 0.5)
x4 = Int.(rand(num_rows) .> 0.5)
X = (x1 = x1, x2 = x2, x3 = x3, x4 = x4)
```

In Soss, the *arguments* of a model are the values it needs to determine a joint distribution over *parameters*. At model definition time, Soss doesn’t consider “the variable that will later be observed” to be special in any way, so the “parameters” include this.

In the model `m`

, the arguments are *features* `X`

and *hyperparameters* `s`

and `t`

. Let’s set our hyperparameters…

```
hyperparams = (s=0.1, t=0.1)
```

Now we can express this as a `SossMLJModel`

,

```
model = SossMLJModel(;
model = m,
hyperparams = hyperparams,
infer = dynamicHMC,
response = :y,
);
```

Here `infer`

can be any function that takes a model with arguments (called a `JointDistribution`

in Soss) and returns samples, and `response`

tells which variable we’ll be observing.

Note that `X`

above is a `NamedTuple`

(generally a `Table`

, as expected by MLJ), but `m`

takes a matrix. There’s another `transform`

argument to `SossMLJModel`

that handles this, with the default

```
@inline function default_transform(tbl)
return (X = MLJModelInterface.matrix(tbl),)
end
```

This allows lots of flexibility in how the Soss model represents the features.

Let’s make some fake data to work with.

```
args = merge(model.transform(X), hyperparams)
truth = rand(m(args))
```

On this run, I get

```
julia> pairs(truth)
pairs(::NamedTuple) with 6 entries:
:p => 4
:η => [-0.0999251, 0.226782, -0.243349, 0.0821754, -0.281185, -0.0371795, 0.103831, -0.016029…
:μ => [-0.0999251, 0.226782, -0.243349, 0.0821754, -0.281185, -0.0371795, 0.103831, -0.016029…
:σ => 0.149136
:β => [0.184801, -0.0151003, 0.135162, -0.25159]
:y => [-0.170022, 0.0709445, -0.197551, -0.0813788, 0.195457, 0.0216646, 0.277798, -0.198569,…
```

Next we’ll give the model `X`

and `truth.y`

, and let it figure out the rest. In MLJ we fit models by creating a `machine`

:

```
mach = MLJBase.machine(model, X, truth.y)
fit!(mach)
```

The `fitresult`

for this model contains a `post`

field that’s an array of named tuples. We can summarize it like this:

```
julia> particles(mach.fitresult.post)
(σ = 0.145 ± 0.0036,
β = Particles{Float64,1000}[0.176 ± 0.005, -0.0169 ± 0.0049, 0.129 ± 0.0084, -0.249 ± 0.0088],)
```

Now for a given `X`

(here we’ll just use the same one), we can get the posterior predictive distribution:

```
predictor_joint = predict_joint(mach, X)
typeof(predictor_joint)
```

`predictor_joint`

is just a distribution we can sample from:

```
single_sample = rand(predictor_joint; response = :y)
```

It’s easy to evaluate the `logpdf`

on any sample like this,

```
logpdf(predictor_joint, single_sample)
```

Instead of just one sample, it’s often useful to get `particles`

, in the sense of MonteCarloMeasurements.jl. Also, note that here and above, we can change the `response`

variable after fitting. So,

```
julia> predict_particles(mach, X; response = :β)
4-element Array{Particles{Float64,1000},1}:
0.176 ± 0.005
-0.0169 ± 0.0049
0.129 ± 0.0084
-0.249 ± 0.0088
```

If we know ground truth, this allow checks like

```
julia> truth.β - predict_particles(mach, X; response = :β)
4-element Array{Particles{Float64,1000},1}:
0.00833 ± 0.005
0.00182 ± 0.0049
0.00666 ± 0.0084
-0.00221 ± 0.0088
```

Or if we don’t, we can do posterior predictive checks against our observed `y`

values. Here are the Bayesian p-values for this model:

```
julia> mean.(truth.y .< predict_particles(mach, X; response = :y)) |> particles
Particles{Float64,1000}
0.498052 ± 0.285
```

There’s lots more we can do, check out the rest of this example.

Many thanks to @tlienart for helping me get up to speed on the MLJ interface and getting us started with a previous version of this library, and to @ablaom and the rest of the MLJ team for helpful discussions, and of course for the excellent MLJ framework.