Basic examples
First time using dynamiqs? Below are a few basic examples to help you get started.
Simulate a lossy quantum harmonic oscillator
This first example shows simulation of a lossy harmonic oscillator with Hamiltonian \(H=\omega a^\dagger a\) and a single jump operator \(L=\sqrt{\kappa} a\).
import dynamiqs as dq
import jax.numpy as jnp
# parameters
n = 128 # Hilbert space dimension
omega = 1.0 # frequency
kappa = 0.1 # decay rate
alpha = 1.0 # initial coherent state amplitude
# initialize operators, initial state and saving times
a = dq.destroy(n)
H = omega * dq.dag(a) @ a
jump_ops = [jnp.sqrt(kappa) * a]
psi0 = dq.coherent(n, alpha)
tsave = jnp.linspace(0, 1.0, 101)
# run simulation
result = dq.mesolve(H, jump_ops, psi0, tsave)
print(result)
Output
|██████████| 100.0% ◆ elapsed 66.94ms ◆ remaining 0.00ms
==== MESolveResult ====
Solver : Tsit5
Infos : 7 steps (7 accepted, 0 rejected)
States : Array complex64 (101, 128, 128) | 12.62 Mb
Compute gradients with respect to some parameters
Suppose that in the above example, we want to compute the gradient of the number of photons in the final state, \(\bar{n} = \mathrm{Tr}[a^\dagger a \rho(t_f)]\), with respect to the decay rate \(\kappa\) and the initial coherent state amplitude \(\alpha\).
import dynamiqs as dq
import jax.numpy as jnp
import jax
# parameters
n = 128 # Hilbert space dimension
omega = 1.0 # frequency
kappa = 0.1 # decay rate
alpha = 1.0 # initial coherent state amplitude
def population(omega, kappa, alpha):
"""Return the oscillator population after time evolution."""
# initialize operators, initial state and saving times
a = dq.destroy(n)
H = omega * dq.dag(a) @ a
jump_ops = [jnp.sqrt(kappa) * a]
psi0 = dq.coherent(n, alpha)
tsave = jnp.linspace(0, 1.0, 101)
# run simulation
result = dq.mesolve(H, jump_ops, psi0, tsave)
return dq.expect(dq.number(n), result.states[-1]).real
# compute gradient with respect to omega, kappa and alpha
grad_population = jax.grad(population, argnums=(0, 1, 2))
grads = grad_population(omega, kappa, alpha)
print(f'Gradient w.r.t. omega={grads[0]:.2f}')
print(f'Gradient w.r.t. kappa={grads[1]:.2f}')
print(f'Gradient w.r.t. alpha={grads[2]:.2f}')
Output
|██████████| 100.0% ◆ elapsed 86.63ms ◆ remaining 0.00ms
Gradient w.r.t. omega=0.00
Gradient w.r.t. kappa=-0.90
Gradient w.r.t. alpha=1.81