def fit(self,
Y_normalized,
embedding,
initialization,
iterations=100,
max_concentration_cw=100,
max_concentration_vmf=500):
"""
Args:
Y_normalized: Mix with shape (..., T, D).
embedding: Embedding from Deep Clustering with shape (..., T, E).
initialization: Shape (..., K, T)
iterations:
max_concentration_cw:
max_concentration_vmf:
Returns:
"""
Y_normalized_for_psd = np.copy(np.swapaxes(Y_normalized, -2, -1), 'C')
Y_normalized_for_pdf = np.copy(Y_normalized, 'C')
affiliations = np.copy(initialization)
D = Y_normalized.shape[-2]
cw = ComplexWatson(D, max_concentration=max_concentration_cw)
for i in range(iterations):
# E step
if i > 0:
affiliations = self.predict(Y_normalized_for_pdf, embedding)
# M step
self.pi = affiliations.mean(axis=-1)
Phi = get_power_spectral_density_matrix(Y_normalized_for_psd,
np.copy(affiliations, 'C'),
sensor_dim=-2,
source_dim=-2,
time_dim=-1)
self.W, eigenvalues = get_pca(Phi)
self.kappa_cw = cw.hypergeometric_ratio_inverse(eigenvalues)
self.mu, self.kappa_vmf = VonMisesFisher.fit(
embedding,
reshape(affiliations, 'fkt->k,t*f'),
max_concentration=max_concentration_vmf)
def fit(self,
Y_normalized,
initialization,
iterations=100,
max_concentration=100):
""" EM for CWMMs with any number of independent dimensions.
Does not support sequence lengths.
Can later be extended to accept more initializations, but for now
only accepts affiliations (masks) as initialization.
Args:
Y_normalized: Mix with shape (..., T, D).
initialization: Shape (..., K, T)
iterations: Most of the time 10 iterations are acceptable.
max_concentration: For numerical stability reasons.
"""
Y_normalized_for_pdf = np.copy(Y_normalized, 'C')
Y_normalized_for_psd = np.copy(np.swapaxes(Y_normalized, -2, -1), 'C')
D = Y_normalized.shape[-2]
cw = ComplexWatson(D, max_concentration=max_concentration)
affiliations = np.copy(initialization)
for i in range(iterations):
# E step
if i > 0:
affiliations = self.predict(Y_normalized_for_pdf)
# M step
self.pi = np.mean(affiliations, axis=-1)
Phi = get_power_spectral_density_matrix(Y_normalized_for_psd,
np.copy(affiliations, 'C'),
sensor_dim=-2,
source_dim=-2,
time_dim=-1)
self.W, eigenvalues = get_pca(Phi)
self.kappa = cw.hypergeometric_ratio_inverse(eigenvalues)
def fit(
self,
Y,
initialization,
source_activity_mask=None,
iterations=100,
max_concentration=100,
) -> ComplexWatsonMixtureModelParameters:
""" EM for CWMMs with any number of independent dimensions.
Does not support sequence lengths.
Can later be extended to accept more initializations, but for now
only accepts affiliations (masks) as initialization.
Args:
Y_normalized: Mix with shape (..., T, D).
initialization: Shape (..., K, T)
iterations: Most of the time 10 iterations are acceptable.
max_concentration: For numerical stability reasons.
"""
*independent, T, D = Y.shape
independent = tuple(independent)
assert D < 20, (D, 'Sure?')
if isinstance(initialization, self.Parameters):
K = initialization.mixture_weights.shape[-1]
assert K < 20, (K, 'Sure?')
else:
K = initialization.shape[-2]
assert K < 20, (K, 'Sure?')
assert initialization.shape[-1] == T, (initialization.shape, T)
assert initialization.shape[:-2] == independent, (
initialization.shape, independent)
Y = _unit_norm(Y, axis=-1, eps=1e-10, eps_style='where')
Y_normalized_for_pdf = np.ascontiguousarray(Y)
Y_normalized_for_psd = np.ascontiguousarray(np.swapaxes(Y, -2, -1))
if isinstance(initialization, self.Parameters):
params = initialization
else:
params = self.Parameters(eps=self.eps)
params.affiliation = np.copy(initialization) # Shape (..., K, T)
cw = ComplexWatson(D, max_concentration=max_concentration)
if isinstance(initialization, self.Parameters):
range_iterations = range(1, 1 + iterations)
else:
range_iterations = range(iterations)
if self.pbar:
import tqdm
range_iterations = tqdm.tqdm(range_iterations, 'cWMM Iteration')
else:
range_iterations = range_iterations
for i in range_iterations:
# E step
if i > 0:
params.affiliation = params._predict(
Y_normalized_for_pdf,
source_activity_mask=source_activity_mask,
)
# M step
params.mixture_weights = np.mean(params.affiliation, axis=-1)
Phi = get_power_spectral_density_matrix(Y_normalized_for_psd,
np.maximum(
params.affiliation,
params.eps),
sensor_dim=-2,
source_dim=-2,
time_dim=-1)
params.complex_watson.mode, eigenvalues = get_pca(Phi)
params.complex_watson.concentration = \
cw.hypergeometric_ratio_inverse(eigenvalues)
return params