---
license: mit
library_name: pytorch
tags:
  - materials-science
  - crystal-generation
  - high-entropy-alloys
  - diffusion-model
  - pytorch-lightning
---

# HEAForge v0.1.0

HEAForge is a multi-property-guided, crystal-aware generative framework for
100-atom high-entropy alloy (HEA) structures. The model represents bounded
relaxations around explicit FCC, BCC, and HCP parent sites and writes final
periodic structures as CIF and VASP CONTCAR files.

Code: https://github.com/AgnonZhang/HEAForge

## Model Bundle

| File | Role | Epoch | Global step | Size |
| --- | --- | ---: | ---: | ---: |
| `models/vae/heaforge_vae_v9.ckpt` | Parent-site residual VAE | 253 | 10,160 | 614,505,820 bytes |
| `models/predictor/heaforge_predictor_v9.ckpt` | Sixteen-property predictor | 40 | 410 | 453,752 bytes |
| `models/ldm/heaforge_ldm_v9.ckpt` | Crystal-conditioned latent diffusion model | 144 | 5,800 | 115,964,558 bytes |

SHA-256 values are provided in `checksums.sha256` and structured metadata is
provided in `manifest.yaml`.

## Training Data

The final crystal-conditioned stage used 790 valid HEA structures, split with
parent-lattice stratification:

- Train: 632
- Validation: 79
- Test: 79

Parent classes are FCC, BCC, and HCP. Every final training structure contains
100 sites. Raw structures and physical-property tables are not distributed in
this model repository.

## Conditioning

The diffusion model uses sixteen physical-property conditions:

1. Magnetic moment per atom
2. Band gap
3. d-band center
4. d-band width
5. Fermi level
6. DOS at the Fermi level
7. Magnetic susceptibility
8. Conductivity divided by relaxation time
9. Seebeck coefficient
10. Hall coefficient
11. Mixing enthalpy
12. Valence electrons per unit cell
13. Electronic thermal conductivity divided by relaxation time
14. Heat capacity
15. Lattice volume distortion
16. Lattice angle distortion

It also conditions on the parent lattice, six cell parameters, and six strain
components. Formation energy is not a target in this release.

## Download

```bash
hf download Aegon567/HEAForge --local-dir checkpoints/heaforge
```

## Inference

Inference requires a locally prepared HEA condition CSV. See the code
repository for data preparation details.

```bash
python scripts/3d_mag/generate_direct_contcars.py \
  --checkpoint checkpoints/heaforge/models/ldm/heaforge_ldm_v9.ckpt \
  --vae-checkpoint checkpoints/heaforge/models/vae/heaforge_vae_v9.ckpt \
  --condition-csv data/3d-mag-crystal/train.csv \
  --total-samples 20 \
  --property-condition-mode sampled-row \
  --crystal-condition-mode sampled \
  --output-dir outputs/heaforge_samples
```

## Intended Use

The bundle is intended for HEA crystal candidate generation, method
development, and research screening. Generated structures should be subjected
to independent geometry relaxation, stability analysis, and electronic-
structure calculations before downstream scientific use.

## Limitations

- CIF labels are requested conditioning targets, not DFT-verified properties
  of the generated structure.
- The current model is specialized to 100-atom FCC/BCC/HCP-derived HEA
  structures and should not be treated as a general inorganic crystal model.
- A valid periodic geometry does not establish thermodynamic or dynamic
  stability.
- Chemical composition cardinality is not guaranteed unless an explicit
  composition constraint or post-generation filter is applied.
- Predictions outside the empirical training-property range are not validated.

## Origin and License

HEAForge is derived from the MIT-licensed
[Chemeleon2](https://github.com/hspark1212/chemeleon2) framework and retains
its VAE, latent diffusion, and reinforcement-learning foundations. HEAForge
adds HEA-specific data processing, multi-property conditioning, parent-site
residual representations, crystal geometry objectives, and direct
CIF/CONTCAR generation.