Read external partition weights

README.md

@@ -52,7 +52,8 @@ pyxdm <wfn_file> [optional]
 - `<wfn_file>`: Path to the wavefunction file (Molden, WFN, etc.)
 
 ##### Optional
-- `--mesh <file>`: Optional postg mesh file for integration grid
+- `--mesh <file>`: Integration mesh read from a file instead of a HORTON grid (see [Mesh and weight files](#mesh-and-weight-files))
+- `--weights <file>`: 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 <schemes>`: 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>`: Path to proatom database file (required for Hirshfeld schemes)
 - `--grid <quality>`: 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:

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,18 +111,24 @@ 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():
                     raise ValueError(
                         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]

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

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)
 

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"]}

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 ("<M1^2>", "<M2^2>", "<M3^2>"):
+                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)])