Gradients: the composed adjoint#

The objective is a scalar function of hundreds of design variables, so the gradient is computed in reverse mode: one adjoint solve per solver per evaluation, whatever the number of variables. Three different “adjoints” are chained, each living where its solver lives.

        flowchart TB
  subgraph r1[" "]
    direction LR
    a["θ → filter H → N(θ)<br/><i>JAX</i>"] -- "N" --> ct["<b>ChargeTransport</b> · Julia<br/>drift-diffusion, 0 V and −5 V<br/><i>adjoint: Jᵀλ = −g</i>"]
    ct -- "n, p" --> b["Soref–Bennett → mesh transfer<br/><i>JAX</i>"]
  end
  subgraph r2[" "]
    direction LR
    gy["<b>gyptis</b> · FEniCS<br/>eigensolve A x = λ B x<br/><i>∂λ/∂ε = yᴴ(∂A − λ∂B)x / yᴴBx</i>"] -- "n_eff²" --> c["J = Δn_eff − w·α<br/><i>JAX</i>"]
    c -. "jax.grad: same chain, reversed" .-> g(["∂J/∂θ"])
  end
  r1 -- "ε per design cell" --> r2
  style r1 fill:none,stroke:none
  style r2 fill:none,stroke:none

jax.grad(J) walks this chain backwards. Each Tesseract call enters the JAX program through tesseract-jax’s apply_tesseract: its forward is the component’s apply endpoint and its backward is the component’s vector_jacobian_product, with abstract_eval reporting the output shapes at trace time; everything between the two solvers is ordinary JAX autodiff.

Discrete adjoint of the drift-diffusion solve#

At each bias the solver finds the state \(x = (\psi, \varphi_n, \varphi_p)\) of the finite-volume residual \(F(x; N) = 0\) and reads out the carrier densities \(n(x), p(x)\) per node. Given the cotangents \(\bar n, \bar p\) arriving from downstream, the Tesseract’s VJP returns \(\bar N = \mathrm{d}J/\mathrm{d}N\):

  1. State gradient. \(g = (\partial n/\partial x)^\top \bar n + (\partial p/\partial x)^\top \bar p\), assembled per DOF by differencing the (smooth, local) density read-out.

  2. Adjoint solve. \(J_F^\top \lambda = -g\) with \(J_F = \partial F/\partial x\) the assembled Jacobian VoronoiFVM already builds for Newton. The Ohmic contacts are enforced by a \(10^{30}\) Dirichlet penalty on their rows, which makes the full Jacobian’s condition number \(\sim 10^{42}\); those rows are eliminated first (the constrained state is known exactly) and the physical block is factorized with partial pivoting. The sensitivity of the pinned DOFs re-enters through the well-scaled coupling column \(\lambda_\mathrm{free}^\top J_F[\mathrm{free}, c]\).

  3. Explicit doping term. Doping enters the bulk residual only through the Poisson source, so \(\bar N_i \mathrel{+}= \lambda_{\psi,i}\, q\, \lambda_1\, |\Omega_i|\) (node volume) — a contracted analytic derivative, not a residual finite difference.

  4. Contact path. Out of equilibrium the contacts are held at \(\psi = \psi_\mathrm{eq} + U\), and \(\psi_\mathrm{eq}\) is itself an output of the equilibrium solve — a local function of the doping at the contact node through charge neutrality, \(\mathrm{d}\psi_\mathrm{eq}/\mathrm{d}N = 1 / (\mathrm{d}/\mathrm{d}\psi \sum_\alpha z_\alpha n_\alpha)\) (\(-U_T/(n+p)\) for Boltzmann statistics). At reverse bias the rib is depleted and this path carries most of the gradient; dropping it leaves the composed gradient wrong by more than half.

The adjoint is taken at the warm-started converged state the forward returned, and the sign/unit conversion between PRISMO’s donor-positive \(\mathrm{cm^{-3}}\) doping and the solver’s acceptor-positive SI convention is undone on the way out. Two forward solves and two adjoint solves per iteration (0 V and −5 V); the loss term rides on the 0 V adjoint at no extra cost.

Hellmann–Feynman eigen-adjoint#

For the generalized eigenproblem \(A(\varepsilon)\, x = \lambda\, B(\varepsilon)\, x\) with right eigenvector \(x\) and left eigenvector \(y\) (\(y^H A = \lambda\, y^H B\)), the first-order change of the eigenvalue with a per-cell permittivity \(\varepsilon_c\) is

\[ \frac{\partial \lambda}{\partial \varepsilon_c} = \frac{y^H \bigl(\partial A/\partial\varepsilon_c - \lambda\, \partial B/\partial\varepsilon_c\bigr) x} {y^H B\, x}, \qquad \frac{\partial (n_\mathrm{eff}^2)}{\partial \varepsilon_c} = \frac{1}{k_0^2}\,\frac{\partial\lambda}{\partial\varepsilon_c}. \]

The gyptis Tesseract recovers \(x\) and \(y\) from one two-sided SLEPc solve (shared by forward and adjoint), and assembles the numerator for all design cells in a single pass by testing the \(\varepsilon\)-derivative density of the weak form against a DG0 test function — no per-cell loop, no re-solve. The result is the field \(w_c = \partial(n_\mathrm{eff}^2)/\partial\varepsilon_c\) on the design cells; the VJP multiplies it by the incoming scalar cotangent. The same \(w_c\), evaluated once at the uniform background, is the mode-overlap weight of the modal loss.

pyadjoint does not tape the SLEPc eigensolve, which is why this adjoint is written by hand. The shift-invert transform behind the solve is factorized with a pivoting direct solver (UMFPACK): the assembled system is an indefinite saddle point, and PETSc’s native non-pivoting LU left a \(2\times10^{-3}\) relative noise floor on \(\Delta n_\mathrm{eff}\) that the adjoint could not explain and the optimizer could not descend through. With pivoting the eigenvalue is smooth along a design line to \(10^{-13}\) and the noise floor is \(2\times10^{-11}\).

The JAX glue#

  • Density filter: a dense \((n_\mathrm{design}, n_\mathrm{design})\) matrix; its VJP is the transpose, written as a custom_vjp so the matrix is not traced.

  • Doping map \(N(\theta)\): custom JVP so the derivative is the even, continuous \(N_\mathrm{ref}\ln 10\, s\, 10^{s|\theta|}\) at the junction too.

  • Soref–Bennett: the signed power \(\operatorname{sign}(x)|x|^p\) gets a custom VJP (\(p|x|^{p-1}\), set to zero at \(x = 0\) where \(p < 1\) would blow up) so depletion differentiates cleanly.

  • Mesh transfer: a constant restriction matrix (1/3 per vertex, zero elsewhere).

  • Design nodes: only silicon nodes carry \(\theta\); the filtered field is scattered into full node order before the doping map, so the transposes are plain gathers.

Validation#

prismo validate-gradient (make validate-gradient-containers) compares the composed adjoint’s directional derivative \(\nabla J \cdot d\) against central finite differences \([J(\theta + h d) - J(\theta - h d)] / 2h\) along sampled unit directions \(d\), for \(h \in \{10^{-2}, 10^{-3}, 10^{-4}\}\), and gates on a stated relative tolerance. Across the real Julia + FEniCS boundary the error follows the \(O(h^2)\) slope to \(\approx 2\times10^{-6}\) at \(h = 10^{-3}\) before finite-difference round-off takes over (the drift-diffusion adjoint reads the density through an internal \(10^{-8}\) difference, and the eigensolver carries its tolerance). --cold resets the ChargeTransport worker before every evaluation so no finite-difference sample inherits another’s Newton start.

prismo probe-objective is the complementary smoothness gauge: it samples \(J(\theta_0 + t\,d)\) at uniform spacing, fits a quadratic, and reports the fit residual, the white-noise amplitude implied by second differences, and the adjoint’s slope against the fitted one. A kink shows as a structured residual that shrinks with the spacing; an evaluation noise floor shows as a white residual that does not.