Skip to content

Commit 9ec7cb4

Browse files
albertomercurioalbertomercurio
committed
Complete documentation
1 parent 51e6e50 commit 9ec7cb4

File tree

1 file changed

+149
-9
lines changed

1 file changed

+149
-9
lines changed

‎docs/src/users_guide/autodiff.md

Lines changed: 149 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
# [Differentiate almost anything in QuantumToolbox](@id doc:autodiff)
1+
# [Differentiate almost anything in QuantumToolbox.jl](@id doc:autodiff)
22

33
Automatic differentiation (AD) has emerged as a key technique in computational science, enabling exact and efficient computation of derivatives for functions defined by code. Unlike symbolic differentiation, which may produce complex and inefficient expressions, or finite-difference methods, which suffer from numerical instability and poor scalability, AD leverages the chain rule at the level of elementary operations to provide machine-precision gradients with minimal overhead.
44

@@ -8,12 +8,6 @@ In `QuantumToolbox.jl`, we have introduced preliminary support for automatic dif
88
At present, this functionality is considered experimental and not all parts of the library are AD-compatible. Here we provide a brief overview of the current state of AD support in `QuantumToolbox.jl` and how to use it.
99

1010

11-
::: warning
12-
13-
At present, this functionality is considered experimental and not all parts of the library are AD-compatible. Here we provide a brief overview of the current state of AD support in `QuantumToolbox.jl` and how to use it.
14-
15-
:::
16-
1711
## [Forward VS Reverse Mode AD](@id doc:autodiff:forward_vs_reverse)
1812

1913
Automatic differentiation can be broadly categorized into two modes: forward mode and reverse mode. The choice between these modes depends on the nature of the function being differentiated and the number of inputs and outputs:
@@ -34,8 +28,14 @@ One of the primary use cases for automatic differentiation in `QuantumToolbox.jl
3428

3529
Hence, the density matrix will evolve according to the master equation
3630

31+
```@raw html
32+
<span id="eq:master-equation"></span>
33+
```
3734
```math
38-
\frac{d \hat{\rho}(\mathbf{p}, t)}{dt} = -i[\hat{H}(\mathbf{p}), \hat{\rho}(\mathbf{p}, t)] + \sum_j \hat{L}_j(\mathbf{p}) \hat{\rho}(\mathbf{p}, t) \hat{L}_j(\mathbf{p})^\dagger - \frac{1}{2} \left\{ \hat{L}_j(\mathbf{p})^\dagger \hat{L}_j(\mathbf{p}), \hat{\rho}(\mathbf{p}, t) \right\} \, ,
35+
\begin{align}
36+
\frac{d \hat{\rho}(\mathbf{p}, t)}{dt} =& -i[\hat{H}(\mathbf{p}), \hat{\rho}(\mathbf{p}, t)] \\
37+
&+ \sum_j \hat{L}_j(\mathbf{p}) \hat{\rho}(\mathbf{p}, t) \hat{L}_j(\mathbf{p})^\dagger - \frac{1}{2} \left\{ \hat{L}_j(\mathbf{p})^\dagger \hat{L}_j(\mathbf{p}), \hat{\rho}(\mathbf{p}, t) \right\} \, ,
38+
\end{align} \tag{1}
3939
```
4040

4141
which depends on the parameters $\mathbf{p}$ and time $t$.
@@ -51,5 +51,145 @@ which will also depend on the parameters $\mathbf{p}$ and time $t$.
5151
Our goal is to compute the derivative of the expectation value with respect to the parameters:
5252

5353
```math
54-
\frac{\partial \langle \hat{O}(\mathbf{p}, t) \rangle}{\partial p_j} = \frac{\partial}{\partial p_j} \text{Tr}[\hat{O} \hat{\rho}(\mathbf{p}, t)] \, .
54+
\frac{\partial \langle \hat{O}(\mathbf{p}, t) \rangle}{\partial p_j} = \frac{\partial}{\partial p_j} \text{Tr}[\hat{O} \hat{\rho}(\mathbf{p}, t)] \, ,
55+
```
56+
57+
and to achieve this, we can use an AD engine like `ForwardDiff.jl` (forward mode) or `Zygote.jl` (reverse mode).
58+
59+
Let's apply this to a simple example of a driven-dissipative quantum harmonic oscillator. The Hamiltonian in the drive frame is given by
60+
61+
```math
62+
\hat{H} = \Delta \hat{a}^\dagger \hat{a} + F \left( \hat{a} + \hat{a}^\dagger \right) \, ,
63+
```
64+
65+
where $\Delta = \omega_0 - \omega_d$ is the cavity-drive detuning, $F$ is the drive strength, and $\hat{a}$ and $\hat{a}^\dagger$ are the annihilation and creation operators, respectively. The system is subject to a single dissipative channel with a Lindblad operator $\hat{L} = \sqrt{\gamma} \hat{a}$, where $\gamma$ is the dissipation rate. If we start from the ground state $\hat{\rho}(0) = \vert 0 \rangle \langle 0 \vert$, the systems evolves according to the master equation in [Eq. (1)](#eq:master-equation).
66+
67+
We now want to study the number of photons at the steady state, and how it varies with $\mathbf{p} = (\Delta, F, \gamma)$, namely $\nabla_\mathbf{p} \langle \hat{a}^\dagger \hat{a} \rangle (\mathbf{p}, t \to \infty)$. We can extract an analytical expression, in order to verify the correctness of the AD implementation:
68+
69+
```math
70+
\langle \hat{a}^\dagger \hat{a} \rangle_\mathrm{ss} = \frac{F^2}{\Delta^2 + \frac{\gamma^2}{4}} \, ,
71+
```
72+
73+
with the gradient given by
74+
75+
```math
76+
\nabla_\mathbf{p} \langle \hat{a}^\dagger \hat{a} \rangle_\mathrm{ss} =
77+
\begin{pmatrix}
78+
\frac{-2 F^2 \Delta}{(\Delta^2 + \frac{\gamma^2}{4})^2} \\
79+
\frac{2 F}{\Delta^2 + \frac{\gamma^2}{4}} \\
80+
\frac{-F^2 \gamma}{2 (\Delta^2 + \frac{\gamma^2}{4})^2}
81+
\end{pmatrix} \, .
82+
```
83+
84+
Although `QuantumToolbox.jl` has the [`steadystate`](@ref) function to directly compute the steady state without explicitly solving the master equation, here we use the [`mesolve`](@ref) function to integrate up to a long time $t_\mathrm{max}$, and then compute the expectation value of the number operator. We will demonstrate how to compute the gradient using both `ForwardDiff.jl` and `Zygote.jl`.
85+
86+
### [Forward Mode AD with ForwardDiff.jl](@id doc:autodiff:forward)
87+
88+
```@setup autodiff
89+
using QuantumToolbox
90+
```
91+
92+
We start by importing `ForwardDiff.jl` and defining the parameters and operators:
93+
94+
```@example autodiff
95+
using ForwardDiff
96+
97+
const N = 20
98+
const a = destroy(N)
99+
const ψ0 = fock(N, 0)
100+
const t_max = 40
101+
const tlist = range(0, t_max, 100)
55102
```
103+
104+
Then, we define a function that take the parameters `p` as an input and returns the expectation value of the number operator at `t_max`. We also define the analytical solution of the steady state photon number and its gradient for comparison:
105+
106+
```@example autodiff
107+
function my_f_mesolve_direct(p)
108+
H = p[1] * a' * a + p[2] * (a + a')
109+
c_ops = [sqrt(p[3]) * a]
110+
sol = mesolve(H, ψ0, tlist, c_ops, progress_bar = Val(false))
111+
return real(expect(a' * a, sol.states[end]))
112+
end
113+
114+
# Analytical solution
115+
function my_f_analytical(p)
116+
Δ, F, γ = p
117+
return F^2 / (Δ^2 + γ^2 / 4)
118+
end
119+
function my_grad_analytical(p)
120+
Δ, F, γ = p
121+
return [
122+
-2 * F^2 * Δ / (Δ^2 + γ^2 / 4)^2,
123+
2 * F / (Δ^2 + γ^2 / 4),
124+
-F^2 * γ / (2 * (Δ^2 + γ^2 / 4)^2)
125+
]
126+
end
127+
```
128+
129+
The gradient can be computed using `ForwardDiff.gradient`:
130+
131+
```@example autodiff
132+
Δ = 1.5
133+
F = 1.5
134+
γ = 1.5
135+
params = [Δ, F, γ]
136+
137+
grad_exact = my_grad_analytical(params)
138+
grad_fd = ForwardDiff.gradient(my_f_mesolve_direct, params)
139+
```
140+
141+
and test if the results match:
142+
143+
```@example autodiff
144+
isapprox(grad_exact, grad_fd; atol = 1e-6)
145+
```
146+
147+
### [Reverse Mode AD with Zygote.jl](@id doc:autodiff:reverse)
148+
149+
Reverse-mode differentiation is significantly more challenging than forward-mode when dealing ODEs, as the complexity arises from the need to propagate gradients backward through the entire time evolution of the quantum state.
150+
151+
`QuantumToolbox.jl` leverages the advanced capabilities of [`SciMLSensitivity.jl`](https://github.com/SciML/SciMLSensitivity.jl) to handle this complexity. `SciMLSensitivity.jl` implements sophisticated methods for computing gradients of ODE solutions, such as the adjoint method, which computes gradients by solving an additional "adjoint" ODE backward in time. For more details on the adjoint method and other sensitivity analysis techniques, please refer to the [`SciMLSensitivity.jl` documentation](https://docs.sciml.ai/SciMLSensitivity/stable/).
152+
153+
In order to reverse-differentiate the master equation, we need to define the operators as [`QuantumObjectEvolution`](@ref) objects, which use [`SciMLOperators.jl`](https://github.com/SciML/SciMLOperators.jl) to represent parameter-dependent operators.
154+
155+
```@example autodiff
156+
using Zygote
157+
using SciMLSensitivity
158+
159+
# For SciMLSensitivity.jl
160+
coef_Δ(p, t) = p[1]
161+
coef_F(p, t) = p[2]
162+
coef_γ(p, t) = sqrt(p[3])
163+
H = QobjEvo(a' * a, coef_Δ) + QobjEvo(a + a', coef_F)
164+
c_ops = [QobjEvo(a, coef_γ)]
165+
const L = liouvillian(H, c_ops)
166+
167+
function my_f_mesolve(p)
168+
sol = mesolve(
169+
L,
170+
ψ0,
171+
tlist,
172+
progress_bar = Val(false),
173+
params = p,
174+
sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()),
175+
)
176+
177+
return real(expect(a' * a, sol.states[end]))
178+
end
179+
```
180+
181+
And the gradient can be computed using `Zygote.gradient`:
182+
183+
```@example autodiff
184+
grad_zygote = Zygote.gradient(my_f_mesolve, params)[1]
185+
```
186+
187+
Finally, we can compare the results from `ForwardDiff.jl` and `Zygote.jl`:
188+
189+
```@example autodiff
190+
isapprox(grad_fd, grad_zygote; atol = 1e-6)
191+
```
192+
193+
## [Conclusion](@id doc:autodiff:conclusion)
194+
195+
In this section, we have explored the integration of automatic differentiation into `QuantumToolbox.jl`, enabling users to compute gradients of observables and cost functionals involving the time evolution of open quantum systems. We demonstrated how to differentiate the master equation using both forward mode with `ForwardDiff.jl` and reverse mode with `Zygote.jl`, showcasing the flexibility and power of automatic differentiation in quantum computing applications. AD can be applied to other functions in `QuantumToolbox.jl`, although the support is still experimental and not all functions are guaranteed to be compatible. We encourage users to experiment with AD in their quantum simulations and contribute to the ongoing development of this feature.

0 commit comments

Comments
 (0)