redback-jax 0.1.0

A lightweight JAX-only version of redback for electromagnetic transient analysis

Redback-JAX

A lightweight JAX-only version of the redback electromagnetic transient analysis package.

Overview

Redback-JAX provides JAX-based implementations for electromagnetic transient modeling and Bayesian inference, focusing on performance and automatic differentiation capabilities. This package is designed to be a lightweight alternative to the full redback package, leveraging JAX for fast computation and gradient-based inference.

Features

• JAX-based models: Fast, differentiable implementations of electromagnetic transient models
• Bayesian inference: Integration with modern probabilistic programming libraries
• Automatic differentiation: Gradient-based optimization and sampling
• GPU/TPU support: Leverage JAX's hardware acceleration capabilities

Installation

From PyPI (when available)

pip install redback-jax

From source

git clone https://github.com/nikhil-sarin/redback-jax.git
cd redback-jax
pip install -e .

With optional dependencies

pip install redback-jax[all]

Quick Start

import redback_jax

# Example usage will be added as the package develops

Dependencies

Core dependencies:

• JAX (>= 0.4.0)
• NumPy (>= 1.20.0)
• SciPy (>= 1.7.0)
• Pandas (>= 1.3.0)
• Matplotlib (>= 3.5.0)
• Astropy (>= 4.0.0)

Optional dependencies (install with [all]):

• NumPyro (>= 0.12.0)
• BlackJAX (>= 1.0.0)
• Optax (>= 0.1.0)

Contributing

Contributions are welcome! Please see the contributing guidelines in the main redback repository for details.

License

This project is licensed under the GNU General Public License v3.0 - see the LICENSE file for details.

Acknowledgments

This package is based on the original redback package. Please cite the original redback paper if you use this software in your research.

Related Projects

• redback - The full-featured electromagnetic transient analysis package
• JAX - The underlying numerical computing library