Instead of fully characterizing the solution of the Cauchy problem of Equation 1, the method of moments (Pérez-García et al., 2007), seeks to provide qualitative description of the solution behavior by studying the evolution of several integral quantities, i.e., the moments, of the solution u ( x , t ) . This approach enables a reduced-order description of the NLS equation by transforming it into a system of (potentially) closed ordinary differential equations (ODEs). More specifically, according to Pérez-García et al. (2007), we define for k = 0,1,2 , ⋯   the moment quantities of solution u ( x , t ) as follows, I k t = ∫ R x k | u x , t | 2 d x , (2) V k t = 2 k − 1 i ∫ R x k u x , t ∂ u ̄ x , t ∂ x − u ̄ x , t ∂ u x , t ∂ x d x , (3) K t = 1 2 ∫ R ∂ u x , t ∂ x 2 d x , (4) J t = ∫ R G ρ x , t , t d x = ∫ R G u x , t 2 , t d x , (5) where u ̄ is the complex conjugate of u , ρ ( x , t ) = u ( x , t ) 2 , and G = G ( ρ , t ) is a function such that ∂ G ∂ ρ ( ρ , t ) = g ( ρ , t ) . Moments Equations 2–5 of the solution u ( x , t ) have intuitive physical meanings; for example, the first moment I 1 ( t ) is associated with the center of mass as described by the (unnormalized) probability density ρ = | u | 2 . Higher moments I k are also associated with this density distribution (i.e., its variance etc.). The V k quantities are the respective ones associated with the momentum density (which is the quantity in the corresponding parenthesis in the right hand side of the V k definition. K stems from (thinking quantum-mechanically) the kinetic part of the Schrödinger problem energy, while J represents the nonlinear part of the corresponding energy. We assume that the IC u 0 ( x ) is regular enough to ensure that all moments are well-defined for t ≥ 0 .

The method of moments aims to extract qualitative information about the solution u ( x , t ) of the PDE Equation 1 by deriving a closed set of evolution ODEs for the moments of u ( x , t ) . Depending on the nonlinearity g ( ρ , t ) , these ODEs can sometimes be determined analytically as is shown in the work of Pérez-García et al. (2007); García-Ripoll and Pérez-García (1999). Below, we provide a few examples.

Example 1

The moments I 1 and V 0 satisfy d I 1 d t = V 0 , d V 0 d t = − I 1 . (6)

This indicates that the evolution of the center of mass I 1 behaves as a harmonic oscillator, independent of the nonlinearity g ( ρ , t ) . More generally, for a parabolic confinement of frequency Ω , this would be reflected in the associated frequency of moment oscillations; this is the so-called dipolar motion (Pitaevskii and Stringari, 2003; Kevrekidis et al., 2015).

Example 2

If g ( ρ , t ) ≡ 0 , i.e., for the linear case of Equation 1, the set of moments I 2 , V 1 , K are closed under d I 2 d t = V 1 , d V 1 d t = 4 K − 2 I 2 , d K d t = − 1 2 V 1 .

Example 3

Assume the nonlinearity g ( ρ , t ) = g ( ρ ) is time-independent and given by g ( ρ ) = g 0 ρ 2 , where g 0 ∈ R is a constant. Although the evolution of the moments I 2 , V 1 , K and J is not closed, it becomes closed under the coordinate transformation E = K + J , i.e., d I 2 d t = V 1 , d V 1 d t = 4 E − 2 I 2 , d E d t = − 1 2 V 1 , (7)

While the examples and conditions under which the method of moments leads to closed equations are well-known, deriving such analytical closure systems requires knowledge of the underlying PDE system, Equation 1 and detailed calculations therewith. This work explores data-driven methods for obtaining analytical or approximate moment closure systems based on empirical observations or simulations of the NLS equation, rather than relying solely on such analytical understanding and derivations. Importantly, the reconstruction of these ODE models can, in principle, take place even for settings where the underlying PDE model is unavailable/has not been specified. Given “experimental” data for the field, one may aspire to utilize the toolboxes presented below in order to obtain these effective, reduced dynamical equations.

For systems with existing analytical closures of the moment equations, such as Examples 1–3, our method seeks to rediscover the moment evolution equations and potentially the necessary coordinate transformations, such as E = K + J in Example 3, in a model-agnostic and data-driven manner. For systems lacking analytical closed moment equations, we seek to derive approximate moment closure equations, providing a principled and reduced-order description of the original PDE, capable of predicting the future evolution of the system. The relevant details will be explained in the following section.

Our first method leverages Sparse Identification of Nonlinear Dynamics (SINDy) by Brunton et al. (2016), a data-driven approach for discovering governing ODEs from simulated or observational data. Consider a nonlinear ODE system of the form: d x d t t = f x t , (8) where x = ( x 1 , … , x n ) ⊤ : [ 0 , ∞ ) → R n represents the state evolution over time, governed by the dynamic constraint f : R n → R n . SINDy aims to identify the unknown dynamics, f ( x ) from a time series of x .

The key assumption is that f ( x ) has a “simple” form and can be expressed or approximated as a linear combination of only a few terms from a suitably chosen library, Θ ( x ) = [ θ 1 ( x ) , … , θ p ( x ) ] . For example, a monomial library of degree up to two, Θ deg ≤ 2 ( x ) , is: Θ deg ≤ 2 x = Θ deg = 1 x , Θ deg = 2 x = x 1 , x 2 , … , x n ⏟ Θ deg = 1 x , x 1 2 , x 1 x 2 , … , x 1 x n , x 2 2 , … , x n 2 ⏟ Θ deg = 2 x = θ 1 x , … , θ p x , (9) with p = n + n 2 . In particular, e.g., the right-hand sides of Equations 6, 7 can all be written as sparse linear combinations of terms from Θ deg ≤ 2 ( x ) . Given a dictionary Θ ( x ) = [ θ 1 ( x ) , … , θ p ( x ) ] of p elements— p is typically larger than n —the sparsity assumption implies the existence of a sparse matrix Ξ = ( ξ 1 , … , ξ n ) ∈ R p × n such that f x ⊤ = f 1 x , ⋯   , f n x ≈ Θ x ⋅ Ξ = θ 1 x , … , θ p x ⋅ | | ⋯ | ξ 1 ξ 2 ⋯ ξ n | | ⋯ | (10) where each sparse column ξ j ∈ R p indicates which nonlinear functions among the library Θ ( x ) = [ θ 1 ( x ) , … , θ p ( x ) ] are used to parsimoniously represent f j ( x ) .

To determine this sparse Ξ , SINDy employs sparse regressions on the data. Specifically, given a time series { x ( t 1 ) , … , x ( t N ) } ⊂ R n of the state x ( t ) at times t 1 , … , t N — N is generally much larger than p , the library size—one can assemble the state matrix X ∈ R N × n and the derivative matrix X ̇ ∈ R N × n : X = | | | X 1 X 2 ⋯ X n | | | ≔ x 1 t 1 x 2 t 1 ⋯ x n t 1 x 1 t 2 x 2 t 2 ⋯ x n t 2 ⋮ ⋮ ⋱ ⋮ x 1 t N x 2 t N ⋯ x n t N ∈ R N × n , d d t X = X ̇ = | | | X ̇ 1 X ̇ 2 ⋯ X ̇ n | | | ≔ x ̇ 1 t 1 x ̇ 2 t 1 ⋯ x ̇ n t 1 x ̇ 1 t 2 x ̇ 2 t 2 ⋯ x ̇ n t 2 ⋮ ⋮ ⋱ ⋮ x ̇ 1 t N x ̇ 2 t N ⋯ x ̇ n t N ∈ R N × n , where X ̇ can be estimated by, e.g., finite differences on X . Define the library matrix Θ ( X ) ∈ R N × p as Θ X = | | | θ 1 X θ 2 X ⋯ θ p X | | | ≔ θ 1 x t 1 θ 2 x t 1 ⋯ θ p x t 1 θ 1 x t 2 θ 2 x t 2 ⋯ θ p x t 2 ⋮ ⋮ ⋱ ⋮ θ 1 x t N θ 2 x t N ⋯ θ p x t N .

Evaluating Equations 8, 9 at all times t = t 1 , … , t N yields | | | X ̇ 1 X ̇ 2 ⋯ X ̇ n | | | = X ̇ ≈ Θ X Ξ = | | θ 1 X ⋯ θ p X | | | | ⋯ | ξ 1 ξ 2 ⋯ ξ n | | ⋯ | (11)

Sparse regression techniques, such as LASSO (Tibshirani, 1996) or sequential thresholded least-squares (Brunton et al., 2016), can then solve the overdetermined system (Equation 11) for the sparse Ξ . This provides an approximation to the governing equation as in Equation 10.

While SINDy can be directly applied to Examples 1–3 to discover closed moment systems based on simulated PDE data (Equation 1), our specific interest lies in the following:

1. Given time series data of moments with a known closed moment system (e.g., Examples (1) and (2)), can we robustly discover the governing dynamics using a correct but potentially oversized dictionary Θ ( x ) , such as polynomials up to a high degree?

These questions will be addressed in Section 4. Of particular interest is the second point, where we demonstrate that, in certain cases, we can gain insight into the appropriate transformation to close the system, even if the moment system is not closed under the originally selected variables. In the following section, we discuss a more principled strategy to discover such transformations in a data-driven fashion.

To illustrate the idea, we focus on Example 3, where the initially selected moments are x = [ I 2 , V 1 , K , J ] ⊤ , and an analytical closure exists only after a coordinate transformation, y = A ⊤ x , y = I 2 V 1 K + J = 1 0 0 0 0 1 0 0 0 0 1 1 I 2 V 1 K J = A ⊤ x ,

Remark 1

Note that the transformation matrix A ∈ R 4 × 3 is not unique. For any full-rank matrix P ∈ R 3 × 3 , the moment system remains closed under the transformation y ̃ = A ̃ ⊤ x , with A ̃ = A P .

Assume we aim to discover A purely from simulated PDE data. We propose the following strategy: Let X ∈ R N × 4 be the state matrix for the original coordinate x = [ I 2 , V 1 , K , J ] ⊤ , sampled at t 1 , … , t N , as defined by Equation 8. Let y = A ⊤ x be the new coordinate, and the associated new state matrix Y ∈ R N × 3 becomes Y = —  y t 1 ⊤  — —  y t 2 ⊤  — ⋮ —  y t N ⊤  — = —  x t 1 ⊤  — —  x t 2 ⊤  — ⋮ —  x t N ⊤  — A = X A

We can then solve for A through the following optimization. min A ∈ R 4 × 3 min Ξ ∈ R p × 3 d d t X A − Θ X A ⋅ Ξ F 2 + μ ‖ Ξ ‖ 1 , (12)  s.t.  A ⊤ A = I 3 × 3 where d d t X A is the derivative matrix, Equation 8, associated with the new state matrix Y = X A , ‖ ⋅ ‖ F and ‖ ⋅ ‖ 1 are the Frobenius norm and l 1 norm, respectively, defined for any m × n matrix B = ( b i j ) as follows: B F = ∑ i = 1 m ∑ j = 1 n b i j 2 , B 1 = max 1 ≤ j ≤ n ∑ i = 1 m b i j .

In addition, μ ≥ 0 is a non-negative weight, and Θ ( X A ) = Θ ( Y ) ∈ R N × p is the library matrix, Equation 8, of a chosen library Θ ( y ) on the new coordinate y = A ⊤ x .

The idea behind Equation 12 is very simple: we search for the transformation matrix A ∈ R 4 × 3 such that the dynamics under the new coordinate y = A ⊤ x can be parsimoniously represented from the library Θ ( y ) . The weight μ ≥ 0 controls the sparsity-promoting L 1 -regularization, and the constraint A ⊤ A = I 3 × 3 prevents the trivial solution A = 0 . The set of A satisfying the constraint A ⊤ A = I 3 × 3 is called the Stiefel manifold (Edelman et al., 1998; Absil et al., 2008).

Remark 2

When the weight μ = 0 , even with the Stiefel manifold constraint A ⊤ A = I 3 × 3 , the solution to Equation 12 is not unique. Just like Remark 1, if ( A * , Ξ * ) is a solution, then for any orthogonal matrix O ∈ O ( 3 ) , the pair ( A ̃ * , Ξ ̃ * ) is also a solution, where A ̃ * = A * O , Ξ ̃ * = O ⊤ Ξ * O

To solve Equation 12, we use alternating optimization, iteratively optimizing A and Ξ while keeping the other fixed. Specifically, when A is fixed, solving for Ξ reduces to a LASSO problem. Conversely, when Ξ is fixed, the problem becomes a Stiefel manifold optimization with a smooth objective function, which can be efficiently solved using methods such as those presented in (Oviedo and Dalmau, 2019; Xiao et al., 2020; Liu et al., 2021). To ensure the algorithm remains unbiased, we implement an annealing strategy for the hyperparameter μ . Initially, μ is kept constant for Iter_scheduled iterations. Following this period, μ is reduced by half every α iterations. Refer to Algorithm 1 for the pseudocode to solve problem Equation 12.

Algorithm 1

Data-driven discovery of coordinate transformation.

We first examine Examples (1) and (2), where analytical (linear) closure exists for the chosen moments x = [ I 1 , V 0 ] (Example 1) and x = [ I 2 , V 1 , K ] (Example 2).

Let the selected moments be x = [ I 1 , V 0 ] . We construct the data matrices X ( 0 ) = [ I 1 ( 0 ) , V 0 ( 0 ) ] ∈ R N × 2 and X ( 1 ) = [ I 1 ( 1 ) , V 0 ( 1 ) ] ∈ R N × 2 from numerically solving the PDE Equation 1 with the ICs: u 1 x , 0 = π − 1 / 4 ⁡ exp − 1 2 x − 5 2 , u 2 x , 0 = 1 2 s e c h 2 x − 5 .

We apply SINDy to this system using monomial libraries of degree up to n ∈ N , Θ deg ≤ n ( x ) , as defined in Equation 9. We find that as long as n ≤ 2 , SINDy applied to either data matrix X ( 0 ) or X ( 1 ) discovers the dynamics nearly perfectly. However, as we gradually expand the library Θ deg ≤ n ( x ) by increasing n , SINDy applied to either X ( 0 ) or X ( 1 ) individually typically produces erroneous dynamics, regardless of how carefully the sparsity-promoting parameter is tuned. A typical negative result from SINDy applied to X ( 0 ) with n = 3 is presented in the Supplementary Material. This outcome is expected, as increasing the library size makes the problem more ill-posed, leading SINDy to overfit the data and produce incorrect dynamics.

To address the overfitting issue, we can enlarge the dataset by concatenating the data matrices X ( 0 ) and X ( 1 ) vertically, forming X = [ X ( 0 ) ⊤ , X ( 1 ) ⊤ ] ⊤ ∈ R 2 N × 2 . When applying SINDy to this new data matrix X (considering boundary issues when taking finite differences), SINDy can now discover the correct dynamics that match Equation 6, even with a much larger library Θ deg ≤ n ( x ) for n up to 16. For example, when n = 16 , the output ODE from SINDy reads d I 1 d t = 1.000 V 0 , d V 0 d t = − 1.000 I 1 , where the coefficients are rounded to three decimal places. This suggests the potential usefulness of concatenating different time series, especially in cases where one may not be familiar with the order of the relevant closure.

For Example 2 with the selected moments x = [ I 2 , V 1 , K ] , our findings are similar to those in Section 4.2. When applying SINDy to a data matrix from a single IC, SINDy discovers erroneous dynamics for the quadratic dictionary Θ deg ≤ 2 . However, using larger data matrices from multiple ICs, SINDy can once again accurately identify the correct dynamics even with the larger dictionaries. Representative negative and positive results are presented in the Supplementary Material, similar to the previous example.

We now turn to Example 3, where the selected moments are x = [ I 2 , V 1 , K , J ] , and the moment system only closes after a coordinate transformation. Specifically, we consider the nonlinearity g ( ρ , t ) = ρ 2 in Equation 1, which satisfies the condition in Example 3. We collect the moment time series data by numerically solving the PDE with the following four distinct ICs, Equations 13–16: u 1 x , 0 = π − 1 / 4 ⁡ exp − 1 2 x − 5 2 , (13) u 2 x , 0 = 1.88 ⁡ exp − 1 2 x − 5 2 , (14) u 3 x , 0 = 1.88 exp − 1 2 x − 5 2 + exp − x − 2 2 , (15) u 4 x , 0 = 1.88 cos 2 x + sin 2 x exp − x 2 . (16)

We consider the following questions:

• What insights can we gain by directly applying SINDy to this system with the selected moments x = [ I 2 , V 1 , K , J ] , where a closure does not exist?

We note that normalized moment time series are used as input for SINDy, as a way of incorporating feature scaling. This is an important aspect that ensures that all the relevant quantities are considered on “equal footing,” when the sparse regression step takes place.

We begin by applying SINDy with a linear library Θ deg = 1 ( x ) to the moment time series data from IC Equation 14. The resulting equations are: d I 2 d t = 1.000 V 1 , d V 1 d t = − 2.000 I 2 + 4.000 K + 3.998 J , d K d t = − 0.569 V 1 , d J d t = 0.069 V 1 . (17)

The coefficients are rounded to three decimal places. The first four panels of Figure 1 compare the “ground-truth” time-evolution data of [ I 2 , V 1 , K , J ] obtained from PDE integration (red curves) and the definition of these quantities through Equations 2–5 with those from integrating the SINDy-predicted ODEs Equation 17 (blue curves). While the SINDy-predicted time evolution of [ I 2 , V 1 ] closely matches the ground truth, there is a significant discrepancy for the moments [ K , J ] in Figure 1. Accordingly, the SINDy-predicted dynamics is not accurate for the original coordinates x = [ I 2 , V 1 , K , J ] .

Nevertheless, interestingly, if we add the ODE for K with that for J in Equation 17, which corresponds to the correct coordinate transformation for closure, we obtain: d I 2 d t = 1.000 V 1 , d V 1 d t = − 2.000 I 2 + 4.000 K + 3.998 J , d K + J d t = − 0.500 V 1 . (18)

Equation 18 closely matches the ground-truth moment system Equation 7. The last panel of Figure 1 shows that the predicted time evolution of K + J aligns perfectly with the ground truth, even though SINDy does not accurately recover those of K and J individually. This finding is intriguing because, even when SINDy is applied to the original coordinates [ I 2 , V 1 , K , J ] without closure, it “strategically chooses” to sacrifice accuracy for K and J but indirectly captures the correct dynamics when considering the proper coordinate transformation E = K + J . Moreover, a similar pattern is consistently observed when applying SINDy to the moments simulated from the other three ICs. The results are shown in the Supplementary Material. We have also provided an analogous example for 2D NLS equations. For details, please refer to Section 3 of the Supplementary Material. Indeed, this clearly illustrates that the methodology presented herein is not in any fundamental way limited to 1D NLS models, but can be extended to higher dimensional ones, as may be deemed desirable.

Next, we investigate the performance of SINDy with an expanded quadratic library Θ deg ≤ 2 ( x ) applied to the moment time series with the same IC Equation 14. The predicted ODEs now become: d I 2 d t = 1.000 V 1 , d V 1 d t = − 2.000 I 2 + 4.000 K + 3.998 J , d K d t = − 0.174 V 1 − 0.003 V 1 K − 0.083 V 1 J , d J d t = − 0.002 I 2 V 1 + 0.082 V 1 J . (19)

As before, if we add the predicted dynamics of K and J in Equation 19, we obtain: d I 2 d t = 1.000 V 1 , d V 1 d t = − 2.000 I 2 + 4.000 K + 3.998 J , d K + J d t = − 0.174 V 1 − 0.003 V 1 K − 0.002 I 2 V 1 − 0.001 V 1 J . (20)

Unlike the previous case with Θ deg = 1 ( x ) , the SINDy-predicted governing Equation 20 using the expanded library Θ deg ≤ 2 ( x ) after the (theoretically motivated) coordinate transformation E = K + J still fails to match the ground truth ODE Equation 7 and remains unclosed (i.e., the third equation in Equation 20 cannot be written using only I 2 , V 1 , and E = K + J ). Specifically, the coefficient of V 1 in Equation 20 is − 0.174 , deviating substantially from that of − 0.5 in the ground truth. Although the coefficients for the additional terms ( V 1 K , I 2 V 1 , and V 1 J ) are relatively small, their presence hinders the accurate recovery of the correct coefficient for V 1 .

Figure 2 compares the ground-truth time evolutions of [ I 2 , V 1 , K , J , E = K + J ] obtained from PDE integration with their SINDy-predicted dynamics after integrating Equations 19, 20. Similar to Figure 1, SINDy with the quadratic library Θ deg ≤ 2 ( x ) applied to the unclosed moments [ I 2 , V 1 , K , J ] sacrifices the exact recovery of the time evolution of J and K , but indirectly captures the correct time evolution of E = J + K . However, unlike the linear library case in Section 4.5, SINDy with the quadratic library can yield the correct time evolution of E = K + J but fails to uncover the correct governing equation for E [Equation 20] due to the expanded library size, leading to overfitting issues discussed above.

Here, we see a key complication within SINDy that has been also present in works such as (Champion et al., 2019) and especially (Bakarji et al., 2022), namely the methodology is likely to result in ODE models that are proximal to the theoretically expected ones, but not identical to them. This, in turn, may result in a nontrivial error outside of the training set. Especially when the library at hand is “richer” than the terms expected to be present, unfortunately, it does not generically seem that the reduced, theoretically expected model will be discovered. Rather, our results suggest that it is possible that the additional “wealth” of the libraries used can be leveraged to approximate the data via different (nonlinear) dependent variable combinations.

As an even more problematic example, applying SINDy with a quadratic library to the moment time series generated from another IC Equation 16 results in a predicted ODE system that not only fails to match the ground truth, even after the theoretically suggested coordinate transformation K + J , but also causes the ODE system to blow up in finite time. For a detailed discussion, we refer the interested reader to the Supplementary Material.

Next, we test the methodology from Section 3.2 to discover the coordinate transformation needed to close the moment system.

Case 1: μ = 0 . We first consider the case where the hyperparameter μ in Equation 12 is set to μ = 0 , i.e., we do not promote sparsity in the matrix Ξ . We use the linear library Θ ( x ) = Θ deg = 1 ( x ) , and set the maximum number of iterations to m a x I t e r = 150 in Algorithm 1. The algorithm produces the following outputs A o u t ≈ − 0.973 0.179 0.144 − 0.169 − 0.982 0.083 − 0.111 − 0.040 − 0.697 − 0.111 − 0.040 − 0.697 , Ξ o u t ≈ − 0.024 − 1.076 0.005 1.016 0.470 − 0.540 0.913 5.778 − 0.446 . (21)

The coefficients are rounded to three decimal places. On the other hand, the “ground-truth” coordinate transformation A g t (after normalization to satisfy the Stiefel manifold constraint) and the corresponding dynamics A g t according to Equation 7 are: A g t = 1 0 0 0 1 0 0 0 1 2 0 0 1 2 ≈ 1 0 0 0 1 0 0 0 0.707 0 0 0.707 ,    Ξ g t = 0 − 2 0 1 0 − 1 2 2 0 4 2 0 ≈ 0 − 2 0 1 0 − 0.354 0 5.657 0 (22)

At first glance, the predicted solution Equation 21 appears different from the ground truth in Equation 22. However, after a further change of coordinates using O = − 0.973 − 0.169 − 0.156 0.179 − 0.982 − 0.0565 0.144 0.0830 − 0.986 ∈ O 3 , we have A o u t O = 1.000 − 0.000 0.000 0.000 1.000 − 0.000 0.000 − 0.000 0.707 0.000 0.000 0.707 ≈ A g t , O ⊤ Ξ o u t O = 0.000 − 2.000 0.000 1.000 0.000 − 0.354 0.000 5.657 0.000 ≈ Ξ g t .

Hence, according to Remark 2, ( A o u t , Ξ o u t ) and ( A g t , Ξ g t ) are equivalent solutions, and our method successfully predicts the correct coordinate transformation A o u t to close the moment system. However, as expected, due to μ = 0 , the linear combination matrix Ξ o u t of the transformed dictionary is not sparse.

Case 2: μ > 0 . We next explore the case when μ ≠ 0 in Equation 12. We use the same linear library Θ ( x ) = Θ deg = 1 ( x ) , and μ is initially set to μ = 1 . In Algorithm 1, we set m a x I t e r = 1000 , I t e r _ s c h e d u l e d = 400 , and α = 20 . The algorithm returns A ̃ o u t ≈ − 1.000 0.000 0.030 0.000 − 0.100 0.002 − 0.021 − 0.001 − 0.707 − 0.021 − 0.001 − 0.707 , Ξ ̃ o u t ≈ 0.000 − 1.830 0.003 1.000 0.011 − 0.383 0.000 5.714 − 0.010 . (23)

Similarly, after a further change of coordinates using: O ̃ = − 1.000 0.000 − 0.030 0.000 − 1.000 − 0.002 0.030 0.002 − 1 ∈ O 3 , we again have A ̃ o u t O ̃ = 1.000 − 0.000 0.000 0.000 1.000 − 0.000 0.000 − 0.000 0.707 0.000 0.000 0.707 ≈ A g t ,    O ̃ ⊤ Ξ ̃ o u t O ̃ = 0.000 − 2.000 0.000 1.000 0.000 − 0.354 0.000 5.657 0.000 ≈ Ξ g t .

Note that the result ( A ̃ o u t , Ξ ̃ o u t ) in Equation 23 for μ > 0 is much closer to the ground truth ( A g t , Ξ g t ) in Equation 22, compared to ( A o u t , Ξ o u t ) in Equation 21 with μ = 0 . This demonstrates that a positive μ with annealed optimization not only discovers the correct coordinate transformation to close the system but also achieves a sparser solution for Ξ , reducing the number of terms on the right-hand side of the moment ODE system.

Finally, we present a case without analytical closure. Our “data-driven closure” aims to provide an accurate, reduced-order description of the PDE by approximating the evolution of the moment systems. In principle, this is the type of problem that we are aiming for, namely the discovery of potential moment closures when these may not be analytically available; the examples presented previously are valuable benchmarks to raise the complications that may emerge when one seeks to use this type of methodology in systems where the answer may be unknown and what credibility one may wish to assign to the obtained results.

Specifically, we consider an NLS Equation 1 with a time-dependent nonlinearity: g ρ , t = sin t + 2 ρ 2 . (24)

Notice that such time-dependent nonlinearities are well-known for some time in atomic physics settings (Donley et al., 2001; Staliunas et al., 2002) (and continue to yield novel insights to this day (Shagalov and Friedland, 2024)) and similar dynamical scenarios have been considered in nonlinear optics (Centurion et al., 2006; Zhang et al., 2021).

Moment systems with such nonlinearity will not close to the best of our knowledge, so we aim to numerically approximate form of the dynamics of the moments x = [ I 2 , V 1 , E ] , where E = K + J . We consider the following three ICs: u 1 x , 0 = 1.88 ⁡ exp − 1 2 x − 5 2 , (25) u 2 x , 0 = 1.88 cos 2 x + sin 2 x exp − x 2 , (26) u 3 x , 0 = exp − 0.1 x 2 exp 0.1 i x 2 . (27)

Notably, the IC Equation 27 includes a quadratic phase, motivated by the quadratic phase approximation (QPA) ansatz for NLS equations discussed by (Pérez-García et al., 2007). It also includes regular, smooth localized initial conditions, as well as one involving Fourier mode oscillations, modulated by the Gaussian term. We expect the SINDy-predicted dynamics to vary with the different ICs of Equations 25–27, hence the relevant choices.

Due to the periodic nature of the nonlinearity g in Equation 24, we expect the moment system to be non-autonomous and exhibit an oscillatory pattern. To capture this, we introduce the following two “auxiliary moments”: H = sin t + 2 , Q = cos t + 2 .

Naturally, one can observe that these are inspired by the nature of g ( ρ , t ) . However, one can envision the use of Fourier modes even if the mathematical model was not known or if the data stemmed from experimental observations.

We then apply SINDy with a linear library Θ deg = 1 ( x ̃ ) to the expanded moment systems x ̃ = [ x , H , Q ] = [ I 2 , V 1 , E , H , Q ] . The time series data of the moments are collected by integrating the PDE up to T = 20 . The exact SINDy-predicted ODEs for all three ICs are presented in the Supplementary Material.

To test the accuracy of the predicted moment systems, we integrate the SINDy-predicted ODEs into the future, up to T = 100 . Figure 3 compares the ground-truth and SINDy-predicted time evolutions for the three ICs Equations 25–27. In all cases, the model closely matches the ground truth up to T = 100 . Indeed, this is well past the training time of T = 20 and thus provides a satisfactory reduced-order description of the underlying PDE effective moment dynamics. Thus, despite the potential shortcomings of the method which we tried to present in an unbiased fashion in case examples where the analytical theory helps assess them, we still find it to be a worthwhile tool to consider. I.e., from a data-driven perspective, it can be seen to potentially provide an effective, low-dimensional dynamical representation of the associated high-dimensional PDE dynamics.