r""" Apply Gaussian Mixture Model (GMM) optimal transport (OT) mapping to input data. The 'barycentric' mapping corresponds to the barycentric projection of the GMM-OT plan, and is called T_bary in [69]. The 'random' mapping takes for each input point a random pair (i,j) of component
(
x, m_s, m_t, C_s, C_t, w_s, w_t, plan=None, method="bary", seed=None
)
| 250 | |
| 251 | |
| 252 | def gmm_ot_apply_map( |
| 253 | x, m_s, m_t, C_s, C_t, w_s, w_t, plan=None, method="bary", seed=None |
| 254 | ): |
| 255 | r""" |
| 256 | Apply Gaussian Mixture Model (GMM) optimal transport (OT) mapping to input |
| 257 | data. The 'barycentric' mapping corresponds to the barycentric projection |
| 258 | of the GMM-OT plan, and is called T_bary in [69]. The 'random' mapping takes |
| 259 | for each input point a random pair (i,j) of components of the GMMs and |
| 260 | applied the Gaussian map, it is called T_rand in [69]. |
| 261 | |
| 262 | Parameters |
| 263 | ---------- |
| 264 | x : array-like, shape (n_samples, d) |
| 265 | Input data points. |
| 266 | m_s : array-like, shape (k_s, d) |
| 267 | Mean vectors of the source GMM components. |
| 268 | m_t : array-like, shape (k_t, d) |
| 269 | Mean vectors of the target GMM components. |
| 270 | C_s : array-like, shape (k_s, d, d) |
| 271 | Covariance matrices of the source GMM components. |
| 272 | C_t : array-like, shape (k_t, d, d) |
| 273 | Covariance matrices of the target GMM components. |
| 274 | w_s : array-like, shape (k_s,) |
| 275 | Weights of the source GMM components. |
| 276 | w_t : array-like, shape (k_t,) |
| 277 | Weights of the target GMM components. |
| 278 | plan : array-like, shape (k_s, k_t), optional |
| 279 | Optimal transport plan between the source and target GMM components. |
| 280 | If not provided, it will be computed internally. |
| 281 | method : {'bary', 'rand'}, optional |
| 282 | Method for applying the GMM OT mapping. 'bary' uses barycentric mapping, |
| 283 | while 'rand' uses random sampling. Default is 'bary'. |
| 284 | seed : int, optional |
| 285 | Seed for the random number generator. Only used when method='rand'. |
| 286 | |
| 287 | Returns |
| 288 | ------- |
| 289 | out : array-like, shape (n_samples, d) |
| 290 | Output data points after applying the GMM OT mapping. |
| 291 | |
| 292 | References |
| 293 | ---------- |
| 294 | .. [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. |
| 295 | |
| 296 | """ |
| 297 | |
| 298 | if plan is None: |
| 299 | plan = gmm_ot_plan(m_s, m_t, C_s, C_t, w_s, w_t) |
| 300 | nx = get_backend(x, m_s, m_t, C_s, C_t, w_s, w_t) |
| 301 | else: |
| 302 | nx = get_backend(x, m_s, m_t, C_s, C_t, w_s, w_t, plan) |
| 303 | |
| 304 | k_s, k_t = m_s.shape[0], m_t.shape[0] |
| 305 | d = m_s.shape[1] |
| 306 | n_samples = x.shape[0] |
| 307 | |
| 308 | if method == "bary": |
| 309 | out = nx.zeros(x.shape) |