Add support for an invariant axis in convolutional sinkhorn

docs/src/time.md

@@ -6,6 +6,12 @@ plotly()
 For non-stationary signals, it is important to consider how the spectrum changes with time. This package has some, so far, basic support for time-frequency representations of non-stationary signals.
 ## Overview
 
+There are two main approaches for time-frequency representations supported
+1. Model based
+2. Spectrogram based
+
+### Model based
+
 We define a custom fit method for fitting time varying spectra, [`TimeWindow`](@ref). It takes as arguments an inner fitmethod, the number of points that form a time window, and the number of points that overlap between two consecutive time windows:
 ```@repl time
 fitmethod = TimeWindow(LS(na=2), 1000, 500)
@@ -66,7 +72,7 @@ evaluate(dist, m, m2)
 ```
 
 
-### Chirp example
+#### Chirp example
 Here we consider the estimation of the distance between two signals containing chirps, where the onset of the chirp differs. We start by creating some signals with different chirp onsets:
 ```@example time
 fs = 100000
@@ -139,6 +145,19 @@ The results are shown below. The figure indicates the cost `log10(c)` using the
 
 In this example, we chose the weight function `simplex_residueweight`, which ensures that each time step has the same amount of spectral mass. Individual poles will still have different masses within each timestep, as determined by the pole's residue.
 
+### Spectrogram based
+The distance [`ConvOptimalTransportDistance`](@ref) can operate on 2D measures on a grid, for example spectrograms. This distance supports calculating barycenters and barycentric coordinates. See [Barycenters between spectrograms](@ref) and [Barycentric coordiantes](@ref).
+
+
+This distance has a special option
+- `invariant_axis::Int = 0`
+
+If this is set to 1 or 2, the distance will be approximately invariant to translations along the invariant axis. As an example, to be invariant to a spectrogram being shifted slightly in time, set `invariant_axis = 2`. If this option is used, the distance is not differentiable and this option will be ignored when calculating barycenters etc. 
+
+```@docs
+ConvOptimalTransportDistance
+```
+
 ## Dynamic Time Warping
 This package interfaces with [DynamicAxisWarping.jl](https://github.com/baggepinnen/DynamicAxisWarping.jl) and provides optimized methods for [`DynamicAxisWarping.dtwnn`](@ref). Below is an example of how to search for a query pattern `Qm` in a much longer pattern `Ym`
 ```julia

src/adjoints.jl

@@ -233,10 +233,10 @@ ZygoteRules.@adjoint function sinkhorn_convolutional(w, A, B; β, kwargs...)
     V2, U2 = copy(V), copy(U) # This is absolutely required!
     function sinkhorn_convolutional_pullback(Δc)
         Δc *= β
-        mV = mean(V2) # If this normalization is not done, one has to normalize the gradients later instead, otherwise most of the gradient will point "into the constraints"
-        mU = mean(U2)
-        @avx @. V2 = (V2 - mV)*Δc
-        @avx @. U2 = (U2 - mU)*Δc
+        # mV = mean(V2) # If this normalization is not done, one has to normalize the gradients later instead, otherwise most of the gradient will point "into the constraints"
+        # mU = mean(U2) # NOTE: the normalization is now done directly in sinkhorn_convolutional
+        @avx @. V2 = (V2)*Δc
+        @avx @. U2 = (U2)*Δc
 
         return (nothing, V2, U2)
     end

src/sinkhorn.jl

@@ -505,7 +505,10 @@ function sinkhorn_convolutional(
     end
     @avx @. V = log(V + ϵ) + alpha
     @avx @. U = log(U + ϵ) + beta
-    β*(dot(A, V) + dot(B, U)), V, U
+    c = β*(dot(A, V) + dot(B, U))
+    V .-= mean(V)
+    U .-= mean(U)
+    c, V, U
 
 end
 
@@ -626,11 +629,13 @@ It's important to tune the two parameters below, see the docstring for [`sinkhor
 
 - `β = 0.001`
 - `dynamic_floor = -10.0`
+- `invariant_axis::Int = 0` If this is set to 1 or 2, the distance will be approximately invariant to translations along the invariant axis. As an example, to be invariant to a spectrogram being shifted slightly in time, set `invariant_axis = 2`.
 """
 ConvOptimalTransportDistance
 Base.@kwdef mutable struct ConvOptimalTransportDistance{T} <: AbstractDistance
     β::T = 0.001
     dynamic_floor::T = NaN
+    invariant_axis::Int = 0
     workspace::Union{Nothing, SCWorkspace{T}} = nothing
 end
 
@@ -644,7 +649,19 @@ function evaluate(d::ConvOptimalTransportDistance, A::AbstractMatrix, B::Abstrac
     if d.workspace === nothing
         d.workspace = SCWorkspace(A,B,d.β)
     end
-    sinkhorn_convolutional(d.workspace::SCWorkspace{T}, A, B; β = d.β, kwargs...)[1]
+    c,V,U = sinkhorn_convolutional(d.workspace::SCWorkspace{T}, A, B; β = d.β, kwargs...)
+    if d.invariant_axis != 0
+        isderiving() && error("Can not differentiate a distance with an invariant axis (I haven't bothered figure out the adjoint).")
+        ia = d.invariant_axis
+        sum_axis = ia == 1 ? 2 : 1
+        v,u = vec(sum(A.*V, dims=sum_axis)), vec(sum(B.*U, dims=sum_axis))
+        v ./= sum(v)
+        u ./= sum(u)
+        Γ = discrete_grid_transportplan(v, u)
+        invariant_cost = sum(Γ[i,j]*abs2((i-j)/2) for i = 1:size(Γ,1), j = 1:size(Γ,2))
+        return c - invariant_cost / size(Γ,1)
+    end
+    return c
 end
 
 """

test/test_convolutional.jl

@@ -44,6 +44,56 @@ end
 
 
 
+
+@testset "Invariant axis" begin
+    @info "Testing Invariant axis"
+    w,h = 5,5
+    for w = 5:6, h = 5:6
+        β = 0.05
+        D = [abs2((i-j)/2) for i = 1:h, j = 1:h] # NOTE: the 1/2 here is important
+        B = [zeros(h,w) for _ in 1:2]
+
+        B[1][1,1] = 1
+        B[1][1,5] = 1
+        B[2][2,1] = 1
+        B[2][2,5] = 1
+
+        B = s1.(B)
+        c,V,U = sinkhorn_convolutional(B[1], B[2], β=β, τ=1e20)
+        c1,V1,U1 = sinkhorn_convolutional(B[1], B[1], β=β, τ=1e20)
+        c2,V2,U2 = sinkhorn_convolutional(B[2], B[2], β=β, τ=1e20)
+        c3 = sqrt((c - 0.5(c1+c2))) # The size of the summed over dimension needs to multiply here
+        v,u = sum(B[1].*V, dims=2)[:], sum(B[2].*U, dims=2)[:] # Integrate over the dimension to be sensitive to (because we need to be invariant to this dimension in this step)
+        Γ = discrete_grid_transportplan(s1(v), s1(u))    # Solve 1D OT between integrated measures
+        invariant_cost = sqrt(dot(Γ, D) / size(B[1],1))
+
+        @test c3 ≈ invariant_cost rtol=0.15
+
+        if isinteractive()
+            Plots.heatmap(Γ, layout=5, title="Transport plan")
+            Plots.heatmap!(V, sp=2, title="V")
+            Plots.heatmap!(U, sp=3, title="U")
+            Plots.plot!(v, sp=4, lab=invariant_cost)
+            Plots.plot!(u, sp=5, lab=c3)
+        end
+
+
+        invariant_cost = dot(Γ, D) / size(B[1],1)
+        di = ConvOptimalTransportDistance(β=β, invariant_axis=1)
+        d = ConvOptimalTransportDistance(β=β, invariant_axis=0)
+        cdist = d(B[1], B[2]) - 0.5*(d(B[1], B[1]) + d(B[2], B[2]))
+        @test cdist ≈ c3^2
+        ci = (di(B[1], B[2]) - 0.5*(di(B[1], B[1]) + di(B[2], B[2])))
+
+        @test ci ≈ cdist-invariant_cost rtol=0.15
+
+        @test di(B[1], B[2]) ≈ d(B[1], B[2])-invariant_cost rtol=0.15
+
+    end
+end
+
+
+
 @testset "Misc convolutional" begin
     @info "Testing Misc convolutional"