python animation.FuncAnimation error : object is not iterable

Question

I'm new to Python and now studying matplotlib to use animation function. I'm following qutip tutorial because of my study. But when I copied and pasted example code of qutip tutorial, it didn't work Error message was Axes3D object is not iterable.
So, I want to check the code I created but don't know whether the problem is my code of is other thing. I don't know what to do and want to know why the example code doesn't work.

This is the example code from the tutorial:

from qutip import *
from scipy import *
def qubit_integrate(w, theta, gamma1, gamma2, psi0, tlist):

    sx = sigmax(); sy = sigmay(); sz = sigmaz(); sm = sigmam()
    H = w * (cos(theta) * sz + sin(theta) * sx)
    c_op_list = []
    n_th = 0.5 # temperature
    rate = gamma1 * (n_th + 1)
    if rate > 0.0: c_op_list.append(sqrt(rate) * sm)
    rate = gamma1 * n_th
    if rate > 0.0: c_op_list.append(sqrt(rate) * sm.dag())
    rate = gamma2
    if rate > 0.0: c_op_list.append(sqrt(rate) * sz)

    output = mesolve(H, psi0, tlist, c_op_list, [sx, sy, sz])  
    return output.expect[0], output.expect[1], output.expect[2]

    w     = 1.0 * 2 * pi   # qubit angular frequency
    theta = 0.2 * pi       # qubit angle from sigma_z axis (toward sigma_x axis)
    gamma1 = 0.5      # qubit relaxation rate
    gamma2 = 0.2      # qubit dephasing rate

    a = 1.0
    psi0 = (a* basis(2,0) + (1-a)*basis(2,1))/(sqrt(a**2 + (1-a)**2))
    tlist = linspace(0,4,250)
    sx, sy, sz = qubit_integrate(w, theta, gamma1, gamma2, psi0, tlist)
from pylab import *
import matplotlib.animation as animation
from mpl_toolkits.mplot3d import Axes3D

fig = figure()
ax = Axes3D(fig,azim=-40,elev=30)
sphere = Bloch(axes=ax)

def animate(i):
    sphere.clear()
    sphere.add_vectors([np.sin(theta),0,np.cos(theta)])
    sphere.add_points([sx[:i+1],sy[:i+1],sz[:i+1]])
    sphere.make_sphere()
    return ax

def init():
    sphere.vector_color = ['r']
    return ax

ani = animation.FuncAnimation(fig, animate, np.arange(len(sx)),
                            init_func=init, blit=True, repeat=False)
ani.save('bloch_sphere.mp4', fps=20, clear_temp=True)

And this is my own code:

import numpy as np
import qutip as q
import scipy as sp

up=q.basis(2,0)
sx=q.sigmax()
sy=q.sigmay()
sz=q.sigmaz()
bz=0.
by=0.
bx=15.
w=np.pi/20

H=w*(sx*bx+sy*by+sz*bz)

def state(t):
    states=[q.Qobj.expm(-(0+1j)*H*t)*up]
    return states

import matplotlib.pyplot as plt
import matplotlib.animation as ani
from mpl_toolkits.mplot3d import Axes3D

fig=plt.figure()
ax=Axes3D(fig,azim=-40,elev=30)
sphere=q.Bloch(axes=ax)
sphere.add_states(up)

def ini():
    sphere.vector_color=("r")
    return ax

t=np.linspace(0,1,256)

def animate(i):
    sphere.clear()
    sphere.add_states(state[i])
    return ax

ani.FuncAnimation(fig,animate,frames=len(t),init_func=ini,blit=True,repeat=False)

plt.show()

Answer 1

Fixing the tutorial

Remove the blit argument to make the tutorial work:

ani = animation.FuncAnimation(fig, animate, np.arange(len(sx)),
                               init_func=init,  repeat=False)
plt.show()
ani.save('bloch_sphere.mp4', fps=20)

Fixing your example

frames must be iterable.

Change:

frames=len(t)

into:

frames=t

I.e., this line:

ani.FuncAnimation(fig, animate, frames=len(t), init_func=ini, blit=True, repeat=False)

should become this:

ani.FuncAnimation(fig, animate, frames=t, init_func=ini, blit=True, repeat=False)

A few more changes.

1. Call your function stare with parasynthesis state(i) not with square brackets state[i]
2. Keep a reference to the animation ani = ani.FuncAnimation

Full code:

def animate(i):
    sphere.clear()
    sphere.add_states(state(i))
    sphere.make_sphere()
    return ax

ani = ani.FuncAnimation(fig, animate, frames=t, init_func=ini, repeat=False)

plt.show()

This is the end state of the animation:

Answer 2

In addition to removing the blit=True argument, I also had to change the writer string value to "ffmpeg" (and of course installing the ffmpeg package, also installed the opencv package via pip but not sure it was required).

anim.save(name + '.mp4', fps=10, writer="ffmpeg", codec=codec)