Simulating Cold Dark Matter with Physics-Informed Kolmogorov-Arnold Networks

SKACH Winter Days, CSCS, Lugano, January 2026







Nicolas Cerardi

with Emma Tolley & Ashutosh Mishra
The dynamics of CDM
  • Cold Dark Matter obeys the Cosmological Vlasov Poisson equations:
    • $$ \frac{\partial f}{\partial t} + \vec{v} \cdot \vec{\nabla}_x f - \vec{\nabla}_x \Phi \cdot \frac{\partial f}{\partial \vec{v}} = 0 $$ $$ \nabla^2_x \Phi = 4\pi G a^2 \bar{\rho}(\tau) \delta(\vec{x}, \tau) $$

1D CDM simulation

Phase-space distribution

  • Cold DM occupies a 1D submanifold in the 1+1D phase-space
  • Multi-stream flow after shell-crossing
Density

  • Singularities arise at shell-crossing locations
  • They affect the CDM dynamics (Poisson equation)
Routes to CDM simulations
  • N-body codes employ particles instead of a fluid
  • This leads to discreteness errors (2-body interactions)


Can we find an alternative to N-body simulations ?

Since CDM occupies a low-dimensional submanifold, we can completely describe the dynamics using the displacement field $\zeta(q,\tau)$: $$ x(q,\tau) = q + \zeta(q,\tau) $$
With the equation of motion: $$ \frac{\partial^2 \zeta}{\partial \tau^2} + \frac{3}{2\tau} \frac{\partial \zeta}{\partial \tau} = - \frac{3}{2\tau} \frac{\partial \phi}{\partial x}$$
Can a neural network learn the solution ?

Physics Informed Neural Networks (PINNs)


PINNs: simulations seen as an optimisation problem
  • Use a neural network to learn the solution of a PDE
  • Network optimisation is guided by the physical constraints (PDE + boundary)
  • Fully unsupervised $\Longrightarrow$ independent from N-body
  • Provides a continuous representation of the solution, not just particles !

Physics Informed Neural Networks (PINNs)

PDE
  • Numerical integration allows to calculate $\nabla_x \phi$
  • Auto-diff provides access to the remaining term of the PDE (velocity and acceleration)

Boundary
  • Initial conditions given from the Zel'dovich Approximation
  • Boundary conditions set to 0

Which neural architecture can we use ?


  • Sharp features in acceleration due to the shell-crossings.
  • Our first attempts with MLPs failed to captures these.
  • We also tested more complex activation functions.
  • We need an architecture that provides high flexibility down to the 2nd derivative of the network output

Kolmogorov-Arnold Networks (KANs)


  • The Kolmogorov-Arnold representation theorem: any multivariate function can be approximated by a combination of (many) flexible univariate functions
  • Allows an efficient parameterization of the target function

Regular MLP

[Liu+25]
Kolmogorov-Arnold Network

Edges are combination of 3rd-order B-splines

$\Rightarrow$ Physics-Informed KANs (PIKANs!)

PIKAN implementation
  • We use a [8, 12, 8, 1] KAN architecture (~4500 parameters)
  • We employ independent KANs on consecutive time chunks
$$ x(q,\tau) = q + \textcolor{cyan}{\zeta_0} + \sum_{i=1}^{k-1} \textcolor{orange}{\zeta_i}(q, \tau_{end,i}) + \textcolor{green}{\zeta_k}(q,\tau)$$
  • $ \zeta_0 $ is the Zel'dovich approximation
  • $ \zeta_i $ are the KANs trained on previous chunks
  • $ \zeta_k $ is the KAN operating at $\tau$

Results

Results : Phase-space


A very good visual agreement, up to 6 shell-crossings !

Results : displacement, velocity and acceleration


Errors w.r.t a high resolution N-body simulation
Error is less than 2% on the displacement

Results : density profiles

  • Density profiles are in good agreement
  • In the core, N-body profiles are noisy
  • The PIKAN provides a smoother solution (too much?)

Caveats
  • Training time is quite long (several hours per time chunk on a single GPU)
  • While the learnt representation is defined continuously, we still need to sample points for the numerical integration of $ \nabla_x\phi $
  • (for now a 1D single halo collapse)

What else can we do with such a tool ?


  • Unlike a N-body sim that advances timestep per timestep, in the PINN framework the optimisation is performed simultaneously over time.
  • $\Rightarrow$ We can set the "initial conditions" at any $\tau$

Results : Backward simulation

  • Test case: reverse time through 2 shell crossings.
  • Start from high-resolution N-body data at $\tau_{end}$, optimise to retrieve the underlying initial conditions.

Input: $ (\zeta, \dot{\zeta}, \ddot{\zeta})$

Input: $ (\zeta, \dot{\zeta}, \delta(x))$

Thanks for your attention !


[Cerardi, Tolley and Mishra, MNRAS, 2025]
Takeaways
  • We presented a new method to simulate CDM, independent from N-body simulations.
  • PIKANs allows a continuous and differentiable modelling of the dynamics.
  • It is possible to run the simulation backward to retrieve the initial conditions.
  • Future work: higher dimensions, complex ICs...