write up gradients for Gaussian PSF

Author: mileslucas

src/PSFModels.jl

@@ -111,9 +111,12 @@ plot(model, axes(other)) # use axes from other array
 """
 module PSFModels
 
+using ChainRulesCore
+import ChainRulesCore: frule, rrule
 using CoordinateTransformations
 using Distances
 using KeywordCalls
+using LinearAlgebra
 using SpecialFunctions
 using StaticArrays
 

src/gaussian.jl

@@ -61,3 +61,51 @@ function (g::Gaussian{T,<:Union{Tuple,AbstractVector}})(point::AbstractVector) w
     val = g.amp * exp(GAUSS_PRE * Δ)
     return convert(T, val)
 end
+
+## gradients
+
+# isotropic
+function fgrad(g::Gaussian, point::AbstractVector)
+    f = g(point)
+
+    xdiff = first(point) - first(g.pos)
+    ydiff = last(point) - last(g.pos)
+    dfdpos = -2 * GAUSS_PRE * f / g.fwhm^2 .* SA[xdiff, ydiff]
+    dfdfwhm = -2 * GAUSS_PRE * f * (xdiff^2 + ydiff^2) / g.fwhm^3
+    dfdamp = f / g.amp
+    return f, dfdpos, dfdfwhm, dfdamp
+end
+
+# short printing
+Base.show(io::IO, g::Gaussian{T}) where {T} = print(io, "Gaussian{$T}(pos=$(g.pos), fwhm=$(g.fwhm), amp=$(g.amp))")
+
+# diagonal
+function fgrad(g::Gaussian{T,<:Union{Tuple,AbstractVector}}, point::AbstractVector) where T
+    f = g(point)
+
+    xdiff = first(point) - first(g.pos)
+    ydiff = last(point) - last(g.pos)
+    dfdpos = -2 * GAUSS_PRE * f .* SA[xdiff / first(g.fwhm)^2, ydiff / last(g.fwhm)^2]
+    dfdfwhm = -2 * GAUSS_PRE * f .* SA[xdiff^2 / first(g.fwhm)^3, ydiff^2 / last(g.fwhm)^3]
+    dfda = f / g.amp
+    return f, dfdpos, dfdfwhm, dfda
+end
+
+function frule((Δpsf, Δp), g::Gaussian, point::AbstractVector)
+    f, dfdpos, dfdfwhm, dfda = fgrad(g, point)
+    Δf = dot(dfdpos, Δpsf.pos) + dot(dfdfwhm, Δpsf.fwhm) + dfda * Δpsf.amp
+    Δf -= dot(dfdpos, Δp)
+    return f, Δf
+end
+
+function rrule(g::G, point::AbstractVector) where {G<:Gaussian}
+    f, dfdpos, dfdfwhm, dfda = fgrad(g, point)
+    function Gaussian_pullback(Δf)
+        ∂pos = dfdpos .* Δf
+        ∂fwhm = dfdfwhm .* Δf
+        ∂g = Tangent{G}(pos=∂pos, fwhm=∂fwhm, amp=dfda * Δf, indices=ZeroTangent())
+        ∂pos = dfdpos .* -Δf
+        return ∂g, ∂pos
+    end
+    return f, Gaussian_pullback
+end