r""" Compute the matrix of the squared Bures distances between the components of two Gaussian Mixture Models (GMMs). Used to compute the GMM Optimal Transport distance [69]. Parameters ---------- m_s : array-like, shape (k_s, d) Mean vectors of the source GMM. m_
(m_s, m_t, C_s, C_t)
| 104 | |
| 105 | |
| 106 | def dist_bures_squared(m_s, m_t, C_s, C_t): |
| 107 | r""" |
| 108 | Compute the matrix of the squared Bures distances between the components of |
| 109 | two Gaussian Mixture Models (GMMs). Used to compute the GMM Optimal |
| 110 | Transport distance [69]. |
| 111 | |
| 112 | Parameters |
| 113 | ---------- |
| 114 | m_s : array-like, shape (k_s, d) |
| 115 | Mean vectors of the source GMM. |
| 116 | m_t : array-like, shape (k_t, d) |
| 117 | Mean vectors of the target GMM. |
| 118 | C_s : array-like, shape (k_s, d, d) |
| 119 | Covariance matrices of the source GMM. |
| 120 | C_t : array-like, shape (k_t, d, d) |
| 121 | Covariance matrices of the target GMM. |
| 122 | |
| 123 | Returns |
| 124 | ------- |
| 125 | dist : array-like, shape (k_s, k_t) |
| 126 | Matrix of squared Bures distances between the components of the source |
| 127 | and target GMMs. |
| 128 | |
| 129 | References |
| 130 | ---------- |
| 131 | .. [69] Delon, J., & Desolneux, A. (2020). A Wasserstein-type distance in the space of Gaussian mixture models. SIAM Journal on Imaging Sciences, 13(2), 936-970. |
| 132 | |
| 133 | """ |
| 134 | nx = get_backend(m_s, C_s, m_t, C_t) |
| 135 | |
| 136 | assert m_s.shape[0] == C_s.shape[0], "Source GMM has different amount of components" |
| 137 | |
| 138 | assert m_t.shape[0] == C_t.shape[0], "Target GMM has different amount of components" |
| 139 | |
| 140 | assert ( |
| 141 | m_s.shape[-1] == m_t.shape[-1] == C_s.shape[-1] == C_t.shape[-1] |
| 142 | ), "All GMMs must have the same dimension" |
| 143 | |
| 144 | D_means = dist(m_s, m_t, metric="sqeuclidean") |
| 145 | |
| 146 | # C2[i, j] = Cs12[i] @ C_t[j] @ Cs12[i], shape (k_s, k_t, d, d) |
| 147 | Cs12 = nx.sqrtm(C_s) # broadcasts matrix sqrt over (k_s,) |
| 148 | C2 = nx.einsum("ikl,jlm,imn->ijkn", Cs12, C_t, Cs12) |
| 149 | C = nx.sqrtm(C2) # broadcasts matrix sqrt over (k_s, k_t) |
| 150 | |
| 151 | # D_covs[i,j] = trace(C_s[i] + C_t[j] - 2C[i,j]) |
| 152 | trace_C_s = nx.einsum("ikk->i", C_s)[:, None] # (k_s, 1) |
| 153 | trace_C_t = nx.einsum("ikk->i", C_t)[None, :] # (1, k_t) |
| 154 | D_covs = trace_C_s + trace_C_t # broadcasts to (k_s, k_t) |
| 155 | D_covs -= 2 * nx.einsum("ijkk->ij", C) |
| 156 | |
| 157 | return nx.maximum(D_means + D_covs, 0) |
| 158 | |
| 159 | |
| 160 | def gmm_ot_loss(m_s, m_t, C_s, C_t, w_s, w_t, log=False): |