furax 0.7__tar.gz

furax-0.7/.gitignore

@@ -0,0 +1,19 @@
+.coverage
+MANIFEST
+build
+dist
+_site
+htmlcov
+__pycache__
+*.egg-info
+/venv*/
+/.vscode/
+/.idea/
+
+/tests/data/planck/
+
+# python
+*.pyc
+*.pyo
+*.pyd
+*.pytest_cache

furax-0.7/.gitlab-ci.yml

@@ -0,0 +1,68 @@
+image: python:3.10
+
+stages:
+  - testing
+  - qa
+  - publish
+  - docs
+
+workflow:
+  rules:
+    - if: '$CI_PIPELINE_SOURCE == "merge_request_event"'
+    - if: '$CI_PIPELINE_SOURCE == "push" && ($CI_COMMIT_TAG != null || $CI_COMMIT_BRANCH == $CI_DEFAULT_BRANCH)'
+
+variables:
+  PIP_CACHE_DIR: "${CI_PROJECT_DIR}/.cache/pip"
+  CONTAINER_TEST_IMAGE: $CI_REGISTRY_IMAGE:$CI_COMMIT_REF_SLUG
+  CONTAINER_RELEASE_IMAGE: $CI_REGISTRY_IMAGE:latest
+
+cache:
+  key: "${CI_JOB_NAME}"
+  paths:
+    - .cache/pip
+
+.testing-unit:
+  stage: testing
+  before_script:
+    - pip install --upgrade pip
+    - pip install jax[cpu]
+    - pip install .[dev]
+  script:
+    - pytest -m "not slow"
+  rules:
+    - if: '$CI_COMMIT_TAG == null'
+
+testing-unit/3.10:
+  extends: .testing-unit
+  image: python:3.10
+  coverage: '/^TOTAL.+?(\d+\%)$/'
+
+testing-unit/3.11:
+  extends: .testing-unit
+  image: python:3.11
+
+testing-unit/3.12:
+  extends: .testing-unit
+  image: python:3.12
+
+pre-commits:
+  stage: qa
+  needs: []
+  before_script:
+    - pip install pre-commit
+    - pre-commit install
+  script:
+    - pre-commit run --all-files
+  rules:
+    - if: '$CI_COMMIT_TAG == null'
+
+publish-release-package:
+  stage: publish
+  interruptible: true
+  before_script:
+    - pip install twine build setuptools_scm
+  script:
+    - python -m build
+    - TWINE_PASSWORD=$CI_JOB_TOKEN TWINE_USERNAME=gitlab-ci-token python -m twine upload --repository-url ${CI_API_V4_URL}/projects/${CI_PROJECT_ID}/packages/pypi dist/*.tar.gz dist/*.whl
+  rules:
+    - if: '$CI_COMMIT_TAG =~ /^v\d+\.\d+/'

furax-0.7/.pre-commit-config.yaml

@@ -0,0 +1,43 @@
+repos:
+- repo: https://github.com/hadialqattan/pycln
+  rev: "v2.4.0"
+  hooks:
+  - id: pycln
+    args:
+    - --all
+
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.7.1
+  hooks:
+    - id: ruff
+    - id: ruff-format
+
+- repo: https://github.com/pre-commit/pre-commit-hooks
+  rev: 'v4.6.0'
+  hooks:
+  - id: trailing-whitespace
+  - id: end-of-file-fixer
+  - id: check-yaml
+  - id: check-merge-conflict
+
+- repo: https://github.com/pre-commit/mirrors-mypy
+  rev: 'v1.11.1'
+  hooks:
+  - id: mypy
+    additional_dependencies:
+    - healpy
+    - jax
+    - jaxtyping
+    - jax-healpy
+    - jax-dataclasses
+    - lineax
+    - types-PyYAML
+    - scipy
+    args:
+    - --strict
+    - --show-error-codes
+    - --enable-error-code=ignore-without-code
+    - --allow-untyped-calls
+    files: ^src/furax/
+
+exclude: '^slurms/'

furax-0.7/PKG-INFO

@@ -0,0 +1,140 @@
+Metadata-Version: 2.1
+Name: furax
+Version: 0.7
+Summary: Operators and solvers for high-performance computing.
+Author-email: Pierre Chanial <chanial@apc.in2p3.fr>
+Maintainer-email: Pierre Chanial <chanial@apc.in2p3.fr>
+Project-URL: homepage, https://scipol.in2p3.fr
+Project-URL: repository, https://github.com/CMBSciPol/furax
+Keywords: scientific computing
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: jaxtyping
+Requires-Dist: healpy>=0.16.6
+Requires-Dist: jax-healpy
+Requires-Dist: jax-dataclasses
+Requires-Dist: lineax
+Requires-Dist: numpy
+Requires-Dist: scipy
+Requires-Dist: typing-extensions; python_version < "3.11"
+Requires-Dist: matplotlib
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: pytest-cov; extra == "dev"
+Requires-Dist: pytest-mock; extra == "dev"
+Requires-Dist: setuptools_scm; extra == "dev"
+Requires-Dist: beartype; extra == "dev"
+
+# Furax
+
+Furax: a Framework for Unified and Robust data Analysis with JAX.
+
+This framework provides building blocks for solving inverse problems, in particular in the astrophysical and cosmological domains.
+
+## Installation of the stable packaged distribution
+
+The GitLab repository is currently private and two pieces of information need to be specified: the URL of the repository and the credentials.
+
+- Create a Personal Access Token by following these [instructions](https://gitlab.in2p3.fr/help/user/profile/personal_access_tokens).
+For the name, choose something like `read scipol` and for the scopes, select `read_api`.
+
+- Edit the file ~/.netrc to store the credentials
+```
+machine gitlab.in2p3.fr
+login __token__
+password <your-personal-access-token>
+```
+Make sure the .netrc file are appropriately set:
+```commandline
+chmod 600 ~/.netrc
+```
+- Edit ~/.config/pip/pip.conf
+```
+[global]
+index-url = https://gitlab.in2p3.fr/api/v4/projects/26092/packages/pypi/simple
+```
+- After these steps are completed: you can install `furax` in the normal way:
+```commandline
+pip install furax
+```
+
+## Installation for development purposes
+
+```bash
+git clone git@github.com:CMBSciPol/furax.git
+cd furax
+python3 -m venv venv
+source venv/bin/activate
+pip install -e .[dev]
+```
+
+Then [Install JAX](https://jax.readthedocs.io/en/latest/installation.html) according to the target architecture.
+
+### Testing
+To check that the package is correctly installed:
+```bash
+pytest
+```
+
+## Running on JeanZay
+
+### Load cuda and and cudnn for JAX
+
+```bash
+module load cuda/11.8.0 cudnn/8.9.7.29-cuda
+```
+
+### Create Python env (only the first time)
+
+```bash
+module load python/3.10.4 && conda deactivate
+python -m venv venv
+source venv/bin/activate
+# install jax
+pip install --upgrade "jax[cuda11_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+# install furax
+pip install -e .[dev]
+```
+### launch script
+
+To launch only the pytests
+
+```bash
+sbatch slurms/astro-sim-v100-testing.slurm
+```
+To launch your own script
+
+```bash
+sbatch slurms/astro-sim-v100-run.slurm yourscript.py
+```
+
+You can also allocate ressources and go into bash mode
+
+```bash
+srun --pty --account=nih@v100 --nodes=1 --ntasks-per-node=1 --cpus-per-task=10 --gres=gpu:1 --hint=nomultithread bash
+module purge
+module load python/3.10.4
+source venv/bin/activate
+module load cuda/11.8.0  cudnn/8.9.7.29-cuda
+# Then do your thing
+python my_script.py
+pytest
+```
+Don't leave the bash running !! (I would suggest running script with sbatch)
+
+### Specific for nih / SciPol project
+
+The repo is already in the commun WORK folder, the data is downloaded and the environment is ready.
+
+You only need to do this
+
+```bash
+cd $ALL_CCFRWORK/furax-main
+```
+
+Then launch scripts as you see fit

furax-0.7/README.md

@@ -0,0 +1,108 @@
+# Furax
+
+Furax: a Framework for Unified and Robust data Analysis with JAX.
+
+This framework provides building blocks for solving inverse problems, in particular in the astrophysical and cosmological domains.
+
+## Installation of the stable packaged distribution
+
+The GitLab repository is currently private and two pieces of information need to be specified: the URL of the repository and the credentials.
+
+- Create a Personal Access Token by following these [instructions](https://gitlab.in2p3.fr/help/user/profile/personal_access_tokens).
+For the name, choose something like `read scipol` and for the scopes, select `read_api`.
+
+- Edit the file ~/.netrc to store the credentials
+```
+machine gitlab.in2p3.fr
+login __token__
+password <your-personal-access-token>
+```
+Make sure the .netrc file are appropriately set:
+```commandline
+chmod 600 ~/.netrc
+```
+- Edit ~/.config/pip/pip.conf
+```
+[global]
+index-url = https://gitlab.in2p3.fr/api/v4/projects/26092/packages/pypi/simple
+```
+- After these steps are completed: you can install `furax` in the normal way:
+```commandline
+pip install furax
+```
+
+## Installation for development purposes
+
+```bash
+git clone git@github.com:CMBSciPol/furax.git
+cd furax
+python3 -m venv venv
+source venv/bin/activate
+pip install -e .[dev]
+```
+
+Then [Install JAX](https://jax.readthedocs.io/en/latest/installation.html) according to the target architecture.
+
+### Testing
+To check that the package is correctly installed:
+```bash
+pytest
+```
+
+## Running on JeanZay
+
+### Load cuda and and cudnn for JAX
+
+```bash
+module load cuda/11.8.0 cudnn/8.9.7.29-cuda
+```
+
+### Create Python env (only the first time)
+
+```bash
+module load python/3.10.4 && conda deactivate
+python -m venv venv
+source venv/bin/activate
+# install jax
+pip install --upgrade "jax[cuda11_local]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+# install furax
+pip install -e .[dev]
+```
+### launch script
+
+To launch only the pytests
+
+```bash
+sbatch slurms/astro-sim-v100-testing.slurm
+```
+To launch your own script
+
+```bash
+sbatch slurms/astro-sim-v100-run.slurm yourscript.py
+```
+
+You can also allocate ressources and go into bash mode
+
+```bash
+srun --pty --account=nih@v100 --nodes=1 --ntasks-per-node=1 --cpus-per-task=10 --gres=gpu:1 --hint=nomultithread bash
+module purge
+module load python/3.10.4
+source venv/bin/activate
+module load cuda/11.8.0  cudnn/8.9.7.29-cuda
+# Then do your thing
+python my_script.py
+pytest
+```
+Don't leave the bash running !! (I would suggest running script with sbatch)
+
+### Specific for nih / SciPol project
+
+The repo is already in the commun WORK folder, the data is downloaded and the environment is ready.
+
+You only need to do this
+
+```bash
+cd $ALL_CCFRWORK/furax-main
+```
+
+Then launch scripts as you see fit

furax-0.7/pyproject.toml

@@ -0,0 +1,114 @@
+[build-system]
+build-backend = 'setuptools.build_meta'
+requires = [
+    'setuptools==69.2.0',
+    'setuptools_scm[toml]>=6.2',
+    'wheel',
+]
+
+[project]
+name = 'furax'
+authors = [
+    {name = 'Pierre Chanial', email = 'chanial@apc.in2p3.fr'},
+]
+maintainers = [
+    {name = 'Pierre Chanial', email = 'chanial@apc.in2p3.fr'},
+]
+description = 'Operators and solvers for high-performance computing.'
+readme = 'README.md'
+keywords = [
+    'scientific computing',
+]
+classifiers = [
+    'Programming Language :: Python',
+    'Programming Language :: Python :: 3',
+    'Intended Audience :: Science/Research',
+    'Operating System :: OS Independent',
+    'Topic :: Scientific/Engineering',
+]
+requires-python = '>=3.10'
+license = {file = 'LICENSE'}
+dependencies = [
+    'jaxtyping',
+    'healpy>=0.16.6',
+    'jax-healpy',
+    'jax-dataclasses',
+    'lineax',
+    'numpy',
+    'scipy',
+    'typing-extensions; python_version < "3.11"',
+    'matplotlib',
+]
+dynamic = ['version']
+
+[project.optional-dependencies]
+dev = [
+    'pytest',
+    'pytest-cov',
+    'pytest-mock',
+    'setuptools_scm',
+    'beartype',
+]
+
+[project.urls]
+homepage = 'https://scipol.in2p3.fr'
+repository = 'https://github.com/CMBSciPol/furax'
+
+[tool.coverage.report]
+exclude_lines = [
+    'pragma: no cover',
+    'except ImportError:',
+    'raise NotImplementedError',
+    'if typing.TYPE_CHECKING:',
+]
+show_missing = true
+skip_covered = true
+
+[[tool.mypy.overrides]]
+module = [
+    'healpy',
+    'jax_healpy',
+    'lineax',
+    'scipy.stats.sampling'
+]
+ignore_missing_imports = true
+
+[tool.pytest.ini_options]
+# addopts = '-ra --cov=furax --jaxtyping-packages=furax,beartype.beartype(conf=beartype.BeartypeConf(strategy=beartype.BeartypeStrategy.On))'
+addopts = '-s -ra'
+testpaths = [
+    'tests',
+]
+markers = [
+    "slow: mark test as slow.",
+]
+
+#[tool.setuptools]
+#packages = ['src/furax']
+
+[tool.setuptools_scm]
+version_scheme = 'post-release'
+fallback_version = '0.0.0'
+
+[tool.ruff]
+line-length = 100
+
+[tool.ruff.lint]
+select = [
+    # pycodestyle
+    'E',
+    # pyflakes
+    'F',
+    # pyupgrade
+    'UP',
+    # flake8-debugger
+    'T10',
+]
+ignore = ['E203', 'E731', 'E741', 'F722']
+
+[tool.ruff.lint.per-file-ignores]
+"src/furax/_base/core.py" = ['E743']
+"src/furax/landscapes.py" = ['UP007']
+
+[tool.ruff.format]
+quote-style = 'single'

furax-0.7/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

furax-0.7/slurms/astro-sim-v100-run.slurm

@@ -0,0 +1,47 @@
+#!/bin/bash
+#SBATCH --account=nih@v100
+#SBATCH --job-name=astro-sim-run # Name of job
+# Other partitions are usable by activating/uncommenting
+# one of the 5 following directives:
+#SBATCH -C v100-16g                 # decommenter pour reserver uniquement des GPU V100 16 Go
+##SBATCH -C v100-32g                 # decommenter pour reserver uniquement des GPU V100 32 Go
+##SBATCH --partition=gpu_p2          # decommenter pour la partition gpu_p2 (GPU V100 32 Go)
+##SBATCH -C a100                     # decommenter pour la partition gpu_p5 (GPU A100 80 Go)
+# Ici, reservation de 8x10=80 CPU (4 taches par noeud) et de 8 GPU (4 GPU par noeud) sur 2 noeuds :
+#SBATCH --nodes=1                  # nombre de noeud
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1      # nombre de tache MPI par noeud (= nombre de GPU par noeud)
+#SBATCH --gres=gpu:1               # nombre de GPU par nœud (max 8 avec gpu_p2, gpu_p4, gpu_p5)
+
+# Le nombre de CPU par tache doit etre adapte en fonction de la partition utilisee. Sachant
+# qu'ici on ne reserve qu'un seul GPU par tache (soit 1/4 ou 1/8 des GPU du noeud suivant
+# la partition), l'ideal est de reserver 1/4 ou 1/8 des CPU du noeud pour chaque tache:
+#SBATCH --cpus-per-task=10           # nombre de CPU par tache (un quart du noeud ici)
+##SBATCH --cpus-per-task=3           # nombre de CPU par tache pour gpu_p2 (1/8 du noeud 8-GPU)
+##SBATCH --cpus-per-task=6           # nombre de CPU par tache pour gpu_p4 (1/8 du noeud 8-GPU)
+##SBATCH --cpus-per-task=8           # nombre de CPU par tache pour gpu_p5 (1/8 du noeud 8-GPU)
+# /!\ Attention, "multithread" fait reference a l'hyperthreading dans la terminologie Slurm
+#SBATCH --hint=nomultithread         # hyperthreading desactive
+#SBATCH --time=01:00:00              # maximum execution time requested (HH:MM:SS)
+## Use this if your code runs fast or you are in dev phase
+##SBATCH --qos=qos_gpu-dev            # quality of service (QOS) dev
+#SBATCH --output=astro-sim-run.out # name of output file
+#SBATCH --error=astro-sim-run.err  # name of error file 
+# Cleans out modules loaded in interactive and inherited by default
+module purge
+
+# Uncomment the following module command if you are using the "gpu_p5" partition
+# to have access to the modules compatible with this partition.
+#module load cpuarch/amd
+
+# Loading modules
+module load python/3.10.4
+source venv/bin/activate
+
+module load cuda/11.8.0  cudnn/8.9.7.29-cuda 
+# Echo of launched commands
+
+set -x
+# For the "gpu_p5" partition, the code must be compiled with the compatible modules.
+# Code execution with binding via bind_gpu.sh : 1 GPU per task
+srun python $1

furax-0.7/slurms/astro-sim-v100-testing.slurm

@@ -0,0 +1,50 @@
+#!/bin/bash
+#SBATCH --account=nih@v100
+#SBATCH --job-name=astro-sim-testing   # Name of job
+# Other partitions are usable by activating/uncommenting
+# one of the 5 following directives:
+#SBATCH -C v100-16g                 # decommenter pour reserver uniquement des GPU V100 16 Go
+##SBATCH -C v100-32g                 # decommenter pour reserver uniquement des GPU V100 32 Go
+##SBATCH --partition=gpu_p2          # decommenter pour la partition gpu_p2 (GPU V100 32 Go)
+##SBATCH -C a100                     # decommenter pour la partition gpu_p5 (GPU A100 80 Go)
+# Ici, reservation de 8x10=80 CPU (4 taches par noeud) et de 8 GPU (4 GPU par noeud) sur 2 noeuds :
+#SBATCH --nodes=1                  # nombre de noeud
+#SBATCH --ntasks=1
+#SBATCH --ntasks-per-node=1      # nombre de tache MPI par noeud (= nombre de GPU par noeud)
+#SBATCH --gres=gpu:4               # nombre de GPU par nœud (max 8 avec gpu_p2, gpu_p4, gpu_p5)
+
+## this does not work
+#SBATCH --gpus-per-task=4           # nombre de GPU par tache
+
+# Le nombre de CPU par tache doit etre adapte en fonction de la partition utilisee. Sachant
+# qu'ici on ne reserve qu'un seul GPU par tache (soit 1/4 ou 1/8 des GPU du noeud suivant
+# la partition), l'ideal est de reserver 1/4 ou 1/8 des CPU du noeud pour chaque tache:
+#SBATCH --cpus-per-task=10           # nombre de CPU par tache (un quart du noeud ici)
+##SBATCH --cpus-per-task=3           # nombre de CPU par tache pour gpu_p2 (1/8 du noeud 8-GPU)
+##SBATCH --cpus-per-task=6           # nombre de CPU par tache pour gpu_p4 (1/8 du noeud 8-GPU)
+##SBATCH --cpus-per-task=8           # nombre de CPU par tache pour gpu_p5 (1/8 du noeud 8-GPU)
+# /!\ Attention, "multithread" fait reference a l'hyperthreading dans la terminologie Slurm
+#SBATCH --hint=nomultithread         # hyperthreading desactive
+#SBATCH --time=01:00:00              # maximum execution time requested (HH:MM:SS)
+## Use this if your code runs fast or you are in dev phase
+##SBATCH --qos=qos_gpu-dev            # quality of service (QOS) dev
+#SBATCH --output=astro-sim-testing.out # name of output file
+#SBATCH --error=astro-sim-testing.err  # name of error file 
+# Cleans out modules loaded in interactive and inherited by default
+module purge
+
+# Uncomment the following module command if you are using the "gpu_p5" partition
+# to have access to the modules compatible with this partition.
+#module load cpuarch/amd
+
+# Loading modules
+module load python/3.10.4
+source venv/bin/activate
+
+module load cuda/11.8.0  cudnn/8.9.7.29-cuda 
+# Echo of launched commands
+
+set -x
+# For the "gpu_p5" partition, the code must be compiled with the compatible modules.
+# Code execution with binding via bind_gpu.sh : 1 GPU per task
+srun python -m pytest -s

furax-0.7/src/furax/__init__.py

@@ -0,0 +1,14 @@
+from . import tree
+from ._base.axes import MoveAxisOperator, RavelOperator, ReshapeOperator
+from ._base.config import Config
+
+__all__ = [
+    # _base.config
+    'Config',
+    # tree
+    'tree',
+    # _base.axes
+    'MoveAxisOperator',
+    'RavelOperator',
+    'ReshapeOperator',
+]