From d7eb0ba29bfe3bebfded5660d16d6e3dc702e29b Mon Sep 17 00:00:00 2001 From: Joao Morado Date: Wed, 11 Nov 2026 00:32:25 +0100 Subject: [PATCH] Read external partition weights --- README.md | 17 ++- pyxdm/cli.py | 15 ++- pyxdm/core/core.py | 2 +- pyxdm/core/session.py | 6 +- pyxdm/partitioning/partitioning.py | 164 ++++++++++++++++++++++++++++- pyxdm/tests/test_cli.py | 47 +++++++++ pyxdm/tests/test_partitioning.py | 36 +++++++ 7 files changed, 281 insertions(+), 6 deletions(-) diff --git a/README.md b/README.md index 22c9bca..57adb0e 100644 --- a/README.md +++ b/README.md @@ -52,7 +52,8 @@ pyxdm [optional] - ``: Path to the wavefunction file (Molden, WFN, etc.) ##### Optional -- `--mesh `: Optional postg mesh file for integration grid +- `--mesh `: Integration mesh read from a file instead of a HORTON grid (see [Mesh and weight files](#mesh-and-weight-files)) +- `--weights `: Atomic weights on the `--mesh` points read from a file, used as the `external` partitioning scheme (see [Mesh and weight files](#mesh-and-weight-files)) - `--scheme `: Partitioning scheme(s) to use, separated by comma. If not specified, all available schemes are calculated. Available schemes: `becke`, `hirshfeld`, `hirshfeld_i`, `is`, `mbis` - `--proatomdb `: Path to proatom database file (required for Hirshfeld schemes) - `--grid `: Grid quality for numerical integration. Options: `coarse`, `medium`, `fine`, `veryfine`, `ultrafine`, `insane` (default: `fine`) @@ -73,6 +74,20 @@ Calculate multiple schemes with custom grid granularity: pyxdm orca.molden.input --scheme mbis,becke --grid ultrafine ``` +#### Mesh and weight files + +`--mesh` and `--weights` are plain text files; lines starting with `#` are ignored. + +- **Mesh:** one integration point per line, `x y z w`, with the Cartesian coordinates in bohr and the integration + weight, so that `sum_i w_i f(r_i)` approximates the integral of `f` over all space. +- **Weights:** one line per mesh point, in the same order as the mesh, and one column per atom, in the order of the + wavefunction file. Each column holds the weight of that atom at each point; at every point the weights of all atoms + usually sum to 1. + +```bash +pyxdm molecule.molden.input --mesh mesh.txt --weights weights.txt +``` + ### Python API You can also use PyXDM as a library in your own Python scripts: diff --git a/pyxdm/cli.py b/pyxdm/cli.py index f7fa21c..75a62f5 100644 --- a/pyxdm/cli.py +++ b/pyxdm/cli.py @@ -32,7 +32,12 @@ def create_argument_parser() -> argparse.ArgumentParser: parser.add_argument("wfn_file", help="Wavefunction file (molden, wfn, etc.)") - parser.add_argument("--mesh", help="Optional postg mesh file") + parser.add_argument("--mesh", help="Optional mesh file: one point per line, x y z (bohr) and integration weight w") + + parser.add_argument( + "--weights", + help="Atomic weights on the --mesh points: one line per point, in mesh order, and one column per atom; selects the 'external' scheme", + ) parser.add_argument( "--xdm-moments", @@ -106,6 +111,11 @@ def main() -> None: session.setup_grid(args.mesh, grid_definition=args.grid) session.setup_calculator() + if args.weights and not args.mesh: + raise ValueError("--weights requires --mesh") + if args.weights and args.scheme is None: + args.scheme = ["external"] + if args.scheme is not None: for scheme in args.scheme: if scheme not in PartitioningSchemeFactory.available_schemes(): @@ -113,11 +123,12 @@ def main() -> None: f"Unknown partitioning scheme: {scheme}. Available schemes are: {PartitioningSchemeFactory.available_schemes()}" ) else: - args.scheme = PartitioningSchemeFactory.available_schemes() + args.scheme = [scheme for scheme in PartitioningSchemeFactory.available_schemes() if scheme != "external"] session.setup_partition_schemes( args.scheme, proatomdb=args.proatomdb, + weights_file=args.weights, ) at_symbols = [get_atomic_symbol(num) for num in session.mol.numbers] diff --git a/pyxdm/core/core.py b/pyxdm/core/core.py index 36abdd5..86b1297 100644 --- a/pyxdm/core/core.py +++ b/pyxdm/core/core.py @@ -16,7 +16,7 @@ class XDMCalculator: """Core calculator for XDM multipole moments.""" - _MOLECULAR_GRID_SCHEMES = ["hirshfeld", "hirshfeld-i", "becke"] + _MOLECULAR_GRID_SCHEMES = ["hirshfeld", "hirshfeld-i", "becke", "external"] def __init__(self, mol: Any) -> None: self.mol = mol diff --git a/pyxdm/core/session.py b/pyxdm/core/session.py index 4ffc23c..c8322c5 100644 --- a/pyxdm/core/session.py +++ b/pyxdm/core/session.py @@ -128,7 +128,7 @@ def setup_calculator(self) -> None: self.calculator = XDMCalculator(self.mol) logger.debug("XDM calculator initialized") - def setup_partition_schemes(self, schemes, proatomdb: Optional[str] = None) -> dict: + def setup_partition_schemes(self, schemes, proatomdb: Optional[str] = None, weights_file: Optional[str] = None) -> dict: """ Setup partitioning schemes for the session. @@ -154,6 +154,8 @@ def setup_partition_schemes(self, schemes, proatomdb: Optional[str] = None) -> d proatomdb : str, optional Path to proatom database for Hirshfeld-based schemes (hirshfeld, hirshfeld-i) + weights_file : str, optional + Path to the atomic weights for the external scheme Returns ------- @@ -173,6 +175,8 @@ def setup_partition_schemes(self, schemes, proatomdb: Optional[str] = None) -> d scheme_kwargs = config.copy() if config else {} if proatomdb: scheme_kwargs["proatom_db"] = proatomdb + if weights_file: + scheme_kwargs["weights_file"] = weights_file partitioning = PartitioningSchemeFactory.create_scheme(scheme, **scheme_kwargs) diff --git a/pyxdm/partitioning/partitioning.py b/pyxdm/partitioning/partitioning.py index 9a1664d..b10f240 100644 --- a/pyxdm/partitioning/partitioning.py +++ b/pyxdm/partitioning/partitioning.py @@ -5,6 +5,8 @@ from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Optional, Union +import numpy as np + from ..grids import CustomGrid if TYPE_CHECKING: @@ -503,6 +505,161 @@ def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None: self._partition_obj.name = self.NAME +class ExternalPartition: + """ + Partition object holding atomic weights read from a file, all on one molecular mesh. + + Attributes + ---------- + name : str + Partition name, "external" + grid : Any + Molecular mesh shared by all atoms + weights : np.ndarray + Atomic weights of shape (n_points, n_atoms) + natom : int + Number of atoms + cache : ExternalPartition + The partition itself, which provides load() and keys() like a HORTON cache + """ + + def __init__(self, grid: Any, weights: np.ndarray) -> None: + """ + Initialize the partition. + + Parameters + ---------- + grid : Any + Molecular mesh shared by all atoms + weights : np.ndarray + Atomic weights of shape (n_points, n_atoms) + + Returns + ------- + None + """ + self.name = ExternalPartitioning.NAME + self.grid = grid + self.weights = weights + self.natom = weights.shape[1] + self.cache = self + + def get_grid(self, index: int) -> Any: + """ + Get the grid of an atom, the molecular mesh. + + Parameters + ---------- + index : int + Atom index + + Returns + ------- + Any + Molecular mesh + """ + return self.grid + + def load(self, key: str, index: int) -> np.ndarray: + """ + Load the weights of an atom. + + Parameters + ---------- + key : str + Cache key, "at_weights" + index : int + Atom index + + Returns + ------- + np.ndarray + Weights of the atom on the mesh + """ + if key != "at_weights": + raise KeyError(key) + return self.weights[:, index] + + def keys(self) -> list[tuple[str, int]]: + """ + List the stored items, as a HORTON cache does. + + Returns + ------- + list[tuple[str, int]] + ("at_weights", index) for each atom + """ + return [("at_weights", i) for i in range(self.natom)] + + def __getitem__(self, key: tuple[str, int]) -> np.ndarray: + """ + Get a stored item by its cache key. + + Parameters + ---------- + key : tuple[str, int] + ("at_weights", index) + + Returns + ------- + np.ndarray + Weights of the atom on the mesh + """ + return self.load(*key) + + +class ExternalPartitioning(PartitioningScheme): + """ + Atomic weights read from a file. + + The file has one line per mesh point, in the order of the mesh (--mesh), and one column per atom, in the order of + the wavefunction file. Lines starting with "#" are ignored. + """ + + NAME: str = "external" + + def __init__(self, weights_file: Optional[str] = None) -> None: + """ + Initialize external partitioning. + + Parameters + ---------- + weights_file : Optional[str], default=None + Path to the weights file + + Returns + ------- + None + """ + super().__init__() + self.weights_file = weights_file + + def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None: + """ + Read the atomic weights and create the partition object. + + Parameters + ---------- + mol : Any + Molecule object + grid : Union[CustomGrid, Any] + Mesh the weights are given on + + Returns + ------- + None + """ + if self.weights_file is None: + raise ValueError("External partitioning requires a weights file. Use --weights argument.") + + weights = np.loadtxt(self.weights_file, comments="#", ndmin=2) + n_points = len(grid.points) + if weights.shape != (n_points, mol.natom): + raise ValueError(f"Weights file has shape {weights.shape}, expected ({n_points}, {mol.natom}) for this mesh and molecule") + + self._partition_obj = ExternalPartition(grid, weights) + + class PartitioningSchemeFactory: """ Factory for creating partitioning schemes. @@ -517,6 +674,7 @@ class PartitioningSchemeFactory: HirshfeldPartitioning.NAME: HirshfeldPartitioning, HirshfeldIPartitioning.NAME: HirshfeldIPartitioning, IterativeStockholderPartitioning.NAME: IterativeStockholderPartitioning, + ExternalPartitioning.NAME: ExternalPartitioning, } @classmethod @@ -529,7 +687,7 @@ def create_scheme(cls, scheme_name: str, **kwargs: Any) -> "PartitioningScheme": scheme_name : str Name of the partitioning scheme to create. Available schemes: 'mbis', 'becke', 'hirshfeld', 'hirshfeld-i', - 'iterative-stockholder' + 'iterative-stockholder', 'external' **kwargs Additional keyword arguments passed to the scheme constructor. Different schemes accept different parameters: @@ -538,6 +696,7 @@ def create_scheme(cls, scheme_name: str, **kwargs: Any) -> "PartitioningScheme": - hirshfeld: proatom_db, lmax - hirshfeld-i: proatom_db, lmax, maxiter, threshold - iterative-stockholder: lmax, maxiter, threshold + - external: weights_file Returns ------- @@ -563,6 +722,9 @@ def create_scheme(cls, scheme_name: str, **kwargs: Any) -> "PartitioningScheme": elif scheme_name in [HirshfeldIPartitioning.NAME]: # Hirshfeld-I accepts proatom_db, lmax, maxiter, threshold filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["proatom_db", "lmax", "maxiter", "threshold"]} + elif scheme_name in [ExternalPartitioning.NAME]: + # External partitioning only accepts 'weights_file' + filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["weights_file"]} elif scheme_name in [IterativeStockholderPartitioning.NAME, "iterstock", "is"]: # Iterative Stockholder accepts lmax, maxiter, threshold filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["lmax", "maxiter", "threshold"]} diff --git a/pyxdm/tests/test_cli.py b/pyxdm/tests/test_cli.py index dfd9ff5..d65c7cb 100644 --- a/pyxdm/tests/test_cli.py +++ b/pyxdm/tests/test_cli.py @@ -63,3 +63,50 @@ def test_invalid_scheme_exits(self, monkeypatch): monkeypatch.setattr(sys, "argv", [MOLDEN_FILE, "--scheme", "invalid_scheme"]) with pytest.raises(SystemExit): main() + + +def _write_mesh_and_weights(tmp_path: Path) -> tuple[Path, Path, np.ndarray, np.ndarray]: + """Write a water mesh and its Becke weights as text files.""" + import horton as ht + from horton.part import BeckeWPart + + mol = ht.IOData.from_file(MOLDEN_FILE) + grid = ht.BeckeMolGrid(mol.coordinates, mol.numbers, mol.pseudo_numbers, agspec="coarse", mode="keep", random_rotate=False) + density = mol.obasis.compute_grid_density_dm(mol.get_dm_full(), grid.points) + part = BeckeWPart(mol.coordinates, mol.numbers, mol.pseudo_numbers, grid, density, local=False) + part.do_all() + weights = np.column_stack([part.cache.load("at_weights", i) for i in range(mol.natom)]) + + mesh_file, weights_file = tmp_path / "mesh.txt", tmp_path / "weights.txt" + np.savetxt(mesh_file, np.column_stack([grid.points, grid.weights]), header="x y z w") + np.savetxt(weights_file, weights, header="weights") + return mesh_file, weights_file, np.column_stack([grid.points, grid.weights]), weights + + +class TestExternalWeights: + def test_moments_match_direct_calculation(self, monkeypatch, tmp_path): + """--mesh with --weights gives the moments of the same weights computed directly.""" + from ..core import XDMCalculator + from ..grids import CustomGrid + from ..partitioning.partitioning import ExternalPartition + + mesh_file, weights_file, mesh, weights = _write_mesh_and_weights(tmp_path) + out = tmp_path / "water.h5" + run_cli(monkeypatch, [MOLDEN_FILE, "--mesh", str(mesh_file), "--weights", str(weights_file), "-o", str(out)]) + + import horton as ht + + mol = ht.IOData.from_file(MOLDEN_FILE) + grid = CustomGrid(np.ascontiguousarray(mesh[:, :3]), mesh[:, 3]) + expected = XDMCalculator(mol).calculate_xdm_moments(ExternalPartition(grid, weights), grid, order=[1, 2, 3]) + with h5py.File(out, "r") as f: + for key in ("", "", ""): + np.testing.assert_allclose( + f[f"external/xdm/external/xdm_results/{key}"][:], expected["external"]["xdm_results"][key], rtol=1e-6 + ) + + def test_weights_require_mesh(self, monkeypatch, tmp_path): + """--weights without --mesh is rejected.""" + _, weights_file, _, _ = _write_mesh_and_weights(tmp_path) + with pytest.raises(SystemExit): + run_cli(monkeypatch, [MOLDEN_FILE, "--weights", str(weights_file)]) diff --git a/pyxdm/tests/test_partitioning.py b/pyxdm/tests/test_partitioning.py index b86ed06..85a1982 100644 --- a/pyxdm/tests/test_partitioning.py +++ b/pyxdm/tests/test_partitioning.py @@ -218,3 +218,39 @@ def test_get_populations_values_are_floats(self): result = scheme.get_populations(mol, MagicMock()) for pop in result["becke"]["populations"]: assert isinstance(pop, float) + + +class TestExternalPartitioning: + def test_factory_creates_external(self): + scheme = PartitioningSchemeFactory.create_scheme("external", weights_file="w.txt", proatom_db="ignored") + assert scheme.weights_file == "w.txt" + + def test_requires_weights_file(self): + with pytest.raises(ValueError, match="weights file"): + PartitioningSchemeFactory.create_scheme("external").compute_weights(MagicMock(), MagicMock()) + + def test_rejects_wrong_shape(self, tmp_path): + """A weights file that does not match the mesh and molecule is rejected.""" + weights_file = tmp_path / "w.txt" + np.savetxt(weights_file, np.ones((5, 2))) + mol, grid = MagicMock(), MagicMock() + mol.natom = 3 + grid.points = np.zeros((5, 3)) + with pytest.raises(ValueError, match="shape"): + PartitioningSchemeFactory.create_scheme("external", weights_file=str(weights_file)).compute_weights(mol, grid) + + def test_partition_object(self, tmp_path): + """The partition object returns each atom's column and lists them like a HORTON cache.""" + weights_file = tmp_path / "w.txt" + weights = np.arange(10.0).reshape(5, 2) + np.savetxt(weights_file, weights) + mol, grid = MagicMock(), MagicMock() + mol.natom = 2 + grid.points = np.zeros((5, 3)) + scheme = PartitioningSchemeFactory.create_scheme("external", weights_file=str(weights_file)) + scheme.compute_weights(mol, grid) + part = scheme.get_partition_object() + assert part.name == "external" + np.testing.assert_array_equal(part.cache.load("at_weights", 1), weights[:, 1]) + assert part.cache.keys() == [("at_weights", 0), ("at_weights", 1)] + np.testing.assert_array_equal(part[("at_weights", 0)], weights[:, 0])