diff --git a/pyxdm/core/session.py b/pyxdm/core/session.py index 6e9911d..2c3193c 100644 --- a/pyxdm/core/session.py +++ b/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 diff --git a/pyxdm/partitioning/partitioning.py b/pyxdm/partitioning/partitioning.py index 202fbf2..9a1664d 100644 --- a/pyxdm/partitioning/partitioning.py +++ b/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,7 +206,7 @@ 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. @@ -203,6 +214,8 @@ def __init__(self, proatom_db: Optional[str] = None) -> None: ---------- 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 ------- @@ -210,6 +223,7 @@ def __init__(self, proatom_db: Optional[str] = None) -> 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)