I'm trying to animate a plot using matplotlib's FuncAnimation, however no frames of the animation are visible until the animation reaches the final frame. If I set repeat = True nothing is ever displayed. When I first run the code a matplotlib icon appears 
but nothing displays when I click on it until it shows me the final frame:
If I save the animation I see the animation display correctly so this leads me to think that my code is mostly correct so I hope this is a simple fix that I'm just missing.
Apologies if I'm dumping too much code but I'm not sure if there's anything that's not needed for the minimum reproducible example.
Here's the main code
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from quantum_custom.constants import spin_down, spin_up, H00, H11, H
import quantum_custom.walk as walk
class QuantumState:
def __init__(self, state):
self.state = state
#"coin flips"
max_N = 100 #this will be the final number of coin flips
positions = 2*max_N + 1
#initial conditions
initial_spin = spin_down
initial_position = np.zeros(positions)
initial_position[max_N] = 1
initial_state = np.kron(np.matmul(H, initial_spin), initial_position) #initial state is Hadamard acting on intial state, tensor product with the initial position
quantum_state = QuantumState(initial_state)
#plot the graph
fig, ax = plt.subplots()
plt.title("N = 0")
x = np.arange(positions)
line, = ax.plot([],[])
loc = range(0, positions, positions // 10)
plt.xticks(loc)
plt.xlim(0, positions)
plt.ylim((0, 1))
ax.set_xticklabels(range(-max_N, max_N + 1, positions // 10))
ax.set_xlabel("x")
ax.set_ylabel("Probability")
def init():
line.set_data([],[])
return line,
def update(N):
next_state = walk.flip_once(quantum_state.state, max_N)
probs = walk.get_prob(next_state, max_N)
quantum_state.state = next_state
start_index = N % 2 + 1
cleaned_probs = probs[start_index::2]
cleaned_x = x[start_index::2]
line.set_data(cleaned_x, cleaned_probs)
if cleaned_probs.max() != 0:
plt.ylim((0, cleaned_probs.max()))
plt.title(f"N = {N}")
return line,
anim = animation.FuncAnimation(
fig,
update,
frames = max_N + 1,
init_func = init,
interval = 20,
repeat = False,
blit = True,
)
anim.save("animated.gif", writer = "ffmpeg", fps = 15)
plt.show()
Here's my quantum_custom.constants module.
#define spin up and spin down vectors as standard basis
spin_up = np.array([1,0])
spin_down = np.array([0,1])
#define our Hadamard operator, H, in terms of ith, jth entries, Hij
H00 = np.outer(spin_up, spin_up)
H01 = np.outer(spin_up, spin_down)
H10 = np.outer(spin_down, spin_up)
H11 = np.outer(spin_down, spin_down)
H = (H00 + H01 + H10 - H11)/np.sqrt(2.0) #matrix representation of Hadamard gate in standard basis
Here's my quantum_custom.walk module.
import numpy as np
from quantum_custom.constants import H00, H11, H
#define walk operators
def walk_operator(max_N):
position_count = 2 * max_N + 1
shift_plus = np.roll(np.eye(position_count), 1, axis = 0)
shift_minus = np.roll(np.eye(position_count), -1, axis = 0)
step_operator = np.kron(H00, shift_plus) + np.kron(H11, shift_minus)
return step_operator.dot(np.kron(H, np.eye(position_count)))
def flip_once(state, max_N):
"""
Flips the Hadamard coin once and acts on the given state appropriately.
Returns the state after the Hadamard coin flip.
"""
walk_op = walk_operator(max_N)
next_state = walk_op.dot(state)
return next_state
def get_prob(state, max_N):
"""
For the given state, calculates the probability of being in any possible position.
Returns an array of probabilities.
"""
position_count = 2 * max_N + 1
prob = np.empty(position_count)
for k in range(position_count):
posn = np.zeros(position_count)
posn[k] = 1
posn_outer = np.outer(posn, posn)
alt_measurement_k = eye_kron(2, posn_outer)
proj = alt_measurement_k.dot(state)
prob[k] = proj.dot(proj.conjugate()).real
return prob
def eye_kron(eye_dim, mat):
"""
Speeds up the calculation of the tensor product of an identity matrix of dimension eye_dim with a given matrix.
This exploits the fact that majority of values in the resulting matrix will be zeroes apart from on the leading diagonal where we simply have copies of the given matrix.
Returns a matrix.
"""
mat_dim = len(mat)
result_dim = eye_dim * mat_dim #dimension of the resulting matrix
result = np.zeros((result_dim, result_dim))
result[0:mat_dim, 0:mat_dim] = mat
result[mat_dim:result_dim, mat_dim:result_dim] = mat
return result
I know that saving the animation is a solution but I'd really like to have the plot display just from running the code as opposed to having to save it. Thanks!
