BorgJaxPM

BorgJaxPM is a Python module that integrates a dark matter N-body forward model, developed using JaxPM. This module leverages the JAX library for acceleration and automatic differentiation, providing efficient and scalable simulations for astrophysical research.

Features

• Dark Matter N-body Simulations: Simulate the evolution of dark matter particles under gravitational interactions.
• Accelerated Computations: Utilizes JAX for GPU acceleration and automatic differentiation.
• Scalability: Designed to handle large-scale simulations with ease.
• Ease of Integration: Seamlessly integrates with existing astrophysical research workflows.

Installation

BorgJaxPM is developed with Python-Poetry support. Follow these steps to install the software:

1. Install Poetry: If you haven't already, install Poetry by following the instructions on the official Poetry website.

2. Clone the Repository: Clone the BorgJaxPM repository from GitHub.

   git clone https://git.aquila-consortium.org/guilhem_lavaux/borgjaxpm
   cd borgjaxpm
   

3. Install Dependencies: Use Poetry to install the required dependencies.

   poetry install
   

4. Activate the Virtual Environment: Activate the virtual environment created by Poetry.

   poetry shell
   

Tutorial / Demonstration

We do have a google colab to showcase the model performances: https://colab.research.google.com/drive/1H-CSxDpnLtJ6d7hW431x58fTL6W-eBNI?usp=sharing

Usage

Once installed, you can use BorgJaxPM in your Python scripts or Jupyter notebooks. Here is a basic example of how to use the module:

import borgjaxpm as bjax
import numpy as np
import aquila_borg as borg
import numpy as np
import matplotlib.pyplot as plt

# Set up the borg
pars = borg.cosmo.CosmologicalParameters()

pars.sigma8 = 0.8
pars.A_s = 0

# Setup simulation box
bx = borg.forward.BoxModel()

bx.L = 500, 500, 500
bx.N = 32,32,32
a_init = 0.1

# Create a simulation chain
chain = borg.forward.ChainForwardModel(bx)
chain @= borg.forward.model_lib.M_PRIMORDIAL(bx, opts=dict(a_final=1.0))
ic_step = borg.forward.model_lib.M_TRANSFER_EHU(bx)
chain @= ic_step
myjax = bjax.BorgJaxPM_v2(
    bx, a_init=a_init, integrator="leapfrog", num_steps=10, supersampling=1, forcesampling=1,
)
chain @= myjax

chain.setCosmoParams(pars)

x = np.random.randn(*bx.N)

# Create white noise
xh = np.fft.rfftn(x, norm='ortho')
delta_m = np.empty(chain.getOutputBoxModel().N)

# Run simulation
chain.forwardModel_v2(xh)
chain.getDensityFinal(delta_m)

# Plot the result
plt.title("BorgJaxPM")
plt.imshow(np.log(2+rho[:,16,:]),vmin=0,vmax=2,cmap=plt.cm.coolwarm)
plt.show()

License

BorgJaxPM is licensed under the CECILL-B license. This license is a permissive free software license initiated in France by INRIA, CEA, and CNRS. It allows for the reuse of the software code with limited obligations, primarily requiring attribution and compliance with French law. For more details, refer to the CECILL-B license agreement.