firefly-activemd 0.3.0__tar.gz

firefly_activemd-0.3.0/PKG-INFO

@@ -0,0 +1,54 @@
+Metadata-Version: 2.4
+Name: firefly-activemd
+Version: 0.3.0
+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 with MACE potentials and DFT reference calculations (SIESTA or Quantum ESPRESSO).
+
+## Installation
+
+```bash
+pip install .
+```
+
+Or in editable mode (for development):
+
+```bash
+pip install -e .
+```
+
+## Usage
+
+```bash
+firefly parameters.json
+```
+
+## Parameters (parameters.json)
+
+| Key | Required | Description |
+|-----|----------|-------------|
+| `file` | yes | Path to the atomic structure file (.xyz, etc.) |
+| `ensemble` | yes | MD ensemble: `iso`, `aniso`, `nptz`, `nptxy`, `bussi`, `ortho` |
+| `code` | yes | DFT code: `siesta` or `qe` |
+| `cmd` | yes | Command to run the DFT code |
+| `thresh` | yes | Uncertainty threshold for active learning |
+| `n_models` | yes | Number of MACE models in the committee |
+| `atoms` | siesta only | Dict mapping element symbols to atomic numbers |
+| `gpus` | no | Number of GPUs (default: 1) |
+| `T_begin` | no | Initial temperature in K (default: 300) |
+| `T_final` | no | Final temperature in K (default: 300) |
+| `pressure` | no | Pressure in bar (default: 0) |
+| `dt` | no | Time step in fs (default: 0.5) |
+| `time` | no | Total simulation time in fs (default: 100000) |
+| `patience` | no | Steps before forcing DFT call (default: 1000) |
+| `active` | no | Uncertainty metric: `energy` or `force` (default: `energy`) |
+| `epoch` | no | Training epochs per retraining step (default: 10) |
+| `premodel` | no | Model name prefix (default: `model`) |

firefly_activemd-0.3.0/README.md

@@ -0,0 +1,43 @@
+# Firefly
+
+Active Learning Molecular Dynamics with MACE potentials and DFT reference calculations (SIESTA or Quantum ESPRESSO).
+
+## Installation
+
+```bash
+pip install .
+```
+
+Or in editable mode (for development):
+
+```bash
+pip install -e .
+```
+
+## Usage
+
+```bash
+firefly parameters.json
+```
+
+## Parameters (parameters.json)
+
+| Key | Required | Description |
+|-----|----------|-------------|
+| `file` | yes | Path to the atomic structure file (.xyz, etc.) |
+| `ensemble` | yes | MD ensemble: `iso`, `aniso`, `nptz`, `nptxy`, `bussi`, `ortho` |
+| `code` | yes | DFT code: `siesta` or `qe` |
+| `cmd` | yes | Command to run the DFT code |
+| `thresh` | yes | Uncertainty threshold for active learning |
+| `n_models` | yes | Number of MACE models in the committee |
+| `atoms` | siesta only | Dict mapping element symbols to atomic numbers |
+| `gpus` | no | Number of GPUs (default: 1) |
+| `T_begin` | no | Initial temperature in K (default: 300) |
+| `T_final` | no | Final temperature in K (default: 300) |
+| `pressure` | no | Pressure in bar (default: 0) |
+| `dt` | no | Time step in fs (default: 0.5) |
+| `time` | no | Total simulation time in fs (default: 100000) |
+| `patience` | no | Steps before forcing DFT call (default: 1000) |
+| `active` | no | Uncertainty metric: `energy` or `force` (default: `energy`) |
+| `epoch` | no | Training epochs per retraining step (default: 10) |
+| `premodel` | no | Model name prefix (default: `model`) |

firefly_activemd-0.3.0/firefly/__init__.py

@@ -0,0 +1,2 @@
+"""ActiveMD: Active Learning Molecular Dynamics with MACE potentials."""
+__version__ = "0.3.0"

firefly_activemd-0.3.0/firefly/cli.py

@@ -0,0 +1,127 @@
+import sys
+import json
+import logging
+from ase import units
+from ase.io import read
+import firefly.simulations.ensemble as sim
+
+logger = logging.getLogger(__name__)
+
+
+class ActiveMD:
+    def __init__(self, config):
+        self.config = config
+
+        erros = []
+
+        if "file" not in self.config or not self.config["file"]:
+            erros.append("The path to the atomic structure file ('file') is mandatory.")
+
+        if "ensemble" not in self.config or not self.config["ensemble"]:
+            erros.append("The simulation ensemble type ('ensemble') is mandatory (e.g. 'iso', 'aniso', 'bussi', etc.).")
+
+        if "code" not in self.config or not self.config["code"]:
+            erros.append("The reference DFT code ('code') is mandatory (use 'siesta' or 'qe').")
+
+        if "cmd" not in self.config or not self.config["cmd"]:
+            erros.append("The execution command for the DFT software ('cmd') is mandatory.")
+
+        if "thresh" not in self.config or self.config.get("thresh") is None:
+            erros.append("The uncertainty threshold ('thresh') is mandatory for Active Learning.")
+
+        if "n_models" not in self.config or self.config.get("n_models") is None:
+            erros.append("The number of models in the ensemble ('n_models') is mandatory.")
+
+        self.code = self.config.get("code")
+        if self.code == 'siesta':
+            if "atoms" not in self.config or not self.config["atoms"]:
+                erros.append("The 'atoms' dictionary mapping symbols to atomic numbers is mandatory when using Siesta.")
+
+        if erros:
+            logger.error("=" * 60)
+            logger.error(" CONFIGURATION ERROR DETECTED IN PARAMETERS.JSON")
+            logger.error("=" * 60)
+            logger.error("The following mandatory options are missing or invalid:")
+            for erro in erros:
+                logger.error(f"  -> {erro}")
+            logger.error("=" * 60)
+            sys.exit(1)
+
+        import os
+        caminho_arquivo = self.config["file"]
+        if not os.path.exists(caminho_arquivo):
+            logger.error("=" * 60)
+            logger.error(" ERROR: ATOMIC STRUCTURE FILE NOT FOUND")
+            logger.error("=" * 60)
+            logger.error(f" The file '{caminho_arquivo}' specified in 'file' does not exist.")
+            logger.error("=" * 60)
+            sys.exit(1)
+
+        self.ensemble = self.config["ensemble"]
+        self.atoms = read(caminho_arquivo, index='0')
+        self.T_beg = self.config.get("T_begin", 300.0)
+        self.T_fin = self.config.get("T_final", 300.0)
+        self.name = self.config.get("premodel", "model")
+        self.n = self.config["n_models"]
+        self.dict = self.config.get("atoms", {})
+        self.patience = self.config.get("patience", 1000)
+        self.thresh = self.config["thresh"]
+        self.cmd = self.config["cmd"]
+        self.time = self.config.get("time", 100000.0) * units.fs
+        self.interval = self.patience
+        self.dt = self.config.get("dt", 0.5) * units.fs
+        self.active = self.config.get("active", "energy")
+        self.epoch = self.config.get("epoch", 10)
+
+        pressao_valor = self.config.get("pressure")
+        if pressao_valor is not None:
+            self.press = pressao_valor * units.bar
+        else:
+            self.press = 0.0 * units.bar
+
+    def iniciar(self):
+        logger.info("Starting molecular dynamics and processing parameters...")
+
+        mapa_ensembles = {
+            'aniso': sim.aniso,
+            'iso': sim.iso,
+            'nptz': sim.NPTz,
+            'nptxy': sim.NPTxy,
+            'bussi': sim.bussi,
+            'ortho': sim.ortho,
+        }
+
+        funcao_simulacao = mapa_ensembles.get(self.ensemble.lower())
+        if funcao_simulacao:
+            funcao_simulacao(data=self)
+        else:
+            logger.error(f"The ensemble '{self.ensemble}' is not supported by the program.")
+            sys.exit(1)
+
+
+def main():
+    logging.basicConfig(
+        level=logging.INFO,
+        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+
+    if len(sys.argv) < 2:
+        logger.error("Correct usage: firefly parameters.json")
+        sys.exit(1)
+
+    try:
+        logger.info(f"Loading configuration file: {sys.argv[1]}")
+        with open(sys.argv[1], 'r') as arquivo:
+            parametros = json.load(arquivo)
+
+        app = ActiveMD(parametros)
+        app.iniciar()
+
+    except FileNotFoundError:
+        logger.error(f"Error: File '{sys.argv[1]}' was not found.")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()

firefly_activemd-0.3.0/firefly/simulations/ensemble.py

@@ -0,0 +1,265 @@
+import time
+import logging
+import numpy as np
+from ase import units
+from ase.cell import Cell
+from ase.io import read, write, Trajectory
+from ase.optimize import LBFGS
+from ase.md.langevinbaoab import LangevinBAOAB
+from ase.md.npt import NPT
+from ase.md.langevin import Langevin
+from ase.md.bussi import Bussi
+from ase.md.nptberendsen import NPTBerendsen
+from ase.md.nptberendsen import Inhomogeneous_NPTBerendsen
+from ase.md.velocitydistribution import MaxwellBoltzmannDistribution, Stationary, ZeroRotation
+from ase.calculators.mixing import AverageCalculator
+from mace.calculators import MACECalculator
+from ase.md.nose_hoover_chain import MaskedMTKNPT
+from ase.md.nose_hoover_chain import MTKNPT
+from ase.md.nose_hoover_chain import IsotropicMTKNPT
+
+import firefly.utils.ThreshComp as thresh
+import firefly.utils.MaceUtils as mu
+
+logger = logging.getLogger(__name__)
+
+amu_A3_to_g_cm3 = 1.660539
+GPa_to_atm = 9869.23
+GPa_to_eVA3 = 1. / 160.21766208
+atm_to_eVA3 = (1. / GPa_to_atm) * GPa_to_eVA3
+
+
+class TemperatureRamp:
+    def __init__(self, dyn, T_start, T_end, Nsteps):
+        self.dyn = dyn
+        self.Nsteps = Nsteps
+        self.T_start = T_start
+        self.T_end = T_end
+
+    def __call__(self):
+        a = (self.Nsteps - self.dyn.nsteps) / self.Nsteps
+        T = self.T_start * a + self.T_end * (1.0 - a)
+        if self.dyn.__class__.__name__ == 'Bussi':
+            self.dyn.temperature_K = T
+            self.dyn.temp = T * units.kB
+            self.dyn.target_kinetic_energy = 0.5 * self.dyn.temp * self.dyn.ndof
+        elif hasattr(self.dyn, 'set_temperature'):
+            self.dyn.set_temperature(temperature_K=T)
+        elif hasattr(self.dyn._thermostat, 'set_temperature'):
+            self.dyn._thermostat.set_temperature(temperature_K=T)
+        else:
+            self.dyn._thermostat._kT = T * units.kB
+
+
+class PrintEnergy:
+    def __init__(self, dyn, interval=1):
+        self.dyn = dyn
+        self.interval = interval
+        self.time_start = time.time()
+        self.time_last = self.time_start
+
+    def __call__(self):
+        a = self.dyn.atoms
+        step = self.dyn.nsteps
+        epot = a.get_potential_energy() / len(a)
+        ekin = a.get_kinetic_energy() / len(a)
+        temp = a.get_temperature()
+        stress = a.get_stress() / GPa_to_eVA3
+        press = ((stress[0] + stress[1] + stress[2]) / 3) * GPa_to_atm
+        dens = (a.get_masses().sum() / a.get_volume()) * amu_A3_to_g_cm3
+        cell = a.cell.cellpar()
+        time_now = time.time()
+        elapsed_time = time_now - self.time_last
+        self.time_last = time_now
+        spcpu = self.interval / elapsed_time
+
+        with open("therm.dat", "a") as f:
+            f.write('%d:  T= %.2f K  P= %.1f atm  K= %.5e  U= %.5e  rho= %.3e  spcpu= %.1f' % (step, temp, press, ekin, epot, dens, spcpu) + ' [%.3f %.3f %.3f %.2f %.2f %.2f]\n' % tuple(cell))
+
+
+def _get_calculators(data):
+    num_gpus = data.config.get("gpus", 1)
+    models = [f"model{i+1}.model" for i in range(data.n)]
+    calculators = []
+    for i, model_path in enumerate(models):
+        gpu_id = i % num_gpus
+        device = f"cuda:{gpu_id}"
+        logger.info(f"Loading {model_path} on device {device}")
+        calculators.append(MACECalculator(model_paths=model_path, device=device))
+    return calculators
+
+
+def aniso(data):
+    atoms = data.atoms
+    data.calculadoras = _get_calculators(data)
+    data.calc = AverageCalculator(data.calculadoras)
+    atoms.calc = data.calc
+
+    press = data.press
+    Tdamp = 100 * data.dt
+    Pdamp = 1000 * data.dt
+    n_steps = int(data.time / data.dt)
+
+    MaxwellBoltzmannDistribution(atoms, temperature_K=data.T_beg)
+    dyn = MTKNPT(data.atoms, timestep=data.dt, temperature_K=data.T_beg, pressure_au=press, tdamp=Tdamp, pdamp=Pdamp)
+    data.dyn = dyn
+
+    if data.active == 'force':
+        dyn.attach(lambda: thresh.thresh_forces(data=data), interval=1)
+    if data.active == 'energy':
+        dyn.attach(lambda: thresh.thresh_energies(data=data), interval=1)
+
+    dyn.attach(TemperatureRamp(dyn, data.T_beg, data.T_fin, n_steps), interval=100)
+    dyn.attach(PrintEnergy(dyn, 1), interval=1)
+    traj = Trajectory('run.traj', 'w', atoms)
+    dyn.attach(traj.write, interval=100)
+
+    logger.info("Starting the simulation")
+    dyn.run(n_steps)
+
+
+def iso(data):
+    atoms = data.atoms
+    data.calculadoras = _get_calculators(data)
+    data.calc = AverageCalculator(data.calculadoras)
+    atoms.calc = data.calc
+
+    press = data.press
+    n_steps = int(data.time / data.dt)
+
+    MaxwellBoltzmannDistribution(data.atoms, temperature_K=data.T_beg)
+    dyn = IsotropicMTKNPT(data.atoms, timestep=data.dt, temperature_K=data.T_beg, pressure_au=press, tdamp=100 * units.fs, pdamp=1000 * units.fs)
+    data.dyn = dyn
+
+    if data.active == 'force':
+        dyn.attach(lambda: thresh.thresh_forces(data=data), interval=1)
+    if data.active == 'energy':
+        dyn.attach(lambda: thresh.thresh_energies(data=data), interval=1)
+
+    dyn.attach(TemperatureRamp(dyn, data.T_beg, data.T_fin, n_steps), interval=100)
+    dyn.attach(PrintEnergy(dyn, 1), interval=1)
+    traj = Trajectory('run.traj', 'w', atoms)
+    dyn.attach(traj.write, interval=100)
+
+    logger.info("Starting the simulation")
+    dyn.run(n_steps)
+
+
+def ortho(data):
+    atoms = data.atoms
+    data.calculadoras = _get_calculators(data)
+    data.calc = AverageCalculator(data.calculadoras)
+    atoms.calc = data.calc
+
+    press = data.press
+    Tdamp = 100 * data.dt
+    Pdamp = 1000 * data.dt
+    n_steps = int(data.time / data.dt)
+
+    MaxwellBoltzmannDistribution(data.atoms, temperature_K=data.T_beg)
+    dyn = MaskedMTKNPT(data.atoms, timestep=data.dt, temperature_K=data.T_beg, pressure_au=press, tdamp=Tdamp, pdamp=Pdamp, mask=(1, 1, 1))
+    data.dyn = dyn
+
+    if data.active == 'force':
+        dyn.attach(lambda: thresh.thresh_forces(data=data), interval=1)
+    if data.active == 'energy':
+        dyn.attach(lambda: thresh.thresh_energies(data=data), interval=1)
+
+    dyn.attach(TemperatureRamp(dyn, data.T_beg, data.T_fin, n_steps), interval=100)
+    dyn.attach(PrintEnergy(dyn, 1), interval=1)
+    traj = Trajectory('run.traj', 'w', atoms)
+    dyn.attach(traj.write, interval=100)
+
+    logger.info("Starting the simulation")
+    dyn.run(n_steps)
+
+
+def NPTz(data):
+    atoms = data.atoms
+    data.calculadoras = _get_calculators(data)
+    data.calc = AverageCalculator(data.calculadoras)
+    atoms.calc = data.calc
+
+    press = data.press
+    Tdamp = 100 * data.dt
+    Pdamp = 1000 * data.dt
+    n_steps = int(data.time / data.dt)
+
+    MaxwellBoltzmannDistribution(data.atoms, temperature_K=data.T_beg)
+    Stationary(data.atoms)
+    ZeroRotation(data.atoms)
+    dyn = MaskedMTKNPT(data.atoms, timestep=data.dt, temperature_K=data.T_beg, pressure_au=press, tdamp=Tdamp, pdamp=Pdamp, mask=[False, False, True])
+    data.dyn = dyn
+
+    if data.active == 'force':
+        dyn.attach(lambda: thresh.thresh_forces(data=data), interval=1)
+    if data.active == 'energy':
+        dyn.attach(lambda: thresh.thresh_energies(data=data), interval=1)
+
+    dyn.attach(TemperatureRamp(dyn, data.T_beg, data.T_fin, n_steps), interval=100)
+    dyn.attach(PrintEnergy(dyn, 1), interval=1)
+    traj = Trajectory('run.traj', 'w', atoms)
+    dyn.attach(traj.write, interval=100)
+
+    logger.info("Starting the simulation")
+    dyn.run(n_steps)
+
+
+def NPTxy(data):
+    atoms = data.atoms
+    data.calculadoras = _get_calculators(data)
+    data.calc = AverageCalculator(data.calculadoras)
+    atoms.calc = data.calc
+
+    press = data.press
+    Tdamp = 100 * data.dt
+    Pdamp = 1000 * data.dt
+    n_steps = int(data.time / data.dt)
+
+    MaxwellBoltzmannDistribution(data.atoms, temperature_K=data.T_beg)
+    Stationary(data.atoms)
+    ZeroRotation(data.atoms)
+    dyn = MaskedMTKNPT(data.atoms, timestep=data.dt, temperature_K=data.T_beg, pressure_au=press, tdamp=Tdamp, pdamp=Pdamp, mask=[True, True, False])
+    data.dyn = dyn
+
+    if data.active == 'force':
+        dyn.attach(lambda: thresh.thresh_forces(data=data), interval=1)
+    if data.active == 'energy':
+        dyn.attach(lambda: thresh.thresh_energies(data=data), interval=1)
+
+    dyn.attach(TemperatureRamp(dyn, data.T_beg, data.T_fin, n_steps), interval=100)
+    dyn.attach(PrintEnergy(dyn, 1), interval=1)
+    traj = Trajectory('run.traj', 'w', atoms)
+    dyn.attach(traj.write, interval=100)
+
+    logger.info("Starting the simulation")
+    dyn.run(n_steps)
+
+
+def bussi(data):
+    atoms = data.atoms
+    data.calculadoras = _get_calculators(data)
+    data.calc = AverageCalculator(data.calculadoras)
+    atoms.calc = data.calc
+
+    Tdamp = 100 * data.dt
+    n_steps = int(data.time / data.dt)
+
+    MaxwellBoltzmannDistribution(data.atoms, temperature_K=data.T_beg)
+    Stationary(data.atoms)
+    ZeroRotation(data.atoms)
+    dyn = Bussi(data.atoms, timestep=data.dt, temperature_K=data.T_beg, taut=Tdamp)
+    data.dyn = dyn
+
+    if data.active == 'force':
+        dyn.attach(lambda: thresh.thresh_forces(data=data), interval=1)
+    if data.active == 'energy':
+        dyn.attach(lambda: thresh.thresh_energies(data=data), interval=1)
+
+    dyn.attach(TemperatureRamp(dyn, data.T_beg, data.T_fin, n_steps), interval=100)
+    dyn.attach(PrintEnergy(dyn, 1), interval=1)
+    traj = Trajectory('run.traj', 'w', atoms)
+    dyn.attach(traj.write, interval=100)
+
+    logger.info("Starting the simulation")
+    dyn.run(n_steps)

firefly_activemd-0.3.0/firefly/utils/DFTCodes.py

@@ -0,0 +1,111 @@
+import os
+import glob
+import logging
+import numpy as np
+from ase import Atoms
+from ase.io import write, read
+from ase.data import atomic_numbers, chemical_symbols
+from ase.calculators.singlepoint import SinglePointCalculator
+from pathlib import Path
+import sisl
+from sisl.io.siesta import stdoutSileSiesta
+
+logger = logging.getLogger(__name__)
+
+
+def create_siesta(data):
+    logger.info("Starting DFT reference calculation with SIESTA...")
+    celli = data.atoms.get_cell()
+    symbols = data.atoms.get_chemical_symbols()
+    posicoes = data.atoms.get_positions()
+
+    with open("geom.in", "w") as arq:
+        for sym, pos2 in zip(symbols, posicoes):
+            x, y, z = pos2
+            Z = data.dict[sym]
+            arq.write(f"{x:12.6f} {y:12.6f} {z:12.6f} {Z:3d}\n")
+
+    with open("lattice.fdf", "w") as arq:
+        arq.write("%block LatticeVectors\n")
+        c = celli.reshape(1, 9)[0]
+        arq.write(f"{c[0]:12.6f} {c[1]:12.6f} {c[2]:12.6f}\n")
+        arq.write(f"{c[3]:12.6f} {c[4]:12.6f} {c[5]:12.6f}\n")
+        arq.write(f"{c[6]:12.6f} {c[7]:12.6f} {c[8]:12.6f}\n")
+        arq.write("%endblock LatticeVectors\n")
+
+    write("last_add.xyz", data.atoms)
+
+    os.system(data.cmd + " < RUN.fdf > saida.out")
+
+    geom = sisl.get_sile("siesta.XV").read_geometry()
+    out = stdoutSileSiesta("saida.out")
+    E = out.read_energy()
+    Etot = E['total']
+    force = np.loadtxt("siesta.FA", skiprows=1)
+    forces = force[:, 1:4]
+    symbols_geom = [chemical_symbols[z] for z in geom.atoms.Z]
+    stress = out.read_stress()
+
+    atoms2 = Atoms(symbols=symbols_geom, positions=geom.xyz, cell=geom.cell, pbc=True)
+    atoms2.info['REF_energy'] = Etot
+    if stress is not None:
+        atoms2.info['REF_stress'] = stress
+    atoms2.arrays['REF_forces'] = forces
+
+    write("train.xyz", atoms2, append=True)
+
+    padroes_temporarios = [
+        "siesta*", "BASIS_*", "NON_TRIMMED_KP_LIST", "FORCE_STRESS", "OUTVARS.yml",
+        "PARALLEL_DIST", "TIMES", "*ion*", "MESSAGES", "CLOCK", "fdf*", "INPUT_TMP.*",
+        "geom.in", "lattice.fdf", "last_add.xyz", "saida.out",
+    ]
+    for padrao in padroes_temporarios:
+        for f in glob.glob(padrao):
+            try:
+                os.remove(f)
+            except OSError:
+                pass
+
+
+def create_qe(data):
+    logger.info("Starting DFT reference calculation with Quantum ESPRESSO (QE)...")
+    celli = data.atoms.get_cell()
+    symbols = data.atoms.get_chemical_symbols()
+    posicoes = data.atoms.get_positions()
+
+    with open("geom.in", "w") as arq:
+        arq.write("\nATOMIC_POSITIONS angstrom\n")
+        for sym, pos2 in zip(symbols, posicoes):
+            x, y, z = pos2
+            arq.write(f"{sym:4s}  {x:12.6f} {y:12.6f} {z:12.6f}\n")
+
+    with open("lattice.in", "w") as arq:
+        arq.write("\nCELL_PARAMETERS angstrom\n")
+        c = celli.reshape(1, 9)[0]
+        arq.write(f"{c[0]:12.6f} {c[1]:12.6f} {c[2]:12.6f}\n")
+        arq.write(f"{c[3]:12.6f} {c[4]:12.6f} {c[5]:12.6f}\n")
+        arq.write(f"{c[6]:12.6f} {c[7]:12.6f} {c[8]:12.6f}\n")
+
+    with open("run.in", "w") as outfile:
+        for fname in ["espresso.in", "geom.in", "lattice.in"]:
+            with open(fname) as infile:
+                outfile.write(infile.read())
+
+    os.system(data.cmd + " < run.in > saida.out")
+
+    atoms3 = read('saida.out', format='espresso-out')
+    atoms3.info['REF_energy'] = atoms3.get_potential_energy()
+    atoms3.arrays['REF_forces'] = atoms3.get_forces()
+    if atoms3.get_stress() is not None:
+        atoms3.info['REF_stress'] = atoms3.get_stress()
+
+    write('train.xyz', atoms3, append=True)
+
+    padroes_temporarios = ["geom.in", "lattice.in", "run.in", "saida.out", "active.save", "active.xml"]
+    for f in padroes_temporarios:
+        if os.path.exists(f):
+            if os.path.isdir(f):
+                import shutil
+                shutil.rmtree(f)
+            else:
+                os.remove(f)

firefly_activemd-0.3.0/firefly/utils/MaceUtils.py

@@ -0,0 +1,138 @@
+import glob
+import re
+import os
+import subprocess
+import time
+import logging
+from mace.calculators import MACECalculator
+from ase.calculators.mixing import AverageCalculator
+
+logger = logging.getLogger(__name__)
+
+
+def _get_calculators(data):
+    num_gpus = data.config.get("gpus", 1)
+    models = [f"model{i+1}.model" for i in range(data.n)]
+    calculators = []
+    for i, model_path in enumerate(models):
+        gpu_id = i % num_gpus
+        device = f"cuda:{gpu_id}"
+        logger.info(f"Loading {model_path} on device {device}")
+        calculators.append(MACECalculator(model_paths=model_path, device=device))
+    return calculators
+
+
+def renew_model(data):
+    data.calculadoras = None
+    data.calculadoras = _get_calculators(data)
+    new_calc = AverageCalculator(data.calculadoras)
+    data.atoms.calc = new_calc
+    logger.info("The model was restarted and distributed across GPUs")
+
+
+def train_parallel(data):
+    n = data.n
+    ep_train = data.epoch
+    ep_swa = data.epoch // 2
+
+    gpus_config = data.config.get("gpus", 1)
+    if isinstance(gpus_config, int):
+        gpus_disponiveis = list(range(gpus_config))
+    elif isinstance(gpus_config, list):
+        gpus_disponiveis = gpus_config
+    else:
+        gpus_disponiveis = [0]
+
+    num_gpus = len(gpus_disponiveis)
+    logger.info(f"[Parallel Training] Starting training of {n} models on {num_gpus} GPU(s): {gpus_disponiveis}")
+
+    jobs_para_executar = []
+    for i in range(n):
+        prefixo = data.name + str(i + 1)
+        arquivos = glob.glob(f"checkpoints/{prefixo}_run-*_epoch-*.pt")
+
+        max_epoch = max((int(re.search(r"epoch-(\d+)", f).group(1)) for f in arquivos), default=0)
+        novo_max = max_epoch + ep_train
+        novo_max_swa = max_epoch + ep_swa
+
+        job_file = f"job_{i+1}.sh"
+        if os.path.exists(job_file):
+            with open(job_file, "r") as f:
+                conteudo = f.read()
+            conteudo = re.sub(r"(--max_num_epochs=)\d+", f"\\1{novo_max}", conteudo)
+            conteudo = re.sub(r"(--start_swa=)\d+", f"\\1{novo_max_swa}", conteudo)
+            with open(job_file, "w") as f:
+                f.write(conteudo)
+
+        jobs_para_executar.append({
+            "id_modelo": i + 1,
+            "cmd": ["sh", job_file],
+            "arquivos_para_limpar": arquivos,
+            "prefixo": prefixo,
+        })
+
+    processos_ativos = {}
+    jobs_restantes = list(jobs_para_executar)
+    jobs_concluidos = []
+    falhas = []
+
+    while jobs_restantes or processos_ativos:
+        for gpu_id in gpus_disponiveis:
+            if gpu_id not in processos_ativos and jobs_restantes:
+                job = jobs_restantes.pop(0)
+                modelo_id = job["id_modelo"]
+                logger.info(f" -> Allocating Model {modelo_id} on GPU {gpu_id}...")
+                env_subprocesso = os.environ.copy()
+                env_subprocesso["CUDA_VISIBLE_DEVICES"] = str(gpu_id)
+                processo = subprocess.Popen(
+                    job["cmd"],
+                    stdout=subprocess.PIPE,
+                    stderr=subprocess.PIPE,
+                    env=env_subprocesso,
+                    text=True,
+                )
+                processos_ativos[gpu_id] = (processo, job)
+
+        time.sleep(1.0)
+
+        gpus_liberadas = []
+        for gpu_id, (proc, job) in processos_ativos.items():
+            status = proc.poll()
+            if status is not None:
+                modelo_id = job["id_modelo"]
+                gpus_liberadas.append(gpu_id)
+                if status == 0:
+                    logger.info(f" [OK] Model {modelo_id} completed on GPU {gpu_id}.")
+                    modelo_saida = f"{data.name}{modelo_id}_stagetwo.model"
+                    modelo_destino = f"model{modelo_id}.model"
+                    if os.path.exists(modelo_saida):
+                        os.rename(modelo_saida, modelo_destino)
+                    else:
+                        logger.warning(f" [Warning] File {modelo_saida} not found post-training.")
+                    for f in job["arquivos_para_limpar"]:
+                        if os.path.exists(f):
+                            os.remove(f)
+                    jobs_concluidos.append(job)
+                else:
+                    _, stderr_output = proc.communicate()
+                    logger.error(f" [ERROR] Model {modelo_id} failed on GPU {gpu_id} with code {status}.")
+                    logger.error(f"Details:\n{stderr_output}")
+                    falhas.append((modelo_id, stderr_output))
+
+        for gpu_id in gpus_liberadas:
+            del processos_ativos[gpu_id]
+
+    for i in range(n):
+        prefixo = data.name + str(i + 1)
+        for temp_file in glob.glob(f"{prefixo}*"):
+            if os.path.exists(temp_file):
+                os.remove(temp_file)
+
+    if falhas:
+        erros_str = ", ".join([f"Model {m_id}" for m_id, _ in falhas])
+        raise RuntimeError(f"Parallel training failed for models: {erros_str}.")
+
+    logger.info("[Parallel Training] All models have been successfully updated!")
+
+
+train_paralell = train_parallel

firefly_activemd-0.3.0/firefly/utils/ThreshComp.py

@@ -0,0 +1,97 @@
+import numpy as np
+import logging
+import firefly.utils.DFTCodes as cd
+import firefly.utils.MaceUtils as mu
+
+logger = logging.getLogger(__name__)
+
+
+def thresh_forces(data):
+    if not hasattr(data, 'calculadoras') or data.calculadoras is None:
+        data.calculadoras = mu._get_calculators(data)
+
+    atoms = data.atoms
+    calc_main = atoms.calc
+    committee_forces = []
+
+    for calc in data.calculadoras:
+        atoms.calc = calc
+        committee_forces.append(atoms.get_forces())
+
+    committee_forces = np.array(committee_forces)
+    std_forces = np.std(committee_forces, axis=0)
+    val_incerteza_estrutura = np.sqrt(np.mean(np.square(std_forces)))
+
+    if val_incerteza_estrutura > data.thresh:
+        logger.warning(f"Force uncertainty exceeded threshold! Value: {val_incerteza_estrutura:.6f} > Threshold: {data.thresh:.6f}")
+
+        if data.interval < data.patience and val_incerteza_estrutura < 3 * data.thresh:
+            logger.info(f"Patience active. Current interval: {data.interval}/{data.patience}")
+            data.atoms.calc = calc_main
+            data.interval += 1
+        else:
+            logger.warning("Patience limit exceeded or uncertainty too high (>= 3 * thresh). Starting DFT calculation and retraining.")
+            data.interval = 0
+            data.otf = val_incerteza_estrutura
+            passo_atual = data.dyn.get_number_of_steps()
+
+            with open("thresh.dat", "a") as f:
+                f.write(f"{passo_atual}\t{val_incerteza_estrutura:.6f}\n")
+
+            if data.code == 'siesta':
+                cd.create_siesta(data=data)
+                mu.train_parallel(data=data)
+                mu.renew_model(data=data)
+
+            if data.code == 'qe':
+                cd.create_qe(data=data)
+                mu.train_parallel(data=data)
+                mu.renew_model(data=data)
+    else:
+        data.atoms.calc = calc_main
+        data.interval += 1
+
+
+def thresh_energies(data):
+    if not hasattr(data, 'calculadoras') or data.calculadoras is None:
+        data.calculadoras = mu._get_calculators(data)
+
+    atoms = data.atoms
+    calc_main = atoms.calc
+    committee_energies = []
+
+    for calc in data.calculadoras:
+        atoms.calc = calc
+        committee_energies.append(atoms.get_potential_energy())
+
+    committee_energies = np.array(committee_energies)
+    std_energy = np.std(committee_energies)
+
+    if std_energy > data.thresh:
+        logger.warning(f"Energy uncertainty exceeded threshold! Value: {std_energy:.6f} > Threshold: {data.thresh:.6f}")
+
+        if data.interval < data.patience and std_energy < 3 * data.thresh:
+            logger.info(f"Patience active. Current interval: {data.interval}/{data.patience}")
+            data.atoms.calc = calc_main
+            data.interval += 1
+        else:
+            logger.warning("Patience limit exceeded or uncertainty too high (>= 3 * thresh). Starting DFT calculation and retraining.")
+            data.interval = 0
+            data.otf = std_energy
+            passo_atual = data.dyn.get_number_of_steps()
+
+            with open("thresh.dat", "a") as f:
+                f.write(f"{passo_atual}\t{std_energy:.6f}\n")
+
+            if data.code == 'siesta':
+                cd.create_siesta(data=data)
+                mu.train_parallel(data=data)
+                mu.renew_model(data=data)
+
+            if data.code == 'qe':
+                cd.create_qe(data=data)
+                mu.train_parallel(data=data)
+                mu.renew_model(data=data)
+    else:
+        data.atoms.calc = calc_main
+        data.interval += 1

firefly_activemd-0.3.0/firefly_activemd.egg-info/PKG-INFO

@@ -0,0 +1,54 @@
+Metadata-Version: 2.4
+Name: firefly-activemd
+Version: 0.3.0
+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 with MACE potentials and DFT reference calculations (SIESTA or Quantum ESPRESSO).
+
+## Installation
+
+```bash
+pip install .
+```
+
+Or in editable mode (for development):
+
+```bash
+pip install -e .
+```
+
+## Usage
+
+```bash
+firefly parameters.json
+```
+
+## Parameters (parameters.json)
+
+| Key | Required | Description |
+|-----|----------|-------------|
+| `file` | yes | Path to the atomic structure file (.xyz, etc.) |
+| `ensemble` | yes | MD ensemble: `iso`, `aniso`, `nptz`, `nptxy`, `bussi`, `ortho` |
+| `code` | yes | DFT code: `siesta` or `qe` |
+| `cmd` | yes | Command to run the DFT code |
+| `thresh` | yes | Uncertainty threshold for active learning |
+| `n_models` | yes | Number of MACE models in the committee |
+| `atoms` | siesta only | Dict mapping element symbols to atomic numbers |
+| `gpus` | no | Number of GPUs (default: 1) |
+| `T_begin` | no | Initial temperature in K (default: 300) |
+| `T_final` | no | Final temperature in K (default: 300) |
+| `pressure` | no | Pressure in bar (default: 0) |
+| `dt` | no | Time step in fs (default: 0.5) |
+| `time` | no | Total simulation time in fs (default: 100000) |
+| `patience` | no | Steps before forcing DFT call (default: 1000) |
+| `active` | no | Uncertainty metric: `energy` or `force` (default: `energy`) |
+| `epoch` | no | Training epochs per retraining step (default: 10) |
+| `premodel` | no | Model name prefix (default: `model`) |

firefly_activemd-0.3.0/firefly_activemd.egg-info/SOURCES.txt

@@ -0,0 +1,16 @@
+README.md
+pyproject.toml
+firefly/__init__.py
+firefly/cli.py
+firefly/simulations/__init__.py
+firefly/simulations/ensemble.py
+firefly/utils/DFTCodes.py
+firefly/utils/MaceUtils.py
+firefly/utils/ThreshComp.py
+firefly/utils/__init__.py
+firefly_activemd.egg-info/PKG-INFO
+firefly_activemd.egg-info/SOURCES.txt
+firefly_activemd.egg-info/dependency_links.txt
+firefly_activemd.egg-info/entry_points.txt
+firefly_activemd.egg-info/requires.txt
+firefly_activemd.egg-info/top_level.txt

firefly_activemd-0.3.0/firefly_activemd.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

firefly_activemd-0.3.0/firefly_activemd.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+firefly = firefly.cli:main

firefly_activemd-0.3.0/firefly_activemd.egg-info/requires.txt

@@ -0,0 +1,4 @@
+ase>=3.28
+mace-torch>=0.3
+numpy>=2
+sisl>=0.16

firefly_activemd-0.3.0/firefly_activemd.egg-info/top_level.txt

@@ -0,0 +1 @@
+firefly

firefly_activemd-0.3.0/pyproject.toml

@@ -0,0 +1,23 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "firefly-activemd"
+version = "0.3.0"
+description = "Active Learning Molecular Dynamics with MACE potentials and DFT reference calculations"
+readme = {file = "README.md", content-type = "text/markdown"}
+requires-python = ">=3.8"
+dependencies = [
+    "ase>=3.28",
+    "mace-torch>=0.3",
+    "numpy>=2",
+    "sisl>=0.16",
+]
+
+[project.scripts]
+firefly = "firefly.cli:main"
+
+[tool.setuptools.packages.find]
+where = ["."]
+include = ["firefly*"]

firefly_activemd-0.3.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+