# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .jl
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Julia 1.6.3
#     language: julia
#     name: julia-1.6
# ---

# # Chainrules: custom AD in julia
#
# In this tutorial, we will look at the automatic differentiation in julia, and in particular how `[ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl)` supports addition of your own differentiation rules into Julia's [Automatic Differentiation](https://juliadiff.org) system. This allows for seamless integration between your code, including its handcoded derivatives, and Julia's native AD systems, e.g. those used by [Flux](https://fluxml.ai/Flux.jl/stable/) Julia's machine learning platform. We use [ChainRules.jl](https://github.com/JuliaDiff/ChainRules.jl) to automatically differentiate codes that involve complex operations implemented by [JUDI.jl](https://github.com/slimgroup/JUDI.jl) (this example) and [JOLI.jl](https://github.com/slimgroup/JOLI.jl).

using Zygote, ChainRulesCore, LinearAlgebra

# ## Simple example
#
# Let's consider a simple example with a basic differentiable function $x -> cos(x) + 1$

mycos(x) = cos(x) + 1

# Now we know that the derivative of this function is $x -> -sin(x)$ from standard functional analysis.e can therefore define through chainrules a new rule for our function `f`

function ChainRulesCore.rrule(::typeof(mycos), x)
    println("Using custom AD rule for mycos")
    y = mycos(x)
    pullback(Δy) = (NoTangent(), -sin(x)*Δy)
    return y, pullback
end

# We now have the rule to compute the directional derivative of our function `mycos` . Let's check the gradient

x0 = randn()
# Standard AD of cos
g1 = gradient(x->norm(cos(x)+1)^2, x0);
# Our definition
g2 = gradient(x->norm(mycos(x))^2, x0);
# Analytical gradient
g3 = -2*sin(x0)*(mycos(x0));
println("True gradient: $(g3) \nStandard AD  : $(g1[1]) \nCustom AD    : $(g2[1])")

# And we see that we get the correct gradient. Nw this is an extremely simple case, now let's look at a more complicated case where we define the AD rule for matrix-free operators defined in JOLI

# # JOLI
#
# We look at how we define automatic differentiation rules involding matrix-free linear operator. We consider operations of the form `A*x` where `A` is a JOLI matrix-free linear operator and we differentiate with respect to the input `x`
#
# In JOLI, the base type of our linear operator is `joAbstractLinearOperator`. If we define the rule for this abstract type, all linear operator should follow. Now in this case the acual operation to be differentiated is the multiplication `*` with two inputs (`A, x`). Because we consider `A` to be fixed,its derivative will be defined as `NoTangent` that is `ChainRules`'s way to state there is no derivative for this input.

using JOLI

function ChainRulesCore.rrule(::typeof(*), A::T, x) where {T<:joAbstractLinearOperator}
    y = A*x
    pullback(Δy) = (NoTangent(), NoTangent(), A'*Δy)
    return y, pullback
end

# With this rule defined we can now use a JOLI operator. Let's solve a simple data fitting problem with a restricted Fourier measuerment
#

using Random
N = 128
# Fourier transform as a linear operator
F = joDFT(N)
# Restriction
R = joRomberg(N; DDT=Complex{Float64}, RDT=Complex{Float64})
# Combine the operators
A = R*F;

# Make data
x = randn(128)
b = A*x;

# Let's create a loss function

loss(x) = .5f0*norm(A*x - b)^2 + norm(A*x - b, 1)

# We can now easily obtain the gradient at any given `x` since the only undefined part would have been the JOLI operator that now has its own differentiation rule

gradient(loss, randn(128))

# ### Optimization
# Let's now slve the problem above with standard gradient descent

using Optim

δloss!(g, x) = begin g.=gradient(loss, x)[1]; return loss(x) end;

summary = optimize(loss, δloss!, randn(N), ConjugateGradient(),
                   Optim.Options(g_tol = 1e-12, iterations = 200, store_trace = true, show_trace = true, show_every=1))

using PyPlot
plot(x, label="true")
plot(summary.minimizer, label="Recovered")
legend()

StarcoderdataJupyterScriptsDedupFiltered