---
jupytext:
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.19.6
  main_language: python
---

# The eikonal equation on a Möbius strip

This demo computes the distance to the boundary of a Möbius strip by
solving the eikonal equation with the latent variable proximal point
(LVPP) algorithm of {cite:t}`eikonal-dokken2025lvpp`. It illustrates
how to

* solve a problem with a pointwise inequality constraint, here
  $|\nabla_\Omega u| \le 1$, without ever forming the constraint,
* use an $H(\mathrm{div})$-conforming latent variable on a surface
  that cannot be oriented.

```{admonition} Difference from the paper
:class: note
{cite:t}`eikonal-dokken2025lvpp` solves the eikonal equation in this
Raviart-Thomas mixed form on flat surfaces, and on manifolds without a
Raviart-Thomas space. What is new here is the combination: a
Raviart-Thomas latent variable on a manifold, and one that cannot be
oriented at that.
```

```{admonition} Download sources
:class: download
* {download}`Python script <./demo_mobius_eikonal.py>`
* {download}`Jupyter notebook <./demo_mobius_eikonal.ipynb>`
```

## Equation and problem definition

The distance to the boundary of a domain $\Omega$ solves the eikonal
equation

$$
  |\nabla_\Omega u| = 1 \text{ in } \Omega, \qquad
  u = 0 \text{ on } \partial\Omega,
$$

where $\nabla_\Omega$ is the gradient along the surface. It is the
limit of a problem with an inequality constraint: among all $u$ with
$|\nabla_\Omega u| \le 1$ and $u = 0$ on the boundary, the distance
function is the largest, so it maximises $\int_\Omega f u$ for
$f \equiv 1$.

### The latent variable

LVPP never imposes the constraint directly. It introduces a *latent*
variable $\psi$, unconstrained, and recovers the gradient from it
through a pointwise map whose range is exactly the feasible set. For
the unit ball that map is

$$
  \nabla_\Omega u = \frac{\psi}{\sqrt{1 + |\psi|^2}},
$$

so $|\nabla_\Omega u| < 1$ holds for every $\psi$, and the constraint
becomes active only in the limit $|\psi| \to \infty$. The
non-smoothness of the constraint is thereby traded for the growth of
$\psi$, which is what makes each subproblem a smooth one that Newton's
method handles.

The proximal point iteration supplies that limit. Given the previous
latent variable $\psi^{k-1}$ and a step size $\alpha_k > 0$, the
$k$-th subproblem is to find $(u^k, \psi^k)$ with

$$
\begin{aligned}
  \nabla_\Omega \cdot \psi^k
    &= \nabla_\Omega \cdot \psi^{k-1} - \alpha_k f, \\
  \nabla_\Omega u^k
    &= \frac{\psi^k}{\sqrt{1 + |\psi^k|^2}}.
\end{aligned}
$$

Increasing $\alpha_k$ drives $|\psi^k|$ up and $|\nabla_\Omega u^k|$
towards one. The iteration stops when $u$ settles.

### Mixed form

Taking $\psi$ in an $H(\mathrm{div})$-conforming space and $u$ in a
discontinuous Lagrange space, and testing the second equation against
$\tau$ with an integration by parts that leaves no boundary term
because $u = 0$ on $\partial\Omega$, the subproblem reads: find
$(u, \psi) \in Q \times V$ with

$$
\begin{aligned}
  \int_\Omega (\nabla_\Omega \cdot \psi) \, v ~\mathrm{d}x
    &= \int_\Omega (\nabla_\Omega \cdot \psi^{k-1}) \, v ~\mathrm{d}x
     - \alpha_k \int_\Omega f \, v ~\mathrm{d}x, \\
  \int_\Omega u \, (\nabla_\Omega \cdot \tau) ~\mathrm{d}x
    + \int_\Omega \frac{\psi \cdot \tau}{\sqrt{1 + |\psi|^2}}
      ~\mathrm{d}x &= 0,
\end{aligned}
$$

for all $(v, \tau) \in Q \times V$. The Dirichlet condition on $u$ is
natural here, so no boundary condition is imposed on the space.

Raviart-Thomas is mapped by the contravariant Piola map, which on a
surface needs the two cells of each facet to agree on its normal. On a
Möbius strip no orientation of the cells achieves that, so the facets
where they disagree carry a sign instead; see {doc}`demo_mobius`.

## Implementation

```python

import contextlib

from mpi4py import MPI

import numpy as np

import basix
import basix.ufl
import dolfinx
import dolfinx.fem.petsc
import ufl
```

The strip is `n` cells around and `m` across, of half-width `width`,
on a mesh of degree `geometry_degree`. The latent variable is of
degree `degree` and $u$ one lower.

```python
n, m = 48, 8
geometry_degree = 2
degree = 3
width = 0.92
```

### The mesh

The strip is built from its parametrisation, the half twist being that
the column closing the ring is column $0$ read upside down. Equispaced
Lagrange nodes lie on a lattice of step $1/k$, so taking every node
from the $k$-times finer parameter grid makes neighbouring cells share
the nodes on their common edge, which is what keeps the geometry
continuous. See {doc}`demo_mobius`, which builds the same surface.

```python
def surface_point(t, s):
    """Map the parameters onto the strip."""
    rho = 1 + s * np.cos(t / 2)
    z = s * np.sin(t / 2)
    return np.column_stack([rho * np.cos(t), rho * np.sin(t), z])
```

```python
def mobius_mesh(n, m, width, k):
    """A Möbius strip, ``n`` cells around and ``m`` across."""
    dtype = dolfinx.default_real_type
    lattice = basix.create_element(
        basix.ElementFamily.P, basix.CellType.triangle, k, basix.LagrangeVariant.equispaced
    ).points
    if MPI.COMM_WORLD.rank == 0:
        ni, nj = k * n, k * m

        def node(i, j):
            return i * (nj + 1) + j if i < ni else nj - j  # the half twist

        i_grid = np.repeat(np.arange(ni), nj + 1)
        j_grid = np.tile(np.arange(nj + 1), ni)
        x = surface_point(2 * np.pi * i_grid / ni, -width + 2 * width * j_grid / nj)
        cells = []
        for i in range(n):
            for j in range(m):
                a = np.array([k * i, k * j])
                b = np.array([k * (i + 1), k * j])
                c = np.array([k * i, k * (j + 1)])
                d = np.array([k * (i + 1), k * (j + 1)])
                for p0, p1, p2 in ((a, b, c), (b, d, c)):
                    ij = np.rint(
                        p0 + np.outer(lattice[:, 0], p1 - p0) + np.outer(lattice[:, 1], p2 - p0)
                    ).astype(np.int64)
                    cells.append([node(int(p), int(q)) for p, q in ij])
        cells = np.asarray(cells, dtype=np.int64)
    else:
        x = np.zeros((0, 3))
        cells = np.zeros((0, len(lattice)), dtype=np.int64)
    domain = ufl.Mesh(
        basix.ufl.element(
            "Lagrange", "triangle", k, basix.LagrangeVariant.equispaced, shape=(3,), dtype=dtype
        )
    )
    return dolfinx.mesh.create_mesh(MPI.COMM_WORLD, cells, domain, x.astype(dtype))
```

```python
mesh = mobius_mesh(n, m, width, geometry_degree)
mesh.topology.create_facet_orientations()

V = dolfinx.fem.functionspace(mesh, ("RT", degree))
Q = dolfinx.fem.functionspace(mesh, ("DG", degree - 1))
W = ufl.MixedFunctionSpace(Q, V)

u = dolfinx.fem.Function(Q, name="u")
psi = dolfinx.fem.Function(V, name="psi")
v, tau = ufl.TestFunctions(W)

# The previous iterate
u0 = dolfinx.fem.Function(Q)
psi0 = dolfinx.fem.Function(V)
```

### The subproblem

Each step solves the two equations above as one nonlinear system, with
the step size $\alpha_k$ carried by a {py:class}`Constant
<dolfinx.fem.Constant>` so that the forms are compiled once.

```python
dx = ufl.Measure("dx", domain=mesh)
f = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1))
alpha = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1))

latent_map = 1 / ufl.sqrt(1 + ufl.dot(psi, psi))
F = ufl.inner(ufl.div(psi) - ufl.div(psi0), v) * dx + alpha * ufl.inner(f, v) * dx
F += ufl.inner(u, ufl.div(tau)) * dx + latent_map * ufl.inner(psi, tau) * dx

F_blocked: list[ufl.Form] = ufl.extract_blocks(F)  # type: ignore[assignment]
tol = max(1e-5, 1000 * float(np.finfo(dolfinx.default_real_type).eps))
problem = dolfinx.fem.petsc.NonlinearProblem(
    F_blocked,
    [u, psi],
    bcs=[],
    petsc_options={
        "snes_type": "newtonls",
        "snes_rtol": tol,
        "snes_atol": tol,
        "snes_max_it": 100,
        "snes_error_if_not_converged": True,
        "ksp_type": "preonly",
        "pc_type": "lu",
        "pc_factor_mat_solver_type": "mumps",
    },
    petsc_options_prefix="lvpp_step_",
)
```

### Output

VTX builds its points from the dofmap, so neither a Raviart-Thomas
function nor the gradient can be written directly. Both are
interpolated into discontinuous Lagrange spaces of the degree of the
geometry first.

```python
S = dolfinx.fem.functionspace(mesh, ("DG", mesh.geometry.cmaps[0].degree))
u_out = dolfinx.fem.Function(S, name="u")

S_vec = dolfinx.fem.functionspace(mesh, ("DG", mesh.geometry.cmaps[0].degree, (mesh.geometry.dim,)))
grad_out = dolfinx.fem.Function(S_vec, name="grad_u")
psi_out = dolfinx.fem.Function(S_vec, name="psi")
grad_expr = dolfinx.fem.Expression(ufl.grad(u), S_vec.element.interpolation_points)
```

### The proximal point loop

The step size is increased geometrically, up to a cap, and the
iteration stops once $u$ stops moving, measured in $L^2$.

```python
change = dolfinx.fem.form(ufl.inner(u - u0, u - u0) * dx)
newton_iterations = []

with contextlib.ExitStack() as stack:
    writers = []
    if dolfinx.has_adios2:
        writers = [
            stack.enter_context(dolfinx.io.VTXWriter(mesh.comm, filename, functions, engine="BP4"))
            for filename, functions in [
                ("mobius_eikonal_u.bp", [u_out]),
                ("mobius_eikonal_psi.bp", [grad_out, psi_out]),
            ]
        ]

    for i in range(1, 100):
        alpha.value = min(2**i, 10)
        problem.solve()
        newton_iterations.append(problem.solver.getIterationNumber())

        diff = np.sqrt(abs(mesh.comm.allreduce(dolfinx.fem.assemble_scalar(change), op=MPI.SUM)))
        if mesh.comm.rank == 0:
            print(
                f"LVPP {i}: alpha={alpha.value.real:g}, "
                f"Newton iterations={newton_iterations[-1]}, |u - u_prev|={diff:.3e}"
            )

        u0.x.array[:] = u.x.array
        psi0.x.array[:] = psi.x.array

        u_out.interpolate(u)
        grad_out.interpolate(grad_expr)
        psi_out.interpolate(psi)
        for writer in writers:
            writer.write(i)

        if diff < 5 * tol:
            break


if mesh.comm.rank == 0:
    print(f"{i} LVPP iterations, {sum(newton_iterations)} Newton iterations in total")
```

The constraint is satisfied throughout, and becomes active away from
the kink of the distance function. We report the mean slope, which
approaches one as the iteration converges.

```python
one = dolfinx.fem.Constant(mesh, dolfinx.default_scalar_type(1))
grad_norm = ufl.sqrt(ufl.inner(ufl.grad(u), ufl.grad(u)))


def integrate(expr):
    """Integrate a real-valued expression over the surface."""
    local = dolfinx.fem.assemble_scalar(dolfinx.fem.form(expr * dx))
    return abs(mesh.comm.allreduce(local, op=MPI.SUM))


mean_slope = integrate(grad_norm) / integrate(one)
if mesh.comm.rank == 0:
    print(f"mean |grad(u)| = {mean_slope:.6f}  (1 in the limit)")
```

## References
```{bibliography}
   :filter: cited
   :labelprefix:
   :keyprefix: eikonal-
```