From 189bb5ea3c4ea55986d2a0a16c860a6f5db387c2 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Tue, 15 Jul 2025 00:25:05 -0700 Subject: [PATCH 1/4] Add new analysis class: potential --- src/mdcraft/analysis/__init__.py | 4 +- src/mdcraft/analysis/potential.py | 2055 +++++++++++++++++++++++++++++ 2 files changed, 2058 insertions(+), 1 deletion(-) create mode 100644 src/mdcraft/analysis/potential.py diff --git a/src/mdcraft/analysis/__init__.py b/src/mdcraft/analysis/__init__.py index 4941b82..9b47472 100644 --- a/src/mdcraft/analysis/__init__.py +++ b/src/mdcraft/analysis/__init__.py @@ -1,7 +1,7 @@ """ Simulation state data and trajectory analysis ============================================= -.. moduleauthor:: Benjamin Ye +.. moduleauthor:: Benjamin Ye , Pierre Walker This module provides a variety of classes for analyzing simulation trajectories. @@ -12,6 +12,7 @@ electrostatics, polymer, profile, + potential, reader, structure, thermodynamics, @@ -23,6 +24,7 @@ "electrostatics", "polymer", "profile", + "potential", "reader", "structure", "thermodynamics", diff --git a/src/mdcraft/analysis/potential.py b/src/mdcraft/analysis/potential.py new file mode 100644 index 0000000..c0cb6f3 --- /dev/null +++ b/src/mdcraft/analysis/potential.py @@ -0,0 +1,2055 @@ +""" +Linear profiles +=============== +.. moduleauthor:: Benjamin Ye + +This module contains classes to compute number density and charge +density profiles and related quantities, like surface charge densities +and potential profiles. +""" + +import logging +from numbers import Real +from typing import Any, Union +import warnings + +import MDAnalysis as mda +import numpy as np +from scipy import integrate, sparse + +from .base import Hash, DynamicAnalysisBase +from .. import FOUND_OPENMM, Q_, ureg +from ..algorithm.accelerated import numba_histogram_bin_edges, numba_histogram +from ..algorithm.molecule import center_of_mass +from ..algorithm.topology import unwrap, wrap +from ..algorithm.unit import is_unitless, strip_unit + +if FOUND_OPENMM: + from openmm import unit + +_ELECTROSTATIC_CONVERSION_FACTORS = ( + (1 * ureg.elementary_charge / (ureg.vacuum_permittivity * ureg.angstrom)).m_as( + ureg.volt + ), + 4 * np.pi, +) + + +def calculate_surface_charge_density( + bins: np.ndarray[float], + charge_density_profile: np.ndarray[float], + dielectric: float = None, + *, + L: float = None, + dV: Union[float, np.ndarray[float]] = None, + reduced: bool = False, +) -> Union[float, np.ndarray[float]]: + r""" + Calculates the surface charge density :math:`\sigma_q` using the + charge density profile :math:`\rho_q(z)`. + + The surface charge density is given by + + .. math:: + + \sigma_q=\frac{1}{L}\left( + \varepsilon_0\varepsilon_\mathrm{r}\Delta V + -\int_0^L z\rho_q(z)\,\mathrm{d}z\right) + + where :math:`\varepsilon_0` and :math:`\varepsilon_\mathrm{r}` are + the vacuum and relative permittivities, respectively, + :math:`\Delta V` is the potential difference, and :math:`L` is the + system dimension. + + The first (static) term accounts for the applied potential across + the dielectric medium, while the second (polarization) term accounts + for the distribution of charged species. + + Parameters + ---------- + bins : array-like + Histogram bin centers :math:`z` for the charge density profile + in `charge_density_profile`. + + **Shape**: :math:`(N_\mathrm{bins},)`. + + **Reference unit**: :math:`\mathrm{Å}`. + + charge_density_profile : array-like + Charge density profile :math:`\rho_q(z)`. + + **Shape**: :math:`(N_\mathrm{bins},)` or + :math:`(N_\mathrm{frames},\,N_\mathrm{bins})`. + + **Reference unit**: :math:`\mathrm{e/Å}^{-3}`. + + dielectric : `float`, optional + Relative permittivity or static dielectric constant + :math:`\varepsilon_\mathrm{r}`. Required if `dV` is provided. + + L : `float`, keyword-only, optional + System size :math:`L` in the dimension that `bins` and + `charge_density_profile` were calculated in. If not specified, + it is determined using the first and last values in `bins`, + assuming that the bin centers are uniformly spaced. + + **Reference unit**: :math:`\mathrm{Å}`. + + dV : `float` or array-like, keyword-only, optional + Potential difference :math:`\Delta V` across the system + dimension specified in `axis`. If not provided, only the + polarizable part (second term) of the surface charge density is + calculated. + + **Shape**: Scalar or :math:`(N_\mathrm{frames},)`. + + **Reference unit**: :math:`\mathrm{V}`. + + reduced : `bool`, keyword-only, default: :code:`False` + Specifies whether the data is in reduced units. + + Returns + ------- + surface_charge_density : `numpy.ndarray` + Surface charge density :math:`\sigma_q`. + + **Shape**: Scalar or :math:`(N_\mathrm{frames},)`. + + **Reference unit**: :math:`\mathrm{e/Å}^2`. + """ + + # Check input array shapes + if bins.ndim != 1: + raise ValueError("'bins' must be a one-dimensional array.") + if charge_density_profile.ndim not in {1, 2}: + emsg = "'charge_density_profile' must be a one- or " "two-dimensional array." + raise ValueError(emsg) + if bins.shape[0] != charge_density_profile.shape[-1]: + emsg = ( + "'bins' and 'charge_density_profile' must have the same" + "length in the last dimension." + ) + raise ValueError(emsg) + if dV is not None: + if dielectric is None: + raise ValueError("'dielectric' must be specified when 'dV' is.") + if not isinstance(dV, Real) and dV.shape[0] != charge_density_profile.shape[0]: + emsg = ( + "The number of potential differences in 'dVs' must " + "be equal to the number of frames in " + "'charge_density_profile'." + ) + raise ValueError(emsg) + + # Determine system size, if not provided + dz = bins[1] - bins[0] + if (uniform := np.allclose(np.diff(bins), dz)) and not np.isclose(bins[0], dz / 2): + wmsg = ( + "'bins' currently does not start at zero, so it will " + f"be shifted left by {(shift := bins[0] - dz / 2):.6g}." + ) + warnings.warn(wmsg) + bins -= shift + if L is None: + if not uniform: + emsg = "'L' must be specified when 'bins' is not uniformly spaced." + raise ValueError(emsg) + L = bins[-1] - bins[0] + dz + + y = -integrate.trapezoid(bins * charge_density_profile, bins) + if dV is not None: + y += dielectric * dV / _ELECTROSTATIC_CONVERSION_FACTORS[reduced] + return y / L + + +def calculate_potential_profile( + bins: np.ndarray[float], + charge_density_profile: np.ndarray[float], + dielectric: float, + *, + L: float = None, + sigma_q: Union[float, np.ndarray[float]] = None, + dV: Union[float, np.ndarray[float]] = None, + threshold: float = 1e-5, + V0: Union[float, np.ndarray[float]] = 0, + method: str = "integral", + pbc: bool = False, + reduced: bool = False, +) -> np.ndarray[float]: + r""" + Calculates the potential profile :math:`\Psi(z)` using the charge + density profile :math:`\rho_q(z)` by numerically solving Poisson's + equation for electrostatics. + + Poisson's equation is given by + + .. math:: + + \varepsilon_0\varepsilon_\mathrm{r}\nabla^2\Psi(z)=-\rho_q(z) + + where :math:`\varepsilon_0` and :math:`\varepsilon_\mathrm{r}` are + the vacuum and relative permittivities, respectively. + + The boundary conditions (BCs) are + + .. math:: + + \left.\frac{\partial\Psi}{\partial z}\right\rvert_{z=0}& + =-\frac{\sigma_q}{\varepsilon_0\varepsilon_\mathrm{r}}\\ + \left.\Psi\right\rvert_{z=0}&=\Psi_0 + + The first BC is used to ensure that the electric field in the bulk + of the solution is zero, while the second BC is used to set the + potential on the left electrode. + + Poisson's equation can be evaluated by using the trapezoidal rule + to numerically integrate the charge density profile twice: + + 1. Integrate the charge density profile. + 2. Apply the first BC by subtracting :math:`\sigma_q` from all points. + 3. Integrate the profile from Step 2. + 4. Divide by :math:`-\varepsilon_0\varepsilon_\mathrm{r}`. + 5. Apply the second BC by adding :math:`\Psi_0` to all points. + + This method is fast but requires many histogram bins to accurately + resolve the potential profile. + + Alternatively, Poisson's equation can be written as a system of + linear equations + + .. math:: + + A\mathbf{\Psi}=\mathbf{b} + + using second-order finite differences. :math:`A` is a tridiagonal + matrix and :math:`\mathbf{b}` is a vector containing the charge + density profile, with boundary conditions applied in the first and + last rows. + + The inner equations are given by + + .. math:: + + \Psi_i^{''}=\frac{\Psi_{i-1}-2\Psi_i+\Psi_{i+1}}{h^2} + =-\frac{1}{\varepsilon_0\varepsilon_\mathrm{r}}\rho_{q,\,i} + + where :math:`i` is the bin index and :math:`h` is the bin width. + + In the case of periodic boundary conditions, the first and last + equations are given by + + .. math:: + + \Psi_0^{''}&=\frac{\Psi_{N-1}-2\Psi_0+\Psi_1}{h^2} + =-\frac{1}{\varepsilon_0\varepsilon_\mathrm{r}}\rho_{q,\,0}\\ + \Psi_{N-1}^{''}&=\frac{\Psi_{N-2}-2\Psi_{N-1}+\Psi_0}{h^2} + =-\frac{1}{\varepsilon_0\varepsilon_\mathrm{r}}\rho_{q,\,N-1} + + When the system has slab geometry, the boundary conditions are + implemented via + + .. math:: + + \Psi_0&=\Psi_0\\ + \Psi_0^\prime&=\frac{-3\Psi_0+4\Psi_1-\Psi_2}{2h} + =\frac{\sigma_q}{\varepsilon_0\varepsilon_\mathrm{r}} + + This method is slower but can be more accurate even with fewer + histogram bins for bulk systems with periodic boundary conditions. + + Parameters + ---------- + bins : array-like + Histogram bin centers :math:`z` for the charge density profile + in `charge_density_profile`. + + **Shape**: :math:`(N_\mathrm{bins},)`. + + **Reference unit**: :math:`\mathrm{Å}`. + + charge_density_profile : array-like + Charge density profile :math:`\rho_q(z)`. + + **Shape**: :math:`(N_\mathrm{bins},)` or + :math:`(N_\mathrm{frames},\,N_\mathrm{bins})`. + + **Reference unit**: :math:`\mathrm{e/Å}^{-3}`. + + dielectric : `float` + Relative permittivity or static dielectric constant + :math:`\varepsilon_\mathrm{r}`. + + L : `float`, keyword-only, optional + System size in the dimension that `bins` and + `charge_density_profile` were calculated in. If not specified, + it is determined using the first and last values in `bins`, + assuming that the bin centers are uniformly spaced. + + **Reference unit**: :math:`\mathrm{Å}`. + + sigma_q : `float` or array-like, keyword-only, optional + Surface charge density :math:`\sigma_q`. Used as a boundary + condition. If not provided, it is determined using `dV` and + `charge_density_profile`. If `dV` is also not provided, the + average value in the center of the integrated charge density + profile is used if :code:`method="integral"`. + + .. note:: + + :math:`\sigma_q` and :math:`\Delta\Psi` should have the same + sign. + + **Shape**: Scalar or :math:`(N_\mathrm{frames},)`. + + **Reference unit**: :math:`\mathrm{e/Å^2}`. + + dV : `float` or array-like, keyword-only, optional + Potential difference :math:`\Delta\Psi` across the system + dimension specified in `axis`. Only used if `sigma_q` was not + provided. + + **Shape**: Scalar or :math:`(N_\mathrm{frames},)`. + + **Reference unit**: :math:`\mathrm{V}`. + + V0 : `float` or array-like, keyword-only, default: :code:`0` + Potential :math:`\Psi_0` at the left boundary. + + **Shape**: Scalar or :math:`(N_\mathrm{frames},)`. + + **Reference unit**: :math:`\mathrm{V}`. + + threshold : `float`, keyword-only, default: :code:`1e-5` + Threshold for determining the plateau regions in the centers of + the integrated charge density profiles to be used as estimates + of `sigma_q`. Only used if `sigma_q` was not provided and + cannot be calculated using `dV` and `charge_density_profile`. + + method : `str`, keyword-only, default: :code:`"integral"` + Method used to calculate the potential profiles. + + **Valid values**: :code:`"integral"`, :code:`"matrix"`. + + pbc : `bool`, keyword-only, default: :code:`False` + Specifies whether the axis has periodic boundary conditions. + Only used when :code:`method="matrix"`. + + reduced : `bool`, keyword-only, default: :code:`False` + Specifies whether the data is in reduced units. + + Returns + ------- + potential : `numpy.ndarray` + Potential profile :math:`\Psi(z)`. + + **Shape**: Same as `charge_density_profile`. + + **Reference unit**: :math:`\mathrm{V}`. + """ + + ecf = _ELECTROSTATIC_CONVERSION_FACTORS[reduced] + + # Check input array shapes + if bins.ndim != 1: + raise ValueError("'bins' must be a one-dimensional array.") + if charge_density_profile.ndim not in {1, 2}: + emsg = "'charge_density_profile' must be a one- or " "two-dimensional array." + raise ValueError(emsg) + if bins.shape[0] != charge_density_profile.shape[-1]: + emsg = ( + "'bins' and 'charge_density_profile' must have the same" + "length in the last dimension." + ) + raise ValueError(emsg) + if ( + sigma_q is not None + and not isinstance(sigma_q, Real) + and sigma_q.shape[0] != charge_density_profile.shape[0] + ): + emsg = ( + "The number of surface charge densities in 'sigmas_q' " + "must be equal to the number of frames in " + "'charge_density_profile'." + ) + raise ValueError(emsg) + if ( + dV is not None + and not isinstance(dV, Real) + and dV.shape[0] != charge_density_profile.shape[0] + ): + emsg = ( + "The number of potential differences in 'dVs' must be " + "equal to the number of frames in " + "'charge_density_profile'." + ) + raise ValueError(emsg) + if ( + V0 is not None + and not isinstance(V0, Real) + and V0.shape[0] != charge_density_profile.shape[0] + ): + emsg = ( + "The number of potentials in 'V0s' must be equal to " + "the number of frames in 'charge_density_profile'." + ) + raise ValueError(emsg) + + # Determine system size, if not provided + dz = bins[1] - bins[0] + if (uniform := np.allclose(np.diff(bins), dz)) and not np.isclose(bins[0], dz / 2): + wmsg = ( + "'bins' currently does not start at zero, so it will " + f"be shifted left by {(shift := bins[0] - dz / 2):.6g}." + ) + warnings.warn(wmsg) + bins -= shift + if L is None: + if not uniform: + emsg = "'L' must be specified when 'bins' is not uniformly spaced." + raise ValueError(emsg) + L = bins[-1] - bins[0] + dz + + # Calculate surface charge density, if not provided + if sigma_q is None and dV is not None: + sigma_q = ( + dielectric * dV / ecf + - integrate.trapezoid(bins * charge_density_profile, bins) + ) / L + + if method == "integral": + potential = integrate.cumulative_trapezoid( + charge_density_profile, bins, initial=0 + ) + + if sigma_q is None: + wmsg = ( + "No surface charge density information. The value " + "will be extracted from the integrated charge " + "density profile, which may be inaccurate due to " + "numerical errors." + ) + warnings.warn(wmsg) + + # Get surface charge density from the integrated charge + # density profile + cut_indices = ( + np.where( + np.diff( + np.abs( + np.gradient( + potential + if potential.ndim == 1 + else potential.mean(axis=0) + ) + ) + < threshold + ) + )[0] + + 1 + ) + if len(cut_indices) == 0: + logging.warning( + "No bulk plateau region found in the integrated " + "charge density profile. The average value over " + "the entire profile will be used." + ) + sigma_q = potential.mean(axis=-1, keepdims=True) + else: + target_index = potential.shape[-1] // 2 + if potential.ndim == 1: + sigma_q = potential[ + cut_indices[cut_indices <= target_index][-1] : cut_indices[ + cut_indices >= target_index + ][0] + ].mean() + else: + sigma_q = potential[ + :, + cut_indices[cut_indices <= target_index][-1] : cut_indices[ + cut_indices >= target_index + ][0], + ].mean(axis=-1, keepdims=True) + + return ( + V0 + - ecf + * (integrate.cumulative_trapezoid(potential - sigma_q, bins, initial=0)) + / dielectric + ) + + elif method == "matrix": + if sigma_q is None: + emsg = ( + "No surface charge density information. Either " + "'sigma_q' or 'dV' must be provided when " + "method='matrix'." + ) + raise ValueError(emsg) + + if not np.allclose(np.diff(bins), dz): + raise ValueError("'bins' must be uniformly spaced.") + + # Set up matrix and load vector for second-order finite + # difference method + N = len(bins) + A = sparse.diags((1, -2, 1), (-1, 0, 1), shape=(N, N), format="csc") + b = charge_density_profile.copy().T + with warnings.catch_warnings(): + warnings.simplefilter("ignore", category=sparse.SparseEfficiencyWarning) + + # Apply boundary conditions and solve + if pbc: + A[0, -1] = A[-1, 0] = 1 + b *= -ecf * dz**2 / dielectric + psi = np.empty_like(b) + psi[1:] = sparse.linalg.spsolve(A[1:, 1:], b[1:]) + psi[0] = psi[-1] + return psi + else: + A[0, :3] = -1.5, 2, -0.5 + A[-1, 0] = 1 + A[-1, -2:] = 0 + b[0] = ecf * dz * sigma_q / dielectric + b[1:-1] *= -ecf * dz**2 / dielectric + b[-1] = V0 + return sparse.linalg.spsolve(A, b) + + +def center_of_charge(group: mda.AtomGroup = None, + grouping: str = None, + *, + charges: Union[np.ndarray[float], list[np.ndarray[float]]] = None, + positions: Union[np.ndarray[float], list[np.ndarray[float]]] = None, + images: Union[np.ndarray[int], list[np.ndarray[int]]] = None, + dimensions: np.ndarray[float] = None, + n_groups: int = None, + raw: bool = False, + ) -> Union[ + np.ndarray[float], tuple[np.ndarray[float], np.ndarray[float], np.ndarray[float]] + ]: + # Check whether grouping is valid + if grouping not in {None, "residues", "segments"}: + emsg = ( + f"Invalid grouping: '{grouping}'. Valid options are " + "None, 'residues', and 'segments'." + ) + raise ValueError(emsg) + + # Get system dimensions if image flags are provided + if images is not None: + if dimensions is None: + try: + dimensions = group.dimensions[:3] + except (NameError, TypeError): + emsg = ( + "Image flags were provided, but no system " + "dimensions were provided or found in the " + "trajectory." + ) + raise ValueError(emsg) + else: + dimensions = np.asarray(dimensions) + + # Get particle masses and positions from the trajectory, if + # necessary + missing = (charges is None, positions is None) + if any(missing): + if group is None: + emsg = ( + "Either a group of atoms or atom positions and " + "charges must be provided." + ) + raise ValueError(emsg) + + # Check whether the groups have equal numbers of atoms + if grouping is None: + same = True + else: + groups = getattr(group, grouping) + + # Calculate and return the centers of mass for different + # groups here if unwrapping and the mass and position arrays + # are not needed + if ( + not ( + same := all( + g.atoms.n_atoms == groups[0].atoms.n_atoms for g in groups + ) + ) + and images is None + and not raw + ): + return np.array([g.atoms.center_of_mass() for g in groups]) + + # Get and unwrap particle positions, if necessary + if missing[1]: + positions = group.positions + if images is not None: + positions += images * dimensions[:3] + + # Get particle charges and ensure correct dimensionality, if + # necessary + if same: + if missing[0]: + charges = group.charges + if grouping is not None or n_groups is not None: + shape = (n_groups or getattr(group, f"n_{grouping}"), -1, 3) + charges = charges.reshape(shape[:-1]) + positions = positions.reshape(shape) + else: + if missing[0]: + charges = [g.atoms.charges for g in groups] + if missing[1]: + positions = [positions[g.atoms.ix] for g in groups] + else: + + # Try to convert arrays to NumPy arrays if they are not already + # to take advantage of vectorized operations later + try: + positions = np.asarray(positions) + charges = np.asarray(charges) + if images is not None: + positions += images + except ValueError: + pass + if type(charges) is not type(positions): + emsg = ( + "The shapes of the arrays containing the charges " + "and positions are incompatible." + ) + raise ValueError(emsg) + if images is not None and type(images) is not type(positions): + emsg = ( + "The shapes of the arrays containing the positions " + "and image flags are incompatible." + ) + raise ValueError(emsg) + + # Calculate the dipole for the specified grouping + if isinstance(positions, np.ndarray): + + # Reshape the mass and position arrays based on the specified + # number of groups + if n_groups is not None: + charges = charges.reshape((n_groups, -1)) + positions = positions.reshape((n_groups, -1, 3)) + + # dipole = sum(q_i * r_i) and should be a vector of shape (3,) + coc = np.einsum("...a,...ad->...d", np.abs(charges), positions) / np.abs(charges).sum( + axis=-1, keepdims=True + ) + else: + if images is not None: + for j, (p, i) in enumerate(zip(positions, images)): + positions[j] = p + i * dimensions + coc = np.array([ + np.sum(np.abs(c) * p, axis=0) / np.sum(np.abs(c)) for c, p in zip(charges, positions) + ]) + + # Return raw charges and positions, if desired + if raw: + return coc, charges, positions + + return coc + +def molecular_dipole(group: mda.AtomGroup = None, + grouping: str = None, + *, + charges: Union[np.ndarray[float], list[np.ndarray[float]]] = None, + positions: Union[np.ndarray[float], list[np.ndarray[float]]] = None, + images: Union[np.ndarray[int], list[np.ndarray[int]]] = None, + dimensions: np.ndarray[float] = None, + n_groups: int = None, + raw: bool = False, + ) -> Union[ + np.ndarray[float], tuple[np.ndarray[float], np.ndarray[float], np.ndarray[float]] + ]: + # Check whether grouping is valid + if grouping not in {None, "residues", "segments"}: + emsg = ( + f"Invalid grouping: '{grouping}'. Valid options are " + "None, 'residues', and 'segments'." + ) + raise ValueError(emsg) + + # Get system dimensions if image flags are provided + if images is not None: + if dimensions is None: + try: + dimensions = group.dimensions[:3] + except (NameError, TypeError): + emsg = ( + "Image flags were provided, but no system " + "dimensions were provided or found in the " + "trajectory." + ) + raise ValueError(emsg) + else: + dimensions = np.asarray(dimensions) + + # Get particle masses and positions from the trajectory, if + # necessary + missing = (charges is None, positions is None) + if any(missing): + if group is None: + emsg = ( + "Either a group of atoms or atom positions and " + "charges must be provided." + ) + raise ValueError(emsg) + + # Check whether the groups have equal numbers of atoms + if grouping is None: + same = True + else: + groups = getattr(group, grouping) + + # Calculate and return the centers of mass for different + # groups here if unwrapping and the mass and position arrays + # are not needed + if ( + not ( + same := all( + g.atoms.n_atoms == groups[0].atoms.n_atoms for g in groups + ) + ) + and images is None + and not raw + ): + return np.array([g.atoms.center_of_mass() for g in groups]) + + # Get and unwrap particle positions, if necessary + if missing[1]: + positions = group.positions + if images is not None: + positions += images * dimensions[:3] + + # Get particle charges and ensure correct dimensionality, if + # necessary + if same: + if missing[0]: + charges = group.charges + if grouping is not None or n_groups is not None: + shape = (n_groups or getattr(group, f"n_{grouping}"), -1, 3) + charges = charges.reshape(shape[:-1]) + positions = positions.reshape(shape) + else: + if missing[0]: + charges = [g.atoms.charges for g in groups] + if missing[1]: + positions = [positions[g.atoms.ix] for g in groups] + else: + + # Try to convert arrays to NumPy arrays if they are not already + # to take advantage of vectorized operations later + try: + positions = np.asarray(positions) + charges = np.asarray(charges) + if images is not None: + positions += images + except ValueError: + pass + if type(charges) is not type(positions): + emsg = ( + "The shapes of the arrays containing the charges " + "and positions are incompatible." + ) + raise ValueError(emsg) + if images is not None and type(images) is not type(positions): + emsg = ( + "The shapes of the arrays containing the positions " + "and image flags are incompatible." + ) + raise ValueError(emsg) + + # Calculate the dipole for the specified grouping + if isinstance(positions, np.ndarray): + + # Reshape the mass and position arrays based on the specified + # number of groups + if n_groups is not None: + charges = charges.reshape((n_groups, -1)) + positions = positions.reshape((n_groups, -1, 3)) + + # dipole = sum(q_i * r_i) and should be a vector of shape (3,) + dipole = np.einsum("ij,ijk->ik", charges, positions) + + else: + if images is not None: + for j, (p, i) in enumerate(zip(positions, images)): + positions[j] = p + i * dimensions + dipole = np.array([np.sum(c*p, axis=0) for c, p in zip(charges, positions)]) + + # Return raw charges and positions, if desired + if raw: + return dipole, charges, positions + + return dipole + +def molecular_quadrupole(group: mda.AtomGroup = None, + grouping: str = None, + *, + charges: Union[np.ndarray[float], list[np.ndarray[float]]] = None, + positions: Union[np.ndarray[float], list[np.ndarray[float]]] = None, + images: Union[np.ndarray[int], list[np.ndarray[int]]] = None, + dimensions: np.ndarray[float] = None, + n_groups: int = None, + raw: bool = False, + ) -> Union[ + np.ndarray[float], tuple[np.ndarray[float], np.ndarray[float], np.ndarray[float]] + ]: + # Check whether grouping is valid + if grouping not in {None, "residues", "segments"}: + emsg = ( + f"Invalid grouping: '{grouping}'. Valid options are " + "None, 'residues', and 'segments'." + ) + raise ValueError(emsg) + + # Get system dimensions if image flags are provided + if images is not None: + if dimensions is None: + try: + dimensions = group.dimensions[:3] + except (NameError, TypeError): + emsg = ( + "Image flags were provided, but no system " + "dimensions were provided or found in the " + "trajectory." + ) + raise ValueError(emsg) + else: + dimensions = np.asarray(dimensions) + + # Get particle charges and positions from the trajectory, if + # necessary + missing = (charges is None, positions is None) + if any(missing): + if group is None: + emsg = ( + "Either a group of atoms or atom positions and " + "charges must be provided." + ) + raise ValueError(emsg) + + # Check whether the groups have equal numbers of atoms + if grouping is None: + same = True + else: + groups = getattr(group, grouping) + + # Calculate and return the centers of mass for different + # groups here if unwrapping and the mass and position arrays + # are not needed + if ( + not ( + same := all( + g.atoms.n_atoms == groups[0].atoms.n_atoms for g in groups + ) + ) + and images is None + and not raw + ): + return np.array([g.atoms.center_of_mass() for g in groups]) + + # Get and unwrap particle positions, if necessary + if missing[1]: + positions = group.positions + if images is not None: + positions += images * dimensions[:3] + + # Get particle charges and ensure correct dimensionality, if + # necessary + if same: + if missing[0]: + charges = group.charges + if grouping is not None or n_groups is not None: + shape = (n_groups or getattr(group, f"n_{grouping}"), -1, 3) + charges = charges.reshape(shape[:-1]) + positions = positions.reshape(shape) + else: + if missing[0]: + charges = [g.atoms.charges for g in groups] + if missing[1]: + positions = [positions[g.atoms.ix] for g in groups] + else: + + # Try to convert arrays to NumPy arrays if they are not already + # to take advantage of vectorized operations later + try: + positions = np.asarray(positions) + charges = np.asarray(charges) + if images is not None: + positions += images + except ValueError: + pass + if type(charges) is not type(positions): + emsg = ( + "The shapes of the arrays containing the charges " + "and positions are incompatible." + ) + raise ValueError(emsg) + if images is not None and type(images) is not type(positions): + emsg = ( + "The shapes of the arrays containing the positions " + "and image flags are incompatible." + ) + raise ValueError(emsg) + + # Calculate the quadrupole for the specified grouping + if isinstance(positions, np.ndarray): + + # Reshape the mass and position arrays based on the specified + # number of groups + if n_groups is not None: + charges = charges.reshape((n_groups, -1)) + positions = positions.reshape((n_groups, -1, 3)) + + + # quadrupole = sum_l q_l(3*r_il*r_jl - |r_l|^2*delta_ij) + quadrupole = np.zeros((positions.shape[0], 3, 3), dtype=float) + rl = np.sum(np.linalg.norm(positions, axis=2)**2 * charges, axis=1) + for i in range(3): + quadrupole[:, i, i] = 3 * np.sum( + charges * positions[:, :, i] * positions[:, :, i], axis=1 + ) - rl + for j in range(i+1, 3): + quadrupole[:, i, j] = np.sum( + charges * (3 * positions[:, :, i] * positions[:, :, j]), axis=1 + ) + quadrupole[:, j, i] = quadrupole[:, i, j] + + else: + if images is not None: + for j, (p, i) in enumerate(zip(positions, images)): + positions[j] = p + i * dimensions + quadrupole = np.zeros((len(positions), 3, 3), dtype=float) + rl = np.sum(np.linalg.norm(positions, axis=2)**2 * charges, axis=1) + for i in range(3): + quadrupole[:, i, i] = np.sum( + charges * positions[:, :, i] * positions[:, :, i], axis=1 + ) - rl + for j in range(i+1, 3): + quadrupole[:, i, j] = np.sum( + charges * (3 * positions[:, :, i] * positions[:, :, j]), axis=1 + ) + quadrupole[:, j, i] = quadrupole[:, i, j] + + # Return raw charges and positions, if desired + if raw: + return quadrupole, charges, positions + + return quadrupole + +class PotentialProfile(DynamicAnalysisBase): + """ + Serial and parallel implementations to calculate the number and + charge density profiles :math:`\\rho_i(z)` and :math:`\\rho_q(z)` of + a constant-volume system along the specified axes. + + The microscopic number density profile of species :math:`i` is + calculated by binning particle positions along an axis :math:`z` + using + + .. math:: + + \\rho_i(z)=\\frac{V}{N_\\mathrm{bins}}\\left\\langle + \\sum_\\alpha\\delta(z-z_\\alpha)\\right\\rangle + + where :math:`V` is the system volume and :math:`N_\\mathrm{bins}` is + the number of bins. The angular brackets denote an ensemble average. + + If the species carry charges, the charge density profile can be + obtained using + + .. math:: + + \\rho_q(z)=e\\sum_i z_i\\rho_i(z) + + where :math:`z_i` is the charge number of species :math:`i` and + :math:`e` is the elementary charge. + + With the charge density profile, the surface charge density is + given by + + .. math:: + + \\sigma_q=\\frac{1}{L}\\left( + \\varepsilon_0\\varepsilon_\\mathrm{r}\\Delta\\Psi + -\\int_0^L z\\rho_q(z)\\,\\mathrm{d}z\\right) + + where :math:`\\varepsilon_0` and :math:`\\varepsilon_\\mathrm{r}` + are the vacuum and relative permittivities, respectively, + :math:`\\Delta\\Psi` is the potential difference, and :math:`L` is + the system dimension, and the potential profile can be computed by + numerically solving Poisson's equation for electrostatics: + + .. math:: + + \\varepsilon_0\\varepsilon_\\mathrm{r}\\nabla^2\\Psi(z)=-\\rho_q(z) + + Parameters + ---------- + groups : `MDAnalysis.AtomGroup` or array-like + Groups of atoms for which density profiles are calculated. + + groupings : `str` or array-like, default: :code:`"atoms"` + Determines whether the centers of mass are used in lieu of + individual atom positions. If `groupings` is a `str`, the same + value is used for all `groups`. + + .. note:: + + If the desired grouping is not :code:`"atoms"`, + + * the trajectory file should have segments (or chains) + containing residues (or molecules) and residues containing + atoms, and + + * residues and segments should be locally unwrapped at the + simulation box edges, if not already, using + :class:`MDAnalysis.transformations.wrap.unwrap`, + :meth:`MDAnalysis.core.groups.AtomGroup.unwrap`, or + :func:`MDAnalysis.lib.mdamath.make_whole`. + + .. container:: + + **Valid values**: + + * :code:`"atoms"`: Atom positions (generally or for + coarse-grained simulations). + * :code:`"residues"`: Residues' centers of mass (for + atomistic simulations). + * :code:`"segments"`: Segments' centers of mass (for + atomistic polymer simulations). + + axes : `int`, `str`, or array-like, default: :code:`"xyz"` + Axes along which to compute the density profiles. + + .. container:: + + **Examples**: + + * :code:`2` for the :math:`z`-direction. + * :code:`"xy"` for the :math:`x`- and :math:`y`-directions. + * :code:`(0, 1)` for the :math:`x`- and :math:`y`-directions. + + n_bins : `int` or array-like + Number of histogram bins :math:`N_\\mathrm{bins}` for each axis + in `axes`. If an `int` is provided, the same value is used for + all axes. + + charges : array-like, `openmm.unit.Quantity`, or `pint.Quantity`, \ + keyword-only, optional + Charges :math:`q_i` for the entities in the + :math:`N_\\mathrm{groups}` atom groups in `groups`. If not + provided, they will be retrieved from the main + :class:`MDAnalysis.core.universe.Universe` object only if it + contains charge information. + + .. note:: + + Depending on the grouping for a specific atom group, all + entities (atoms, residues, or segments) must carry the same + charge. Otherwise, the charge density contribution for that + atom group would not make sense. If this condition does not + hold, change how the atoms are grouped in the atom groups so + that all entities share the same charge. + + **Shape**: :math:`(N_\\mathrm{groups},)`. + + **Reference unit**: :math:`\\mathrm{e}`. + + dimensions : array-like, `openmm.unit.Quantity`, or \ + `pint.Quantity`, keyword-only, optional + System dimensions :math:`(L_x,\\,L_y,\\,L_z)`. If the + :class:`MDAnalysis.core.universe.Universe` object that the + atom groups in `groups` belong to does not contain + dimensionality information, provide it here. Affected by + `dim_scales`. + + **Shape**: :math:`(3,)`. + + **Reference unit**: :math:`\\mathrm{Å}`. + + dim_scales : `float` or array-like, keyword-only, optional + Scale factors for the system dimensions. If an `int` is + provided, the same value is used for all axes. + + **Shape**: :math:`(3,)`. + + average : `bool`, keyword-only, default: :code:`True` + Determines whether the density profiles are averaged over the + :math:`N_\\mathrm{frames}` analysis frames. + + recenter : `int`, `list`, `MDAnalysis.AtomGroup`, or `tuple`, \ + keyword-only, optional + Constrains the center of mass of an atom group by adjusting the + particle coordinates every analysis frame. Either specify an + :class:`MDAnalysis.core.groups.AtomGroup`, its index within + `groups`, a list of atom groups or their indices, or a tuple + containing the aforementioned information and the fixed center + of mass coordinates, in that order. To avoid recentering in a + specific dimension, set the coordinate to :code:`numpy.nan`. If + the center of mass is not specified, the center of the + simulation box is used. + + **Shape**: :math:`(3,)` for the fixed center of mass. + + reduced : `bool`, keyword-only, default: :code:`False` + Specifies whether the data is in reduced units. Affects + `results.number_densities`, `results.charge_densities`, etc. + + parallel : `bool`, keyword-only, default: :code:`False` + Determines whether the analysis is performed in parallel. + + .. note:: + + The Joblib threading backend generally provides the best + performance. + + verbose : `bool`, keyword-only, default: :code:`True` + Determines whether detailed progress is shown. + + **kwargs + Additional keyword arguments to pass to + :class:`MDAnalysis.analysis.base.AnalysisBase`. + + Attributes + ---------- + universe : `MDAnalysis.Universe` + :class:`MDAnalysis.core.universe.Universe` object containing all + information describing the simulation system. + + axes : `tuple` + Axes along which the density profiles are calculated. + + results.units : `dict` + Reference units for the results. For example, to get the + reference units for `results.bins`, call + :code:`results.units["bins"]`. + + results.times : `numpy.ndarray` + Times :math:`t`. Only available if :code:`average=False`. + + **Shape**: :math:`(N_\\mathrm{frames},)`. + + **Reference unit**: :math:`\\mathrm{ps}`. + + results.bins : `dict` + Bin centers :math:`z` corresponding to the density profiles in + each dimension. The key is the axis, e.g., + :code:`results.bins["z"]` for the :math:`z`-axis. + + **Shape**: Each array has shape :math:`(N_\\mathrm{bins},)`. + + **Reference unit**: :math:`\\mathrm{Å}`. + + results.bin_edges : `dict` + Bin edges corresponding to the density profiles in each + dimension. The key is the axis, e.g., + :code:`results.bin_edges["z"]` for the :math:`z`-axis. + + **Shape**: Each array has shape :math:`(N_\\mathrm{bins}+1,)`. + + **Reference unit**: :math:`\\mathrm{Å}`. + + results.number_densities : `dict` + Number density profiles :math:`\\rho(z)`. The key is the axis, + e.g., :code:`results.number_densities["z"]` for the + :math:`z`-axis. + + **Shape**: Each array has shape + :math:`(N_\\mathrm{groups},\\,N_\\mathrm{bins})`. If + :code:`average=False`, an additional second dimension of + length :math:`N_\\mathrm{frames}` is present. + + **Reference unit**: :math:`\\mathrm{Å}^{-3}`. + + results.charge_densities : `dict` + Charge density profiles :math:`\\rho_q(z)`. Only available if + charge information was found or provided. The key is the axis, + e.g., :code:`results.charge_densities["z"]` for the + :math:`z`-axis. + + **Shape**: Each array has shape :math:`(N_\\mathrm{bins},)`. If + :code:`average=False`, an additional first dimension of length + :math:`N_\\mathrm{frames}` is present. + + **Reference unit**: :math:`\\mathrm{e/Å}^{-3}`. + + results.surface_charge_densities : `numpy.ndarray` + Surface charge densities :math:`\\sigma_q`. Only available after + running :meth:`calculate_surface_charge_densities`. + + **Shape**: :math:`(N_\\mathrm{axes},)` or + :math:`(N_\\mathrm{axes},\\,N_\\mathrm{frames})`. + + results.potentials : `dict` + Potential profiles :math:`\\Psi(z)`. Only available after + running :meth:`calculate_potential_profiles`. The key is the + axis, e.g., :code:`results.potentials["z"]` for the + :math:`z`-axis. + + **Shape**: Each array has shape :math:`(N_\\mathrm{bins},)`. If + :code:`average=False`, an additional second dimension of + length :math:`N_\\mathrm{frames}` is present. + + **Reference unit**: :math:`\\mathrm{V}`. + + Example + -------- + First, this analysis class must be imported: + + >>> from mdcraft.analysis.profile import PotentialProfile + + Then, after loading a simulation trajectory: + + >>> universe = mda.Universe("simulation.nc", "topology.cif") + + We must then select the atom-groups to be analyzed: + + >>> ag1 = universe.select_atoms("resname CAT") + >>> ag2 = universe.select_atoms("resname ANI") + + The `PotentialProfile` class can be instantiated with the selected atom-groups and the axes to calculate the density profiles along (charges can also be provided to obtain the charge density profiles): + + >>> prof = PotentialProfile([ag1, ag2], groupings="residues", axes="x", charges=[1, -1]) + >>> prof.run() + + The results can be obtained under the `results` attribute: + + >>> prof.results.number_densities["x"] + + To calculate the surface charge densities, the `calculate_surface_charge_densities` method can be used: + + >>> prof.calculate_surface_charge_densities(dV=1, dielectric=80) + + Further, to calculate the potential profiles, the `calculate_potential_profiles` method can be used: + + >>> prof.calculate_potential_profiles(dielectric=80) + + These results can be saved to a file using the `save` method: + + >>> prof.save("density_profiles") + """ + + def __init__( + self, + groups: Union[mda.AtomGroup, tuple[mda.AtomGroup]], + groupings: Union[str, tuple[str]] = "atoms", + axes: Union[int, str, tuple[Union[int, str]]] = "xyz", + dipole_axes: Union[int, str, tuple[Union[int, str]]] = "xyz", + dipole_center: str = "center_of_mass", + n_bins: Union[int, tuple[int]] = 201, + *, + charges: Union[np.ndarray[float], "unit.Quantity", Q_] = None, + dimensions: Union[np.ndarray[float], "unit.Quantity", Q_] = None, + recenter: Union[ + mda.AtomGroup, + int, + list[mda.AtomGroup, int], + tuple[ + Union[mda.AtomGroup, int, list[mda.AtomGroup, int]], np.ndarray[float] + ], + ] = None, + dim_scales: Union[float, tuple[float]] = 1, + average: bool = True, + reduced: bool = False, + parallel: bool = False, + verbose: bool = True, + **kwargs, + ) -> None: + + self._groups = [groups] if isinstance(groups, mda.AtomGroup) else groups + self.universe = self._groups[0].universe + super().__init__(self.universe.trajectory, parallel, verbose, **kwargs) + + GROUPINGS = {"atoms", "residues", "segments"} + self._n_groups = len(self._groups) + if isinstance(groupings, str): + if groupings not in GROUPINGS: + emsg = ( + f"Invalid grouping '{groupings}'. Valid values: " + "'" + "', '".join(GROUPINGS) + "'." + ) + raise ValueError(emsg) + self._groupings = self._n_groups * [groupings] + else: + if self._n_groups != len(groupings): + emsg = ( + "The shape of 'groupings' is incompatible with " "that of 'groups'." + ) + raise ValueError(emsg) + for gr in groupings: + if gr not in GROUPINGS: + emsg = ( + f"Invalid grouping '{gr}'. Valid values: " + "'" + "', '".join(GROUPINGS) + "'." + ) + raise ValueError(emsg) + self._groupings = groupings + + if isinstance(axes, int): + self._axis_indices = np.array((axes,), dtype=int) + else: + self._axis_indices = np.fromiter( + (ord(ax.lower()) - 120 if isinstance(ax, str) else ax for ax in axes), + count=len(axes), + dtype=int, + ) + + self.axes = tuple(chr(120 + i) for i in self._axis_indices) + if not all(ax in "xyz" for ax in self.axes): + raise ValueError("Invalid axis passed in 'axes'.") + self._n_axes = len(self.axes) + + if isinstance(dipole_axes, int): + self._dipole_axis_indices = np.array((dipole_axes,), dtype=int) + else: + self._dipole_axis_indices = np.fromiter( + (ord(ax.lower()) - 120 if isinstance(ax, str) else ax for ax in dipole_axes), + count=len(dipole_axes), + dtype=int, + ) + self.dipole_axes = tuple(chr(120 + i) for i in self._dipole_axis_indices) + if not all(ax in "xyz" for ax in self.dipole_axes): + raise ValueError("Invalid axis passed in 'dipole_axes'.") + + if dipole_center not in {"center_of_mass", "center_of_charge"}: + emsg = ( + f"Invalid dipole center '{dipole_center}'. " + "Valid values: 'center_of_mass', 'center_of_charge'." + ) + raise ValueError(emsg) + self._dipole_center = dipole_center + + if isinstance(n_bins, int): + self._n_bins = n_bins * np.ones(len(self.axes), dtype=int) + elif isinstance(n_bins, str): + emsg = "'n_bins' must be an integer or an iterable object." + raise ValueError(emsg) + else: + if len(n_bins) != len(self.axes): + emsg = ( + "The shape of 'n_bins' is incompatible with " + "the number of axes to calculate density " + "profiles along." + ) + raise ValueError(emsg) + if not all(isinstance(n, int) for n in n_bins): + emsg = "All bin counts in 'n_bins' must be integers." + raise ValueError(emsg) + self._n_bins = n_bins + + if charges is not None: + if len(charges) != self._n_groups: + emsg = ( + "The shape of 'charges' is incompatible with " "that of 'groups'." + ) + raise ValueError(emsg) + if reduced and not is_unitless(charges): + emsg = "'charges' cannot have units when 'reduced=True'." + raise TypeError(emsg) + self._charges = np.asarray(strip_unit(charges, "e")[0]) + elif hasattr(self.universe.atoms, "charges"): + self._charges = np.empty(self._n_groups) + for i, (ag, gr) in enumerate(zip(self._groups, self._groupings)): + qs = getattr(ag, gr).charges + if not np.allclose(qs, q := qs[0]): + self._charges = None + wmsg = ( + f"Not all {gr} in `groups[{i}]` share the " + "same charge. The charge density profile will " + "not be calculated." + ) + warnings.warn(wmsg) + break + self._charges[i] = q + else: + self._charges = None + + if dimensions is not None: + if len(dimensions) != 3: + raise ValueError("'dimensions' must have length 3.") + if reduced and not is_unitless(dimensions): + emsg = "'dimensions' cannot have units when 'reduced=True'." + raise TypeError(emsg) + self._dimensions = np.asarray(strip_unit(dimensions, "Å")[0]) + elif self.universe.dimensions is not None: + self._dimensions = self.universe.dimensions[:3].copy() + else: + raise ValueError("No system dimensions found or provided.") + + if isinstance(dim_scales, Real) or ( + len(dim_scales) == 3 and all(isinstance(f, Real) for f in dim_scales) + ): + self._dimensions *= dim_scales + else: + emsg = ( + "'dim_scales' must be a floating-point number " + "or an array with shape (3,)." + ) + raise ValueError(emsg) + + self._Ns = np.fromiter( + (getattr(ag, f"n_{gr}") for (ag, gr) in zip(self._groups, self._groupings)), + dtype=int, + count=self._n_groups, + ) + self._N = self._Ns.sum() + self._slices = [] + _ = 0 + for N in self._Ns: + self._slices.append(slice(_, _ + N)) + _ += N + + if recenter is None: + self._recenter = recenter + else: + if isinstance(recenter, (int, mda.AtomGroup)): + grp = recenter + com = self._dimensions / 2 + elif isinstance(recenter, tuple) and len(recenter) == 2: + grp, com = recenter + if len(com) != 3: + emsg = ( + "The target center of mass in 'recenter' " "must have length 3." + ) + raise ValueError(emsg) + com = np.asarray(com) + else: + emsg = ( + "Invalid value passed to 'recenter'. The " + "argument must either be an atom group, its " + "index in 'groups', multiple groups/indices, or " + "a tuple containing the aforementioned " + "information and a target center of mass, in " + "that order." + ) + raise ValueError(emsg) + + if isinstance(grp, int): + if not 0 <= grp < self._n_groups: + emsg = "Invalid group index passed to 'recenter'." + raise ValueError(emsg) + grp = self._slices[grp] + elif isinstance(grp, mda.AtomGroup): + try: + grp = self._slices[self._groups.index(grp)] + except ValueError: + emsg = "The atom group passed to 'recenter' is not " "in 'groups'." + raise ValueError(emsg) + else: + try: + grp = np.r_[ + tuple( + self._slices[ + g if isinstance(g, int) else self._groups.index(g) + ] + for g in grp + ) + ] + except ValueError: + emsg = "Invalid atom group or index passed to 'recenter'." + raise ValueError(emsg) + self._recenter = (grp, com) + + self._dielectrics = np.empty(self._n_axes) + self._dVs = np.empty(self._n_axes) + self._dielectrics[:] = self._dVs[:] = np.nan + + self._average = average + self._reduced = reduced + self._verbose = verbose + + def _prepare(self) -> None: + + # Specify bin centers and edges for each axis + self.results.bins = Hash() + self.results.bin_edges = Hash() + for ia, ax, n in zip(self._axis_indices, self.axes, self._n_bins): + self.results.bin_edges[ax] = numba_histogram_bin_edges( + np.asarray((0.0, self._dimensions[ia])), n + ) + self.results.bins[ax] = ( + self.results.bin_edges[ax][:-1] + self.results.bin_edges[ax][1:] + ) / 2 + + # Store entity masses for center of mass calculations + self._masses = np.empty(self._N) + for ag, gr, s in zip(self._groups, self._groupings, self._slices): + self._masses[s] = getattr(ag, gr).masses + + if self._recenter is not None: + + # Navigate to first frame in analysis + self._sliced_trajectory[0] + + # Preallocate arrays to determine the number of periodic + # boundary crossings for each entity + self._positions_old = np.empty((self._N, 3)) + for ag, gr, s in zip(self._groups, self._groupings, self._slices): + self._positions_old[s] = ( + ag.positions if gr == "atoms" else center_of_mass(ag, gr) + ) + self._images = np.zeros((self._N, 3), dtype=int) + self._thresholds = self._dimensions / 2 + + if self._parallel: + + # Preallocate array to hold processed entity positions + # for all frames to be analyzed in parallel + self._positions = np.empty((self.n_frames, self._N, 3)) + for i, _ in enumerate(self._sliced_trajectory): + + # Get raw entity positions for current frame + for ag, gr, s in zip(self._groups, self._groupings, self._slices): + self._positions[i, s] = ( + ag.positions if gr == "atoms" else center_of_mass(ag, gr) + ) + + # Globally unwrap entity positions for correct + # center of mass + unwrap( + self._positions[i], + self._positions_old, + self._dimensions, + thresholds=self._thresholds, + images=self._images, + ) + + # Shift entity positions by the difference between + # the group and target centers of mass + self._positions[i] -= np.fromiter( + ( + 0 if np.isnan(tx) else gx - tx + for gx, tx in zip( + center_of_mass( + positions=self._positions[i, self._recenter[0]], + masses=self._masses[self._recenter[0]], + ), + self._recenter[1], + ) + ), + dtype=float, + count=3, + ) + + # Wrap entity positions back into the simulation box so that + # they belong to a histogram bin + wrap(self._positions, self._dimensions) + + # Preallocate arrays to hold entity positions for a given frame + # (so that one doesn't have to be recreated each frame) and + # number density profiles for serial analysis + if not self._parallel: + self._positions = np.empty((self._N, 3)) + self._dipoles = np.empty((self._N, 3), dtype=float) + self._quadrupoles = np.empty((self._N, 3, 3), dtype=float) + shape = [self._n_groups] + if not self._average: + shape.append(self.n_frames) + self.results.number_densities = Hash( + {ax: np.zeros((*shape, n)) for ax, n in zip(self.axes, self._n_bins)} + ) + self.results.dipoles = Hash( + {(ax, ax_d): np.zeros((*shape, n)) for ax, n in zip(self.axes, self._n_bins) for ax_d in self.dipole_axes} + ) + self.results.quadrupoles = Hash( + {(ax, ax_d1, ax_d2): np.zeros((*shape, n)) for ax, n in zip(self.axes, self._n_bins) for ax_d1 in self.dipole_axes for ax_d2 in self.dipole_axes} + ) + + # Store reference units + self.results.units = Hash( + {"bins": ureg.angstrom, "number_densities": ureg.angstrom**-3, + "dipoles": ureg.debye} + ) + + # Preallocate dictionary to hold charge density profiles, if + # necessary + if self._charges is not None: + self.results.charge_densities = Hash() + self.results.units["charge_densities"] = ( + ureg.elementary_charge / ureg.angstrom**3 + ) + + # Store time information, if necessary + if not self._average: + self.results.times = np.fromiter( + (ts.time for ts in self._sliced_trajectory), + dtype=float, + count=self.n_frames, + ) + self.results.units["times"] = ureg.picosecond + + def _single_frame(self) -> None: + + # Store atom or center-of-mass positions in the current frame + for ag, gr, s in zip(self._groups, self._groupings, self._slices): + if self._dipole_center == "center_of_mass": + self._positions[s] = ( + ag.positions if gr == "atoms" else center_of_mass(ag, gr) + ) + # elif self._dipole_center == "center_of_geometry": + # self._positions[s] = ( + # ag.positions if gr == "atoms" else center_of_geometry(ag, gr) + # ) + elif self._dipole_center == "center_of_charge": + self._positions[s] = ( + ag.positions if gr == "atoms" else center_of_charge(ag, gr) + ) + self._dipoles[s] = molecular_dipole(ag, gr) + self._quadrupoles[s] = molecular_quadrupole(ag, gr) + + if self._recenter is not None: + + # Globally unwrap entity positions for correct centers of + # mass + unwrap( + self._positions, + self._positions_old, + self._dimensions, + thresholds=self._thresholds, + images=self._images, + ) + + # Shift entity positions by the difference between the group + # and target centers of mass + self._positions -= np.fromiter( + ( + 0 if np.isnan(tx) else gx - tx + for gx, tx in zip( + center_of_mass( + positions=self._positions[self._recenter[0]], + masses=self._masses[self._recenter[0]], + ), + self._recenter[1], + ) + ), + dtype=float, + count=3, + ) + + # Wrap entity positions back into the simulation box so that + # they belong to a histogram bin + wrap(self._positions, self._dimensions) + + # Compute and tally the bin counts for the entity positions + for ax, ia, n in zip(self.axes, self._axis_indices, self._n_bins): + for ig, (gr, s) in enumerate(zip(self._groupings, self._slices)): + if self._average: + self.results.number_densities[ax][ig] += numba_histogram( + self._positions[s, ia], n, self.results.bin_edges[ax] + ) + for ax_d, id in zip(self.dipole_axes, self._dipole_axis_indices): + + self.results.dipoles[ax, ax_d][ig] += np.histogram( + self._positions[s, ia], bins=n, range=self.results.bin_edges[ax][[0,-1]], weights=self._dipoles[s, id] + )[0] + for ax_d1, id1 in zip(self.dipole_axes, self._dipole_axis_indices): + for ax_d2, id2 in zip(self.dipole_axes, self._dipole_axis_indices): + self.results.quadrupoles[ax, ax_d1, ax_d2][ig] += np.histogram( + self._positions[s, ia], bins=n, range=self.results.bin_edges[ax][[0,-1]], weights=self._quadrupoles[s, id1, id2] + )[0] + else: + self.results.number_densities[ax][ig, self._frame_index] = ( + numba_histogram( + self._positions[s, ia], n, self.results.bin_edges[ax] + ) + ) + for ax_d in self.dipole_axes: + self.results.dipoles[ax, ax_d][ig, self._frame_index] = np.histogram( + self._positions[s, ia], bins=n, range=self.results.bin_edges[ax][[0,-1]], weights=self._dipoles[s, ax_d] + )[0] + for ax_d1, id1 in zip(self.dipole_axes, self._dipole_axis_indices): + for ax_d2, id2 in zip(self.dipole_axes, self._dipole_axis_indices): + self.results.quadrupoles[ax, ax_d1, ax_d2][ig, self._frame_index] = np.histogram( + self._positions[s, ia], bins=n, range=self.results.bin_edges[ax][[0,-1]], weights=self._quadrupoles[s, id1, id2] + )[0] + + def _conclude(self) -> None: + + # Consolidate parallel results and clean up memory by deleting + # arrays that will not be reused + if self._parallel: + self._results = sorted(self._results) + self.results.number_densities = Hash() + for ax in self.axes: + self.results.number_densities[ax] = np.stack( + [r[1][ax] for r in self._results], axis=1 + ) + + if self._average: + self.results.number_densities[ax] = self.results.number_densities[ + ax + ].sum(axis=1) + for ax_d in self.dipole_axes: + self.results.dipoles[ax, ax_d] = np.stack( + [r[1][ax, ax_d] for r in self._results], axis=1 + ) + if self._average: + self.results.dipoles[ax, ax_d] = self.results.dipoles[ + ax, ax_d + ].sum(axis=1) + + del self._results + if self._recenter is not None: + del self._positions + else: + del self._positions + if self._recenter is not None: + del self._positions_old, self._images, self._thresholds + + # Normalize histograms by bin volume + volume = np.prod(self._dimensions) + for ax, n in zip(self.axes, self._n_bins): + denom = volume / n + if self._average: + denom *= self.n_frames + self.results.number_densities[ax] /= denom + for ax_d in self.dipole_axes: + self.results.dipoles[ax, ax_d] /= denom + + if self._charges is not None: + self.results.charge_densities[ax] = np.einsum( + "g,g...b->...b", self._charges, self.results.number_densities[ax] + ) + + def _validate_input( + self, value: Any, unit_: str, name: str, n_axes: int + ) -> Union[Real, np.ndarray[Real]]: + """ + Validates and processes input values to the calculation methods. + + Parameters + ---------- + value : `Any` + Value to validate and process. + + unit_ : `str` + Reference unit for the value. + + name : `str` + Name of the value being validated. + + n_axes : `int` + Number of axes along which to calculate the value. + + Returns + ------- + `Real` or `numpy.ndarray` + Processed value. + """ + + if value is None: + value = n_axes * [value] + else: + if self._reduced and not is_unitless(value): + emsg = f"'{name}' cannot have units when 'reduced=True'." + raise TypeError(emsg) + value = strip_unit(value, unit_)[0] + if isinstance(value, Real): + value = n_axes * [value] + elif len(value) != n_axes: + emsg = f"The length of '{name}' must match the " "number of axes." + raise ValueError(emsg) + return value + + def calculate_surface_charge_densities( + self, + axes: Union[str, tuple[str]] = None, + dielectrics: Union[float, tuple[float]] = None, + *, + dVs: Union[float, np.ndarray[float], "unit.Quantity", Q_] = None, + ) -> None: + """ + Calculates the surface charge densities :math:`\\sigma_q` for + the specified system dimensions using the charge density + profiles :math:`\\rho_q(z)`. + + Parameters + ---------- + axes : `str` or array-like, optional + Axes along which to compute the potential profiles. If not + specified, all axes for which charge density profiles were + calculated will be used. + + **Examples**::code:`"xy"` or :code:`("x", "y")` + for the :math:`x`- and :math:`y`-directions. + + dielectrics : `float`, optional + Relative permittivities or dielectric constants + :math:`\\varepsilon_\\mathrm{r}`. Only optional if + previously provided to another calculation method in this + class. + + dVs : `float`, array-like, `openmm.unit.Quantity`, or \ + `pint.Quantity`, keyword-only, optional + Potential differences :math:`\\Delta\\Psi` across the system + dimensions specified in `axes`. Can be retrieved if + previously provided to another calculation method in this + class. + + **Shapes**: :math:`(N_\\mathrm{axes},)` or + :math:`(N_\\mathrm{axes},\\,N_\\mathrm{frames})`. + + **Reference unit**: :math:`\\mathrm{V}`. + """ + + # Ensure charge density profiles have already been calculated + if not hasattr(self.results, "charge_densities"): + emsg = ( + "Either call PotentialProfile.run() before " + "PotentialProfile.calculate_potential_profiles() or " + "provide charge information when initializing the " + "PotentialProfile object." + ) + raise RuntimeError(emsg) + + # Validate inputs + if axes is None: + axes = self.axes + axis_indices = self._axis_indices + else: + if isinstance(axes, Real) or any(not isinstance(ax, str) for ax in axes): + raise ValueError("'axes' must only contain strings.") + axes = tuple(axes) + axis_indices = [ord(ax.lower()) - 120 for ax in axes] + try: + relative_axes = [self.axes.index(ax) for ax in axes] + except ValueError: + raise ValueError("Invalid axis passed in 'axes'.") + n_axes = len(axes) + dimensions = self._dimensions[axis_indices] + + if dielectrics is not None: + self._dielectrics[relative_axes] = self._validate_input( + dielectrics, "", "dielectrics", n_axes + ) + elif np.any(np.isnan(self._dielectrics)): + raise ValueError("No dielectric constants found or provided.") + + if dVs is not None: + self._dVs[relative_axes] = self._validate_input(dVs, "V", "dVs", n_axes) + + # Preallocate dictionary to hold potential profiles + if not hasattr(self.results, "surface_charge_densities"): + shape = [self._n_axes] + if not self._average: + shape.append(self.n_frames) + self.results.surface_charge_densities = np.empty(shape) + self.results.surface_charge_densities[:] = np.nan + self.results.units["surface_charge_densities"] = ( + ureg.elementary_charge / ureg.angstrom**2 + ) + + # Calculate surface charge densities + for i, (ax, dielectric, L, dV) in enumerate( + zip(axes, self._dielectrics, dimensions, self._dVs) + ): + self.results.surface_charge_densities[i] = calculate_surface_charge_density( + self.results.bins[ax], + self.results.charge_densities[ax], + dielectric, + L=L, + dV=None if np.isnan(dV) else dV, + reduced=self._reduced, + ) + + def calculate_potential_profiles( + self, + axes: Union[str, tuple[str]] = None, + dielectrics: Union[float, tuple[float]] = None, + *, + sigmas_q: Union[float, np.ndarray[float], "unit.Quantity", Q_] = None, + dVs: Union[float, np.ndarray[float], "unit.Quantity", Q_] = None, + thresholds: Union[float, np.ndarray[float]] = 1e-5, + V0s: Union[float, np.ndarray[float], "unit.Quantity", Q_] = 0, + methods: Union[str, tuple[str]] = "integral", + pbcs: Union[bool, tuple[bool]] = False, + ) -> None: + """ + Calculates the potential profiles in the specified dimensions + using the charge density profiles by numerically solving + Poisson's equation for electrostatics. + + Parameters + ---------- + axes : `str` or array-like, optional + Axes along which to compute the potential profiles. If not + specified, all axes for which charge density profiles were + calculated will be used. + + **Examples**::code:`"xy"` or :code:`("x", "y")` + for the :math:`x`- and :math:`y`-directions. + + dielectrics : `float`, optional + Relative permittivities or dielectric constants + :math:`\\varepsilon_\\mathrm{r}`. Only optional if + previously provided to another calculation method in this + class. + + sigmas_q : `float`, array-like, `openmm.unit.Quantity`, or \ + `pint.Quantity`, keyword-only, optional + Surface charge densities :math:`\\sigma_q`. Used to ensure + that the electric field in the bulk of the solution is zero. + If not provided, it is determined using `dVs` and the charge + density profiles, or the average values in the centers of + the integrated charge density profiles. + + .. note:: + + :math:`\\sigma_q` and :math:`\\Delta\\Psi` should have + the same sign. + + **Shapes**: :math:`(N_\\mathrm{axes},)` or + :math:`(N_\\mathrm{axes},\\,N_\\mathrm{frames})`. + + **Reference unit**: :math:`\\mathrm{e/Å^2}`. + + dVs : `float`, array-like, `openmm.unit.Quantity`, or \ + `pint.Quantity`, keyword-only, optional + Potential differences :math:`\\Delta\\Psi` across the system + dimensions specified in `axes`. Can be retrieved if + previously provided to another calculation method in this + class. Has no effect if `sigmas_q` is provided since this + value is used solely to calculate `sigmas_q`. + + **Shapes**: :math:`(N_\\mathrm{axes},)` or + :math:`(N_\\mathrm{axes},\\,N_\\mathrm{frames})`. + + **Reference unit**: :math:`\\mathrm{V}`. + + thresholds : `float` or array-like, keyword-only, \ + default: :code:`1e-5` + Thresholds for determining the plateau regions of the + integrals of the charge density profiles to calculate + `sigmas_q`. Has no effect if `sigmas_q` is provided, or if + `sigmas_q` can be calculated using `dVs` and + `charge_density_profiles`. + + V0s : `float`, array-like, `openmm.unit.Quantity`, or \ + `pint.Quantity`, keyword-only, default: :code:`0` + Potentials :math:`\\Psi_0` at the left boundary. + + **Shapes**: :math:`(N_\\mathrm{axes},)` or + :math:`(N_\\mathrm{axes},\\,N_\\mathrm{frames})`. + + **Reference unit**: :math:`\\mathrm{V}`. + + methods : `str` or array-like, keyword-only, \ + default: :code:`"integral"` + Methods to use to calculate the potential profiles. + + **Valid values**: :code:`"integral"`, :code:`"matrix"`. + + pbcs : `bool`, keyword-only, default: :code:`False` + Specifies whether the system has periodic boundary + conditions in each of the axes. Only used when + :code:`method="matrix"`. + """ + + # Ensure charge density profiles have already been calculated + if not hasattr(self.results, "charge_densities"): + emsg = ( + "Either call PotentialProfile.run() before " + "PotentialProfile.calculate_potential_profiles() or " + "provide charge information when initializing the " + "PotentialProfile object." + ) + raise RuntimeError(emsg) + + # Preallocate dictionary to hold potential profiles + if not hasattr(self.results, "potentials"): + self.results.potentials = Hash() + self.results.units["potentials"] = ureg.volt + + # Validate inputs + if axes is None: + axes = self.axes + axis_indices = self._axis_indices + else: + if isinstance(axes, Real) or any(not isinstance(ax, str) for ax in axes): + raise ValueError("'axes' must only contain strings.") + axes = tuple(axes) + axis_indices = [ord(ax.lower()) - 120 for ax in axes] + try: + relative_axes = [self.axes.index(ax) for ax in axes] + except ValueError: + raise ValueError("Invalid axis passed in 'axes'.") + n_axes = len(axes) + dimensions = self._dimensions[axis_indices] + + if dielectrics is not None: + self._dielectrics[relative_axes] = self._validate_input( + dielectrics, "", "dielectrics", n_axes + ) + elif np.any(np.isnan(self._dielectrics)): + raise ValueError("No dielectric constants found or provided.") + + if sigmas_q is None: + if hasattr(self.results, "surface_charge_densities") and np.all( + np.isfinite(self._dVs[relative_axes]) + ): + sigmas_q = -self.results.surface_charge_densities + else: + sigmas_q = self._validate_input(sigmas_q, "e/Å^2", "sigmas_q", n_axes) + if dVs is not None: + self._dVs[relative_axes] = self._validate_input( + dVs, "V", "dVs", n_axes + ) + else: + sigmas_q = self._validate_input(sigmas_q, "e/Å^2", "sigmas_q", n_axes) + + V0s = self._validate_input(V0s, "V", "V0s", n_axes) + thresholds = self._validate_input(thresholds, "", "thresholds", n_axes) + + METHODS = {"integral", "matrix"} + if isinstance(methods, str) and methods in METHODS: + methods = n_axes * [methods] + elif len(methods) != n_axes: + emsg = "The length of 'methods' must match the number of axes." + raise ValueError(emsg) + else: + for method in methods: + if method not in METHODS: + emsg = ( + f"Invalid method '{method}'. Valid values: " + "'" + "', '".join(METHODS) + "'." + ) + raise ValueError(emsg) + + if isinstance(pbcs, bool): + pbcs = n_axes * [pbcs] + elif len(pbcs) != n_axes: + emsg = "The length of 'pbcs' must match the number of axes." + raise ValueError(emsg) + elif not all(isinstance(pbc, bool) for pbc in pbcs): + raise ValueError("All values in 'pbcs' must be booleans.") + + # Calculate potential profiles + for ax, dielectric, L, sigma_q, dV, threshold, V0, method, pbc in zip( + axes, + self._dielectrics, + dimensions, + sigmas_q, + self._dVs, + thresholds, + V0s, + methods, + pbcs, + ): + self.results.potentials[ax] = calculate_potential_profile( + self.results.bins[ax], + self.results.charge_densities[ax], + dielectric, + L=L, + sigma_q=sigma_q, + dV=None if np.isnan(dV) else dV, + threshold=threshold, + V0=V0, + method=method, + pbc=pbc, + reduced=self._reduced, + ) \ No newline at end of file From 9e2ebae77418b618f55bc63fceef70df2762ec19 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Tue, 15 Jul 2025 00:56:31 -0700 Subject: [PATCH 2/4] Add weights option to numba_histogram --- src/mdcraft/algorithm/accelerated.py | 11 ++++++++--- src/mdcraft/analysis/potential.py | 24 ++++++++++++------------ 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/src/mdcraft/algorithm/accelerated.py b/src/mdcraft/algorithm/accelerated.py index a1222a9..6555dc7 100644 --- a/src/mdcraft/algorithm/accelerated.py +++ b/src/mdcraft/algorithm/accelerated.py @@ -45,7 +45,7 @@ def numba_histogram_bin_edges( @numba.njit(fastmath=True) def numba_histogram( - array: np.ndarray[float], n_bins: int, bin_edges: np.ndarray[float] + array: np.ndarray[float], n_bins: int, bin_edges: np.ndarray[float], weights: np.ndarray[float] | None = None ) -> np.ndarray[int]: r""" Serial Numba-accelerated function to compute the histogram of a @@ -63,6 +63,11 @@ def numba_histogram( bin_edges : `np.ndarray` Bin edges. + weights : `np.ndarray`, optional + Weights for each element in the array. If provided, the histogram + will be weighted by these values. If `None`, each element is counted + equally. + Returns ------- histogram_ : `np.ndarray` @@ -71,13 +76,13 @@ def numba_histogram( min_, max_ = bin_edges[0], bin_edges[-1] histogram_ = np.zeros(n_bins, dtype=np.intp) - for x in array: + for i, x in enumerate(array): if x == max_: bin_ = n_bins - 1 else: bin_ = int(n_bins * (x - min_) / (max_ - min_)) if 0 <= bin_ < n_bins: - histogram_[bin_] += 1 + histogram_[bin_] += 1 if weights is None else weights[i] return histogram_ diff --git a/src/mdcraft/analysis/potential.py b/src/mdcraft/analysis/potential.py index c0cb6f3..c921221 100644 --- a/src/mdcraft/analysis/potential.py +++ b/src/mdcraft/analysis/potential.py @@ -1644,14 +1644,14 @@ def _single_frame(self) -> None: ) for ax_d, id in zip(self.dipole_axes, self._dipole_axis_indices): - self.results.dipoles[ax, ax_d][ig] += np.histogram( - self._positions[s, ia], bins=n, range=self.results.bin_edges[ax][[0,-1]], weights=self._dipoles[s, id] - )[0] + self.results.dipoles[ax, ax_d][ig] += numba_histogram( + self._positions[s, ia], n, self.results.bin_edges[ax], self._dipoles[s, id] + ) for ax_d1, id1 in zip(self.dipole_axes, self._dipole_axis_indices): for ax_d2, id2 in zip(self.dipole_axes, self._dipole_axis_indices): - self.results.quadrupoles[ax, ax_d1, ax_d2][ig] += np.histogram( - self._positions[s, ia], bins=n, range=self.results.bin_edges[ax][[0,-1]], weights=self._quadrupoles[s, id1, id2] - )[0] + self.results.quadrupoles[ax, ax_d1, ax_d2][ig] += numba_histogram( + self._positions[s, ia], n, self.results.bin_edges[ax], self._quadrupoles[s, id1, id2] + ) else: self.results.number_densities[ax][ig, self._frame_index] = ( numba_histogram( @@ -1659,14 +1659,14 @@ def _single_frame(self) -> None: ) ) for ax_d in self.dipole_axes: - self.results.dipoles[ax, ax_d][ig, self._frame_index] = np.histogram( - self._positions[s, ia], bins=n, range=self.results.bin_edges[ax][[0,-1]], weights=self._dipoles[s, ax_d] - )[0] + self.results.dipoles[ax, ax_d][ig, self._frame_index] = numba_histogram( + self._positions[s, ia], n, self.results.bin_edges[ax], self._dipoles[s, ax_d] + ) for ax_d1, id1 in zip(self.dipole_axes, self._dipole_axis_indices): for ax_d2, id2 in zip(self.dipole_axes, self._dipole_axis_indices): - self.results.quadrupoles[ax, ax_d1, ax_d2][ig, self._frame_index] = np.histogram( - self._positions[s, ia], bins=n, range=self.results.bin_edges[ax][[0,-1]], weights=self._quadrupoles[s, id1, id2] - )[0] + self.results.quadrupoles[ax, ax_d1, ax_d2][ig, self._frame_index] = numba_histogram( + self._positions[s, ia], n, self.results.bin_edges[ax], self._quadrupoles[s, id1, id2] + ) def _conclude(self) -> None: From 85c5175b1005aa59c636f258a675527ff96d9a59 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Wed, 16 Jul 2025 09:52:38 -0700 Subject: [PATCH 3/4] typos --- src/mdcraft/algorithm/accelerated.py | 2 +- src/mdcraft/analysis/potential.py | 12 ++++++------ 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/mdcraft/algorithm/accelerated.py b/src/mdcraft/algorithm/accelerated.py index 6555dc7..d0688c3 100644 --- a/src/mdcraft/algorithm/accelerated.py +++ b/src/mdcraft/algorithm/accelerated.py @@ -75,7 +75,7 @@ def numba_histogram( """ min_, max_ = bin_edges[0], bin_edges[-1] - histogram_ = np.zeros(n_bins, dtype=np.intp) + histogram_ = np.zeros(n_bins, dtype=np.float64) for i, x in enumerate(array): if x == max_: bin_ = n_bins - 1 diff --git a/src/mdcraft/analysis/potential.py b/src/mdcraft/analysis/potential.py index c921221..6e4f032 100644 --- a/src/mdcraft/analysis/potential.py +++ b/src/mdcraft/analysis/potential.py @@ -1658,9 +1658,9 @@ def _single_frame(self) -> None: self._positions[s, ia], n, self.results.bin_edges[ax] ) ) - for ax_d in self.dipole_axes: + for ax_d, id in zip(self.dipole_axes, self._dipole_axis_indices): self.results.dipoles[ax, ax_d][ig, self._frame_index] = numba_histogram( - self._positions[s, ia], n, self.results.bin_edges[ax], self._dipoles[s, ax_d] + self._positions[s, ia], n, self.results.bin_edges[ax], self._dipoles[s, id] ) for ax_d1, id1 in zip(self.dipole_axes, self._dipole_axis_indices): for ax_d2, id2 in zip(self.dipole_axes, self._dipole_axis_indices): @@ -1694,10 +1694,10 @@ def _conclude(self) -> None: ].sum(axis=1) del self._results - if self._recenter is not None: - del self._positions - else: - del self._positions + # if self._recenter is not None: + # del self._positions + # else: + # del self._positions if self._recenter is not None: del self._positions_old, self._images, self._thresholds From 94153c19f27decb396157faf06ef2bc28f376eb1 Mon Sep 17 00:00:00 2001 From: pw0908 Date: Wed, 16 Jul 2025 13:10:05 -0700 Subject: [PATCH 4/4] Add center of geometry --- src/mdcraft/analysis/potential.py | 127 +++++++++++++++++++++++++++++- 1 file changed, 125 insertions(+), 2 deletions(-) diff --git a/src/mdcraft/analysis/potential.py b/src/mdcraft/analysis/potential.py index 6e4f032..a91b009 100644 --- a/src/mdcraft/analysis/potential.py +++ b/src/mdcraft/analysis/potential.py @@ -652,6 +652,125 @@ def center_of_charge(group: mda.AtomGroup = None, return coc +def center_of_geometry(group: mda.AtomGroup = None, + grouping: str = None, + *, + positions: Union[np.ndarray[float], list[np.ndarray[float]]] = None, + images: Union[np.ndarray[int], list[np.ndarray[int]]] = None, + dimensions: np.ndarray[float] = None, + n_groups: int = None, + raw: bool = False, + ) -> Union[ + np.ndarray[float], tuple[np.ndarray[float], np.ndarray[float], np.ndarray[float]] + ]: + # Check whether grouping is valid + if grouping not in {None, "residues", "segments"}: + emsg = ( + f"Invalid grouping: '{grouping}'. Valid options are " + "None, 'residues', and 'segments'." + ) + raise ValueError(emsg) + + # Get system dimensions if image flags are provided + if images is not None: + if dimensions is None: + try: + dimensions = group.dimensions[:3] + except (NameError, TypeError): + emsg = ( + "Image flags were provided, but no system " + "dimensions were provided or found in the " + "trajectory." + ) + raise ValueError(emsg) + else: + dimensions = np.asarray(dimensions) + + # Get particle masses and positions from the trajectory, if + # necessary + missing = (positions is None) + if missing: + if group is None: + emsg = ( + "Either a group of atoms or atom positions" + "must be provided." + ) + raise ValueError(emsg) + + # Check whether the groups have equal numbers of atoms + if grouping is None: + same = True + else: + groups = getattr(group, grouping) + + # Calculate and return the centers of mass for different + # groups here if unwrapping and the mass and position arrays + # are not needed + if ( + not ( + same := all( + g.atoms.n_atoms == groups[0].atoms.n_atoms for g in groups + ) + ) + and images is None + and not raw + ): + return np.array([g.atoms.center_of_mass() for g in groups]) + + # Get and unwrap particle positions, if necessary + if missing: + positions = group.positions + if images is not None: + positions += images * dimensions[:3] + + # Ensure correct dimensionality, if + # necessary + if same: + if grouping is not None or n_groups is not None: + shape = (n_groups or getattr(group, f"n_{grouping}"), -1, 3) + positions = positions.reshape(shape) + else: + if missing: + positions = [positions[g.atoms.ix] for g in groups] + else: + + # Try to convert arrays to NumPy arrays if they are not already + # to take advantage of vectorized operations later + try: + positions = np.asarray(positions) + if images is not None: + positions += images + except ValueError: + pass + if images is not None and type(images) is not type(positions): + emsg = ( + "The shapes of the arrays containing the positions " + "and image flags are incompatible." + ) + raise ValueError(emsg) + + # Calculate the dipole for the specified grouping + if isinstance(positions, np.ndarray): + + # Reshape the mass and position arrays based on the specified + # number of groups + if n_groups is not None: + positions = positions.reshape((n_groups, -1, 3)) + + # dipole = sum(q_i * r_i) and should be a vector of shape (3,) + cog = np.sum(positions, axis=-2) / positions.shape[-2] + else: + if images is not None: + for j, (p, i) in enumerate(zip(positions, images)): + positions[j] = p + i * dimensions + cog = np.array([np.sum(p, axis=0) / p.shape[0] for p in positions]) + + # Return raw charges and positions, if desired + if raw: + return cog, positions + + return cog + def molecular_dipole(group: mda.AtomGroup = None, grouping: str = None, *, @@ -1316,10 +1435,10 @@ def __init__( if not all(ax in "xyz" for ax in self.dipole_axes): raise ValueError("Invalid axis passed in 'dipole_axes'.") - if dipole_center not in {"center_of_mass", "center_of_charge"}: + if dipole_center not in {"center_of_mass", "center_of_charge", "center_of_geometry"}: emsg = ( f"Invalid dipole center '{dipole_center}'. " - "Valid values: 'center_of_mass', 'center_of_charge'." + "Valid values: 'center_of_mass', 'center_of_charge', 'center_of_geometry'." ) raise ValueError(emsg) self._dipole_center = dipole_center @@ -1599,6 +1718,10 @@ def _single_frame(self) -> None: self._positions[s] = ( ag.positions if gr == "atoms" else center_of_charge(ag, gr) ) + elif self._dipole_center == "center_of_geometry": + self._positions[s] = ( + ag.positions if gr == "atoms" else center_of_geometry(ag, gr) + ) self._dipoles[s] = molecular_dipole(ag, gr) self._quadrupoles[s] = molecular_quadrupole(ag, gr)