Ermakov–Based embeddings as a parameter-free layer for physics-informed neural networks

Introduction

Physics-informed neural networks (PINNs) [1] struggle when a solution develops multiple scales or when coefficients vary rapidly [2, 3]. Supplying the network with structure that the equation already possesses is a natural remedy, and symmetry has been injected into neural PDE solvers as data augmentation [4], as a soft loss penalty [5], and architecturally through equivariant moving frames [6].

This paper takes the architectural route for one specific classical object. Ermakov [7] and Lewis [8] showed that an auxiliary function $\rho$ obeying the Pinney equation [9]

generates an exact invariant for the parametric oscillator. We observe that the same $\rho$ organises three superficially unrelated problems, that the sign of $c$ is what distinguishes them, and that writing the associated change of variables into the architecture costs no parameters. We call the resulting model an Ermakov-based network (EB-Net).

Contributions.

1. A moment identity that makes the embedding practical (lemmas 1 and 4): for the linear conservative families, $\rho$ follows in closed form from the initial data and the coefficients, with no auxiliary network and no knowledge of the solution.

2. One embedding, three families (sections 3–5): the same $(\xi,\theta)$ construction, differing only in the sign of $c$ and in which auxiliary factor — amplitude or phase — is attached.

3. A comparative anatomy (table 6): each family needs a different component of the same construction, and we quantify each component’s share.

4. Ablation as a requirement, not a courtesy (remark 3): family I is degenerate, in that a trivial time reparametrisation reproduces most of the reported gain and the whole of its transfer property. Any symmetry-informed comparison on such a problem is uninterpretable without that control. We regard this as the most portable lesson of the paper.

5. A negative result we consider load-bearing (remark 5): extrapolation beyond the training horizon helps only when the reduced solution converges to a fixed point. When it beats instead, no model extrapolates, ours included.

What we do not claim. The scaling-variable reduction is classical (remark 2); so are the diffusive-time substitution and the Lewis–Riesenfeld transformation. We claim none of them, nor a new architecture class, nor superiority over neural operators, nor applicability beyond linear first-kind self-similar problems.

The Ermakov–Pinney backbone

The invariant, for either sign of $c$

Given $\ddot{y}+\omega^{2}(t)y=0$ and $\rho$ solving (1), the Ermakov–Lewis quantity

satisfies $\dot{I}=0$ identically. We stress for either sign of $c$: the familiar $c>0$ branch is oscillatory and underlies the Lewis–Riesenfeld theory of the time-dependent quantum oscillator [10]; the $c<0$ branch describes monotone expansion and is the one relevant to diffusion. Both are used below. Pinney’s formula gives the general solution of (1) as $\rho^{2}=Ay_{1}^{2}+2By_{1}y_{2}+Cy_{2}^{2}$ with $c=AC-B^{2}$ when the Wronskian of $y_{1},y_{2}$ is unity.

The generic embedding

All three families use the same three ingredients:

where $\kappa$ is the diffusivity ($\kappa= D$) or $\hbar/m$ in the quantum case, $\alpha$ is the scaling dimension fixed by the conserved quantity, and $\Lambda$ is an auxiliary factor: $\Lambda\equiv1$ in the diffusive families and $\Lambda= e^{iS}$ with $S = m\dot{\rho}x^2/(2\hbar\rho)$ in the quantum one. Which of these ingredients carries the method turns out to depend on the family, which is the comparative result of this paper.

Three clocks

It pays to be explicit about what each candidate input measures time with.

clock	dimensions	what the solution does in it
$t$	time	evolves at a rate that varies with the coefficients
$\tau= \int_0^t D$	length$^2$	still spreads, but at a uniform rate: $m_2$ is linear in $\tau$
$\theta= \int_0^t \kappa/\rho^2$	dimensionless	does not spread: $\Phi(\cdot,\theta)$ is stationary or slowly varying

The middle row invites a misreading that we have found easy to make. The diffusive time $\tau$ is not “time in units of the spreading”; it is time in units of accumulated diffusivity, with dimensions of length squared. Only $\theta$ is the solution’s own clock: $\mathrm{d}\theta$ measures elapsed time against the time the solution currently needs to change by its own width, so one unit of $\theta$ is one e-folding of $\rho$ wherever one stands on the trajectory. That distinction is the content of the method, and it explains the extrapolation results below: extrapolating in $\tau$ asks the network about scales it never saw, whereas extrapolating in $\theta$ asks about a profile it has already watched converge.

Family I: linear diffusion ($c=-1$, $\omega\equiv0$)

Proof. $\frac{\mathrm{d}}{\mathrm{d}t}\int x^2u=D\int x^2u_{xx}=D[x^2u_x-2xu]_{-\infty}^{\infty}+2D\int u=2D$.

Setting $\rho=\sqrt{m_2(0)+2\tau}$, $\xi=x/\rho$, $\theta=\log(\rho/\rho_0)$ and $u=\rho^{-1}v$ gives, by a short computation (appendix A),

so $v$ obeys the autonomous Ornstein–Uhlenbeck equation whose unique decaying unit-mass steady state is $v_\infty=(2\pi)^{-1/2}e^{-\xi^2/2}$. In diffusive time $\rho$ solves $\rho''=-\rho^{-3}$, i.e. (1) with $\omega\equiv0$, $c=-1$.

Remark 2 (The reduction is classical). For constant $D$ this is the scaling-variable change of the dynamical-systems approach to self-similar asymptotics. Gallay and Wayne [11] use $\xi=x/\sqrt{1+t}$, $\tau=\log(1+t)$ and $\omega=(1+t)^{-1}w$, obtaining $\partial_\tau w=\Delta_\xi w+\frac{1}{2}(\xi\cdot\nabla_\xi)w+w$ with a Gaussian kernel — the same operator as (5) up to the factor $\frac{1}{2}$, a convention in the normalisation of $\rho$. The similarity ansatz is older still [12], and $\tau=\int D$ is textbook [13]. Lemma 1 is what extends it to $D=D(t)$ and to unnormalised initial data, and is therefore the part that makes an architectural layer possible.

Remark 3 (This family is degenerate, and that is the lesson). The map $t\mapsto\tau$ alone already removes $D(t)$ from (4), and the exact solution depends on $t$ only through $\tau$. A model fed $(x,\tau)$ therefore inherits both a large share of the accuracy gain and the entire zero-shot transfer property, with no Ermakov structure whatsoever. Table 1 shows PINN-$\tau$ recovering a factor 2.2 of the 4.0 total. We conclude that family I cannot support a claim about symmetry, and we present it not as evidence but as the reason the rest of the paper is organised around controls.

Model	Input	Gaussian	bimodal
PINN	$(x,t)$	$3.84 \times10^{-3} \ \pm4.7 \times10^{-4}$	$9.63 \times10^{-3} \ \pm2.6 \times10^{-3}$
PINN-$\tau$	$(x,\tau)$	$1.76 \times10^{-3} \ \pm2.5 \times10^{-4}$	$8.16 \times10^{-3} \ \pm2.7 \times10^{-3}$
EB-coords	$(\xi,\theta)$	$9.56 \times10^{-4} \ \pm1.4 \times10^{-4}$	$6.33 \times10^{-3} \ \pm3.3 \times10^{-3}$
EB-Net	$(\xi,\theta),\ \rho^{-\alpha}v$	$1.16 \times10^{-3} \ \pm4.7 \times10^{-4}$	$5.53 \times10^{-3} \ \pm2.0 \times10^{-3}$

Transfer across coefficient schedules certifies the trivial step

Because (5) contains no trace of $D$, the reduced profile is universal across the family $\{D(\cdot)\}$: changing the coefficient only changes the path $t \mapsto(\rho(t),\theta(t))$ traced in the $(\xi,\theta)$ plane. One is tempted to present the resulting zero-shot transfer as the decisive argument for the embedding. Table 2 shows why that would be wrong.

Evaluation $D(t)$	PINN	PINN-$\tau$	EB-Net
$D_1$ (training)	$3.32 \times10^{-3}$	$1.83 \times10^{-3}$	$1.05 \times10^{-3}$
$D_2$ (phase $+\pi/2$)	$6.12 \times10^{-2}$	$1.83 \times10^{-3}$	$8.84 \times10^{-4}$
$D_3$ ($f=3$, $a=0.9$)	$5.35 \times10^{-2}$	$1.85 \times10^{-3}$	$1.01 \times10^{-3}$
$D_4$ ($D_0=1.6$, fast) †	$1.26 \times10^{-1}$	$1.14 \times10^{-2}$	$1.81 \times10^{-3}$
$D_5$ ($D_0=0.4$, slow)	$2.26 \times10^{-1}$	$1.62 \times10^{-3}$	$6.56 \times10^{-4}$
† only $D_4$ leaves the trained diffusive-time interval.

On $D_2$, $D_3$ and $D_5$ — which stay inside the trained diffusive-time interval — PINN-$\tau$ transfers as well as EB-Net, for the reason given in remark 3. The plain PINN has no mechanism at all and degrades by one to two orders of magnitude. The informative column is $D_4$, which overshoots the interval, together with the temporal extrapolation to $t \in[T,2T]$: there EB-Net reaches $3.91 \times10^{-3}$ against $4.61 \times10^{-2}$ for PINN-$\tau$ and $1.71 \times10^{-1}$ for the baseline. Two effects compound. The logarithmic time contracts overshoot — a 100% excess in $\tau$ is only 28.9% in $\theta$ — and the function being extrapolated has already collapsed onto a constant profile.

The physical residual carries an implicit $\rho^{-3}$ weighting

Identity (5) has a consequence that any follow-up work in this direction will have to account for. Minimising the physical residual, as a standard PINN does, is equivalent to minimising the reduced residual weighted by $(D/\rho^3)^2$. Over our test horizon $\rho$ grows by a factor 3, so early times are weighted up to 27 times more heavily than late ones. This bias interacts with the amplitude ansatz: the model carrying the $\rho^{-\alpha}$ prefactor inherits a steeper weighting than the one that does not, which is why the two are indistinguishable in Table 1 even though only the former makes the trace exactly stationary.

Reweighting makes every model roughly an order of magnitude worse in physical $L^2$ ($3.84 \times10^{-3} \rightarrow3.64 \times10^{-2}$ for the baseline): the implicit bias is not a defect for this metric, since early times are where the profile is narrow and hardest to resolve. The result we care about is that the ordering survives ($3.64 \times10^{-2} > 1.99 \times10^{-2} > 1.05 \times10^{-2} \approx1.04 \times10^{-2}$). Since the ablation ladder is invariant under this rather drastic change of loss weighting, the gain is attributable to the coordinates rather than to a fortunate weighting.

Model	physical residual	reduced residual
PINN	$3.84 \times10^{-3}$	$3.64 \times10^{-2}$
PINN-$\tau$	$1.76 \times10^{-3}$	$1.99 \times10^{-2}$
EB-coords	$9.56 \times10^{-4}$	$1.05 \times10^{-2}$
EB-Net	$1.16 \times10^{-3}$	$1.04 \times10^{-2}$

Model	Relative $L^2$ error Gaussian	Relative $L^2$ error bimodal	Extrapolation $t \in[T,2T]$ Gaussian	Extrapolation $t \in[T,2T]$ bimodal
PINN	$1.19 \times10^{-2}\ \pm1.2 \times10^{-3}$	$1.74 \times10^{-2}\ \pm1.4 \times10^{-3}$	$1.21 \times10^{0}$	$1.14 \times10^{0}$
PINN-$\tau$	$8.00 \times10^{-2}\ \pm4.5 \times10^{-2}$	$1.13 \times10^{-1}\ \pm3.6 \times10^{-2}$	$4.16 \times10^{-1}$	$2.89 \times10^{-1}$
EB-coords	$4.13 \times10^{-2}\ \pm4.9 \times10^{-2}$	$9.80 \times10^{-2}\ \pm2.0 \times10^{-2}$	$6.31 \times10^{-1}$	$2.86 \times10^{0}$
EB-Net	$1.01 \times10^{-3}\ \pm9.7 \times10^{-5}$	$6.32 \times10^{-3}\ \pm6.3 \times10^{-4}$	$6.81 \times10^{-3}$	$7.09 \times10^{-2}$

Family II: diffusion in a breathing trap ($c=-1$, $\omega\ne0$)

The drift contains $x$ explicitly, so no reparametrisation $s=s(t)$ can remove it: the degeneracy of family I is gone. Physically this is diffusion inside a trap whose stiffness is modulated — optical tweezers at varying power, a colloid in an oscillating field.

Lemma 4 (Moment identity, generalised). $\mathrm{d}m_2/\mathrm{d}t = 2D(t) - 2\gamma(t)m_2$, a linear ODE; hence $\rho=\sqrt{m_2}$ is again available in closed form from $m_2(0)$ alone.

Writing $u=\rho^{-1}v(\xi,\theta)$ with $\xi=x/\rho$ and $\theta'=D/\rho^2$ reproduces exactly the operator of (5) provided $\rho\dot{\rho}=D-\gamma\rho^2$; differentiating that relation gives

which for constant $D$ is (1) with $\omega^2=\dot{\gamma}-\gamma^2$ and $c=-D^2$. Two consequences matter. First, $\omega\ne0$: the Ermakov structure is no longer in its degenerate branch. Second, $\rho$ oscillates rather than growing as a power of $t$, so the classical scaling variables $\xi=x/\sqrt{t}$, $\theta=\log t$ of remark 2 do not describe this case at all.

The trivial control does not merely fail; it hurts. PINN-$\tau$ reaches $8.00 \times10^{-2}$ against the baseline’s $1.19 \times10^{-2}$, with a seed-to-seed spread an order of magnitude larger than any other model’s. Feeding $\tau$ supplies a clock that absorbs $D$ but is inconsistent with the drift, and the network is worse off than with physical time. We regard this as the cleanest available refutation of the objection that family I invites.

The amplitude prefactor switches from neutral to essential. It was indistinguishable from omission in family I; here removing it costs $40.8\times$ (Gaussian) and $15.5\times$ (bimodal), and in the bimodal case EB-coords is worse than the plain PINN. The trap compresses and expands the density, so the amplitude genuinely oscillates and the ansatz must absorb it.

Extrapolation. With weights trained only on $[0,T]$, EB-Net reaches $6.81 \times10^{-3}$ on $t \in[T,2T]$ against the baseline’s $1.21 \times10^{0}$, a factor 178. Figure 1(c) shows EB-Net crossing $t=T$ without visible incident.

Family III: the quantum oscillator ($c=+1$)

With $\hbar=m=1$,

The Lewis–Riesenfeld transformation [10]

with $\rho$ solving (1) at $c=+1$, yields the autonomous oscillator $i\phi_\theta=-\frac{1}{2}\phi_{\xi\xi}+\frac{1}{2}\xi^2\phi$. We verified this symbolically; the residual difference is exactly zero (appendix B). Here $\alpha=1/2$ because $|\psi|^2$ is conserved, and $\Lambda=e^{iS}$ is not optional — omitting it breaks the reduction, a claim we measure rather than assert.

The honest control is adiabatic, not trivial. No time reparametrisation absorbs $\omega(t)$. The baseline a physicist would reach for is the adiabatic one: rescale by the instantaneous oscillator length $1/\sqrt{\omega}$ and use the adiabatic phase $\int\omega\,\mathrm{d}t$. It is exact in the slow limit. We therefore choose a regime where it has no right to work: $\omega\in[0.3,1.7]$, $\rho$ breathing by a factor 2.4, and $|\dot{\omega}|/\omega^2=2.2$.

The invariant does physical work. Evaluating (2) as a quantum expectation on the exact solution gives $\langle I\rangle=0.500000000000$ at every time to $2.2\times10^{-16}$, while $\langle H\rangle$ runs from 0.500 to

caption: Family III. Mean $\pm$ s.d. over three seeds. The adiabatic control is worse than the plain PINN in both cases. Omitting the quadratic phase costs $5.7\times$ on the ground state.

:::

0.960. The energy is not conserved; the Ermakov–Lewis invariant is, and its value is exactly the ground-state energy of the reduced oscillator.

Remark 5 (When the method does not help). For the $(n=0)+(n=2)$ superposition the reduced solution is not stationary: it beats at $E_2-E_0=2$ in $\theta$. There is then no fixed point to extrapolate toward, and Figure 2(d) shows every model, ours included, degrading to $\sim0.8$ past $t=T$. The extrapolation advantage of the embedding is therefore conditional on convergence of the reduced solution, not a general property. We prefer this sharp, falsifiable statement to a blanket claim.

Synthesis: which component carries the method

Read across the rows, table 6 says three things. The control degrades from sufficient to actively harmful as one moves away from the degenerate branch. The gain over the plain baseline grows in the same direction. And the component that matters changes: bare coordinates, then amplitude, then phase. A reader who saw only family I would conclude the prefactor is decorative; one who saw only family III would conclude the same of the amplitude. Both conclusions would be artefacts of the chosen test problem, which is our main methodological point.

Family	Pinney branch	Trivial control	Load-bearing piece	vs PINN	piece’s share
I. diffusion $u_t = Du_{xx}$	$c=-1$, $\omega=0$	$\tau$ (suffices)	coordinates only	4.0	prefactor: 1.22×
II. trap $+\gamma(xu)_x$	$c=-1$, $\omega\ne0$	$\tau$ (fails)	+ prefactor $\rho^{-\alpha}$	11.7	prefactor: 41×
III. quantum $\omega(t)$	$c=+1$	adiabatic (fails)	+ phase $e^{iS}$	13.7	phase: 5.7×

Architecture and protocol

EB-Net has three blocks: a parameter-free embedding $(x,t)\mapsto(\xi,\theta)$; an MLP backbone; and a reconstruction $\rho^{-\alpha}\Lambda\Phi$. Because the embedding carries no weights, every model in every ladder has an identical parameter count (12 737 for the real families) and the same per-iteration cost to within 4%.

Fairness. Within each family all models (i) minimise the same loss functional, evaluated by automatic differentiation in physical coordinates — no model is trained on the reduced equation; (ii) share architecture (4 $\times$ 64, tanh, Xavier initialisation), collocation budget and optimiser schedule (Adam [14] with cosine decay, then L-BFGS); (iii) run over the same seeds; and (iv) normalise inputs to $O(1)$ in their own natural parametrisation. The initial-condition weight was selected per model by a sweep on a held-out seed, which selected the same value ($\lambda_{\mathrm{ic}}=10$) for all of them.

References and precision. The scale functions are obtained by integrating the relevant ODE system to rtol = $10^{-12}$ and exposed to automatic differentiation through a $C^1$ cubic Hermite interpolant built from exact nodal derivatives. Exact solutions satisfy their PDEs to $10^{-10}$ or better on the evaluation grid, so the reported errors are not reference-limited.

Related work

Scaling variables. The reduction of section 3 is classical [11, 12], as is the diffusive-time substitution [13] and the Lewis–Riesenfeld transformation [10]. Our role is to move them from the analysis into the architecture, which lemmas 1–4 make possible for time-dependent coefficients.

Symmetry in neural solvers. Lie point symmetries have been used for augmentation [4], as a loss penalty [5], and structurally for ODEs [6]; conservative projections appear in [15]. Transformations of PINN inputs to smooth the target [16] and architectures for variable-coefficient PDEs [17] are the nearest PINN-side neighbours.

Dynamic rescaling. Self-similar blow-up profiles have been found by fixing a similarity ansatz [18]. The closest work is Kavousanakis et al. [19], who learn the rescaling with an auxiliary network under algebraic constraints, producing an index-2 differential-algebraic system. That method applies when the scaling is unknown, including second-kind similarity; ours applies when it is known, which by our lemmas is the whole linear conservative class, and costs no extra parameters.

Limitations

• Linearity and conservation. Both lemmas use them; for nonlinear problems (porous medium, Fisher–KPP) the moment hierarchy does not close and $\rho$ is no longer explicit.

• First-kind similarity only. Second-kind exponents, fixed by a nonlinear eigenvalue problem [12], are outside the construction.

• Conditional extrapolation. Remark 5.

• Operator-learning baselines are a different comparison. A neural operator [20, 21] trained across the coefficient family would be the natural rival for the transfer claim of § 3.1, and we have not run it. We note that the comparison is not like-for-like: an operator network amortises training over a distribution of coefficients, whereas every model compared here solves a single instance from its residual alone, with no data. The operator setting is the subject of the companion $N$-dimensional work (§ 10), where it is the main object rather than a baseline.

• Absolute accuracy in family III. All models remain near $10^{-1}$ on the superposition at our training budget; the comparison is informative but the absolute numbers are not competitive with a spectral solver.

• Scope. One spatial dimension throughout.

Outlook: the $N$-dimensional case

The one-dimensional construction generalises, but not by replacing scalars with matrices. For a general symmetric frequency matrix $\Omega^2(t)$ whose eigenvectors rotate, the correct frame follows from the linear equation $\ddot{Q}+\Omega^2Q=0$ rather than from a symmetric matrix Pinney equation, which fails to preserve symmetry as soon as $[\Omega^2(t),\Omega^2(t')] \ne0$; the width and phase matrices then obey a matrix Ermakov–Pinney system closed by a Lyapunov equation. The scalar Ermakov clock $\theta$ becomes a $U(N)$-valued clock carrying a rotational term, so no single latent time exists unless the frame is isotropic, and the $N$ Williamson invariants generalise (2). The frame remains observable from phase-space moments, which additionally permits reconstruction of an unknown $\Omega^2(t)$. These results, their operator-learning consequences and their verification against a spectral solution of the two-dimensional Schrödinger equation with a chirped, cross-coupled trap are developed in a companion paper in preparation. The present work establishes the one-dimensional case that the companion assumes.

Conclusion

A single Pinney function organises linear diffusion, diffusion in a breathing trap, and the quantum oscillator with time-dependent frequency; the sign of $c$ is what separates them. Writing the associated change of variables into a neural solver’s input layer costs no parameters and is made practical for time-dependent coefficients by an exact moment identity. What the construction buys depends on the family in a way we have quantified: coordinates alone where a time reparametrisation would have sufficed, the amplitude prefactor where a trap makes the density oscillate, the quadratic phase where the dynamics are quantum. The honest control degrades from sufficient to harmful along the same axis. We regard the resulting table, rather than any single error figure, as the contribution.

Acknowledgements

The authors thank the transformative agreements of Universidad Autónoma Metropolitana.

Data availability statement

No new data were created in this study beyond those generated by the accompanying code. All results, figures and tables are reproduced from scratch by the software archived at [DOI on acceptance]; no numerical value in this paper is transcribed by hand.

Declaration of generative AI use

During the preparation of this work the authors used a large language model to assist with code implementation, numerical experimentation and manuscript drafting. The authors reviewed and edited all content and take full responsibility for the published work.

Conflicts of interest

The authors declare no conflicts of interest.

Derivation for family I

In diffusive time $s=\tau$, (4) reads $u_s=u_{xx}$ and $\rho\rho'=1$ since $\rho^2=m_2(0)+2s$. With $u=\rho^{-1}v(\xi,\theta)$, $\xi=x/\rho$, $\theta'=\rho^{-2}$:

Multiplying $u_s-u_{xx}$ by $\rho^3$ and using $\rho\rho'=\rho^2\theta'=1$ gives $v_\theta-v_{\xi\xi}-\xi v_\xi-v$; multiplying by $D/\rho^3$ and using $u_t=Du_s$ gives (5). Integrating $\theta'=\rho^{-2}$ yields $\theta=\log(\rho/\rho_0)$. The exponent $\alpha=1$ is fixed by mass conservation: $\int u\,\mathrm{d}x=\rho^{1-\alpha}\int v\,\mathrm{d}\xi$ is $t$-independent only for $\alpha=1$; in $d$ dimensions $\alpha=d$, and in the quantum case $\alpha=1/2$ because $\int|\psi|^2$ is the conserved quantity.

Code

Implementation: ebnet.py (core and family I), trap.py and run_trap.py (family II), quantum.py and run_quantum.py (family III), plus verify_*.py for the symbolic and numerical checks and make_tables_final.py, which regenerates every number in this paper from the stored results.

Listing 1: Family II: the scale function comes from a linear moment ODE, so no auxiliary network is needed.

def _build(self, n=6001):
    def rhs(t, y):
        m2, _ = y
        return [2*self._Dnp(t) - 2*self._gnp(t)*m2, # {xref}`Lemma 2`
                self._Dnp(t) / m2] # dtheta/dt
    ts = np.linspace(0.0, self.t_end, n)
    sol = solve_ivp(rhs, (0.0, self.t_end), [self.m2_0, 0.0], t_eval=ts,
                    rtol=1e-12, atol=1e-14, method="DOP853")
    m2, th = sol.y
    rho = np.sqrt(m2)
    drho = (self._Dnp(ts) - self._gnp(ts)*m2) / rho # exact nodal derivative
    self._rho = Hermite(ts, rho, drho) # C^1, differentiable
    self._theta = Hermite(ts, th, self._Dnp(ts)/m2)

Listing 2: Family III: the full Lewis–Riesenfeld reconstruction. The ablated model drops cS/sS, which is the phase-factor comparison of table 5.

class QEBNet(_QSolver):
    def forward(self, x, t):
        rho, drho = self.p.rho(t), self.p.drho(t)
        f = self.mlp(torch.cat([x/rho, self.p.theta(t)/self._thmax], 1))

fr, fi = f[:, :1], f[:, 1:]
S = drho * x**2 / (2*rho) # quadratic phase
cS, sS = torch.cos(S), torch.sin(S)
pre = rho**(-0.5)
return torch.cat([pre*(fr*cS - fi*sS),
                  pre*(fr*sS + fi*cS)], dim=1)

Listing 3: Symbolic verification of the Lewis–Riesenfeld transformation: the printed difference is exactly zero.

psi = rho**sp.Rational(-1,2) * sp.exp(I*rp*x**2/(2*rho)) * phi(x/rho, th)
expr = I*sp.diff(psi,t) + sp.Rational(1,2)*sp.diff(psi,x,2) \
     - sp.Rational(1,2)*w**2*x**2*psi
expr = expr.subs(sp.Derivative(rho,(t,2)), 1/rho**3 - w**2*rho).doit()
expr = expr.subs(sp.Derivative(th,t), 1/rho**2)
# -> i phi_theta + 1/2 phi_xixi - 1/2 xi**2 phi (difference = 0)