Skip to main content

UMA machine-learned force-field integrated with ASE workflows

Project description

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, and vibrational/thermochemical analysis 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.
  • -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).
  • -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} [{sp,geoopt,freqs} ...]] [-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} [{sp,geoopt,freqs} ...]
                        Run type(s) to execute: 'sp', 'geoopt', 'freqs', 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.

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.

uma-ase web interface uma-ase web interface uma-ase web interface uma-ase web interface uma-ase web interface

  • 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.
  • Single Run tab. Upload one geometry, set charge/spin/run types (sp, geoopt, freqs) and launch a job. The panel streams the full stdout/stderr log, exposes download buttons for log/trajectory/optimized XYZ, and automatically pushes optimized structures into the Visualize tab when geoopt completes.
  • 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.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

uma_ase-0.7.1.tar.gz (7.6 MB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

uma_ase-0.7.1-py3-none-any.whl (8.2 MB view details)

Uploaded Python 3

File details

Details for the file uma_ase-0.7.1.tar.gz.

File metadata

  • Download URL: uma_ase-0.7.1.tar.gz
  • Upload date:
  • Size: 7.6 MB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.0rc3

File hashes

Hashes for uma_ase-0.7.1.tar.gz
Algorithm Hash digest
SHA256 7b9beb6df8624c7b357a3c62c924693952391443edbddd5f3842150759251f3c
MD5 6865b66695610add97267ec5f1813daf
BLAKE2b-256 dd8f89e4795d15c2b992df6a76a0c1189e739c46e162728558e8b87e0fc2ad55

See more details on using hashes here.

File details

Details for the file uma_ase-0.7.1-py3-none-any.whl.

File metadata

  • Download URL: uma_ase-0.7.1-py3-none-any.whl
  • Upload date:
  • Size: 8.2 MB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.12.0rc3

File hashes

Hashes for uma_ase-0.7.1-py3-none-any.whl
Algorithm Hash digest
SHA256 8fcc23378bee2635aed23b9a22b8cca0b077e1364ac9b9a1ad0065453521d532
MD5 dfd063b0af0f38f033ef9f3097043e2b
BLAKE2b-256 ca9139bfbc31734928447ada02c07268cf5a5fb7602d18aa5044e2602159bc42

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page