Metadata-Version: 2.4
Name: foldkit
Version: 1.0.0
Summary: Toolkit for working with and storing AlphaFold3 co-folding results
Author-email: Jonathan Levine <jonalevine1@gmail.com>
License: MIT
Project-URL: Documentation, https://jonlevi.github.io/foldkit/
Project-URL: Source, https://github.com/jonlevi/foldkit
Keywords: protein,protein design,alphafold,bioinformatics
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: pandas
Requires-Dist: biopython
Requires-Dist: tqdm
Requires-Dist: argcomplete
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Provides-Extra: docs
Requires-Dist: sphinx; extra == "docs"
Requires-Dist: pydata-sphinx-theme; extra == "docs"
Requires-Dist: sphinx-copybutton; extra == "docs"
Requires-Dist: sphinx-autodoc-typehints; extra == "docs"
Requires-Dist: sphinx-design; extra == "docs"
Dynamic: license-file

# foldkit

foldkit is a Python toolkit for working with and efficiently storing AlphaFold3 (AF3) co-folding results.

It provides:

🐍 A Python API for easily accessing AF3 confidence metrics and structural ensembles

🧬 Convenient access to ensemble-level metrics across seeds and samples

📦 An efficient storage format that substantially reduces the size of AF3 output directories

🖥️ A command-line interface (CLI) for converting raw AF3 results into the compressed FoldKit format


foldkit is particularly useful for large-scale protein–protein and protein–peptide modeling campaigns where hundreds or thousands of AF3 predictions need to be stored and analyzed.

You can find the full documentation here: https://jonlevi.github.io/foldkit/index.html 

## Installation
`pip install foldkit`

### Bash autocompletion
FoldKit's CLI supports Bash autocompletion through `argcomplete`.

To enable it:

`activate-global-python-argcomplete`

On shared systems where you do not have permission to modify the system-wide configuration:

`activate-global-python-argcomplete --user` 

You may need to restart your shell after enabling autocompletion.

## foldkit API
foldkit has two primary use cases.

(1) **Convenient access AF3 confidence metrics for single structures and ensembles** 

foldkit provides a Python interface for accessing the confidence metrics generated by AF3, including metrics that describe interactions between chains, and for aggregating these metrics over ensembles of predicted structures across multiple seeds and samples.

(2) **Efficient storage and retrieval of AlphaFold3 results.**

 The default JSON formats for AF3 confidence results are large, and can take up a lot of unnecessary space. foldkit has a CLI for exporting the AF3 confidence JSONs to space-efficient .npz files, removing other unnecessary files, and copying over the rest. The resulting foldkit files can be loaded directly through the same Python API described below.

# Python Interface Tutorial

## 1. Loading a single AF3 prediction

Suppose you have an AF3 output directory containing a single predicted protein complex.

For example, consider a TCR–pMHC complex with four chains:

- `A` — TCRα
- `B` — TCRβ
- `M` — MHC
- `P` — peptide

The AF3 results are stored in:

```text
tutorial_example/single_result/
```

The directory contains:

```text
single_result/
├── confidences.json
├── model.cif
└── summary_confidences.json
```

Load the result with:

```python
import foldkit

result_obj = foldkit.AF3Result.load_af3_result(
    "tutorial_example/single_result"
)
```

The resulting `AF3Result` object provides access to the confidence metadata and methods for calculating statistics from it.

For example, you can inspect the chains:

```python
>>> result_obj.chains
[np.str_('A'), np.str_('B'), np.str_('M'), np.str_('P')]
```

### Confidence metrics

FoldKit provides convenient access to common AF3 confidence metrics.

For example:

```python
>>> result_obj.get_ptm()
0.81
```

Get the mean pTM for a specific chain:

```python
>>> result_obj.get_ptm("A")
0.82
```

Calculate the inter-chain pAE between TCRβ and the peptide:

```python
>>> result_obj.get_ipae(chain1="B", chain2="P")
np.float64(6.245691056910569)
```

Calculate ipSAE between the same chains:

```python
>>> result_obj.get_ipsae(chain1="B", chain2="P")
np.float64(0.292483968491584)
```

FoldKit's ipSAE implementation follows the methodology described by the [Dunbrack Lab IPSAE package](https://github.com/DunbrackLab/IPSAE).

---

## Custom aggregation functions

By default, FoldKit aggregates residue-level confidence metrics using the mean.

You can provide a custom aggregation function with the `agg` argument.

For example, to retrieve the maximum inter-chain pAE:

```python
>>> result_obj.get_ipae(
...     chain1="B",
...     chain2="P",
...     agg=max,
... )
np.float64(29.8)
```

This allows the same interface to be used for different ways of summarizing residue-level confidence matrices.

---

# Working with AF3 Ensembles

AF3 can generate multiple predictions of the same complex using different seeds and samples.

FoldKit provides an `AF3Ensemble` object for working with these predictions collectively.

Suppose your ensemble is organized as:

```text
tutorial_example/ensemble_result/
├── ranking_scores.csv
├── seed-1_sample-0/
├── seed-1_sample-1/
├── seed-1_sample-2/
├── ...
├── seed-5_sample-0/
├── seed-5_sample-1/
├── ...
└── seed-10_sample-4/
```

Load the entire ensemble with:

```python
>>> ensemble_obj = foldkit.AF3Ensemble.load_af3_result(
...     "tutorial_example/ensemble_result"
... )
```

FoldKit will load the individual predictions and construct an `AF3Ensemble` object.

Inspect the ensemble:

```python
>>> ensemble_obj.size
20

>>> ensemble_obj.seeds
[1, 2, 5, 10]

>>> ensemble_obj.samples
[0, 1, 2, 3, 4]
```

---

## Accessing individual predictions

The individual `AF3Result` objects are stored in:

```python
ensemble_obj.af3_results
```

You can retrieve a specific prediction directly using its seed and sample:

```python
>>> ensemble_obj.get_result_by_seed_and_sample(
...     seed=5,
...     sample=2,
... )
<foldkit.af3_result.AF3Result object at ...>
```

---

## Accessing the top-ranked prediction

AF3 ranking scores are stored in:

```python
ensemble_obj.af3_ranking_scores
```

The highest-ranked prediction can be retrieved directly:

```python
>>> ensemble_obj.get_top_ranked_result()
<foldkit.af3_result.AF3Result object at ...>
```

---

## Calculating metrics across an ensemble

FoldKit can calculate metrics either for individual structures or across the entire ensemble.

For example, retrieve pLDDT for chain `M` for every prediction:

```python
>>> ensemble_obj.get_all_plddt("M")
{
    "seed-1_sample-3": np.float64(78.71985788561527),
    "seed-2_sample-4": np.float64(77.89603812824957),
    "seed-1_sample-4": np.float64(78.20608318890814),
    ...
}
```

To calculate a single value aggregated across the ensemble:

```python
>>> ensemble_obj.get_ensemble_plddt("M")
np.float64(78.2713937608319)
```

By default, the ensemble-level aggregation is the mean:

```python
>>> ensemble_obj.get_ensemble_plddt(
...     "M",
...     ensemble_agg=max,
... )
np.float64(78.91594800693241)
```

### Matrix aggregation vs. ensemble aggregation

There are two separate levels of aggregation:

1. **Matrix aggregation (`agg`)** — how residue-level values within an individual prediction are summarized
2. **Ensemble aggregation (`ensemble_agg`)** — how values from individual predictions are summarized across the ensemble

For example:

```python
# Maximum value within each pLDDT matrix,
# followed by the mean across predictions
>>> ensemble_obj.get_ensemble_plddt(
...     "M",
...     agg=max,
...     ensemble_agg=np.mean,
... )
np.float64(98.647)
```

versus:

```python
# Mean within each pLDDT matrix,
# followed by the maximum across predictions
>>> ensemble_obj.get_ensemble_plddt(
...     "M",
...     agg=np.mean,
...     ensemble_agg=max,
... )
np.float64(78.91594800693241)
```

These operations are intentionally separate, allowing flexible analysis of AF3 ensembles.

---

# Supported Metrics

foldkit currently provides access to:

1. **pLDDT**
2. **pAE and iPAE**
3. **pTM and ipTM**
4. **Contact probabilities**
5. **ipSAE**

See the [IPSAE implementation](https://github.com/DunbrackLab/IPSAE) for additional information about ipSAE.

---

# Compressed foldkit Format

FoldKit can convert raw AF3 output into a substantially more space-efficient representation.

## Loading a compressed single result

Suppose you have exported the example above using the FoldKit CLI and now have:

```text
tutorial_example/single_result_export/
```

The directory contains a compressed `.npz` file instead of the original confidence JSON files.

You can load it directly:

```python
>>> foldkit.AF3Result.load_compressed_result(
...     "tutorial_example/single_result_export"
... )
<foldkit.af3_result.AF3Result object at ...>
```

The resulting object has the same interface as an `AF3Result` loaded directly from the original AF3 output.

---

## Loading a compressed ensemble

Compressed ensembles can be loaded in the same way:

```python
>>> foldkit.AF3Ensemble.load_compressed_result(
...     "tutorial_example/ensemble_result_export"
... )
<foldkit.af3_ensemble.AF3Ensemble object at ...>
```

This provides the same ensemble-level interface while avoiding the need to retain the original large JSON confidence files.

---

# Loading AF3 Server Results

AF3 Server results have a slightly different directory structure from locally generated AF3 results.

FoldKit can load these results using a separate function:

For example:

```python
foldkit.AF3Ensemble.load_webserver_result('tutorial_example/server')
```

---

# Command-Line Interface for Exporting

foldkit provides a CLI for converting AF3 output directories into the compressed FoldKit format.

Get help with:

```bash
foldkit -h
```

The main commands are:

| Command | Description |
|---|---|
| `export-single-result` | Export one AF3 prediction |
| `export-ensemble-result` | Export an ensemble of predictions |
| `webserver-export` | Export AF3 Server results |
| `batch-export` | Export multiple ensembles |

The general workflow is:

```text
Raw AF3 output
      │
      ▼
    foldkit
      │
      ▼
Compressed FoldKit output
      │
      ▼
Load with AF3Result / AF3Ensemble
```

After successfully exporting a result, the original AF3 output directory can be safely deleted if it is no longer needed.

---

## 1. Export a single AF3 prediction

Use `export-single-result` for one prediction corresponding to a single seed/sample.

```bash
foldkit export-single-result \
    <input_directory> \
    <output_directory>
```

For example:

```bash
foldkit -v export-single-result \
    tutorial_example/single_result \
    tutorial_example/single_result_export
```

Output:

```text
✅ Exported Data to : tutorial_example/single_result_export
```

---

## 2. Export an AF3 ensemble

Use `export-ensemble-result` for a directory containing multiple predictions of the same complex across seeds and/or samples.

```bash
foldkit export-ensemble-result \
    <input_directory> \
    <output_directory>
```

For example:

```bash
foldkit -v export-ensemble-result \
    tutorial_example/ensemble_result \
    tutorial_example/ensemble_result_export
```

This exports each prediction independently while preserving the ensemble directory structure:

```text
ensemble_result_export/
├── seed-1_sample-0/
├── seed-1_sample-1/
├── ...
├── seed-5_sample-0/
├── ...
└── seed-10_sample-4/
```

---

## 3. Export AF3 Server results

Use `webserver-export` for AF3 Server output.

```bash
foldkit -v webserver-export \
    tutorial_example/server \
    tutorial_example/server_export
```

The resulting directory can then be loaded using foldkit's regular compressed-result ensemble interface.

---

## 4. Batch export multiple ensembles

Use `batch-export` when you have a directory containing many AF3 ensemble directories.

For example:

```text
af3_results/
├── complex_1/
│   ├── seed-1_sample-0/
│   ├── seed-1_sample-1/
│   └── ...
├── complex_2/
│   ├── seed-1_sample-0/
│   ├── seed-1_sample-1/
│   └── ...
└── ...
```

Run:

```bash
foldkit batch-export \
    <input_directory> \
    <output_directory>
```

This is useful for large-scale co-folding campaigns.

---

# Storage Efficiency

The primary motivation for FoldKit's storage format is the substantial amount of disk space consumed by AF3 confidence JSON files.

As an initial benchmark, a single AF3 output directory for a four-chain complex occupies approximately:

| Format | Storage |
|---|---:|
| Raw AF3 | ~7.8 MB |
| FoldKit | ~1.9 MB |

This corresponds to approximately a **4× reduction in storage** for a single prediction.

The savings become much more substantial for large co-folding campaigns.

For example, consider a dataset containing:

- ~1,000 complexes
- 4 seeds per complex
- 5 samples per seed
- 20,000 total predictions

The total storage requirement is approximately:

```text
Raw AF3       157 GB
FoldKit        38 GB
              ───────
Savings       119 GB
```

The storage advantage becomes increasingly important as both **ensemble size** and **dataset size** increase.

---

# Contributing

Clone the repository and install the development dependencies:

```bash
pip install -e ".[dev]"
```

Run the test suite:

```bash
PYTHONPATH=src python -m pytest tests/ -vvv
```

## Build the package

```bash
python -m build
```

## Publish to PyPI

```bash
pip install --upgrade build twine
python -m build
twine check dist/*
twine upload dist/* -u __token__ -p <API TOKEN>
```

---

# Building the Documentation

FoldKit documentation is built using Sphinx and deployed from main through the `gh-pages` branch.

### 1. Build the HTML documentation to test it

```bash
cd docs
make html
cd ..
```


### 2. Commit the changes

```bash
git add .
git commit -m "Update docs"
```

### 3. Push to GitHub

```bash
git push
```

### 3. Check Actions tab to make sure deploy was successful
