diff --git a/.gitignore b/.gitignore index 8ee24b55b..8cc016e5e 100644 --- a/.gitignore +++ b/.gitignore @@ -16,3 +16,4 @@ dist styles tests/regression/**/results uv.lock +healpy-data diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 000000000..ca2afec40 --- /dev/null +++ b/benchmarks/README.md @@ -0,0 +1,28 @@ +# GLASS Benchmarks + +These benchmarks are intended to allow benchmarking GLASS on various machines, +architectures and node configurations. + +Before running any benchmarks, you must first setup your venv with the required +dependencies. For all backends, this will require installing the dependency +group `benchmarks`, like so, + +```sh +uv sync --group benchmarks +``` + +However, it may also be necessary to install other dependencies / dependency +groups. For example, read the +[prerequisites for Archer2](./archer2/README.md#prerequisites) + +##  healpy-data + +glass depends on data from the healpy-data repo. If not found locally, glass +downloads this from the internet. Since most clusters will not allow internet +access from a worker node. We must provide a local copy of this data before +submitting a job. To download this data we can use git: + +```sh +git clone --depth 1 https://github.com/healpy/healpy-data +export HEALPY_DATAPATH="$(pwd)/healpy-data" +``` diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 000000000..96c79f0df --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1 @@ +"""Benchmarks from GLASS.""" diff --git a/benchmarks/archer2/README.md b/benchmarks/archer2/README.md new file mode 100644 index 000000000..351daad67 --- /dev/null +++ b/benchmarks/archer2/README.md @@ -0,0 +1,90 @@ +# Archer2 + +This document contains instructions specific to running the glass benchmarks on +Archer2 + +## Prerequisites + +As Archer2 is a managed cluster, there are several steps to setting up your +environment. First you should [setup uv](#setting-up-uv-on-archer2). Then you +will need to install you python virtual environment. Finally, you will need to +load specific modules for the GPU benchmark. + +Note that for the GPU benchmarks Archer2 only support rocm up to v0.6.x. +Therefore, we are restricted to using `jax==0.4.35`. This in turn restricts us +to using python 3.12. Thus, to produce a useful CPU vs GPU comparison create +your venv using the following command + +```sh +uv venv --python 3.12 +``` + +### CPU prerequisites + +To setup your python environment for running the cpu benchmark on Archer2, run +the following commands. + +```sh +uv sync --group benchmarks +``` + +### GPU prerequisites + +For the gpu benchmark, there is an additional dependency group `archer2-gpu`: + +```sh +uv sync --group benchmarks --group archer2-gpu +``` + +Once your python environment is setup you must load the relevant modules via the +provided script [setup-gpu-env.sh](./setup-gpu-env.sh). + +> Note that setup-gpu-env.sh is called automatically by the submission script +> [submit-gpu.sh](./submit-gpu.sh) + +## Running the benchmarks + +Benchmarks should be submitted as a batch job to slurm via the provided script. + +For example to benchmark using jax with amd/rocm, run the following from the +root of the glass repo on Archer2: + +```sh +sbatch benchmarks/archer2/submit-gpu.sh -d "$(pwd)" -x jax --healpy-datapath "$HEALPY_DATAPATH" +``` + +> To understand what HEALPY_DATAPATH is, read an explanation in +> [benchmarks/README.md#healpy-data](../README.md#healpy-data) + +## Setting up UV on Archer2 + +Firstly, install uv via curl onto the `/work` partition + +```sh +cd "${HOME/home/work}" +curl -LsSf https://astral.sh/uv/install.sh | sh +``` + +Then we must make sure uv is available on the login node and the worker node. To +do this we can update our start up scripts (`.profile`) on both partitions. +Therefore, execute the following + +```sh +cat <<'EOF' >> "$HOME/.profile" +WORK_DIR="${HOME/home/work}" +cd "$WORK_DIR" +source "$WORK_DIR/.profile" +EOF +``` + +and similarly + +```sh +cat <<'EOF' >> "${HOME/home/work}/.profile" +export HOME="${HOME/home/work}" +source "$HOME/.local/bin/env" +EOF +``` + +Now when you next login to archer2, uv will be in your path and you will be on +the `/work` partition as your `HOME` dir. diff --git a/benchmarks/archer2/setup-gpu-env.sh b/benchmarks/archer2/setup-gpu-env.sh new file mode 100755 index 000000000..2c322bf53 --- /dev/null +++ b/benchmarks/archer2/setup-gpu-env.sh @@ -0,0 +1,11 @@ +#!/bin/bash --login + +# Load GPU modules +module load PrgEnv-amd/8.6.0 +module load rocm +module load craype-accel-amd-gfx90a +module load craype-x86-milan + +# Ensure the rocm library and build is know to jax +export LD_LIBRARY_PATH="/opt/rocm/lib:$LD_LIBRARY_PATH" +export ROCM_PATH="/opt/rocm" diff --git a/benchmarks/archer2/submit-cpu.sh b/benchmarks/archer2/submit-cpu.sh new file mode 100755 index 000000000..cac3f6e06 --- /dev/null +++ b/benchmarks/archer2/submit-cpu.sh @@ -0,0 +1,103 @@ +#!/bin/bash --login +# shellcheck disable=SC1091 + +#SBATCH --job-name=glass_benchmark_cpu +#SBATCH --output=%x-%j.out +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 +#SBATCH --cpus-per-task=1 +#SBATCH --time=0:30:0 +#SBATCH --partition=standard +#SBATCH --qos=standard + +GLASS_DIR="" +ARRAY_BACKEND="numpy" +HEALPY_DATAPATH="" + +help() { + echo "Usage:" + echo " $0 -d [-x ] [--healpy-datapath ] [-h|--help]" + echo "" + echo "ARGS:" + echo " -h | --help Display this help message." + echo " -d | --glass-dir Path to the cloned glass directory." + echo " -x | --array-backend The array backend to use for the benchmarks." + echo " Defaults to NumPy." + echo " --healpy-datapath The path to the healpy-data repo to allow" + echo " running offline. Defaults to /healpy-data" +} + +# Ensure uv is available +source "${HOME/home/work}/.profile" # HOME starts as /home/... but uv needs to be on /work/... + +# Stop unintentional multi-threading within software libraries +export OMP_NUM_THREADS=1 +# Ensure the cpus-per-task option is propagated to srun commands +export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK + +# check for no input arguments and show help +if [ $# -eq 0 ]; +then + help + exit 1 +fi + +while [ $# -gt 0 ] ; do + case $1 in + -h | --help) + help + exit 0 + ;; + -d | --glass-dir) + GLASS_DIR="$2" + shift 2 + continue + ;; + -x | --array-backend) + ARRAY_BACKEND="$2" + shift 2 + continue + ;; + --healpy-datapath) + HEALPY_DATAPATH="$2" + shift 2 + continue + ;; + *) + echo "Invalid option: $1" >&2; + help + exit 1 + ;; + esac + shift 1 +done + +# Ensure GLASS_DIR is provided +if [[ "$GLASS_DIR" == "" ]] +then + echo "GLASS_DIR must be provided" + help + exit 1 +fi + +# Set HEALPY_DATAPATH default +if [[ "$HEALPY_DATAPATH" == "" ]]; then + HEALPY_DATAPATH="$GLASS_DIR/healpy-data" +fi + +# Get execution method +PYTHON_COMMAND="$GLASS_DIR/.venv/bin/python" +if [[ -n "$SLURM_JOB_ID" || -n "$SLURM_BATCH_SCRIPT" ]]; then + echo "Running under SLURM (batch). SLURM_JOB_ID=${SLURM_JOB_ID:-unknown}" + PYTHON_COMMAND="srun $PYTHON_COMMAND" +else + echo "Running directly from CLI" +fi + +for n in {128,256,512,1024} +do + echo "Running benchmark with nside/lmax = $n" + sed -i -E "s/nside = lmax = [0-9]+/nside = lmax = $n/g" benchmarks/lensing.py + + ARRAY_BACKEND="$ARRAY_BACKEND" HEALPY_DATAPATH="$HEALPY_DATAPATH" $PYTHON_COMMAND benchmarks/lensing.py +done diff --git a/benchmarks/archer2/submit-gpu.sh b/benchmarks/archer2/submit-gpu.sh new file mode 100755 index 000000000..af6b4afab --- /dev/null +++ b/benchmarks/archer2/submit-gpu.sh @@ -0,0 +1,104 @@ +#!/bin/bash --login +# shellcheck disable=SC1091 + +#SBATCH --job-name=glass_benchmark_gpu +#SBATCH --output=%x-%j.out +#SBATCH --cpus-per-task=1 +#SBATCH --nodes=1 +#SBATCH --gpus=1 +#SBATCH --time=8:00:0 +#SBATCH --partition=gpu +#SBATCH --qos=gpu-shd + +# Recommended environment settings +# Stop unintentional multi-threading within software libraries +export OMP_NUM_THREADS=1 +# Ensure the cpus-per-task option is propagated to srun commands +export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK + +GLASS_DIR="" +ARRAY_BACKEND="numpy" +HEALPY_DATAPATH="" + +help() { + echo "Usage:" + echo " $0 -d [-x ] [--healpy-datapath ] [-h|--help]" + echo "" + echo "ARGS:" + echo " -h | --help Display this help message." + echo " -d | --glass-dir Path to the cloned glass directory." + echo " -x | --array-backend The array backend to use for the benchmarks." + echo " Defaults to NumPy." + echo " --healpy-datapath The path to the healpy-data repo to allow" + echo " running offline. Defaults to /healpy-data" +} + +# Ensure uv is available +source "${HOME/home/work}/.profile" # HOME starts as /home/... but uv needs to be on /work/... + +# Stop unintentional multi-threading within software libraries +export OMP_NUM_THREADS=1 +# Ensure the cpus-per-task option is propagated to srun commands +export SRUN_CPUS_PER_TASK=$SLURM_CPUS_PER_TASK + +# check for no input arguments and show help +if [ $# -eq 0 ]; +then + help + exit 1 +fi + +while [ $# -gt 0 ] ; do + case $1 in + -h | --help) + help + exit 0 + ;; + -d | --glass-dir) + GLASS_DIR="$2" + shift 2 + continue + ;; + -x | --array-backend) + ARRAY_BACKEND="$2" + shift 2 + continue + ;; + --healpy-datapath) + HEALPY_DATAPATH="$2" + shift 2 + continue + ;; + *) + echo "Invalid option: $1" >&2; + help + exit 1 + ;; + esac + shift 1 +done + +# Ensure GLASS_DIR is provided +if [[ "$GLASS_DIR" == "" ]] +then + echo "GLASS_DIR must be provided" + help + exit 1 +fi + +# Set HEALPY_DATAPATH default +if [[ "$HEALPY_DATAPATH" == "" ]]; then + HEALPY_DATAPATH="$GLASS_DIR/healpy-data" +fi + +# Setup environment +source "$GLASS_DIR/benchmarks/archer2/setup-gpu-env.sh" + +for n in {128,256,512,1024} +do + echo "Running benchmark with nside/lmax = $n" + sed -i -E "s/nside = lmax = [0-9]+/nside = lmax = $n/g" benchmarks/lensing.py + + # Run benchmark via slurm + HEALPY_DATAPATH="$HEALPY_DATAPATH" ARRAY_BACKEND="$ARRAY_BACKEND" srun "$GLASS_DIR/.venv/bin/python" benchmarks/lensing.py +done diff --git a/benchmarks/benchmark_utils.py b/benchmarks/benchmark_utils.py new file mode 100644 index 000000000..c8dadc186 --- /dev/null +++ b/benchmarks/benchmark_utils.py @@ -0,0 +1,178 @@ +"""Helper functions for running glass benchmarks.""" + +from __future__ import annotations + +import cProfile +import io +import os +from pstats import SortKey, Stats +from timeit import timeit +from typing import TYPE_CHECKING + +import jax +import numpy as np +import array_api_strict + +if TYPE_CHECKING: + from types import FunctionType, ModuleType + from typing import Any + + from glass._types import FloatArray + from glass.cosmology import Cosmology + +xp_available_backends: dict[str, ModuleType] = {} + +# environment variable to specify array backends for testing +# can be: +# - a particular array library (numpy, jax, array_api_strict, ...) +# - all (try finding every supported array library available in the environment) +ARRAY_BACKEND: str = os.environ.get("ARRAY_BACKEND", "") + +# if no backend passed, use numpy by default +if not ARRAY_BACKEND or ARRAY_BACKEND == "numpy": + xp_available_backends["numpy"] = np +elif ARRAY_BACKEND == "array_api_strict": + xp_available_backends["array_api_strict"] = array_api_strict +elif ARRAY_BACKEND == "jax": + xp_available_backends["jax.numpy"] = jax.numpy +# if all, try importing every backend +elif ARRAY_BACKEND == "all": + xp_available_backends["numpy"] = np + xp_available_backends["array_api_strict"] = array_api_strict + xp_available_backends["jax.numpy"] = jax.numpy +else: + msg = f"unsupported array backend: {ARRAY_BACKEND}" + raise ValueError(msg) + +print("Running benchmarks for backends: ", ", ".join(xp_available_backends.keys())) # noqa: T201 + +# Configure backends +array_api_strict.set_array_api_strict_flags(api_version="2025.12") +jax.config.update("jax_enable_x64", val=True) + +RUN_PROFILE: str = os.environ.get("RUN_PROFILE") + + +def run_benchmark( + function_to_benchmark: FunctionType, + *args: tuple[Any, ...], + xp: ModuleType, + **kwargs: dict[str, Any], +) -> None: + """ + Run a benchmark of the provided function and its arguments. + + Parameters + ---------- + function_to_benchmark + The function which should be benchmarked. Note that the function must accept + all values passed via `args`, `xp` and `kwargs` + args + Positional arguments to be passed to `function_to_benchmark` + xp + The array backend to benchmark with. Will also be passed as a named argument to + `function_to_benchmark` + kwargs + Extra named arguments to be passed to `function_to_benchmark` + """ + if RUN_PROFILE: + pr = cProfile.Profile() + pr.enable() + + function_to_benchmark(*args, xp=xp, **kwargs) + + pr.disable() + sortby = SortKey.CUMULATIVE + s = io.StringIO() + ps = Stats(pr, stream=s).sort_stats(sortby) + ps.print_stats(0.05) + else: + # benchmark the task + result = timeit( + lambda: function_to_benchmark(*args, xp=xp, **kwargs), number=1 + ) + # report the result + print(f"Took {result:.3f} seconds with {xp.__name__}") # noqa: T201 + + +class CosmologyWrapper: + """An Python Array API compatible wrapper for a Cosmology instance.""" + + cosmo: Cosmology + cosmo_xp: ModuleType + xp: ModuleType + + def __init__( + self, *, cosmo: Cosmology, cosmo_xp: ModuleType = np, xp: ModuleType + ) -> None: + self.cosmo = cosmo + self.cosmo_xp = cosmo_xp + self.xp = xp + + @property + def Omega_m0(self) -> FloatArray: # noqa: N802 + """Matter density parameter at redshift 0.""" + return self.xp.asarray(self.cosmo.Omega_m0) + + @property + def critical_density0(self) -> FloatArray: + """Critical density at redshift 0 in Msol Mpc-3.""" + return self.xp.asarray(self.cosmo.critical_density0) + + @property + def hubble_distance(self) -> FloatArray: + """Hubble distance in Mpc.""" + return self.xp.asarray(self.cosmo.hubble_distance) + + def H_over_H0(self, z: FloatArray) -> FloatArray: # noqa: N802 + """Standardised Hubble function :math:`E(z) = H(z)/H_0`.""" + return self.xp.asarray(self.cosmo.H_over_H0(self.cosmo_xp.asarray(z))) + + def xm( + self, + z: FloatArray, + z2: FloatArray | None = None, + ) -> FloatArray: + """ + Dimensionless transverse comoving distance. + + :math:`x_M(z) = d_M(z)/d_H` + """ + if z2 is not None: + z2 = self.cosmo_xp.asarray(z2) + return self.xp.asarray( + self.cosmo.xm( + self.cosmo_xp.asarray(z), + z2, + ) + ) + + def rho_m_z(self, z: FloatArray) -> FloatArray: + """Redshift-dependent matter density in Msol Mpc-3.""" + return self.xp.asarray(self.cosmo.rho_m_z(self.cosmo_xp.asarray(z))) + + def comoving_distance( + self, + z: float, + z2: float | None = None, + ) -> FloatArray: + """Comoving distance :math:`d_c(z)` in Mpc.""" + return self.xp.asarray(self.cosmo.comoving_distance(z, z2)) + + def inv_comoving_distance(self, dc: FloatArray) -> FloatArray: + """Inverse function for the comoving distance in Mpc.""" + return self.xp.asarray( + self.cosmo.inv_comoving_distance(self.cosmo_xp.asarray(dc)) + ) + + def Omega_m(self, z: FloatArray) -> FloatArray: # noqa: N802 + """Matter density parameter at redshift z.""" + return self.xp.asarray(self.cosmo.Omega_m(self.cosmo_xp(z))) + + def transverse_comoving_distance( + self, + z: float, + z2: float | None = None, + ) -> FloatArray: + """Transverse comoving distance :math:`d_M(z)` in Mpc.""" + return self.xp.asarray(self.cosmo.transverse_comoving_distance(z, z2)) diff --git a/benchmarks/lensing.py b/benchmarks/lensing.py new file mode 100644 index 000000000..1673139e2 --- /dev/null +++ b/benchmarks/lensing.py @@ -0,0 +1,118 @@ +"""Benchmark for a realistic example lensing simulation.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +from benchmark_utils import CosmologyWrapper, run_benchmark, xp_available_backends + +# use the CAMB cosmology that generated the matter power spectra +import camb # ty: ignore[unresolved-import] +from cosmology.compat.camb import Cosmology # ty: ignore[unresolved-import] + +# almost all GLASS functionality is available from the `glass` namespace +import glass +import glass.ext.camb # ty: ignore[unresolved-import] +from glass import _rng + +if TYPE_CHECKING: + from collections.abc import Sequence + from types import ModuleType + + from glass._types import AngularPowerSpectra, FloatArray, UnifiedGenerator + from glass.shells import RadialWindow + + +# Run benchmarks for each requested backend +for xp in xp_available_backends.values(): + # cosmology for the simulation + h = 0.7 + Oc = 0.25 + Ob = 0.05 + + # basic parameters of the simulation + nside = lmax = 256 + + # set up CAMB parameters for matter angular power spectrum + pars = camb.set_params( + H0=100 * h, + omch2=Oc * h**2, + ombh2=Ob * h**2, + NonLinear=camb.model.NonLinear_both, + ) + results = camb.get_background(pars) + + # get the cosmology from CAMB + cosmo = CosmologyWrapper(cosmo=Cosmology(results), cosmo_xp=np, xp=xp) + + # shells of 200 Mpc in comoving distance spacing + zb = glass.distance_grid(cosmo, 0.0, 1.0, dx=200.0) + + # linear radial window functions + shells = glass.linear_windows(zb) + + # linear radial window functions using numpy to allow calling camb + shells_np = glass.linear_windows(np.asarray(zb)) + + # compute the angular matter power spectra of the shells with CAMB + cls = [xp.asarray(cl) for cl in glass.ext.camb.matter_cls(pars, lmax, shells_np)] # ty: ignore[unresolved-attribute] + + # apply discretisation to the full set of spectra: + # - HEALPix pixel window function (`nside=nside`) + # - maximum angular mode number (`lmax=lmax`) + # - number of correlated shells (`ncorr=3`) + cls = glass.discretized_cls(cls, nside=nside, lmax=lmax, ncorr=3) + + # set up lognormal fields for simulation + fields = glass.lognormal_fields(shells) + + # compute Gaussian spectra for lognormal fields from discretised spectra + gls = glass.solve_gaussian_spectra(fields, cls) + + # localised redshift distribution + # the actual density per arcmin2 does not matter here, it is never used + z = xp.linspace(0.0, 1.0, 101) + dndz = xp.exp(-((z - 0.5) ** 2) / (0.1) ** 2) + + # distribute dN/dz over the radial window functions + ngal = glass.partition(z, dndz, shells) + + shape = 12 * nside**2 + + def lensing_benchmark( # noqa: PLR0913 + *, + cosmo: CosmologyWrapper, + fields: Sequence[glass.grf.Lognormal], + gls: AngularPowerSpectra, + nside: int, + shells: list[RadialWindow], + xp: ModuleType, + ) -> tuple[FloatArray, FloatArray, FloatArray]: + """Realistic lensing simulation benchmark.""" + urng: UnifiedGenerator = _rng.rng_dispatcher(xp=xp) + + # this will compute the convergence field iteratively + convergence = glass.MultiPlaneConvergence(cosmo) + + # generator for lognormal matter fields + matter = glass.generate(fields, gls, nside, ncorr=3, rng=urng) + + # main loop to simulate the matter fields iterative + for i, delta_i in enumerate(matter): + # add lensing plane from the window function of this shell + convergence.add_window(delta_i, shells[i]) + + # compute shear field + glass.from_convergence(convergence.kappa, shear=True) + + # Run benchmark passing convergence and matter + run_benchmark( + lensing_benchmark, + cosmo=cosmo, + fields=fields, + gls=gls, + nside=nside, + shells=shells, + xp=xp, + ) diff --git a/glass/lensing.py b/glass/lensing.py index 566f48474..38a9a8d1b 100644 --- a/glass/lensing.py +++ b/glass/lensing.py @@ -34,8 +34,6 @@ import typing from typing import TYPE_CHECKING, Literal -import numpy as np - import array_api_compat import array_api_extra as xpx @@ -394,6 +392,8 @@ def shear_from_convergence( The shear map. """ + xp = kappa.__array_namespace__() + nside = hp.get_nside(kappa) if lmax is None: lmax = 3 * nside - 1 @@ -402,21 +402,21 @@ def shear_from_convergence( alm = hp.map2alm(kappa, lmax=lmax, pol=False, use_pixel_weights=True) # zero B-modes - blm = np.zeros_like(alm) + blm = xp.zeros_like(alm) # factor to convert convergence alm to shear alm - ell = np.arange(lmax + 1) - fl = np.sqrt((ell + 2) * (ell + 1) * ell * (ell - 1)) - fl /= np.clip(ell * (ell + 1), 1, None) + ell = xp.arange(lmax + 1, dtype=xp.float64) + fl = xp.sqrt((ell + 2) * (ell + 1) * ell * (ell - 1)) + fl /= xp.clip(ell * (ell + 1), 1, None) fl *= -1 # if discretised, factor out spin-0 kernel and apply spin-2 kernel if discretized: - pw0, pw2 = hp.pixwin(nside, lmax=lmax, pol=True, xp=np) - fl *= pw2 / pw0 + pw0, pw2 = hp.pixwin(nside, lmax=lmax, pol=True, xp=xp) + fl = fl * (pw2 / pw0) # apply correction to E-modes - hp.almxfl(alm, fl, inplace=True) + hp.almxfl(alm, fl, inplace=xp.__name__ != "jax.numpy") # transform to shear maps return hp.alm2map_spin([alm, blm], nside, 2, lmax) # ty: ignore[invalid-return-type] @@ -558,9 +558,9 @@ def add_plane( # lensing weight of mass plane to be added f = 3 * self.cosmo.Omega_m0 / 2 - f *= x2 * self.r23 - f *= (1 + self.z2) / self.cosmo.H_over_H0(self.z2) - f *= w2 + f = f * (x2 * self.r23) + f = f * ((1 + self.z2) / self.cosmo.H_over_H0(self.z2)) + f = f * w2 # create kappa planes on first iteration if self.kappa2 is None: diff --git a/glass/shells.py b/glass/shells.py index 2fbe65a4a..6c1162038 100644 --- a/glass/shells.py +++ b/glass/shells.py @@ -404,7 +404,7 @@ def linear_windows( ) if weight is not None: w *= weight(z) - ws.append(RadialWindow(z, w, zmid)) + ws.append(RadialWindow(z, w, float(zmid))) return ws diff --git a/pyproject.toml b/pyproject.toml index 0c7e2eb84..bf28c9ca9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,6 +7,20 @@ requires = [ ] [dependency-groups] +archer2-gpu = [ + "jax-rocm60-pjrt==0.4.35", + "jax-rocm60-plugin==0.4.35", +] +benchmarks = [ + "array_api_strict>=2.3.1", + "camb", + "cosmology-api", + "cosmology-compat-camb", + "glass-ext-camb", + "jax==0.4.35; python_version < '3.13'", + "jax>=0.6.2; python_version >= '3.13'", + "numpy<2.4", +] build = [ "build", ] @@ -244,6 +258,12 @@ lint.per-file-ignores = {"__init__.py" = [ ], "tests/*" = [ "INP001", # implicit-namespace-package "S101", # assert +], "tests/benchmarks/*" = [ + "D100", # undocumented-public-module + "N806", # non-lowercase-variable-in-function + "PLR2004", # magic-value-comparison + "S311", # suspicious-non-cryptographic-random-usage + "SLF001", # private-member-access ], "tests/core/*" = [ "ANN001", # missing-type-function-argument "ANN202", # missing-return-type-private-function diff --git a/tests/core/test_shells.py b/tests/core/test_shells.py index 831773ce4..232d09084 100644 --- a/tests/core/test_shells.py +++ b/tests/core/test_shells.py @@ -167,7 +167,7 @@ def test_linear_windows(xp: ModuleType) -> None: # check values of zeff - xpx.testing.assert_equal(xp.stack([w.zeff for w in ws]), zgrid[1:-1]) + xpx.testing.assert_equal(xp.asarray([w.zeff for w in ws]), zgrid[1:-1]) # check weight function input diff --git a/tests/regression/archer2/README.md b/tests/regression/archer2/README.md index c964f1afd..d783e6ef6 100644 --- a/tests/regression/archer2/README.md +++ b/tests/regression/archer2/README.md @@ -5,36 +5,8 @@ using [Archer2](https://www.archer2.ac.uk/). ## Setting up -1. **Install uv:** Firstly, install uv via curl onto the `/work` partition - - ```sh - cd "${HOME/home/work}" - curl -LsSf https://astral.sh/uv/install.sh | sh - ``` - - Then we must make sure uv is available on the login node and the worker node. - To do this we can update our startup scripts (`.profile`) on both partitions. - Therefore, execute the following - - ```sh - cat <<'EOF' >> "$HOME/.profile" - WORK_DIR="${HOME/home/work}" - cd "$WORK_DIR" - source "$WORK_DIR/.profile" - EOF - ``` - - and similarly - - ```sh - cat <<'EOF' >> "${HOME/home/work}/.profile" - export HOME="${HOME/home/work}" - source "$HOME/.local/bin/env" - EOF - ``` - - Now when you next log in to archer2, uv will be in your path, and you will be - on the `/work` partition as your `HOME` dir. +1. **Install uv:** Follow the instructions for + [Setting up UV on Archer2](../../benchmarks/README.md#setting-up-uv-on-archer2). 2. **Clone GLASS:** Clone the glass repo into the `/work` partition of Archer2: