Metadata-Version: 2.4
Name: proxide
Version: 0.1.0a4
Requires-Dist: jax
Requires-Dist: numpy
Requires-Dist: hydride
Requires-Dist: grain
Requires-Dist: jax-md
Requires-Dist: msgpack
Requires-Dist: msgpack-numpy
Requires-Dist: equinox
Requires-Dist: jaxtyping
Requires-Dist: array-record
Requires-Dist: typer>=0.12.0
Requires-Dist: rich>=13.0.0
Requires-Dist: tqdm>=4.66.0
Requires-Dist: jax ; extra == 'cpu'
Requires-Dist: jax[cuda12] ; extra == 'cuda'
Requires-Dist: ruff ; extra == 'dev'
Requires-Dist: pytest ; extra == 'dev'
Requires-Dist: pytest-cov ; extra == 'dev'
Requires-Dist: ty ; extra == 'dev'
Requires-Dist: chex ; extra == 'dev'
Requires-Dist: sphinx ; extra == 'dev'
Requires-Dist: sphinx-book-theme ; extra == 'dev'
Requires-Dist: myst-nb ; extra == 'dev'
Requires-Dist: sphinx-design ; extra == 'dev'
Requires-Dist: sphinx-copybutton ; extra == 'dev'
Requires-Dist: sphinx-autodoc-typehints ; extra == 'dev'
Requires-Dist: sphinxext-rediraffe ; extra == 'dev'
Requires-Dist: maturin ; extra == 'dev'
Requires-Dist: psutil ; extra == 'dev'
Requires-Dist: matplotlib ; extra == 'dev'
Requires-Dist: expaloma>=0.1.0 ; extra == 'espaloma'
Requires-Dist: rdkit ; extra == 'espaloma'
Requires-Dist: espaloma-charge>=0.0.8 ; extra == 'espaloma-openff'
Requires-Dist: foldcomp ; extra == 'foldcomp'
Requires-Dist: huggingface-hub ; extra == 'ml'
Requires-Dist: optax ; extra == 'ml'
Requires-Dist: rdkit ; extra == 'molecules'
Requires-Dist: openmm ; extra == 'openmm'
Requires-Dist: jax[tpu] ; extra == 'tpu'
Requires-Dist: mdtraj ; extra == 'trajectories'
Requires-Dist: h5py ; extra == 'trajectories'
Requires-Dist: biotite ; extra == 'trajectories'
Requires-Dist: tables ; extra == 'trajectories'
Provides-Extra: cpu
Provides-Extra: cuda
Provides-Extra: dev
Provides-Extra: espaloma
Provides-Extra: espaloma-openff
Provides-Extra: foldcomp
Provides-Extra: ml
Provides-Extra: molecules
Provides-Extra: openmm
Provides-Extra: tpu
Provides-Extra: trajectories
License-File: LICENSE
Summary: Protein I/O for JAX with high-performance Rust backend
Author-email: Marielle Russo <67157875+maraxen@users.noreply.github.com>
License-Expression: MIT
Requires-Python: >=3.11
Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
Project-URL: Homepage, https://github.com/maraxen/proxide
Project-URL: Issues, https://github.com/maraxen/proxide/issues
Project-URL: Repository, https://github.com/maraxen/proxide

# Proxide

**Proxide** is a high-performance library for Protein I/O and Physics bridging in JAX. It combines a flexible Python/JAX frontend with a highly optimized Rust backend (`_proxider`) to provide fast structure parsing, force field parameterization, and seamless integration with JAX MD.

**NOTE**: This is a research library in active development.

---

## 🚀 Features

- **Hybrid Architecture**:
  - **Rust Backend**: 25x faster parsing (PDB/mmCIF), 50x faster topology generation, and robust force field parameterization.
  - **JAX Frontend**: Differentiable physics, geometric deep learning utilities, and seamless GPU integration.
- **Robust I/O**: Load PDB, mmCIF, and PQR files with automatic error handling and corrections.
- **Molecular Dynamics**: Parse OpenMM XML force fields, assign GAFF parameters, and generate fully parameterized `AtomicSystem` objects for JAX MD.
- **Trajectory Support**: High-performance parsing of XTC, DCD, and TRR trajectories.

---

## 📦 Installation

Proxide requires a Rust toolchain to build the backend extension.

### Prerequisites

- **Python**: 3.11+
- **Rust**: 1.75+ (Install via [rustup.rs](https://rustup.rs))
- **C++ Compiler**: For compiling HDF5/chemfiles dependencies if needed.

### From Source

```bash
# Clone the repository
git clone https://github.com/maraxen/proxide.git
cd proxide

# Install with uv (recommended) or pip
uv pip install .

# For development (includes test dependencies)
uv pip install -e ".[dev]"
```

The installation process will automatically compile the Rust `_proxider` extension using `maturin`.

### Optional: Espaloma Charge (ML partial charges)

Inference uses the JAX/Equinox port [`expaloma`](https://github.com/maraxen/expaloma) (bundled weights; no PyTorch/DGL at runtime). The **`[espaloma]`** extra pulls **`expaloma`** from **PyPI** (and `rdkit`).

```bash
uv pip install -e ".[espaloma]"
```

For **local development** against an editable `expaloma` checkout, use a path override (e.g. `uv.sources` in your workspace) or:

```bash
uv pip install -e ../expaloma
uv pip install -e ".[dev]"
```

**OpenFF toolkit wrapper** (legacy upstream `espaloma_charge`): `uv pip install -e ".[espaloma-openff]"`.

API: `proxide.chem.partial_charges` (`assign_espaloma_charges_rdkit`, etc.). Validation uses **golden** vectors under `tests/data/espaloma_golden/`; refresh those files when bumping `expaloma`.

```bash
pytest -m espaloma
```

---

## 🛠️ Usage

### Loading a Structure

Use the high-level `parse_structure` function for fast, robust parsing:

```python
from proxide import parse_structure

# Load a PDB file to a unified Protein object
protein = parse_structure("path/to/structure.pdb")

# Access data as JAX arrays
print(protein.coordinates.shape)  # (N_residues, 37, 3)
```

### Force Field Parameterization

Proxide can automatically assign force field parameters (charges, radii, etc.) via the Rust backend:

```python
from proxide import parse_structure, OutputSpec

# Configure parsing options
spec = OutputSpec(
    add_hydrogens=True,             # Add missing geometric hydrogens
    infer_bonds=True,               # Infer connectivity if missing
    parameterize_md=True,           # Compute MD parameters
    force_field="protein.ff14SB.xml" # Use standard AMBER force field
)

# Parse and parameterize in one step
protein = parse_structure("path/to/structure.pdb", spec)

# Access MD parameters
print(protein.charges)  # Partial charges
print(protein.sigma)    # Lennard-Jones sigma
```

### Trajectory Parsing

```python
from proxide import parse_xtc

# Fast Rust-based XTC parsing
traj_data = parse_xtc("path/to/trajectory.xtc")
coords = traj_data["coordinates"]  # (N_frames, N_atoms, 3)
```

### Fast Partial Charges (Espaloma)

Proxide provides access to high-performance Expaloma graph-neural-network inference for charge assignment using native Rust embedded weights:

```bash
# Assign charges with the fast Rust backend
proxide charges molecule.sdf --backend rust
```

---

## ⚡ Performance

The migration to a Rust backend has yielded significant performance improvements compared to the pure Python implementation:

| Operation | Speedup | Example Metric |
|:----------|:--------|:---------------|
| **PDB Parsing** | **25x** | - |
| **mmCIF Parsing** | **25x** | - |
| **Topology Generation** | **50x** | - |
| **Force Field Loading** | **10x** | - |
| **Espaloma Inference** | **7.5x**| 13.6ms (JAX) ➔ 1.8ms (Rust) |

---

## 🚀 Deployment & CI/CD

Proxide uses **Maturin** to seamlessly build Python wheels and publish them to PyPI. To publish `proxide` releases securely without secrets, we utilize [PyPI Trusted Publishing](https://docs.pypi.org/trusted-publishers/) (OIDC).
Do **not** configure a `PYPI_API_TOKEN` GitHub Secret. Instead, create a Trusted Publisher in the PyPI dashboard bound to this repository to allow `.github/workflows/publish.yml` to authenticate.

---

## ⚠️ Migration Notes

If you are migrating from older versions of Proxide:

1. **Biotite Removal**: Direct dependency on `biotite` for parsing has been removed. All parsing is now handled by `_proxider`.
2. **API Changes**:
    - `proxide.io.parsing.biotite` -> `proxide.parse_structure`
    - `proxide.physics.force_fields` -> `proxide.load_forcefield`
3. **JAX by Default**: Most I/O functions now return JAX arrays by default. Use `use_jax=False` if you specifically need NumPy arrays.

---

## 🔧 Development

### Running Tests

```bash
# Run all tests
uv run pytest

# Run fast smoke tests
uv run pytest -m smoke
```

### Building Wheels Locally

Before pushing a release to GitHub, test the wheel build locally to catch platform-specific issues early:

```bash
# Test the build locally (validates HDF5 static compilation and workflow config)
bash scripts/test-wheel-build.sh

# This will:
# 1. Check build prerequisites (cargo, rustc, python3)
# 2. Build the Rust extension with static HDF5 compilation
# 3. Verify the GitHub workflow is correctly configured
```

**Note**: The Rust extension uses static HDF5 compilation via the `hdf5-metno` crate with the `static` feature enabled. This means HDF5 is compiled from source during the build and included in the binary, eliminating platform-specific library dependencies.

### Linting and Typing

```bash
# Linting
uv run ruff check src/proxide/ --fix

# Type Checking
uv run ty check
```

