Hello, so I am currently working on code for both the LCHS method and the Spectral split-step method using PennyLane resources to solve PDEs, specifically for my case I am attempting to solve the Fokker-Planck equation for the Wiener process with mu= 0 and D=1/2 and how it evolves. I have had many issues when it comes to solving my PDE, and am still having concerns about my solution and how it compares to the classical version. I think if I get the correct solution for my spectral split-step and what I could be doing wrong, I will be able to deduce why my LCHS code isn’t working (or you’ll see another post by me).
Firstly, if you are to run this code now and observe the graph, you’ll notice that as you increase the qubit number, the probabilities for all except the initial case decrease. However, if you are to remove the comment for # current_state /= scipy.integrate.simps(current_state,x_grid) # Normalize (line 65) and run it again, you’ll notice that the probabilities remain constant with changing the qubits, however, the probabilities are too high for the evolved distributions.
I’m unsure if this normalization, bringing it closer to the answer, is just a fluke, and that my circuit is evolving the equation too much, or if I am on the right track.
# Put code here
import numpy as np
import pennylane as qml
import matplotlib.pyplot as plt
import scipy
import time
def fokker_planck_spectral_split_step(n_qubits, dt, total_steps, mu, D):
"""
Solve the Fokker-Planck equation using the spectral split-step method.
Args:
n_qubits: Number of qubits for spatial discretization
dt: Time step size
total_steps: Number of time steps to simulate
x_min, x_max: Domain boundaries
mu: Drift coefficient (can be a constant or function of x)
D: Diffusion coefficient (can be a constant or function of x)
"""
dev = qml.device("default.qubit", wires=n_qubits)
n_points = 2**n_qubits #amount of nodes
x_min, x_max = -10,10
dx = (x_max - x_min) / (n_points-1)
x_grid = np.linspace(x_min, x_max, n_points)
#Calculate momentum grid values
k_grid = 2*np.pi*np.fft.fftfreq(n_points, dx)
#Gaussian distribution
def initial_condition(x):
return scipy.stats.norm.pdf(x, scale=0.5) #scale = width
#position space: e^(-i*H_drift*dt/2) = 0 (for Wiener Process)
def calculate_drift_phases():
phases = np.zeros(n_points)
return phases # handle drift in momentum space
#momentum space: e^(-i*H_diffusion*dt)
def diffusion_phases(k):
return -D * (k**2)
@qml.qnode(dev)
def fokker_planck_step(state_vector):
qml.StatePrep(state_vector, wires=range(n_qubits), normalize=True)
#QFT to momentum space
qml.QFT(wires=range(n_qubits))
# Apply diffusion in momentum space
generator = diffusion_phases(k_grid) #only -D *k^2 for Wiener process
diagonal_elements = np.exp(1j * generator*dt) #the exponential of FT , dt is time slice
qml.DiagonalQubitUnitary(diagonal_elements, wires=range(n_qubits)) #allows operator to be applied
#Inverse QFT to position space
qml.adjoint(qml.QFT(wires=range(n_qubits)))
return qml.state()
# Time evolution
psi_0 = initial_condition(x_grid) #making the gaussian depn on qubits
psi_0 /= scipy.integrate.simps(psi_0,x_grid) # Normalize (gaussian is already normalized)
current_state = psi_0.copy()
q_states = [current_state]
starttime = time.time()
for step in range(total_steps):
# Quantum evolution
current_state = fokker_planck_step(current_state)
# current_state /= scipy.integrate.simps(current_state,x_grid) # Normalize
q_states.append(np.real(abs(current_state)))
endtime = time.time()
print(f'Time for algorithm: {endtime-starttime} seconds')
return {
'x_grid': x_grid,
'quantum_states': q_states,
'times': np.arange(0, (total_steps+1) * dt, dt)
}
#plotting function
def plot_fokker_planck_results(results, t_cut):
'''Args:
results: the results from running the quantum algorithm
t_cut: the time cutoff for observations on the graph
'''
x_grid = results['x_grid']
quantum_states = results['quantum_states']
times = results['times']
t_list = np.arange(0, t_cut, 0.4)
t_fixed = []
for t_target in t_list:
idx = np.argmin(np.abs(np.array(times) - t_target))
t_fixed.append(idx)
print(f"Target time: {t_target}, Using time: {times[idx]:.3f}")
for i, t in enumerate(t_fixed):
plt.plot(x_grid, quantum_states[t], 'o-',
label=f't = {t_list[i]:.2f}')
plt.title("SpectralSplit")
plt.xlabel('x')
plt.ylabel('P(x,t)')
plt.legend()
plt.grid(True)
plt.savefig('fokker_planck_spectral.png', dpi=300)
plt.show()
results = fokker_planck_spectral_split_step(
n_qubits=8,
dt=0.1,
total_steps=16,
mu=0.0, # drift (pure diffusion case)
D=0.5 #Diffusion coefficient
)
time_cutoff = 2
plot_fokker_planck_results(results,time_cutoff)
For quick comparison I have also attached the code here for the classical solution to the Fokker-Planck equation using FD method:
import matplotlib.pyplot as plt
import scipy
def fokker_planck_classical(bits, D, t_max,time_steps):
# Parameters
x_min, x_max = -10, 10 # Spatial domain
Nx = 2**bits # Number of spatial grid points
Nt = time_steps # Number of time steps
Dx = (x_max - x_min) / Nx # Spatial step
tau = t_max / Nt # Time step
#Discretized space and timea
t = np.linspace(0, t_max, Nt + 1) # Include final time
x_grid = np.linspace(x_min, x_max, Nx)
#Gaussian distributuiio
P = scipy.stats.norm.pdf(x_grid, scale=0.5)
#Finite difference solution
P_all = [P.copy()]
for _ in range(Nt):
P_new = P.copy()
P_new[1:-1] = P[1:-1] + D * tau / Dx**2 * (P[2:] - 2*P[1:-1] + P[:-2])
P = P_new /scipy.integrate.simps(P_new,x_grid)
P_all.append(P.copy())
return {
'x_grid': x_grid,
'classical_states': P_all,
'times': t
}
def plot_results(results, t_cut):
x_grid = results['x_grid']
states = results['classical_states']
times = results['times']
t_list = np.arange(0, t_cut, 0.4)
t_fixed = []
for t_target in t_list:
idx = np.argmin(np.abs(np.array(times) - t_target))
t_fixed.append(idx)
print(f"Target time: {t_target}, Using time: {times[idx]:.2f}")
for i, t in enumerate(t_fixed):
plt.plot(x_grid, states[t], 'o-',
label=f't = {t_list[i]:.2f}')
plt.title("Classical Finite Difference")
plt.xlabel('x')
plt.ylabel('P(x,t)')
plt.legend()
plt.grid(True)
plt.savefig('fokker_planck_classical.png', dpi=300)
plt.show()
results = fokker_planck_classical(
bits=8, # just match with qubits from quantum runs
D=0.5, #diffusion coefficient
t_max=2, #time cutoff
time_steps=500
)
time_cutoff = 2
plot_results(results, time_cutoff)