Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 27 additions & 6 deletions pyxdm/core/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
-------
Expand All @@ -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

Expand Down
65 changes: 47 additions & 18 deletions pyxdm/partitioning/partitioning.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
"""
Expand Down Expand Up @@ -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()

Expand 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:
"""
Expand Down Expand Up @@ -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()

Expand 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.
Expand All @@ -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
-------
Expand All @@ -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:
"""
Expand Down Expand Up @@ -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,
)
Expand Down Expand Up @@ -364,6 +384,7 @@ def __init__(
self,
maxiter: int = 500,
threshold: float = 1e-6,
lmax: int = 3,
) -> None:
"""
Initialize Iterative Stockholder partitioning.
Expand All @@ -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:
"""
Expand All @@ -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,
)
Expand All @@ -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
Expand All @@ -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.
Expand All @@ -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,
)
Expand Down Expand Up @@ -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
-------
Expand All @@ -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)
Expand Down
Loading