Metadata-Version: 2.4
Name: firefly-activemd
Version: 0.3.2
Summary: Active Learning Molecular Dynamics with MACE potentials and DFT reference calculations
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: ase>=3.28
Requires-Dist: mace-torch>=0.3
Requires-Dist: numpy>=2
Requires-Dist: sisl>=0.16

# Firefly

Active Learning Molecular Dynamics (ALMD) with MACE machine learning potentials and on-the-fly DFT reference calculations. Firefly runs an MD simulation using a committee of MACE models and automatically triggers DFT calculations and model retraining whenever the ensemble uncertainty exceeds a defined threshold.

---

## How it works

At every MD step, Firefly evaluates the disagreement among the MACE committee (standard deviation of energies or forces). If the uncertainty exceeds `thresh`, it enters a patience window. If uncertainty persists beyond `patience` steps — or exceeds `3 × thresh` — it:

1. Pauses the dynamics
2. Runs a single-point DFT calculation (SIESTA or Quantum ESPRESSO)
3. Appends the result to `train.xyz`
4. Retrains all committee models in parallel across GPUs
5. Reloads the updated models and resumes the simulation

---

## Installation

### 1. Install PyTorch (required first)

Install PyTorch with the CUDA version that matches your drivers. For CUDA 12.6:

```bash
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu126
```

For other CUDA versions or CPU-only, see https://pytorch.org/get-started/locally/

### 2. Install Firefly

**From PyPI:**

```bash
pip install firefly-activemd
```

**From source:**

```bash
pip install -e .
```

> PyTorch must be installed **before** Firefly regardless of the installation method, as it is not pulled automatically with the correct CUDA build.

---

## Usage

```bash
firefly parameters.json
```

---

## Configuration file (parameters.json)

### Mandatory

| Key | Type | Description |
|-----|------|-------------|
| `file` | string | Path to the atomic structure file (`.xyz`, `.cif`, etc.) |
| `ensemble` | string | MD ensemble type (see [Ensembles](#ensembles)) |
| `code` | string | DFT reference code: `"siesta"` or `"qe"` |
| `cmd` | string | Command used to execute the DFT code |
| `thresh` | float | Uncertainty threshold that triggers a DFT call |
| `n_models` | int | Number of MACE models in the committee |
| `atoms` | dict | Element symbol → atomic number mapping **(SIESTA only)** |

### Optional

| Key | Type | Default | Description |
|-----|------|---------|-------------|
| `gpus` | int or list | `1` | Number of GPUs or list of GPU IDs (e.g. `[0,1,2,3]`) |
| `T_begin` | float | `300.0` | Initial temperature (K) |
| `T_final` | float | `300.0` | Final temperature (K) — linear ramp from `T_begin` |
| `pressure` | float | `0.0` | Target pressure (bar) — only used in NPT ensembles |
| `dt` | float | `0.5` | Time step (fs) |
| `time` | float | `100000.0` | Total simulation time (fs) |
| `patience` | int | `1000` | Steps to wait before forcing a DFT call after threshold is exceeded |
| `active` | string | `"energy"` | Uncertainty metric: `"energy"` or `"force"` |
| `epoch` | int | `10` | Training epochs added per retraining event |
| `premodel` | string | `"model"` | Prefix for checkpoint and model filenames |

### Example

```json
{
    "file": "estrutura.xyz",
    "ensemble": "nptz",
    "code": "siesta",
    "cmd": "siesta",
    "thresh": 0.05,
    "n_models": 4,
    "gpus": 4,
    "T_begin": 300.0,
    "T_final": 300.0,
    "pressure": 1.0,
    "dt": 0.5,
    "time": 100000.0,
    "patience": 1000,
    "atoms": {
        "H": 1,
        "O": 8
    }
}
```

---

## Ensembles

| Name | Description |
|------|-------------|
| `aniso` | Anisotropic NPT — cell shape and volume free (MTK barostat) |
| `iso` | Isotropic NPT — uniform scaling of all cell vectors (MTK barostat) |
| `nptz` | NPT along Z only — XY plane is frozen |
| `nptxy` | NPT along XY only — Z axis is frozen |
| `ortho` | Orthorhombic NPT — diagonal cell components free |
| `bussi` | NVT with Bussi–Donadio–Parrinello thermostat (no barostat) |

All NPT ensembles use the Nosé–Hoover chain (MTK) thermostat with:
- Thermostat damping: `100 × dt`
- Barostat damping: `1000 × dt`

Temperature ramps linearly from `T_begin` to `T_final` over the full simulation.

---

## Required files in the working directory

Before running, the working directory must contain:

| File | Description |
|------|-------------|
| `model1.model` … `modelN.model` | Trained MACE models (N = `n_models`) |
| `job_1.sh` … `job_N.sh` | Training scripts for each model |
| `checkpoints/` | MACE checkpoint directory |
| `RUN.fdf` | SIESTA input template **(SIESTA only)** — the SIESTA label inside this file must be set to `siesta` |
| `espresso.in` | QE input template **(QE only)** — the `prefix` inside this file must be set to `active` |

---

## Output files

| File | Description |
|------|-------------|
| `therm.dat` | Thermodynamic log: step, T (K), P (atm), Ekin, Epot, density, cell parameters |
| `run.traj` | ASE trajectory file written every 100 steps |
| `thresh.dat` | Log of steps where uncertainty exceeded the threshold |
| `train.xyz` | Accumulated DFT reference structures for retraining |

---

## Uncertainty threshold logic

```
Each MD step:
  std = std_dev(committee energies or forces)

  if std > thresh:
      if interval < patience and std < 3 × thresh:
          wait (patience mode)
      else:
          → run DFT
          → retrain all models in parallel
          → reload models
          → resume MD
  else:
      continue
```

The factor `3 × thresh` acts as a hard cutoff: if disagreement is extreme, patience is skipped and DFT is called immediately.

---

## Parallel training

Models are retrained in parallel, one per GPU. GPU assignment follows a round-robin scheme based on `gpus`. Training epochs and SWA checkpoints are updated automatically in the job scripts before each retraining event.
