animation.py 71 KB

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