Add full HORTON partitioning options support

Author: JMorado

pyxdm/core/session.py

@@ -127,16 +127,32 @@ def setup_calculator(self) -> None:
         self.calculator = XDMCalculator(self.mol)
         logger.debug("XDM calculator initialized")
 
-    def setup_partition_schemes(self, schemes: list[str], proatomdb: Optional[str] = None) -> dict:
+    def setup_partition_schemes(self, schemes, proatomdb: Optional[str] = None) -> dict:
         """
         Setup partitioning schemes for the session.
 
         Parameters
         ----------
-        schemes : list of str
-            List of partitioning scheme names to use
+        schemes : list of str or dict
+            List of partitioning scheme names to use, or a dictionary mapping
+            scheme names to their configuration options.
+            
+            If dict, keys are scheme names and values are dicts of kwargs:
+            - mbis: lmax (int, default=3), maxiter (int, default=500), threshold (float, default=1e-6)
+            - becke: lmax (int, default=3), k (int, default=3)
+            - hirshfeld: lmax (int, default=3)
+            - hirshfeld-i: lmax (int, default=3), maxiter (int, default=500), threshold (float, default=1e-6)
+            - iterative-stockholder: lmax (int, default=3), maxiter (int, default=500), threshold (float, default=1e-6)
+            
+            Example:
+                schemes = {
+                    "mbis": {"lmax": 4, "maxiter": 500, "threshold": 1e-6},
+                    "hirshfeld-i": {"lmax": 3, "maxiter": 1000, "threshold": 1e-5},
+                    "becke": {"lmax": 3, "k": 3}
+                }
+                
         proatomdb : str, optional
-            Path to proatom database for Hirshfeld schemes
+            Path to proatom database for Hirshfeld-based schemes (hirshfeld, hirshfeld-i)
 
         Returns
         -------
@@ -146,9 +162,14 @@ def setup_partition_schemes(self, schemes: list[str], proatomdb: Optional[str] =
         self.partitions = {}
         self.partition_schemes = {}
 
-        for scheme in schemes:
+        if isinstance(schemes, dict):
+            scheme_configs = schemes
+        else:
+            scheme_configs = {scheme: {} for scheme in schemes}
+
+        for scheme, config in scheme_configs.items():
             try:
-                scheme_kwargs = {}
+                scheme_kwargs = config.copy() if config else {}
                 if proatomdb:
                     scheme_kwargs["proatom_db"] = proatomdb
 

pyxdm/partitioning/partitioning.py

@@ -138,15 +138,24 @@ class BeckePartitioning(PartitioningScheme):
 
     NAME: str = "becke"
 
-    def __init__(self) -> None:
+    def __init__(self, lmax: int = 3, k: int = 3) -> None:
         """
         Initialize Becke partitioning.
 
+        Parameters
+        ----------
+        lmax : int, default=3
+            Maximum angular momentum for multipole expansion
+        k : int, default=3
+            Order of the polynomials used in the Becke switching function
+
         Returns
         -------
         None
         """
         super().__init__()
+        self.lmax = lmax
+        self.k = k
 
     def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None:
         """
@@ -178,6 +187,8 @@ def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None:
             grid,
             rho_total,
             local=False,
+            lmax=self.lmax,
+            k=self.k,
         )
         becke.do_all()
 
@@ -195,21 +206,24 @@ class HirshfeldPartitioning(PartitioningScheme):
 
     NAME: str = "hirshfeld"
 
-    def __init__(self, proatom_db: Optional[str] = None) -> None:
+    def __init__(self, proatom_db: Optional[str] = None, lmax: int = 3) -> None:
         """
         Initialize Hirshfeld partitioning.
 
         Parameters
         ----------
         proatom_db : Optional[str], default=None
             Path to pro-atom database. If None, uses default database.
+        lmax : int, default=3
+            Maximum angular momentum for multipole expansion
 
         Returns
         -------
         None
         """
         super().__init__()
         self.proatom_db = proatom_db
+        self.lmax = lmax
 
     def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None:
         """
@@ -249,6 +263,7 @@ def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None:
             rho_total,
             proatomdb,
             local=False,
+            lmax=self.lmax,
         )
         hirshfeld.do_all()
 
@@ -272,6 +287,7 @@ def __init__(
         proatom_db: Optional[str] = None,
         maxiter: int = 500,
         threshold: float = 1e-6,
+        lmax: int = 3,
     ) -> None:
         """
         Initialize Hirshfeld-I partitioning.
@@ -284,6 +300,8 @@ def __init__(
             Maximum number of iterations for convergence
         threshold : float, default=1e-6
             Convergence threshold for iterative process
+        lmax : int, default=3
+            Maximum angular momentum for multipole expansion
 
         Returns
         -------
@@ -293,6 +311,7 @@ def __init__(
         self.proatom_db = proatom_db
         self.maxiter = maxiter
         self.threshold = threshold
+        self.lmax = lmax
 
     def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None:
         """
@@ -328,6 +347,7 @@ def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None:
             rho_total,
             proatomdb,
             local=False,
+            lmax=self.lmax,
             maxiter=self.maxiter,
             threshold=self.threshold,
         )
@@ -364,6 +384,7 @@ def __init__(
         self,
         maxiter: int = 500,
         threshold: float = 1e-6,
+        lmax: int = 3,
     ) -> None:
         """
         Initialize Iterative Stockholder partitioning.
@@ -374,10 +395,13 @@ def __init__(
             Maximum number of iterations for self-consistent procedure
         threshold : float, default=1e-6
             Convergence threshold for density changes between iterations
+        lmax : int, default=3
+            Maximum angular momentum for multipole expansion
         """
         super().__init__()
         self.maxiter = maxiter
         self.threshold = threshold
+        self.lmax = lmax
 
     def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None:
         """
@@ -404,6 +428,7 @@ def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None:
             mol.pseudo_numbers,
             grid,
             rho_total,
+            lmax=self.lmax,
             maxiter=self.maxiter,
             threshold=self.threshold,
         )
@@ -425,7 +450,7 @@ class MBISPartitioning(PartitioningScheme):
 
     NAME: str = "mbis"
 
-    def __init__(self, maxiter: int = 500, threshold: float = 1e-6) -> None:
+    def __init__(self, maxiter: int = 500, threshold: float = 1e-6, lmax: int = 3) -> None:
         """Initialize MBIS partitioning.
 
         Parameters
@@ -434,10 +459,13 @@ def __init__(self, maxiter: int = 500, threshold: float = 1e-6) -> None:
             Maximum number of iterations for convergence
         threshold : float, default=1e-6
             Convergence threshold for iterative process
+        lmax : int, default=3
+            Maximum angular momentum for multipole expansion
         """
         super().__init__()
         self.maxiter = maxiter
         self.threshold = threshold
+        self.lmax = lmax
 
     def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None:
         """Create MBIS partition object for grid projection.
@@ -463,6 +491,7 @@ def compute_weights(self, mol: Any, grid: Union[CustomGrid, Any]) -> None:
             mol.pseudo_numbers,
             grid,
             rho_total,
+            lmax=self.lmax,
             maxiter=self.maxiter,
             threshold=self.threshold,
         )
@@ -504,11 +533,11 @@ def create_scheme(cls, scheme_name: str, **kwargs: Any) -> "PartitioningScheme":
         **kwargs
             Additional keyword arguments passed to the scheme constructor.
             Different schemes accept different parameters:
-            - mbis: maxiter, threshold, agspec
-            - becke: (no parameters)
-            - hirshfeld: proatom_db
-            - hirshfeld-i: proatom_db, maxiter, threshold
-            - iterstock/iterative-stockholder/is: maxiter, threshold
+            - mbis: lmax, maxiter, threshold
+            - becke: lmax, k
+            - hirshfeld: proatom_db, lmax
+            - hirshfeld-i: proatom_db, lmax, maxiter, threshold
+            - iterative-stockholder: lmax, maxiter, threshold
 
         Returns
         -------
@@ -526,20 +555,20 @@ def create_scheme(cls, scheme_name: str, **kwargs: Any) -> "PartitioningScheme":
 
         # Filter kwargs based on scheme requirements
         if scheme_name in [BeckePartitioning.NAME]:
-            # Becke doesn't accept any special parameters
-            filtered_kwargs = {}
+            # Becke accepts lmax, k
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["lmax", "k"]}
         elif scheme_name in [HirshfeldPartitioning.NAME]:
-            # Hirshfeld only accepts 'proatom_db'
-            filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["proatom_db"]}
+            # Hirshfeld accepts proatom_db, lmax
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["proatom_db", "lmax"]}
         elif scheme_name in [HirshfeldIPartitioning.NAME]:
-            # Hirshfeld-I accepts proatom_db, maxiter, threshold
-            filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["proatom_db", "maxiter", "threshold"]}
+            # 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 [IterativeStockholderPartitioning.NAME, "iterstock", "is"]:
-            # Iterative Stockholder accepts maxiter, threshold
-            filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["maxiter", "threshold"]}
+            # Iterative Stockholder accepts lmax, maxiter, threshold
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["lmax", "maxiter", "threshold"]}
         else:
-            # MBIS accepts maxiter, threshold, agspec
-            filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["maxiter", "threshold"]}
+            # MBIS accepts lmax, maxiter, threshold
+            filtered_kwargs = {k: v for k, v in kwargs.items() if k in ["lmax", "maxiter", "threshold"]}
 
         scheme_class = cls._schemes[scheme_name]
         result = scheme_class(**filtered_kwargs)