47 lines
1.5 KiB
Python
47 lines
1.5 KiB
Python
import numpy as np
|
|
|
|
def smartmean(x, dim=None):
|
|
"""
|
|
Python implementation of the MATLAB smartmean function.
|
|
Calculates the mean while ignoring NaN and Inf values.
|
|
|
|
Parameters:
|
|
x (ndarray): Input array
|
|
dim (int, optional): Dimension along which to compute the mean
|
|
|
|
Returns:
|
|
ndarray: Mean values along the specified dimension, ignoring NaN/Inf
|
|
"""
|
|
if isinstance(x, list):
|
|
x = np.array(x)
|
|
|
|
# If dimension is not specified, find the first non-singleton dimension
|
|
if dim is None:
|
|
# Find the first dimension with size > 1
|
|
if len(x.shape) == 1:
|
|
dim = 0
|
|
else:
|
|
non_singleton_dims = [i for i, size in enumerate(x.shape) if size > 1]
|
|
dim = non_singleton_dims[0] if non_singleton_dims else 0
|
|
|
|
# Create a mask of finite values
|
|
mask = np.isfinite(x)
|
|
|
|
# Replace non-finite values with zeros
|
|
x_clean = np.where(mask, x, 0)
|
|
|
|
# Sum the finite values and divide by the count of finite values
|
|
with np.errstate(divide='ignore', invalid='ignore'):
|
|
result = np.sum(x_clean, axis=dim) / np.sum(mask, axis=dim)
|
|
|
|
# Set result to NaN where all values are non-finite
|
|
if dim == 0:
|
|
result[np.all(~mask, axis=dim)] = np.nan
|
|
elif dim == 1:
|
|
result[np.all(~mask, axis=dim)] = np.nan
|
|
else:
|
|
all_nan_indices = np.where(np.all(~mask, axis=dim))
|
|
if all_nan_indices[0].size > 0:
|
|
result[all_nan_indices] = np.nan
|
|
|
|
return result |