ml-generative-adit

Files (13)

• SKILL.md
• README.md · examples/molecules/
• generation_metadata.json · examples/molecules/
• molecule_0000.xyz · examples/molecules/
• molecule_0001.xyz · examples/molecules/
• molecule_0002.xyz · examples/molecules/
• molecule_0003.xyz · examples/molecules/
• molecule_0004.xyz · examples/molecules/
• molecule_0005.xyz · examples/molecules/
• molecule_0006.xyz · examples/molecules/
• molecule_0007.xyz · examples/molecules/
• molecule_0008.xyz · examples/molecules/
• molecule_0009.xyz · examples/molecules/

SKILL.md (first lines)

# ADiT Structure Generation Skill

## Goal

Generate novel crystal structures and molecules using ADiT (All-atom Diffusion Transformers, ICML 2025), a unified latent diffusion framework from Meta FAIR Chemistry that jointly generates both periodic materials and non-periodic molecular systems from a shared latent space.

## 1. Prerequisites

> [!IMPORTANT]
> **GPU Required**: ADiT requires a CUDA-compatible GPU. CPU inference is extremely slow.

- Runs as the `adit` MCP server and its scripts run in the `adit` environment: on x86_64 a uv environment created on first use (CUDA 12.6 or 13 by driver), on aarch64 the `generative` container image.
- The AADT repository cloned next to this project as `../adit`, or anywhere with `ADIT_REPO` pointing to it, at the verified commit. `venv/run` mounts it into the container on aarch64.

  ```bash
  git clone https://github.com/facebookresearch/all-atom-diffusion-transformer ${CLAUDE_SKILL_DIR}/../../../adit
  git -C ${CLAUDE_SKILL_DIR}/../../../adit checkout b9ce505f170597a7c8ca50d13ce8e15df21cf8c9
  ```
- Pre-trained weights are automatically downloaded from HuggingFace on first use.

## 2. Available Models

ADiT provides a joint pre-trained model trained on:

- **MP20**: Materials Project 2020 dataset (inorganic crystals, ~45K structures)
- **QM9**: Small organic molecules (~134K molecules)

The single checkpoint handles both crystal and molecule generation, selected via the `generation_type` parameter.

## 3. MCP Tool Usage

### Crystal Generation

Generate novel periodic crystal structures (saved as CIF files):

```bash
adit.generate_structures(
    generation_type="crystals",    # Generate periodic crystals
    num_structures=10,             # Number of structures to generate
    batch_size=100,                # Batch size for GPU efficiency
    cfg_scale=2.0,                 # Classifier-free guidance scale
    output_dir="research/my_project/crystals"
)
```

### Molecule Generation

Generate novel non-periodic molecules (saved as XYZ files):

```bash
adit.generate_structures(
    generation_type="molecules",   # Generate molecules
    num_structures=10,
    batch_size=100,
    cfg_scale=2.0,
    output_dir="research/my_project/molecules"
)
```

## 4. Parameters

| Parameter | Default | Description |
|-----------|---------|-------------|
| `generation_type` | `"crystals"` | `"crystals"` for periodic structures (CIF), `"molecules"` for non-periodic (XYZ) |
| `num_structures` | `10` | Total number of structures to generate |
| `batch_size` | `100` | Batch size (larger = faster on GPU) |
| `cfg_scale` | `2.0` | Classifier-free guidance scale. Higher = more typical but less diverse |
| `device` | `"auto"` | Device: `"auto"`, `"cpu"`, or `"cuda"` |
| `output_dir` | auto | Output directory. Auto-creates under research dir |

## 5. Output Files

### Crystal Generation
- `crystal_XXXX.cif`: Generated crystal structure files (pymatgen CIF format)
- `generation_metadata.json`: Generation parameters and statistics

### Molecule Generation
- `molecule_XXXX.xyz`: Generated molecule files (ASE XYZ format)
- `generation_metadata.json`: Generation parameters and statistics

## 6. Limitations

> [!WARNING]
> **No Conditional Generation**: The public checkpoint is unconditional only — you cannot
> condition on specific compositions, space groups, or properties.
> To get specific compositions: generate many structures and filter.

> [!WARNING]
> **No Fine-Tuning via MCP**: Fine-tuning requires the full AADT training pipeline
> with multi-GPU setup and wandb logging. Use the raw codebase for training.

> [!NOTE]
> **Atom Count Distribution**: The number of atoms per generated structure is sampled from
> the training dataset distribution. For crystals (MP20), this peaks around 8-20 atoms.
> For molecules (QM9), this peaks around 18 atoms including hydrogens.

## 7. Best Practices

> [!TIP]
> - **Start with crystals**: Crystal generation on MP20 tends to produce more valid structures
> - **Guidance scale**: Use 2.0 (default) for balanced diversity/quality. Increase to 3.0-4.0 for more "typical" structures
> - **Validate outputs**: Always validate generated structures via relaxation and stability analysis
> - **Batch size**: Use batch_size=100 for best GPU throughput

## 8. Workflow Integration

ADiT works well in combination with:
- **Structure relaxation**: Use MLIP tools (MACE, FairChem, MatGL) to optimize generated structures
- **Stability analysis**: Calculate E_hull to identify thermodynamically stable phases
- **Property prediction**: Use MLIPs or DFT to calculate properties of generated structures
- **Comparison with MatterGen**: Generate structures with both ADiT and MatterGen for diversity

## 9. Architecture

ADiT uses a two-stage latent diffusion approach:
1. **VAE Autoencoder**: Maps all-atom representations (atoms, coords, lattice) to a shared latent space
2. **DiT Denoiser**: Trained via flow matching to generate new latent embeddings
3. **Decoder**: Converts latent embeddings back to atomic structures

From learningmatter-mit/AtomisticSkills, MIT.

More from learningmatter-mit/AtomisticSkills