def interpolate(interpolate_n, probs, radii):
"""Interpoleate the probailites to evenly spaced radii
Using the max and min radii with `interpolate_n` elements
interpolate the previously caculated probabilites to new spacing.
Args:
interpolate_n: the number of elements
probs: the probabilites
radii: the radii
Returns:
Tuple of new probabilites and radii
"""
if interpolate_n is None:
return probs, radii
radii_interp = np.linspace(radii.min(), radii.max(), interpolate_n)
interp = lambda x: np.interp(radii_interp, radii, x)
return (
sequence(np.transpose, fmap(interp), list, np.array,
np.transpose)(probs),
radii_interp,
)
def np_paircorr(x_data, cutoff_r=None, interpolate_n=None):
"""Numpy only version of `paircorr_from_twopoint`"""
return sequence(
calc_radii,
lambda x: (calc_probs(x_data, x[2], x[1]), x[0]),
star(calc_cutoff(cutoff_r)),
star(interpolate(interpolate_n)),
lambda x: (np.transpose(x[0]), x[1]),
)(x_data[0].shape)
def calc_radii(shape):
"""Computer the radii with unique distances from the center
Args:
shape: the shape of the domain
Returns:
a tuple of (radii, index_split, index_sort) where the
index_split are groupings of indices at the same distance
and index_sort are the sorted indices
"""
flatten = lambda x: x.reshape(-1)
return sequence(
dist_mesh,
flatten,
sort_array,
lambda x: np.unique(x[0], return_index=True) + (x[1], ),
)(shape)
def calc_probs(x_data, ind_sort, ind_split):
"""Compute the two point probabilities averages radially
Args:
x_data: the two point statistics
ind_sort: the sorted index of distances from the center
ind_split: the index groupings
Returns:
the calculated probabilites averaged over each radii
grouping
"""
return sequence(
lambda x: x.reshape(x.shape[0], -1)[:, ind_sort],
lambda x: np.split(x, ind_split[1:], axis=1),
fmap(lambda x: np.average(x, axis=1)),
list,
np.array,
)(x_data)