Merge pull request #32 from slimgroup/integrate-rrule

ChainRules

Author: mloubout

Diff for: .github/workflows/ci-joli.yaml

@@ -19,14 +19,7 @@ jobs:
       fail-fast: false
 
       matrix:
-        version:
-          - '1.1'
-          - '1.2'
-          - '1.3'
-          - '1.4'
-          - '1.5'
-          - '1.6'
-
+        version: ['1.6', '1.7']
         os:
           - ubuntu-latest
           - macos-latest

Diff for: .travis.yml

@@ -4,7 +4,7 @@ os:
   - linux
   - osx
 julia:
-  - 1.3
+  - 1.6
 #  - nightly
 branches:
   only:

Diff for: Project.toml

@@ -1,7 +1,7 @@
 name = "JOLI"
 uuid = "bb331ad6-a1cf-11e9-23da-9bcb53c69f6f"
 authors = ["Henryk Modzelewski <[email protected]>"]
-version = "0.7.16"
+version = "0.8.0"
 
 [deps]
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
@@ -20,17 +20,24 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SharedArrays = "1a1011a3-84de-559e-8e89-a11a2f7dc383"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Wavelets = "29a6e085-ba6d-5f35-a997-948ac2efa89a"
 
 [compat]
 DistributedArrays = "0.5, 0.6"
 FFTW = "1"
+Flux = "0.12"
 InplaceOps = "0.3.0"
 IterativeSolvers = "0.8, 0.9"
-NFFT = "0.4, 0.5, 0.6"
+NFFT = "0.6 - 0.12"
 Nullables = "1"
 PyCall = "1.18, 1.90, 1.91, 1.62"
-SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1, 2"
+SpecialFunctions = "1.2, 2"
 Wavelets = "0.8, 0.9"
 julia = "1"
+
+[extras]
+Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[targets]
+test = ["Test", "Flux"]

Diff for: src/JOLI.jl

@@ -56,6 +56,7 @@ using NFFT
 using Wavelets
 using PyCall
 using SpecialFunctions
+using SpecialFunctions.ChainRulesCore
 
 # what's imported from Base
 import Base.eltype
@@ -84,6 +85,9 @@ import DistributedArrays.SPMD: scatter
 # what's imported from IterativeSolvers
 import IterativeSolvers.Adivtype
 
+# what's imported from ChainRulesCore
+import  SpecialFunctions.ChainRulesCore.rrule
+
 # extra exported methods
 export deltype, reltype
 export elements, hasinverse, issquare, istall, iswide, iscomplex, islinear, isadjoint, isequiv
@@ -134,4 +138,8 @@ include("joLinearFunctionConstructors.jl")
 include("joLinearOperatorConstructors.jl")
 include("joMixedConstructors.jl")
 
+
+# ChainRules
+include("rrule.jl")
+
 end # module

Diff for: src/joLinearFunctionConstructors/joNFFT.jl

@@ -7,27 +7,27 @@ module joNFFT_etc
     using JOLI: jo_convert
     function apply_nfft_centered(pln,n,v::Vector{vdt},rdt::DataType) where vdt<:Union{AbstractFloat,Complex}
         iv=jo_convert(ComplexF64,v,false)
-        rv=nfft(pln,iv)/sqrt(n)
+        rv= (pln*iv)/sqrt(n)
         rv=fftshift(rv)
         rv=jo_convert(rdt,rv,false)
         return rv
     end
     function apply_infft_centered(pln,n,v::Vector{vdt},rdt::DataType) where vdt<:Union{AbstractFloat,Complex}
         iv=jo_convert(ComplexF64,v,false)
         iv=ifftshift(iv)
-        rv=nfft_adjoint(pln,iv)/sqrt(n)
+        rv=(adjoint(pln)*iv)/sqrt(n)
         rv=jo_convert(rdt,rv,false)
         return rv
     end
     function apply_nfft(pln,n,v::Vector{vdt},rdt::DataType) where vdt<:Union{AbstractFloat,Complex}
         iv=jo_convert(ComplexF64,v,false)
-        rv=nfft(pln,iv)/sqrt(n)
+        rv= (pln*iv)/sqrt(n)
         rv=jo_convert(rdt,rv,false)
         return rv
     end
     function apply_infft(pln,n,v::Vector{vdt},rdt::DataType) where vdt<:Union{AbstractFloat,Complex}
         iv=jo_convert(ComplexF64,v,false)
-        rv=nfft_adjoint(pln,iv)/sqrt(n)
+        rv=(adjoint(pln)*iv)/sqrt(n)
         rv=jo_convert(rdt,rv,false)
         return rv
     end
@@ -36,21 +36,21 @@ using .joNFFT_etc
 
 export joNFFT
 """
-    julia> op = joNFFT(N,nodes[,m=...][,sigma=...][,window=...][,K=...];
+    julia> op = joNFFT(N,pos[,m=...][,sigma=...][,window=...];
                 [centered=...,][DDT=...,][RDT=...,][name=...])
 
 1D NFFT transform over fast dimension (wrapper to https://github.com/tknopp/NFFT.jl)
 
 # Signature
 
     function joNFFT(N::Integer,pos::Vector{joFloat},
-        m=4,sigma=2.0,window=:kaiser_bessel,K=2000; centered::Bool=false,
+        m=4,sigma=2.0,window=:kaiser_bessel; centered::Bool=false,
         DDT::DataType=joComplex,RDT::DataType=DDT,name::String="joNFFT")
 
 # Arguments
 
 - `N`: size
-- `nodes`: nodes' positions
+- `pos`: nodes positions
 - optional
     - see https://github.com/tknopp/NFFT.jl for info about optional parameters to NFFTplan: `m`, `sigma`, `window`, and `K`
 - keywords
@@ -67,24 +67,24 @@ export joNFFT
 
 1D NFFT
 
-    joNFFT(N,nodes)
+    joNFFT(N,pos)
 
 centered coefficients
 
-    joNFFT(N,nodes; centered=true)
+    joNFFT(N,pos; centered=true)
 
 examples with DDT/RDT
 
-    % joNFFT(N,nodes; DDT=ComplexF32)
-    % joNFFT(N,nodes; DDT=ComplexF32,RDT=ComplexF64)
+    % joNFFT(N,pos; DDT=ComplexF32)
+    % joNFFT(N,pos; DDT=ComplexF32,RDT=ComplexF64)
 
 """
 function joNFFT(N::Integer,pos::Vector{joFloat},
-    m=4,sigma=2.0,window=:kaiser_bessel,K=2000; centered::Bool=false,
+    m=4,sigma=2.0,window=:kaiser_bessel; centered::Bool=false,
     DDT::DataType=joComplex,RDT::DataType=DDT,name::String="joNFFT")
 
     M = length(pos)
-    p = try NFFTPlan(pos,N,m,sigma,window,K) catch; plan_nfft(pos,N,m,sigma,window,K) end
+    p = plan_nfft(pos, N; m=m, σ=sigma, window=window)
     if centered
         return joLinearFunctionFwd_A(M,N,
             v1->joNFFT_etc.apply_nfft_centered(p,N,v1,RDT),

Diff for: src/joMixedConstructors/joGaussian.jl

@@ -122,7 +122,7 @@ examples with DDT/RDT
 """
 function joGaussian(M::Integer,N::Integer=M;
          implicit::Bool=false,normalized::Bool=false,orthonormal::Bool=false,
-         RNG::AbstractRNG=Random.seed!(),
+         RNG::AbstractRNG=MersenneTwister(),
          DDT::DataType=joFloat,RDT::DataType=DDT,
          name::String="joGaussian")
 

Diff for: src/rrule.jl

@@ -0,0 +1,8 @@
+function rrule(::typeof(*), A::joAbstractLinearOperator{ADDT,ARDT}, v) where {ADDT,ARDT}
+  Y = A*v
+  function time_pullback(dy)
+    DY = unthunk(Vector(dy))
+    return NoTangent(), NoTangent(), @thunk(A' * DY)
+  end
+  return Y, time_pullback
+end

Diff for: test/runtests.jl

@@ -46,6 +46,7 @@ include("test_joCurvelet2DnoFFT.jl")
 include("test_joExtend.jl")
 include("test_joGaussian.jl")
 include("test_joOuterProd.jl")
+include("test_rrules.jl")
 etime=time()
 dtime=etime-stime
 println("\nTest Total elapsed time: ",round(dtime,digits=1),"s")

Diff for: test/test_joNFFT.jl

@@ -4,7 +4,7 @@ tsname="joNFFT"
 for t=1:T # start test loop
     m=4^t
     n=3^t
-    nodes=sort(rand(joFloat,n))
+    nodes=sort(rand(joFloat,n)) .- joFloat(.5)
 
     verbose && println("$tsname ($m[,$m]) - not centered")
     @testset "$m [x $m]" begin

Diff for: test/test_rrules.jl

@@ -0,0 +1,11 @@
+using Flux
+
+x = rand(Float32, 3)
+
+W = randn(Float32, 3, 3)
+W_JOLI = joMatrix(W)
+
+g = gradient(x -> sum(W*x), x)[1]
+g_JOLI = gradient(x -> sum(W_JOLI*x), x)[1]
+
+@test isapprox(sum((g - g_JOLI).^2)/sum(g_JOLI.^2), 0f0; atol=1f-8)