From ae73b327edfde7fa1a6969181152d462d09a3df8 Mon Sep 17 00:00:00 2001 From: Timotej Bernat Date: Wed, 24 Apr 2024 17:53:50 -0600 Subject: [PATCH] Added optional weighting to centroid calculation --- polymerist/maths/lattices/coordops.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/polymerist/maths/lattices/coordops.py b/polymerist/maths/lattices/coordops.py index 99ed6ce..f313be3 100644 --- a/polymerist/maths/lattices/coordops.py +++ b/polymerist/maths/lattices/coordops.py @@ -6,12 +6,17 @@ import numpy as np from itertools import product as cartesian_product -from ...genutils.typetools.numpytypes import Shape, N, M, L, DType +from ...genutils.typetools.numpytypes import Shape, N, M, DType -def mean_coord(coords : np.ndarray[Shape[M, N], DType]) -> np.ndarray[Shape[N], DType]: - '''The average (center-of-mass) coordinate of a vector of coordinates''' - return coords.mean(axis=0) +def mean_coord(coords : np.ndarray[Shape[M, N], DType], weights : Optional[np.ndarray[Shape[M, 1], DType]]=None) -> np.ndarray[Shape[N], DType]: + '''The average (center-of-mass) coordinate of a vector of coordinates, with optional array of weightsfor each coordinates''' + if weights is None: + weights = np.ones((coords.shape[0], 1), dtype=coords.dtype) + assert(weights.size == coords.shape[0]) + weights = weights.reshape((coords.shape[0], 1)) + + return (coords * weights).mean(axis=0) center_of_mass = COM = mean_coord def dists_to_point(coords : np.ndarray[Shape[M, N], DType], point : np.ndarray[Shape[N], DType], norm_order : Optional[int]=None) -> np.ndarray[Shape[M], DType]: