add chainrules-extension for basic nfft-functions

AbstractNFFTs/Project.toml

@@ -7,5 +7,11 @@ version = "0.8.2"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+
+[extensions]
+AbstractNFFTsChainRulesCoreExt = "ChainRulesCore"
+
 [compat]
 julia = "1.6"

AbstractNFFTs/ext/AbstractNFFTsChainRulesCoreExt.jl

@@ -0,0 +1,73 @@
+module AbstractNFFTsChainRulesCoreExt
+
+using AbstractNFFTs
+import ChainRulesCore
+
+###############
+# mul-interface
+###############
+function ChainRulesCore.frule((_, Δx, _), ::typeof(*), plan::AbstractFTPlan, x::AbstractArray)
+  y = plan*x
+  Δy = plan*Δx
+  return y, Δy
+end
+function ChainRulesCore.rrule(::typeof(*), plan::AbstractFTPlan, x::AbstractArray)
+  y = plan*x
+  project_x = ChainRulesCore.ProjectTo(x)
+  function mul_pullback(ȳ)
+      x̄ = project_x( adjoint(plan)*ChainRulesCore.unthunk(ȳ) )
+      return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), x̄
+  end
+  return y, mul_pullback
+end
+
+##################
+# NFFT, NFCT, NFST
+##################
+for (op,trans) in zip([:nfft, :nfct, :nfst], [:adjoint, :transpose, :transpose])
+
+  func_trans = Symbol("$(op)_$(trans)")
+  pbfunc = Symbol("$(op)_pullback")
+  pbfunc_trans = Symbol("$(op)_$(trans)_pullback")
+  @eval begin
+
+    # direct trafo
+    function ChainRulesCore.frule((_, Δx, _), ::typeof($(op)), k::AbstractMatrix, x::AbstractArray)
+      y = $(op)(k,x)
+      Δy = $(op)(k,Δx)
+      return y, Δy
+    end
+    function ChainRulesCore.rrule(::typeof($(op)), k::AbstractMatrix, x::AbstractArray)
+      y = $(op)(k,x)
+      project_x = ChainRulesCore.ProjectTo(x)
+      function $(pbfunc)(ȳ)
+        x̄ = project_x($(func_trans)(k, size(y), ChainRulesCore.unthunk(ȳ)))
+        return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), x̄
+      end
+      return y, nfft_pullback
+    end
+
+    # adjoint trafo
+    function ChainRulesCore.frule((_, Δx, _), ::typeof($(func_trans)), k::AbstractMatrix, N, x::AbstractArray)
+      y = $(func_trans)(k,N,x)
+      Δy = $(func_trans)(k,N,Δx)
+      return y, Δy
+    end
+    function ChainRulesCore.rrule(::typeof($(func_trans)), k::AbstractMatrix, N, x::AbstractArray)
+      y = $(func_trans)(k,N,x)
+      project_x = ChainRulesCore.ProjectTo(x)
+      function $(pbfunc_trans)(ȳ)
+        x̄ = project_x($(op)(k, ChainRulesCore.unthunk(ȳ)))
+        return ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), ChainRulesCore.NoTangent(), x̄
+      end
+      return y, $(pbfunc_trans)
+    end
+
+  end
+
+end
+
+
+
+
+end # module