added ot_reg_cost

Author: zsteve

src/OptimalTransport.jl

@@ -14,7 +14,7 @@ export sinkhorn, sinkhorn2
 export emd, emd2
 export sinkhorn_stabilized, sinkhorn_stabilized_epsscaling, sinkhorn_barycenter
 export sinkhorn_unbalanced, sinkhorn_unbalanced2
-export ot_reg_plan
+export ot_reg_plan, ot_reg_cost
 
 const MOI = MathOptInterface
 
@@ -567,6 +567,25 @@ function ot_reg_plan(mu, nu, C, eps; reg_func="L2", method="lorenz", kwargs...)
     end
 end
 
+"""
+    ot_reg_cost(mu, nu, C, eps; reg_func = "L2", method = "lorenz", kwargs...)
+
+Compute the optimal transport cost between `mu` and `nu` for optimal transport with a 
+general choice of regulariser `math Ω(γ)`. 
+
+See also: [`ot_reg_plan`](@ref)
+
+"""
+function ot_reg_cost(mu, nu, C, eps; reg_func="L2", method="lorenz", kwargs...)
+    γ = if (reg_func == "L2") && (method == "lorenz")
+        quadreg(mu, nu, C, eps; kwargs...)
+    else
+        @warn "Unimplemented"
+        nothing 
+    end
+    return dot(γ, C)
+end
+
 
 """
     quadreg(mu, nu, C, ϵ; θ = 0.1, tol = 1e-5,maxiter = 50,κ = 0.5,δ = 1e-5)

test/runtests.jl

@@ -201,6 +201,9 @@ end
         γ_pot = sparse(POT.smooth_ot_dual(μ, ν, C, eps; max_iter=5000))
         # need to use a larger tolerance here because of a quirk with the POT solver 
         @test norm(γ - γ_pot, Inf) < 1e-4
+        c = ot_reg_cost(μ, ν, C, eps; reg_func="L2", method="lorenz")
+        c_pot = dot(γ_pot, C)
+        @test c ≈ c_pot atol = 1e-4
     end
 end