You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

1771 lines
66 KiB

4 years ago
  1. # TODO:
  2. # * Documentation -- this will need a new section of the User's Guide.
  3. # Both for Animations and just timers.
  4. # - Also need to update http://www.scipy.org/Cookbook/Matplotlib/Animations
  5. # * Blit
  6. # * Currently broken with Qt4 for widgets that don't start on screen
  7. # * Still a few edge cases that aren't working correctly
  8. # * Can this integrate better with existing matplotlib animation artist flag?
  9. # - If animated removes from default draw(), perhaps we could use this to
  10. # simplify initial draw.
  11. # * Example
  12. # * Frameless animation - pure procedural with no loop
  13. # * Need example that uses something like inotify or subprocess
  14. # * Complex syncing examples
  15. # * Movies
  16. # * Can blit be enabled for movies?
  17. # * Need to consider event sources to allow clicking through multiple figures
  18. import abc
  19. import base64
  20. import contextlib
  21. from io import BytesIO
  22. import itertools
  23. import logging
  24. import os
  25. from pathlib import Path
  26. import platform
  27. import shutil
  28. import subprocess
  29. import sys
  30. from tempfile import TemporaryDirectory
  31. import uuid
  32. import numpy as np
  33. from matplotlib._animation_data import (DISPLAY_TEMPLATE, INCLUDED_FRAMES,
  34. JS_INCLUDE)
  35. from matplotlib import cbook, rcParams, rcParamsDefault, rc_context
  36. _log = logging.getLogger(__name__)
  37. # Process creation flag for subprocess to prevent it raising a terminal
  38. # window. See for example:
  39. # https://stackoverflow.com/questions/24130623/using-python-subprocess-popen-cant-prevent-exe-stopped-working-prompt
  40. if platform.system() == 'Windows':
  41. subprocess_creation_flags = CREATE_NO_WINDOW = 0x08000000
  42. else:
  43. # Apparently None won't work here
  44. subprocess_creation_flags = 0
  45. # Other potential writing methods:
  46. # * http://pymedia.org/
  47. # * libmng (produces swf) python wrappers: https://github.com/libming/libming
  48. # * Wrap x264 API:
  49. # (http://stackoverflow.com/questions/2940671/
  50. # how-to-encode-series-of-images-into-h264-using-x264-api-c-c )
  51. def adjusted_figsize(w, h, dpi, n):
  52. '''Compute figure size so that pixels are a multiple of n
  53. Parameters
  54. ----------
  55. w, h : float
  56. Size in inches
  57. dpi : float
  58. The dpi
  59. n : int
  60. The target multiple
  61. Returns
  62. -------
  63. wnew, hnew : float
  64. The new figure size in inches.
  65. '''
  66. # this maybe simplified if / when we adopt consistent rounding for
  67. # pixel size across the whole library
  68. def correct_roundoff(x, dpi, n):
  69. if int(x*dpi) % n != 0:
  70. if int(np.nextafter(x, np.inf)*dpi) % n == 0:
  71. x = np.nextafter(x, np.inf)
  72. elif int(np.nextafter(x, -np.inf)*dpi) % n == 0:
  73. x = np.nextafter(x, -np.inf)
  74. return x
  75. wnew = int(w * dpi / n) * n / dpi
  76. hnew = int(h * dpi / n) * n / dpi
  77. return (correct_roundoff(wnew, dpi, n), correct_roundoff(hnew, dpi, n))
  78. # A registry for available MovieWriter classes
  79. class MovieWriterRegistry(object):
  80. '''Registry of available writer classes by human readable name.'''
  81. def __init__(self):
  82. self.avail = dict()
  83. self._registered = dict()
  84. self._dirty = False
  85. def set_dirty(self):
  86. """Sets a flag to re-setup the writers."""
  87. self._dirty = True
  88. def register(self, name):
  89. """Decorator for registering a class under a name.
  90. Example use::
  91. @registry.register(name)
  92. class Foo:
  93. pass
  94. """
  95. def wrapper(writerClass):
  96. self._registered[name] = writerClass
  97. if writerClass.isAvailable():
  98. self.avail[name] = writerClass
  99. return writerClass
  100. return wrapper
  101. def ensure_not_dirty(self):
  102. """If dirty, reasks the writers if they are available"""
  103. if self._dirty:
  104. self.reset_available_writers()
  105. def reset_available_writers(self):
  106. """Reset the available state of all registered writers"""
  107. self.avail = {}
  108. for name, writerClass in self._registered.items():
  109. if writerClass.isAvailable():
  110. self.avail[name] = writerClass
  111. self._dirty = False
  112. def list(self):
  113. '''Get a list of available MovieWriters.'''
  114. self.ensure_not_dirty()
  115. return list(self.avail)
  116. def is_available(self, name):
  117. '''Check if given writer is available by name.
  118. Parameters
  119. ----------
  120. name : str
  121. Returns
  122. -------
  123. available : bool
  124. '''
  125. self.ensure_not_dirty()
  126. return name in self.avail
  127. def __getitem__(self, name):
  128. self.ensure_not_dirty()
  129. if not self.avail:
  130. raise RuntimeError("No MovieWriters available!")
  131. try:
  132. return self.avail[name]
  133. except KeyError:
  134. raise RuntimeError(
  135. 'Requested MovieWriter ({}) not available'.format(name))
  136. writers = MovieWriterRegistry()
  137. class AbstractMovieWriter(abc.ABC):
  138. '''
  139. Abstract base class for writing movies. Fundamentally, what a MovieWriter
  140. does is provide is a way to grab frames by calling grab_frame().
  141. setup() is called to start the process and finish() is called afterwards.
  142. This class is set up to provide for writing movie frame data to a pipe.
  143. saving() is provided as a context manager to facilitate this process as::
  144. with moviewriter.saving(fig, outfile='myfile.mp4', dpi=100):
  145. # Iterate over frames
  146. moviewriter.grab_frame(**savefig_kwargs)
  147. The use of the context manager ensures that setup() and finish() are
  148. performed as necessary.
  149. An instance of a concrete subclass of this class can be given as the
  150. ``writer`` argument of `Animation.save()`.
  151. '''
  152. @abc.abstractmethod
  153. def setup(self, fig, outfile, dpi=None):
  154. '''
  155. Perform setup for writing the movie file.
  156. Parameters
  157. ----------
  158. fig: `matplotlib.figure.Figure` instance
  159. The figure object that contains the information for frames
  160. outfile: string
  161. The filename of the resulting movie file
  162. dpi: int, optional
  163. The DPI (or resolution) for the file. This controls the size
  164. in pixels of the resulting movie file. Default is ``fig.dpi``.
  165. '''
  166. @abc.abstractmethod
  167. def grab_frame(self, **savefig_kwargs):
  168. '''
  169. Grab the image information from the figure and save as a movie frame.
  170. All keyword arguments in savefig_kwargs are passed on to the `savefig`
  171. command that saves the figure.
  172. '''
  173. @abc.abstractmethod
  174. def finish(self):
  175. '''Finish any processing for writing the movie.'''
  176. @contextlib.contextmanager
  177. def saving(self, fig, outfile, dpi, *args, **kwargs):
  178. '''
  179. Context manager to facilitate writing the movie file.
  180. ``*args, **kw`` are any parameters that should be passed to `setup`.
  181. '''
  182. # This particular sequence is what contextlib.contextmanager wants
  183. self.setup(fig, outfile, dpi, *args, **kwargs)
  184. try:
  185. yield self
  186. finally:
  187. self.finish()
  188. class MovieWriter(AbstractMovieWriter):
  189. '''Base class for writing movies.
  190. This class is set up to provide for writing movie frame data to a pipe.
  191. See examples for how to use these classes.
  192. Attributes
  193. ----------
  194. frame_format : str
  195. The format used in writing frame data, defaults to 'rgba'
  196. fig : `~matplotlib.figure.Figure`
  197. The figure to capture data from.
  198. This must be provided by the sub-classes.
  199. '''
  200. def __init__(self, fps=5, codec=None, bitrate=None, extra_args=None,
  201. metadata=None):
  202. '''MovieWriter
  203. Parameters
  204. ----------
  205. fps: int
  206. Framerate for movie.
  207. codec: string or None, optional
  208. The codec to use. If ``None`` (the default) the ``animation.codec``
  209. rcParam is used.
  210. bitrate: int or None, optional
  211. The bitrate for the saved movie file, which is one way to control
  212. the output file size and quality. The default value is ``None``,
  213. which uses the ``animation.bitrate`` rcParam. A value of -1
  214. implies that the bitrate should be determined automatically by the
  215. underlying utility.
  216. extra_args: list of strings or None, optional
  217. A list of extra string arguments to be passed to the underlying
  218. movie utility. The default is ``None``, which passes the additional
  219. arguments in the ``animation.extra_args`` rcParam.
  220. metadata: Dict[str, str] or None
  221. A dictionary of keys and values for metadata to include in the
  222. output file. Some keys that may be of use include:
  223. title, artist, genre, subject, copyright, srcform, comment.
  224. '''
  225. self.fps = fps
  226. self.frame_format = 'rgba'
  227. if codec is None:
  228. self.codec = rcParams['animation.codec']
  229. else:
  230. self.codec = codec
  231. if bitrate is None:
  232. self.bitrate = rcParams['animation.bitrate']
  233. else:
  234. self.bitrate = bitrate
  235. if extra_args is None:
  236. self.extra_args = list(rcParams[self.args_key])
  237. else:
  238. self.extra_args = extra_args
  239. if metadata is None:
  240. self.metadata = dict()
  241. else:
  242. self.metadata = metadata
  243. @property
  244. def frame_size(self):
  245. '''A tuple ``(width, height)`` in pixels of a movie frame.'''
  246. w, h = self.fig.get_size_inches()
  247. return int(w * self.dpi), int(h * self.dpi)
  248. def _adjust_frame_size(self):
  249. if self.codec == 'h264':
  250. wo, ho = self.fig.get_size_inches()
  251. w, h = adjusted_figsize(wo, ho, self.dpi, 2)
  252. if not (wo, ho) == (w, h):
  253. self.fig.set_size_inches(w, h, forward=True)
  254. _log.info('figure size (inches) has been adjusted '
  255. 'from %s x %s to %s x %s', wo, ho, w, h)
  256. else:
  257. w, h = self.fig.get_size_inches()
  258. _log.debug('frame size in pixels is %s x %s', *self.frame_size)
  259. return w, h
  260. def setup(self, fig, outfile, dpi=None):
  261. '''
  262. Perform setup for writing the movie file.
  263. Parameters
  264. ----------
  265. fig : matplotlib.figure.Figure
  266. The figure object that contains the information for frames
  267. outfile : string
  268. The filename of the resulting movie file
  269. dpi : int, optional
  270. The DPI (or resolution) for the file. This controls the size
  271. in pixels of the resulting movie file. Default is fig.dpi.
  272. '''
  273. self.outfile = outfile
  274. self.fig = fig
  275. if dpi is None:
  276. dpi = self.fig.dpi
  277. self.dpi = dpi
  278. self._w, self._h = self._adjust_frame_size()
  279. # Run here so that grab_frame() can write the data to a pipe. This
  280. # eliminates the need for temp files.
  281. self._run()
  282. def _run(self):
  283. # Uses subprocess to call the program for assembling frames into a
  284. # movie file. *args* returns the sequence of command line arguments
  285. # from a few configuration options.
  286. command = self._args()
  287. output = subprocess.PIPE
  288. _log.info('MovieWriter.run: running command: %s', command)
  289. self._proc = subprocess.Popen(command, shell=False,
  290. stdout=output, stderr=output,
  291. stdin=subprocess.PIPE,
  292. creationflags=subprocess_creation_flags)
  293. def finish(self):
  294. '''Finish any processing for writing the movie.'''
  295. self.cleanup()
  296. def grab_frame(self, **savefig_kwargs):
  297. '''
  298. Grab the image information from the figure and save as a movie frame.
  299. All keyword arguments in savefig_kwargs are passed on to the `savefig`
  300. command that saves the figure.
  301. '''
  302. _log.debug('MovieWriter.grab_frame: Grabbing frame.')
  303. try:
  304. # re-adjust the figure size in case it has been changed by the
  305. # user. We must ensure that every frame is the same size or
  306. # the movie will not save correctly.
  307. self.fig.set_size_inches(self._w, self._h)
  308. # Tell the figure to save its data to the sink, using the
  309. # frame format and dpi.
  310. self.fig.savefig(self._frame_sink(), format=self.frame_format,
  311. dpi=self.dpi, **savefig_kwargs)
  312. except (RuntimeError, IOError) as e:
  313. out, err = self._proc.communicate()
  314. _log.info('MovieWriter -- Error running proc:\n%s\n%s', out, err)
  315. raise IOError('Error saving animation to file (cause: {0}) '
  316. 'Stdout: {1} StdError: {2}. It may help to re-run '
  317. 'with logging level set to '
  318. 'DEBUG.'.format(e, out, err))
  319. def _frame_sink(self):
  320. '''Returns the place to which frames should be written.'''
  321. return self._proc.stdin
  322. def _args(self):
  323. '''Assemble list of utility-specific command-line arguments.'''
  324. return NotImplementedError("args needs to be implemented by subclass.")
  325. def cleanup(self):
  326. '''Clean-up and collect the process used to write the movie file.'''
  327. out, err = self._proc.communicate()
  328. self._frame_sink().close()
  329. _log.debug('MovieWriter -- Command stdout:\n%s', out)
  330. _log.debug('MovieWriter -- Command stderr:\n%s', err)
  331. @classmethod
  332. def bin_path(cls):
  333. '''
  334. Returns the binary path to the commandline tool used by a specific
  335. subclass. This is a class method so that the tool can be looked for
  336. before making a particular MovieWriter subclass available.
  337. '''
  338. return str(rcParams[cls.exec_key])
  339. @classmethod
  340. def isAvailable(cls):
  341. '''
  342. Check to see if a MovieWriter subclass is actually available.
  343. '''
  344. return shutil.which(cls.bin_path()) is not None
  345. class FileMovieWriter(MovieWriter):
  346. '''`MovieWriter` for writing to individual files and stitching at the end.
  347. This must be sub-classed to be useful.
  348. '''
  349. def __init__(self, *args, **kwargs):
  350. MovieWriter.__init__(self, *args, **kwargs)
  351. self.frame_format = rcParams['animation.frame_format']
  352. def setup(self, fig, outfile, dpi=None, frame_prefix='_tmp',
  353. clear_temp=True):
  354. '''Perform setup for writing the movie file.
  355. Parameters
  356. ----------
  357. fig : matplotlib.figure.Figure
  358. The figure to grab the rendered frames from.
  359. outfile : str
  360. The filename of the resulting movie file.
  361. dpi : number, optional
  362. The dpi of the output file. This, with the figure size,
  363. controls the size in pixels of the resulting movie file.
  364. Default is fig.dpi.
  365. frame_prefix : str, optional
  366. The filename prefix to use for temporary files. Defaults to
  367. ``'_tmp'``.
  368. clear_temp : bool, optional
  369. If the temporary files should be deleted after stitching
  370. the final result. Setting this to ``False`` can be useful for
  371. debugging. Defaults to ``True``.
  372. '''
  373. self.fig = fig
  374. self.outfile = outfile
  375. if dpi is None:
  376. dpi = self.fig.dpi
  377. self.dpi = dpi
  378. self._adjust_frame_size()
  379. self.clear_temp = clear_temp
  380. self.temp_prefix = frame_prefix
  381. self._frame_counter = 0 # used for generating sequential file names
  382. self._temp_names = list()
  383. self.fname_format_str = '%s%%07d.%s'
  384. @property
  385. def frame_format(self):
  386. '''
  387. Format (png, jpeg, etc.) to use for saving the frames, which can be
  388. decided by the individual subclasses.
  389. '''
  390. return self._frame_format
  391. @frame_format.setter
  392. def frame_format(self, frame_format):
  393. if frame_format in self.supported_formats:
  394. self._frame_format = frame_format
  395. else:
  396. self._frame_format = self.supported_formats[0]
  397. def _base_temp_name(self):
  398. # Generates a template name (without number) given the frame format
  399. # for extension and the prefix.
  400. return self.fname_format_str % (self.temp_prefix, self.frame_format)
  401. def _frame_sink(self):
  402. # Creates a filename for saving using the basename and the current
  403. # counter.
  404. fname = self._base_temp_name() % self._frame_counter
  405. # Save the filename so we can delete it later if necessary
  406. self._temp_names.append(fname)
  407. _log.debug('FileMovieWriter.frame_sink: saving frame %d to fname=%s',
  408. self._frame_counter, fname)
  409. self._frame_counter += 1 # Ensures each created name is 'unique'
  410. # This file returned here will be closed once it's used by savefig()
  411. # because it will no longer be referenced and will be gc-ed.
  412. return open(fname, 'wb')
  413. def grab_frame(self, **savefig_kwargs):
  414. '''
  415. Grab the image information from the figure and save as a movie frame.
  416. All keyword arguments in savefig_kwargs are passed on to the `savefig`
  417. command that saves the figure.
  418. '''
  419. # Overloaded to explicitly close temp file.
  420. _log.debug('MovieWriter.grab_frame: Grabbing frame.')
  421. try:
  422. # Tell the figure to save its data to the sink, using the
  423. # frame format and dpi.
  424. with self._frame_sink() as myframesink:
  425. self.fig.savefig(myframesink, format=self.frame_format,
  426. dpi=self.dpi, **savefig_kwargs)
  427. except RuntimeError:
  428. out, err = self._proc.communicate()
  429. _log.info('MovieWriter -- Error running proc:\n%s\n%s', out, err)
  430. raise
  431. def finish(self):
  432. # Call run here now that all frame grabbing is done. All temp files
  433. # are available to be assembled.
  434. self._run()
  435. MovieWriter.finish(self) # Will call clean-up
  436. # Check error code for creating file here, since we just run
  437. # the process here, rather than having an open pipe.
  438. if self._proc.returncode:
  439. try:
  440. stdout = [s.decode() for s in self._proc._stdout_buff]
  441. stderr = [s.decode() for s in self._proc._stderr_buff]
  442. _log.info("MovieWriter.finish: stdout: %s", stdout)
  443. _log.info("MovieWriter.finish: stderr: %s", stderr)
  444. except Exception as e:
  445. pass
  446. raise RuntimeError('Error creating movie, return code: {}'
  447. .format(self._proc.returncode))
  448. def cleanup(self):
  449. MovieWriter.cleanup(self)
  450. # Delete temporary files
  451. if self.clear_temp:
  452. _log.debug('MovieWriter: clearing temporary fnames=%s',
  453. self._temp_names)
  454. for fname in self._temp_names:
  455. os.remove(fname)
  456. @writers.register('pillow')
  457. class PillowWriter(MovieWriter):
  458. @classmethod
  459. def isAvailable(cls):
  460. try:
  461. import PIL
  462. except ImportError:
  463. return False
  464. return True
  465. def __init__(self, *args, **kwargs):
  466. if kwargs.get("extra_args") is None:
  467. kwargs["extra_args"] = ()
  468. super().__init__(*args, **kwargs)
  469. def setup(self, fig, outfile, dpi=None):
  470. self._frames = []
  471. self._outfile = outfile
  472. self._dpi = dpi
  473. self._fig = fig
  474. def grab_frame(self, **savefig_kwargs):
  475. from PIL import Image
  476. buf = BytesIO()
  477. self._fig.savefig(buf, **dict(savefig_kwargs, format="rgba"))
  478. renderer = self._fig.canvas.get_renderer()
  479. # Using frombuffer / getbuffer may be slightly more efficient, but
  480. # Py3-only.
  481. self._frames.append(Image.frombytes(
  482. "RGBA",
  483. (int(renderer.width), int(renderer.height)),
  484. buf.getvalue()))
  485. def finish(self):
  486. self._frames[0].save(
  487. self._outfile, save_all=True, append_images=self._frames[1:],
  488. duration=int(1000 / self.fps))
  489. # Base class of ffmpeg information. Has the config keys and the common set
  490. # of arguments that controls the *output* side of things.
  491. class FFMpegBase(object):
  492. '''Mixin class for FFMpeg output.
  493. To be useful this must be multiply-inherited from with a
  494. `MovieWriterBase` sub-class.
  495. '''
  496. exec_key = 'animation.ffmpeg_path'
  497. args_key = 'animation.ffmpeg_args'
  498. @property
  499. def output_args(self):
  500. args = ['-vcodec', self.codec]
  501. # For h264, the default format is yuv444p, which is not compatible
  502. # with quicktime (and others). Specifying yuv420p fixes playback on
  503. # iOS,as well as HTML5 video in firefox and safari (on both Win and
  504. # OSX). Also fixes internet explorer. This is as of 2015/10/29.
  505. if self.codec == 'h264' and '-pix_fmt' not in self.extra_args:
  506. args.extend(['-pix_fmt', 'yuv420p'])
  507. # The %dk adds 'k' as a suffix so that ffmpeg treats our bitrate as in
  508. # kbps
  509. if self.bitrate > 0:
  510. args.extend(['-b', '%dk' % self.bitrate])
  511. if self.extra_args:
  512. args.extend(self.extra_args)
  513. for k, v in self.metadata.items():
  514. args.extend(['-metadata', '%s=%s' % (k, v)])
  515. return args + ['-y', self.outfile]
  516. @classmethod
  517. def isAvailable(cls):
  518. return (
  519. super().isAvailable()
  520. # Ubuntu 12.04 ships a broken ffmpeg binary which we shouldn't use.
  521. # NOTE: when removed, remove the same method in AVConvBase.
  522. and b'LibAv' not in subprocess.run(
  523. [cls.bin_path()], creationflags=subprocess_creation_flags,
  524. stdout=subprocess.DEVNULL, stderr=subprocess.PIPE).stderr)
  525. # Combine FFMpeg options with pipe-based writing
  526. @writers.register('ffmpeg')
  527. class FFMpegWriter(FFMpegBase, MovieWriter):
  528. '''Pipe-based ffmpeg writer.
  529. Frames are streamed directly to ffmpeg via a pipe and written in a single
  530. pass.
  531. '''
  532. def _args(self):
  533. # Returns the command line parameters for subprocess to use
  534. # ffmpeg to create a movie using a pipe.
  535. args = [self.bin_path(), '-f', 'rawvideo', '-vcodec', 'rawvideo',
  536. '-s', '%dx%d' % self.frame_size, '-pix_fmt', self.frame_format,
  537. '-r', str(self.fps)]
  538. # Logging is quieted because subprocess.PIPE has limited buffer size.
  539. # If you have a lot of frames in your animation and set logging to
  540. # DEBUG, you will have a buffer overrun.
  541. if _log.getEffectiveLevel() > logging.DEBUG:
  542. args += ['-loglevel', 'quiet']
  543. args += ['-i', 'pipe:'] + self.output_args
  544. return args
  545. # Combine FFMpeg options with temp file-based writing
  546. @writers.register('ffmpeg_file')
  547. class FFMpegFileWriter(FFMpegBase, FileMovieWriter):
  548. '''File-based ffmpeg writer.
  549. Frames are written to temporary files on disk and then stitched
  550. together at the end.
  551. '''
  552. supported_formats = ['png', 'jpeg', 'ppm', 'tiff', 'sgi', 'bmp',
  553. 'pbm', 'raw', 'rgba']
  554. def _args(self):
  555. # Returns the command line parameters for subprocess to use
  556. # ffmpeg to create a movie using a collection of temp images
  557. return [self.bin_path(), '-r', str(self.fps),
  558. '-i', self._base_temp_name(),
  559. '-vframes', str(self._frame_counter)] + self.output_args
  560. # Base class of avconv information. AVConv has identical arguments to FFMpeg.
  561. class AVConvBase(FFMpegBase):
  562. '''Mixin class for avconv output.
  563. To be useful this must be multiply-inherited from with a
  564. `MovieWriterBase` sub-class.
  565. '''
  566. exec_key = 'animation.avconv_path'
  567. args_key = 'animation.avconv_args'
  568. # NOTE : should be removed when the same method is removed in FFMpegBase.
  569. isAvailable = classmethod(MovieWriter.isAvailable.__func__)
  570. # Combine AVConv options with pipe-based writing
  571. @writers.register('avconv')
  572. class AVConvWriter(AVConvBase, FFMpegWriter):
  573. '''Pipe-based avconv writer.
  574. Frames are streamed directly to avconv via a pipe and written in a single
  575. pass.
  576. '''
  577. # Combine AVConv options with file-based writing
  578. @writers.register('avconv_file')
  579. class AVConvFileWriter(AVConvBase, FFMpegFileWriter):
  580. '''File-based avconv writer.
  581. Frames are written to temporary files on disk and then stitched
  582. together at the end.
  583. '''
  584. # Base class for animated GIFs with convert utility
  585. class ImageMagickBase(object):
  586. '''Mixin class for ImageMagick output.
  587. To be useful this must be multiply-inherited from with a
  588. `MovieWriterBase` sub-class.
  589. '''
  590. exec_key = 'animation.convert_path'
  591. args_key = 'animation.convert_args'
  592. @property
  593. def delay(self):
  594. return 100. / self.fps
  595. @property
  596. def output_args(self):
  597. return [self.outfile]
  598. @classmethod
  599. def _init_from_registry(cls):
  600. if sys.platform != 'win32' or rcParams[cls.exec_key] != 'convert':
  601. return
  602. import winreg
  603. for flag in (0, winreg.KEY_WOW64_32KEY, winreg.KEY_WOW64_64KEY):
  604. try:
  605. hkey = winreg.OpenKeyEx(winreg.HKEY_LOCAL_MACHINE,
  606. r'Software\Imagemagick\Current',
  607. 0, winreg.KEY_QUERY_VALUE | flag)
  608. binpath = winreg.QueryValueEx(hkey, 'BinPath')[0]
  609. winreg.CloseKey(hkey)
  610. break
  611. except Exception:
  612. binpath = ''
  613. if binpath:
  614. for exe in ('convert.exe', 'magick.exe'):
  615. path = os.path.join(binpath, exe)
  616. if os.path.exists(path):
  617. binpath = path
  618. break
  619. else:
  620. binpath = ''
  621. rcParams[cls.exec_key] = rcParamsDefault[cls.exec_key] = binpath
  622. @classmethod
  623. def isAvailable(cls):
  624. '''
  625. Check to see if a ImageMagickWriter is actually available.
  626. Done by first checking the windows registry (if applicable) and then
  627. running the commandline tool.
  628. '''
  629. bin_path = cls.bin_path()
  630. if bin_path == "convert":
  631. cls._init_from_registry()
  632. return super().isAvailable()
  633. # Note: the base classes need to be in that order to get
  634. # isAvailable() from ImageMagickBase called and not the
  635. # one from MovieWriter. The latter is then called by the
  636. # former.
  637. @writers.register('imagemagick')
  638. class ImageMagickWriter(ImageMagickBase, MovieWriter):
  639. '''Pipe-based animated gif.
  640. Frames are streamed directly to ImageMagick via a pipe and written
  641. in a single pass.
  642. '''
  643. def _args(self):
  644. return ([self.bin_path(),
  645. '-size', '%ix%i' % self.frame_size, '-depth', '8',
  646. '-delay', str(self.delay), '-loop', '0',
  647. '%s:-' % self.frame_format]
  648. + self.output_args)
  649. # Note: the base classes need to be in that order to get
  650. # isAvailable() from ImageMagickBase called and not the
  651. # one from MovieWriter. The latter is then called by the
  652. # former.
  653. @writers.register('imagemagick_file')
  654. class ImageMagickFileWriter(ImageMagickBase, FileMovieWriter):
  655. '''File-based animated gif writer.
  656. Frames are written to temporary files on disk and then stitched
  657. together at the end.
  658. '''
  659. supported_formats = ['png', 'jpeg', 'ppm', 'tiff', 'sgi', 'bmp',
  660. 'pbm', 'raw', 'rgba']
  661. def _args(self):
  662. return ([self.bin_path(), '-delay', str(self.delay), '-loop', '0',
  663. '%s*.%s' % (self.temp_prefix, self.frame_format)]
  664. + self.output_args)
  665. # Taken directly from jakevdp's JSAnimation package at
  666. # http://github.com/jakevdp/JSAnimation
  667. def _included_frames(frame_list, frame_format):
  668. """frame_list should be a list of filenames"""
  669. return INCLUDED_FRAMES.format(Nframes=len(frame_list),
  670. frame_dir=os.path.dirname(frame_list[0]),
  671. frame_format=frame_format)
  672. def _embedded_frames(frame_list, frame_format):
  673. """frame_list should be a list of base64-encoded png files"""
  674. template = ' frames[{0}] = "data:image/{1};base64,{2}"\n'
  675. return "\n" + "".join(
  676. template.format(i, frame_format, frame_data.replace('\n', '\\\n'))
  677. for i, frame_data in enumerate(frame_list))
  678. @writers.register('html')
  679. class HTMLWriter(FileMovieWriter):
  680. supported_formats = ['png', 'jpeg', 'tiff', 'svg']
  681. args_key = 'animation.html_args'
  682. @classmethod
  683. def isAvailable(cls):
  684. return True
  685. def __init__(self, fps=30, codec=None, bitrate=None, extra_args=None,
  686. metadata=None, embed_frames=False, default_mode='loop',
  687. embed_limit=None):
  688. self.embed_frames = embed_frames
  689. self.default_mode = default_mode.lower()
  690. # Save embed limit, which is given in MB
  691. if embed_limit is None:
  692. self._bytes_limit = rcParams['animation.embed_limit']
  693. else:
  694. self._bytes_limit = embed_limit
  695. # Convert from MB to bytes
  696. self._bytes_limit *= 1024 * 1024
  697. if self.default_mode not in ['loop', 'once', 'reflect']:
  698. self.default_mode = 'loop'
  699. _log.warning("unrecognized default_mode: using 'loop'")
  700. super().__init__(fps, codec, bitrate, extra_args, metadata)
  701. def setup(self, fig, outfile, dpi, frame_dir=None):
  702. root, ext = os.path.splitext(outfile)
  703. if ext not in ['.html', '.htm']:
  704. raise ValueError("outfile must be *.htm or *.html")
  705. self._saved_frames = []
  706. self._total_bytes = 0
  707. self._hit_limit = False
  708. if not self.embed_frames:
  709. if frame_dir is None:
  710. frame_dir = root + '_frames'
  711. if not os.path.exists(frame_dir):
  712. os.makedirs(frame_dir)
  713. frame_prefix = os.path.join(frame_dir, 'frame')
  714. else:
  715. frame_prefix = None
  716. super().setup(fig, outfile, dpi, frame_prefix, clear_temp=False)
  717. def grab_frame(self, **savefig_kwargs):
  718. if self.embed_frames:
  719. # Just stop processing if we hit the limit
  720. if self._hit_limit:
  721. return
  722. f = BytesIO()
  723. self.fig.savefig(f, format=self.frame_format,
  724. dpi=self.dpi, **savefig_kwargs)
  725. imgdata64 = base64.encodebytes(f.getvalue()).decode('ascii')
  726. self._total_bytes += len(imgdata64)
  727. if self._total_bytes >= self._bytes_limit:
  728. _log.warning(
  729. "Animation size has reached %s bytes, exceeding the limit "
  730. "of %s. If you're sure you want a larger animation "
  731. "embedded, set the animation.embed_limit rc parameter to "
  732. "a larger value (in MB). This and further frames will be "
  733. "dropped.", self._total_bytes, self._bytes_limit)
  734. self._hit_limit = True
  735. else:
  736. self._saved_frames.append(imgdata64)
  737. else:
  738. return super().grab_frame(**savefig_kwargs)
  739. def _run(self):
  740. # make a duck-typed subprocess stand in
  741. # this is called by the MovieWriter base class, but not used here.
  742. class ProcessStandin(object):
  743. returncode = 0
  744. def communicate(self):
  745. return '', ''
  746. self._proc = ProcessStandin()
  747. # save the frames to an html file
  748. if self.embed_frames:
  749. fill_frames = _embedded_frames(self._saved_frames,
  750. self.frame_format)
  751. Nframes = len(self._saved_frames)
  752. else:
  753. # temp names is filled by FileMovieWriter
  754. fill_frames = _included_frames(self._temp_names,
  755. self.frame_format)
  756. Nframes = len(self._temp_names)
  757. mode_dict = dict(once_checked='',
  758. loop_checked='',
  759. reflect_checked='')
  760. mode_dict[self.default_mode + '_checked'] = 'checked'
  761. interval = 1000 // self.fps
  762. with open(self.outfile, 'w') as of:
  763. of.write(JS_INCLUDE)
  764. of.write(DISPLAY_TEMPLATE.format(id=uuid.uuid4().hex,
  765. Nframes=Nframes,
  766. fill_frames=fill_frames,
  767. interval=interval,
  768. **mode_dict))
  769. class Animation(object):
  770. '''This class wraps the creation of an animation using matplotlib.
  771. It is only a base class which should be subclassed to provide
  772. needed behavior.
  773. This class is not typically used directly.
  774. Parameters
  775. ----------
  776. fig : matplotlib.figure.Figure
  777. The figure object that is used to get draw, resize, and any
  778. other needed events.
  779. event_source : object, optional
  780. A class that can run a callback when desired events
  781. are generated, as well as be stopped and started.
  782. Examples include timers (see :class:`TimedAnimation`) and file
  783. system notifications.
  784. blit : bool, optional
  785. controls whether blitting is used to optimize drawing. Defaults
  786. to ``False``.
  787. See Also
  788. --------
  789. FuncAnimation, ArtistAnimation
  790. '''
  791. def __init__(self, fig, event_source=None, blit=False):
  792. self._fig = fig
  793. # Disables blitting for backends that don't support it. This
  794. # allows users to request it if available, but still have a
  795. # fallback that works if it is not.
  796. self._blit = blit and fig.canvas.supports_blit
  797. # These are the basics of the animation. The frame sequence represents
  798. # information for each frame of the animation and depends on how the
  799. # drawing is handled by the subclasses. The event source fires events
  800. # that cause the frame sequence to be iterated.
  801. self.frame_seq = self.new_frame_seq()
  802. self.event_source = event_source
  803. # Instead of starting the event source now, we connect to the figure's
  804. # draw_event, so that we only start once the figure has been drawn.
  805. self._first_draw_id = fig.canvas.mpl_connect('draw_event', self._start)
  806. # Connect to the figure's close_event so that we don't continue to
  807. # fire events and try to draw to a deleted figure.
  808. self._close_id = self._fig.canvas.mpl_connect('close_event',
  809. self._stop)
  810. if self._blit:
  811. self._setup_blit()
  812. def _start(self, *args):
  813. '''
  814. Starts interactive animation. Adds the draw frame command to the GUI
  815. handler, calls show to start the event loop.
  816. '''
  817. # First disconnect our draw event handler
  818. self._fig.canvas.mpl_disconnect(self._first_draw_id)
  819. self._first_draw_id = None # So we can check on save
  820. # Now do any initial draw
  821. self._init_draw()
  822. # Add our callback for stepping the animation and
  823. # actually start the event_source.
  824. self.event_source.add_callback(self._step)
  825. self.event_source.start()
  826. def _stop(self, *args):
  827. # On stop we disconnect all of our events.
  828. if self._blit:
  829. self._fig.canvas.mpl_disconnect(self._resize_id)
  830. self._fig.canvas.mpl_disconnect(self._close_id)
  831. self.event_source.remove_callback(self._step)
  832. self.event_source = None
  833. def save(self, filename, writer=None, fps=None, dpi=None, codec=None,
  834. bitrate=None, extra_args=None, metadata=None, extra_anim=None,
  835. savefig_kwargs=None):
  836. """
  837. Save the animation as a movie file by drawing every frame.
  838. Parameters
  839. ----------
  840. filename : str
  841. The output filename, e.g., :file:`mymovie.mp4`.
  842. writer : :class:`MovieWriter` or str, optional
  843. A `MovieWriter` instance to use or a key that identifies a
  844. class to use, such as 'ffmpeg'. If ``None``, defaults to
  845. :rc:`animation.writer` = 'ffmpeg'.
  846. fps : number, optional
  847. Frames per second in the movie. Defaults to ``None``, which will use
  848. the animation's specified interval to set the frames per second.
  849. dpi : number, optional
  850. Controls the dots per inch for the movie frames. This combined with
  851. the figure's size in inches controls the size of the movie. If
  852. ``None``, defaults to :rc:`savefig.dpi`.
  853. codec : str, optional
  854. The video codec to be used. Not all codecs are supported
  855. by a given :class:`MovieWriter`. If ``None``, default to
  856. :rc:`animation.codec` = 'h264'.
  857. bitrate : number, optional
  858. Specifies the number of bits used per second in the compressed
  859. movie, in kilobits per second. A higher number means a higher
  860. quality movie, but at the cost of increased file size. If ``None``,
  861. defaults to :rc:`animation.bitrate` = -1.
  862. extra_args : list, optional
  863. List of extra string arguments to be passed to the underlying movie
  864. utility. If ``None``, defaults to :rc:`animation.extra_args`.
  865. metadata : Dict[str, str], optional
  866. Dictionary of keys and values for metadata to include in
  867. the output file. Some keys that may be of use include:
  868. title, artist, genre, subject, copyright, srcform, comment.
  869. extra_anim : list, optional
  870. Additional `Animation` objects that should be included
  871. in the saved movie file. These need to be from the same
  872. `matplotlib.figure.Figure` instance. Also, animation frames will
  873. just be simply combined, so there should be a 1:1 correspondence
  874. between the frames from the different animations.
  875. savefig_kwargs : dict, optional
  876. Is a dictionary containing keyword arguments to be passed
  877. on to the `savefig` command which is called repeatedly to
  878. save the individual frames.
  879. Notes
  880. -----
  881. *fps*, *codec*, *bitrate*, *extra_args* and *metadata* are used to
  882. construct a `.MovieWriter` instance and can only be passed if
  883. *writer* is a string. If they are passed as non-*None* and *writer*
  884. is a `.MovieWriter`, a `RuntimeError` will be raised.
  885. """
  886. # If the writer is None, use the rc param to find the name of the one
  887. # to use
  888. if writer is None:
  889. writer = rcParams['animation.writer']
  890. elif (not isinstance(writer, str) and
  891. any(arg is not None
  892. for arg in (fps, codec, bitrate, extra_args, metadata))):
  893. raise RuntimeError('Passing in values for arguments '
  894. 'fps, codec, bitrate, extra_args, or metadata '
  895. 'is not supported when writer is an existing '
  896. 'MovieWriter instance. These should instead be '
  897. 'passed as arguments when creating the '
  898. 'MovieWriter instance.')
  899. if savefig_kwargs is None:
  900. savefig_kwargs = {}
  901. # Need to disconnect the first draw callback, since we'll be doing
  902. # draws. Otherwise, we'll end up starting the animation.
  903. if self._first_draw_id is not None:
  904. self._fig.canvas.mpl_disconnect(self._first_draw_id)
  905. reconnect_first_draw = True
  906. else:
  907. reconnect_first_draw = False
  908. if fps is None and hasattr(self, '_interval'):
  909. # Convert interval in ms to frames per second
  910. fps = 1000. / self._interval
  911. # Re-use the savefig DPI for ours if none is given
  912. if dpi is None:
  913. dpi = rcParams['savefig.dpi']
  914. if dpi == 'figure':
  915. dpi = self._fig.dpi
  916. if codec is None:
  917. codec = rcParams['animation.codec']
  918. if bitrate is None:
  919. bitrate = rcParams['animation.bitrate']
  920. all_anim = [self]
  921. if extra_anim is not None:
  922. all_anim.extend(anim
  923. for anim
  924. in extra_anim if anim._fig is self._fig)
  925. # If we have the name of a writer, instantiate an instance of the
  926. # registered class.
  927. if isinstance(writer, str):
  928. if writer in writers.avail:
  929. writer = writers[writer](fps, codec, bitrate,
  930. extra_args=extra_args,
  931. metadata=metadata)
  932. else:
  933. _log.warning("MovieWriter {} unavailable. Trying to use {} "
  934. "instead.".format(writer, writers.list()[0]))
  935. try:
  936. writer = writers[writers.list()[0]](fps, codec, bitrate,
  937. extra_args=extra_args,
  938. metadata=metadata)
  939. except IndexError:
  940. raise ValueError("Cannot save animation: no writers are "
  941. "available. Please install ffmpeg to "
  942. "save animations.")
  943. _log.info('Animation.save using %s', type(writer))
  944. if 'bbox_inches' in savefig_kwargs:
  945. _log.warning("Warning: discarding the 'bbox_inches' argument in "
  946. "'savefig_kwargs' as it may cause frame size "
  947. "to vary, which is inappropriate for animation.")
  948. savefig_kwargs.pop('bbox_inches')
  949. # Create a new sequence of frames for saved data. This is different
  950. # from new_frame_seq() to give the ability to save 'live' generated
  951. # frame information to be saved later.
  952. # TODO: Right now, after closing the figure, saving a movie won't work
  953. # since GUI widgets are gone. Either need to remove extra code to
  954. # allow for this non-existent use case or find a way to make it work.
  955. with rc_context():
  956. if rcParams['savefig.bbox'] == 'tight':
  957. _log.info("Disabling savefig.bbox = 'tight', as it may cause "
  958. "frame size to vary, which is inappropriate for "
  959. "animation.")
  960. rcParams['savefig.bbox'] = None
  961. with writer.saving(self._fig, filename, dpi):
  962. for anim in all_anim:
  963. # Clear the initial frame
  964. anim._init_draw()
  965. for data in zip(*[a.new_saved_frame_seq() for a in all_anim]):
  966. for anim, d in zip(all_anim, data):
  967. # TODO: See if turning off blit is really necessary
  968. anim._draw_next_frame(d, blit=False)
  969. writer.grab_frame(**savefig_kwargs)
  970. # Reconnect signal for first draw if necessary
  971. if reconnect_first_draw:
  972. self._first_draw_id = self._fig.canvas.mpl_connect('draw_event',
  973. self._start)
  974. def _step(self, *args):
  975. '''
  976. Handler for getting events. By default, gets the next frame in the
  977. sequence and hands the data off to be drawn.
  978. '''
  979. # Returns True to indicate that the event source should continue to
  980. # call _step, until the frame sequence reaches the end of iteration,
  981. # at which point False will be returned.
  982. try:
  983. framedata = next(self.frame_seq)
  984. self._draw_next_frame(framedata, self._blit)
  985. return True
  986. except StopIteration:
  987. return False
  988. def new_frame_seq(self):
  989. """Return a new sequence of frame information."""
  990. # Default implementation is just an iterator over self._framedata
  991. return iter(self._framedata)
  992. def new_saved_frame_seq(self):
  993. """Return a new sequence of saved/cached frame information."""
  994. # Default is the same as the regular frame sequence
  995. return self.new_frame_seq()
  996. def _draw_next_frame(self, framedata, blit):
  997. # Breaks down the drawing of the next frame into steps of pre- and
  998. # post- draw, as well as the drawing of the frame itself.
  999. self._pre_draw(framedata, blit)
  1000. self._draw_frame(framedata)
  1001. self._post_draw(framedata, blit)
  1002. def _init_draw(self):
  1003. # Initial draw to clear the frame. Also used by the blitting code
  1004. # when a clean base is required.
  1005. pass
  1006. def _pre_draw(self, framedata, blit):
  1007. # Perform any cleaning or whatnot before the drawing of the frame.
  1008. # This default implementation allows blit to clear the frame.
  1009. if blit:
  1010. self._blit_clear(self._drawn_artists, self._blit_cache)
  1011. def _draw_frame(self, framedata):
  1012. # Performs actual drawing of the frame.
  1013. raise NotImplementedError('Needs to be implemented by subclasses to'
  1014. ' actually make an animation.')
  1015. def _post_draw(self, framedata, blit):
  1016. # After the frame is rendered, this handles the actual flushing of
  1017. # the draw, which can be a direct draw_idle() or make use of the
  1018. # blitting.
  1019. if blit and self._drawn_artists:
  1020. self._blit_draw(self._drawn_artists, self._blit_cache)
  1021. else:
  1022. self._fig.canvas.draw_idle()
  1023. # The rest of the code in this class is to facilitate easy blitting
  1024. def _blit_draw(self, artists, bg_cache):
  1025. # Handles blitted drawing, which renders only the artists given instead
  1026. # of the entire figure.
  1027. updated_ax = []
  1028. for a in artists:
  1029. # If we haven't cached the background for this axes object, do
  1030. # so now. This might not always be reliable, but it's an attempt
  1031. # to automate the process.
  1032. if a.axes not in bg_cache:
  1033. bg_cache[a.axes] = a.figure.canvas.copy_from_bbox(a.axes.bbox)
  1034. a.axes.draw_artist(a)
  1035. updated_ax.append(a.axes)
  1036. # After rendering all the needed artists, blit each axes individually.
  1037. for ax in set(updated_ax):
  1038. ax.figure.canvas.blit(ax.bbox)
  1039. def _blit_clear(self, artists, bg_cache):
  1040. # Get a list of the axes that need clearing from the artists that
  1041. # have been drawn. Grab the appropriate saved background from the
  1042. # cache and restore.
  1043. axes = {a.axes for a in artists}
  1044. for a in axes:
  1045. if a in bg_cache:
  1046. a.figure.canvas.restore_region(bg_cache[a])
  1047. def _setup_blit(self):
  1048. # Setting up the blit requires: a cache of the background for the
  1049. # axes
  1050. self._blit_cache = dict()
  1051. self._drawn_artists = []
  1052. self._resize_id = self._fig.canvas.mpl_connect('resize_event',
  1053. self._handle_resize)
  1054. self._post_draw(None, self._blit)
  1055. def _handle_resize(self, *args):
  1056. # On resize, we need to disable the resize event handling so we don't
  1057. # get too many events. Also stop the animation events, so that
  1058. # we're paused. Reset the cache and re-init. Set up an event handler
  1059. # to catch once the draw has actually taken place.
  1060. self._fig.canvas.mpl_disconnect(self._resize_id)
  1061. self.event_source.stop()
  1062. self._blit_cache.clear()
  1063. self._init_draw()
  1064. self._resize_id = self._fig.canvas.mpl_connect('draw_event',
  1065. self._end_redraw)
  1066. def _end_redraw(self, evt):
  1067. # Now that the redraw has happened, do the post draw flushing and
  1068. # blit handling. Then re-enable all of the original events.
  1069. self._post_draw(None, False)
  1070. self.event_source.start()
  1071. self._fig.canvas.mpl_disconnect(self._resize_id)
  1072. self._resize_id = self._fig.canvas.mpl_connect('resize_event',
  1073. self._handle_resize)
  1074. def to_html5_video(self, embed_limit=None):
  1075. """
  1076. Convert the animation to an HTML5 ``<video>`` tag.
  1077. This saves the animation as an h264 video, encoded in base64
  1078. directly into the HTML5 video tag. This respects the rc parameters
  1079. for the writer as well as the bitrate. This also makes use of the
  1080. ``interval`` to control the speed, and uses the ``repeat``
  1081. parameter to decide whether to loop.
  1082. Parameters
  1083. ----------
  1084. embed_limit : float, optional
  1085. Limit, in MB, of the returned animation. No animation is created
  1086. if the limit is exceeded.
  1087. Defaults to :rc:`animation.embed_limit` = 20.0.
  1088. Returns
  1089. -------
  1090. video_tag : str
  1091. An HTML5 video tag with the animation embedded as base64 encoded
  1092. h264 video.
  1093. If the *embed_limit* is exceeded, this returns the string
  1094. "Video too large to embed."
  1095. """
  1096. VIDEO_TAG = r'''<video {size} {options}>
  1097. <source type="video/mp4" src="data:video/mp4;base64,{video}">
  1098. Your browser does not support the video tag.
  1099. </video>'''
  1100. # Cache the rendering of the video as HTML
  1101. if not hasattr(self, '_base64_video'):
  1102. # Save embed limit, which is given in MB
  1103. if embed_limit is None:
  1104. embed_limit = rcParams['animation.embed_limit']
  1105. # Convert from MB to bytes
  1106. embed_limit *= 1024 * 1024
  1107. # Can't open a NamedTemporaryFile twice on Windows, so use a
  1108. # TemporaryDirectory instead.
  1109. with TemporaryDirectory() as tmpdir:
  1110. path = Path(tmpdir, "temp.m4v")
  1111. # We create a writer manually so that we can get the
  1112. # appropriate size for the tag
  1113. Writer = writers[rcParams['animation.writer']]
  1114. writer = Writer(codec='h264',
  1115. bitrate=rcParams['animation.bitrate'],
  1116. fps=1000. / self._interval)
  1117. self.save(str(path), writer=writer)
  1118. # Now open and base64 encode.
  1119. vid64 = base64.encodebytes(path.read_bytes())
  1120. vid_len = len(vid64)
  1121. if vid_len >= embed_limit:
  1122. _log.warning(
  1123. "Animation movie is %s bytes, exceeding the limit of %s. "
  1124. "If you're sure you want a large animation embedded, set "
  1125. "the animation.embed_limit rc parameter to a larger value "
  1126. "(in MB).", vid_len, embed_limit)
  1127. else:
  1128. self._base64_video = vid64.decode('ascii')
  1129. self._video_size = 'width="{}" height="{}"'.format(
  1130. *writer.frame_size)
  1131. # If we exceeded the size, this attribute won't exist
  1132. if hasattr(self, '_base64_video'):
  1133. # Default HTML5 options are to autoplay and display video controls
  1134. options = ['controls', 'autoplay']
  1135. # If we're set to repeat, make it loop
  1136. if hasattr(self, 'repeat') and self.repeat:
  1137. options.append('loop')
  1138. return VIDEO_TAG.format(video=self._base64_video,
  1139. size=self._video_size,
  1140. options=' '.join(options))
  1141. else:
  1142. return 'Video too large to embed.'
  1143. def to_jshtml(self, fps=None, embed_frames=True, default_mode=None):
  1144. """Generate HTML representation of the animation"""
  1145. if fps is None and hasattr(self, '_interval'):
  1146. # Convert interval in ms to frames per second
  1147. fps = 1000 / self._interval
  1148. # If we're not given a default mode, choose one base on the value of
  1149. # the repeat attribute
  1150. if default_mode is None:
  1151. default_mode = 'loop' if self.repeat else 'once'
  1152. if not hasattr(self, "_html_representation"):
  1153. # Can't open a NamedTemporaryFile twice on Windows, so use a
  1154. # TemporaryDirectory instead.
  1155. with TemporaryDirectory() as tmpdir:
  1156. path = Path(tmpdir, "temp.html")
  1157. writer = HTMLWriter(fps=fps,
  1158. embed_frames=embed_frames,
  1159. default_mode=default_mode)
  1160. self.save(str(path), writer=writer)
  1161. self._html_representation = path.read_text()
  1162. return self._html_representation
  1163. def _repr_html_(self):
  1164. '''IPython display hook for rendering.'''
  1165. fmt = rcParams['animation.html']
  1166. if fmt == 'html5':
  1167. return self.to_html5_video()
  1168. elif fmt == 'jshtml':
  1169. return self.to_jshtml()
  1170. class TimedAnimation(Animation):
  1171. ''':class:`Animation` subclass for time-based animation.
  1172. A new frame is drawn every *interval* milliseconds.
  1173. Parameters
  1174. ----------
  1175. fig : matplotlib.figure.Figure
  1176. The figure object that is used to get draw, resize, and any
  1177. other needed events.
  1178. interval : number, optional
  1179. Delay between frames in milliseconds. Defaults to 200.
  1180. repeat_delay : number, optional
  1181. If the animation in repeated, adds a delay in milliseconds
  1182. before repeating the animation. Defaults to ``None``.
  1183. repeat : bool, optional
  1184. Controls whether the animation should repeat when the sequence
  1185. of frames is completed. Defaults to ``True``.
  1186. blit : bool, optional
  1187. Controls whether blitting is used to optimize drawing. Defaults
  1188. to ``False``.
  1189. '''
  1190. def __init__(self, fig, interval=200, repeat_delay=None, repeat=True,
  1191. event_source=None, *args, **kwargs):
  1192. # Store the timing information
  1193. self._interval = interval
  1194. self._repeat_delay = repeat_delay
  1195. self.repeat = repeat
  1196. # If we're not given an event source, create a new timer. This permits
  1197. # sharing timers between animation objects for syncing animations.
  1198. if event_source is None:
  1199. event_source = fig.canvas.new_timer()
  1200. event_source.interval = self._interval
  1201. Animation.__init__(self, fig, event_source=event_source,
  1202. *args, **kwargs)
  1203. def _step(self, *args):
  1204. '''
  1205. Handler for getting events.
  1206. '''
  1207. # Extends the _step() method for the Animation class. If
  1208. # Animation._step signals that it reached the end and we want to
  1209. # repeat, we refresh the frame sequence and return True. If
  1210. # _repeat_delay is set, change the event_source's interval to our loop
  1211. # delay and set the callback to one which will then set the interval
  1212. # back.
  1213. still_going = Animation._step(self, *args)
  1214. if not still_going and self.repeat:
  1215. self._init_draw()
  1216. self.frame_seq = self.new_frame_seq()
  1217. if self._repeat_delay:
  1218. self.event_source.remove_callback(self._step)
  1219. self.event_source.add_callback(self._loop_delay)
  1220. self.event_source.interval = self._repeat_delay
  1221. return True
  1222. else:
  1223. return Animation._step(self, *args)
  1224. else:
  1225. return still_going
  1226. def _stop(self, *args):
  1227. # If we stop in the middle of a loop delay (which is relatively likely
  1228. # given the potential pause here, remove the loop_delay callback as
  1229. # well.
  1230. self.event_source.remove_callback(self._loop_delay)
  1231. Animation._stop(self)
  1232. def _loop_delay(self, *args):
  1233. # Reset the interval and change callbacks after the delay.
  1234. self.event_source.remove_callback(self._loop_delay)
  1235. self.event_source.interval = self._interval
  1236. self.event_source.add_callback(self._step)
  1237. Animation._step(self)
  1238. class ArtistAnimation(TimedAnimation):
  1239. '''Animation using a fixed set of `Artist` objects.
  1240. Before creating an instance, all plotting should have taken place
  1241. and the relevant artists saved.
  1242. Parameters
  1243. ----------
  1244. fig : matplotlib.figure.Figure
  1245. The figure object that is used to get draw, resize, and any
  1246. other needed events.
  1247. artists : list
  1248. Each list entry a collection of artists that represent what
  1249. needs to be enabled on each frame. These will be disabled for
  1250. other frames.
  1251. interval : number, optional
  1252. Delay between frames in milliseconds. Defaults to 200.
  1253. repeat_delay : number, optional
  1254. If the animation in repeated, adds a delay in milliseconds
  1255. before repeating the animation. Defaults to ``None``.
  1256. repeat : bool, optional
  1257. Controls whether the animation should repeat when the sequence
  1258. of frames is completed. Defaults to ``True``.
  1259. blit : bool, optional
  1260. Controls whether blitting is used to optimize drawing. Defaults
  1261. to ``False``.
  1262. '''
  1263. def __init__(self, fig, artists, *args, **kwargs):
  1264. # Internal list of artists drawn in the most recent frame.
  1265. self._drawn_artists = []
  1266. # Use the list of artists as the framedata, which will be iterated
  1267. # over by the machinery.
  1268. self._framedata = artists
  1269. TimedAnimation.__init__(self, fig, *args, **kwargs)
  1270. def _init_draw(self):
  1271. # Make all the artists involved in *any* frame invisible
  1272. figs = set()
  1273. for f in self.new_frame_seq():
  1274. for artist in f:
  1275. artist.set_visible(False)
  1276. artist.set_animated(self._blit)
  1277. # Assemble a list of unique figures that need flushing
  1278. if artist.get_figure() not in figs:
  1279. figs.add(artist.get_figure())
  1280. # Flush the needed figures
  1281. for fig in figs:
  1282. fig.canvas.draw_idle()
  1283. def _pre_draw(self, framedata, blit):
  1284. '''
  1285. Clears artists from the last frame.
  1286. '''
  1287. if blit:
  1288. # Let blit handle clearing
  1289. self._blit_clear(self._drawn_artists, self._blit_cache)
  1290. else:
  1291. # Otherwise, make all the artists from the previous frame invisible
  1292. for artist in self._drawn_artists:
  1293. artist.set_visible(False)
  1294. def _draw_frame(self, artists):
  1295. # Save the artists that were passed in as framedata for the other
  1296. # steps (esp. blitting) to use.
  1297. self._drawn_artists = artists
  1298. # Make all the artists from the current frame visible
  1299. for artist in artists:
  1300. artist.set_visible(True)
  1301. class FuncAnimation(TimedAnimation):
  1302. """
  1303. Makes an animation by repeatedly calling a function *func*.
  1304. Parameters
  1305. ----------
  1306. fig : matplotlib.figure.Figure
  1307. The figure object that is used to get draw, resize, and any
  1308. other needed events.
  1309. func : callable
  1310. The function to call at each frame. The first argument will
  1311. be the next value in *frames*. Any additional positional
  1312. arguments can be supplied via the *fargs* parameter.
  1313. The required signature is::
  1314. def func(frame, *fargs) -> iterable_of_artists
  1315. If ``blit == True``, *func* must return an iterable of all artists
  1316. that were modified or created. This information is used by the blitting
  1317. algorithm to determine which parts of the figure have to be updated.
  1318. The return value is unused if ``blit == False`` and may be omitted in
  1319. that case.
  1320. frames : iterable, int, generator function, or None, optional
  1321. Source of data to pass *func* and each frame of the animation
  1322. - If an iterable, then simply use the values provided. If the
  1323. iterable has a length, it will override the *save_count* kwarg.
  1324. - If an integer, then equivalent to passing ``range(frames)``
  1325. - If a generator function, then must have the signature::
  1326. def gen_function() -> obj
  1327. - If *None*, then equivalent to passing ``itertools.count``.
  1328. In all of these cases, the values in *frames* is simply passed through
  1329. to the user-supplied *func* and thus can be of any type.
  1330. init_func : callable, optional
  1331. A function used to draw a clear frame. If not given, the
  1332. results of drawing from the first item in the frames sequence
  1333. will be used. This function will be called once before the
  1334. first frame.
  1335. The required signature is::
  1336. def init_func() -> iterable_of_artists
  1337. If ``blit == True``, *init_func* must return an iterable of artists
  1338. to be re-drawn. This information is used by the blitting
  1339. algorithm to determine which parts of the figure have to be updated.
  1340. The return value is unused if ``blit == False`` and may be omitted in
  1341. that case.
  1342. fargs : tuple or None, optional
  1343. Additional arguments to pass to each call to *func*.
  1344. save_count : int, optional
  1345. The number of values from *frames* to cache.
  1346. interval : number, optional
  1347. Delay between frames in milliseconds. Defaults to 200.
  1348. repeat_delay : number, optional
  1349. If the animation in repeated, adds a delay in milliseconds
  1350. before repeating the animation. Defaults to *None*.
  1351. repeat : bool, optional
  1352. Controls whether the animation should repeat when the sequence
  1353. of frames is completed. Defaults to *True*.
  1354. blit : bool, optional
  1355. Controls whether blitting is used to optimize drawing. Note: when using
  1356. blitting any animated artists will be drawn according to their zorder.
  1357. However, they will be drawn on top of any previous artists, regardless
  1358. of their zorder. Defaults to *False*.
  1359. """
  1360. def __init__(self, fig, func, frames=None, init_func=None, fargs=None,
  1361. save_count=None, **kwargs):
  1362. if fargs:
  1363. self._args = fargs
  1364. else:
  1365. self._args = ()
  1366. self._func = func
  1367. # Amount of framedata to keep around for saving movies. This is only
  1368. # used if we don't know how many frames there will be: in the case
  1369. # of no generator or in the case of a callable.
  1370. self.save_count = save_count
  1371. # Set up a function that creates a new iterable when needed. If nothing
  1372. # is passed in for frames, just use itertools.count, which will just
  1373. # keep counting from 0. A callable passed in for frames is assumed to
  1374. # be a generator. An iterable will be used as is, and anything else
  1375. # will be treated as a number of frames.
  1376. if frames is None:
  1377. self._iter_gen = itertools.count
  1378. elif callable(frames):
  1379. self._iter_gen = frames
  1380. elif cbook.iterable(frames):
  1381. self._iter_gen = lambda: iter(frames)
  1382. if hasattr(frames, '__len__'):
  1383. self.save_count = len(frames)
  1384. else:
  1385. self._iter_gen = lambda: iter(range(frames))
  1386. self.save_count = frames
  1387. if self.save_count is None:
  1388. # If we're passed in and using the default, set save_count to 100.
  1389. self.save_count = 100
  1390. else:
  1391. # itertools.islice returns an error when passed a numpy int instead
  1392. # of a native python int (http://bugs.python.org/issue30537).
  1393. # As a workaround, convert save_count to a native python int.
  1394. self.save_count = int(self.save_count)
  1395. self._init_func = init_func
  1396. # Needs to be initialized so the draw functions work without checking
  1397. self._save_seq = []
  1398. TimedAnimation.__init__(self, fig, **kwargs)
  1399. # Need to reset the saved seq, since right now it will contain data
  1400. # for a single frame from init, which is not what we want.
  1401. self._save_seq = []
  1402. def new_frame_seq(self):
  1403. # Use the generating function to generate a new frame sequence
  1404. return self._iter_gen()
  1405. def new_saved_frame_seq(self):
  1406. # Generate an iterator for the sequence of saved data. If there are
  1407. # no saved frames, generate a new frame sequence and take the first
  1408. # save_count entries in it.
  1409. if self._save_seq:
  1410. # While iterating we are going to update _save_seq
  1411. # so make a copy to safely iterate over
  1412. self._old_saved_seq = list(self._save_seq)
  1413. return iter(self._old_saved_seq)
  1414. else:
  1415. if self.save_count is not None:
  1416. return itertools.islice(self.new_frame_seq(), self.save_count)
  1417. else:
  1418. frame_seq = self.new_frame_seq()
  1419. def gen():
  1420. try:
  1421. for _ in range(100):
  1422. yield next(frame_seq)
  1423. except StopIteration:
  1424. pass
  1425. else:
  1426. cbook.warn_deprecated(
  1427. "2.2", "FuncAnimation.save has truncated your "
  1428. "animation to 100 frames. In the future, no such "
  1429. "truncation will occur; please pass 'save_count' "
  1430. "accordingly.")
  1431. return gen()
  1432. def _init_draw(self):
  1433. # Initialize the drawing either using the given init_func or by
  1434. # calling the draw function with the first item of the frame sequence.
  1435. # For blitting, the init_func should return a sequence of modified
  1436. # artists.
  1437. if self._init_func is None:
  1438. self._draw_frame(next(self.new_frame_seq()))
  1439. else:
  1440. self._drawn_artists = self._init_func()
  1441. if self._blit:
  1442. if self._drawn_artists is None:
  1443. raise RuntimeError('The init_func must return a '
  1444. 'sequence of Artist objects.')
  1445. for a in self._drawn_artists:
  1446. a.set_animated(self._blit)
  1447. self._save_seq = []
  1448. def _draw_frame(self, framedata):
  1449. # Save the data for potential saving of movies.
  1450. self._save_seq.append(framedata)
  1451. # Make sure to respect save_count (keep only the last save_count
  1452. # around)
  1453. self._save_seq = self._save_seq[-self.save_count:]
  1454. # Call the func with framedata and args. If blitting is desired,
  1455. # func needs to return a sequence of any artists that were modified.
  1456. self._drawn_artists = self._func(framedata, *self._args)
  1457. if self._blit:
  1458. if self._drawn_artists is None:
  1459. raise RuntimeError('The animation function must return a '
  1460. 'sequence of Artist objects.')
  1461. self._drawn_artists = sorted(self._drawn_artists,
  1462. key=lambda x: x.get_zorder())
  1463. for a in self._drawn_artists:
  1464. a.set_animated(self._blit)