The code

Architecture overview

QMLearn is entirely Python. PySCF (or Psi4) supplies GTO integrals and reference RDMs. scikit-learn supplies the regressors. ASE handles geometries and molecular dynamics. The design centers on three objects:

• QMMol: one geometry in a consistent AO frame (Kabsch alignment to a reference). Exposes external potentials, overlap, 1-RDMs, 2-RDMs, energies, and forces.
• QMModel: sklearn wrappers that learn maps from V_ext to γ, δγ, Γ, Γ^c, or Δ, plus optional second-learn models for energy and forces.
• QMLCalculator: ASE calculator. method='gamma' for 1-RDM workflows; method='gamma2' / 'gamma2cum' for 2-RDM AIMD.

qmlearn.drivers.mol: QMMol

Construct a molecule with a basis, method, and optional active space, then duplicate it at new geometries while keeping the reference AO frame:

from ase.build import molecule
from qmlearn.drivers.mol import QMMol

atoms = molecule('H2O')
qmmol = QMMol(atoms=atoms, method='rks', basis='6-31g', xc='lda,vwn')
qmmol.run()

q2 = qmmol.duplicate(atoms)
print(q2.vext.shape)          # external potential in AO basis
energy = q2.calc_etotal(gamma)  # energy from a 1-RDM

Supported PySCF methods include dft/rks, hf/rhf, mp2, cisd, fci, casci, and ccsd. CAS methods take ncas and nelecas.

qmlearn.model: QMModel

A QMModel holds a dictionary of sklearn estimators. The primary method is the object learned from V_ext:

from sklearn.kernel_ridge import KernelRidge
from sklearn.linear_model import LinearRegression
from qmlearn.model import QMModel

mmodels = {
    'delta_gamma': KernelRidge(alpha=0.0, kernel='rbf'),
    'd_energy': LinearRegression(),
    'd_forces': LinearRegression(),
}
model = QMModel(mmodels=mmodels, method='delta_gamma',
                refqmmol=qmmol, purify_gamma=True)
model.fit(vext_list, delta_gamma_list)

MQMModel extends this to fragment (many-body expansion) models for large systems.

qmlearn.io: databases and db2qmmodel

DBHDF5 stores a reference QMMol, training atoms, and properties (V_ext, γ, Γ, energies, forces). db2qmmodel reads a file and fits the requested maps, including optional delta-learning of energy and forces:

from qmlearn.io.model import db2qmmodel
from sklearn.kernel_ridge import KernelRidge

qmmodel = db2qmmodel(
    'train.hdf5',
    names='*',
    mmodels={'gamma2c': KernelRidge(alpha=0.0, kernel='rbf')},
    target='gamma2c',
    method='gamma2c',
    purify_gamma=True,
)

Related helpers: read_db, merge_db, and get_train_atoms for normal-mode sampling of training geometries.

qmlearn.api.api4ase: QMLCalculator

ASE-compatible calculator. Implemented properties include energy, forces, dipole, gamma, and gamma2.

from qmlearn.api.api4ase import QMLCalculator

# 1-RDM (+ optional energy/force corrections)
atoms.calc = QMLCalculator(
    qmmodel=qmmodel,
    second_learn={'energy': 'd_energy', 'forces': 'd_forces'},
    method='gamma',
    properties=('energy', 'forces'),
)

# 2-RDM from Γᶜ
atoms.calc = QMLCalculator(qmmodel=qmmodel, method='gamma2',
                           properties=('energy',))

# 2-RDM from cumulant (needs two fitted models)
atoms.calc = QMLCalculator(qmmodel=qmmodel, qmmodel2=qmmodel2,
                           method='gamma2cum', properties=('energy',))

Package map

Worked examples are on the tutorials page.