acteval 0.1.0__py3-none-any.whl

acteval/distance/wasserstein.py

@@ -0,0 +1,200 @@
+import numpy as np
+from ot import dist, emd2, emd2_1d
+from scipy.stats import wasserstein_distance
+
+try:
+    import torch
+
+    _TORCH_AVAILABLE = True
+except ImportError:
+    _TORCH_AVAILABLE = False
+
+
+def emd(a: tuple[np.array, np.array], b: tuple[np.array, np.array]) -> float:
+    if a[0].ndim == 1:
+        return emd1d(a, b)
+    elif a[0].ndim == 2:
+        return emd2d(a, b)
+    else:
+        raise ValueError("Only 1d and 2d features are supported")
+
+
+def emd2d(a: tuple[np.array, np.array], b: tuple[np.array, np.array]) -> float:
+    ak, aw = a
+    bk, bw = b
+    aw = aw / aw.sum()
+    bw = bw / bw.sum()
+    d = dist(ak, bk, metric="cityblock")
+    return emd2(aw, bw, d, check_marginals=False)
+
+
+def emd1d(a: tuple[np.array, np.array], b: tuple[np.array, np.array]) -> float:
+    ak, aw = a
+    bk, bw = b
+    if (
+        aw.sum() == 0
+    ):  # avoid division by zero but also has to assume distribution of just [0]
+        ak = np.array([0.0])
+        aw = np.array([1.0])
+    if bw.sum() == 0:
+        bk = np.array([0.0])
+        bw = np.array([1.0])
+    aw = aw / aw.sum()
+    bw = bw / bw.sum()
+    return emd2_1d(ak, bk, aw, bw, metric="cityblock")
+
+
+def wasserstein(x: list[list], y: list[list]):
+    return wasserstein_distance(x, y)
+
+
+def sliced_wasserstein(x: list[list], y: list[list], num_proj=100):
+    """
+    https://stats.stackexchange.com/questions/404775/calculate-earth-movers-distance-for-two-grayscale-images
+    """
+    x = np.array(x)
+    y = np.array(y)
+    dim = x.shape[1]
+    ests = []
+    for _ in range(num_proj):
+        # sample uniformly from the unit sphere
+        dir = np.random.randn(dim)
+        dir /= np.linalg.norm(dir)
+
+        # project the data
+        x_proj = x @ dir
+        y_proj = y @ dir
+
+        # compute 1d wasserstein
+        ests.append(wasserstein_distance(x_proj, y_proj))
+    return np.mean(ests)
+
+
+def sinkhorn(x: list[list], y: list[list], eps=0.01, max_iter=10, reduction=None):
+    if not _TORCH_AVAILABLE:
+        raise ImportError("torch is required for sinkhorn")
+    x = torch.tensor(x, dtype=torch.float)
+    y = torch.tensor(y, dtype=torch.float)
+    model = SinkhornDistance(eps=eps, max_iter=max_iter, reduction=reduction)
+    return model(x, y)[0].item()
+
+
+# Adapted from https://github.com/gpeyre/SinkhornAutoDiff
+class SinkhornDistance:
+    r"""
+    https://dfdazac.github.io/sinkhorn.html
+    Given two empirical measures each with :math:`P_1` locations
+    :math:`x\in\mathbb{R}^{D_1}` and :math:`P_2` locations :math:`y\in\mathbb{R}^{D_2}`,
+    outputs an approximation of the regularized OT cost for point clouds.
+
+    Args:
+        eps (float): regularization coefficient
+        max_iter (int): maximum number of Sinkhorn iterations
+        reduction (string, optional): Specifies the reduction to apply to the output:
+            'none' | 'mean' | 'sum'. 'none': no reduction will be applied,
+            'mean': the sum of the output will be divided by the number of
+            elements in the output, 'sum': the output will be summed. Default: 'none'
+
+    Shape:
+        - Input: :math:`(N, P_1, D_1)`, :math:`(N, P_2, D_2)`
+        - Output: :math:`(N)` or :math:`()`, depending on `reduction`
+    """
+
+    def __init__(self, eps, max_iter, reduction="none"):
+        if not _TORCH_AVAILABLE:
+            raise ImportError("torch is required for SinkhornDistance")
+        super(SinkhornDistance, self).__init__()
+        self.eps = eps
+        self.max_iter = max_iter
+        self.reduction = reduction
+
+    def forward(self, x, y):
+        # The Sinkhorn algorithm takes as input three variables :
+        print("calc cost matrix")
+        C = self._cost_matrix(x, y)  # Wasserstein cost function
+        x_points = x.shape[-2]
+        y_points = y.shape[-2]
+        if x.dim() == 2:
+            batch_size = 1
+        else:
+            batch_size = x.shape[0]
+
+        # both marginals are fixed with equal weights
+        print("mu")
+        mu = (
+            torch.empty(batch_size, x_points, dtype=torch.float, requires_grad=False)
+            .fill_(1.0 / x_points)
+            .squeeze()
+        )
+        print("nu")
+        nu = (
+            torch.empty(batch_size, y_points, dtype=torch.float, requires_grad=False)
+            .fill_(1.0 / y_points)
+            .squeeze()
+        )
+
+        u = torch.zeros_like(mu)
+        v = torch.zeros_like(nu)
+        # To check if algorithm terminates because of threshold
+        # or max iterations reached
+        actual_nits = 0
+        # Stopping criterion
+        thresh = 1e-1
+
+        # Sinkhorn iterations
+        for i in range(self.max_iter):
+            print(i, "start")
+            u1 = u  # useful to check the update
+            u = (
+                self.eps
+                * (torch.log(mu + 1e-8) - torch.logsumexp(self.M(C, u, v), dim=-1))
+                + u
+            )
+            v = (
+                self.eps
+                * (
+                    torch.log(nu + 1e-8)
+                    - torch.logsumexp(self.M(C, u, v).transpose(-2, -1), dim=-1)
+                )
+                + v
+            )
+            err = (u - u1).abs().sum(-1).mean()
+            print(i, err.item())
+
+            actual_nits += 1
+            if err.item() < thresh:
+                print(f"Sinkhorn converged at iteration {i}")
+                break
+
+        U, V = u, v
+        # Transport plan pi = diag(a)*K*diag(b)
+        pi = torch.exp(self.M(C, U, V))
+        # Sinkhorn distance
+        cost = torch.sum(pi * C, dim=(-2, -1))
+
+        if self.reduction == "mean":
+            cost = cost.mean()
+        elif self.reduction == "sum":
+            cost = cost.sum()
+
+        return cost, pi, C
+
+    def __call__(self, x, y):
+        return self.forward(x, y)
+
+    def M(self, C, u, v):
+        "Modified cost for logarithmic updates"
+        return (-C + u.unsqueeze(-1) + v.unsqueeze(-2)) / self.eps
+
+    @staticmethod
+    def _cost_matrix(x, y, p=1):
+        "Returns the matrix of $|x_i-y_j|^p$."
+        x_col = x.unsqueeze(-2)
+        y_lin = y.unsqueeze(-3)
+        C = torch.sum((torch.abs(x_col - y_lin)) ** p, -1)
+        return C
+
+    @staticmethod
+    def ave(u, u1, tau):
+        "Barycenter subroutine, used by kinetic acceleration through extrapolation."
+        return tau * u + (1 - tau) * u1