Cosmo-SPINN: Fuzzy Dark Matter Simulations with Physics-Informed Generative NetworksJournal: ApJ

Fuzzy DM consists of an ultralight boson with mass $m\ll 30$ eV resulting in a macroscopic de Broglie wavelength. An axion-like particle is one possible candidate for FDM. The axion $\phi$ is a real angular field with a periodicity of f, f being an axion-decay constant. Taking a primordial value of order f, the axion transitions from being frozen early in the universe to oscillating in the late times, resulting in a relic density of (1; 32; 18):

Assuming non-relativisitic regime and t as the proper time, we introduce complex-valued scalar field or wavefunction $\psi(\bm{x},t)\in\mathbb{C}$ with $\bm{x}\in\mathbb{R}^{3}$ that relates to the real axion field $\phi$ as:

where, $\hbar$ is the reduced Planck’s constant and $c$ is the speed of light in vacuum. Under the condition $|\dot{\psi}|\ll m|\psi|$, and adopting the perturbed Friedmann-Lemaître-Robertson-Walker (FLRW) metric,

the Klein-Gordon equation for the axion $\phi$ gives the following Schrödinger-Poisson equations:

where $a$ is the cosmological scale factor, $V$ is the Newtonian gravitational potential, the angle brackets in $\langle|\psi|^{2}\rangle$ indicate the spatial average, and in the above equations, all the quantities and coordinates are in the “comoving” form which are related to “physical” quantities as:

Despite the appearance of Schrödinger equation, the wavefunction $\psi$ should be interpreted as a classical field. $\psi$ is subject to wave effects such as interference, similar to interference of waves in classical electromagnetism. It should be noted that the Eqs. (4 & 5) describes the evolution of a single macroscopic wavefunction of a Bose-Einstein condensate, with a mass density $\rho=m|\psi|^{2}$, rather than the wavefunction of an individual particle. One can also introduce super-comoving time $d\tilde{t}=a^{-2}dt$ (33) and $\tilde{V}=aV$, to obtain the more familiar Schrödinger equation without scale factors:

This formulation can be useful because it removes scale factors in the denominator which makes controlling error at low redshifts very difficult.

The Schrödinger equation implies a conservation law described by the continuity equation

with the density current

The (comoving) velocity field $\mathbf{v}$ dictates the peculiar velocity of matter at each point. Using the 31 transformation:

with absolute value $\sqrt{\rho/m}$ and phase S, where

is indeed the same mass density as in Eq. (7). Inserting this into the expression for the current in Eq. (8) yields

that is, the gradient of the wave function’s phase determines the velocity field. The wave function written using the Madelung transformation (Eq. 9) , can be used to rewrite the SP Equations yielding the continuity Eq. (7) along with a modified Euler equation:

(cf. e.g. 37). This allows for a hydrodynamical interpretation of the density and velocity fields $\rho$ and $\mathbf{v}$. When the quantum pressure term vanishes, which happens on large scales (3; 18), the Madelung equations recover the pressureless Euler–Poisson equations, similar to the Vlasov-Poisson equations governing CDM in the limit $\hbar/m\to 0$ and before shell crossing. Despite the fact that the Madelung equations are highly non-linear, they become tractable with existing hydrodynamic codes. However, they are inherently ill-behaved due to the presence of quantum pressure term which blows up as $\rho\to 0$, which is common in the DM overdensity evolution where the density vanishes, especially in the voids between halos and filaments. Therefore we choose to evolve the SP-system of equations (4 & 5) instead of the Madelung formalism in our work, in particular with the super-comoving time formalism as it allows us to control the errors better at lower redshifts, permitting the long-time evolution.

It should be noted that Eqs.( 4 & 5 ) only have a single parameter given by the constant $\hbar/m$, which is related to the de Broglie wavelength as

Another important scale set by $\hbar/m$ is the FDM Jeans length. In FDM, quantum pressure arising from the Heisenberg uncertainty principle counteracts gravitational collapse below a characteristic scale. As a result, at least to linear order the density perturbations larger than the Jeans length can grow gravitationally, while smaller perturbations remain oscillatory and do not collapse. The corresponding comoving Jeans wavenumber $k_{J}$ at redshift $z$ (correspondingly at scale factor $a$) is given by (16; 26):

The simulations are performed within a standard cosmological $\Lambda$CDM background, where CDM is replaced by FDM inside a periodic cubic domain of side length $L$. The simulation volume samples the large-scale matter distribution of the Universe and is initialized with the mean comoving matter density

To evolve the system, we solve the Schrödinger–Poisson (SP) equations using the second-order unitary split-step pseudo-spectral method described in 37. The method combines operator splitting with Fast Fourier Transforms (FFT), allowing the kinetic and gravitational terms to be evolved independently while preserving unitary time evolution.

The formal evolution of the wavefunction over a timestep $\Delta t$ is given by

where $\mathcal{T}$ denotes the time-ordering operator. Using the Baker–Campbell–Hausdorff expansion, the evolution operator is approximated to second-order accuracy by separating the kinetic and potential contributions,

This decomposition corresponds to the standard “kick–drift–kick” scheme,

The potential “kick” updates are performed in configuration space, while the kinetic “drift” step is evaluated in Fourier space using FFT methods. Spatial derivatives are therefore computed spectrally, which improves numerical accuracy and computational efficiency.

At each timestep, the gravitational potential is obtained by solving the Poisson equation in Fourier space,

The numerical evolution over one timestep proceeds as follows: (i) compute the gravitational potential from the density field, (ii) apply a half-step potential update (kick), (iii) evolve the kinetic term for a full timestep in Fourier space (drift), (iv) recompute the gravitational potential, and (v) apply a second half-step potential update. The complete algorithm is summarized below.

This numerical scheme preserves second-order accuracy and unitary time evolution. The pseudo-spectral formulation efficiently computes spatial derivatives and gravitational interactions, making it well suited for simulations of wave-like dark matter and Bose–Einstein condensate systems. The kick operator acts locally in position space and only modifies the phase of the wavefunction. To avoid phase aliasing, the phase change per timestep is required to remain below $2\pi$, giving

where $|V|_{\max}$ is the maximum potential amplitude and $C_{\mathrm{kick}}\leq 1$ is a numerical safety factor.

The drift operator is evaluated in Fourier space and contributes a phase proportional to $k^{2}$. Requiring the maximum phase change to remain below $2\pi$ gives

where $\Delta x$ is the grid spacing and $k_{\max}=\pi/\Delta x$ is the Nyquist wavenumber. The factor $C_{\mathrm{drift}}\leq 1$ is an empirical stability coefficient whose value depends on the implementation scheme and the specific application at hand.

The scaling $\Delta t\propto\Delta x^{2}$ is characteristic of diffusion-type equations and corresponds to requiring that the maximum resolved velocity does not propagate more than one grid cell per timestep. This is the main reason why generating large scale FDM simulations are computationally extremely expensive and time consuming.

Physics Informed Generative Networks (PIGNs) extend physics-informed machine learning frameworks by combining generative neural architectures with physical constraints derived from the governing equations of the system. Unlike traditional Physics Informed Neural Networks (PINNs), which directly learn the solution of differential equations through collocation points in space-time, PIGNs are designed to generate physically consistent realizations of complex fields while preserving the statistical and dynamical properties imposed by the underlying physics. Such approaches are particularly useful for high-dimensional cosmological simulations and super-resolution tasks, where reconstructing high-resolution fields from low-resolution inputs is an inherently ill-posed inverse problem. Physics-informed generative models address this ambiguity by generating realizations that satisfy the governing equations while reproducing the correct physical structures and statistical distributions.

In this work, the generative networks are trained to learn a physics-informed mapping between cosmological field realizations while being constrained by the governing SP equations. Depending on the task, the architecture is adapted to either predict the cosmological fields at a specified cosmological scale factor from the initial conditions (evolution model) or reconstruct high-resolution realizations from low-resolution inputs (super-resolution model). In both cases, the networks take as input the relevant field variables together with the scale factor and predict physically consistent realizations of the wavefunction and gravitational potential fields across the simulation domain. Physical information is incorporated through additional loss terms constructed from the residuals of the governing equations, thereby regularizing the learning process and promoting consistency with the underlying dynamics.

The general form of the governing differential equations can be written as

with associated boundary and initial conditions,

where $G(X,\theta)$ denotes the generative neural network parameterized by $\theta$, $\mathcal{D}$ represents the non-linear differential operator governing the dynamics, $\Lambda$ denotes the physical parameters entering the equations, $f(X)$ is the source term, and $\mathcal{B}$ specifies the boundary or initial condition operator and $g(X)$ specifies the corresponding prescribed boundary or initial condition over the computational domain $\Omega\subset\mathbb{R}^{d+1}$ with boundary $\partial\Omega$. In our case, the residuals for the SP equations would be:

where $\psi$ is the complex wavefunction and $V$ is the gravitational potential. Periodic boundary conditions are expected to be learned purely from the data loss by the network.

The training objective of the PIGN combines the generative reconstruction loss with physics-based regularization terms. In addition to matching the target high-resolution realizations, the network is penalized for violating the SP equations and the imposed additional physical consistencies such as mass conservation and so on (if any). This enables the model to generate super-resolved realizations that remain consistent with the underlying cosmological dynamics while accurately reproducing the statistical properties of the simulated fields, whether for cosmological field evolution or super-resolution.

Initial conditions are generated at z=127 (corrresponding to a scale factor of $a\approx 0.0078)$ using AxionCAMB (15) with an axion mass of $m=2.5\times 10^{-22}$eV, assuming all matter is axion dark matter. The resulting power spectra are passed to JaxPM¹¹ 1 https://github.com/DifferentiableUniverseInitiative/JaxPM, which generates particle positions and velocities using second-order Lagrangian Perturbation Theory (2LPT). A uniform-grid density field is then constructed via Cloud-in-Cell (CIC) interpolation of the displacement field. The 2LPT velocity field is used to initialise the phase of the wavefunction. Taking the divergence of Eq. (10) gives

which is a Poisson equation solved spectrally. The initial wavefunction is then fully determined by

All our simulations were generated using the cosmological parameters $\Omega_{m}=0.27$, $\Omega_{b}=0$, $\Omega_{\Lambda}=0.73$, $H_{0}=100$ $\textrm{km}\ \textrm{s}^{-1}\textrm{Mpc}^{-1}$ ($h=1$; set to 1 so that units are in $h^{-1}$), and   $\sigma_{8}=0.8$ with ICs as described in section III.1. Multiple realizations of the initial conditions at z = 127 were generated using different random seeds while fixing the underlying cosmology and initial matter power spectrum. The simulation box size is set to $L=1\,h^{-1}\,\textrm{Mpc}$, a scale at which contributions from the quantum pressure term are non-negligible and must be accounted for. All the simulations are generated using the Jaxion (38) code.

Here we introduce Cosmo-SPINN, a physics-informed generative framework for FDM inside the periodic cubic box of length L = 1$\,h^{-1}$Mpc. As illustrated in Fig. 1, the framework is applied to two related tasks: (i) the evolution of cosmological fields from initial conditions to a specified scale factor and (ii) the super-resolution of FDM simulations. In both cases, the network predicts the real and imaginary components of the wavefunction together with the gravitational potential while enforcing consistency with the underlying SP dynamics through a physics-informed loss function.

We use $\lambda=\hbar/m$ to simplify notation. Writing out the real and imaginary components of the wavefunction $\psi$ explicitly as $\psi=\text{Re}(\psi)+i\,\text{Im}(\psi)=\textrm{R}+i\,\textrm{I}$, we can rewrite Eqs. (5 & 6) as:

As discussed in Section II.1, our choice to train the Cosmo-SPINN using the SP equations is motivated by their ability to capture the wave-like behavior inherent to FDM dynamics. Consequently, we adopt the SP formulation for training our physics-informed model. Specifically, we design the network to predict the real $\mathrm{R}$ and imaginary $\mathrm{I}$ components of the wavefunction $\Psi$, as well as the gravitational potential $V$, which is treated as an auxiliary output rather than computed through direct numerical solution. This strategy can significantly reduce the computational cost with automatic differentiation by circumventing the need to solve Poisson’s equation at every training step, as this is typically solved with FFTs and computing the Fourier transform for an $N\times N$ grid requires $\mathcal{O}(N^{2}logN)$ operations and approximately $\mathcal{O}(N^{2})$ memory, which can quickly be restrictive for larger grids such as in the super-resolution task. However, automatic differentiation requires the whole grid to be passed to the model which can be memory-intensive in case of a single GPU. As a result, we still resort to FFTs for derivatives’ calculations in our case.

We define a neural network $G(X;\theta)$ that predicts the real and imaginary components of the wavefunction together with the gravitational potential,

where $(\cdot)_{\theta}$ denotes the approximation realized by the network parameters $\theta$. The precise form of the input $X$ depends on the task considered. For the evolution model, the input consists of the initial conditions together with a target scale factor,

and the network predicts the fields at scale factor $a$. For the super-resolution model, the input consists of low-resolution fields and the corresponding scale factor,

and the network predicts the corresponding high-resolution fields. A schematic overview of both tasks is shown in Fig. 1. We predict real $\mathrm{R}$ and imaginary $\mathrm{I}$ parts of $\psi$ directly along with the potential $V$. This approach ensures mass conservation and correct evolution, eliminating the need to explicitly implement a continuity equation.

The data loss is defined as the mean squared error (MSE) between the predicted and target fields,

where $y=(\mathrm{R},\mathrm{I},V)$ contains the real and imaginary components of the wavefunction together with the gravitational potential, and $N_{\mathrm{elem}}$ is the total number of tensor elements (all spatial grid points and all field channels). The total loss is defined as

where $\lambda_{w}$ is a weighting coefficient, which is chosen such that both the losses contribute equally in the gradient space. We explore two formulations of the physics loss $\mathcal{L}_{\mathrm{physics}}$.

Both models share a common set of building blocks. Scalar conditions (scale factor $a$) are encoded via a Fourier feature embedding,

with $K=8$, producing a 16-dimensional embedding that is passed through a two-layer MLP to yield a $c$-dimensional conditioning vector (c = 64 for us). This vector is spatially broadcast and concatenated channel-wise to the input volume before the first convolution. Spatial feature extraction is performed by 3D ResNet blocks, each consisting of two $3\times 3\times 3$ convolutions with LeakyReLU activations and a skip connection; when the number of channels changes, the skip connection uses a $1\times 1\times 1$ convolution, otherwise it is the identity.

We first evaluate the convergence behaviour of the Evolution Model shown in Fig. 1 (left). Figure 2 shows the evolution of the training and test data losses (left panel) alongside the normalized SP loss (right panel). We consider two different traning strategies: (i) using only data loss, and (ii) using data loss and SP loss after epoch 250. After the physics-informed term is introduced at epoch 250, the optimization dynamics change as the network is required to satisfy both data fidelity and physical consistency. While the normalized SP loss decreases by nearly an order of magnitude for the physics-informed model, it remains nearly constant for the data-only baseline throughout the remainder of training, indicating that supervised optimization alone is insufficient to improve consistency with the underlying physical dynamics described by the SP equations. At the same time, both the training and test data loss continue to decrease while remaining comparable to those obtained using the data-only objective. Furthermore, the reduced train–test gap suggests that the SP constraint acts as a useful regularizer, guiding the model toward physically admissible solutions that also generalize better to unseen data.

The combined objective here is defined as

where $\lambda_{w}$ is an adaptive gradient-balancing factor that equalizes the gradient magnitudes of the data and SP losses, while $\alpha(t)$ controls the overall strength of the physics constraint. Rather than introducing the physics loss abruptly, its contribution is gradually increased according to

where $\alpha_{\max}=0.02$ and $N_{\mathrm{ramp}}=1000$ epochs. This warm-up strategy allows the network to first learn an accurate data-driven representation before progressively enforcing the physical constraint, thereby avoiding the optimization instabilities associated with suddenly introducing a strong additional objective.

Although the adaptive weighting balances the optimization in gradient space, the SP loss remains a considerably stronger constraint since it is evaluated in physical space and directly enforces the underlying nonlinear dynamics. Moreover, the evolution model is trained using a single realization of the initial conditions, for which 796 simulation snapshots are available over the interval $a\in[0.0798,\,0.0803]$. Only $20\%$ of these redshifts are used for supervised training, whereas the SP loss is evaluated on randomly sampled redshifts spanning the entire interval. Consequently, the physics objective constrains a substantially larger portion of the evolution than the supervised data loss. Without the additional scaling and gradual ramp-up, the optimization becomes biased toward satisfying the governing equations at the expense of accurately reconstructing the available training snapshots. We therefore set $\alpha_{\max}=0.02$, which provides a suitable balance between physical consistency and data fidelity. Empirically, larger values of $\alpha_{\max}$ further reduce the SP residual and improve the test loss, but lead to a noticeable deterioration in the reconstruction of the training data.

Unlike first-order time-stepping constraints such as KDK, the SP loss directly constrains the underlying dynamics rather than enforcing agreement over a single numerical integration step. Our objective is therefore to incorporate the full physical dynamics into the learning process whenever possible. Nevertheless, the increased complexity of the SP constraint makes the optimization more challenging, motivating us to investigate alternative formulations of the physics loss for more demanding learning tasks.

We therefore compare two different formulations of the physics loss for the super-resolution model shown in Fig. 1 (right): the SP residual and a first-order KDK approximation, as described in Section III.3. Unlike the evolution model, no additional scaling parameter $\alpha$ is introduced in this case. The super-resolution model is trained using all available redshifts, with the train–test split performed over independent realizations rather than temporal snapshots. Consequently, both the supervised and physics losses are evaluated over the same range of redshifts, eliminating the imbalance present in the evolution model. Our objective here is therefore to investigate whether an appropriate formulation of the physics loss can provide stable optimization without requiring an additional stabilization parameter such as $\alpha$.

For this, the super-resolution model is first trained for the initial 20 epochs using only the MSE data loss before introducing the respective physics losses. A baseline model trained with only MSE objective (data loss) is also included to isolate the effect of the physical regularization, as illustrated in Fig. 3 with the blue curves. Figure 3 shows the evolution of the normalized SP (left panel) and KDK (right panel) losses for both the baseline model (blue curves) and the physics-informed model (orange curves), with training (solid) and test (dashed) curves shown for each. In both cases, the inclusion of the physics loss leads to a significant reduction in the normalized physics losses on both the training and test sets, demonstrating that the model learns solutions that better satisfy the underlying dynamics. However, the two formulations of the physics loss exhibit noticeably different convergence behaviour. For the SP-based loss, the training residual decreases steadily, but after approximately 40 epochs the test residual begins to diverge from the training residual, regardless of whether the physics loss is included. The primary effect of the physics loss in this case is to reduce the overall magnitude of the residual rather than eliminate the growing train–test gap. This behaviour likely reflects the increased complexity of the coupled SP equations. In contrast, the KDK-based constraint produces a much smoother and more stable convergence. The training and test residuals closely track one another throughout training, indicating improved stability and better generalization to unseen realizations and redshift slices. Although the final training loss is slightly higher than that obtained with the SP formulation, the smaller train–test gap suggests that the KDK approximation provides a more robust physics regularization for this super-resolution task.

Overall, both approaches improve the physical consistency of the generated fields relative to the data-only model. The SP residual achieves a larger reduction in the physics error, while the KDK formulation yields more stable training and better agreement between training and test performance.

We now assess the performance of both models on unseen test data. We first examine the evolution model, evaluating its ability to reproduce the wavefunction and density fields. Additional results demonstrating its generalization to previously unseen realizations of the initial conditions are presented in Appendix B. We then consider the super-resolution model, which constitutes the primary contribution of this work, comparing reconstructions obtained using the SP-residual and KDK-based physics losses (throughout the remainder of this section, we denote the super-resolution model trained with the SP-residual loss as the SP model and the model trained with the KDK-based loss as the KDK model). All super-resolution results presented in this section are obtained using the model checkpoints selected according to the minimum physics test loss during training shown in Fig. 3. Note that, while all models used in this work were trained on the training set discussed in Section III.2, the snapshots presented here were taken from the reserved test set. That is, the simulation data used to produce the plots that follow were not used to train any of the models.

In most FDM simulations, initial conditions are generated exactly as described in this paper: first, the power spectrum is constructed using AxionCAMB for the given axion mass, then the initial conditions are usually obtained from the N-body particle distribution via density assignment such as the cloud-in-cell algorithm. The second step is especially convenient due to the availaibility of public codes (such as JaxPM for our case) for an arbitrary power spectrum.

However, this is not an ideal way to generate the initial conditions for field-based simulations. The reasons are twofold: first, the velocities of these N-body particles are typically computed using 2LPT method or Zeldovich approximation which are suited for CDM perturbations and secondly, converting from positions and velocities to fields incur some mesh error for FDM simulations. In this study, we nevertheless adopt this conventional particle-based initialization procedure because of its straightforward implementation and because we wanted to evaluate the performance of the ML models used here in regard to the inclusion of physics loss. In future work, we plan to generate more reliable initial conditions, following the methodology presented in the appendix of 30.

The approach adopted in this work is semi-supervised, combining simulation-derived ground-truth data with physics-based constraints during training. A natural question is whether accurate cosmological simulations can be learned solely from the governing physical laws, without any supervised data. For relatively simple systems, such as gravitational evolution of a single Fourier-mode density perturbation, our previous work on SPINN demonstrated that fully physics-driven learning is feasible. However, extending this approach to realistic cosmological settings, characterized by complex initial conditions, an expanding background, and long evolutionary timescales, remains challenging. In these regimes, the optimization landscape becomes highly non-convex, and current optimization strategies are often insufficient for training fully unsupervised PINNs. An autoregressive framework that predicts the system evolution one step at a time using only physical constraints could be a promising alternative, though controlling the accumulation of errors over long time horizons would be a significant challenge.

Our implementation of the super-resolution task relies on high-resolution (HR) FDM numerical simulations for the training objective. However, these simulations become prohibitively computationally expensive with increasing box sizes. An alternative is to train a generative model using only the physics loss (without any data loss), which leaves the solution space largely unconstrained and allows the network to generate multiple high-resolution realizations that satisfy the governing physics, rather than converging to a unique solution. Including a low-resolution data loss would further constrain the admissible solution space by requiring the generated high-resolution fields, when downsampled, to agree with the corresponding low-resolution inputs. The primary challenge is that these two losses often define competing optimization objectives, resulting in conflicting gradient updates that can hinder convergence and lead to suboptimal solutions.