Bring Your Own Model

MAPLE is not limited to its shipped backends (UMA, ANI, AIMNet2, MACE). You can plug in your own machine-learning potential — trained in your group, exported as a TorchScript .pt/.jpt or an ordinary PyTorch .pt/.model checkpoint — without editing MAPLE's source code. A custom backend is just a small Python class that exposes a few methods; once registered, it works with every task (single point, optimization, frequency, TS search, IRC, scan, MD).

How it works

Two header keys do all the wiring:

• module= — the import path of your plugin file/package. MAPLE imports it, which runs the @register_calculator decorator and adds your model name to the registry.
• model_path= — the checkpoint file MAPLE validates (it must exist) and hands to your class. Your code decides how to load it — this is where your .pt / .jpt / .model comes in.

#model=mymodel(module=my_lab.maple_plugin, model_path=/abs/path/model.jpt)
#sp
#device=cpu

Step 1 — Write the backend class

A backend inherits CalcABC and declares its capabilities as class attributes. You fill in three things: (a) the capability attributes, (b) how to load your model in __init__, and (c) the forward pass in _forward. Everything else — unit conversion, implicit-solvent hook, results bookkeeping, the analytic-Hessian loop — is handled for you.

import numpy as np
import torch
from ase.calculators.calculator import all_changes

from maple.function.calculator.calculator_base import (
    CalcABC, EV2HARTREE, hessian_via_double_autograd, register_calculator,
)


@register_calculator
class MyCalculator(CalcABC):
    # (a) Capability declaration, read by the factory before construction
    MODEL_NAMES = ('mymodel',)            # input-header routing keys, lowercase
    MODEL_ENERGY_UNIT = 'eV'              # 'eV' or 'hartree' -- declare honestly
    SUPPORTED_HESSIAN_MODES = ('analytic', 'numerical')
    SUPPORTS_CHARGE_MULT = False
    SUPPORTS_PBC = False                  # fail-fast on periodic atoms
    CHECKPOINT_FILENAME = None            # no HuggingFace auto-download
    REQUIRES_LOCAL_MODEL_FILE = False
    OPTION_KEYS = ()                      # extra model_options keys you accept
    MODEL_PATH_OPTION = 'model_path'      # kwarg that consumes model_path

    implemented_properties = ['energy', 'forces', 'free_energy', 'hessian']

    @classmethod
    def build_kwargs_from_options(cls, model, options, *, resolved_model_path=None):
        # Route the factory-resolved checkpoint path into __init__ kwargs
        kwargs = {}
        if resolved_model_path is not None:
            kwargs['model_path'] = resolved_model_path
        return kwargs

    def __init__(self, device, model='mymodel', *, model_path=None,
                 implicit='none', solvent='none', **_ignored):
        super().__init__()
        self.device = device
        self.dtype = torch.float64
        self.hessian = 'analytic'
        # (b) Load YOUR model here -- see Step 2
        self.model = load_my_model(model_path, device)
        # Keep this hook (no-op unless implicit='gbsa')
        self.implicit_solv_init(implicit=implicit, solvent=solvent)

    # (c) Single private forward: atoms -> scalar energy (in MODEL_ENERGY_UNIT)
    def _forward(self, atoms, *, requires_grad):
        positions = torch.tensor(
            atoms.get_positions(), dtype=self.dtype, device=self.device,
            requires_grad=requires_grad,
        )
        numbers = torch.tensor(atoms.get_atomic_numbers(), dtype=torch.long, device=self.device)
        energy_eV = self.model(numbers, positions)   # YOUR forward signature
        return energy_eV, positions

    def calculate(self, atoms=None, properties=['energy'], system_changes=all_changes):
        properties = self._normalize_properties(properties)
        atoms = super().calculate(atoms, properties, system_changes)  # PBC/solvent guards

        needs_forces = 'forces' in properties
        energy_eV, positions = self._forward(atoms, requires_grad=needs_forces)

        forces_np = None
        if needs_forces:
            forces = -torch.autograd.grad(energy_eV, positions)[0]
            forces_np = forces.detach().cpu().numpy()

        hessian = None
        if 'hessian' in properties:
            hessian = self.get_hessian(atoms)

        # _finalize_results owns eV->Hartree + solvent + writing self.results.
        # Do NOT multiply by EV2HARTREE yourself in this path.
        self._finalize_results(atoms, energy=energy_eV.item(), forces=forces_np, hessian=hessian)

    def _analytic_hessian(self, atoms) -> np.ndarray:
        # The Hessian path converts units itself: multiply by EV2HARTREE here.
        energy_eV, positions = self._forward(atoms, requires_grad=True)
        return hessian_via_double_autograd(lambda: energy_eV * EV2HARTREE, positions)

Step 2 — Load your model file

MAPLE hands you a validated, existing file path; how you turn it into a model object is entirely yours. The two common formats:

TorchScript (.pt / .jpt)

If your model was saved with torch.jit.script/trace, load it directly — this matches MAPLE's shipped ANI/MACE/AIMNet2 backends:

def load_my_model(model_path, device):
    model = torch.jit.load(model_path, map_location=device).eval()
    for p in model.parameters():
        p.requires_grad_(False)
    return model

Plain checkpoint / state_dict (.pt / .model)

For an ordinary (non-scripted) PyTorch model, rebuild the architecture and load weights:

def load_my_model(model_path, device):
    ckpt = torch.load(model_path, map_location=device, weights_only=False)
    net = MyNetwork(**ckpt['hparams'])      # your nn.Module
    net.load_state_dict(ckpt['state_dict'])
    net = net.to(device).double().eval()
    for p in net.parameters():
        p.requires_grad_(False)
    return net

Step 3 — Run it

Make your plugin importable (put its directory on PYTHONPATH or install it as a package), then point an input file at it. The example below mirrors examples/calculator_plugin/sp_jit.inp:

#model=mymodel(module=my_calculator_plugin, model_path=model.jpt)
#sp
#device=cpu

0 1
O   0.000   0.000   0.000
H   0.960   0.000   0.000
H  -0.240   0.930   0.000

export PYTHONPATH="$PWD:$PYTHONPATH"
maple sp.inp          # -> Energy: ... Hartree

The same backend now works for any task — just change the task header:

#model=mymodel(module=my_calculator_plugin, model_path=model.jpt)
#opt(method=lbfgs)
#device=gpu0

XYZ /path/to/molecule.xyz

Alternative: register process-wide

Instead of repeating module= in every input, register the plugin once via an environment variable. MAPLE imports each listed module at start-up:

export MAPLE_CALCULATOR_PLUGINS=my_calculator_plugin,other_lab.plugin

After this, #model=mymodel(model_path=...) resolves without the module= key.

Capability checklist

Set each attribute to what your model actually supports — MAPLE enforces them so misuse fails loudly instead of producing wrong numbers:

See also

• Runnable example: examples/calculator_plugin/ — two complete backends (TorchScript and torch.load), a toy-model builder, and ready input files.
• Calculator Backends — the full developer contract, capability matrix, and plugin-discovery layers.
• maple/function/calculator/AUTHORING.md and CALCULATOR_REVIEW.md in the source tree.