tight_bbox.py 2.9 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788
  1. """
  2. Helper module for the *bbox_inches* parameter in `.Figure.savefig`.
  3. """
  4. from matplotlib.transforms import Bbox, TransformedBbox, Affine2D
  5. def adjust_bbox(fig, bbox_inches, fixed_dpi=None):
  6. """
  7. Temporarily adjust the figure so that only the specified area
  8. (bbox_inches) is saved.
  9. It modifies fig.bbox, fig.bbox_inches,
  10. fig.transFigure._boxout, and fig.patch. While the figure size
  11. changes, the scale of the original figure is conserved. A
  12. function which restores the original values are returned.
  13. """
  14. origBbox = fig.bbox
  15. origBboxInches = fig.bbox_inches
  16. orig_tight_layout = fig.get_tight_layout()
  17. _boxout = fig.transFigure._boxout
  18. fig.set_tight_layout(False)
  19. old_aspect = []
  20. locator_list = []
  21. sentinel = object()
  22. for ax in fig.axes:
  23. locator_list.append(ax.get_axes_locator())
  24. current_pos = ax.get_position(original=False).frozen()
  25. ax.set_axes_locator(lambda a, r, _pos=current_pos: _pos)
  26. # override the method that enforces the aspect ratio on the Axes
  27. if 'apply_aspect' in ax.__dict__:
  28. old_aspect.append(ax.apply_aspect)
  29. else:
  30. old_aspect.append(sentinel)
  31. ax.apply_aspect = lambda pos=None: None
  32. def restore_bbox():
  33. for ax, loc, aspect in zip(fig.axes, locator_list, old_aspect):
  34. ax.set_axes_locator(loc)
  35. if aspect is sentinel:
  36. # delete our no-op function which un-hides the original method
  37. del ax.apply_aspect
  38. else:
  39. ax.apply_aspect = aspect
  40. fig.bbox = origBbox
  41. fig.bbox_inches = origBboxInches
  42. fig.set_tight_layout(orig_tight_layout)
  43. fig.transFigure._boxout = _boxout
  44. fig.transFigure.invalidate()
  45. fig.patch.set_bounds(0, 0, 1, 1)
  46. if fixed_dpi is None:
  47. fixed_dpi = fig.dpi
  48. tr = Affine2D().scale(fixed_dpi)
  49. dpi_scale = fixed_dpi / fig.dpi
  50. _bbox = TransformedBbox(bbox_inches, tr)
  51. fig.bbox_inches = Bbox.from_bounds(0, 0,
  52. bbox_inches.width, bbox_inches.height)
  53. x0, y0 = _bbox.x0, _bbox.y0
  54. w1, h1 = fig.bbox.width * dpi_scale, fig.bbox.height * dpi_scale
  55. fig.transFigure._boxout = Bbox.from_bounds(-x0, -y0, w1, h1)
  56. fig.transFigure.invalidate()
  57. fig.bbox = TransformedBbox(fig.bbox_inches, tr)
  58. fig.patch.set_bounds(x0 / w1, y0 / h1,
  59. fig.bbox.width / w1, fig.bbox.height / h1)
  60. return restore_bbox
  61. def process_figure_for_rasterizing(fig, bbox_inches_restore, fixed_dpi=None):
  62. """
  63. A function that needs to be called when figure dpi changes during the
  64. drawing (e.g., rasterizing). It recovers the bbox and re-adjust it with
  65. the new dpi.
  66. """
  67. bbox_inches, restore_bbox = bbox_inches_restore
  68. restore_bbox()
  69. r = adjust_bbox(fig, bbox_inches, fixed_dpi)
  70. return bbox_inches, r