Automatic Differentiation of quantum master equation using Zygote

Specific Domains
, ,
1

Hello,

I need to differentiate a quantum master equation over a parameter, using QuantumToolbox.jl. Under the hood it is simply an Ordinary Differential Equation involving SciMLOperators.

I’m able to differentiate the Schrödinger equation with the following code

using QuantumToolbox
using OrdinaryDiffEq
using SciMLSensitivity
using Zygote
using Enzyme
using BenchmarkTools

##

const N = 16
F = 1
γ = 1

const a = destroy(N)

const ψ0 = fock(N, 0)

# coef1(p, t) = p.Δ
coef1(p, t) = p[1]

const H = QobjEvo(a' * a, coef1) + F * (a' + a)
const c_ops = [sqrt(γ) * a]

const L = liouvillian(H, c_ops)

function my_f_sesolve(p)
  tlist = range(0, 1.5, 100)

  sol = sesolve(H, ψ0, tlist, progress_bar=Val(false), params=p, saveat=[tlist[end]], sensealg = BacksolveAdjoint(autojacvec=EnzymeVJP()))

  # return real(abs2(sol.states[end].data[end]))
  return real(expect(a' * a, sol.states[end]))
end

Δ = 1.0
params = [Δ]
my_f_sesolve(params) # 1.8585255970284336

##

Zygote.gradient(my_f_sesolve, params) # ([-2.7374779226913737],)

The internal code is probably too complicated to be explained here, but the working principle is basic linear algebra between matrices and vectors.

Current limitations of autodiff for this situation can be found in this issue, especially in the last comment where I summarize everything.

However, if I try to differentiate the master equation it fails. This case should be equivalent to the previous one, where the SciMLOperator involved in the ODE is only bigger than the Schrödinger case. Here the code

function my_f_mesolve(p)
  tlist = range(0, 1.5, 100)

  sol = mesolve(H, ψ0, tlist, c_ops, progress_bar=Val(false), params=p, saveat=[tlist[end]], sensealg = BacksolveAdjoint(autojacvec=EnzymeVJP()))

  # return real(abs2(sol.states[end].data[end]))
  return real(expect(a' * a, sol.states[end]))
end

##

Δ = 1.0
params = [Δ]
my_f_mesolve(params) # 0.925041907046927

##

Zygote.gradient(my_f_mesolve, params)

And the Error:

Summary 
ERROR: Need an adjoint for constructor QuantumObject{OperatorQuantumObject, Dimensions{1, Tuple{Space}}, Matrix{ComplexF64}}. Gradient is of type Diagonal{Float64, Vector{Float64}}
Stacktrace:
  [1] error(s::String)
    @ Base ./error.jl:35
  [2] (::Zygote.Jnew{QuantumObject{…}, Nothing, false})(Δ::Diagonal{Float64, Vector{…}})
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/lib/lib.jl:324
  [3] (::Zygote.var"#2230#back#327"{Zygote.Jnew{QuantumObject{…}, Nothing, false}})(Δ::Diagonal{Float64, Vector{Float64}})
    @ Zygote ~/.julia/packages/ZygoteRules/CkVIK/src/adjoint.jl:72
  [4] QuantumObject
    @ ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/quantum_object.jl:59 [inlined]
  [5] (::Zygote.Pullback{Tuple{…}, Any})(Δ::Diagonal{Float64, Vector{…}})
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
  [6] *
    @ ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/arithmetic_and_attributes.jl:72 [inlined]
  [7] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Diagonal{Float64, Vector{…}})
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
  [8] expect
    @ ~/.julia/packages/QuantumToolbox/YmsAS/src/qobj/functions.jl:58 [inlined]
  [9] my_f_mesolve
    @ ~/GitHub/Research/Undef/Autodiff QuantumToolbox/autodiff.jl:65 [inlined]
 [10] (::Zygote.Pullback{Tuple{…}, Tuple{…}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface2.jl:0
 [11] (::Zygote.var"#88#89"{Zygote.Pullback{Tuple{…}, Tuple{…}}})(Δ::Float64)
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface.jl:97
 [12] gradient(f::Function, args::Vector{Float64})
    @ Zygote ~/.julia/packages/Zygote/HdT4O/src/compiler/interface.jl:154
 [13] top-level scope
    @ ~/GitHub/Research/Undef/Autodiff QuantumToolbox/autodiff.jl:76
Some type information was truncated. Use `show(err)` to see complete types.

The differentiation of the ODE itself seem to work, since if I return real(abs2(sol.states[end].data[end])) instead of real(expect(a' * a, sol.states[end])) it works. So, apparently the expect function doesn’t work in the case of mesolve, but it works in the case of sesolve.

To give more context, a QuantumObject like the a variable here is a struct containing three fields

  • The data field: the actual matrix or vector representing the state or the operator
  • The type and dimensions fields: some structs to store the information of the quantum system

Conclusions

It seems that I have to define some adjoint for the QuantumObject struct, whatever it means. If so, I don’t understand why I don’t need it for the Schrödinger equation case, but I need it for the master equation.

2

Yes I need to get to this… I know, it’s still in my email.

3

Don’t worry, in this thread I’m asking something different from the issue on GitHub. Here I get the error

ERROR: Need an adjoint for constructor QuantumObject{OperatorQuantumObject, Dimensions{1, Tuple{Space}}, Matrix{ComplexF64}}. Gradient is of type Diagonal{Float64, Vector{Float64}}

When I return real(expect(a' * a, sol.states[end])) rather than a simple matrix element real(abs2(sol.states[end].data[end])).

So it is something apparently not related to the ODE itself, but to my custom type.

4

Oh you’re missing an adjoint for this constructor function.

5

Ok, but how is it possible that with sesolve (the first ODE) it doesn’t need it? It only happens if I call the expect function, which is basically

expect(A::QuantumObject, B::QuantumObject) = tr(A.data * B.data)

where the data field is a simple matrix.

How can I fix it?

6

Because the generic adjoint of getfield uses the constructor adjoint in the backpass. You just need something like this:

I.e., the derivative of getfield is basically just an identity, so the return is A.data, but then for the backwards pass if J = I then J' = I, so the derivative is just the identity, so you just build a new QuantumObject with the pullback value.

7

I tried first this

Zygote.@adjoint function Zygote.literal_getproperty(A::QuantumObject, ::Val{:data})
  A.data, p̄ -> (QuantumObject(p̄, A.type, A.dimensions),)
end

but I still get the same error.

I then tried

Zygote.@adjoint function QuantumObject(A::AbstractArray, type, dimensions)
  QuantumObject(A, type, dimensions), p̄ -> (QuantumObject(p̄, type, dimensions),)
end

But the I get the error

BoundsError: attempt to access Tuple{Nothing, QuantumObject{OperatorQuantumObject, Dimensions{1, Tuple{Space}}, Diagonal{Float64, Vector{Float64}}}} at index [3]

I really don’t know why it works for sesolve but not for mesolve. They are both ODE with SciMLOperators, where in the second case the operator is just slightly different but still of the same form.

The error seems to be related to the last line and not to the ODE itself.

8

I tried a simpler function to test Zygote.jl on QuantumObjects, and it works without defining custom adjoints for the QuantumObject constructor.

const ψgg = tensor(basis(2, 1), basis(2, 1))
const ψee = tensor(basis(2, 0), basis(2, 0))

function my_fidelity(p)
  ψ = normalize(p[1] * ψgg + p[2] * ψee)

  return fidelity(ψ, ψgg)
end

p = [0.3, 0.4]
my_fidelity(p) # 0.6

Zygote.gradient(my_fidelity, p) # ([1.28, -0.96],)

Here ψgg and ψee are QuantumObject types containing a Vector as data. The fidelity function just computes dot(ψ.data, ψgg.data).

9

You’re missing the derivative w.r.t. the other parts. I’d normally do that one with ChainRules though. But:

Zygote.@adjoint function QuantumObject(A::AbstractArray, type, dimensions)
  QuantumObject(A, type, dimensions), p̄ -> (QuantumObject(p̄, type, dimensions), ChainRules.ZeroTangent(), ChainRules.ZeroTangent())
end

since the derivative w.r.t. the other two parts is structurally zero since they are discrete.

10

It depends on how the values are accessed.

11

Ok, fair, thanks.

I tried and now the error is

ERROR: MethodError: no method matching *(::QuantumObject{OperatorQuantumObject, Dimensions{…}, Diagonal{…}}, ::Adjoint{ComplexF64, Matrix{…}})

It is weird, since it is trying to do the product between a QuantumObject and a matrix.

However, the sesolve case no longer works with this adjoint definition, and it gives

ERROR: MethodError: no method matching QuantumObject(::@NamedTuple{data::SparseMatrixCSC{ComplexF64, Int64}, type::Nothing, dimensions::Nothing}, ::OperatorQuantumObject, ::Dimensions{1, Tuple{Space}})

I tried to change the function to something like

Zygote.@adjoint function QuantumObject(A, type, dimensions)
  back(p̄) = (QuantumObject(p̄, type, dimensions), ChainRules.ZeroTangent(), ChainRules.ZeroTangent())
  back(p̄::NamedTuple{(:data, )}) = (QuantumObject(p̄.data, type, dimensions), ChainRules.ZeroTangent(), ChainRules.ZeroTangent())
  QuantumObject(A, type, dimensions), back
end

But it is still calling the first method of the back function.

EDIT

By changing the definition to

Zygote.@adjoint function QuantumObject(A, type, dimensions)
  back(p̄) = (QuantumObject(p̄, type, dimensions), ChainRules.ZeroTangent(), ChainRules.ZeroTangent())
  back(p̄::NamedTuple{(:data, :type, :dimensions)}) = (QuantumObject(p̄.data, type, dimensions), ChainRules.ZeroTangent(), ChainRules.ZeroTangent())
  QuantumObject(A, type, dimensions), back
end

now sesolve has the same error as for mesolve.

12

back(p̄::NamedTuple{(:data, :type, :dimensions)}) that’s a missing projection / means that the constructor adjoint wasn’t called.

13

Why this is happening? How to fix it?

EDIT

I tried to play a bit. I found the following case working

Zygote.@adjoint function QuantumObject(data, type, dims)
  obj = QuantumObject(data, type, dims)
  function pullback(Δobj)
      # Return gradient for data, and ignore type and dims
      return (Δobj.data, nothing, nothing)
  end
  function pullback(Δobj::AbstractArray)
    # Return gradient for data, and ignore type and dims
    return (Δobj, nothing, nothing)
  end
  return obj, pullback
end

I need the definition of the pullback also for standard AbstractArrays.

Is it reasonable? If yes, how can I convert it to a ChainRulesCore.jl rrule?

14

In the direction of defining the adjoint with ChainRulesCore.jl, I tried this implementation

function ChainRulesCore.rrule(::Type{QuantumObject}, data, type, dimensions)
  obj = QuantumObject(data, type, dimensions)
  f_pullback(Δobj) = (NoTangent(), Δobj.data, NoTangent(), NoTangent())
  f_pullback(Δobj_data::AbstractArray) = (NoTangent(), Δobj_data, NoTangent(), NoTangent())
  return obj, f_pullback
end

It seems to work. I had to change typeof(QuantumObject) with Type{QuantumObject}, maybe because it is a constructor? Also, the pullback function is not using the Tangent type. But I didn’t managed to use it.

With this implementation, my code is working also with ChainRulesCore.jl.

Do you think that these two definitions (Zygote and ChainRulesCore) are correct and reasonable?

15

yes

1 Like