I'm a PhD candidate in computational cosmology at the AstroParticule & Cosmology Laboratory (APC, CNRS/IN2P3) in Paris, building differentiable and distributed cosmological simulations for the next generation of cosmological surveys.

Research Focus

• Full-field weak lensing inference - Developing scalable simulation pipelines for LSST
• CMB component separation - Part of the Simons Observatory collaboration
• Cross-survey cosmology - Combining CMB and large-scale structure datasets for enhanced constraints
• High-performance computing - Leveraging JAX, CUDA, and distributed computing for cosmological modeling

Main Projects

Core Packages

Project	Description
jax-fli	JAX Based Field Level Inference package with fully distributed and high resolution nbody and lensing simulations
jaxDecomp	JAX bindings to NVIDIA cuDecomp for distributed 3D FFTs and multi-GPU domain decomposition
FURAX_CS	Component separation pipeline for the Simons Observatory and LiteBird using FURAX
CADRE	Robust GPU-accelerated minimizer for parametric CMB component separation under low signal-to-noise and spatially varying foreground SEDs
JaxPM	Particle-mesh cosmological simulation toolkit in JAX with multi-accelerator support
FURAX	Framework for Unified and Robust data Analysis with JAX for inverse problems in cosmology
jax-healpy	JAX-native HEALPix utilities - GPU-ready and differentiable
jax-grid-search	Distributed grid search + gradient-based optimization built on JAX/Optax

Contributions

• S2FFT — Differentiable spherical/Wigner transforms (JAX & PyTorch). Contribution: translated part of the spherical harmonics algorithm to CUDA to avoid long JAX JIT compile times.

Tech Stack

Specialties: Distributed Computing | Automatic Differentiation | Bayesian Inference | HPC | CMB Analysis | Weak Lensing

Website: askabalan.github.io