Powered by AppSignal & Oban Pro
Notesclub
Sign in with GitHub
created by hec & contributors
terms privacy

Level Set Bayesian Inverse Problem

notebooks/08_level_set.livemd

avatar
Igor O'sten
avatar
avatar
Share to X
Share to Bluesky
More notebooks

Level Set Bayesian Inverse Problem


# CPU only — no GPU required
System.put_env("EXLA_CPU_ONLY", "true")
System.put_env("CUDA_VISIBLE_DEVICES", "")
Mix.install([
  {:exmc, path: Path.expand("../", __DIR__)},
  {:exla, "~> 0.10"},
  {:kino_vega_lite, "~> 0.1"}
])
Application.put_env(:exla, :clients, host: [platform: :host])
Application.put_env(:exla, :default_client, :host)
Nx.default_backend(Nx.BinaryBackend)
Nx.Defn.default_options(compiler: EXLA, client: :host)

Why This Matters

You are drilling a borehole. Somewhere below the surface is a boundary — intact rock on one side, a fault or void on the other. You can’t see it. You can only measure indirect signals: seismic reflections, gravity anomalies, resistivity profiles. Each measurement constrains the boundary’s location, but none pins it down.

The level set method represents this boundary as the zero contour of a continuous function. Bayesian inference over the level set function gives you not a best-guess boundary, but a distribution of plausible boundaries — each consistent with the measurements, each telling a different story about what lies beneath. The spread of those boundaries is the honest answer to “how well do we know where the fault is?”

alias VegaLite, as: Vl
alias Exmc.{Builder, NUTS.Sampler, Dist}
alias Exmc.Physics.{LevelSet, Heat2D}
[Exmc.Physics.LevelSet, Exmc.Physics.Heat2D]

What is a Level Set?

A level set function $\phi(x, y)$ classifies space into two regions:

  • $\phi > 0$ → material A (e.g., intact rock, healthy tissue)
  • $\phi < 0$ → material B (e.g., void, crack, reservoir)

The zero contour $\phi = 0$ defines the boundary between materials. This representation is powerful because:

  1. Topology changes (merging/splitting boundaries) happen naturally
  2. The function lives in $\mathbb{R}^{n}$ — no constrained geometry
  3. Smooth Heaviside approximation makes it differentiable through EXLA JIT

Industrial applications:

  • Subsurface reservoir characterization (oil & gas)
  • Non-destructive testing (crack detection in turbine blades)
  • Electrical impedance tomography (geophysical imaging)
  • Medical imaging (tumor boundary detection) 
# Visualize: a circular level set on an 8x8 grid
ny = 8
nx = 8

phi_demo =
  for i <- 0..(ny - 1), j <- 0..(nx - 1) do
    r = :math.sqrt(:math.pow(i - 3.5, 2) + :math.pow(j - 3.5, 2))
    %{"row" => i, "col" => j, "phi" => 2.0 - r}
  end

Vl.new(width: 300, height: 300, title: "Level Set: Circular Inclusion")
|> Vl.data_from_values(phi_demo)
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "col", type: :ordinal, title: "x")
|> Vl.encode_field(:y, "row", type: :ordinal, title: "y", sort: :descending)
|> Vl.encode_field(:color, "phi",
  type: :quantitative,
  scale: [scheme: "redblue", domain_mid: 0],
  title: "phi"
)
{"$schema":"https://vega.github.io/schema/vega-lite/v5.json","data":{"values":[{"col":0,"phi":-2.9497474683058327,"row":0},{"col":1,"phi":-2.3011626335213133,"row":0},{"col":2,"phi":-1.8078865529319543,"row":0},{"col":3,"phi":-1.5355339059327378,"row":0},{"col":4,"phi":-1.5355339059327378,"row":0},{"col":5,"phi":-1.8078865529319543,"row":0},{"col":6,"phi":-2.3011626335213133,"row":0},{"col":7,"phi":-2.9497474683058327,"row":0},{"col":0,"phi":-2.3011626335213133,"row":1},{"col":1,"phi":-1.5355339059327378,"row":1},{"col":2,"phi":-0.9154759474226504,"row":1},{"col":3,"phi":-0.5495097567963922,"row":1},{"col":4,"phi":-0.5495097567963922,"row":1},{"col":5,"phi":-0.9154759474226504,"row":1},{"col":6,"phi":-1.5355339059327378,"row":1},{"col":7,"phi":-2.3011626335213133,"row":1},{"col":0,"phi":-1.8078865529319543,"row":2},{"col":1,"phi":-0.9154759474226504,"row":2},{"col":2,"phi":-0.12132034355964239,"row":2},{"col":3,"phi":0.41886116991581024,"row":2},{"col":4,"phi":0.41886116991581024,"row":2},{"col":5,"phi":-0.12132034355964239,"row":2},{"col":6,"phi":-0.9154759474226504,"row":2},{"col":7,"phi":-1.8078865529319543,"row":2},{"col":0,"phi":-1.5355339059327378,"row":3},{"col":1,"phi":-0.5495097567963922,"row":3},{"col":2,"phi":0.41886116991581024,"row":3},{"col":3,"phi":1.2928932188134525,"row":3},{"col":4,"phi":1.2928932188134525,"row":3},{"col":5,"phi":0.41886116991581024,"row":3},{"col":6,"phi":-0.5495097567963922,"row":3},{"col":7,"phi":-1.5355339059327378,"row":3},{"col":0,"phi":-1.5355339059327378,"row":4},{"col":1,"phi":-0.5495097567963922,"row":4},{"col":2,"phi":0.41886116991581024,"row":4},{"col":3,"phi":1.2928932188134525,"row":4},{"col":4,"phi":1.2928932188134525,"row":4},{"col":5,"phi":0.41886116991581024,"row":4},{"col":6,"phi":-0.5495097567963922,"row":4},{"col":7,"phi":-1.5355339059327378,"row":4},{"col":0,"phi":-1.8078865529319543,"row":5},{"col":1,"phi":-0.9154759474226504,"row":5},{"col":2,"phi":-0.12132034355964239,"row":5},{"col":3,"phi":0.41886116991581024,"row":5},{"col":4,"phi":0.41886116991581024,"row":5},{"col":5,"phi":-0.12132034355964239,"row":5},{"col":6,"phi":-0.9154759474226504,"row":5},{"col":7,"phi":-1.8078865529319543,"row":5},{"col":0,"phi":-2.3011626335213133,"row":6},{"col":1,"phi":-1.5355339059327378,"row":6},{"col":2,"phi":-0.9154759474226504,"row":6},{"col":3,"phi":-0.5495097567963922,"row":6},{"col":4,"phi":-0.5495097567963922,"row":6},{"col":5,"phi":-0.9154759474226504,"row":6},{"col":6,"phi":-1.5355339059327378,"row":6},{"col":7,"phi":-2.3011626335213133,"row":6},{"col":0,"phi":-2.9497474683058327,"row":7},{"col":1,"phi":-2.3011626335213133,"row":7},{"col":2,"phi":-1.8078865529319543,"row":7},{"col":3,"phi":-1.5355339059327378,"row":7},{"col":4,"phi":-1.5355339059327378,"row":7},{"col":5,"phi":-1.8078865529319543,"row":7},{"col":6,"phi":-2.3011626335213133,"row":7},{"col":7,"phi":-2.9497474683058327,"row":7}]},"encoding":{"color":{"field":"phi","scale":{"domainMid":0,"scheme":"redblue"},"title":"phi","type":"quantitative"},"x":{"field":"col","title":"x","type":"ordinal"},"y":{"field":"row","sort":"descending","title":"y","type":"ordinal"}},"height":300,"mark":"rect","title":"Level Set: Circular Inclusion","width":300}

The Forward Model

We use a 2D steady-state heat equation as the forward model:

$$-\nabla \cdot (\kappa \nabla T) = 0$$

where $\kappa(x,y)$ is the thermal conductivity field determined by the level set:

$$\kappa = \kappaa \cdot H\varepsilon(\phi) + \kappab \cdot (1 - H\varepsilon(\phi))$$

$H\varepsilon$ is a smooth Heaviside function: $H\varepsilon(\phi) = \frac{1}{2}(1 + \tanh(\phi/\varepsilon))$.

Boundary conditions: $T = 1$ at top, $T = 0$ at bottom. We observe temperature at the bottom row (simulating sensors).

# Show how the level set maps to a conductivity field
phi_grid =
  for i <- 0..(ny - 1), j <- 0..(nx - 1) do
    r = :math.sqrt(:math.pow(i - 3.5, 2) + :math.pow(j - 3.5, 2))
    2.0 - r
  end
  |> Nx.tensor()
  |> Nx.reshape({ny, nx})

kappa_a = 5.0
kappa_b = 1.0

kappa_field = LevelSet.material_field(phi_grid, Nx.tensor(kappa_a), Nx.tensor(kappa_b))

kappa_data =
  for i <- 0..(ny - 1), j <- 0..(nx - 1) do
    %{"row" => i, "col" => j, "kappa" => Nx.to_number(kappa_field[i][j])}
  end

Vl.new(width: 300, height: 300, title: "Conductivity Field kappa(x,y)")
|> Vl.data_from_values(kappa_data)
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "col", type: :ordinal, title: "x")
|> Vl.encode_field(:y, "row", type: :ordinal, title: "y", sort: :descending)
|> Vl.encode_field(:color, "kappa",
  type: :quantitative,
  scale: [scheme: "viridis"],
  title: "kappa"
)
{"$schema":"https://vega.github.io/schema/vega-lite/v5.json","data":{"values":[{"col":0,"kappa":1.0109333992004395,"row":0},{"col":1,"kappa":1.0397156476974487,"row":0},{"col":2,"kappa":1.1047667264938354,"row":0},{"col":3,"kappa":1.1772663593292236,"row":0},{"col":4,"kappa":1.1772663593292236,"row":0},{"col":5,"kappa":1.1047667264938354,"row":0},{"col":6,"kappa":1.0397156476974487,"row":0},{"col":7,"kappa":1.0109333992004395,"row":0},{"col":0,"kappa":1.0397156476974487,"row":1},{"col":1,"kappa":1.1772663593292236,"row":1},{"col":2,"kappa":1.5524996519088745,"row":1},{"col":3,"kappa":1.999694585800171,"row":1},{"col":4,"kappa":1.999694585800171,"row":1},{"col":5,"kappa":1.5524996519088745,"row":1},{"col":6,"kappa":1.1772663593292236,"row":1},{"col":7,"kappa":1.0397156476974487,"row":1},{"col":0,"kappa":1.1047667264938354,"row":2},{"col":1,"kappa":1.5524996519088745,"row":2},{"col":2,"kappa":2.758542776107788,"row":2},{"col":3,"kappa":3.7919411659240723,"row":2},{"col":4,"kappa":3.7919411659240723,"row":2},{"col":5,"kappa":2.758542776107788,"row":2},{"col":6,"kappa":1.5524996519088745,"row":2},{"col":7,"kappa":1.1047667264938354,"row":2},{"col":0,"kappa":1.1772663593292236,"row":3},{"col":1,"kappa":1.999694585800171,"row":3},{"col":2,"kappa":3.7919411659240723,"row":3},{"col":3,"kappa":4.719764709472656,"row":3},{"col":4,"kappa":4.719764709472656,"row":3},{"col":5,"kappa":3.7919411659240723,"row":3},{"col":6,"kappa":1.999694585800171,"row":3},{"col":7,"kappa":1.1772663593292236,"row":3},{"col":0,"kappa":1.1772663593292236,"row":4},{"col":1,"kappa":1.999694585800171,"row":4},{"col":2,"kappa":3.7919411659240723,"row":4},{"col":3,"kappa":4.719764709472656,"row":4},{"col":4,"kappa":4.719764709472656,"row":4},{"col":5,"kappa":3.7919411659240723,"row":4},{"col":6,"kappa":1.999694585800171,"row":4},{"col":7,"kappa":1.1772663593292236,"row":4},{"col":0,"kappa":1.1047667264938354,"row":5},{"col":1,"kappa":1.5524996519088745,"row":5},{"col":2,"kappa":2.758542776107788,"row":5},{"col":3,"kappa":3.7919411659240723,"row":5},{"col":4,"kappa":3.7919411659240723,"row":5},{"col":5,"kappa":2.758542776107788,"row":5},{"col":6,"kappa":1.5524996519088745,"row":5},{"col":7,"kappa":1.1047667264938354,"row":5},{"col":0,"kappa":1.0397156476974487,"row":6},{"col":1,"kappa":1.1772663593292236,"row":6},{"col":2,"kappa":1.5524996519088745,"row":6},{"col":3,"kappa":1.999694585800171,"row":6},{"col":4,"kappa":1.999694585800171,"row":6},{"col":5,"kappa":1.5524996519088745,"row":6},{"col":6,"kappa":1.1772663593292236,"row":6},{"col":7,"kappa":1.0397156476974487,"row":6},{"col":0,"kappa":1.0109333992004395,"row":7},{"col":1,"kappa":1.0397156476974487,"row":7},{"col":2,"kappa":1.1047667264938354,"row":7},{"col":3,"kappa":1.1772663593292236,"row":7},{"col":4,"kappa":1.1772663593292236,"row":7},{"col":5,"kappa":1.1047667264938354,"row":7},{"col":6,"kappa":1.0397156476974487,"row":7},{"col":7,"kappa":1.0109333992004395,"row":7}]},"encoding":{"color":{"field":"kappa","scale":{"scheme":"viridis"},"title":"kappa","type":"quantitative"},"x":{"field":"col","title":"x","type":"ordinal"},"y":{"field":"row","sort":"descending","title":"y","type":"ordinal"}},"height":300,"mark":"rect","title":"Conductivity Field kappa(x,y)","width":300}
# Solve heat equation and show temperature field
temperature = Heat2D.solve(kappa_field, bc_top: 1.0, bc_bottom: 0.0, iterations: 100)

temp_data =
  for i <- 0..(ny - 1), j <- 0..(nx - 1) do
    %{"row" => i, "col" => j, "T" => Nx.to_number(temperature[i][j])}
  end

Vl.new(width: 300, height: 300, title: "Temperature Field T(x,y)")
|> Vl.data_from_values(temp_data)
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "col", type: :ordinal, title: "x")
|> Vl.encode_field(:y, "row", type: :ordinal, title: "y", sort: :descending)
|> Vl.encode_field(:color, "T",
  type: :quantitative,
  scale: [scheme: "inferno"],
  title: "T"
)
{"$schema":"https://vega.github.io/schema/vega-lite/v5.json","data":{"values":[{"T":1.0,"col":0,"row":0},{"T":1.0,"col":1,"row":0},{"T":1.0,"col":2,"row":0},{"T":1.0,"col":3,"row":0},{"T":1.0,"col":4,"row":0},{"T":1.0,"col":5,"row":0},{"T":1.0,"col":6,"row":0},{"T":1.0,"col":7,"row":0},{"T":0.8571428656578064,"col":0,"row":1},{"T":0.7700230479240417,"col":1,"row":1},{"T":0.7072517275810242,"col":2,"row":1},{"T":0.6818984150886536,"col":3,"row":1},{"T":0.6818984150886536,"col":4,"row":1},{"T":0.7072517275810242,"col":5,"row":1},{"T":0.7700230479240417,"col":6,"row":1},{"T":0.8571428656578064,"col":7,"row":1},{"T":0.7142857313156128,"col":0,"row":2},{"T":0.6204333901405334,"col":1,"row":2},{"T":0.5815988779067993,"col":2,"row":2},{"T":0.5727589130401611,"col":3,"row":2},{"T":0.5727589130401611,"col":4,"row":2},{"T":0.5815988779067993,"col":5,"row":2},{"T":0.6204333901405334,"col":6,"row":2},{"T":0.7142857313156128,"col":7,"row":2},{"T":0.5714285373687744,"col":0,"row":3},{"T":0.5340873599052429,"col":1,"row":3},{"T":0.5230944156646729,"col":2,"row":3},{"T":0.5213513970375061,"col":3,"row":3},{"T":0.5213513970375061,"col":4,"row":3},{"T":0.5230944156646729,"col":5,"row":3},{"T":0.5340873599052429,"col":6,"row":3},{"T":0.5714285373687744,"col":7,"row":3},{"T":0.4285714030265808,"col":0,"row":4},{"T":0.46591296792030334,"col":1,"row":4},{"T":0.4769061207771301,"col":2,"row":4},{"T":0.4786492586135864,"col":3,"row":4},{"T":0.4786492586135864,"col":4,"row":4},{"T":0.4769061207771301,"col":5,"row":4},{"T":0.46591296792030334,"col":6,"row":4},{"T":0.4285714030265808,"col":7,"row":4},{"T":0.2857142686843872,"col":0,"row":5},{"T":0.3795669674873352,"col":1,"row":5},{"T":0.4184015095233917,"col":2,"row":5},{"T":0.42724159359931946,"col":3,"row":5},{"T":0.42724159359931946,"col":4,"row":5},{"T":0.4184015095233917,"col":5,"row":5},{"T":0.3795669674873352,"col":6,"row":5},{"T":0.2857142686843872,"col":7,"row":5},{"T":0.1428571343421936,"col":0,"row":6},{"T":0.22997714579105377,"col":1,"row":6},{"T":0.2927486300468445,"col":2,"row":6},{"T":0.31810182332992554,"col":3,"row":6},{"T":0.31810182332992554,"col":4,"row":6},{"T":0.2927486300468445,"col":5,"row":6},{"T":0.22997714579105377,"col":6,"row":6},{"T":0.1428571343421936,"col":7,"row":6},{"T":0.0,"col":0,"row":7},{"T":0.0,"col":1,"row":7},{"T":0.0,"col":2,"row":7},{"T":0.0,"col":3,"row":7},{"T":0.0,"col":4,"row":7},{"T":0.0,"col":5,"row":7},{"T":0.0,"col":6,"row":7},{"T":0.0,"col":7,"row":7}]},"encoding":{"color":{"field":"T","scale":{"scheme":"inferno"},"title":"T","type":"quantitative"},"x":{"field":"col","title":"x","type":"ordinal"},"y":{"field":"row","sort":"descending","title":"y","type":"ordinal"}},"height":300,"mark":"rect","title":"Temperature Field T(x,y)","width":300}

Synthetic Data

We generate synthetic sensor readings from a known circular inclusion, then add Gaussian noise. The inverse problem: given only the noisy bottom-row sensors, recover the hidden geometry.

# True level set: circular inclusion at center
true_phi =
  for i <- 0..(ny - 1), j <- 0..(nx - 1) do
    r = :math.sqrt(:math.pow(i - 3.5, 2) + :math.pow(j - 3.5, 2))
    2.0 - r
  end
  |> Nx.tensor()
  |> Nx.reshape({ny, nx})

# Forward solve with true geometry
true_kappa = LevelSet.material_field(true_phi, Nx.tensor(kappa_a), Nx.tensor(kappa_b))
true_temp = Heat2D.solve(true_kappa, bc_top: 1.0, bc_bottom: 0.0, iterations: 100)
true_sensors = Heat2D.read_sensors(true_temp, :bottom_row)

# Add observation noise
sigma_obs_true = 0.02
rng = :rand.seed_s(:exsss, 42)

{noise_list, _rng} =
  Enum.map_reduce(1..nx, rng, fn _, rng ->
    {n, rng} = :rand.normal_s(rng)
    {sigma_obs_true * n, rng}
  end)

sensor_data = Nx.add(true_sensors, Nx.tensor(noise_list))

# Plot sensor data vs truth
sensor_chart_data =
  for j <- 0..(nx - 1) do
    [
      %{"x" => j, "T" => Nx.to_number(true_sensors[j]), "series" => "True"},
      %{"x" => j, "T" => Nx.to_number(sensor_data[j]), "series" => "Observed (noisy)"}
    ]
  end
  |> List.flatten()

Vl.new(width: 400, height: 250, title: "Bottom-Row Sensor Readings")
|> Vl.data_from_values(sensor_chart_data)
|> Vl.mark(:line, point: true, stroke_width: 2)
|> Vl.encode_field(:x, "x", type: :quantitative, title: "Sensor position")
|> Vl.encode_field(:y, "T", type: :quantitative, title: "Temperature")
|> Vl.encode_field(:color, "series", type: :nominal)
{"$schema":"https://vega.github.io/schema/vega-lite/v5.json","data":{"values":[{"T":0.0,"series":"True","x":0},{"T":0.011677770875394344,"series":"Observed (noisy)","x":0},{"T":0.0,"series":"True","x":1},{"T":-0.03801949322223663,"series":"Observed (noisy)","x":1},{"T":0.0,"series":"True","x":2},{"T":2.9738631565123796e-4,"series":"Observed (noisy)","x":2},{"T":0.0,"series":"True","x":3},{"T":0.03911164775490761,"series":"Observed (noisy)","x":3},{"T":0.0,"series":"True","x":4},{"T":-0.0361650176346302,"series":"Observed (noisy)","x":4},{"T":0.0,"series":"True","x":5},{"T":0.023496273905038834,"series":"Observed (noisy)","x":5},{"T":0.0,"series":"True","x":6},{"T":0.03217975050210953,"series":"Observed (noisy)","x":6},{"T":0.0,"series":"True","x":7},{"T":0.017942722886800766,"series":"Observed (noisy)","x":7}]},"encoding":{"color":{"field":"series","type":"nominal"},"x":{"field":"x","title":"Sensor position","type":"quantitative"},"y":{"field":"T","title":"Temperature","type":"quantitative"}},"height":250,"mark":{"point":true,"strokeWidth":2,"type":"line"},"title":"Bottom-Row Sensor Readings","width":400}

Bayesian Model

The model has three components:

  1. Level set prior: $\phi_i \sim N(0, 1)$ with Laplacian smoothness penalty $\log p(\phi) = -\frac{1}{2}\sum \phi_i^2 - \frac{\lambda}{2}\sum(\nabla^2\phi)^2$
  2. Observation noise: $\sigma_\text{obs} \sim \text{HalfCauchy}(0.1)$
  3. Likelihood: $y\text{obs} \sim N(f(\phi), \sigma\text{obs})$ where $f$ is the forward model 
n = ny * nx

ir = Builder.new_ir()

# Phi: level set values with combined Normal + Laplacian prior
# Custom dist returns scalar logp (required by compiler contract)
laplacian_fn = LevelSet.laplacian_prior_logpdf(ny, nx)

phi_prior_logpdf = fn x, params ->
  lambda = params.lambda
  normal_logp = Nx.negate(Nx.sum(Nx.multiply(0.5, Nx.pow(x, 2))))
  laplacian_logp = laplacian_fn.(Nx.tensor(0.0), %{phi: x, lambda: lambda})
  Nx.add(normal_logp, laplacian_logp)
end

phi_dist = Dist.Custom.new(phi_prior_logpdf, support: :real)

ir =
  Dist.Custom.rv(ir, "phi", phi_dist, %{
    lambda: Nx.tensor(1.0)
  }, shape: {n})

# Observation noise
ir = Builder.rv(ir, "sigma_obs", Dist.HalfCauchy, %{scale: Nx.tensor(0.1)}, transform: :log)

# Forward model likelihood
forward_fn = fn kappa_field ->
  Heat2D.solve(kappa_field, bc_top: 1.0, bc_bottom: 0.0, iterations: 50)
  |> Heat2D.read_sensors(:bottom_row)
end

likelihood_logpdf = fn _x, params ->
  phi_flat = params.phi
  sigma = params.sigma_obs
  obs = params.data

  phi_2d = Nx.reshape(phi_flat, {ny, nx})
  kappa = LevelSet.material_field(phi_2d, kappa_a, kappa_b, eps: 1.0)
  predictions = forward_fn.(kappa)

  resid = Nx.subtract(obs, predictions)
  n_obs = Nx.tensor(Nx.size(obs) * 1.0)

  Nx.subtract(
    Nx.sum(Nx.negate(Nx.divide(Nx.pow(resid, 2), Nx.multiply(2.0, Nx.pow(sigma, 2))))),
    Nx.multiply(n_obs, Nx.log(sigma))
  )
end

ll_dist = Dist.Custom.new(likelihood_logpdf, support: :real)

ir =
  Dist.Custom.rv(ir, "ll", ll_dist, %{
    phi: "phi",
    sigma_obs: "sigma_obs",
    data: sensor_data
  })

ir = Builder.obs(ir, "ll_obs", "ll", Nx.tensor(0.0))

IO.puts("Model: #{n + 1} free parameters (#{n} phi + 1 sigma_obs)")
Model: 65 free parameters (64 phi + 1 sigma_obs)
:ok

MCMC Sampling

# Init at prior mode: flat phi (zero Laplacian penalty), sigma_obs near true
init = %{"phi" => Nx.broadcast(0.0, {n}), "sigma_obs" => 0.03}

{trace, stats} =
  Sampler.sample(ir, init,
    num_warmup: 500,
    num_samples: 500,
    seed: 42,
    ncp: false
  )

IO.puts("Step size: #{Float.round(stats.step_size, 4)}")
IO.puts("Divergences: #{stats.divergences}")
IO.puts("sigma_obs posterior mean: #{Float.round(Nx.to_number(Nx.mean(trace["sigma_obs"])), 4)}")
IO.puts("sigma_obs true: #{sigma_obs_true}")
Step size: 0.2113
Divergences: 8
sigma_obs posterior mean: 0.0338
sigma_obs true: 0.02
:ok

Results

True vs Recovered Level Set

mean_phi = Nx.mean(trace["phi"], axes: [0]) |> Nx.reshape({ny, nx})

true_phi_data =
  for i <- 0..(ny - 1), j <- 0..(nx - 1) do
    %{"row" => i, "col" => j, "phi" => Nx.to_number(true_phi[i][j])}
  end

recovered_phi_data =
  for i <- 0..(ny - 1), j <- 0..(nx - 1) do
    %{"row" => i, "col" => j, "phi" => Nx.to_number(mean_phi[i][j])}
  end

heatmap_spec = fn data, title ->
  Vl.new(width: 250, height: 250, title: title)
  |> Vl.data_from_values(data)
  |> Vl.mark(:rect)
  |> Vl.encode_field(:x, "col", type: :ordinal, title: "x")
  |> Vl.encode_field(:y, "row", type: :ordinal, title: "y", sort: :descending)
  |> Vl.encode_field(:color, "phi",
    type: :quantitative,
    scale: [scheme: "redblue", domain_mid: 0],
    title: "phi"
  )
end

Vl.new()
|> Vl.concat(
  [
    heatmap_spec.(true_phi_data, "True Level Set"),
    heatmap_spec.(recovered_phi_data, "Posterior Mean Level Set")
  ],
  :horizontal
)
{"$schema":"https://vega.github.io/schema/vega-lite/v5.json","hconcat":[{"data":{"values":[{"col":0,"phi":-2.9497475624084473,"row":0},{"col":1,"phi":-2.3011627197265625,"row":0},{"col":2,"phi":-1.8078866004943848,"row":0},{"col":3,"phi":-1.5355339050292969,"row":0},{"col":4,"phi":-1.5355339050292969,"row":0},{"col":5,"phi":-1.8078866004943848,"row":0},{"col":6,"phi":-2.3011627197265625,"row":0},{"col":7,"phi":-2.9497475624084473,"row":0},{"col":0,"phi":-2.3011627197265625,"row":1},{"col":1,"phi":-1.5355339050292969,"row":1},{"col":2,"phi":-0.9154759645462036,"row":1},{"col":3,"phi":-0.5495097637176514,"row":1},{"col":4,"phi":-0.5495097637176514,"row":1},{"col":5,"phi":-0.9154759645462036,"row":1},{"col":6,"phi":-1.5355339050292969,"row":1},{"col":7,"phi":-2.3011627197265625,"row":1},{"col":0,"phi":-1.8078866004943848,"row":2},{"col":1,"phi":-0.9154759645462036,"row":2},{"col":2,"phi":-0.12132034450769424,"row":2},{"col":3,"phi":0.41886118054389954,"row":2},{"col":4,"phi":0.41886118054389954,"row":2},{"col":5,"phi":-0.12132034450769424,"row":2},{"col":6,"phi":-0.9154759645462036,"row":2},{"col":7,"phi":-1.8078866004943848,"row":2},{"col":0,"phi":-1.5355339050292969,"row":3},{"col":1,"phi":-0.5495097637176514,"row":3},{"col":2,"phi":0.41886118054389954,"row":3},{"col":3,"phi":1.2928931713104248,"row":3},{"col":4,"phi":1.2928931713104248,"row":3},{"col":5,"phi":0.41886118054389954,"row":3},{"col":6,"phi":-0.5495097637176514,"row":3},{"col":7,"phi":-1.5355339050292969,"row":3},{"col":0,"phi":-1.5355339050292969,"row":4},{"col":1,"phi":-0.5495097637176514,"row":4},{"col":2,"phi":0.41886118054389954,"row":4},{"col":3,"phi":1.2928931713104248,"row":4},{"col":4,"phi":1.2928931713104248,"row":4},{"col":5,"phi":0.41886118054389954,"row":4},{"col":6,"phi":-0.5495097637176514,"row":4},{"col":7,"phi":-1.5355339050292969,"row":4},{"col":0,"phi":-1.8078866004943848,"row":5},{"col":1,"phi":-0.9154759645462036,"row":5},{"col":2,"phi":-0.12132034450769424,"row":5},{"col":3,"phi":0.41886118054389954,"row":5},{"col":4,"phi":0.41886118054389954,"row":5},{"col":5,"phi":-0.12132034450769424,"row":5},{"col":6,"phi":-0.9154759645462036,"row":5},{"col":7,"phi":-1.8078866004943848,"row":5},{"col":0,"phi":-2.3011627197265625,"row":6},{"col":1,"phi":-1.5355339050292969,"row":6},{"col":2,"phi":-0.9154759645462036,"row":6},{"col":3,"phi":-0.5495097637176514,"row":6},{"col":4,"phi":-0.5495097637176514,"row":6},{"col":5,"phi":-0.9154759645462036,"row":6},{"col":6,"phi":-1.5355339050292969,"row":6},{"col":7,"phi":-2.3011627197265625,"row":6},{"col":0,"phi":-2.9497475624084473,"row":7},{"col":1,"phi":-2.3011627197265625,"row":7},{"col":2,"phi":-1.8078866004943848,"row":7},{"col":3,"phi":-1.5355339050292969,"row":7},{"col":4,"phi":-1.5355339050292969,"row":7},{"col":5,"phi":-1.8078866004943848,"row":7},{"col":6,"phi":-2.3011627197265625,"row":7},{"col":7,"phi":-2.9497475624084473,"row":7}]},"encoding":{"color":{"field":"phi","scale":{"domainMid":0,"scheme":"redblue"},"title":"phi","type":"quantitative"},"x":{"field":"col","title":"x","type":"ordinal"},"y":{"field":"row","sort":"descending","title":"y","type":"ordinal"}},"height":250,"mark":"rect","title":"True Level Set","width":250},{"data":{"values":[{"col":0,"phi":-0.01833565941881902,"row":0},{"col":1,"phi":-0.006772024340194442,"row":0},{"col":2,"phi":-0.009725158085565984,"row":0},{"col":3,"phi":0.0531424806009578,"row":0},{"col":4,"phi":0.02953791253620989,"row":0},{"col":5,"phi":-0.07374609945153453,"row":0},{"col":6,"phi":-0.061221099655842916,"row":0},{"col":7,"phi":-0.04548406614757544,"row":0},{"col":0,"phi":-0.02803658428671883,"row":1},{"col":1,"phi":0.009135571357477563,"row":1},{"col":2,"phi":0.012627265742253483,"row":1},{"col":3,"phi":0.011770831446388621,"row":1},{"col":4,"phi":-0.0017396443292260054,"row":1},{"col":5,"phi":-0.02283336110464358,"row":1},{"col":6,"phi":-0.0034225661106555946,"row":1},{"col":7,"phi":0.055671217831798486,"row":1},{"col":0,"phi":-0.018916393430209014,"row":2},{"col":1,"phi":-0.008563524703628932,"row":2},{"col":2,"phi":-0.018517281582354803,"row":2},{"col":3,"phi":-0.01346827070169582,"row":2},{"col":4,"phi":-0.009509899151768557,"row":2},{"col":5,"phi":-0.018920411933661306,"row":2},{"col":6,"phi":0.027941711466774027,"row":2},{"col":7,"phi":0.054875542821723494,"row":2},{"col":0,"phi":-0.03742576412996145,"row":3},{"col":1,"phi":-0.02960412403852148,"row":3},{"col":2,"phi":-0.042146424165880235,"row":3},{"col":3,"phi":-0.03796174418209242,"row":3},{"col":4,"phi":-0.020407209473143926,"row":3},{"col":5,"phi":0.004880158055642537,"row":3},{"col":6,"phi":0.008996063162202465,"row":3},{"col":7,"phi":0.040862196155613366,"row":3},{"col":0,"phi":-0.004151732292035829,"row":4},{"col":1,"phi":-0.022684020609726695,"row":4},{"col":2,"phi":-0.0189544581841116,"row":4},{"col":3,"phi":-0.037664252846993145,"row":4},{"col":4,"phi":-0.0077244478828666675,"row":4},{"col":5,"phi":0.012159283427447974,"row":4},{"col":6,"phi":-0.004339633056587987,"row":4},{"col":7,"phi":0.025271345115490482,"row":4},{"col":0,"phi":-0.02444537180384507,"row":5},{"col":1,"phi":-0.020525264457856435,"row":5},{"col":2,"phi":-0.014314286121927765,"row":5},{"col":3,"phi":-0.022514808991454523,"row":5},{"col":4,"phi":-0.013110889776687095,"row":5},{"col":5,"phi":-0.009234429517602855,"row":5},{"col":6,"phi":0.009370030201282903,"row":5},{"col":7,"phi":0.030408765680723492,"row":5},{"col":0,"phi":0.06110781390607839,"row":6},{"col":1,"phi":0.005510370511430867,"row":6},{"col":2,"phi":0.00376774125594721,"row":6},{"col":3,"phi":-0.0026478701130358915,"row":6},{"col":4,"phi":-0.012331139819452228,"row":6},{"col":5,"phi":0.02346245615913636,"row":6},{"col":6,"phi":0.022577805707278003,"row":6},{"col":7,"phi":-0.028480607370488787,"row":6},{"col":0,"phi":-0.0017400420140259473,"row":7},{"col":1,"phi":-0.07338566208570206,"row":7},{"col":2,"phi":0.03220742308750258,"row":7},{"col":3,"phi":-0.031157225037728765,"row":7},{"col":4,"phi":-0.003188427791417097,"row":7},{"col":5,"phi":0.07503912616035704,"row":7},{"col":6,"phi":0.016407495622195047,"row":7},{"col":7,"phi":0.038072241365545045,"row":7}]},"encoding":{"color":{"field":"phi","scale":{"domainMid":0,"scheme":"redblue"},"title":"phi","type":"quantitative"},"x":{"field":"col","title":"x","type":"ordinal"},"y":{"field":"row","sort":"descending","title":"y","type":"ordinal"}},"height":250,"mark":"rect","title":"Posterior Mean Level Set","width":250}]}

Posterior Uncertainty

The standard deviation of $\phi$ across MCMC samples shows where the model is uncertain. High uncertainty near the boundary means the data is informative about the inclusion but the exact boundary location is hard to pin down.

std_phi = Nx.standard_deviation(trace["phi"], axes: [0]) |> Nx.reshape({ny, nx})

std_data =
  for i <- 0..(ny - 1), j <- 0..(nx - 1) do
    %{"row" => i, "col" => j, "std" => Nx.to_number(std_phi[i][j])}
  end

Vl.new(width: 300, height: 300, title: "Posterior Std Dev of phi")
|> Vl.data_from_values(std_data)
|> Vl.mark(:rect)
|> Vl.encode_field(:x, "col", type: :ordinal, title: "x")
|> Vl.encode_field(:y, "row", type: :ordinal, title: "y", sort: :descending)
|> Vl.encode_field(:color, "std",
  type: :quantitative,
  scale: [scheme: "oranges"],
  title: "std(phi)"
)
{"$schema":"https://vega.github.io/schema/vega-lite/v5.json","data":{"values":[{"col":0,"row":0,"std":0.9754265780173516},{"col":1,"row":0,"std":1.0249681086136715},{"col":2,"row":0,"std":1.0140549884707744},{"col":3,"row":0,"std":0.9571385381384335},{"col":4,"row":0,"std":0.9577519816849301},{"col":5,"row":0,"std":0.9471892998139539},{"col":6,"row":0,"std":0.9754597063617423},{"col":7,"row":0,"std":0.9509176887168177},{"col":0,"row":1,"std":0.9463752510785335},{"col":1,"row":1,"std":0.5043948713302359},{"col":2,"row":1,"std":0.44751520328104816},{"col":3,"row":1,"std":0.4275587212183085},{"col":4,"row":1,"std":0.42083302970692527},{"col":5,"row":1,"std":0.4150397150108576},{"col":6,"row":1,"std":0.4659587502675418},{"col":7,"row":1,"std":0.9663413993822745},{"col":0,"row":2,"std":0.9286188709015799},{"col":1,"row":2,"std":0.4783377343500926},{"col":2,"row":2,"std":0.4092286594269943},{"col":3,"row":2,"std":0.3795930815562371},{"col":4,"row":2,"std":0.3781694540569417},{"col":5,"row":2,"std":0.3701885403956607},{"col":6,"row":2,"std":0.39940213643488587},{"col":7,"row":2,"std":0.980509962102642},{"col":0,"row":3,"std":0.9674675329822943},{"col":1,"row":3,"std":0.45729638227033675},{"col":2,"row":3,"std":0.4090136451951416},{"col":3,"row":3,"std":0.3904849464350839},{"col":4,"row":3,"std":0.39339655106814153},{"col":5,"row":3,"std":0.38972369946432855},{"col":6,"row":3,"std":0.41739283705082586},{"col":7,"row":3,"std":1.0008491342194354},{"col":0,"row":4,"std":0.8846906630629503},{"col":1,"row":4,"std":0.4153075881024623},{"col":2,"row":4,"std":0.3943403777623395},{"col":3,"row":4,"std":0.386074959322377},{"col":4,"row":4,"std":0.3995268392584451},{"col":5,"row":4,"std":0.4042959570853862},{"col":6,"row":4,"std":0.46022676770712884},{"col":7,"row":4,"std":0.9290937247281053},{"col":0,"row":5,"std":0.9794881131027802},{"col":1,"row":5,"std":0.4514655853278928},{"col":2,"row":5,"std":0.41281524803695874},{"col":3,"row":5,"std":0.3970162150853084},{"col":4,"row":5,"std":0.39148658397018826},{"col":5,"row":5,"std":0.40243535008788595},{"col":6,"row":5,"std":0.439250260756653},{"col":7,"row":5,"std":0.9480443488359539},{"col":0,"row":6,"std":1.0002681247306187},{"col":1,"row":6,"std":0.4576430277579079},{"col":2,"row":6,"std":0.44555891356016697},{"col":3,"row":6,"std":0.423800847981942},{"col":4,"row":6,"std":0.4233664215445668},{"col":5,"row":6,"std":0.43073053635608477},{"col":6,"row":6,"std":0.459622829211377},{"col":7,"row":6,"std":0.9647407840348505},{"col":0,"row":7,"std":1.0433768472132647},{"col":1,"row":7,"std":0.9560835570168408},{"col":2,"row":7,"std":0.977005874120023},{"col":3,"row":7,"std":0.8912381648911331},{"col":4,"row":7,"std":0.9426549083553669},{"col":5,"row":7,"std":0.9642387007406867},{"col":6,"row":7,"std":1.0178775086314231},{"col":7,"row":7,"std":0.8902948226300121}]},"encoding":{"color":{"field":"std","scale":{"scheme":"oranges"},"title":"std(phi)","type":"quantitative"},"x":{"field":"col","title":"x","type":"ordinal"},"y":{"field":"row","sort":"descending","title":"y","type":"ordinal"}},"height":300,"mark":"rect","title":"Posterior Std Dev of phi","width":300}

True vs Recovered Conductivity

recovered_kappa =
  LevelSet.material_field(mean_phi, Nx.tensor(kappa_a), Nx.tensor(kappa_b))

true_kappa_data =
  for i <- 0..(ny - 1), j <- 0..(nx - 1) do
    %{"row" => i, "col" => j, "kappa" => Nx.to_number(true_kappa[i][j])}
  end

recovered_kappa_data =
  for i <- 0..(ny - 1), j <- 0..(nx - 1) do
    %{"row" => i, "col" => j, "kappa" => Nx.to_number(recovered_kappa[i][j])}
  end

kappa_spec = fn data, title ->
  Vl.new(width: 250, height: 250, title: title)
  |> Vl.data_from_values(data)
  |> Vl.mark(:rect)
  |> Vl.encode_field(:x, "col", type: :ordinal, title: "x")
  |> Vl.encode_field(:y, "row", type: :ordinal, title: "y", sort: :descending)
  |> Vl.encode_field(:color, "kappa",
    type: :quantitative,
    scale: [scheme: "viridis", domain: [kappa_b, kappa_a]],
    title: "kappa"
  )
end

Vl.new()
|> Vl.concat(
  [
    kappa_spec.(true_kappa_data, "True Conductivity"),
    kappa_spec.(recovered_kappa_data, "Recovered Conductivity")
  ],
  :horizontal
)
{"$schema":"https://vega.github.io/schema/vega-lite/v5.json","hconcat":[{"data":{"values":[{"col":0,"kappa":1.0109333992004395,"row":0},{"col":1,"kappa":1.0397156476974487,"row":0},{"col":2,"kappa":1.1047667264938354,"row":0},{"col":3,"kappa":1.1772663593292236,"row":0},{"col":4,"kappa":1.1772663593292236,"row":0},{"col":5,"kappa":1.1047667264938354,"row":0},{"col":6,"kappa":1.0397156476974487,"row":0},{"col":7,"kappa":1.0109333992004395,"row":0},{"col":0,"kappa":1.0397156476974487,"row":1},{"col":1,"kappa":1.1772663593292236,"row":1},{"col":2,"kappa":1.5524996519088745,"row":1},{"col":3,"kappa":1.999694585800171,"row":1},{"col":4,"kappa":1.999694585800171,"row":1},{"col":5,"kappa":1.5524996519088745,"row":1},{"col":6,"kappa":1.1772663593292236,"row":1},{"col":7,"kappa":1.0397156476974487,"row":1},{"col":0,"kappa":1.1047667264938354,"row":2},{"col":1,"kappa":1.5524996519088745,"row":2},{"col":2,"kappa":2.758542776107788,"row":2},{"col":3,"kappa":3.7919411659240723,"row":2},{"col":4,"kappa":3.7919411659240723,"row":2},{"col":5,"kappa":2.758542776107788,"row":2},{"col":6,"kappa":1.5524996519088745,"row":2},{"col":7,"kappa":1.1047667264938354,"row":2},{"col":0,"kappa":1.1772663593292236,"row":3},{"col":1,"kappa":1.999694585800171,"row":3},{"col":2,"kappa":3.7919411659240723,"row":3},{"col":3,"kappa":4.719764709472656,"row":3},{"col":4,"kappa":4.719764709472656,"row":3},{"col":5,"kappa":3.7919411659240723,"row":3},{"col":6,"kappa":1.999694585800171,"row":3},{"col":7,"kappa":1.1772663593292236,"row":3},{"col":0,"kappa":1.1772663593292236,"row":4},{"col":1,"kappa":1.999694585800171,"row":4},{"col":2,"kappa":3.7919411659240723,"row":4},{"col":3,"kappa":4.719764709472656,"row":4},{"col":4,"kappa":4.719764709472656,"row":4},{"col":5,"kappa":3.7919411659240723,"row":4},{"col":6,"kappa":1.999694585800171,"row":4},{"col":7,"kappa":1.1772663593292236,"row":4},{"col":0,"kappa":1.1047667264938354,"row":5},{"col":1,"kappa":1.5524996519088745,"row":5},{"col":2,"kappa":2.758542776107788,"row":5},{"col":3,"kappa":3.7919411659240723,"row":5},{"col":4,"kappa":3.7919411659240723,"row":5},{"col":5,"kappa":2.758542776107788,"row":5},{"col":6,"kappa":1.5524996519088745,"row":5},{"col":7,"kappa":1.1047667264938354,"row":5},{"col":0,"kappa":1.0397156476974487,"row":6},{"col":1,"kappa":1.1772663593292236,"row":6},{"col":2,"kappa":1.5524996519088745,"row":6},{"col":3,"kappa":1.999694585800171,"row":6},{"col":4,"kappa":1.999694585800171,"row":6},{"col":5,"kappa":1.5524996519088745,"row":6},{"col":6,"kappa":1.1772663593292236,"row":6},{"col":7,"kappa":1.0397156476974487,"row":6},{"col":0,"kappa":1.0109333992004395,"row":7},{"col":1,"kappa":1.0397156476974487,"row":7},{"col":2,"kappa":1.1047667264938354,"row":7},{"col":3,"kappa":1.1772663593292236,"row":7},{"col":4,"kappa":1.1772663593292236,"row":7},{"col":5,"kappa":1.1047667264938354,"row":7},{"col":6,"kappa":1.0397156476974487,"row":7},{"col":7,"kappa":1.0109333992004395,"row":7}]},"encoding":{"color":{"field":"kappa","scale":{"domain":[1.0,5.0],"scheme":"viridis"},"title":"kappa","type":"quantitative"},"x":{"field":"col","title":"x","type":"ordinal"},"y":{"field":"row","sort":"descending","title":"y","type":"ordinal"}},"height":250,"mark":"rect","title":"True Conductivity","width":250},{"data":{"values":[{"col":0,"kappa":2.9633327901983217,"row":0},{"col":1,"kappa":2.9864561583605864,"row":0},{"col":2,"kappa":2.9805502970008755,"row":0},{"col":3,"kappa":3.106185020155576,"row":0},{"col":4,"kappa":3.0590586500781107,"row":0},{"col":5,"kappa":2.852774598865966,"row":0},{"col":6,"kappa":2.8777105437600676,"row":0},{"col":7,"kappa":2.9090945474689134,"row":0},{"col":0,"kappa":2.943941518914351,"row":1},{"col":1,"kappa":3.0182706344368695,"row":1},{"col":2,"kappa":3.0252531893099435,"row":1},{"col":3,"kappa":3.02354057570116,"row":1},{"col":4,"kappa":2.9965207148514064,"row":1},{"col":5,"kappa":2.954341212439635,"row":1},{"col":6,"kappa":2.9931548945064295,"row":1},{"col":7,"kappa":3.1112275507910323,"row":1},{"col":0,"kappa":2.9621717250617388,"row":2},{"col":1,"kappa":2.9828733692452203,"row":2},{"col":2,"kappa":2.9629696691784306,"row":2},{"col":3,"kappa":2.9730650871902733,"row":2},{"col":4,"kappa":2.9809807750477146,"row":2},{"col":5,"kappa":2.9621636909306406,"row":2},{"col":6,"kappa":3.055868884013596,"row":2},{"col":7,"kappa":3.1096410527761376,"row":2},{"col":0,"kappa":2.925183400045412,"row":3},{"col":1,"kappa":2.9408090426464932,"row":3},{"col":2,"kappa":2.9157570266184822,"row":3},{"col":3,"kappa":2.9241129615865575,"row":3},{"col":4,"kappa":2.959191245888766,"row":3},{"col":5,"kappa":3.0097602386283135,"row":3},{"col":6,"kappa":3.017991640977605,"row":3},{"col":7,"kappa":3.0816789370777378,"row":3},{"col":0,"kappa":2.99169658312421,"row":4},{"col":1,"kappa":2.9546397387781083,"row":4},{"col":2,"kappa":2.962095622843681,"row":4},{"col":3,"kappa":2.9247070943428346,"row":4},{"col":4,"kappa":2.984551411490511,"row":4},{"col":5,"kappa":3.024317368440539,"row":4},{"col":6,"kappa":2.9913207883702606,"row":4},{"col":7,"kappa":3.0505319334360754,"row":4},{"col":0,"kappa":2.951118992713342,"row":5},{"col":1,"kappa":2.9589552347906447,"row":5},{"col":2,"kappa":2.9713733829158113,"row":5},{"col":3,"kappa":2.954977989228581,"row":5},{"col":4,"kappa":2.9737797228113565,"row":5},{"col":5,"kappa":2.981531665922289,"row":5},{"col":6,"kappa":3.018739511978554,"row":5},{"col":7,"kappa":3.0607987924433933,"row":5},{"col":0,"kappa":3.1220637302556344,"row":6},{"col":1,"kappa":3.01102062947895,"row":6},{"col":2,"kappa":3.0075354468545097,"row":6},{"col":3,"kappa":2.99470427215042,"row":6},{"col":4,"kappa":2.9753389703092314,"row":6},{"col":5,"kappa":3.0469163036981333,"row":6},{"col":6,"kappa":3.0451479401777295,"row":6},{"col":7,"kappa":2.9430541815317026,"row":6},{"col":0,"kappa":2.9965199194842143,"row":7},{"col":1,"kappa":2.8534915862928476,"row":7},{"col":2,"kappa":3.064392582517434,"row":7},{"col":3,"kappa":2.9377057064860055,"row":7},{"col":4,"kappa":2.9936231660262687,"row":7},{"col":5,"kappa":3.1497971949438437,"row":7},{"col":6,"kappa":3.0328120468982225,"row":7},{"col":7,"kappa":3.0761077136862918,"row":7}]},"encoding":{"color":{"field":"kappa","scale":{"domain":[1.0,5.0],"scheme":"viridis"},"title":"kappa","type":"quantitative"},"x":{"field":"col","title":"x","type":"ordinal"},"y":{"field":"row","sort":"descending","title":"y","type":"ordinal"}},"height":250,"mark":"rect","title":"Recovered Conductivity","width":250}]}

Sigma_obs Posterior

sigma_samples = Nx.to_flat_list(trace["sigma_obs"])

sigma_data =
  Enum.map(sigma_samples, fn s -> %{"sigma_obs" => s} end)

Vl.new(width: 400, height: 250, title: "Posterior: sigma_obs")
|> Vl.layers([
  Vl.new()
  |> Vl.data_from_values(sigma_data)
  |> Vl.mark(:bar, opacity: 0.7, color: "steelblue")
  |> Vl.encode_field(:x, "sigma_obs", type: :quantitative, bin: [maxbins: 30], title: "sigma_obs")
  |> Vl.encode(:y, aggregate: :count),
  Vl.new()
  |> Vl.data_from_values([%{"v" => sigma_obs_true}])
  |> Vl.mark(:rule, color: "red", stroke_width: 2, stroke_dash: [6, 3])
  |> Vl.encode_field(:x, "v", type: :quantitative)
])
{"$schema":"https://vega.github.io/schema/vega-lite/v5.json","height":250,"layer":[{"data":{"values":[{"sigma_obs":0.035212084427790606},{"sigma_obs":0.030645218081237505},{"sigma_obs":0.03228959123571987},{"sigma_obs":0.027322270649506743},{"sigma_obs":0.03203294815247456},{"sigma_obs":0.03139593652467564},{"sigma_obs":0.0349249622449424},{"sigma_obs":0.03552568626166169},{"sigma_obs":0.025886457596498772},{"sigma_obs":0.02443327091707623},{"sigma_obs":0.028934206320007202},{"sigma_obs":0.02450775721338328},{"sigma_obs":0.03625752341172144},{"sigma_obs":0.016866639865592975},{"sigma_obs":0.05810687503567036},{"sigma_obs":0.018576647591797633},{"sigma_obs":0.023062385418556073},{"sigma_obs":0.03390894640809098},{"sigma_obs":0.026270464526428654},{"sigma_obs":0.034879777155513075},{"sigma_obs":0.030239442054014125},{"sigma_obs":0.03251314103195065},{"sigma_obs":0.023071282141715877},{"sigma_obs":0.029711440630708547},{"sigma_obs":0.03176446872202455},{"sigma_obs":0.027586185421348674},{"sigma_obs":0.037469967734160836},{"sigma_obs":0.024461489592081},{"sigma_obs":0.040112280684841115},{"sigma_obs":0.0281640306968565},{"sigma_obs":0.031366270190334784},{"sigma_obs":0.03391763072969118},{"sigma_obs":0.02881907878180352},{"sigma_obs":0.027246015207668863},{"sigma_obs":0.022471946596364506},{"sigma_obs":0.03289928489741731},{"sigma_obs":0.026902579386107767},{"sigma_obs":0.030705878223010765},{"sigma_obs":0.03456272828584684},{"sigma_obs":0.04597069427406549},{"sigma_obs":0.021099574397908594},{"sigma_obs":0.03468856610630709},{"sigma_obs":0.05031111597289247},{"sigma_obs":0.022720102751858682},{"sigma_obs":0.04287198200465464},{"sigma_obs":0.032545898973952714},{"sigma_obs":0.027713121228527086},{"sigma_obs":0.037151933966668584},{"sigma_obs":0.023198304159537064},{"sigma_obs":0.032198739728889404},{"sigma_obs":0.022311593336633993},{"sigma_obs":0.04205874603460646},{"sigma_obs":0.03895740252884293},{"sigma_obs":0.03798696593796949},{"sigma_obs":0.02778575841583433},{"sigma_obs":0.028677099391233734},{"sigma_obs":0.031461833530831186},{"sigma_obs":0.03614185884660064},{"sigma_obs":0.026594156873747304},{"sigma_obs":0.03266468227806486},{"sigma_obs":0.02725656559462514},{"sigma_obs":0.03034342337461943},{"sigma_obs":0.027896754015730833},{"sigma_obs":0.03168834467296852},{"sigma_obs":0.03629574867032669},{"sigma_obs":0.05030201727861521},{"sigma_obs":0.0344276399743702},{"sigma_obs":0.04611263172738623},{"sigma_obs":0.02032553515540931},{"sigma_obs":0.031703456078174284},{"sigma_obs":0.02888916068636781},{"sigma_obs":0.022596327201829195},{"sigma_obs":0.040404417270868534},{"sigma_obs":0.02349047885517754},{"sigma_obs":0.07655179921742869},{"sigma_obs":0.07621601579555885},{"sigma_obs":0.026292926756414862},{"sigma_obs":0.030698221987549172},{"sigma_obs":0.02367226593316712},{"sigma_obs":0.02492052457700871},{"sigma_obs":0.023538084125926007},{"sigma_obs":0.04257597959326102},{"sigma_obs":0.02725229951579712},{"sigma_obs":0.02946848958913038},{"sigma_obs":0.030243448854112242},{"sigma_obs":0.029564999555352172},{"sigma_obs":0.028303657079868767},{"sigma_obs":0.02087088703159705},{"sigma_obs":0.05352743070646408},{"sigma_obs":0.03308113402207565},{"sigma_obs":0.04331103945068667},{"sigma_obs":0.05756894554934363},{"sigma_obs":0.02588475046440007},{"sigma_obs":0.036427297716715386},{"sigma_obs":0.033828963946261004},{"sigma_obs":0.03382992992748585},{"sigma_obs":0.0315071048332602},{"sigma_obs":0.028289252039390284},{"sigma_obs":0.027876646941758898},{"sigma_obs":0.0378670728997022},{"sigma_obs":0.029503220649086993},{"sigma_obs":0.045869337675050706},{"sigma_obs":0.027633303755353722},{"sigma_obs":0.029954410562460856},{"sigma_obs":0.03155082562174318},{"sigma_obs":0.02138428721808222},{"sigma_obs":0.05072113334793068},{"sigma_obs":0.018802668594131357},{"sigma_obs":0.04726825957893882},{"sigma_obs":0.05260466844978561},{"sigma_obs":0.05355093912831128},{"sigma_obs":0.021683411219804233},{"sigma_obs":0.03066015965268079},{"sigma_obs":0.04340974863674857},{"sigma_obs":0.045952377394455214},{"sigma_obs":0.019145145134776893},{"sigma_obs":0.05567231120553712},{"sigma_obs":0.02971163863764864},{"sigma_obs":0.030047446313235776},{"sigma_obs":0.06245034274037182},{"sigma_obs":0.06593860993715712},{"sigma_obs":0.061582403535043274},{"sigma_obs":0.057424832416842014},{"sigma_obs":0.04693802962032946},{"sigma_obs":0.03805259372344054},{"sigma_obs":0.0260626072101258},{"sigma_obs":0.02963097051476659},{"sigma_obs":0.033853402461739304},{"sigma_obs":0.026331906490241627},{"sigma_obs":0.03596455729062005},{"sigma_obs":0.025283414497674062},{"sigma_obs":0.02674180568967867},{"sigma_obs":0.026599370116788135},{"sigma_obs":0.03669157528033218},{"sigma_obs":0.028265591799817472},{"sigma_obs":0.02754535211376625},{"sigma_obs":0.05183130888777708},{"sigma_obs":0.01914864323147639},{"sigma_obs":0.03869051916625912},{"sigma_obs":0.042684631874354534},{"sigma_obs":0.024655511902585858},{"sigma_obs":0.02730922469748963},{"sigma_obs":0.02728103307675518},{"sigma_obs":0.028790613910693444},{"sigma_obs":0.03627298508700914},{"sigma_obs":0.024895438172864645},{"sigma_obs":0.03690337897076735},{"sigma_obs":0.026228343939894955},{"sigma_obs":0.039533788930137634},{"sigma_obs":0.06897781153077381},{"sigma_obs":0.021775255187509147},{"sigma_obs":0.02137014483712422},{"sigma_obs":0.030105047775070923},{"sigma_obs":0.03557979183641031},{"sigma_obs":0.03544295368943481},{"sigma_obs":0.04096481424934708},{"sigma_obs":0.029246232284218878},{"sigma_obs":0.022310588621084364},{"sigma_obs":0.05552094797089668},{"sigma_obs":0.0328935008691691},{"sigma_obs":0.021294379253476322},{"sigma_obs":0.045891231326964085},{"sigma_obs":0.0225391539043091},{"sigma_obs":0.028921831802230936},{"sigma_obs":0.024082489782144306},{"sigma_obs":0.02942716926372557},{"sigma_obs":0.030289519098011686},{"sigma_obs":0.025978970475686383},{"sigma_obs":0.034090417337878084},{"sigma_obs":0.03035151824999972},{"sigma_obs":0.03314969091075969},{"sigma_obs":0.032356918638405896},{"sigma_obs":0.025598570097283356},{"sigma_obs":0.03376331598090721},{"sigma_obs":0.0727376138681626},{"sigma_obs":0.024339511937001897},{"sigma_obs":0.024254119680878074},{"sigma_obs":0.028630738955745545},{"sigma_obs":0.02943538496439946},{"sigma_obs":0.02006816700384143},{"sigma_obs":0.031723647033279984},{"sigma_obs":0.04006880254030902},{"sigma_obs":0.02172246267724434},{"sigma_obs":0.0203114332783034},{"sigma_obs":0.02684699674012207},{"sigma_obs":0.04884427560908249},{"sigma_obs":0.05464120179910081},{"sigma_obs":0.04754876084491793},{"sigma_obs":0.021212390661216395},{"sigma_obs":0.06734738208419071},{"sigma_obs":0.019811954655730345},{"sigma_obs":0.03798381131271214},{"sigma_obs":0.045485591652170955},{"sigma_obs":0.03069022059237808},{"sigma_obs":0.026598295232127332},{"sigma_obs":0.03566468146463673},{"sigma_obs":0.021923221485691924},{"sigma_obs":0.04354717071047622},{"sigma_obs":0.04606070107151182},{"sigma_obs":0.03693325268253541},{"sigma_obs":0.01805406540165239},{"sigma_obs":0.048653870617240726},{"sigma_obs":0.04681366413416316},{"sigma_obs":0.040767238998785975},{"sigma_obs":0.02846712259839391},{"sigma_obs":0.03000036146411571},{"sigma_obs":0.036816416406612734},{"sigma_obs":0.023901452615273842},{"sigma_obs":0.03669926501134393},{"sigma_obs":0.033413413918824464},{"sigma_obs":0.021129845982608215},{"sigma_obs":0.0469852622003145},{"sigma_obs":0.043384864780672376},{"sigma_obs":0.05104177885850591},{"sigma_obs":0.02958306721686658},{"sigma_obs":0.030989070888535017},{"sigma_obs":0.02900510680246758},{"sigma_obs":0.02711473259766581},{"sigma_obs":0.03978273437126748},{"sigma_obs":0.02296058944901126},{"sigma_obs":0.03222132461620836},{"sigma_obs":0.05631886584244675},{"sigma_obs":0.06022755632983116},{"sigma_obs":0.02817862143394441},{"sigma_obs":0.02862880903886774},{"sigma_obs":0.03052323859573353},{"sigma_obs":0.02536019366598146},{"sigma_obs":0.055633074162852764},{"sigma_obs":0.034672326903036364},{"sigma_obs":0.029251343548188445},{"sigma_obs":0.025383724896285233},{"sigma_obs":0.0405964716204559},{"sigma_obs":0.037664408066597946},{"sigma_obs":0.02742703180423496},{"sigma_obs":0.028407862006936107},{"sigma_obs":0.030568166171627744},{"sigma_obs":0.02800045295038306},{"sigma_obs":0.029088518711465356},{"sigma_obs":0.02936815685082059},{"sigma_obs":0.031725630824681744},{"sigma_obs":0.04450648376121401},{"sigma_obs":0.02006388148676087},{"sigma_obs":0.0500796879745298},{"sigma_obs":0.028412953125800947},{"sigma_obs":0.023955203902840318},{"sigma_obs":0.037026709421524896},{"sigma_obs":0.03273214732383727},{"sigma_obs":0.02433470888781895},{"sigma_obs":0.042666196136675515},{"sigma_obs":0.04258884849364962},{"sigma_obs":0.04866454615353696},{"sigma_obs":0.025024961680300606},{"sigma_obs":0.028949785308025736},{"sigma_obs":0.034008818157112784},{"sigma_obs":0.028960115480311858},{"sigma_obs":0.02384456560031853},{"sigma_obs":0.031108864627097284},{"sigma_obs":0.028190735329888573},{"sigma_obs":0.02889775738817985},{"sigma_obs":0.04136121377190112},{"sigma_obs":0.04337834833348149},{"sigma_obs":0.03454718689171794},{"sigma_obs":0.026510225541875453},{"sigma_obs":0.036728474975334754},{"sigma_obs":0.0349826268504171},{"sigma_obs":0.03496439118661874},{"sigma_obs":0.0269855843245864},{"sigma_obs":0.045978100094094025},{"sigma_obs":0.038124870043195924},{"sigma_obs":0.02471742343869572},{"sigma_obs":0.025276332671280443},{"sigma_obs":0.02977172433013306},{"sigma_obs":0.032747371089713394},{"sigma_obs":0.04717409068878374},{"sigma_obs":0.0481001554826275},{"sigma_obs":0.045416909199107285},{"sigma_obs":0.022064586270435047},{"sigma_obs":0.04036084074371313},{"sigma_obs":0.02220152448108633},{"sigma_obs":0.04290921288641775},{"sigma_obs":0.019525688470315085},{"sigma_obs":0.034220627858054985},{"sigma_obs":0.027532576757659147},{"sigma_obs":0.021950137031564062},{"sigma_obs":0.05794543107785173},{"sigma_obs":0.06944579168472702},{"sigma_obs":0.07245293266580888},{"sigma_obs":0.05481114556137345},{"sigma_obs":0.05061184393248017},{"sigma_obs":0.021426484532521977},{"sigma_obs":0.03635813240959556},{"sigma_obs":0.025025139838497724},{"sigma_obs":0.029581895272524087},{"sigma_obs":0.029599081509747358},{"sigma_obs":0.029161043360334888},{"sigma_obs":0.03145054284085134},{"sigma_obs":0.02448210446621562},{"sigma_obs":0.025784520449831982},{"sigma_obs":0.035155701507314584},{"sigma_obs":0.03122319712004632},{"sigma_obs":0.03254259398318232},{"sigma_obs":0.028572878480460723},{"sigma_obs":0.027256141590466413},{"sigma_obs":0.04018974303047435},{"sigma_obs":0.031035339002977887},{"sigma_obs":0.05509145161169383},{"sigma_obs":0.052373528289742614},{"sigma_obs":0.04165121839961745},{"sigma_obs":0.030481776508178775},{"sigma_obs":0.033820275147110904},{"sigma_obs":0.02768288272634888},{"sigma_obs":0.025744188737382635},{"sigma_obs":0.034867575148441216},{"sigma_obs":0.029860952566983182},{"sigma_obs":0.025017084336152063},{"sigma_obs":0.027049542685414873},{"sigma_obs":0.021088990507505883},{"sigma_obs":0.030853238725255815},{"sigma_obs":0.03257631455522385},{"sigma_obs":0.031561926892592525},{"sigma_obs":0.02515383129313323},{"sigma_obs":0.027430689582255367},{"sigma_obs":0.02209348523191999},{"sigma_obs":0.0228173411798834},{"sigma_obs":0.0361024506534797},{"sigma_obs":0.03203744333459546},{"sigma_obs":0.035394424748623886},{"sigma_obs":0.03621579293499682},{"sigma_obs":0.04632757905431398},{"sigma_obs":0.01956403045136559},{"sigma_obs":0.06514743291545261},{"sigma_obs":0.03168722853662327},{"sigma_obs":0.047551099727348785},{"sigma_obs":0.034559869490269166},{"sigma_obs":0.028042874423681693},{"sigma_obs":0.029658723291580404},{"sigma_obs":0.02504421967300826},{"sigma_obs":0.02496658707147291},{"sigma_obs":0.024721336710076418},{"sigma_obs":0.0266768834624742},{"sigma_obs":0.032879638147848736},{"sigma_obs":0.023718331662153577},{"sigma_obs":0.039293956560559486},{"sigma_obs":0.033046984746219824},{"sigma_obs":0.027368853924337733},{"sigma_obs":0.028648390907769647},{"sigma_obs":0.03157530852706185},{"sigma_obs":0.024472368914309},{"sigma_obs":0.02881733467478921},{"sigma_obs":0.034691700499716115},{"sigma_obs":0.028840517568294484},{"sigma_obs":0.03334456831869451},{"sigma_obs":0.03822693697571863},{"sigma_obs":0.023837391036488274},{"sigma_obs":0.04250275753172172},{"sigma_obs":0.041500034827288305},{"sigma_obs":0.028564137995059413},{"sigma_obs":0.03503670404974648},{"sigma_obs":0.023040278583819924},{"sigma_obs":0.040180203463291105},{"sigma_obs":0.020082213387734877},{"sigma_obs":0.027059860872582848},{"sigma_obs":0.024626230918867384},{"sigma_obs":0.03442389199132028},{"sigma_obs":0.046532715656124705},{"sigma_obs":0.045616576403864566},{"sigma_obs":0.023526059833504132},{"sigma_obs":0.02455500113468858},{"sigma_obs":0.051382363134732593},{"sigma_obs":0.021489916491881162},{"sigma_obs":0.02474225923567807},{"sigma_obs":0.03379617661787294},{"sigma_obs":0.02694695466489432},{"sigma_obs":0.032212467559711036},{"sigma_obs":0.048229278523590595},{"sigma_obs":0.02052628586898114},{"sigma_obs":0.020086526153387834},{"sigma_obs":0.04602163383757957},{"sigma_obs":0.04967840039410833},{"sigma_obs":0.051123348911302526},{"sigma_obs":0.02632935391585984},{"sigma_obs":0.0354363151247383},{"sigma_obs":0.028167128698060803},{"sigma_obs":0.03187277539310265},{"sigma_obs":0.030723245021570336},{"sigma_obs":0.022966234447688107},{"sigma_obs":0.051947577840438255},{"sigma_obs":0.05500557025380343},{"sigma_obs":0.036001116879309396},{"sigma_obs":0.03809664691789402},{"sigma_obs":0.04849400769165091},{"sigma_obs":0.02876626192009953},{"sigma_obs":0.05842627464476342},{"sigma_obs":0.03318055612342812},{"sigma_obs":0.024606641410086744},{"sigma_obs":0.028727321997625708},{"sigma_obs":0.03317686256223143},{"sigma_obs":0.02869523723422928},{"sigma_obs":0.04910395316945205},{"sigma_obs":0.03329299594581242},{"sigma_obs":0.03265472776839922},{"sigma_obs":0.02577814551043577},{"sigma_obs":0.04340632360574244},{"sigma_obs":0.033777473698573736},{"sigma_obs":0.02925043980816175},{"sigma_obs":0.03348077721088499},{"sigma_obs":0.03543038481279242},{"sigma_obs":0.01982364205000445},{"sigma_obs":0.041502077366860776},{"sigma_obs":0.02687051940751864},{"sigma_obs":0.03231468060698278},{"sigma_obs":0.0428659347182932},{"sigma_obs":0.023441208030556937},{"sigma_obs":0.041893896091824236},{"sigma_obs":0.022508611249847547},{"sigma_obs":0.022820532791259916},{"sigma_obs":0.03145090681895451},{"sigma_obs":0.04175204408893009},{"sigma_obs":0.027642564101906117},{"sigma_obs":0.029631256697623856},{"sigma_obs":0.026313110323602215},{"sigma_obs":0.06694381588020508},{"sigma_obs":0.07483873573501482},{"sigma_obs":0.028053345787694185},{"sigma_obs":0.030741207008502254},{"sigma_obs":0.042975248567321316},{"sigma_obs":0.025905170831294587},{"sigma_obs":0.034772575170458384},{"sigma_obs":0.04026562078206785},{"sigma_obs":0.023184868849733933},{"sigma_obs":0.04402436562231941},{"sigma_obs":0.021433315884234378},{"sigma_obs":0.023906414229756836},{"sigma_obs":0.025922266318283434},{"sigma_obs":0.032737284776073655},{"sigma_obs":0.04573296325743993},{"sigma_obs":0.02163848694100101},{"sigma_obs":0.029219571935360247},{"sigma_obs":0.03052453195078693},{"sigma_obs":0.032202124235338504},{"sigma_obs":0.027883511595361656},{"sigma_obs":0.0262059093867792},{"sigma_obs":0.033888290118326496},{"sigma_obs":0.02663964455639899},{"sigma_obs":0.03275589439901578},{"sigma_obs":0.021619500303868448},{"sigma_obs":0.03209469748910511},{"sigma_obs":0.024647445372064355},{"sigma_obs":0.03773115020590215},{"sigma_obs":0.037211752188187994},{"sigma_obs":0.024911707575283815},{"sigma_obs":0.025060466708429364},{"sigma_obs":0.03516489666427177},{"sigma_obs":0.03653465093684448},{"sigma_obs":0.024321458661248966},{"sigma_obs":0.023435131297596955},{"sigma_obs":0.04117948945923724},{"sigma_obs":0.046712604020086076},{"sigma_obs":0.020594495395614438},{"sigma_obs":0.03698678949664288},{"sigma_obs":0.018011841540969777},{"sigma_obs":0.044483969243482566},{"sigma_obs":0.025154901420297426},{"sigma_obs":0.026041183773840494},{"sigma_obs":0.03195428398088878},{"sigma_obs":0.028484818451079377},{"sigma_obs":0.040645938344662964},{"sigma_obs":0.023222178387371575},{"sigma_obs":0.025446920870641324},{"sigma_obs":0.04163731407143495},{"sigma_obs":0.0298422075993797},{"sigma_obs":0.029452084963807978},{"sigma_obs":0.03553178092190235},{"sigma_obs":0.030376147139283525},{"sigma_obs":0.017995528782148818},{"sigma_obs":0.059544158181951655},{"sigma_obs":0.06686471444952014},{"sigma_obs":0.02903270392689151},{"sigma_obs":0.03490885130693214},{"sigma_obs":0.0255678446645002},{"sigma_obs":0.027765577339943767},{"sigma_obs":0.04915889805015963},{"sigma_obs":0.0473172925694933},{"sigma_obs":0.02142942252339385},{"sigma_obs":0.02495952422769016},{"sigma_obs":0.05109522403615791},{"sigma_obs":0.01967369592125994},{"sigma_obs":0.042944457851781755},{"sigma_obs":0.030399126247605605},{"sigma_obs":0.020635913181089278},{"sigma_obs":0.06727772517755398},{"sigma_obs":0.03950671747607791},{"sigma_obs":0.03654747137319289},{"sigma_obs":0.03662607825145946},{"sigma_obs":0.0358564032289995},{"sigma_obs":0.02167448971669709},{"sigma_obs":0.04463096683368453},{"sigma_obs":0.04741562918296545},{"sigma_obs":0.04242496972013627},{"sigma_obs":0.02215317933408811}]},"encoding":{"x":{"bin":{"maxbins":30},"field":"sigma_obs","title":"sigma_obs","type":"quantitative"},"y":{"aggregate":"count"}},"mark":{"color":"steelblue","opacity":0.7,"type":"bar"}},{"data":{"values":[{"v":0.02}]},"encoding":{"x":{"field":"v","type":"quantitative"}},"mark":{"color":"red","strokeDash":[6,3],"strokeWidth":2,"type":"rule"}}],"title":"Posterior: sigma_obs","width":400}

Scaling Up

This 8x8 (65-parameter) demo runs in under a minute on GPU. For real applications:

  • 16x16 grid (257 params): feasible with 1000 warmup + 1000 draws, ~5-10 min on GPU
  • 32x32 grid (1025 params): needs careful step-size tuning and possibly dense mass matrix
  • GPU acceleration: device: :cuda amortizes the per-leapfrog cost of the forward PDE solve
  • Distributed sampling: Sampler.sample_chains/4 runs independent chains on separate BEAM nodes
  • Better priors: Replace Laplacian with a full Gaussian Random Field (Matern kernel via SPDE) for anisotropic smoothness control

Other notebooks:

Back