discrete1d_benchmarks.jl

using BenchmarkTools
using Distances
using Distributions
using StatsBase

using LinearAlgebra
using SparseArrays

function f(
    c, μ::DiscreteNonParametric, ν::DiscreteNonParametric, plan::SparseMatrixCSC
)
    # extract non-zero flows
    I, J, W = findnz(plan)

    # compute the cost
    μsupport = support(μ)
    νsupport = support(ν)
    cost = sum(w * c(μsupport[i], νsupport[j]) for (i, j, w) in zip(I, J, W))

    return cost
end

function g(c, μ::DiscreteNonParametric, ν::DiscreteNonParametric, plan)
    return dot(plan, StatsBase.pairwise(c, support(μ), support(ν)))
end

function benchmark_f_and_g(c, nμ, nν)
    μprobs = rand(nμ)
    μprobs ./= sum(μprobs)
    μsupport = randn(nμ)
    μ = DiscreteNonParametric(μsupport, μprobs)

    νprobs = rand(nν)
    νprobs ./= sum(νprobs)
    νsupport = randn(nν)
    ν = DiscreteNonParametric(νsupport, νprobs)

    ntotal = nμ * nν
    nnonzero = nμ + nν
    nonzero_indices = sample(1:ntotal, nnonzero)
    I = mod1.(nonzero_indices, nμ)
    J = fld1.(nonzero_indices, nμ)
    W = randn(nnonzero)
    plan = sparse(I, J, W, nμ, nν)

    bf = @benchmark f($c, $μ, $ν, $plan)
    bg = @benchmark g($c, $μ, $ν, $plan)

    println("Benchmark f:")
    display(bf)
    println("Benchmark g:")
    display(bg)

    return nothing
end

#############
# SqEuclidean
#############

benchmark_f_and_g(SqEuclidean(), 10, 10)
#=
Benchmark f:
BenchmarkTools.Trial:
  memory estimate:  624 bytes
  allocs estimate:  3
  --------------
  minimum time:     162.058 ns (0.00% GC)
  median time:      185.211 ns (0.00% GC)
  mean time:        237.313 ns (7.71% GC)
  maximum time:     3.936 μs (89.04% GC)
  --------------
  samples:          10000
  evals/sample:     796
Benchmark g:
BenchmarkTools.Trial:
  memory estimate:  896 bytes
  allocs estimate:  1
  --------------
  minimum time:     133.809 ns (0.00% GC)
  median time:      165.768 ns (0.00% GC)
  mean time:        227.333 ns (9.31% GC)
  maximum time:     3.840 μs (87.38% GC)
  --------------
  samples:          10000
  evals/sample:     873
=#

benchmark_f_and_g(SqEuclidean(), 15, 15)
#=
Benchmark f:
BenchmarkTools.Trial:
  memory estimate:  912 bytes
  allocs estimate:  3
  --------------
  minimum time:     199.355 ns (0.00% GC)
  median time:      231.807 ns (0.00% GC)
  mean time:        305.065 ns (7.42% GC)
  maximum time:     5.929 μs (87.69% GC)
  --------------
  samples:          10000
  evals/sample:     595
Benchmark g:
BenchmarkTools.Trial:
  memory estimate:  1.98 KiB
  allocs estimate:  1
  --------------
  minimum time:     256.086 ns (0.00% GC)
  median time:      349.171 ns (0.00% GC)
  mean time:        482.246 ns (9.48% GC)
  maximum time:     8.767 μs (83.37% GC)
  --------------
  samples:          10000
  evals/sample:     337
=#

benchmark_f_and_g(SqEuclidean(), 20, 20)
#=
Benchmark f:
BenchmarkTools.Trial:
  memory estimate:  1.17 KiB
  allocs estimate:  3
  --------------
  minimum time:     241.719 ns (0.00% GC)
  median time:      263.763 ns (0.00% GC)
  mean time:        361.468 ns (7.92% GC)
  maximum time:     7.070 μs (88.14% GC)
  --------------
  samples:          10000
  evals/sample:     417
Benchmark g:
BenchmarkTools.Trial:
  memory estimate:  3.25 KiB
  allocs estimate:  1
  --------------
  minimum time:     447.765 ns (0.00% GC)
  median time:      614.967 ns (0.00% GC)
  mean time:        968.838 ns (18.51% GC)
  maximum time:     28.940 μs (92.76% GC)
  --------------
  samples:          10000
  evals/sample:     196
=#

benchmark_f_and_g(SqEuclidean(), 200, 200)
#=
Benchmark f:
BenchmarkTools.Trial:
  memory estimate:  9.75 KiB
  allocs estimate:  3
  --------------
  minimum time:     2.021 μs (0.00% GC)
  median time:      2.465 μs (0.00% GC)
  mean time:        3.604 μs (14.80% GC)
  maximum time:     539.172 μs (92.39% GC)
  --------------
  samples:          10000
  evals/sample:     9
Benchmark g:
BenchmarkTools.Trial:
  memory estimate:  312.58 KiB
  allocs estimate:  2
  --------------
  minimum time:     32.552 μs (0.00% GC)
  median time:      65.388 μs (0.00% GC)
  mean time:        101.859 μs (10.65% GC)
  maximum time:     4.258 ms (92.83% GC)
  --------------
  samples:          10000
  evals/sample:     1
=#

###########
# Euclidean
###########

benchmark_f_and_g(Euclidean(), 10, 10)
#=
Benchmark f:
BenchmarkTools.Trial:
  memory estimate:  672 bytes
  allocs estimate:  3
  --------------
  minimum time:     174.918 ns (0.00% GC)
  median time:      193.881 ns (0.00% GC)
  mean time:        245.403 ns (6.50% GC)
  maximum time:     4.798 μs (0.00% GC)
  --------------
  samples:          10000
  evals/sample:     741
Benchmark g:
BenchmarkTools.Trial:
  memory estimate:  896 bytes
  allocs estimate:  1
  --------------
  minimum time:     193.638 ns (0.00% GC)
  median time:      218.862 ns (0.00% GC)
  mean time:        283.115 ns (7.99% GC)
  maximum time:     5.541 μs (89.91% GC)
  --------------
  samples:          10000
  evals/sample:     600
=#

benchmark_f_and_g(Euclidean(), 15, 15)
#=
Benchmark f:
BenchmarkTools.Trial:
  memory estimate:  912 bytes
  allocs estimate:  3
  --------------
  minimum time:     206.579 ns (0.00% GC)
  median time:      286.637 ns (0.00% GC)
  mean time:        376.781 ns (8.03% GC)
  maximum time:     7.974 μs (88.27% GC)
  --------------
  samples:          10000
  evals/sample:     570
Benchmark g:
BenchmarkTools.Trial:
  memory estimate:  1.98 KiB
  allocs estimate:  1
  --------------
  minimum time:     389.847 ns (0.00% GC)
  median time:      428.923 ns (0.00% GC)
  mean time:        573.644 ns (8.20% GC)
  maximum time:     12.382 μs (88.28% GC)
  --------------
  samples:          10000
  evals/sample:     202
=#

benchmark_f_and_g(Euclidean(), 20, 20)
#=
Benchmark f:
BenchmarkTools.Trial:
  memory estimate:  1.17 KiB
  allocs estimate:  3
  --------------
  minimum time:     261.805 ns (0.00% GC)
  median time:      295.536 ns (0.00% GC)
  mean time:        431.306 ns (7.95% GC)
  maximum time:     14.292 μs (92.01% GC)
  --------------
  samples:          10000
  evals/sample:     349
Benchmark g:
BenchmarkTools.Trial:
  memory estimate:  3.25 KiB
  allocs estimate:  1
  --------------
  minimum time:     762.697 ns (0.00% GC)
  median time:      892.545 ns (0.00% GC)
  mean time:        1.265 μs (14.18% GC)
  maximum time:     53.744 μs (92.13% GC)
  --------------
  samples:          10000
  evals/sample:     99
=#

benchmark_f_and_g(Euclidean(), 200, 200)
#=
Benchmark f:
BenchmarkTools.Trial:
  memory estimate:  9.75 KiB
  allocs estimate:  3
  --------------
  minimum time:     2.121 μs (0.00% GC)
  median time:      2.742 μs (0.00% GC)
  mean time:        4.139 μs (14.51% GC)
  maximum time:     601.941 μs (92.60% GC)
  --------------
  samples:          10000
  evals/sample:     9
Benchmark g:
BenchmarkTools.Trial:
  memory estimate:  312.58 KiB
  allocs estimate:  2
  --------------
  minimum time:     38.262 μs (0.00% GC)
  median time:      53.612 μs (0.00% GC)
  mean time:        74.147 μs (9.54% GC)
  maximum time:     2.482 ms (94.25% GC)
  --------------
  samples:          10000
  evals/sample:     1
=#