Metadata-Version: 2.5
Name: zatom-nucleus
Version: 1.0.3
Summary: Shared utilities and training infrastructure for atomistic foundation models.
Project-URL: Homepage, https://zatomai.github.io/nucleus
Project-URL: Repository, https://github.com/Zatom-AI/nucleus
Project-URL: Issues, https://github.com/Zatom-AI/nucleus/issues
Author-email: Zatom-AI <acmwhb@lbl.gov>
License: Nucleus Copyright (c) 2026, The Regents of the University of California,
        through Lawrence Berkeley National Laboratory (subject to receipt of
        any required approvals from the U.S. Dept. of Energy). All rights reserved.
        
        Redistribution and use in source and binary forms, with or without
        modification, are permitted provided that the following conditions are met:
        
        (1) Redistributions of source code must retain the above copyright notice,
        this list of conditions and the following disclaimer.
        
        (2) Redistributions in binary form must reproduce the above copyright
        notice, this list of conditions and the following disclaimer in the
        documentation and/or other materials provided with the distribution.
        
        (3) Neither the name of the University of California, Lawrence Berkeley
        National Laboratory, U.S. Dept. of Energy nor the names of its contributors
        may be used to endorse or promote products derived from this software
        without specific prior written permission.
        
        
        THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
        ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
        CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
        
        You are under no obligation whatsoever to provide any bug fixes, patches,
        or upgrades to the features, functionality or performance of the source
        code ("Enhancements") to anyone; however, if you choose to make your
        Enhancements available either publicly, or directly to Lawrence Berkeley
        National Laboratory, without imposing a separate written license agreement
        for such Enhancements, then you hereby grant the following license: a
        non-exclusive, royalty-free perpetual license to install, use, modify,
        prepare derivative works, incorporate into other computer software,
        distribute, and sublicense such enhancements or derivative works thereof,
        in binary and source code form.
License-File: LICENSE.md
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: BSD License
Classifier: Natural Language :: English
Classifier: Operating System :: MacOS
Classifier: Operating System :: Microsoft :: Windows
Classifier: Operating System :: POSIX :: Linux
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: Implementation :: CPython
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: <3.13,>=3.12
Requires-Dist: atomworks[ase,ml,openbabel]<3.0.0,>=2.2.1
Requires-Dist: beartype<1,>=0.18.0
Requires-Dist: cuequivariance-ops-cu13==0.10.0; sys_platform == 'linux'
Requires-Dist: cuequivariance-ops-torch-cu13==0.10.0; sys_platform == 'linux'
Requires-Dist: cuequivariance-torch==0.10.0; sys_platform == 'linux'
Requires-Dist: cuequivariance==0.10.0; sys_platform == 'linux'
Requires-Dist: dm-tree<1,>=0.1.6
Requires-Dist: einops<1,>=0.8.0
Requires-Dist: einx<1,>=0.1.0
Requires-Dist: environs<12,>=11.0.0
Requires-Dist: huggingface-hub<2,>=1.0.0
Requires-Dist: hydra-colorlog>=1.2.0
Requires-Dist: hydra-core<1.4,>=1.3.0
Requires-Dist: jaxtyping<1,>=0.2.17
Requires-Dist: lightning>=2.5.0
Requires-Dist: loralib>=0.1.1
Requires-Dist: opt-einsum<4,>=3.4.0
Requires-Dist: pandas
Requires-Dist: pyarrow
Requires-Dist: pydantic
Requires-Dist: rc-foundry<0.3.0,>=0.2.0
Requires-Dist: rich>=13.9.4
Requires-Dist: rootutils<1.1,>=1.0.7
Requires-Dist: toolz
Requires-Dist: torch==2.13.0
Requires-Dist: torchdata
Requires-Dist: typer<1,>=0.20.0
Requires-Dist: wandb<1,>=0.15.10
Requires-Dist: zstandard
Provides-Extra: all
Requires-Dist: accelerate; extra == 'all'
Requires-Dist: biopython; extra == 'all'
Requires-Dist: biotite; extra == 'all'
Requires-Dist: fpsim2; extra == 'all'
Requires-Dist: huggingface-hub; extra == 'all'
Requires-Dist: matminer>=0.10; extra == 'all'
Requires-Dist: posebusters; extra == 'all'
Requires-Dist: prody; extra == 'all'
Requires-Dist: pydantic>=2.8; extra == 'all'
Requires-Dist: pymatgen; extra == 'all'
Requires-Dist: rdkit; extra == 'all'
Requires-Dist: riboseek==0.1.0; extra == 'all'
Requires-Dist: smact; extra == 'all'
Requires-Dist: torch-geometric; extra == 'all'
Requires-Dist: tqdm; extra == 'all'
Requires-Dist: transformers>=4.51; extra == 'all'
Provides-Extra: dev
Requires-Dist: accelerate; extra == 'dev'
Requires-Dist: atomworks[ase,dev,ml,openbabel]<3.0.0,>=2.2.1; extra == 'dev'
Requires-Dist: biopython; extra == 'dev'
Requires-Dist: biotite; extra == 'dev'
Requires-Dist: debugpy<2,>=1.8.5; extra == 'dev'
Requires-Dist: fpsim2; extra == 'dev'
Requires-Dist: huggingface-hub; extra == 'dev'
Requires-Dist: ipdb; extra == 'dev'
Requires-Dist: ipykernel<7,>=6.29.4; extra == 'dev'
Requires-Dist: matminer>=0.10; extra == 'dev'
Requires-Dist: posebusters; extra == 'dev'
Requires-Dist: pre-commit; extra == 'dev'
Requires-Dist: prody; extra == 'dev'
Requires-Dist: pydantic>=2.8; extra == 'dev'
Requires-Dist: pymatgen; extra == 'dev'
Requires-Dist: pytest-benchmark<6,>=5.0.0; extra == 'dev'
Requires-Dist: pytest-cov<5,>=4.1.0; extra == 'dev'
Requires-Dist: pytest-dotenv<1,>=0.5.2; extra == 'dev'
Requires-Dist: pytest-testmon<3,>=2.1.1; extra == 'dev'
Requires-Dist: pytest-xdist<4,>=3.6.1; extra == 'dev'
Requires-Dist: pytest<9,>=8.2.0; extra == 'dev'
Requires-Dist: rdkit; extra == 'dev'
Requires-Dist: riboseek==0.1.0; extra == 'dev'
Requires-Dist: ruff==0.8.3; extra == 'dev'
Requires-Dist: smact; extra == 'dev'
Requires-Dist: torch-geometric; extra == 'dev'
Requires-Dist: tqdm; extra == 'dev'
Requires-Dist: transformers>=4.51; extra == 'dev'
Provides-Extra: zatom2
Requires-Dist: pydantic>=2.8; extra == 'zatom2'
Provides-Extra: zevals
Requires-Dist: accelerate; extra == 'zevals'
Requires-Dist: biopython; extra == 'zevals'
Requires-Dist: biotite; extra == 'zevals'
Requires-Dist: fpsim2; extra == 'zevals'
Requires-Dist: huggingface-hub; extra == 'zevals'
Requires-Dist: matminer>=0.10; extra == 'zevals'
Requires-Dist: posebusters; extra == 'zevals'
Requires-Dist: prody; extra == 'zevals'
Requires-Dist: pydantic>=2.8; extra == 'zevals'
Requires-Dist: pymatgen; extra == 'zevals'
Requires-Dist: rdkit; extra == 'zevals'
Requires-Dist: riboseek==0.1.0; extra == 'zevals'
Requires-Dist: smact; extra == 'zevals'
Requires-Dist: torch-geometric; extra == 'zevals'
Requires-Dist: tqdm; extra == 'zevals'
Requires-Dist: transformers>=4.51; extra == 'zevals'
Description-Content-Type: text/markdown

<p align="center">
  <img src="docs/assets/nucleus_header.jpg" alt="Nucleus — Atomistic Foundation Modeling across Molecules, Materials & Proteins" width="100%">
</p>

# Nucleus

Nucleus is the shared monorepo for **Zatom-2**, an atomistic foundation model for
molecules, materials, and protein generation, and **Zevals**, its evaluation
pipelines. It includes training code, paper configurations, data preparation
tools, the inference CLI, and Zatom-2's project page.

[Project page](https://zatomai.github.io/nucleus) ·
[Preprint PDF](models/zatom2/docs/static/zatom2_paper.pdf) ·
[arXiv (placeholder; submission pending)](https://arxiv.org/abs/XXXX.XXXXX) ·
[Weights](https://huggingface.co/Zatom-AI/Zatom-2) ·
[PyPI](https://pypi.org/project/zatom-nucleus/)

## Install

### From PyPI

`zatom-nucleus` and `zatom2` can be installed using [uv](https://docs.astral.sh/uv/) (recommended) or `pip`:

```bash
uv venv
source .venv/bin/activate
uv pip install zatom-nucleus zatom2
```

*(or `uv pip install "zatom-nucleus[all]"` to include `zevals`)*

> [!TIP]
> When managing environments with `uv venv`, install packages with `uv pip install` rather than bare `pip install`. Because `uv venv` does not install `pip` inside the virtual environment by default, running `pip` directly can invoke an outer Python or Conda environment on your `PATH` and cause dynamic library conflicts (such as duplicate OpenMP runtime initializations on macOS). Alternatively, create your virtual environment with `uv venv --seed` if you require standard `pip`.

### From source

Use Python 3.12 and [uv](https://docs.astral.sh/uv/). The paper environment uses
Linux, NVIDIA A100 80 GB GPUs, PyTorch 2.13, and CUDA 13.2. A compatible NVIDIA
driver is required for GPU execution. The lockfile pins public AtomWorks and
Foundry revisions.

```bash
git clone https://github.com/Zatom-AI/nucleus.git
cd nucleus
uv sync --locked --extra all
source .venv/bin/activate
```

One root installation provides `nucleus`, `zatom2`, and `zevals`; separate
model-package installs are unnecessary. For training, copy `.env.example` to
`.env` and follow the [data preparation guide](docs/data.md). Inference on the
included molecule example does not require training datasets.

## Download weights and generate a molecule

```bash
nucleus list-available
nucleus install zatom2
zatom2 design inputs=models/zatom2/docs/examples/demo.json out_dir=logs/demo
```

`zatom2` selects the checksummed large joint-property model (configuration VI).
The downloader verifies SHA-256 and stores weights in `~/.nucleus/checkpoints`.
Use a paper identifier from `nucleus list-available` to select another model:

```bash
nucleus install zatom2-table3-scope-from-joint-properties-force-conditioned
zatom2 design \
  ckpt_path=zatom2-table3-scope-from-joint-properties-force-conditioned \
  inputs=models/zatom2/docs/examples/protein.json out_dir=logs/protein
```

See [inference](docs/inference.md) for force conditioning, custom checkpoint
directories, output formats, and evaluation. Weight availability is tracked in
the [release status](docs/release_status.md).

## Reproduce the paper

- [Paper-to-command guide](docs/reproducibility.md): pretraining, ablations,
  scaling, generation, structure/property prediction, and protein transfer.
- [Data preparation](docs/data.md): sources, licenses, frozen splits, LMDB
  loading, molecular bonds, and evaluation references.
- [Zevals](models/zevals/README.md): metrics and separate external-oracle setup.
- [Preprint source](models/zatom2/docs/manuscript/arxiv_preprint.tex) and
  [project page source](models/zatom2/docs/index.html).

## Repository layout

| Path | Contents |
| --- | --- |
| `src/nucleus/` | Shared architecture, trainers, transforms, metrics, and loaders |
| `src/nucleus_cli/` | Weight discovery and verified downloads |
| `models/zatom2/` | Zatom-2, Hydra configurations, examples, paper, and page |
| `models/zevals/` | Zevals pipelines, metrics, and oracle adapters |
| `scripts/` | Portable data preparation and paper analyses |
| `data/` | Small reproducibility artifacts and frozen selection indices |
| `tests/` | Regression tests for released functionality |

Code is distributed under the [Nucleus license](LICENSE.md). Dataset, model,
oracle, and webpage licenses are described in
[third-party notices](THIRD_PARTY_NOTICES.md). Citation metadata is in
[CITATION.cff](CITATION.cff); the arXiv identifier will be updated after submission.
