qrunch.chemistry.ground_state_problem.builders.tools.localization.weighted_pipek_mezey

Weighted Pipek-Mezey localization.

This module implements a weighted variant of the Pipek-Mezey localization functional:

\[\begin{split}F_{\\text{weighted}} = \\sum_{i} \\left( \\sum_{A} w_{A} P_{A}^{(i)} \\right)^2\end{split}\]

where \(w_A\) are atom-specific weights that encourage orbitals to localize preferentially within the embedded region.

Module Attributes

`PopulationCostFunction`	`(mo_coeff, U) -> float`.
`PopulationGradientFunction`	`(mo_coeff, U) -> NDArray`.
`CostFunction`	Callable that takes a unitary rotation matrix U and returns the cost function value.
`GradientFunction`	Callable that takes a unitary rotation U and returns the packed gradient.

Functions

`optimize_orbital_rotation`(n_orbitals, ...[, ...])	Optimize an orbital rotation on the unitary manifold using re-centered L-BFGS-B.
`rotation_from_parameters`(x, n)	Convert packed antisymmetric parameters to a unitary rotation matrix via matrix exponential.
`run_pipek_mezey_optimization`(mo_coeff, *, ...)	Run the orbital-rotation optimizer for a Pipek-Mezey-style localizer.

Classes

`AvoidRegionMixingPipekMezey`	Pipek-Mezey-style localizer that discourages cross-region orbital mixing.
`AvoidRegionMixingPipekMezeyLocalizer`	Pipek-Mezey localizer that discourages cross-region mixing.
`AvoidRegionMixingPipekMezeyOptions`	Configuration for the avoid-cross-region-mixing Pipek-Mezey localizer.
`OrbitalRotationResult`	Result of an orbital rotation optimization.
`PopulationProjectionTensorCalculator`	Population-projection tensor calculator for Pipek-Mezey-style localizers.
`RestrictedAvoidRegionMixingPipekMezeyLocalizer`	Restricted localizer that uses the avoid-cross-region-mixing Pipek-Mezey functional.
`UnrestrictedAvoidRegionMixingPipekMezeyLocalizer`	Unrestricted localizer that uses the avoid-cross-region-mixing Pipek-Mezey functional.

class AvoidRegionMixingPipekMezey

Bases: object

Pipek-Mezey-style localizer that discourages cross-region orbital mixing.

Let

\[E_i = \sum_{A \in \text{embedded}} P_A^{(i)}, \qquad N_i = \sum_{B \notin \text{embedded}} P_B^{(i)},\]

where \(P_A^{(i)}\) is the population of orbital \(i\) on atom \(A\).

Because the populations sum to one per orbital we have \(E_i + N_i = 1\), so the cross-region product \(M_i = E_i N_i = E_i (1 - E_i)\) is zero exactly when the orbital is fully localized either on the embedded region or on the environment. To discourage cross-region mixing this localizer therefore maximizes the functional

\[F_{\text{region}} = \sum_i \left( (1 + \epsilon)\, E_i^{p} + N_i^{p} \right),\]

with exponent \(p \in \{2, 4\}\) and a small positive bias \(\epsilon\) (embedded_bias).

Region separation alone does not require the orbitals to be localized within each region. To additionally reward atomic localization we add a standard per-atom Pipek-Mezey term

\[F_{\text{PM}} = \sum_i \sum_A (P_A^{(i)})^{p}\]

and maximize the convex combination

\[F = (1 - \lambda)\, F_{\text{region}} + \lambda\, F_{\text{PM}},\]

where \(\lambda\) is localization_weight. Setting \(\lambda = 0\) recovers the pure cross-region functional and \(\lambda = 1\) recovers standard Pipek-Mezey.

__init__(molecular_configuration: MolecularConfiguration, options: AvoidRegionMixingPipekMezeyOptions | None = None) → None

Initialize the avoid-cross-region-mixing Pipek-Mezey localizer.

Parameters:

molecular_configuration (MolecularConfiguration) – The molecular configuration with embedded_atoms defined.
options (AvoidRegionMixingPipekMezeyOptions | None) – Configuration options for convergence, population method, and algorithm parameters.

Return type:

None

cost_function(mo_coeff: ndarray[tuple[Any, ...], dtype[float64]], rotation_matrix: ndarray[tuple[Any, ...], dtype[float64]] | None = None) → float

Evaluate the combined cross-region / per-atom Pipek-Mezey cost function.

\(F = (1-\lambda) \sum_i ((1+\epsilon) E_i^{p} + N_i^{p}) + \lambda \sum_i \sum_A (P_A^{(i)})^{p}\) (to be maximized).

Parameters:

mo_coeff (ndarray[tuple[Any, ...], dtype[float64]]) – Base MO coefficients of shape (n_ao, n_mo).
rotation_matrix (ndarray[tuple[Any, ...], dtype[float64]] | None) – Unitary rotation matrix. Defaults to the identity.

Return type:

float

get_gradient(mo_coeff: ndarray[tuple[Any, ...], dtype[float64]], rotation_matrix: ndarray[tuple[Any, ...], dtype[float64]] | None = None) → ndarray[tuple[Any, ...], dtype[float64]]

Compute the packed Riemannian ascent gradient of the combined cost function.

The gradient is derived from the same population tensor used by cost_function(), with the two region groups stacked on top of the per-atom populations and the weights accumulated in _group_weights. This guarantees the analytic gradient matches the numerical gradient of cost_function().

Parameters:

mo_coeff (ndarray[tuple[Any, ...], dtype[float64]]) – Base MO coefficients of shape (n_ao, n_mo).
rotation_matrix (ndarray[tuple[Any, ...], dtype[float64]] | None) – Current unitary rotation matrix.

Return type:

ndarray[tuple[Any, …], dtype[float64]]

localize(mo_coeff: ndarray[tuple[Any, ...], dtype[float64]]) → ndarray[tuple[Any, ...], dtype[float64]]

Run the localization optimization on the unitary manifold.

Parameters:: mo_coeff (ndarray[tuple[Any, ...], dtype[float64]]) – Molecular orbital coefficients of shape (n_ao, n_mo).
Return type:: ndarray[tuple[Any, …], dtype[float64]]

class AvoidRegionMixingPipekMezeyLocalizer

Bases: OrbitalLocalizer

Pipek-Mezey localizer that discourages cross-region mixing.

Maximizes

\[F = \sum_i \left( E_i^{p} + N_i^{p} \right),\]

where \(E_i\) (resp. \(N_i\)) is the total population of orbital \(i\) on the embedded atoms (resp. on the environment atoms) defined by molecular_configuration.embedded_atoms. This is equivalent to running a standard Pipek-Mezey procedure with the atoms coarse-grained into two regions, and produces orbitals that are localized either fully on the embedded region or fully on the environment.

__init__(options: AvoidRegionMixingPipekMezeyOptions | None = None) → None

Initialize the avoid-cross-region-mixing Pipek-Mezey localizer.

Parameters:: options (AvoidRegionMixingPipekMezeyOptions | None) – The algorithmic options.
Return type:: None

run(molecular_configuration: MolecularConfiguration, occupied_molecular_orbitals: RestrictedOccupiedMolecularOrbitals) → RestrictedOccupiedMolecularOrbitals

run(molecular_configuration: MolecularConfiguration, occupied_molecular_orbitals: UnrestrictedOccupiedMolecularOrbitals) → UnrestrictedOccupiedMolecularOrbitals

Run the localization algorithm that discourages cross-region orbital mixing.

Parameters:

molecular_configuration (MolecularConfiguration) – Molecular configuration with embedded_atoms defined.
occupied_molecular_orbitals (RestrictedOccupiedMolecularOrbitals | UnrestrictedOccupiedMolecularOrbitals) – The occupied molecular orbitals to localize.

Return type:

RestrictedOccupiedMolecularOrbitals | UnrestrictedOccupiedMolecularOrbitals

class AvoidRegionMixingPipekMezeyOptions

Bases: DataclassPublicAPI

Configuration for the avoid-cross-region-mixing Pipek-Mezey localizer.

Parameters:

population_method – Population analysis method for computing atomic charges. (default=PopulationMethod.LOWDIN)
embedded_bias – A small positive bias that breaks the symmetry between embedded and environment groups so that ambiguous orbitals (equal population on both regions) preferentially localize on the embedded region. The embedded group’s contribution to the cost function is multiplied by (1 + embedded_bias). (default=0.01)
localization_weight – Convex-combination weight \(\lambda \in [0, 1]\) that mixes a standard per-atom Pipek-Mezey localization term into the cross-region functional. The overall functional becomes \(F = (1-\lambda) F_{\text{region}} + \lambda F_{\text{PM}}\), where \(F_{\text{region}}\) is the embedded/environment two-group term and \(F_{\text{PM}} = \sum_i \sum_A (P_A^{(i)})^p\) is the standard Pipek-Mezey functional. localization_weight=0 recovers the pure cross-region-mixing functional; localization_weight=1 recovers standard Pipek-Mezey. (default=0.2)
convergence_tolerance – Convergence threshold for cost function change. (default=1e-6)
convergence_tolerance_gradient – Convergence threshold for orbital rotation gradient norm. (default=None)
max_cycle – The maximum number of macro iterations. (default=200)
exponent – The power used to define the localization norm. It can be 2 or 4. (default=2)
initial_guess – Initial guess strategy. "molecular" uses mo_coeff, "atomic" uses atomic orbitals. For (near-)symmetric systems (e.g. a homodimer), the atomic initial guess pre-rotates the MOs to align with atomic orbitals and breaks the symmetry that would prevent gradient-based convergence. (default=”atomic”)

__init__(*, population_method: PopulationMethod = PopulationMethod.LOWDIN, embedded_bias: float = 0.01, localization_weight: float = 0.2, convergence_tolerance: float = 1e-06, convergence_tolerance_gradient: float | None = None, max_cycle: int = 200, exponent: Literal[2, 4] = 2, initial_guess: Literal['molecular', 'atomic'] = 'atomic') → None

Parameters:

population_method (PopulationMethod)
embedded_bias (float)
localization_weight (float)
convergence_tolerance (float)
convergence_tolerance_gradient (float | None)
max_cycle (int)
exponent (Literal[2, 4])
initial_guess (Literal['molecular', 'atomic'])

Return type:

None

convergence_tolerance: float = 1e-06

convergence_tolerance_gradient: float | None = None

embedded_bias: float = 0.01

exponent: Literal[2, 4] = 2

initial_guess: Literal['molecular', 'atomic'] = 'atomic'

localization_weight: float = 0.2

max_cycle: int = 200

population_method: PopulationMethod = 'lowdin'

CostFunction

Callable that takes a unitary rotation matrix U and returns the cost function value.

alias of Callable[[ndarray[tuple[Any, …], dtype[float64]]], float]

GradientFunction

Callable that takes a unitary rotation U and returns the packed gradient.

alias of Callable[[ndarray[tuple[Any, …], dtype[float64]]], ndarray[tuple[Any, …], dtype[float64]]]

class OrbitalRotationResult

Bases: object

Result of an orbital rotation optimization.

Parameters:

rotation_matrix – The optimized unitary rotation matrix.
converged – Whether the optimization converged.
cost_function_value – Final cost function value.
gradient_norm – Final gradient norm.
n_iterations – Number of iterations performed.

__init__(rotation_matrix: ndarray[tuple[Any, ...], dtype[float64]], converged: bool, cost_function_value: float, gradient_norm: float, n_iterations: int) → None

Parameters:

rotation_matrix (ndarray[tuple[Any, ...], dtype[float64]])
converged (bool)
cost_function_value (float)
gradient_norm (float)
n_iterations (int)

Return type:

None

converged: bool

cost_function_value: float

gradient_norm: float

n_iterations: int

rotation_matrix: ndarray[tuple[Any, ...], dtype[float64]]

PopulationCostFunction

(mo_coeff, U) -> float.

Type:: Cost function of a Pipek-Mezey-style localizer

alias of Callable[[ndarray[tuple[Any, …], dtype[float64]], ndarray[tuple[Any, …], dtype[float64]]], float]

PopulationGradientFunction

(mo_coeff, U) -> NDArray.

Type:: Packed-gradient function of a Pipek-Mezey-style localizer

alias of Callable[[ndarray[tuple[Any, …], dtype[float64]], ndarray[tuple[Any, …], dtype[float64]]], ndarray[tuple[Any, …], dtype[float64]]]

class PopulationProjectionTensorCalculator

Bases: object

Population-projection tensor calculator for Pipek-Mezey-style localizers.

__init__(molecular_configuration: MolecularConfiguration, population_method: PopulationMethod) → None

Initialize the Population-projection tensor calculator.

Parameters:

molecular_configuration (MolecularConfiguration) – The molecular configuration.
population_method (PopulationMethod) – Population analysis method used to compute atomic populations.

Return type:

None

get_atomic_populations(mo_coeff: ndarray[tuple[Any, ...], dtype[float64]]) → ndarray[tuple[Any, ...], dtype[float64]]

Compute the atomic population projection tensor using precomputed quantities.

For Lowdin, it uses the cached S @ ortho_ao matrix to avoid recomputing the orthogonalization each call.

Parameters:: mo_coeff (ndarray[tuple[Any, ...], dtype[float64]]) – Molecular orbital coefficients of shape (n_ao, n_mo).
Return type:: ndarray[tuple[Any, …], dtype[float64]]

property number_of_atoms: int: Total number of atoms in the molecular configuration.

property overlap: ndarray[tuple[Any, ...], dtype[float64]]: Cached AO overlap matrix of shape (n_ao, n_ao).

class RestrictedAvoidRegionMixingPipekMezeyLocalizer

Bases: object

Restricted localizer that uses the avoid-cross-region-mixing Pipek-Mezey functional.

__init__(options: AvoidRegionMixingPipekMezeyOptions | None = None) → None

Initialize the restricted avoid-cross-region-mixing localizer.

Parameters:: options (AvoidRegionMixingPipekMezeyOptions | None) – Configuration options for the localization.
Return type:: None

run(molecular_configuration: MolecularConfiguration, occupied_molecular_orbitals: RestrictedOccupiedMolecularOrbitals) → RestrictedOccupiedMolecularOrbitals

Localize occupied MOs while penalizing cross-region mixing.

Parameters:

molecular_configuration (MolecularConfiguration) – Molecular configuration with embedded_atoms defined.
occupied_molecular_orbitals (RestrictedOccupiedMolecularOrbitals) – The occupied molecular orbitals to localize.

Return type:

RestrictedOccupiedMolecularOrbitals

class UnrestrictedAvoidRegionMixingPipekMezeyLocalizer

Bases: object

Unrestricted localizer that uses the avoid-cross-region-mixing Pipek-Mezey functional.

__init__(options: AvoidRegionMixingPipekMezeyOptions | None = None) → None

Initialize the unrestricted avoid-cross-region-mixing localizer.

Parameters:: options (AvoidRegionMixingPipekMezeyOptions | None) – Configuration options for the localization.
Return type:: None

run(molecular_configuration: MolecularConfiguration, occupied_molecular_orbitals: UnrestrictedOccupiedMolecularOrbitals) → UnrestrictedOccupiedMolecularOrbitals

Localize occupied alpha/beta MOs while penalizing cross-region mixing.

Parameters:

molecular_configuration (MolecularConfiguration) – Molecular configuration with embedded_atoms defined.
occupied_molecular_orbitals (UnrestrictedOccupiedMolecularOrbitals) – The unrestricted occupied molecular orbitals to localize.

Return type:

UnrestrictedOccupiedMolecularOrbitals

optimize_orbital_rotation(n_orbitals: int, cost_function: Callable[[ndarray[tuple[Any, ...], dtype[float64]]], float], gradient_function: Callable[[ndarray[tuple[Any, ...], dtype[float64]]], ndarray[tuple[Any, ...], dtype[float64]]], initial_rotation: ndarray[tuple[Any, ...], dtype[float64]] | None = None, convergence_tolerance: float = 1e-06, convergence_tolerance_gradient: float | None = None, max_iterations: int = 200) → OrbitalRotationResult

Optimize an orbital rotation on the unitary manifold using re-centered L-BFGS-B.

The algorithm performs repeated short L-BFGS-B optimizations, each starting from x=0. At each macro step:

Define cost/gradient as functions of a small perturbation x around U_current: f(x) = -cost_function(U_current @ expm(A(x))) g(x) = -gradient_function(U_current @ expm(A(x)))
Run L-BFGS-B for a limited number of iterations.
Update U_current = U_current @ expm(A(x_opt)) and restart from x=0.

This re-centering keeps x small, ensuring that the Riemannian gradient (which is exact at x=0) remains a good approximation of the Euclidean gradient w.r.t. x.

Parameters:

n_orbitals (int) – Number of orbitals (dimension of the rotation matrix).
cost_function (Callable[[ndarray[tuple[Any, ...], dtype[float64]]], float]) – Callable that evaluates f(U) given rotation matrix U. Should be maximized.
gradient_function (Callable[[ndarray[tuple[Any, ...], dtype[float64]]], ndarray[tuple[Any, ...], dtype[float64]]]) – Callable that returns the packed gradient given rotation matrix U. The gradient should point in the ascent direction of cost_function (i.e. positive gradient). It will be negated internally for minimization.
initial_rotation (ndarray[tuple[Any, ...], dtype[float64]] | None) – Initial unitary rotation matrix. Defaults to identity.
convergence_tolerance (float) – Convergence threshold for cost function change.
convergence_tolerance_gradient (float | None) – Convergence threshold for gradient norm. If None, set to sqrt(convergence_tolerance * 0.1).
max_iterations (int) – Maximum total number of macro iterations.

Return type:

OrbitalRotationResult

rotation_from_parameters(x: ndarray[tuple[Any, ...], dtype[float64]], n: int) → ndarray[tuple[Any, ...], dtype[float64]]

Convert packed antisymmetric parameters to a unitary rotation matrix via matrix exponential.

The rotation is U = expm(A) where A is the antisymmetric matrix reconstructed from x.

Parameters:

x (ndarray[tuple[Any, ...], dtype[float64]]) – Flat array of rotation parameters, length n*(n-1)/2.
n (int) – Dimension of the resulting rotation matrix.

Return type:

ndarray[tuple[Any, …], dtype[float64]]

run_pipek_mezey_optimization(mo_coeff: ndarray[tuple[Any, ...], dtype[float64]], *, overlap: ndarray[tuple[Any, ...], dtype[float64]], cost_function: Callable[[ndarray[tuple[Any, ...], dtype[float64]], ndarray[tuple[Any, ...], dtype[float64]]], float], gradient_function: Callable[[ndarray[tuple[Any, ...], dtype[float64]], ndarray[tuple[Any, ...], dtype[float64]]], ndarray[tuple[Any, ...], dtype[float64]]], initial_guess: Literal['molecular', 'atomic'], max_cycle: int, convergence_tolerance: float, convergence_tolerance_gradient: float | None) → ndarray[tuple[Any, ...], dtype[float64]]

Run the orbital-rotation optimizer for a Pipek-Mezey-style localizer.

This function decouples the orbital-rotation optimizer from any specific cost functional: the caller injects a cost_function and a gradient_function of signature (mo_coeff, U) -> float and (mo_coeff, U) -> NDArray respectively, where U is the current unitary rotation matrix in MO space.

Parameters:

mo_coeff (ndarray[tuple[Any, ...], dtype[float64]]) – Molecular orbital coefficients of shape (n_ao, n_mo).
overlap (ndarray[tuple[Any, ...], dtype[float64]]) – AO overlap matrix of shape (n_ao, n_ao), required only for the "atomic" initial guess.
cost_function (Callable[[ndarray[tuple[Any, ...], dtype[float64]], ndarray[tuple[Any, ...], dtype[float64]]], float]) – Functional to be maximized on the unitary manifold.
gradient_function (Callable[[ndarray[tuple[Any, ...], dtype[float64]], ndarray[tuple[Any, ...], dtype[float64]]], ndarray[tuple[Any, ...], dtype[float64]]]) – Packed Riemannian ascent gradient of cost_function.
initial_guess (Literal['molecular', 'atomic']) – Initial rotation strategy. "molecular" starts from the identity, "atomic" uses an atomic-orbital projection guess.
max_cycle (int) – Maximum number of macro iterations of the optimizer.
convergence_tolerance (float) – Convergence threshold for the cost-function change.
convergence_tolerance_gradient (float | None) – Convergence threshold for the gradient norm. If None, set to sqrt(convergence_tolerance * 0.1).

Return type:

ndarray[tuple[Any, …], dtype[float64]]