add the definition of r to the description

Author: fegounna

ot/low_rank/_factor_relaxation.py

@@ -105,11 +105,11 @@ def solve_balanced_FRLC(
     .. math::
         \textbf{P} = \mathop{\arg \min}_P \quad \langle \textbf{P}, \mathbf{M} \rangle_F
 
-            \text{s.t.} \ \textbf{P} = \textbf{Q} \operatorname{diag}(1/g_Q)\textbf{T}\operatorname{diag}(1/g_R)\textbf{R}^T
+            \text{s.t.}  \textbf{P} = \textbf{Q} \operatorname{diag}(1/g_Q)\textbf{T}\operatorname{diag}(1/g_R)\textbf{R}^T
 
-                \textbf{Q} &\in \Pi_{a,\cdot}, \quad \textbf{R} \in \Pi_{b,\cdot}, \quad \textbf{T} \in \Pi_{g_Q,g_R}
+                \textbf{Q} \in \Pi_{a,\cdot}, \quad \textbf{R} \in \Pi_{b,\cdot}, \quad \textbf{T} \in \Pi_{g_Q,g_R}
 
-                \textbf{Q}, \textbf{R}, \textbf{T} &\geq 0
+                \textbf{Q}  \in \mathbb{R}^+_{n,r},\textbf{R} \in \mathbb{R}^+_{m,r},\textbf{T} \in \mathbb{R}^+_{r,r}
 
     where:
 
@@ -215,3 +215,42 @@ def solve_balanced_FRLC(
             return nx.dot(nx.dot(Q_new, X_new), R_new.T)  # Shape (n, m)
 
         Q, R, T, X = Q_new, R_new, T_new, X_new
+
+
+if __name__ == "__main__":
+    import torch
+
+    grid_size = 4
+    torch.manual_seed(42)
+    x_vals = torch.linspace(0, 3, grid_size)
+    y_vals = torch.linspace(0, 3, grid_size)
+    X, Y = torch.meshgrid(x_vals, y_vals, indexing="ij")
+    source_points = torch.stack([X.ravel(), Y.ravel()], dim=-1)  # (16, 2)
+    a = torch.ones(len(source_points)) / len(source_points)  # Uniform distribution
+
+    # Generate Target Distribution (Gaussian Samples)
+    mean = torch.tensor([2.0, 2.0])
+    cov = torch.tensor([[1.0, 0.5], [0.5, 1.0]])
+    target_points = torch.distributions.MultivariateNormal(
+        mean, covariance_matrix=cov
+    ).sample((len(source_points),))  # (16, 2)
+    b = torch.ones(len(target_points)) / len(target_points)  # Uniform distribution
+
+    # Compute Cost Matrix (Squared Euclidean Distance)
+    C = torch.cdist(source_points, target_points, p=2) ** 2
+
+    # Solve OT problem (assuming you have PyTorch versions of these functions)
+    print(type(a.numpy()))
+    P = solve_balanced_FRLC(
+        a.to(torch.float64),
+        b.to(torch.float64),
+        C.to(torch.float64),
+        10,
+        tau=1e2,
+        gamma=1e2,
+        stopThr=1e-7,
+        numItermax=100,
+        log=True,
+    )
+    P = sinkhorn(a, b, C, reg=1)
+    print(torch.sum(P * C))