Current section
Files
Jump to
Current section
Files
src/viva_math@transport.erl
-module(viva_math@transport).
-compile([no_auto_import, nowarn_unused_vars, nowarn_unused_function, nowarn_nomatch, inline]).
-define(FILEPATH, "src/viva_math/transport.gleam").
-export([wasserstein_1_empirical/2, wasserstein_2_empirical/2, wasserstein_2_gaussian/2, wasserstein_pad/2, wasserstein_2_multivariate/4]).
-if(?OTP_RELEASE >= 27).
-define(MODULEDOC(Str), -moduledoc(Str)).
-define(DOC(Str), -doc(Str)).
-else.
-define(MODULEDOC(Str), -compile([])).
-define(DOC(Str), -compile([])).
-endif.
?MODULEDOC(
" Optimal transport distances for empirical affective distributions.\n"
"\n"
" 1D empirical Wasserstein distances and PAD componentwise aggregation\n"
" following Villani (2008), *Optimal Transport*, and Peyré & Cuturi (2019),\n"
" *Computational Optimal Transport*.\n"
"\n"
" **Complexity**: O((n+m)·log(n+m)) dominated by sort. The post-sort\n"
" integral walks both quantile sequences in linear time without random\n"
" indexing — see `walk_quantile_*` helpers below.\n"
).
-file("src/viva_math/transport.gleam", 623).
?DOC(
" O(n+m) quantile-integral walk for W_1 over `{(i+1)/n} ∪ {(j+1)/m}`.\n"
" Slab contribution `|p_i − q_j| · Δu`. Twin of `walk_quantile_squared`.\n"
).
-spec walk_quantile_abs(
list(float()),
list(float()),
float(),
float(),
integer(),
integer(),
float(),
float()
) -> float().
walk_quantile_abs(P_sorted, Q_sorted, N, M, I, J, U_prev, Acc) ->
case {P_sorted, Q_sorted} of
{[P_val | P_tail], [Q_val | Q_tail]} ->
U_p = case N of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> erlang:float(I + 1) / Gleam@denominator
end,
U_q = case M of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator@1 -> erlang:float(J + 1) / Gleam@denominator@1
end,
U_next = gleam@float:min(U_p, U_q),
New_acc = Acc + (gleam@float:absolute_value(P_val - Q_val) * (U_next
- U_prev)),
Advance_p = U_p =< U_q,
Advance_q = U_q =< U_p,
Next_p = case Advance_p of
true ->
P_tail;
false ->
P_sorted
end,
Next_q = case Advance_q of
true ->
Q_tail;
false ->
Q_sorted
end,
Next_i = case Advance_p of
true ->
I + 1;
false ->
I
end,
Next_j = case Advance_q of
true ->
J + 1;
false ->
J
end,
walk_quantile_abs(
Next_p,
Next_q,
N,
M,
Next_i,
Next_j,
U_next,
New_acc
);
{_, _} ->
Acc
end.
-file("src/viva_math/transport.gleam", 560).
?DOC(" Equal-size W_1 numerator `Σ |p_i − q_i|` via single-pass cons traversal.\n").
-spec walk_pair_abs(list(float()), list(float()), float()) -> float().
walk_pair_abs(P, Q, Acc) ->
case {P, Q} of
{[A | Rest_p], [B | Rest_q]} ->
walk_pair_abs(
Rest_p,
Rest_q,
Acc + gleam@float:absolute_value(A - B)
);
{_, _} ->
Acc
end.
-file("src/viva_math/transport.gleam", 683).
-spec sort_samples(list(float())) -> list(float()).
sort_samples(Samples) ->
gleam@list:sort(Samples, fun(A, B) -> case {A < B, A > B} of
{true, _} ->
lt;
{_, true} ->
gt;
{_, _} ->
eq
end end).
-file("src/viva_math/transport.gleam", 27).
?DOC(
" Wasserstein-1 (Earth Mover) between 1D empirical samples.\n"
"\n"
" `W_1(P, Q) = ∫_0^1 |F_P⁻¹(u) − F_Q⁻¹(u)| du`.\n"
"\n"
" Equivalent to `∫_ℝ |F_P(x) − F_Q(x)| dx` (the `W_1` duality holds via\n"
" integration by parts; the absolute-value identity is symmetric). For\n"
" equal sample sizes reduces to `(1/n)·Σ |p_(i) − q_(i)|`. For unequal sizes\n"
" integrates `|p_i − q_j|·Δu` across the union of quantile breakpoints\n"
" `{i/n} ∪ {j/m}` in linear time after sorting.\n"
).
-spec wasserstein_1_empirical(list(float()), list(float())) -> {ok, float()} |
{error, nil}.
wasserstein_1_empirical(P, Q) ->
case {P, Q} of
{[], _} ->
{error, nil};
{_, []} ->
{error, nil};
{_, _} ->
P_sorted = sort_samples(P),
Q_sorted = sort_samples(Q),
P_len = erlang:length(P_sorted),
Q_len = erlang:length(Q_sorted),
case P_len =:= Q_len of
true ->
{ok, case erlang:float(P_len) of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> walk_pair_abs(
P_sorted,
Q_sorted,
+0.0
)
/ Gleam@denominator
end};
false ->
{ok,
walk_quantile_abs(
P_sorted,
Q_sorted,
erlang:float(P_len),
erlang:float(Q_len),
0,
0,
+0.0,
+0.0
)}
end
end.
-file("src/viva_math/transport.gleam", 571).
?DOC(
" O(n+m) quantile-integral walk for W_2² over `{(i+1)/n} ∪ {(j+1)/m}`.\n"
" Slab contribution `(p_i − q_j)² · Δu`. Specialised to avoid a `Bool`\n"
" branch in the hot loop (BEAM doesn't specialise on literal `True`/`False`).\n"
).
-spec walk_quantile_squared(
list(float()),
list(float()),
float(),
float(),
integer(),
integer(),
float(),
float()
) -> float().
walk_quantile_squared(P_sorted, Q_sorted, N, M, I, J, U_prev, Acc) ->
case {P_sorted, Q_sorted} of
{[P_val | P_tail], [Q_val | Q_tail]} ->
U_p = case N of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> erlang:float(I + 1) / Gleam@denominator
end,
U_q = case M of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator@1 -> erlang:float(J + 1) / Gleam@denominator@1
end,
U_next = gleam@float:min(U_p, U_q),
Delta = P_val - Q_val,
New_acc = Acc + ((Delta * Delta) * (U_next - U_prev)),
Advance_p = U_p =< U_q,
Advance_q = U_q =< U_p,
Next_p = case Advance_p of
true ->
P_tail;
false ->
P_sorted
end,
Next_q = case Advance_q of
true ->
Q_tail;
false ->
Q_sorted
end,
Next_i = case Advance_p of
true ->
I + 1;
false ->
I
end,
Next_j = case Advance_q of
true ->
J + 1;
false ->
J
end,
walk_quantile_squared(
Next_p,
Next_q,
N,
M,
Next_i,
Next_j,
U_next,
New_acc
);
{_, _} ->
Acc
end.
-file("src/viva_math/transport.gleam", 549).
?DOC(
" Equal-size W_2 numerator `Σ (p_i − q_i)²` via single-pass cons traversal,\n"
" no `list.zip` intermediate.\n"
).
-spec walk_pair_squared(list(float()), list(float()), float()) -> float().
walk_pair_squared(P, Q, Acc) ->
case {P, Q} of
{[A | Rest_p], [B | Rest_q]} ->
D = A - B,
walk_pair_squared(Rest_p, Rest_q, Acc + (D * D));
{_, _} ->
Acc
end.
-file("src/viva_math/transport.gleam", 73).
?DOC(
" Wasserstein-2 between 1D empirical samples.\n"
"\n"
" `W_2²(P, Q) = ∫_0^1 (F_P⁻¹(u) − F_Q⁻¹(u))² du`.\n"
"\n"
" For equal sample sizes the inverse-CDF integral reduces to the sorted\n"
" pairwise mean-square difference. For unequal sizes we integrate over the\n"
" union of quantile breakpoints `{i/n} ∪ {j/m}` in linear time: in each\n"
" slab `[u_prev, u_next]` both inverse CDFs are constant, contributing\n"
" `(p_sorted[i] − q_sorted[j])² · (u_next − u_prev)`.\n"
"\n"
" Returns `W_2`, not `W_2²`. References: Villani (2008); Peyré & Cuturi (2019).\n"
"\n"
" **Note**: a CDF-based path `∫(F_P−F_Q)² dx` is **incorrect** for unequal\n"
" sample sizes — the `W_1`/`W_2` duality via integration by parts only\n"
" holds for `p=1` (absolute value), not for the quadratic kernel.\n"
).
-spec wasserstein_2_empirical(list(float()), list(float())) -> {ok, float()} |
{error, nil}.
wasserstein_2_empirical(P, Q) ->
case {P, Q} of
{[], _} ->
{error, nil};
{_, []} ->
{error, nil};
{_, _} ->
P_sorted = sort_samples(P),
Q_sorted = sort_samples(Q),
P_len = erlang:length(P_sorted),
Q_len = erlang:length(Q_sorted),
case P_len =:= Q_len of
true ->
{ok, math:sqrt(case erlang:float(P_len) of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> walk_pair_squared(
P_sorted,
Q_sorted,
+0.0
)
/ Gleam@denominator
end)};
false ->
{ok,
math:sqrt(
walk_quantile_squared(
P_sorted,
Q_sorted,
erlang:float(P_len),
erlang:float(Q_len),
0,
0,
+0.0,
+0.0
)
)}
end
end.
-file("src/viva_math/transport.gleam", 116).
?DOC(
" Closed form W_2 for scalar Gaussians.\n"
"\n"
" `W_2(N(μ₁, σ₁²), N(μ₂, σ₂²)) = √((μ₁ − μ₂)² + (σ₁ − σ₂)²)`.\n"
"\n"
" `stddev` is normalised via `abs` before subtracting — a negative `stddev`\n"
" has no Gaussian meaning (`N(μ, σ²)` depends only on `σ²`), and silently\n"
" using the signed value would yield `W_2(N(0, 1), N(0, 1))` of `2.0` for\n"
" `Gaussian(0, -1) vs Gaussian(0, 1)`, which is wrong.\n"
).
-spec wasserstein_2_gaussian(
viva_math@distributions:gaussian(),
viva_math@distributions:gaussian()
) -> float().
wasserstein_2_gaussian(G1, G2) ->
Mean_delta = erlang:element(2, G1) - erlang:element(2, G2),
Stddev_delta = gleam@float:absolute_value(erlang:element(3, G1)) - gleam@float:absolute_value(
erlang:element(3, G2)
),
math:sqrt((Mean_delta * Mean_delta) + (Stddev_delta * Stddev_delta)).
-file("src/viva_math/transport.gleam", 665).
?DOC(" Single-pass PAD projection — avoids three `list.map` passes.\n").
-spec split_pad(
list(viva_math@vector:vec3()),
list(float()),
list(float()),
list(float())
) -> {list(float()), list(float()), list(float())}.
split_pad(Xs, P_acc, A_acc, D_acc) ->
case Xs of
[] ->
{lists:reverse(P_acc), lists:reverse(A_acc), lists:reverse(D_acc)};
[V | Rest] ->
split_pad(
Rest,
[viva_math@vector:pleasure(V) | P_acc],
[viva_math@vector:arousal(V) | A_acc],
[viva_math@vector:dominance(V) | D_acc]
)
end.
-file("src/viva_math/transport.gleam", 140).
?DOC(
" Component-wise (marginal) Wasserstein-2 over PAD dimensions.\n"
"\n"
" `D(P, Q) = √( W_2²(P_P, Q_P) + W_2²(P_A, Q_A) + W_2²(P_D, Q_D) )`.\n"
"\n"
" This is **not** the multivariate W_2 (which requires solving a full\n"
" optimal-transport assignment with cost `‖x − y‖²`). It is the Euclidean\n"
" norm of the per-axis marginal Wasserstein distances — equivalent to the\n"
" Sliced Wasserstein along the canonical PAD basis.\n"
"\n"
" Triangle inequality holds (Minkowski over the marginal vector), so this\n"
" is a **pseudo-metric** on PAD distributions: `D(P, Q) = 0` does **not**\n"
" imply `P = Q` as joints — two distributions with identical marginals but\n"
" different correlations are tied. Useful as a fast lower bound on the true\n"
" multivariate W_2; tight when marginals are product distributions.\n"
).
-spec wasserstein_pad(
list(viva_math@vector:vec3()),
list(viva_math@vector:vec3())
) -> {ok, float()} | {error, nil}.
wasserstein_pad(P, Q) ->
case {P, Q} of
{[], _} ->
{error, nil};
{_, []} ->
{error, nil};
{_, _} ->
{P_p, P_a, P_d} = split_pad(P, [], [], []),
{Q_p, Q_a, Q_d} = split_pad(Q, [], [], []),
case {wasserstein_2_empirical(P_p, Q_p),
wasserstein_2_empirical(P_a, Q_a),
wasserstein_2_empirical(P_d, Q_d)} of
{{ok, Pleasure}, {ok, Arousal}, {ok, Dominance}} ->
{ok,
math:sqrt(
((Pleasure * Pleasure) + (Arousal * Arousal)) + (Dominance
* Dominance)
)};
{_, _, _} ->
{error, nil}
end
end.
-file("src/viva_math/transport.gleam", 518).
-spec transport_cost_log_row(
list(float()),
float(),
list(float()),
float(),
float()
) -> float().
transport_cost_log_row(Costs, Alpha_i, Beta, Epsilon, Acc) ->
case {Costs, Beta} of
{[Cost | Cost_tail], [Beta_j | Beta_tail]} ->
Plan_mass = math:exp(case Epsilon of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> ((Alpha_i + Beta_j) - Cost) / Gleam@denominator
end),
transport_cost_log_row(
Cost_tail,
Alpha_i,
Beta_tail,
Epsilon,
Acc + (Plan_mass * Cost)
);
{_, _} ->
Acc
end.
-file("src/viva_math/transport.gleam", 504).
-spec transport_cost_log(
list(list(float())),
list(float()),
list(float()),
float()
) -> float().
transport_cost_log(Costs, Alpha, Beta, Epsilon) ->
case {Costs, Alpha} of
{[Cost_row | Cost_tail], [Alpha_i | Alpha_tail]} ->
transport_cost_log_row(Cost_row, Alpha_i, Beta, Epsilon, +0.0) + transport_cost_log(
Cost_tail,
Alpha_tail,
Beta,
Epsilon
);
{_, _} ->
+0.0
end.
-file("src/viva_math/transport.gleam", 365).
-spec drop_column_head(list(list(float())), list(list(float()))) -> list(list(float())).
drop_column_head(Rows, Acc) ->
case Rows of
[] ->
lists:reverse(Acc);
[[_ | Rest] | Tail] ->
drop_column_head(Tail, [Rest | Acc]);
[[] | Tail@1] ->
drop_column_head(Tail@1, [[] | Acc])
end.
-file("src/viva_math/transport.gleam", 484).
-spec plan_column_head_sum(
list(list(float())),
list(float()),
float(),
float(),
float()
) -> float().
plan_column_head_sum(Rows, Alpha, Beta_j, Epsilon, Acc) ->
case {Rows, Alpha} of
{[[Cost | _] | Rest_rows], [Alpha_i | Rest_alpha]} ->
plan_column_head_sum(
Rest_rows,
Rest_alpha,
Beta_j,
Epsilon,
Acc + math:exp(case Epsilon of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> ((Alpha_i + Beta_j) - Cost) / Gleam@denominator
end)
);
{_, _} ->
Acc
end.
-file("src/viva_math/transport.gleam", 436).
-spec max_column_marginal_violation_walk(
list(float()),
list(list(float())),
list(float()),
list(float()),
float(),
float(),
float()
) -> float().
max_column_marginal_violation_walk(
First_row,
Costs,
Alpha,
Beta,
Epsilon,
Mass,
Acc
) ->
case {First_row, Beta} of
{[], _} ->
Acc;
{[_ | Rest], [Beta_j | Beta_tail]} ->
Column_sum = plan_column_head_sum(
Costs,
Alpha,
Beta_j,
Epsilon,
+0.0
),
Tails = drop_column_head(Costs, []),
max_column_marginal_violation_walk(
Rest,
Tails,
Alpha,
Beta_tail,
Epsilon,
Mass,
gleam@float:max(
Acc,
gleam@float:absolute_value(Column_sum - Mass)
)
);
{_, _} ->
Acc
end.
-file("src/viva_math/transport.gleam", 414).
-spec max_column_marginal_violation(
list(list(float())),
list(float()),
list(float()),
float(),
float()
) -> float().
max_column_marginal_violation(Costs, Alpha, Beta, Epsilon, Mass) ->
case Costs of
[] ->
+0.0;
[First | _] ->
max_column_marginal_violation_walk(
First,
Costs,
Alpha,
Beta,
Epsilon,
Mass,
+0.0
)
end.
-file("src/viva_math/transport.gleam", 464).
-spec plan_row_sum(list(float()), float(), list(float()), float(), float()) -> float().
plan_row_sum(Costs, Alpha_i, Beta, Epsilon, Acc) ->
case {Costs, Beta} of
{[Cost | Cost_tail], [Beta_j | Beta_tail]} ->
plan_row_sum(
Cost_tail,
Alpha_i,
Beta_tail,
Epsilon,
Acc + math:exp(case Epsilon of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> ((Alpha_i + Beta_j) - Cost) / Gleam@denominator
end)
);
{_, _} ->
Acc
end.
-file("src/viva_math/transport.gleam", 390).
-spec max_row_marginal_violation(
list(list(float())),
list(float()),
list(float()),
float(),
float(),
float()
) -> float().
max_row_marginal_violation(Costs, Alpha, Beta, Epsilon, Mass, Acc) ->
case {Costs, Alpha} of
{[Row | Rest], [Alpha_i | Alpha_tail]} ->
Row_sum = plan_row_sum(Row, Alpha_i, Beta, Epsilon, +0.0),
max_row_marginal_violation(
Rest,
Alpha_tail,
Beta,
Epsilon,
Mass,
gleam@float:max(Acc, gleam@float:absolute_value(Row_sum - Mass))
);
{_, _} ->
Acc
end.
-file("src/viva_math/transport.gleam", 376).
-spec marginal_violation(
list(list(float())),
list(float()),
list(float()),
float(),
float(),
float()
) -> float().
marginal_violation(Costs, Alpha, Beta, Epsilon, A, B) ->
gleam@float:max(
max_row_marginal_violation(Costs, Alpha, Beta, Epsilon, A, +0.0),
max_column_marginal_violation(Costs, Alpha, Beta, Epsilon, B)
).
-file("src/viva_math/transport.gleam", 349).
-spec logsumexp_column_head(
list(list(float())),
list(float()),
float(),
list(float())
) -> float().
logsumexp_column_head(Rows, Alpha, Epsilon, Acc) ->
case {Rows, Alpha} of
{[[Cost | _] | Rest_rows], [Alpha_i | Rest_alpha]} ->
logsumexp_column_head(
Rest_rows,
Rest_alpha,
Epsilon,
[case Epsilon of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> ((+0.0 - Cost) + Alpha_i) / Gleam@denominator
end | Acc]
);
{_, _} ->
viva_math@scalar:logsumexp(lists:reverse(Acc))
end.
-file("src/viva_math/transport.gleam", 311).
-spec update_beta_columns(
list(float()),
list(list(float())),
list(float()),
float(),
float(),
list(float())
) -> list(float()).
update_beta_columns(First_row, Costs, Alpha, Epsilon, Log_b, Acc) ->
case First_row of
[] ->
lists:reverse(Acc);
[_ | Rest] ->
Beta_j = (Epsilon * Log_b) - (Epsilon * logsumexp_column_head(
Costs,
Alpha,
Epsilon,
[]
)),
Tails = drop_column_head(Costs, []),
update_beta_columns(
Rest,
Tails,
Alpha,
Epsilon,
Log_b,
[Beta_j | Acc]
)
end.
-file("src/viva_math/transport.gleam", 299).
-spec update_beta(list(list(float())), list(float()), float(), float()) -> list(float()).
update_beta(Costs, Alpha, Epsilon, Log_b) ->
case Costs of
[] ->
[];
[First | _] ->
update_beta_columns(First, Costs, Alpha, Epsilon, Log_b, [])
end.
-file("src/viva_math/transport.gleam", 333).
-spec logsumexp_row(list(float()), list(float()), float(), list(float())) -> float().
logsumexp_row(Costs, Beta, Epsilon, Acc) ->
case {Costs, Beta} of
{[Cost | Cost_tail], [Beta_j | Beta_tail]} ->
logsumexp_row(Cost_tail, Beta_tail, Epsilon, [case Epsilon of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> ((+0.0 - Cost) + Beta_j) / Gleam@denominator
end | Acc]);
{_, _} ->
viva_math@scalar:logsumexp(lists:reverse(Acc))
end.
-file("src/viva_math/transport.gleam", 282).
-spec update_alpha(
list(list(float())),
list(float()),
float(),
float(),
list(float())
) -> list(float()).
update_alpha(Costs, Beta, Epsilon, Log_a, Acc) ->
case Costs of
[] ->
lists:reverse(Acc);
[Row | Rest] ->
Alpha_i = (Epsilon * Log_a) - (Epsilon * logsumexp_row(
Row,
Beta,
Epsilon,
[]
)),
update_alpha(Rest, Beta, Epsilon, Log_a, [Alpha_i | Acc])
end.
-file("src/viva_math/transport.gleam", 243).
-spec sinkhorn_log_iter(
list(list(float())),
float(),
integer(),
float(),
float(),
float(),
float(),
float(),
list(float()),
list(float())
) -> {list(float()), list(float())}.
sinkhorn_log_iter(
Costs,
Epsilon,
Remaining,
A,
B,
Log_a,
Log_b,
Tol,
Alpha,
Beta
) ->
case Remaining =< 0 of
true ->
{Alpha, Beta};
false ->
Next_alpha = update_alpha(Costs, Beta, Epsilon, Log_a, []),
Next_beta = update_beta(Costs, Next_alpha, Epsilon, Log_b),
case marginal_violation(Costs, Next_alpha, Next_beta, Epsilon, A, B)
=< Tol of
true ->
{Next_alpha, Next_beta};
false ->
sinkhorn_log_iter(
Costs,
Epsilon,
Remaining - 1,
A,
B,
Log_a,
Log_b,
Tol,
Next_alpha,
Next_beta
)
end
end.
-file("src/viva_math/transport.gleam", 540).
-spec fill(integer(), float(), list(float())) -> list(float()).
fill(Count, Value, Acc) ->
case Count =< 0 of
true ->
Acc;
false ->
fill(Count - 1, Value, [Value | Acc])
end.
-file("src/viva_math/transport.gleam", 216).
-spec sinkhorn_log_domain(
list(list(float())),
float(),
integer(),
float(),
float(),
float()
) -> {list(float()), list(float())}.
sinkhorn_log_domain(Costs, Epsilon, Max_iter, A, B, Tol) ->
Alpha0 = fill(erlang:length(Costs), +0.0, []),
Beta0 = case Costs of
[] ->
[];
[First | _] ->
fill(erlang:length(First), +0.0, [])
end,
sinkhorn_log_iter(
Costs,
Epsilon,
Max_iter,
A,
B,
math:log(A),
math:log(B),
Tol,
Alpha0,
Beta0
).
-file("src/viva_math/transport.gleam", 209).
-spec cost_matrix(list(viva_math@vector:vec3()), list(viva_math@vector:vec3())) -> list(list(float())).
cost_matrix(Xs, Ys) ->
gleam@list:map(
Xs,
fun(X) ->
gleam@list:map(
Ys,
fun(Y) -> viva_math@vector:distance_squared(X, Y) end
)
end
).
-file("src/viva_math/transport.gleam", 181).
?DOC(
" Multivariate Wasserstein-2 distance via Sinkhorn-Knopp entropic\n"
" regularization (Cuturi 2013).\n"
"\n"
" Solves the entropic OT problem\n"
" `min_π ⟨π, C⟩ + ε·H(π)` subject to `π·1 = a`, `πᵀ·1 = b`,\n"
" where `C[i,j] = ‖x_i − y_j‖²` and `π` is the transport plan.\n"
"\n"
" Returns `√(⟨π, C⟩)`, the W₂ distance induced by the regularized plan,\n"
" dropping the entropic term itself.\n"
"\n"
" - `xs`, `ys`: empirical PAD samples with uniform weights.\n"
" - `epsilon`: regularization strength, typically `0.01` to `1.0`.\n"
" - `max_iter`: upper bound on Sinkhorn iterations, typically `100` to\n"
" `1000`; convergence may stop earlier once marginal violation is below\n"
" the internal tolerance.\n"
" - Returns `Error(Nil)` if either list is empty.\n"
).
-spec wasserstein_2_multivariate(
list(viva_math@vector:vec3()),
list(viva_math@vector:vec3()),
float(),
integer()
) -> {ok, float()} | {error, nil}.
wasserstein_2_multivariate(Xs, Ys, Epsilon, Max_iter) ->
case {Xs, Ys} of
{[], _} ->
{error, nil};
{_, []} ->
{error, nil};
{_, _} ->
N = erlang:float(erlang:length(Xs)),
M = erlang:float(erlang:length(Ys)),
A = case N of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator -> 1.0 / Gleam@denominator
end,
B = case M of
+0.0 -> +0.0;
-0.0 -> -0.0;
Gleam@denominator@1 -> 1.0 / Gleam@denominator@1
end,
Costs = cost_matrix(Xs, Ys),
Epsilon_safe = gleam@float:max(Epsilon, 1.0e-12),
{Alpha, Beta} = sinkhorn_log_domain(
Costs,
Epsilon_safe,
Max_iter,
A,
B,
1.0e-9
),
Cost = transport_cost_log(Costs, Alpha, Beta, Epsilon_safe),
{ok, math:sqrt(gleam@float:max(Cost, +0.0))}
end.