Skip to content

Commit

Permalink
Add support for an invariant axis in convolutional sinkhorn
Browse files Browse the repository at this point in the history
  • Loading branch information
baggepinnen committed Jun 23, 2020
1 parent 5374e5e commit eac7d2d
Show file tree
Hide file tree
Showing 4 changed files with 93 additions and 7 deletions.
21 changes: 20 additions & 1 deletion docs/src/time.md
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
8 changes: 4 additions & 4 deletions src/adjoints.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
21 changes: 19 additions & 2 deletions src/sinkhorn.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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

Expand All @@ -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

"""
Expand Down
50 changes: 50 additions & 0 deletions test/test_convolutional.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down

0 comments on commit eac7d2d

Please sign in to comment.