Metadata-Version: 2.4
Name: atom-jepa
Version: 0.3.0
Summary: Atom-JEPA: Joint-Embedding Predictive Architecture for 3D Atomistic Systems
Author: The Atom-JEPA authors
License-Expression: MIT
Project-URL: Homepage, https://github.com/khelverskovp/atom-jepa
Project-URL: Repository, https://github.com/khelverskovp/atom-jepa
Project-URL: Models, https://huggingface.co/atom-jepa/atom-jepa
Keywords: self-supervised learning,JEPA,equivariant neural networks,molecules,crystals,representation learning
Classifier: Programming Language :: Python :: 3
Classifier: Intended Audience :: Science/Research
Classifier: Topic :: Scientific/Engineering :: Chemistry
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: atom_jepa/models/equiformer_v3/LICENSE
License-File: atom_jepa/models/equiformer_v3/LICENSE-fairchem
Requires-Dist: torch>=2.1
Requires-Dist: numpy
Requires-Dist: e3nn
Requires-Dist: torch_geometric
Requires-Dist: ase
Requires-Dist: huggingface_hub
Requires-Dist: safetensors
Provides-Extra: rdkit
Requires-Dist: rdkit; extra == "rdkit"
Provides-Extra: pymatgen
Requires-Dist: pymatgen; extra == "pymatgen"
Provides-Extra: all
Requires-Dist: rdkit; extra == "all"
Requires-Dist: pymatgen; extra == "all"
Provides-Extra: cu12
Requires-Dist: cuequivariance<0.13,>=0.10; extra == "cu12"
Requires-Dist: cuequivariance-torch<0.13,>=0.10; extra == "cu12"
Requires-Dist: cuequivariance-ops-cu12<0.13,>=0.10; extra == "cu12"
Requires-Dist: cuequivariance-ops-torch-cu12<0.13,>=0.10; extra == "cu12"
Provides-Extra: cu13
Requires-Dist: cuequivariance<0.13,>=0.10; extra == "cu13"
Requires-Dist: cuequivariance-torch<0.13,>=0.10; extra == "cu13"
Requires-Dist: cuequivariance-ops-cu13<0.13,>=0.10; extra == "cu13"
Requires-Dist: cuequivariance-ops-torch-cu13<0.13,>=0.10; extra == "cu13"
Dynamic: license-file

# Atom-JEPA

Embeddings from the pretrained encoders of **Atom-JEPA: Joint-Embedding Predictive
Architecture for 3D Atomistic Systems**.

Atom-JEPA is a self-supervised pretraining method for 3D atomistic systems. It learns
by predicting the latent representations of one part of a structure from another, and
the same recipe applies to both molecules and crystals. This package loads the
pretrained encoders and turns 3D structures into embeddings. Pretraining and
fine-tuning code is in the [GitHub repository](https://github.com/khelverskovp/atom-jepa).

## Installation

Install PyTorch first ([pytorch.org](https://pytorch.org/get-started/locally/)), then:

```bash
pip install atom-jepa            # structures as arrays or ase.Atoms
pip install "atom-jepa[rdkit]"   # + SMILES input
pip install "atom-jepa[all]"     # + SMILES and pymatgen input
pip install "atom-jepa[cu12]"    # + cuEquivariance GPU kernels for CUDA 12 ([cu13] for CUDA 13)
```

Extras combine, e.g. `"atom-jepa[all,cu12]"`; `python -m atom_jepa.info` shows what your
machine supports and which kernel extra matches your PyTorch.

## Usage

```python
from atom_jepa import AtomJEPA

model = AtomJEPA.from_pretrained("molecules")    # or "crystals"

emb = model.embed("CC(=O)Oc1ccccc1C(=O)O")       # one structure -> [256]
embs = model.embed(["CCO", "c1ccncc1"])          # a list -> [2, 256]
```

The weights (`molecules`, `crystals`) are downloaded from
[huggingface.co/atom-jepa/atom-jepa](https://huggingface.co/atom-jepa/atom-jepa) on
first use and cached locally. A path to a local checkpoint also works.

`embed` accepts any of these as one structure, or a list of them:

| input | notes |
|---|---|
| SMILES string | one RDKit ETKDGv3 + MMFF94 conformer, hydrogens added (needs `[rdkit]`) |
| `ase.Atoms` | periodic if any of its `pbc` flags is set |
| pymatgen `Structure` / `Molecule` | crystals / molecules (needs `[pymatgen]`) |
| `(atomic_numbers, positions[, cell])` | positions in Angstrom; cell with lattice vectors as rows |
| `{"atomic_numbers", "node_coordinates"[, "cell"]}` | the same as a dict |

For several conformers per molecule, use
`atom_jepa.conformers.smiles_to_samples(smiles, num_conformers=10)` and embed the list.

### Feature options

```python
model.embed(smiles, layers="all")                    # [S, 8, 256]     output of each of the 8 blocks
model.embed(smiles, layers=[4, 8])                   # [S, 2, 256]
model.embed(smiles, degrees="all")                   # [S, 9, 256]     l=0, l=1 (3), l=2 (5) components
model.embed(smiles, degrees="all", invariant=True)   # [S, 3, 256]     l=0 and per-channel norms of l=1, l=2
model.embed(smiles, layers="all", degrees="all")     # [S, 8, 9, 256]
model.embed(smiles, per_atom=True)                   # list of [n_atoms, 256]
```

- `layers="last"` (default) is the last block after the encoder's final norm; block
  numbers 1-8 and `"all"` are the block outputs before it.
- Structure features are the mean over atoms. The l>0 components are equivariant: l=1
  rotates like an (x, y, z) vector. `invariant=True` makes them rotation invariant.

### Fine-tuning

`AtomJEPA` is a `torch.nn.Module`: calling it on a batch returns differentiable features
(with the same options), so you can train your own head and fine-tune the encoder:

```python
import torch
from torch.utils.data import DataLoader
from atom_jepa import AtomJEPA, to_sample

model = AtomJEPA.from_pretrained("molecules").train()
head = torch.nn.Linear(model.embedding_dim, 1).to(model.device)

data = [{**to_sample(s), "y": torch.tensor([y])} for s, y in zip(train_smiles, train_y)]
loader = DataLoader(data, batch_size=32, shuffle=True, collate_fn=model.collate)
opt = torch.optim.AdamW([{"params": model.parameters(), "lr": 1e-5},
                         {"params": head.parameters(), "lr": 1e-3}])
for batch in loader:
    loss = torch.nn.functional.mse_loss(head(model(batch)), batch["y"].to(model.device))
    opt.zero_grad(); loss.backward(); opt.step()
```

`model.set_grad_checkpointing(True)` trades compute for memory on large structures. The
paper's benchmark recipes are the fine-tuning scripts in the
[GitHub repository](https://github.com/khelverskovp/atom-jepa#fine-tuning).

### Execution

On a CUDA GPU the encoder runs in bf16 and, if the kernels are installed, with
cuEquivariance; otherwise it runs in plain fp32 PyTorch. The settings are printed when
loading and can be set explicitly:

```python
model = AtomJEPA.from_pretrained("molecules")                        # auto
# [atom-jepa] molecules: device=cuda, precision=bf16, cuequivariance=on, compile=off
#   we recommend compile=True for larger jobs
model = AtomJEPA.from_pretrained("molecules", compile=True)          # compile the blocks (CUDA)
model = AtomJEPA.from_pretrained("molecules", precision="fp32", cuequivariance=False)  # plain
```

bf16 changes the embeddings slightly; use `precision="fp32"` for exact
fp32 features. Compiling adds a one-off cost on the first batches.

## License

The code is released under the MIT License. The pretrained models are released under
CC BY 4.0.
