Add add_library_charges_to_forcefield fn

docs/changelog.md

@@ -1,9 +1,14 @@
 # Changelog
 
-## Unreleased
+## 0.9.0
+
+### Features
+
+- Add `presto.create_types.add_library_charges_to_forcefield` to write custom partial charges from OpenFF `Molecule` objects into a force field as library charges, addressing [#64](https://github.com/cole-group/presto/issues/64).
 
 ### Documentation
 
+- Add [Use custom charges](how-to/use-custom-charges.md) how-to guide.
 - Add a `CITATION.cff` file and cite the presto preprint in the README and docs, and point users to the OpenFF publications to cite, in [#76](https://github.com/cole-group/presto/pull/76).
 - Document installing `presto` from `conda-forge` as `presto-fit` (note this comes without the MLP dependencies, which must be installed separately) in [#70](https://github.com/cole-group/presto/pull/70).
 

docs/how-to/index.md

@@ -7,6 +7,7 @@ Short, task-oriented recipes. Each page assumes you've finished **[Get started](
 - **[Fit a single molecule](fit-single-molecule.md)** — defaults you may want to change.
 - **[Fit a congeneric series](fit-congeneric-series.md)** — share parameters across related molecules using `max_extend_distance`.
 - **[Use SDF inputs](use-sdf-inputs.md)** — switch from SMILES to one or more `.sdf` files.
+- **[Use custom charges](use-custom-charges.md)** — bake your own partial charges into the force field as library charges.
 - **[Wipe output and rerun](clean-rerun.md)**
 
 ## By component

docs/how-to/use-custom-charges.md

@@ -0,0 +1,76 @@
+# Use custom charges
+
+By default a `presto` fit takes partial charges from the base force field's charge model (usually Ash-GC). If you have your own partial charges — from a higher level of theory, a custom charge model, or an external pipeline — you can bake them into the force field as **library charges** so they are used in place of the default model.
+
+`presto` provides [`add_library_charges_to_forcefield`](../reference/api/create_types.md#presto.create_types.add_library_charges_to_forcefield) for this. It takes an OpenFF `Molecule` (or list of molecules) that already has `partial_charges` set, plus a `ForceField`, and returns a copy of the force field with `LibraryCharges` written for those charges.
+
+This is a standalone helper: it is **not** part of the automatic fitting pipeline. You call it yourself to produce a force field (or `.offxml` file) that you then use as your starting point.
+
+## Set partial charges on the molecule
+
+Any method that populates `Molecule.partial_charges` works — assign them with the toolkit, or set them directly from your own data:
+
+```python
+import numpy as np
+import openff.toolkit
+from openff.units import unit
+
+mol = openff.toolkit.Molecule.from_smiles("CCO")
+
+# Option A: assign with a toolkit charge method
+mol.assign_partial_charges("am1bcc")
+
+# Option B: set your own charges directly (must be ordered by atom index)
+mol.partial_charges = np.array([...]) * unit.elementary_charge
+```
+
+The charges must sum to the molecule's formal charge. If they do not, `add_library_charges_to_forcefield` raises a `ValueError` up front rather than letting the error surface later when the system is built.
+
+## Write the library charges
+
+```python
+from presto.create_types import add_library_charges_to_forcefield
+
+ff = openff.toolkit.ForceField("openff_unconstrained-2.3.0.offxml")
+ff_with_charges = add_library_charges_to_forcefield(mol, ff)
+
+# Persist for reuse / inspection
+ff_with_charges.to_file("custom_charges.offxml")
+```
+
+The original `ff` is left unchanged; a modified copy is returned. The new parameters are tagged with `l-bespoke-*` IDs under the `LibraryCharges` handler.
+
+Pass a list to write charges for several molecules at once:
+
+```python
+ff_with_charges = add_library_charges_to_forcefield([mol_a, mol_b], ff)
+```
+
+## How the charges are matched
+
+The library charges reproduce your input charges and always keep the molecule net-neutral (or at its formal charge). This is achieved by reusing `presto`'s [type generation](../concepts/type-generation.md) machinery:
+
+- One library charge is generated per atom, using a whole-molecule SMARTS with a single tagged atom (the non-tagged hydrogens are merged onto their heavy atoms, which keeps SMARTS matching fast).
+- Because each pattern spans the whole molecule, it only matches that atom's symmetry class, so every atom is covered exactly once.
+- Symmetry-equivalent atoms collapse onto one library charge whose value is the **mean** of their input charges. Averaging both symmetrises equivalent atoms and preserves the total charge exactly.
+
+A whole-molecule SMARTS is used (rather than a smaller `max_extend_distance`) to ensure that charges are only applied to the molecule they were derived for: a `LibraryCharges` handler is only applied if it covers every atom of the molecule, and the assigned charges must sum to the formal charge. Whole-molecule patterns guarantee both.
+
+## Use the charges in a fit
+
+Point `initial_force_field` at the `.offxml` you wrote:
+
+```yaml
+param_settings:
+    molecule_input_type: smiles
+    molecules:
+        - CCO
+    initial_force_field: custom_charges.offxml
+```
+
+`presto` only retrains valence parameters (bonds, angles, torsions), so your custom charges are carried through the fit unchanged. Make sure the molecules you fit are the same ones you wrote charges for, so the library charges match.
+
+## Reference
+
+- API reference: [`add_library_charges_to_forcefield`](../reference/api/create_types.md#presto.create_types.add_library_charges_to_forcefield).
+- Concept: [Type generation and SMIRKS specificity](../concepts/type-generation.md).

mkdocs.yml

@@ -33,6 +33,7 @@ nav:
   - Fit a single molecule: how-to/fit-single-molecule.md
   - Fit a congeneric series: how-to/fit-congeneric-series.md
   - Use SDF inputs: how-to/use-sdf-inputs.md
+  - Use custom charges: how-to/use-custom-charges.md
   - Clean and rerun a fit: how-to/clean-rerun.md
   - Choose an MLP: how-to/choose-an-mlp.md
   - Use an ASE calculator: how-to/use-ase-calculator.md

presto/create_types.py

@@ -1,13 +1,15 @@
 """Create new tagged SMARTS parameter types for molecules of interest."""
 
+from __future__ import annotations
+
 import copy
 import warnings
 from collections import defaultdict
 from collections.abc import Mapping
 
 import openff.toolkit
 from loguru import logger
-from openff.units import Quantity
+from openff.units import Quantity, unit
 from rdkit import Chem
 
 from .settings import TypeGenerationSettings
@@ -248,7 +250,7 @@ def add_types_to_forcefield(
     openff.toolkit.ForceField
         Force field with bespoke parameters added, deduplicated across all molecules
     """
-    # Convert single molecule to list for backward compatibility
+    # Convert single molecule to list
     if isinstance(mols, openff.toolkit.Molecule):
         mols = [mols]
 
@@ -328,3 +330,111 @@ def add_types_to_forcefield(
     ff_copy = _remove_redundant_smarts(mols_for_typing, ff_copy, id_substring="bespoke")
 
     return ff_copy
+
+
+def add_library_charges_to_forcefield(
+    mols: openff.toolkit.Molecule | list[openff.toolkit.Molecule],
+    force_field: openff.toolkit.ForceField,
+) -> openff.toolkit.ForceField:
+    """Write per-atom ``LibraryCharges`` from molecules' partial charges into a force field.
+
+    For each atom of each molecule a bespoke single-tagged-atom SMARTS spanning the
+    whole molecule is generated using the same machinery as the valence types (see
+    :func:`add_types_to_forcefield`). The charge assigned to each SMARTS is the mean
+    of the partial charges of all atoms that produce it (i.e. symmetry-equivalent
+    atoms), which both symmetrises equivalent atoms and preserves the total molecular
+    charge exactly.
+
+    Because the SMARTS spans the whole molecule, each one only matches its own
+    symmetry class, so every atom is covered by exactly one library charge and the net
+    charge of the molecule is reproduced. This is required because OpenFF interchange
+    does not renormalise charges: it only applies a ``LibraryCharges`` handler if it
+    covers every atom of the molecule, and otherwise raises if the assigned charges do
+    not sum to the formal charge. The non-tagged hydrogens are merged onto their heavy
+    atoms by ``MergeQueryHs`` for fast SMARTS matching.
+
+    Parameters
+    ----------
+    mols : openff.toolkit.Molecule | list[openff.toolkit.Molecule]
+        Molecule or molecules with ``partial_charges`` set.
+    force_field : openff.toolkit.ForceField
+        The base force field to add the library charges to.
+
+    Returns:
+    -------
+    openff.toolkit.ForceField
+        A copy of the force field with bespoke library charges added, deduplicated
+        across all molecules.
+    """
+    # Convert single molecule to list
+    if isinstance(mols, openff.toolkit.Molecule):
+        mols = [mols]
+
+    # Validate partial charges before doing any work
+    charges_per_mol: list[list[Quantity]] = []
+    for mol in mols:
+        if mol.partial_charges is None:
+            raise ValueError(
+                f"Molecule {mol.to_smiles(explicit_hydrogens=False)} is missing "
+                "partial charges. Set Molecule.partial_charges before generating "
+                "library charges."
+            )
+
+        charges = list(mol.partial_charges)
+        charge_sum = sum(c.m_as(unit.elementary_charge) for c in charges)
+        formal_sum = mol.total_charge.m_as(unit.elementary_charge)
+        if abs(charge_sum - formal_sum) > 0.01:
+            raise ValueError(
+                f"Partial charges of molecule {mol.to_smiles(explicit_hydrogens=False)} "
+                f"sum to {charge_sum:.4f} e, which differs from its formal charge of "
+                f"{formal_sum:.4f} e by more than 0.01 e. Library charges would not "
+                "produce an integral net charge. Please ensure that the partial charges "
+                "are consistent with the formal charge of the molecule."
+            )
+        charges_per_mol.append(charges)
+
+    # Strip stereochemistry so generated SMARTS match either stereoisomer (atom index
+    # order is preserved by the copy, so the captured charges stay aligned).
+    mols_for_typing = [_remove_stereochemical_information(mol) for mol in mols]
+
+    ff_copy = copy.deepcopy(force_field)
+    parameter_handler = ff_copy.get_parameter_handler("LibraryCharges")
+
+    # Collect, for each unique whole-molecule SMARTS, the charges of every atom that
+    # produces it (across all molecules) so they can be averaged. The dict preserves
+    # insertion order, giving deterministic parameter ordering.
+    charges_by_smarts: dict[str, list[float]] = defaultdict(list)
+
+    for mol, charges in zip(mols_for_typing, charges_per_mol, strict=True):
+        for atom_index in range(mol.n_atoms):
+            smarts = _create_smarts(mol, (atom_index,), max_extend_distance=-1)
+            charges_by_smarts[smarts].append(
+                float(charges[atom_index].m_as(unit.elementary_charge))
+            )
+
+    logger.info(
+        f"Generated {len(charges_by_smarts)} bespoke library charge SMARTS patterns "
+        f"across {len(mols_for_typing)} molecules."
+    )
+
+    handler_copy = copy.deepcopy(parameter_handler)
+
+    for smarts, atom_charges in charges_by_smarts.items():
+        mean_charge = sum(atom_charges) / len(atom_charges)
+
+        counter = len(handler_copy.parameters) + 1
+        new_param_dict = {
+            "smirks": smarts,
+            "charge1": mean_charge * unit.elementary_charge,
+            "id": f"l-bespoke-{counter}",  # l for library charge
+        }
+
+        _add_parameter_with_overwrite(handler_copy, new_param_dict)
+
+    ff_copy.deregister_parameter_handler("LibraryCharges")
+    ff_copy.register_parameter_handler(handler_copy)
+
+    # Remove any redundant parameters that are not used by any molecule
+    ff_copy = _remove_redundant_smarts(mols_for_typing, ff_copy, id_substring="bespoke")
+
+    return ff_copy