uma-ase 0.8.0

UMA machine-learned force-field integrated with ASE workflows

uma-ase

uma-ase bundles UMA (Universal Model for Atoms (https://huggingface.co/facebook/UMA) machine-learned force-field (MLFF) with the Atomic Simulation Environment (ASE) methods (https://ase-lib.org/). It supports basic single-point energy calculations, geometry optimisation, vibrational/thermochemical analysis, and transition-state searches via ASE's dimer algorithm from a single command-line entry point or an optional web GUI service.

Requirements

The project currently ships and has been validated with Python 3.12. The runtime stack is:

• Python ≥ 3.10 (tested with 3.12)
• ASE ≥ 3.26.0
• numpy ≥ 2.2
• torch ≥ 2.6
• fairchem-core ≥ 2.10 for UMA checkpoints and calculators
• flask ≥ 3.0 when you want the optional web GUI

Install these packages with:

pip install -r requirements.txt

If you build your own environment, make sure fairchem is present—geometry optimisations rely on it to compute per-atom reference energies. On Apple Silicon you may also want to limit OpenMP threads when driving torch (e.g. export OMP_NUM_THREADS=1) to avoid shared-memory warnings.

Installation

Released package

pip install uma-ase[server]

The server extra installs the optional GUI web interface. Omit it when you only need the command-line tooling. Afterwards, install FairChem explicitly if your environment does not already ship it:

pip install fairchem-core

FairChem is repidly evolving. Check version changes and compatibility with uma-ase.

From source

python -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
pip install -r requirements.txt
pip install -e ".[server]"

The editable install registers the uma-ase and uma-ase-server console scripts for local development.

Command-line (CLI) usage

• Installed via pip (PyPI, wheel, or editable install)

  # Set up where your local checkpoint UMA MLFF files are located.  
  export UMA_CUSTOM_CHECKPOINT_DIR=/abs/path/to/checkpoints
  (include this your .bashrc or .bash_profile conf files).
  
  # Single run
  uma-ase -input geometry.xyz -chg 0 -spin 1 -run-type sp
  
  # Using local checkpoints (optional)
  uma-ase -input geometry.xyz -chg +1 -run-type sp freqs -temp 300 -press 130000.0
  

• Working from a cloned repository without installing

  export PYTHONPATH=src
  python -m uma_ase.py -input geometry.xyz -chg 0 -spin 1 -run-type geoopt
  

The CLI always emits a consolidated log (e.g. molecule-SP-OPT.log), a trajectory (*.traj), an optimised XYZ (*-geoopt-OPT.xyz), and frequency outputs when requested. Run uma-ase -h (or python -m uma_ase.cli -h) for the full reference.

FairChem’s UMA models can be loaded directly from the official distribution (see the FairChem/UMA docs), or you can download the checkpoint files yourself for extra stability and keep them wherever you prefer. Point the shell environment variable UMA_CUSTOM_CHECKPOINT_DIR at your directory so -mlff-chk uma-s-1p1 and similar shortcuts resolve against your local files. If the variable is unset, uma-ase looks for checkpoints under ~/.uma_ase/checkpoints. When an XYZ comment line embeds an string and two signed integers (e.g. https://… -1 2), uma-ase reads them before invoking ASE and uses the first as the default charge and the optional second one as the default spin multiplicity, unless you supplied -chg/-spin explicitly.

CLI option summary

• -input (required): input geometry file readable by ASE, typically .xyz. A standard XYZ can omit -chg/-spin when the second-line comment ends with signed integers (optionally alongside a source URL), e.g.

  199
  https://iochem-bd.iciq.es/browse/handle/100/69781 -1 2
  O  0.0 0.0 0.0
  H  0.0 0.0 0.96
  H  0.0 0.75 -0.48
  C .....
  

  uma-ase reads the charge -1 and spin 2 before invoking ASE and adopts them as defaults.
• -chg (default 0): total molecular charge; omitted values default to 0 or to the signed integer embedded in the XYZ comment line if present.
• -spin (default 1): spin multiplicity; omitted values default to 1 or to the second signed integer embedded in the XYZ comment line when present.
• -run-type (default sp): workflow steps to run; choose any sequence of sp, geoopt, freqs, ts. A ts run first executes a frequency calculation (discarding thermochemistry) to extract the softest mode, launches ASE’s dimer method seeded with that mode, and finally recomputes full vibrational data on the converged saddle point.
• -cpu: force CPU execution even when CUDA is available (helpful on shared HPC nodes or to avoid GPU startup costs).
• -iter (default 250): maximum geometry optimisation cycles.
• -grad (default 0.01 eV/Å): convergence threshold on the maximum force component.
• -optimizer (default LBFGS): ASE optimiser to use (BFGS, LBFGS, FIRE, BFGSLineSearch, MDMin).
• --ts-displacement / --ts-dimer-separation / --ts-trial-step / --ts-max-step / --ts-max-rot: tune the initial mode displacement and dimer parameters used by the ts workflow. These settings also control the preparatory and final frequency calculations, so every TS run delivers both the dimer trajectory and the validated saddle-point frequencies in a single invocation.
• -mlff-chk (default uma-s-1p1): UMA checkpoint identifier.
• -mlff-task (default omol): UMA task/model name passed to the calculator.
• -temp/--temperature (default 298.15 K): vibrational/thermochemistry temperature.
• -press/--pressure (default 101325 Pa): vibrational/thermochemistry pressure.
• -visualize: open the geoopt trajectory in ASE’s viewer after completion.

CLI help (uma-ase -h)

usage: uma-ase [-h] -input INPUT [-chg CHG] [-spin SPIN]
               [-run-type {sp,geoopt,freqs,ts} [{sp,geoopt,freqs,ts} ...]] [-cpu]
               [-iter ITER] [-grad GRAD] [-optimizer NAME]
               [-mlff-chk CHECKPOINT] [-mlff-task TASK] [-temp T]
               [-press P] [-visualize]

Provide the required -input value. Charge defaults to 0 (or the first signed
integer embedded in the XYZ comment line) and spin defaults to 1 (or the second
integer when present).

options:
  -h, --help            show this help message and exit
  -input INPUT          Input geometry readable by ASE (XYZ comment may append
                        signed charge/spin metadata).
  -chg CHG              Molecular charge override (default 0; inferred from
                        XYZ comment when present).
  -spin SPIN            Spin multiplicity (default 1; inferred from XYZ comment
                        when a second integer is present).
  -run-type {sp,geoopt,freqs,ts} [{sp,geoopt,freqs,ts} ...]
                        Run type(s) to execute: 'sp', 'geoopt', 'freqs',
                        'ts', or any sequence thereof (default: sp).
  -cpu                  Force UMA calculations to run on CPU even when CUDA is
                        available.
  -iter ITER            Max number of geometry optimization cycles.
                        Default=250
  -grad GRAD            Max grad for convergence. Default=0.01 eV/A
  -optimizer NAME       ASE optimizer (e.g. BFGS, LBFGS, FIRE,
                        BFGSLineSearch,MDMin). Default='LBFGS'.
  -mlff-chk CHECKPOINT  UMA checkpoint identifier. Default='uma-s-1p1'.
  -mlff-task TASK       UMA task/model identifier. Default='omol'.
  -temp   T             Temperature in Kelvin for vibrational analysis
                        (default 298.15 K).
  -press   P            Pressure in Pascals for vibrational analysis (default
                        101325.0 Pa).
  -visualize            Open the trajectory of a geoopt run in an interactive
                        viewer.

Extracting vibrational mode matrices

Frequency runs leave one multi-frame .xyz per normal mode under freqs/<job>/ and a packaged archive <input>-FREQS-modes.zip. Use the standalone helper to convert those animations into raw N_atoms x 3 displacement matrices:

# Convert modes that were already unzipped by the workflow
uma-ase-modes freqs/Nb6O19-geoopt-OPT-FREQS -o Nb6O19-modes

# Or point it directly at the <input>-FREQS-modes.zip bundle
uma-ase-modes Nb6O19-FREQS-modes.zip --normalize

Each mode is saved as <job>-mode-XXX.npy with optional normalisation so the largest component equals 1. A companion modes.json lists the source file, atom count, and phase parameters used to reconstruct each displacement for downstream scripts.

Web interface (optional)

uma-ase-server

then visit http://127.0.0.1:8000. The webapp (uma-ase.html) is bundled with the package and submits jobs to /api/uma-ase/run. The backend stores each uploaded geometry in a temporary directory, delegates to the CLI, returns the generated log, and removes temporary files automatically. The page focuses on job submission, showing a live summary of the uploaded structure, and exposing UMA checkpoint/task selectors.

Working directly from the source tree without installing? Prefix the module path:

PYTHONPATH=src python -m uma_ase.server
# or export PYTHONPATH=src once, then:
python -m uma_ase.server

Each run stores the returned log under ~/.uma_ase/results/ (configurable via UMA_RESULTS_DIR), and the interface enables a Download Log button once a job finishes.

• Visualize tab. Opens first by default so you can inspect structures immediately. The bundled JSmol canvas synchronises with every geometry you upload in Single Run, keeps charge/spin indicators aligned, and still lets you drag-and-drop alternate files straight into the viewer panel. The tab now shares space with a dedicated 3Dmol.js Trajectory Viewer: drag in any .traj/.xyz/.xyz.gz file (or load the auto-generated optimization/MD trajectories) and scrub through frames, play/pause, adjust FPS, or capture live MD runs. A Vibrational Normal Modes panel accepts the packaged <input>-FREQS-modes.zip plus its log file, parses the “Vibrational Summary”, and builds a selector that animates each normal mode as a multi-frame trajectory (palindrome loop) complete with a stop button. When a freqs/TS job finishes, the server auto-downloads the ZIP + log, populates the selectors, reloads the initial geometry in the viewer, and jumps straight into the Trajectory tab so you can inspect the modes without manual uploads.
• Single Run tab. Upload one geometry, set charge/spin/run types (sp, geoopt, freqs, ts) and launch a job. The panel streams the full stdout/stderr log, exposes download buttons for log/trajectory/optimized XYZ/FREQS bundles, automatically pushes optimized structures into the Visualize tab when geoopt completes, and auto-loads the traj output (including an energy vs. step chart) whenever a GeoOpt/TS step finishes.
• Multi Run tab. Point at a folder containing multiple .xyz files (non-OPT/geoopt names) and the UI builds a summary table with the detected charge/spin for each entry. Choose the shared run settings once, then fan out to as many jobs as needed. The Results panel shows the complete output of every job (including inline download links per file) and a “Download Files” button bundles every artifact from the batch into a single ZIP.
• Analyze tab. Post-process folders of uma-ase results directly in the browser.
  • RMSD Report reads uma-ase .log files and generates PDF/LaTeX/DOCX artifacts plus an inline PDF preview powered by styled_rmsd_report.py.
  • XYZ Pair Analysis drives scripts-to-share_v2/driver.py, matching .xyz files with their -geoopt-OPT counterparts, collecting TXT/PDF/LaTeX/DOCX outputs, and exposing download buttons for each format.
• Simulate tab. Reuse the same UMA checkpoints to launch ASE MD jobs by choosing steps, timestep, target temperature, ensemble-specific knobs, and sampling intervals. The UI mirrors the Single Run layout, streams the MD log in real time, and exposes downloads for the trajectory, multi-frame XYZ, and final snapshot once the simulation completes.

Simulate tab (Molecular Dynamics)

The Simulate tab is a front-end for ASE’s molecular dynamics drivers coupled to UMA MLFF calculators. It consumes whichever structure is currently loaded in the Visualize tab and lets you tweak every relevant thermostat/barostat parameter without touching Python. Controls are grouped as follows.

Section	Options
System	Charge, spin multiplicity, UMA checkpoint (uma-s-1p1 default), UMA task (omol default).
Integrator	Dynamics engine selector covering Langevin (NVT), VelocityVerlet (NVE), NVT Berendsen, Nosé-Hoover Chain, Bussi, Andersen, Langevin-Hoover BAOAB (NVT/NPT), NPT Berendsen, Isotropic MTK, Full MTK, and Melchionna NPT.
Core MD settings	Steps, timestep (fs), target temperature (K), per-engine extras (friction for Langevin, thermostat relaxation for Berendsen).
NPT controls	Pressure (bar), compressibility (1/bar), barostat relaxation (fs), MTK damping and chain lengths, BAOAB T_tau/P_tau/cell-mass factor/hydrostatic toggle, Melchionna thermostat time, barostat time, bulk modulus, and mask vector.
PBC / cell	Checkbox to enable periodic boundary conditions and numeric fields for cell lengths a/b/c (Å) plus angles α/β/γ (°). When enabled the app scales the uploaded geometry into that lattice before the run.
Sampling	Trajectory write interval, log interval.

Each control only appears when relevant (e.g., MTK sliders only unlock when you pick an MTK engine), mirroring ASE’s own API. During a run the Simulate output panel shows a live tail of the job log, and once the job finishes the full log stays visible while the JSmol viewer auto-loads md-final.xyz. Download buttons become active for the MD log, ASE .traj, multi-frame .xyz, and the final snapshot. Any backend error (for example, picking BAOAB when your ASE build predates ase.md.langevinbaoab) gets surfaced inline with the current log.

Package layout

src/uma_ase/
├── __init__.py          # Version metadata
├── __main__.py          # Enables `python -m uma_ase`
├── cli.py               # Console entry point
├── server.py            # Flask application (optional)
├── utils.py             # CLI parser and helper utilities
├── workflows.py         # Core UMA/ASE workflow orchestration
└── static/uma-ase.html  # Single-page frontend served by the Flask app

Development workflow

1. Create a virtual environment and install the package in editable mode (pip install -e .[server]).
2. Run unit or integration tests as desired (add your preferred framework).
3. Build distributions for publishing:

   python -m build
   

4. Upload to a package index (e.g., GitLab Package Registry or PyPI):

   python -m twine upload dist/*
   

License

(c) CC BY This code has been generated by ChatGPT Codex agent under the supervision of Carles Bo.