ASKabalan/jax-fli-experiments
jax-fli experiments Data, samples, and reference catalogs for the jax-fli forward-modelling experiments. Each experiment is exposed as one or more HuggingFace dataset configs; load a config with datasets.load_dataset. Experiment 00 — CosmoGrid reference A single CosmoGrid simulation (cosmo_000001) packaged as jax-fli Catalog parquet files, used as reference truth for the lensing / masked-shear experiments. Config Field Shape NSIDE Description… See the full description on the dataset page: https://huggingface.co/datasets/ASKabalan/jax-fli-experiments.
jax-fli experiments
Data, samples, and reference catalogs for the `jax-fli` forward-modelling experiments. Each experiment is exposed as one or more HuggingFace dataset configs; load a config with datasets.load_dataset.
Experiment 00 — CosmoGrid reference
A single CosmoGrid simulation (cosmo_000001) packaged as jax-fli Catalog parquet files, used as reference truth for the lensing / masked-shear experiments.
import jax_fli as jfli
from datasets import load_dataset
from huggingface_hub import snapshot_download
from jax_fli.io import Catalog
REPO = "ASKabalan/jax-fli-experiments"
# --- Forecast convergence (kappa), nside 512: one small config, plain (non-streaming) load ---
kappa = Catalog.from_dataset(
load_dataset(REPO, "00-cosmogrid-000001-kappa", split="train").with_format("numpy")
).field[0] # SphericalKappaField (4, npix)
# --- Born convergence, nside 2048, 4 bins: one ~0.8 GB file per config, plain load ---
born_des = Catalog.from_dataset(
load_dataset(REPO, "00-cosmogrid-000001-born-des", split="train").with_format("numpy")
).field[0] # SphericalKappaField (4, npix) — DES Y3 source n(z); use "00-cosmogrid-000001-born-s3" for Stage-3
# --- Density lightcone, nside 2048, 56 shells: MUST be streamed ---
# A plain load_dataset concatenates the 56 (npix,) shell rows into one Arrow column and overflows its
# INT32 list-offset. Snapshot the per-shell parquets (+ README, which carries the config glob), then
# stream from the local dir and stack the shells into the (56, npix) lightcone.
local = snapshot_download(
REPO,
repo_type="dataset",
allow_patterns=["00-cosmogrid/cosmo_000001/density/cosmogrid_density_nside2048_shell*.parquet", "README.md"],
)
shells = [
Catalog.from_dataset(row).field[0]
for row in load_dataset(local, "00-cosmogrid-000001-density", split="train", streaming=True)
]
density = jfli.SphericalDensity.stack(shells) # SphericalDensity (56, npix) float32Experiment 01 — Resolution convergence
PM mesh-resolution study (512³–3072³, fixed box/steps/cosmology, independent IC realisation per mesh) painting 10 comoving HEALPix shells (nside 512, z≈0.017–0.35). See `docs/5-experiments/01-resolution-convergence`.
Three configs, each bundling all five resolutions m∈{512,1024,2048,2560,3072} as rows (the 01-resolution-density config additionally carries pencil-decomposition variants at m2560 and m3072, tagged with a _pencil suffix on each row's name):
The raw perf_pm.csv, the per-rung perf markdown reports, and the m<m>.args.log launch logs sit alongside it in 01-resolution/perf/ (plain files). Rows within a parquet config are ordered by filename, not mesh, so index them by each field's mesh_size:
from datasets import load_dataset
from jax_fli.io.catalog import Catalog
REPO = "ASKabalan/jax-fli-experiments"
cat = Catalog.from_dataset(
load_dataset(REPO, "01-resolution-spectra", split="train").with_format("numpy")
)
spectra = {int(f.mesh_size[0]): f for f in cat.field} # {512: PowerSpectrum(10,1536), 1024: ..., ...}