Select the y-limits of the mu panels automatically

The fixed limits hid the markers of f'(mu*) and f'(mu**). The limits now
frame the finite features of each curve: its values at both ends of the
mu range, its local extrema, the markers and the reference lines, while
divergences at the poles still run off the axes.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

plot_mu_trajectories.py

@@ -35,6 +35,7 @@ N_TAIL = 100          # extra samples extending the outer branches geometrically
 EPS = 0.01            # gap left on each side of a pole
 MU_HALF_WIDTH = 3     # mu is sampled in [m - MU_HALF_WIDTH, m + MU_HALF_WIDTH], m the mean of the poles
 N_DOTTED_OUTER = 2    # points drawn dotted at the ends of the outer branches
+Y_PADDING = 0.1       # fraction of the y-range of the features added below and above them
 
 COLORMAP = 'RdYlBu_r'
 ELLIPSE_COLOR = 'royalblue'
@@ -103,6 +104,27 @@ def ellipse_boundary(Q, n=1000):
     return Q.d + np.sqrt(abs(Q.gamma)) * u @ Q.V.T
 
 
+def auto_ylim(mu, values, branches, xlim, extra=(), floor=None):
+    """Return y-limits framing the finite features of a curve plotted over xlim.
+
+    The features are the values at both ends of xlim, the local extrema of each branch and
+    `extra` (markers, reference lines). Divergences at the poles are not features: the curve
+    runs off the axes there.
+    """
+    visible = (mu >= xlim[0]) & (mu <= xlim[1])
+    features = [values[visible][0], values[visible][-1], *extra]
+    for branch in branches:
+        v = values[branch[visible[branch]]]
+        slopes = np.diff(v)
+        features.extend(v[1:-1][slopes[:-1] * slopes[1:] < 0])
+    low, high = min(features), max(features)
+    padding = Y_PADDING * (high - low) if high > low else 1
+    low -= padding
+    if floor is not None:
+        low = max(low, floor)
+    return low, high + padding
+
+
 def colored_line(ax, x, y, idx, cmap, norm, n_dotted=0):
     """Draw y(x) colored by idx, with the first and last n_dotted points joined by dotted segments."""
     points = np.column_stack((x, y))
@@ -159,12 +181,18 @@ def plot_trajectories(Q, x0):
     fig, axs = plt.subplots(2, 2, figsize=(10, 7))
     ax_x, ax_f, ax_dist, ax_df = axs.flat
 
+    # (point, direction, degenerate point or None) for each pole cancelled by a zero component of x0
+    inactive = [(*pole_line(Q, x0, k), degenerate_point(Q, x0, k)) for k in np.flatnonzero(~active)]
+    xd_dists = [np.linalg.norm(xd - x0) for _, _, xd in inactive if xd is not None]
+
+    # (axis, function, label, dotted points of the middle branch, extra values to show, lower bound)
     panels = [
-        (ax_f, F.f, r'$f(\mu)$', (-1, 1), 5),
-        (ax_dist, lambda m: F.dist(F.x_not_std(m)), r'$||\mathbf{x}^0 - \mathbf{x}(\mu)||_2$', (0, 2), 6),
-        (ax_df, F.d_f, r'$f^\prime(\mu)$', (-2, 2), 50),
+        (ax_f, F.f, r'$f(\mu)$', 5, [0], None),
+        (ax_dist, lambda m: F.dist(F.x_not_std(m)), r'$||\mathbf{x}^0 - \mathbf{x}(\mu)||_2$', 6,
+         [dist + 0.3 for dist in xd_dists], 0),  # green line and its label above it
+        (ax_df, F.d_f, r'$f^\prime(\mu)$', 50, [], None),
     ]
-    for ax, fun, ylabel, ylim, n_dotted_middle in panels:
+    for ax, fun, ylabel, n_dotted_middle, extra, floor in panels:
         values = np.array([fun(m) for m in mu])
         for i, branch in enumerate(branches):
             n_dotted = 0 if degenerate else (n_dotted_middle if i == 1 else N_DOTTED_OUTER)
@@ -173,6 +201,8 @@ def plot_trajectories(Q, x0):
             ax.axvline(pole, **POLE_STYLE)
         for root, style in markers:
             ax.scatter(root, fun(root), **style)
+        root_values = [fun(root) for root, _ in markers]
+        ylim = auto_ylim(mu, values, branches, xlim, extra + root_values, floor)
         ax.set(xlabel=r'$\mu$', ylabel=ylabel, xlim=xlim, ylim=ylim)
         ax.set_xticks(ticks, tick_labels)
     ax_f.axhline(0, color='red', alpha=0.5, zorder=1)
@@ -195,9 +225,7 @@ def plot_trajectories(Q, x0):
             point, direction = pole_line(Q, x0, k)
             ax_x.axline(point, point + direction, linestyle='dotted', color='grey')
 
-    for k in np.flatnonzero(~active):
-        point, direction = pole_line(Q, x0, k)
-        xd = degenerate_point(Q, x0, k)
+    for point, direction, xd in inactive:
         if xd is None:
             ax_x.plot(*np.column_stack((point - direction, point + direction)), color='grey', zorder=0)
             continue