Automatic differentiation through inverse of coordinate transforms

CoordinateSystemInversionProblem.jl

using ForwardDiff: Dual, Tag, jacobian, value, tagtype, partials, gradient
using LinearAlgebra
using NLsolve
using Random
using Test

Random.seed!(0)

const jac = jacobian

valueise(x) = x
valueise(x::T) where {T<:Dual} = valueise(value(x))

function changevalue!(x::AbstractVector, y, s)
  for i in eachindex(x)
    x[i] == s[i] && (x[i] = y[i])
  end
  return x
end
function changevalue!(x::AbstractVector{<:Dual}, y, s)
  for i in eachindex(x)
    x[i] = changevalue!(x[i], y[i], s[i])
  end
  return x
end
function changevalue!(x::Dual{<:Tag}, y, s) where {T<:AbstractFloat, D}
  return Dual{tagtype(x)}(changevalue!(x.value, y, s), partials(x))
end
function changevalue!(x::Dual{<:Tag{<:Function,T}, T, D}, y::T, s::T
                     ) where {T<:AbstractFloat, D}
  return x.value == s ? Dual{tagtype(x)}(y, partials(x)) : x
end

changevalue(x, y) = changevalue!(deepcopy(x), y, valueise.(x))
"""
Given coordinate transform f(y) -> x, find y for given x.
"""
function inversion(f::F, x::AbstractVector{T}, initial_y=ones(size(x))
    ) where {F, T}
  f!(m, y) = (m .= f(y) .- valueise.(x); return nothing)
  y = NLsolve.nlsolve(f!, initial_y, autodiff=:forward, ftol=1e-12).zero
  # What is the proper way? This only works for a single AD pass.
  return changevalue(x, y)
end

"""Cylindrical to Cartesian"""
function f(rθ)
  r, θ = rθ
  return [r * cos(θ), r * sin(θ)]
end

"""Cartesian to Cylindrical"""
function f⁻¹(xy)
  x, y = xy
  return [sqrt(x^2 + y^2), atan(y, x)]
end

# I want to work in curvlinear coordinates R
# For the real problem it's only tractable to write down `f⁻¹`
# so I use an inversion function to go from `R` to `X` to get useful
# things out of `f⁻¹` in lieu of having an `f`
@testset "ForwardDiff" begin
  X = 5 * (rand(2) .- 0.5)
  R = f⁻¹(X)
  @assert f(R) ≈ X  # otherwise it'll never work
  @assert f⁻¹(X) ≈ R  # otherwise it'll never work
  @assert R[1] > 0
  @assert -π <= R[2] <= π

  @test inversion(f⁻¹, R) ≈ X

  @test inv(jac(f, R)) ≈ jac(f⁻¹, X) # yep

  @test inv(jac(f, R)) ≈ jac(r->f⁻¹(inversion(f⁻¹, r)), R)

  # now apply AD on top of the jacobian calculation
  # create arbitrary-ish function
  foo(R) = det(jac(f, R)) # get a single number so can do gradient
  bar(R) = det(inv(jac(r->f⁻¹(inversion(f⁻¹, r)), R)))
  expected = gradient(foo, R)
  @test expected ≈ [1, 0] # easy answer
  # gradient(bar, R) actually calculates d bar / dX so need dX/dR' d bar/dX
  @test jac(r->f⁻¹(inversion(f⁻¹, r)), R) * gradient(bar, R) ≈ expected # nearly
end