Skip to content

Commit

Permalink
feat: cap max scaling iterations in yscale* fcns
Browse files Browse the repository at this point in the history
  • Loading branch information
andrzejnovak committed Apr 6, 2024
1 parent 5ef07a1 commit 1348573
Showing 1 changed file with 31 additions and 4 deletions.
35 changes: 31 additions & 4 deletions src/mplhep/plot.py
Original file line number Diff line number Diff line change
Expand Up @@ -992,7 +992,7 @@ def _draw_text_bbox(ax):
textboxes = [k for k in ax.get_children() if isinstance(k, AnchoredText)]
fig.canvas.draw()
if len(textboxes) > 1:
print("Warning: More than one textbox found")
logging.warning("More than one textbox found")
for box in textboxes:
if box.loc in [1, 2]:
bbox = box.get_tightbbox(fig.canvas.renderer)
Expand All @@ -1002,36 +1002,63 @@ def _draw_text_bbox(ax):
return bbox


def yscale_legend(ax=None, otol=0):
def yscale_legend(ax=None, otol=0, soft_fail=False):
"""
Automatically scale y-axis up to fit in legend().
Set `otol > 0` for less strict scaling.
Set `soft_fail=True` to return even if it could not fit the legend.
"""
if ax is None:
ax = plt.gca()

scale_factor = 10 ** (1.05) if ax.get_yscale() == "log" else 1.05
max_scales = 0
while overlap(ax, _draw_leg_bbox(ax)) > otol:
logging.debug(
f"Legend overlap with other artists is {overlap(ax, _draw_leg_bbox(ax))}."
)
logging.info("Scaling y-axis by 5% to fit legend")
ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[-1] * scale_factor)
ax.figure.canvas.draw()
if max_scales > 10:
if not soft_fail:
raise RuntimeError(
"Could not fit legend in 10 iterations, return anyway by passing `soft_fail=True`."
)
else:
logging.warning("Could not fit legend in 10 iterations")
break
max_scales += 1
return ax


def yscale_anchored_text(ax=None, otol=0):
def yscale_anchored_text(ax=None, otol=0, soft_fail=False):
"""
Automatically scale y-axis up to fit AnchoredText
Set `otol > 0` for less strict scaling.
Set `soft_fail=True` to return even if it could not fit the AnchoredText.
"""
if ax is None:
ax = plt.gca()

scale_factor = 10 ** (1.05) if ax.get_yscale() == "log" else 1.05
max_scales = 0
while overlap(ax, _draw_text_bbox(ax)) > otol:
ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[-1] * 1.1)
logging.debug(
f"AnchoredText overlap with other artists is {overlap(ax, _draw_text_bbox(ax))}."
)
logging.info("Scaling y-axis by 5% to fit legend")
ax.set_ylim(ax.get_ylim()[0], ax.get_ylim()[-1] * scale_factor)
ax.figure.canvas.draw()
if max_scales > 10:
if not soft_fail:
raise RuntimeError(
"Could not fit AnchoredText in 10 iterations, return anyway by passing `soft_fail=True`."
)
else:
logging.warning("Could not fit AnchoredText in 10 iterations")
break
max_scales += 1
return ax


Expand Down

0 comments on commit 1348573

Please sign in to comment.