Skip to main content

Section 1.6 Physics-Informed Machine Learning

The physics constraints and model architecture used in Scriber Labs physics-informed machine learning projects are discussed in detail in the following sections.

Subsection 1.6.1 Physics constraints and model architecture

Subsubsection 1.6.1.1 L2 Normalization of Learned Wavefunctions

In quantum mechanics, the Born rule requires that the total probability of finding a particle anywhere in space equals unity. For a one-dimensional wavefunction \(\psi(x)\text{,}\) this requires
\begin{equation*} \int_{-\infty}^{\infty}{ | \psi(x) |^2 \, dx} = 1 \, . \end{equation*}
On a uniform spatial grid with spacing \(\Delta x\text{,}\) the integral is replaced by a discrete sum, yielding the constraint
\begin{equation*} \sum_{i}{ |\psi(x_i) |^2 \, \Delta x} = 1 \, . \end{equation*}
Any neural network that parameterizes a wavefunction must respect this constraint for its output to be physically interpretable as a quantum state.

Remark 1.6.1. Normalization in Scriber Labs PIML Project.

Note that the following remark pertains specifically to the Scriber Labs project, low-fidelity inverse Scrödinger solver via physics-informed machine learning.
In the context of the Rather than enforcing normalization as a soft penalty in the loss function — which would compete with other loss terms and never guarantee exact satisfaction—we enforce it by construction inside the forward pass of the network. Given the raw output \(\psi_{\mathrm{raw}}(x)\) of a multilayer perceptron, we rescale it to unit norm:
\begin{equation*} \psi(x) = \frac{\psi_{\mathrm{raw}}(x)}{\sqrt{\sum_{i} \psi_{\mathrm{raw}}(x_i)^2} \, \Delta x + \epsilon} \, , \end{equation*}
where \(\epsilon = 10^{-8}\) is a small constant added for numerical stability during early training, when the network output may be near zero.
This approach, sometimes called hard normalization or normalization by construction ensures that every wavefunction emitted by the model is a valid quantum state at every training step. Moreover, it eliminates the need for a dedicated normalization loss term and prevents the optimizer from trading physical validity against other objectives such as the Schrödinger residual or data mismatch.

Note 1.6.2.

L2 normalization guarantees a unit norm but does not enforce orthogonality between distinct eigenstates. For multi-state problems where the eigenstates must satisfy \(\langle \psi_m | \psi_n \rangle = \delta_{mn}\text{.}\) This can be achieved via a Gram–Schmidt procedure after the normalization step.
Listing 1.6.3. Hard L2 normalization enforced inside the forward pass.
class NormalizedWavefunctionNet(nn.Module):
    """Wraps an MLP to enforce L2 normalization by construction."""

def __init__(self, base_net: nn.Module, dx: float) -> None:
    super().__init__()
    self.base_net = base_net
    self.dx = dx

def forward(self, x: torch.Tensor) -> torch.Tensor:
    psi_raw = self.base_net(x)
    norm_sq = torch.sum(psi_raw ** 2) * self.dx
    norm = torch.sqrt(norm_sq + 1e-8)
    return psi_raw / norm