How to assign variable fig for animation function in Python 3

Question

I want to make an animation of multiple plots whose rendering evolves in time.

The files that I need are under the format, for example for one :

DD0043/DD0043. So I use the trick : f'{43:04}' to fill the zeros leading for each file (the files go from DD0000/DD0000 to DD0922/DD0922.

Here the script, warning, the plot is done with yt-project tool :

import yt
import os, sys
import numpy as np
from matplotlib.animation import FuncAnimation
from matplotlib import rc_context
from matplotlib import pyplot as plt

# animate must accept an integer frame number. We use the frame number
# to identify which dataset in the time series we want to load
def animate(i):
  plot._switch_ds(array_data[i])

# Number of files
numFiles = int(os.popen('ls -dl DD* | wc -l').read())

# Array for each data directory
array_data = np.array(numFiles)

for i in range(numFiles):
  data = yt.load('DD'+str(f'{i:04}')+'/DD'+str(f'{i:04}'))
  sc = yt.create_scene(data, lens_type='perspective')

  source = sc[0]

  source.set_field('density')
  source.set_log(True)

  # Set up the camera parameters: focus, width, resolution, and image orientation
  sc.camera.focus = ds.domain_center
  sc.camera.resolution = 1024
  sc.camera.north_vector = [0, 0, 1]
  sc.camera.position = [1.7, 1.7, 1.7]

  # You may need to adjust the alpha values to get an image with good contrast.
  # For the annotate_domain call, the fourth value in the color tuple is the
  # alpha value.
  sc.annotate_axes(alpha=.02)
  sc.annotate_domain(ds, color=[1, 1, 1, .01])

  text_string = "T = {} Gyr".format(float(array_data[i].current_time.to('Gyr')))

fig = plt.figure()
animation = FuncAnimation(fig, animate, frames=numFiles)

# Override matplotlib's defaults to get a nicer looking font
with rc_context({'mathtext.fontset': 'stix'}):
    animation.save('animation.mp4')

But at the execution, I get the following error :

923
Traceback (most recent call last):
  File "vol-annotated.py", line 52, in <module>
    animation.save('animation.mp4')
  File "/Users/fab/Library/Python/3.7/lib/python/site-packages/matplotlib/animation.py", line 1135, in save
    anim._init_draw()
  File "/Users/fab/Library/Python/3.7/lib/python/site-packages/matplotlib/animation.py", line 1743, in _init_draw
    self._draw_frame(next(self.new_frame_seq()))
StopIteration

I don't know if I do the things correctly, especially for the variable fig that I initialize with :

fig = plt.figure()

Actually, I am trying to adapt to my case this script which creates a movie :

make animation

i.e :

import yt
from matplotlib.animation import FuncAnimation
from matplotlib import rc_context

ts = yt.load('GasSloshingLowRes/sloshing_low_res_hdf5_plt_cnt_*')

plot = yt.SlicePlot(ts[0], 'z', 'density')
plot.set_zlim('density', 8e-29, 3e-26)

fig = plot.plots['density'].figure

# animate must accept an integer frame number. We use the frame number
# to identify which dataset in the time series we want to load
def animate(i):
    ds = ts[i]
    plot._switch_ds(ds)

animation = FuncAnimation(fig, animate, frames=len(ts))

# Override matplotlib's defaults to get a nicer looking font
with rc_context({'mathtext.fontset': 'stix'}):
    animation.save('animation.mp4')

UPDATE 1: I didn't find a way to use animation.save correctly to generate an animation: always this issue about the fig variable.

But I managed to generate all the images corresponding for each one to an output file DDxxxx/DDxxxx. I have proceeded like this:

import yt
import os, sys
import numpy as np
from matplotlib.animation import FuncAnimation
from matplotlib import rc_context

# Number of files
numFiles = int(os.popen('ls -dl DD* | wc -l').read())

# Loop to load input files
ts = []
for j in range(numFiles):
  ts = np.append(ts, yt.load('DD'+str(f'{j:04}')+'/DD'+str(f'{j:04}')))

plot = yt.SlicePlot(ts[0], 'z', 'density')
plot.set_zlim('density', 8e-29, 3e-26)

# create plotting figure
fig = plot.plots['density'].figure

# animate must accept an integer frame number. We use the frame number
# to identify which dataset in the time series we want to load
def animate(i):
  ds = ts[i]
  sc = yt.create_scene(ds, lens_type='perspective')

  source = sc[0]

  source.set_field('density')
  source.set_log(True)

  # Set up the camera parameters: focus, width, resolution, and image orientation
  sc.camera.focus = ds.domain_center
  sc.camera.resolution = 1024
  sc.camera.north_vector = [0, 0, 1]
  sc.camera.position = [1.7, 1.7, 1.7]

  # You may need to adjust the alpha values to get an image with good contrast.
  # For the annotate_domain call, the fourth value in the color tuple is the
  # alpha value.
  sc.annotate_axes(alpha=.02)
  sc.annotate_domain(ds, color=[1, 1, 1, .01])

  text_string = "T = {} Gyr".format(float(ds.current_time.to('Gyr')))

  ## Here the scene needs to be painted into my figure / plot. 
  sc.save('rendering_'+str(i)+'.png')

animation = FuncAnimation(fig, animate, frames=numFiles)

# Override matplotlib's defaults to get a nicer looking font
with rc_context({'mathtext.fontset': 'stix'}):
    animation.save('animation.mp4')

If I open a single .png, I get a correct image representing a 3D scene.

Unfortunately, the animation function is not working, I get just a 2D heatmap plot showing the density projected: I would like to get an animation of the 3D scene figures (rendering_xxx.png).

It seems that I have to use ffmpeg to generate this animation from the multiple .png image, excepted if I find a way to know how to use Python FuncAnimation function (included in yt library ? or in Python by default ?).

UPDATE 2: here an example of figure (a frame actually) of animation I would like to get (this is a figure which represents gas density inside a box, i.e. in 3D) :

Unfortunately, @NightTrain's script produces this kind of plot :

As you can see, I don't understand why I get a 2D heatmap with NightTrain's solution instead of a 3D scene.

Moreover, there is no animation in this 2D heatmap, the movie displays always this same figure.

UPDATE3 : the last solution suggested by @Night train produces the following error :

  Traceback (most recent call last):
      File "plot_3D_enzo_with_animation_LAST.py", line 30, in <module>
        plot = yt.SlicePlot(ts[0], 'z', 'density')
      File "/Users/henry/Library/Python/3.7/lib/python/site-packages/yt/data_objects/time_series.py", line 201, in __getitem__
        o = self._pre_outputs[key]
    IndexError: list index out of range

I don't understand why this error occurs.

Answer 1

If you could provide more information it would be easier to help. I fixed your code and it is running now. You also forgot to use the text_string variable. Since the array_data variable isn't used I removed it.

import yt
import os, sys
import numpy as np
from matplotlib.animation import FuncAnimation
from matplotlib import rc_context
from matplotlib import pyplot as plt

import pathlib
import glob

base_path = "enzo_tiny_cosmology"
paths = sorted(glob.glob(base_path + "/DD*/DD[0-9][0-9][0-9][0-9]"))
# paths = [x.joinpath(x.name).as_posix() for x in sorted(pathlib.Path(base_path).glob("DD*"))]

# Array for each data directory
# array_data = np.zeros(len(paths))
# array_data = [None for x in range(len(paths))]

ts = yt.load(paths)
# ts = yt.load(base_path + "/DD*/DD[0-9][0-9][0-9][0-9]")
# print(ts.outputs)

plot = yt.SlicePlot(ts[0], 'z', 'density')
fig = plot.plots['density'].figure

# animate must accept an integer frame number. We use the frame number
# to identify which dataset in the time series we want to load
def animate(i):

  data = ts[i]
  sc = yt.create_scene(data, lens_type='perspective')

  source = sc[0]

  source.set_field('density')
  source.set_log(True)

  # Set up the camera parameters: focus, width, resolution, and image orientation
  sc.camera.focus = data.domain_center
  sc.camera.resolution = 1024
  sc.camera.north_vector = [0, 0, 1]
  sc.camera.position = [1.7, 1.7, 1.7]

  # You may need to adjust the alpha values to get an image with good contrast.
  # For the annotate_domain call, the fourth value in the color tuple is the
  # alpha value.
  sc.annotate_axes(alpha=.02)
  sc.annotate_domain(data, color=[1, 1, 1, .01])

  text_string = "T = {} Gyr".format(float(data.current_time.to('Gyr')))

  plot._switch_ds(data)

animation = FuncAnimation(fig, animate, frames = len(paths))

# Override matplotlib's defaults to get a nicer looking font
with rc_context({'mathtext.fontset': 'stix'}):
    animation.save('animation.mp4')

Instead of counting the lines of ls -dlyou might want to use a python solution. which also lets you use the paths directly without contructing them later. You can use either pathlib or the os module.

import pathlib
import glob

base_path = "enzo_tiny_cosmology"
paths = sorted(glob.glob(base_path + "/DD*/DD[0-9][0-9][0-9][0-9]"))
paths = [x.joinpath(x.name).as_posix() for x in sorted(pathlib.Path(base_path).glob("DD*"))]

For testing I downloaded these datasets:

curl -sSO https://yt-project.org/data/enzo_tiny_cosmology.tar.gz
tar xzf enzo_tiny_cosmology.tar.gz

curl -sSO https://yt-project.org/data/GasSloshingLowRes.tar.gz
tar xzf GasSloshingLowRes.tar.gz

UPDATE:

If you want to save the rendered scenes as video you could e.g. use imageio or opencv:

import yt, glob, imageio

# animate must accept an integer frame number. We use the frame number
# to identify which dataset in the time series we want to load
def animate(data):
  sc = yt.create_scene(data, lens_type='perspective')

  source = sc[0]
  source.set_field('density')
  source.set_log(True)

  # Set up the camera parameters: focus, width, resolution, and image orientation
  sc.camera.focus = data.domain_center
  sc.camera.resolution = 1024
  sc.camera.north_vector = [0, 0, 1]
  sc.camera.position = [1.7, 1.7, 1.7]

  # You may need to adjust the alpha values to get an image with good contrast.
  # For the annotate_domain call, the fourth value in the color tuple is the
  # alpha value.
  sc.annotate_axes(alpha=.02)
  sc.annotate_domain(data, color=[1, 1, 1, .01])

  plot._switch_ds(data)
  sc.save(f'rendering_{i:04d}.png')
  return sc.render()

paths = sorted(glob.glob("/DD*/DD[0-9][0-9][0-9][0-9]"))
ts = yt.load(paths)
plot = yt.SlicePlot(ts[0], 'z', 'density')
plot.set_zlim('density', 8e-29, 3e-26)

vid_writer = imageio.get_writer("animation.mp4", fps = 10)
for frame in ts:
    rendered_image = animate(frame)
    vid_writer.append_data(rendered_image)
vid_writer.close()