diff --git a/ot/gmm.py b/ot/gmm.py index 3cb6cea6a..778f83b2a 100644 --- a/ot/gmm.py +++ b/ot/gmm.py @@ -249,6 +249,36 @@ def gmm_ot_plan(m_s, m_t, C_s, C_t, w_s, w_t, log=False): return emd(w_s, w_t, D, log=log) +def logsumexp(a, scaling_factor, axis=None): + """ + Computes log(sum(scaling_factor * exp(x))) stably using the log-sum-exp trick + with per-element weight. The backend nx.logsumexp does not allow passing + scaling weights. + + Parameters + ---------- + a : array-like + Log-values to sum. + scaling_factor : array-like + Weights for each term, must be of the same shape as a. + + Returns + ------- + float + log(sum(scaling_factor * exp(a))), computed stably. + + References + ---------- + Gundersen, G. (2020). The Log-Sum-Exp trick. Blog Post. Retrieved from https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/ + """ + nx = get_backend(a, scaling_factor) + if scaling_factor is None: + scaling_factor = 1 + shift = nx.max(a) + y = shift + nx.log(nx.sum(scaling_factor * nx.exp(a - shift), axis=axis)) + return y + + def gmm_ot_apply_map( x, m_s, m_t, C_s, C_t, w_s, w_t, plan=None, method="bary", seed=None ): @@ -322,10 +352,14 @@ def gmm_ot_apply_map( # gaussian mapping between components i and j applied to x T_ij_x = x @ A + b - z = w_s[:, None, None] * nx.exp(logpdf - logpdf[i][None, :, :]) - denom = nx.sum(z, axis=0) - out = out + plan[i, j] * T_ij_x / denom + log_g_i_x = logpdf[i] + # Could be optimized, that's not too smart to compute denom here at each iteration + denom = logsumexp( + logpdf.squeeze(), scaling_factor=w_s.reshape((-2, 1)), axis=0 + ) + p_ij_x = plan[i, j] * nx.exp(log_g_i_x - denom.reshape((-2, 1))) + out = out + p_ij_x * T_ij_x return out @@ -334,8 +368,8 @@ def gmm_ot_apply_map( # i and j, b[i, j] is the translation part rng = np.random.RandomState(seed) - A = nx.zeros((k_s, k_t, d, d)) - b = nx.zeros((k_s, k_t, d)) + A = nx.zeros((k_s, k_t, d, d), type_as=C_s) + b = nx.zeros((k_s, k_t, d), type_as=A) # only need to compute for non-zero plan entries for i, j in zip(*nx.where(plan > 0)): @@ -354,10 +388,10 @@ def gmm_ot_apply_map( for i_sample in range(n_samples): log_g = logpdf[i_sample] - log_diff = log_g[:, None] - log_g[None, :] - weighted_exp = w_s[:, None] * nx.exp(log_diff) - denom = nx.sum(weighted_exp, axis=0)[:, None] * nx.ones(plan.shape[1]) - p_mat = plan / denom + log_denom = logsumexp(log_g, scaling_factor=w_s) + p_mat = plan * nx.exp( + log_g.reshape((k_s, 1)) - log_denom + ) # shape (k_s, k_t): p_mat[i,j] = plan[i,j]*g_i(x)/D(x) p = p_mat.reshape(k_s * k_t) # stack line-by-line # sample between 0 and k_s * k_t - 1 diff --git a/test/test_gmm.py b/test/test_gmm.py index 629a68d57..bd9324343 100644 --- a/test/test_gmm.py +++ b/test/test_gmm.py @@ -6,6 +6,7 @@ # # License: MIT License +import warnings import numpy as np import pytest from ot.utils import proj_simplex @@ -18,6 +19,7 @@ gmm_ot_apply_map, gmm_ot_plan_density, gmm_barycenter_fixed_point, + logsumexp, ) try: @@ -160,6 +162,47 @@ def test_gmm_apply_map(): gmm_ot_apply_map(x, m_s, m_t, C_s, C_t, w_s, w_t, plan=plan) +@pytest.skip_backend("tf") # skips because of array assignment +@pytest.skip_backend("jax") +def test_gmm_apply_map_overflow(nx): + x_coord = 12.0 + d = 1 + k = 2 + x = nx.from_numpy(np.array([x_coord]).reshape((-1, d))) + + m_s = nx.from_numpy(np.array([x_coord, 0.0], dtype=np.float64).reshape((k, d))) + m_t = nx.from_numpy(np.array([x_coord, 0.0], dtype=np.float64).reshape((k, d))) + + C_s = nx.from_numpy(np.ones((k, d, d)) / 10.0) + C_t = nx.from_numpy(np.ones((k, d, d)) / 10.0) + + w_s = nx.from_numpy(np.array([0.5, 0.5])) + w_t = nx.from_numpy(np.array([0.5, 0.5])) + + with warnings.catch_warnings(): + warnings.simplefilter(action="error") + gmm_ot_apply_map( + x, + m_s, + m_t, + C_s, + C_t, + w_s, + w_t, + method="rand", + seed=0, + ) + + +def test_logsumexp_overflow_safety(nx): + """Test that large exponents don't overflow.""" + x = nx.from_numpy(np.array([0.0, 710.0])) + result = logsumexp(x, scaling_factor=nx.from_numpy(np.array([1, 1]))) + # Should be equal to exp(log(0)) + exp(log(710)) = 710 + assert nx.isfinite(result) + assert nx.allclose(result, nx.from_numpy(np.array([710.0]))) + + @pytest.mark.skipif(not torch, reason="No torch available") def test_gradient_gmm_ot_loss_pytorch(): m_s, m_t, C_s, C_t, w_s, w_t = get_gmms()