Level Set Bayesian Inverse Problem

backend_dep =
  case :os.type() do
    {:unix, :darwin} -> {:emlx, "~> 0.2"}
    _ -> {:exla, "~> 0.10"}
  end

Mix.install([
  {:exmc, path: Path.expand("../", __DIR__)},
  backend_dep,
  {:kino_vega_lite, "~> 0.1"}
])

alias VegaLite, as: Vl
alias Exmc.{Builder, NUTS.Sampler, Dist}
alias Exmc.Physics.{LevelSet, 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"
)

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"
)

# 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"
)

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)

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)")

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}")

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
)

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)"
)

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
)

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)
])

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