diff --git a/timflow/calibrate.py b/timflow/calibrate.py index 64693f4..bcbbaee 100644 --- a/timflow/calibrate.py +++ b/timflow/calibrate.py @@ -1447,6 +1447,7 @@ def plot_transient_results( obs_kwargs: dict | None = None, model_kwargs: dict | None = None, sharey: bool = False, + units: dict | None = None, ) -> tuple[plt.Figure, np.ndarray]: """Plot modeled vs observed transient head time series. @@ -1455,6 +1456,11 @@ def plot_transient_results( method before calibration to inspect the initial fit, or after calibration to inspect the calibrated fit. + For a cross-section model (``ModelXsection``), subplots are ordered + by ``(x, layer)`` and titled with the observation's location and + layer. For other model types, where ``x`` alone is not a meaningful + ordering, subplots are ordered and titled by observation name instead. + Parameters ---------- tmin : float, optional @@ -1474,6 +1480,10 @@ def plot_transient_results( sharey : bool, optional If ``True``, all subplots share the same y-axis limits. Default is ``False``. + units : dict, optional + Optional units to append to the axis labels and, for + cross-section models, the per-subplot ``x`` value in the title, + e.g. ``{"head": "m", "time": "days", "x": "m"}``. Returns ------- @@ -1498,6 +1508,25 @@ def plot_transient_results( if n_obs == 0: raise ValueError("No transient observations to plot.") + # Cross-section models have a meaningful x-coordinate to sort/title + # on; for other model types (e.g. areal models) fall back to the + # observation name instead. + _model_for_type = ( + self.transient_model + if self.transient_model is not None + else self.steady_model + ) + is_xsection = type(_model_for_type).__name__ == "ModelXsection" + if is_xsection: + all_items = sorted(all_items, key=lambda item: (item[1].x, item[1].layer)) + else: + all_items = sorted(all_items, key=lambda item: item[0]) + + units = units or {} + head_unit = f" [{units['head']}]" if "head" in units else "" + time_unit = f" [{units['time']}]" if "time" in units else "" + x_unit = f" {units['x']}" if "x" in units else "" + # Default styling obs_kw: dict = {"color": "k", "marker": ".", "linestyle": "none"} obs_kw.update(obs_kwargs or {}) @@ -1515,7 +1544,7 @@ def plot_transient_results( # Create subplots fig, ax_array = plt.subplots( - n_obs, 1, sharex=True, sharey=sharey, figsize=figsize + n_obs, 1, sharex=True, sharey=sharey, figsize=figsize, layout="constrained" ) axes = np.atleast_1d(ax_array) @@ -1575,11 +1604,20 @@ def plot_transient_results( model_kw["color"] = f"C{i}" # cycle through colors for each subplot ax.plot(t_plot[mask], h_obs_plot[mask], label=obs_label, **obs_kw) ax.plot(t_plot[mask], h_mod[mask], **{**model_kw, "label": model_label}) - ax.set_ylabel("head") + + if is_xsection: + ax.set_title( + f"x: {obs.x:.1f}{x_unit}, layer: {obs.layer}", + loc="right", + fontsize="medium", + ) + else: + ax.set_title(f"layer: {obs.layer}", loc="right", fontsize="medium") + ax.set_ylabel(f"head{head_unit}") ax.grid(True) ax.legend(loc=(0, 1), frameon=False, ncol=2) ax.set_xlim(left=t_plot[mask][0], right=t_plot[mask][-1]) i += 1 - axes[-1].set_xlabel("time") - fig.tight_layout() + axes[-1].set_xlabel(f"time{time_unit}") + fig.align_ylabels(axes) return fig, axes