Examples | NVIDIA Alchemi NIMs

Use the Quickstart Guide to launch the NIM before running the following examples.

For the request format and required atomic data fields, refer to API Reference.

Install the Python libraries needed for the examples on this page:

pip install requests ase numpy aiohttp

Using ASE for Structure I/O

Use the NIM for BGR to process a chemistry dataset by combining it with the I/O capabilities from the Atomic Simulation Environment (ASE). With a file structures.extxyz containing multiple structures in ASE extended XYZ format, and NIM for BGR running at http://localhost:8000/v1/infer, the following script runs batched optimization. Refer to the example structures.extxyz file.

Download structures.extxyz

The container-wide ALCHEMI_NIM_PBC setting determines the accepted request form. In the default periodic mode, each ASE Atoms object must define a cell, and the request must include cell. To process isolated molecules, first launch a separate container with a compatible model and ALCHEMI_NIM_PBC=false, then ensure that each request omits cell. Setting Atoms.pbc to false does not switch the running container to non-periodic mode. Refer to the AIMNet2 launch example.

The following Python script reads structures from structures.extxyz, performs batched optimization, and produces a structures_opt.extxyz file with optimized structures.

import requests
import ase.io
import numpy as np
EXTXYZ_FILE = "structures.extxyz"
INFER_URL = "http://localhost:8000/v1/infer"
# Read atoms from file
atoms_list = ase.io.read(EXTXYZ_FILE, index=":")
if not isinstance(atoms_list, list):
    atoms_list = [atoms_list]
# Convert ASE Atoms to API format
def ase_to_api(atoms, structure_id=None):
    data = {
        'coord': atoms.positions.flatten().tolist(),
        'numbers': atoms.numbers.tolist(),
        'charge': atoms.info.get('charge', 0),
        'mult': atoms.info.get('mult', 1),
    }
    if atoms.cell.volume > 0:
        data['cell'] = atoms.cell.array.flatten().tolist()
        data['pbc'] = atoms.pbc.tolist()
    if structure_id:
        data['structure_id'] = structure_id
    return data
# Tag structures with initial index for tracking
atoms_data = [ase_to_api(a, f"struct_{i}") for i, a in enumerate(atoms_list)]
# Create input request
input_data = {'atoms': atoms_data}
# Submit request and wait for results
response = requests.post(
    INFER_URL,
    json=input_data,
    headers={"Content-Type": "application/json"},
)
response.raise_for_status()
result = response.json()
optimized = result['atoms']
# Convert results back to ASE and write
optimized_atoms = []
for opt in optimized:
    a = ase.Atoms(
        positions=np.array(opt['coord']).reshape(-1, 3),
        numbers=opt['numbers']
    )
    if opt.get('cell'):
        a.set_cell(np.array(opt['cell']).reshape(3, 3))
    if opt.get('pbc'):
        a.set_pbc(opt['pbc'])
    a.info['energy'] = opt['energy']
    a.info['converged'] = opt['converged']
    a.info['optimizer_nsteps'] = opt['optimizer_nsteps']
    a.arrays['forces'] = np.array(opt['forces']).reshape(-1, 3)
    optimized_atoms.append(a)
# Check results
not_converged = sum(1 for a in optimized_atoms if not a.info['converged'])
if not_converged > 0:
    print(f'Warning: {not_converged} structures were not converged!')
# Write results
ase.io.write('structures_opt.extxyz', optimized_atoms)
print(f"Optimized {len(optimized_atoms)} structures")

Concurrent Requests

The inference API accepts multiple concurrent requests; the server processes them and returns responses in order. There is no dedicated async/polling API. Achieve concurrent behavior by issuing multiple requests and waiting on the responses.

Concurrent behavior example (Python)

import concurrent.futures
import requests
def run_infer(payload):
    return requests.post("http://localhost:8000/v1/infer", json=payload)
# Example: run three inference requests concurrently
payloads = [
    {"atoms": [{"coord": [0, 0, 0, 0, 0, 1], "numbers": [1, 1], "cell": [10, 0, 0, 0, 10, 0, 0, 0, 10]}]}
    for _ in range(3)
]
with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
    results = list(executor.map(run_infer, payloads))
for r in results:
    r.raise_for_status()
    print(r.json()["status"])

Asynchronous example (Python)

Use asyncio with an async HTTP client to send multiple inference requests without blocking. Install aiohttp (for example, pip install aiohttp) and run the following:

import asyncio
import aiohttp
async def infer(session, payload):
    async with session.post("http://localhost:8000/v1/infer", json=payload) as resp:
        resp.raise_for_status()
        return await resp.json()
async def main():
    payloads = [
        {"atoms": [{"coord": [0, 0, 0, 0, 0, 1], "numbers": [1, 1], "cell": [10, 0, 0, 0, 10, 0, 0, 0, 10]}]}
        for _ in range(3)
    ]
    async with aiohttp.ClientSession() as session:
        results = await asyncio.gather(*[infer(session, p) for p in payloads])
    for r in results:
        print(r["status"])
asyncio.run(main())

Partial Optimization With Active Masks

Freeze specific atoms during optimization by using the active_mask field:

import requests
# Optimize only the top two atoms, freeze the bottom one
package = {
    'atoms': [{
        'coord': [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 2.0],
        'numbers': [1, 1, 1],
        'cell': [10, 0, 0, 0, 10, 0, 0, 0, 10],
        'active_mask': [False, True, True]  # Freeze first atom
    }],
}
result = requests.post("http://localhost:8000/v1/infer", json=package)
print(result.json())