def mix_normals(means,
sigmas,
weights=None,
normalize=False,
axis=None,
keepdims=False):
"""
Return the mixture distribution of a sum of random variables.
See https://en.wikipedia.org/wiki/Mixture_distribution#Moments.
Arguments:
means (numpy.ndarray): Variable means
sigmas (numpy.ndarray): Variable standard deviations
weights (numpy.ndarray): Variable weights
normalize (bool): Whether to normalize weights so that they sum to 1
for non-missing values
"""
isnan_mean, isnan_sigmas, weights = prepare_normals(
means, sigmas, weights, normalize, axis)
wmeans = np.nansum(weights * means, axis=axis, keepdims=True)
variances = np.nansum(
weights * (means**2 + sigmas**2), axis=axis, keepdims=True) - wmeans**2
# np.nansum interprets sum of nans as 0
isnan = isnan_mean.all(axis=axis, keepdims=True)
wmeans[isnan] = np.nan
variances[isnan] = np.nan
return (numpy_dropdims(wmeans, axis=axis, keepdims=keepdims),
numpy_dropdims(np.sqrt(variances), axis=axis, keepdims=keepdims))
def select_repeat_tracks(runs):
"""
Return Tracks composed of the best track for each initial point.
Selects the track for each point that minimizes the temporal mean standard
deviation for vx + vy, i.e. mean(sqrt(vx_sigma**2 + vy_sigma**2)).
Arguments:
runs (iterable): Tracks objects with identical point and time dimensions
"""
# Compute metric for each track
metric = np.row_stack([
np.nanmean(np.sqrt(np.nansum(run.sigmas[..., 3:5]**2, axis=2)), axis=1)
for run in runs
])
# Choose run with the smallest metric
selected = np.argmin(metric, axis=0)
# Merge runs
means = runs[0].means.copy()
sigmas = runs[0].sigmas.copy()
for i, run in enumerate(runs[1:], start=1):
mask = selected == i
means[mask, ...] = run.means[mask, ...]
sigmas[mask, ...] = run.sigmas[mask, ...]
return glimpse.Tracks(datetimes=runs[0].datetimes,
means=means,
sigmas=sigmas)
def flatten_tracks_doug(runs):
# Join together second forward and backward runs
f, r = runs['fv'], runs['rv']
means = np.column_stack((f.means[..., 3:], r.means[..., 3:]))
sigmas = np.column_stack((f.sigmas[..., 3:], r.sigmas[..., 3:]))
# Flatten joined runs
# Mean: Inverse-variance weighted mean
# Sigma: Linear combination of weighted correlated random variables
# (approximation using the weighted mean of the variances)
weights = sigmas**-2
weights *= 1 / np.nansum(weights, axis=1, keepdims=True)
allnan = np.isnan(means).all(axis=1, keepdims=True)
means = np.nansum(weights * means, axis=1, keepdims=True)
sigmas = np.sqrt(np.nansum(weights * sigmas**2, axis=1, keepdims=True))
# np.nansum interprets sum of nans as 0
means[allnan] = np.nan
sigmas[allnan] = np.nan
return means.squeeze(axis=1), sigmas.squeeze(axis=1)
def sum_normals(means,
sigmas,
weights=None,
normalize=False,
correlation=0,
axis=None,
keepdims=False):
"""
Return the mean and sigma of the sum of random variables.
See https://en.wikipedia.org/wiki/Propagation_of_uncertainty#Linear_combinations.
Arguments:
means (numpy.ndarray): Variable means
sigmas (numpy.ndarray): Variable standard deviations
weights (numpy.ndarray): Variable weights
normalize (bool): Whether to normalize weights so that they sum to 1
for non-missing values
correlation (float): Correlation to assume between pairs of different variables
"""
isnan_mean, isnan_sigmas, weights = prepare_normals(
means, sigmas, weights, normalize, axis)
wmeans = np.nansum(weights * means, axis=axis, keepdims=True)
# Initialize variance as sum of diagonal elements
variances = np.nansum(weights**2 * sigmas**2, axis=axis, keepdims=True)
# np.nansum interprets sum of nans as 0
allnan = isnan_mean.all(axis=axis, keepdims=True)
wmeans[allnan] = np.nan
variances[allnan] = np.nan
if correlation:
# Add off-diagonal elements
n = means.size if axis is None else means.shape[axis]
pairs = np.triu_indices(n=n, k=1)
variances += 2 * np.nansum(
correlation * np.take(weights, pairs[0], axis=axis) *
np.take(weights, pairs[1], axis=axis) *
np.take(sigmas, pairs[0], axis=axis) *
np.take(sigmas, pairs[1], axis=axis),
axis=axis,
keepdims=True)
return (numpy_dropdims(wmeans, axis=axis, keepdims=keepdims),
numpy_dropdims(np.sqrt(variances), axis=axis, keepdims=keepdims))
def prepare_normals(means, sigmas, weights, normalize, axis):
isnan_mean = np.isnan(means)
isnan_sigmas = np.isnan(sigmas)
if np.any(isnan_mean != isnan_sigmas):
raise ValueError('mean and sigma NaNs do not match')
if np.any(sigmas == 0):
raise ValueError('sigmas cannot be 0')
if weights is None:
weights = np.ones(means.shape)
if normalize:
weights = weights * (
1 / np.nansum(weights * ~isnan_mean, axis=axis, keepdims=True))
return isnan_mean, isnan_sigmas, weights
# Single median: Select mean, sigma of median neighbor
is_single_median = is_median & (np.sum(is_median, axis=1, keepdims=True)
== 1)
mask = is_single_median.any(axis=1)
fmeans[mask, dim] = m[is_single_median]
fsigmas[mask, dim] = s[is_single_median]
# Multiple median: Take unweighted average (correlation = 1)
is_multiple_median = is_median & ~is_single_median
mask = is_multiple_median.any(axis=1)
fmeans[mask, dim] = medians.squeeze(axis=1)[mask]
if exact_median_variance:
# "Exact" variance
n = s.shape[1]
pairs = np.triu_indices(n=n, k=1)
sqweights = 1 / np.sum(is_multiple_median, axis=1, keepdims=False)**2
variances = sqweights * np.nansum(s**2, axis=1, keepdims=False)
variances += 2 * sqweights * np.nansum(
np.take(s, pairs[0], axis=1) * np.take(s, pairs[1], axis=1),
axis=1,
keepdims=False)
fsigmas[mask, dim] = np.sqrt(variances[mask])
else:
# Approximate variance using the mean of the variances
fsigmas[mask, dim] = np.sqrt(np.nanmean(s**2, axis=1)[mask])
# Write files
glimpse.helpers.write_pickle(ti, os.path.join(arrays_path, 'ti.pkl'))
glimpse.helpers.write_pickle(fmeans, os.path.join(arrays_path, 'means.pkl'))
glimpse.helpers.write_pickle(fsigmas, os.path.join(arrays_path, 'sigmas.pkl'))
nobservers[np.isnan(means[..., 0])] = 0
glimpse.helpers.write_pickle(nobservers,