Contributing to Furax#

We welcome contributions to Furax! This guide will help you get started with contributing code, documentation, or bug reports.

Development Setup#

Fork and Clone#

Fork the Furax repository on GitHub
Clone your fork locally:

git clone https://github.com/your-username/furax.git
cd furax

Development Installation#

Install Furax in development mode with all dependencies:

pip install -e .[dev]

This installs:

Core dependencies (JAX, Lineax, etc.)
Development tools (pytest, mypy, ruff, pre-commit)
Documentation tools (sphinx, etc.)

Pre-commit Hooks#

Set up pre-commit hooks to ensure code quality:

pre-commit install

This will automatically run linting, formatting, and type checking on every commit.

Code Quality Standards#

Formatting and Linting#

Furax uses Ruff for both linting and formatting:

# Check formatting and style
ruff check src/

# Auto-format code
ruff format src/

# Fix auto-fixable issues
ruff check --fix src/

Configuration:

Line length: 100 characters
String quotes: Single quotes preferred
Import sorting: Automatic

Type Checking#

We use MyPy for static type checking:

# Type check the core package
mypy src/furax/

Type checking is enforced only on the src/furax/ directory. External dependencies like healpy and jax-healpy are ignored.

Key requirements:

All public functions should have type annotations
Use jaxtyping for array type annotations
Complex types should be documented

Testing#

Test Structure#

Tests are organized in the tests/ directory:

tests/
├── conftest.py              # Global fixtures
├── core/                    # Linear operator tests
├── obs/                     # Observation framework tests
├── interfaces/              # External interface tests
└── data/                    # Test data files

Running Tests#

# Run all tests
pytest

# Run with coverage and colored output
pytest -s -ra --color=yes

# Run specific test file
pytest tests/core/test_diagonal.py

# Run tests matching pattern
pytest -k "test_diagonal"

# Run tests with specific markers
pytest -m "slow"

Test Features#

JAX x64 precision: Automatically enabled for numerical accuracy
Parametrized fixtures: Tests run with different Stokes combinations (I, QU, IQU, IQUV)
Data fixtures: Cached test data with automatic downloads
Custom assertions: Specialized checks for Furax data types

Writing Tests#

Use parametrized fixtures for comprehensive testing:

import pytest
from furax.obs.stokes import Stokes

@pytest.mark.parametrize("stokes_fixture", ["I", "QU", "IQU"], indirect=True)
def test_stokes_arithmetic(stokes_fixture):
    stokes_data = stokes_fixture

    # Test addition
    result = stokes_data + stokes_data
    assert isinstance(result, type(stokes_data))

    # Test scalar multiplication
    scaled = 2.0 * stokes_data
    assert isinstance(scaled, type(stokes_data))

HPC Testing#

For GPU testing on HPC clusters:

# Submit to SLURM queue (JeanZay example)
sbatch slurms/astro-sim-v100-testing.slurm

Code Architecture#

Core Principles#

Composability: Linear operators should compose naturally with @ and +
JAX Integration: All data structures are PyTrees compatible with JAX transformations
Type Safety: Extensive use of type hints and jaxtyping
Mathematical Clarity: Code should reflect mathematical operations clearly

Operator Development#

When creating new operators, inherit from AbstractLinearOperator:

from dataclasses import field

from furax import AbstractLinearOperator, symmetric
from jaxtyping import Array, Float

@symmetric
class MyCustomOperator(AbstractLinearOperator):
    """A custom homothety operator with static scaling factor.

    Example:

        >>> op = MyCustomOperator(10, in_structure=jax.ShapeDtypeStruct((2,), jnp.float32))
        >>> op(jnp.array([1., 2]))
        Array([10., 20.], dtype=float32)
        >>> op.I(jnp.array([1., 2]))
        Array([0.1, 0.2], dtype=float32)

    """
    scaling_factor: float = field(metadata={'static': True})

    def mv(self, x: Float[Array, "n"]) -> Float[Array, "n"]:
        # Implement the linear operation
        return self.scaling_factor * x

    def inverse(self) -> AbstractLinearOperator:
        # Overrides the default implementation
        return MyCustomOperator(1/self.scaling_factor, in_structure=self.in_structure)

    @property
    def is_negative_semidefinite(self) -> bool:
        return self.scaling_factor <= 0

    @property
    def is_positive_semidefinite(self) -> bool:
        return self.scaling_factor >= 0

Key requirements:

Implement mv for matrix-vector multiplication
Specify mathematical properties when known
Include comprehensive docstrings with examples

Data Structure Development#

New Stokes classes should follow the established pattern:

from furax.obs.stokes import Stokes
from jaxtyping import Array, Float

class StokesXY(Stokes):
    """Custom Stokes parameters for X and Y polarization."""

    X: Float[Array, "n_pix"]
    Y: Float[Array, "n_pix"]

    @classmethod
    def from_stokes(cls, x: Array, y: Array) -> "StokesXY":
        return cls(X=jnp.asarray(x), Y=jnp.asarray(y))

    @property
    def stokes(self) -> str:
        return "XY"

Documentation#

Docstring Style#

Use Google-style docstrings with type information:

def my_function(
    data: Float[Array, "n_pix"],
    scale: float = 1.0
) -> Float[Array, "n_pix"]:
    """Process CMB data with scaling.

    Args:
        data: Input CMB map with shape (n_pix,)
        scale: Scaling factor to apply

    Returns:
        Scaled CMB map with same shape as input

    Example:
        >>> import jax.numpy as jnp
        >>> data = jnp.array([1., 2., 3.])
        >>> result = my_function(data, scale=2.0)
        >>> print(result)
        [2. 4. 6.]
    """
    return scale * data

Building Documentation#

# Build HTML documentation
cd docs
make html

# View in browser
open build/html/index.html

Mathematical Notation#

Use proper LaTeX for mathematical expressions:

The maximum likelihood estimator is:

.. math::

   \\hat{m} = (P^T N^{-1} P)^{-1} P^T N^{-1} d

where :math:`P` is the pointing matrix.

Pull Request Process#

Before Submitting#

Ensure all tests pass: pytest
Check code quality: ruff check src/ and mypy src/furax/
Update documentation if needed
Add tests for new functionality

Pull Request Guidelines#

Clear Description: Explain what the PR does and why
Small, Focused Changes: One feature or fix per PR
Test Coverage: Include tests for new code
Documentation: Update docs for user-facing changes
Backwards Compatibility: Avoid breaking existing APIs without discussion

Example PR Description:

## Summary

Adds support for non-uniform noise in ToeplitzOperator

## Changes

- Modified SymmetricBandToeplitzOperator to accept per-pixel noise scaling
- Added unit tests for new functionality
- Updated documentation with usage examples

## Testing

- All existing tests pass
- New tests added in test_toeplitz.py
- Verified with realistic CMB noise simulation

Review Process#

Automated checks run on all PRs (tests, linting, type checking)
Code review by maintainers
Address feedback and update PR
Merge once approved and all checks pass

Issue Reporting#

Bug Reports#

Include:

Clear description of the problem
Minimal code example that reproduces the issue
System information (OS, Python version, JAX version)
Expected vs actual behavior

Feature Requests#