import numpy as np
from smonitor import signal
from molsysmt import pyunitwizard as puw
from molsysmt._private import rust_backend as _kernels
from molsysmt._private.argdigest import arg_digest
from molsysmt._private.execution import Reducer
from molsysmt._private.smonitor import NotImplementedMethodError
from molsysmt._private.variables import is_all, is_iterable_of_iterables
from molsysmt._private.weighted_geometry import prepare_weights
from molsysmt.lib.structure._kernel_inputs import extract_coordinates_value_and_unit
class _CenterReducer(Reducer):
"""
Reducer for get_center heavy path.
Accumulates per-chunk center arrays and concatenates them on finalize.
"""
def __init__(self, weights, atoms_per_group=None):
self._weights = weights
self._atoms_per_group = atoms_per_group
self._n_groups = 1 if atoms_per_group is None else len(atoms_per_group)
self._chunks = []
def initialize(self, metadata):
self._chunks = []
def consume(self, chunk):
coords = chunk["coordinates"] # (chunk_size, n_atoms, 3), float64, read-only
coords_w = np.array(coords, dtype=np.float64) # writable copy for kernel
if self._atoms_per_group is None:
result = _kernels.get_center(coords_w, self._weights)
else:
result = _kernels.get_center_groups_of_atoms(
coords_w, self._atoms_per_group, self._weights
)
self._chunks.append(result)
def finalize(self):
return np.concatenate(self._chunks, axis=0)
# --- optional extensions ---
def checkpoint(self):
return {"chunks": [c.tolist() for c in self._chunks]}
def restore(self, state):
self._chunks = [np.array(c, dtype=np.float64) for c in state["chunks"]]
def merge(self, other):
self._chunks.extend(other._chunks)
# Keep this decorator import beside the public function below.
from molsysmt.configure import with_configure_overrides # noqa: E402
[docs]
@signal(tags=["api", "structure"])
@arg_digest()
@with_configure_overrides
def get_center(
molecular_system,
selection="all",
weights=None,
structure_indices="all",
syntax="MolSysMT",
engine="MolSysMT",
heavy_mode="auto",
parallel=None,
num_threads=None,
skip_digestion=False,
):
"""
Computing centers (centroids or weighted centers) of atom selections.
Parameters
----------
molecular_system : molecular system
Molecular system in any supported MolSysMT format.
selection : str, list, tuple, or numpy.ndarray, default='all'
Selection string or boolean/integer array specifying elements.
weights : numpy.ndarray, list, or tuple, default=None
Atomic mass weights array for center calculation.
structure_indices : int, list, tuple, or numpy.ndarray, default='all'
Structure indices (0-based) to include or process.
syntax : str, default='MolSysMT'
Selection syntax used to evaluate `selection` (e.g., 'MolSysMT', 'MDTraj').
engine : str, default='MolSysMT'
Backend used to perform the calculation.
heavy_mode : {'auto', 'force', 'off'}, default='auto'
Policy controlling the chunked path for memory-intensive calculations.
parallel : bool or None, default=None
Whether to use CPU parallelism, or `None` to use the configured default.
num_threads : int or None, default=None
Number of CPU worker threads, or `None` to use the configured default.
skip_digestion : bool, default=False
Whether to skip MolSysMT's internal argument digestion mechanism.
Returns
-------
quantity
Centers as a PyUnitWizard quantity in length units.
Raises
------
ArgumentError
If the atom or frame selection is empty, or weights are non-finite,
negative, or have zero total weight.
ArgumentLengthError
If the number of weights does not match the selected atoms.
NotImplementedMethodError
If an unsupported engine is requested.
Notes
-----
Chunked calculations use nm internally and return a quantity standardized
to the session's length unit. Native iteration reads the existing structures
domain and copies only selected coordinate blocks; accumulated output still
occupies memory.
Examples
--------
>>> import molsysmt as msm
>>> from molsysmt.structure.get_center import get_center
>>> molsys = msm.convert(msm.systems['alanine dipeptide']['alanine_dipeptide.h5msm'], to_form='molsysmt.MolSys')
>>> get_center(molsys, weights='masses').shape
(1, 1, 3)
.. versionadded:: 1.0.0
"""
from molsysmt._private.structure_indices import ensure_nonempty_structure_indices
from molsysmt.basic import get, select
ensure_nonempty_structure_indices(
structure_indices,
caller="molsysmt.structure.get_center",
)
if engine == "MolSysMT":
atom_indices = select(molecular_system, selection=selection)
if not is_iterable_of_iterables(atom_indices):
n_atoms = (
len(np.atleast_1d(atom_indices))
if not is_all(atom_indices)
else get(molecular_system, element="system", n_atoms=True)
)
n_structures = get(molecular_system, element="system", n_structures=True)
weights_arr = prepare_weights(
weights,
n_atoms,
molecular_system=molecular_system,
selection=atom_indices,
syntax=syntax,
caller="molsysmt.structure.get_center",
)
from molsysmt._private.execution import ChunkedExecutor
from molsysmt._private.execution.memory_policy import (
decide_mode,
estimate_footprint,
)
from molsysmt.basic import get_form
form = get_form(molecular_system)
footprint = estimate_footprint(n_atoms, n_structures)
mode = decide_mode(footprint, heavy_mode)
if mode == "heavy":
reducer = _CenterReducer(weights=weights_arr)
executor = ChunkedExecutor(
molecular_system=molecular_system,
form=form,
operation="get_center",
reducer=reducer,
atom_indices=atom_indices,
structure_indices=None
if is_all(structure_indices)
else structure_indices,
heavy_mode=heavy_mode,
attributes=["coordinates"],
)
center_val = executor.execute() # (n_structures, 1, 3), float64, nm
center = puw.quantity(center_val, "nm")
else:
coordinates = get(
molecular_system,
element="atom",
selection=atom_indices,
structure_indices=structure_indices,
coordinates=True,
)
coordinates, length_unit = extract_coordinates_value_and_unit(
coordinates
)
center = _kernels.get_center(coordinates, weights_arr)
center = puw.quantity(center, length_unit)
del coordinates, length_unit
else:
atoms_per_group = np.array(
[len(group) for group in atom_indices], dtype=np.int64
)
groups_of_atoms = np.concatenate(atom_indices)
n_atoms_flat = len(groups_of_atoms)
n_structures = get(molecular_system, element="system", n_structures=True)
if weights is not None and not isinstance(weights, str):
if is_iterable_of_iterables(weights):
weights = np.concatenate(weights)
weights_arr = prepare_weights(
weights,
n_atoms_flat,
molecular_system=molecular_system,
selection=groups_of_atoms,
syntax=syntax,
group_sizes=atoms_per_group,
caller="molsysmt.structure.get_center",
)
from molsysmt._private.execution.memory_policy import (
decide_mode,
estimate_footprint,
)
from molsysmt.basic import get_form
form = get_form(molecular_system)
footprint = estimate_footprint(n_atoms_flat, n_structures)
mode = decide_mode(footprint, heavy_mode)
if mode == "heavy":
reducer = _CenterReducer(
weights=weights_arr, atoms_per_group=atoms_per_group
)
from molsysmt._private.execution import ChunkedExecutor
executor = ChunkedExecutor(
molecular_system=molecular_system,
form=form,
operation="get_center",
reducer=reducer,
atom_indices=groups_of_atoms,
structure_indices=None
if is_all(structure_indices)
else structure_indices,
heavy_mode=heavy_mode,
attributes=["coordinates"],
)
center_val = executor.execute() # (n_structures, n_groups, 3)
center = puw.quantity(center_val, "nm")
else:
coordinates = get(
molecular_system,
element="atom",
selection=groups_of_atoms,
structure_indices=structure_indices,
coordinates=True,
)
coordinates, length_unit = extract_coordinates_value_and_unit(
coordinates
)
center = _kernels.get_center_groups_of_atoms(
coordinates, atoms_per_group, weights_arr
)
center = puw.quantity(center, length_unit)
del coordinates, length_unit, groups_of_atoms, weights_arr
center = puw.standardize(center)
return center
else:
raise NotImplementedMethodError()
