{ "cells": [ { "cell_type": "markdown", "id": "be6fbd9c-1268-4dab-9309-170c47ed85f0", "metadata": {}, "source": [ "# Modelforge ASE Calculator: Simple Walkthrough\n", "\n", "This notebook demonstrates how to use a potential trained with modelforge as a calculator in the Atomic Simulation Environment (ASE).\n", "It is written for first-time users of both modelforge and ASE.\n", "\n", "What you will do in this tutorial:\n", "1. Load a trained NNP model checkpoint.\n", "2. Wrap it with `ModelForgeCalculator`.\n", "3. Build a molecule from a SMILES string.\n", "4. Compute single-point energy and forces.\n", "5. Run geometry optimization and a short molecular dynamics simulation." ] }, { "cell_type": "code", "execution_count": 1, "id": "1c30380a-e276-4356-a3be-3601a8717b6e", "metadata": { "scrolled": true }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "\u001b[32m2026-04-20 15:24:41.993\u001b[0m | \u001b[34m\u001b[1mDEBUG \u001b[0m | \u001b[36mmodelforge.potential.potential\u001b[0m:\u001b[36mgenerate_potential\u001b[0m:\u001b[36m860\u001b[0m - \u001b[34m\u001b[1mtraining_parameter=None\u001b[0m\n", "\u001b[32m2026-04-20 15:24:41.994\u001b[0m | \u001b[34m\u001b[1mDEBUG \u001b[0m | \u001b[36mmodelforge.potential.potential\u001b[0m:\u001b[36mgenerate_potential\u001b[0m:\u001b[36m861\u001b[0m - \u001b[34m\u001b[1mpotential_parameter=SchNetParameters(potential_name='SchNet', only_unique_pairs=False, core_parameter=CoreParameter(number_of_radial_basis_functions=128, maximum_interaction_radius=0.49999999999999994, number_of_interaction_modules=9, number_of_filters=256, shared_interactions=True, activation_function_parameter=ActivationFunctionConfig(activation_function_name='ShiftedSoftplus', activation_function_arguments=None, activation_function=ShiftedSoftplus()), featurization=Featurization(properties_to_featurize=['atomic_number', 'per_system_total_charge'], atomic_number=AtomicNumber(maximum_atomic_number=101, number_of_per_atom_features=512), atomic_period=AtomicPeriod(maximum_period=8, number_of_per_period_features=32), atomic_group=AtomicGroup(maximum_group=18, number_of_per_group_features=32)), predicted_properties=['per_atom_energy', 'per_atom_charge'], predicted_dim=[1, 1]), postprocessing_parameter=PostProcessingParameter(properties_to_process=['per_atom_energy'], per_atom_energy=PerAtomEnergy(normalize=False, from_atom_to_system_reduction=True, keep_per_atom_property=True), per_atom_charge=PerAtomCharge(conserve=True, conserve_strategy='default'), per_system_electrostatic_energy=None, per_system_zbl_energy=None, per_system_vdw_energy=None, sum_per_system_energy=None, general_postprocessing_operation=GeneralPostProcessingOperation(calculate_molecular_self_energy=True, calculate_atomic_self_energy=False)), potential_seed=-1)\u001b[0m\n", "\u001b[32m2026-04-20 15:24:41.994\u001b[0m | \u001b[34m\u001b[1mDEBUG \u001b[0m | \u001b[36mmodelforge.potential.potential\u001b[0m:\u001b[36mgenerate_potential\u001b[0m:\u001b[36m862\u001b[0m - \u001b[34m\u001b[1mdataset_parameter=None\u001b[0m\n", "\u001b[32m2026-04-20 15:24:41.995\u001b[0m | \u001b[34m\u001b[1mDEBUG \u001b[0m | \u001b[36mmodelforge.potential.potential\u001b[0m:\u001b[36msetup_potential\u001b[0m:\u001b[36m705\u001b[0m - \u001b[34m\u001b[1mpotential_seed None\u001b[0m\n", "\u001b[32m2026-04-20 15:24:41.995\u001b[0m | \u001b[34m\u001b[1mDEBUG \u001b[0m | \u001b[36mmodelforge.potential.schnet\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m64\u001b[0m - \u001b[34m\u001b[1mInitializing the SchNet architecture.\u001b[0m\n", "\u001b[32m2026-04-20 15:24:42.008\u001b[0m | \u001b[34m\u001b[1mDEBUG \u001b[0m | \u001b[36mmodelforge.potential.potential\u001b[0m:\u001b[36msetup_potential\u001b[0m:\u001b[36m724\u001b[0m - \u001b[34m\u001b[1mOnly unique pairs: False\u001b[0m\n", "\u001b[32m2026-04-20 15:24:42.008\u001b[0m | \u001b[34m\u001b[1mDEBUG \u001b[0m | \u001b[36mmodelforge.potential.potential\u001b[0m:\u001b[36msetup_potential\u001b[0m:\u001b[36m758\u001b[0m - \u001b[34m\u001b[1mCutoffs: local_cutoff=0.49999999999999994, vdw_cutoff=-1, electrostatic_cutoff=-1\u001b[0m\n", "\u001b[32m2026-04-20 15:24:42.010\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mmodelforge.potential.neighbors\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m480\u001b[0m - \u001b[1mNeighborlistForInference initialized\u001b[0m\n", "\u001b[32m2026-04-20 15:24:42.103\u001b[0m | \u001b[34m\u001b[1mDEBUG \u001b[0m | \u001b[36mmodelforge.potential.potential\u001b[0m:\u001b[36mload_state_dict\u001b[0m:\u001b[36m672\u001b[0m - \u001b[34m\u001b[1mRemoved prefixes: {'potential.'}\u001b[0m\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "Loaded checkpoint: /Users/jenniferclark/bin/modelforge/modelforge-ase/modelforge/ase/tests/data/model.ckpt\n" ] } ], "source": [ "from modelforge.potential.potential import load_inference_model_from_checkpoint\n", "\n", "# Helper utilities to load the example model checkpoint bundled with modelforge.\n", "from modelforge.utils.io import get_path_string\n", "from modelforge.ase.tests import data\n", "\n", "checkpoint_file_path = f\"{get_path_string(data)}/model.ckpt\" # This is an example model used in testing \n", "potential = load_inference_model_from_checkpoint(checkpoint_file_path, jit=False)\n", "print(f\"Loaded checkpoint: {checkpoint_file_path}\")" ] }, { "cell_type": "markdown", "id": "f966aefd-ccf3-4a6f-bd91-0333c4e639af", "metadata": {}, "source": [ "To evaluate energies and forces in ASE, create a `ModelForgeCalculator` with the loaded potential and attach it to an ASE `Atoms` object." ] }, { "cell_type": "code", "execution_count": 2, "id": "6c1b834b-2e71-4100-abb4-dd34e541dd92", "metadata": {}, "outputs": [], "source": [ "# Install ASE\n", "from modelforge.ase import ModelForgeCalculator\n", "\n", "modelforge_calculator = ModelForgeCalculator(potential)\n" ] }, { "cell_type": "markdown", "id": "b8ce0fa1-d91a-46fa-ba98-62b4ea75108f", "metadata": {}, "source": [ "ASE includes tools for building molecules, and modelforge-ase also provides helper functions for a simple SMILES-based workflow.\n", "\n", "An ASE `Atoms` object stores the system definition: element identities, 3D positions, and optional simulation metadata (for example cell vectors and periodic boundary settings). By itself, `Atoms` is only a container for structure.\n", "\n", "To compute potential energy or forces, you must attach a calculator first (here, `ModelForgeCalculator`). Without a calculator, calls like `atoms.get_potential_energy()` and `atoms.get_forces()` cannot run." ] }, { "cell_type": "code", "execution_count": 3, "id": "b31a0d3b-e412-435e-adda-a554d705d7d2", "metadata": {}, "outputs": [], "source": [ "from modelforge.ase import smiles_to_ase, ase_to_rdkit\n", "\n", "# Set optimize=True to run an MMFF94 geometry optimization in RDKit before conversion.\n", "smiles = \"NCCCCCCO\"\n", "atoms = smiles_to_ase(smiles, optimize=False)\n", "\n", "# Attach a calculator before requesting energy/forces from ASE.\n", "atoms.calc = modelforge_calculator" ] }, { "cell_type": "markdown", "id": "7415213c-8e71-4e26-b4f8-e6c00acefdf5", "metadata": {}, "source": [ "You can view the 3D structure directly in ASE. (A separate RDKit view is also possible using the helper conversion utilities.)" ] }, { "cell_type": "code", "execution_count": 4, "id": "e0ba739d-dd83-494f-bdc1-e5dab0b84cfa", "metadata": {}, "outputs": [ { "data": { "text/html": [ "\n", " \n", " ASE atomic visualization\n", " \n", " \n", " \n", " \n", " \n", "\n", "\n", "\n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", " \n", "\n", "\n", "\n", " \n", " \n", "\n", "\n" ], "text/plain": [ "" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from ase.visualize import view\n", "view(atoms, viewer='x3d')" ] }, { "cell_type": "code", "execution_count": 5, "id": "80095751", "metadata": {}, "outputs": [], "source": [ "# Optional alternative viewer (install nglview first).\n", "# from nglview import show_ase\n", "# show_ase(atoms)" ] }, { "cell_type": "markdown", "id": "a1c77a27-03cb-410e-ace6-2af7c7c12b18", "metadata": {}, "source": [ "### Single-point energy and force computation\n", "With the calculator attached, ASE can evaluate potential energy and per-atom forces. `ModelForgeCalculator` handles unit conversion from modelforge internal units to ASE units (eV and angstrom)." ] }, { "cell_type": "code", "execution_count": 6, "id": "ddebdd47-a48f-477e-ba36-a62066a799c7", "metadata": { "scrolled": true }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Potential energy: -87.677811 eV\n", "Forces (eV/angstrom):\n", "[[ 9.35617831e+00 -1.04941746e+01 -2.46965417e+00]\n", " [-4.73898635e+00 2.68741981e+00 -1.49133123e+00]\n", " [ 1.76531205e+00 -4.81745170e-01 1.01300176e+00]\n", " [-6.05733491e-01 6.38226477e-01 7.25054235e-02]\n", " [-8.13248220e-03 -3.55702940e-01 -7.89400711e-01]\n", " [ 1.21058832e+00 8.26806163e-01 -6.81625482e-01]\n", " [ 6.57577191e-01 4.57251602e-01 -7.60445774e-01]\n", " [-1.75304901e+00 1.35702911e+00 -3.41923599e+00]\n", " [-8.86226398e+00 -1.88798054e+00 4.35927589e+00]\n", " [ 3.55498647e+00 7.28130734e+00 4.92473686e-01]\n", " [ 6.69284696e-01 6.24919797e-01 -1.99190706e+00]\n", " [ 2.48271806e-01 2.60230954e+00 7.58894359e-01]\n", " [-2.14171833e-01 -1.01460840e+00 -7.52640876e-01]\n", " [-9.50183786e-01 -6.46874168e-01 9.01118020e-01]\n", " [-2.31531027e-01 5.94798309e-01 1.97134656e-01]\n", " [ 3.06845990e-01 4.83842126e-01 -1.15194275e+00]\n", " [-2.01529873e-01 8.74319249e-02 1.19573265e-01]\n", " [-2.83709665e-01 5.89578509e-02 6.02928039e-01]\n", " [-3.62168497e-01 6.41423562e-01 6.89800633e-01]\n", " [ 3.00789842e-01 -1.08247846e+00 4.95305264e-01]\n", " [ 3.53056884e-02 -2.18362045e+00 -1.30540650e+00]\n", " [-3.37182905e-01 -8.14991585e-03 2.12953493e+00]\n", " [ 4.43502713e-01 -1.86389677e-01 2.98204480e+00]]\n" ] } ], "source": [ "pe = atoms.get_potential_energy()\n", "forces = atoms.get_forces()\n", "print(f\"Potential energy: {pe:.6f} eV\")\n", "print(\"Forces (eV/angstrom):\")\n", "print(forces)" ] }, { "cell_type": "markdown", "id": "b2637fa1-5424-4825-951a-d673e47aac91", "metadata": {}, "source": [ "### Geometry optimization\n", "ASE offers several structure optimizers. Here we use BFGS and stop when the maximum force falls below 0.05 eV/angstrom." ] }, { "cell_type": "code", "execution_count": 7, "id": "7034be70-e174-49a0-87aa-96f1dd7d4aff", "metadata": { "scrolled": true }, "outputs": [ { "data": { "text/plain": [ "True" ] }, "execution_count": 7, "metadata": {}, "output_type": "execute_result" } ], "source": [ "from ase.optimize import BFGS\n", "\n", "opt = BFGS(atoms, logfile=f\"{smiles}_opt.log\", trajectory=f\"{smiles}_opt.traj\")\n", "opt.run(fmax=0.05)" ] }, { "cell_type": "markdown", "id": "943126ef-8106-437b-970e-57b1c1d2943f", "metadata": {}, "source": [ "## Molecular dynamics simulation\n", "You can also run MD with ASE. This example runs short Langevin dynamics at 298 K and writes trajectory/log files." ] }, { "cell_type": "code", "execution_count": 8, "id": "ec574833-26eb-4293-80de-7771643eb367", "metadata": { "scrolled": true }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "/Users/jenniferclark/mamba/envs/modelforge/lib/python3.11/site-packages/ase/md/langevin.py:110: FutureWarning: The implementation of `fixcm=True` in `Langevin` does not strictly sample the correct NVT distributions. The deviations are typically small for large systems but can be more pronounced for small systems. Use `fixcm=False` together with `ase.constraints.FixCom`. `fixcm` is deprecated since ASE 3.28.0 and will be removed in a future release.\n", " warnings.warn(msg, FutureWarning)\n" ] }, { "data": { "text/plain": [ "True" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import ase.units as ase_units\n", "from ase.io.trajectory import Trajectory\n", "from ase.md import Langevin\n", "from ase.io import write, read\n", "\n", "trajectory_name = f\"{smiles}_sim.traj\"\n", "write(f\"{smiles}_system.pdb\", atoms)\n", "\n", "traj = Trajectory(trajectory_name, \"w\", atoms)\n", "\n", "dyn = Langevin(\n", " atoms,\n", " timestep=1.0 * ase_units.fs,\n", " temperature_K=298.0, # Kelvin\n", " friction=1.0 / ase_units.fs,\n", " trajectory=traj,\n", " logfile=f\"{smiles}_md.log\",\n", ")\n", "dyn.run(100)" ] }, { "cell_type": "markdown", "id": "6871520f", "metadata": {}, "source": [ "Convert ASE trajectory to XYZ so it can be opened in many viewers." ] }, { "cell_type": "code", "execution_count": 9, "id": "0e6e8ffb", "metadata": {}, "outputs": [], "source": [ "trj = read(trajectory_name, \":\")\n", "write(f\"{smiles}_sim.xyz\", trj, format=\"xyz\")" ] }, { "cell_type": "markdown", "id": "f88767f9", "metadata": {}, "source": [ "Now let's visualize the trajectory.\n", "\n", "First try the nglview widget inline. If widget rendering is unavailable in your frontend, the code falls back to ASE's inline x3d viewer." ] }, { "cell_type": "code", "execution_count": 10, "id": "645c351d", "metadata": {}, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "view(trj)" ] }, { "cell_type": "code", "execution_count": 11, "id": "1773bc8d", "metadata": {}, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "2026-04-20 15:24:48.133 python[1375:35253137] +[IMKClient subclass]: chose IMKClient_Modern\n" ] } ], "source": [ "## Notebook widget viewer (requires nglview + ipywidgets).\n", "#from nglview import show_asetraj\n", "#show_asetraj(trj)" ] } ], "metadata": { "kernelspec": { "display_name": "modelforge", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.15" } }, "nbformat": 4, "nbformat_minor": 5 }