diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..3ebec44 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,103 @@ +# psPlotKit — Agent Instructions + +## Project Overview + +psPlotKit is a Python plotting/analysis toolkit for data generated by the **Parameter Sweep (PS) tool** and **Loop tool** — primarily HDF5 (`.h5`) and JSON files produced by WaterTAP/IDAES process simulations. It imports hierarchical simulation results, manages unit conversions, computes levelized costs, and produces publication-quality matplotlib figures. + +## Architecture — Two-Layer Pipeline + +``` +data_manager/ → data_plotter/ + PsDataImport FigureGenerator (low-level matplotlib wrapper) + PsData linePlotter, boxPlotter, BreakDownPlotter, MapPlotter + PsDataManager (dict) (each plotter consumes a PsDataManager) + PsCosting + PsExpression + PsDataExporter +``` + +1. **`PsDataImport`** — reads `.h5` / `.json`, auto-discovers directories (parameter sweep runs), indexes data keys, and produces `PsData` objects. +2. **`PsData`** — wraps a numpy array with physical units (via the `quantities` library), supports unit conversion (`to_units`), unit assignment (`assign_units`), feasibility masking, and label generation for matplotlib. +3. **`PsDataManager(dict)`** — central data store. Keys are `(directory_tuple, data_key)` tuples. Provides `load_data`, `register_data_key`, `reduce_data` (stacking/min-reduction across sweep axes), `normalize_data`, `eval_function`, `get_costing`, and data selection helpers. +4. **`PsCosting`** — computes CAPEX/OPEX breakdowns and levelized costs per device group from costing blocks in the simulation data. +5. **`PsExpression`** — evaluates mathematical expressions on data stored in `PsDataManager`. +6. **`PsDataExporter`** — exports data from `PsDataManager` to CSV, JSON, or HDF5. +7. **Plotters** (`linePlotter`, `BreakDownPlotter`, `boxPlotter`, `MapPlotter`) — high-level plotting classes that accept a `PsDataManager`, select data, and delegate rendering to `FigureGenerator`. +8. **`FigureGenerator`** — thin matplotlib wrapper providing `plot_line`, `plot_bar`, `plot_map`, `plot_area`, axis styling, colorbars, and figure save/show. + +## Key Conventions + +### Naming & Deprecation Pattern +Every class has a **PascalCase canonical name** and a **legacy camelCase/lowercase subclass alias** that logs a deprecation warning. Always use the canonical names: +- `PsData` (not `psData`), `PsDataManager` (not `psDataManager`) +- `FigureGenerator` (not `figureGenerator`), `BreakDownPlotter` (not `breakDownPlotter`) +- `PsCosting` (not `psCosting`), `PsExpression` (not `psExpression`) +- `PsDataExporter` (not `psDataExporter`) + +### Data Key Dictionaries +When specifying data to import, keys are passed as dicts with specific fields: +```python +{"filekey": "fs.costing.LCOW", "return_key": "LCOW", "units": "USD/m**3"} +# Optional fields: "assign_units", "conversion_factor", "directories" +``` + +### Unit Handling +Units flow through `quantities` (aliased as `qs`). Custom units (`USD`, `PPM`) are registered in `CustomUnits`. String unit normalization happens in `PsData._convert_string_unit` (e.g., `"gal"` → `"US_liquid_gallon"`, `"°C"` → `"*degC"`, `"USD/a"` → `"USD/year"`). + +### Logging +Every module creates its logger via `psPlotKit.util.logger.define_logger(__name__, "DisplayName")`. Use this pattern — do not use `logging.getLogger` directly. + +### Composite Tuple Keys +`PsDataManager` stores data under composite tuple keys built from directory labels and data keys (e.g., `(("erd_type", "pressure_exchanger"), "membrane_cost", "LCOW")`). The `add_data(dir_key, data_key, PsData)` method constructs these automatically. + +## Development Setup + +```bash +conda env create -f psplotkit.yml # Python 3.11, pytest, editable install +conda activate psPlotKit +``` + +Dependencies: `h5py`, `quantities`, `matplotlib`, `numpy`, `scipy`, `sigfig`, `pyyaml`. + +## Testing + +```bash +pytest src/psPlotKit/data_manager/tests/ +``` + +Tests live alongside the modules they cover (e.g., `data_manager/tests/`). Test fixtures load `.h5` files from the same test directory using `os.path.dirname(os.path.abspath(__file__))`. + +## File Layout + +- `src/psPlotKit/data_manager/` — data import, storage, costing logic, expressions, exporter +- `src/psPlotKit/data_plotter/` — all plotting classes and figure generation +- `src/psPlotKit/util/` — shared logger and filesystem helpers +- `src/psPlotKit/examples/` — runnable example scripts with sample data in `examples/data/` +- `docs/` — mkdocs documentation (API reference, guides, examples) +- `legacy_data_collator.py` — legacy monolithic class; do not extend, prefer the modular pipeline above + +## Common Tasks + +### Adding a New Data Key +1. Create data key dict with `filekey`, `return_key`, `units` +2. Register with `PsDataManager.register_data_key(data_key_dict)` +3. Load data with `PsDataManager.load_data()` + +### Adding a New Plotter +1. Create new plotter class in `src/psPlotKit/data_plotter/` +2. Inherit from base plotter pattern (accept `PsDataManager` in constructor) +3. Use `FigureGenerator` for low-level matplotlib operations +4. Add to `src/psPlotKit/data_plotter/__init__.py` + +### Adding a New Costing Package +1. Create new module in `src/psPlotKit/data_manager/costing_packages/` +2. Implement costing computation functions +3. Register in `PsCosting` or use directly + +## Important Notes + +- **Never modify `legacy_data_collator.py`** — it is deprecated +- **Always use canonical PascalCase class names** in new code +- **Unit handling** must go through `quantities` library +- **Logging** must use `define_logger` from `psPlotKit.util.logger` +- **Tests** should be placed in `tests/` subdirectory alongside the module \ No newline at end of file diff --git a/src/psPlotKit/data_manager/costing_packages/tests/test_watertap_costing.py b/src/psPlotKit/data_manager/costing_packages/tests/test_watertap_costing.py index bcdf43f..8a15230 100644 --- a/src/psPlotKit/data_manager/costing_packages/tests/test_watertap_costing.py +++ b/src/psPlotKit/data_manager/costing_packages/tests/test_watertap_costing.py @@ -299,7 +299,9 @@ def test_complex(): # lets create costing pacakage package = WaterTapCostingPackage( - costing_block="costing", validation_key="costing.LCOW" + costing_block="costing", + validation_key="costing.LCOW", + sec_validation_key="costing.SEC", ) package.register_product_flow("avg_product_flow_rate") diff --git a/src/psPlotKit/data_manager/costing_packages/watertap_costing.py b/src/psPlotKit/data_manager/costing_packages/watertap_costing.py index f99552d..6a8220e 100644 --- a/src/psPlotKit/data_manager/costing_packages/watertap_costing.py +++ b/src/psPlotKit/data_manager/costing_packages/watertap_costing.py @@ -2,7 +2,12 @@ class WaterTapCostingPackage(PsCostingPackage): - def __init__(self, costing_block="fs.costing", validation_key="fs.costing.LCOW"): + def __init__( + self, + costing_block="fs.costing", + validation_key="fs.costing.LCOW", + sec_validation_key="fs.costing.SEC", + ): super().__init__(costing_block=costing_block) self.costing_package_name = "watertap" @@ -41,7 +46,9 @@ def __init__(self, costing_block="fs.costing", validation_key="fs.costing.LCOW") lambda ek: ek.total_capital_cost * ek.capital_recovery_factor + ek.total_operating_cost, ) + self.add_validation("LCOW", file_key=validation_key, rtol=1e-4) + self.add_validation("SEC", file_key=sec_validation_key, rtol=1e-4) self.register_fraction("LCOW") self.register_fraction("LCOW_opex", "LCOW") self.register_fraction("LCOW_capex", "LCOW") @@ -50,6 +57,11 @@ def register_product_flow( self, file_key="fs.product.properties[0.0].flow_vol_phase[Liq]" ): self.add_parameter("product_flow", file_key=file_key) + self.add_formula( + "SEC", + lambda ek: (ek.aggregate_electricity_flow) / (ek.product_flow), + units="kWh/m**3", + ) self.add_formula( "LCOW", lambda ek: (ek.total_annualized_cost) diff --git a/src/psPlotKit/data_manager/ps_costing.py b/src/psPlotKit/data_manager/ps_costing.py index 182a5e0..217af51 100644 --- a/src/psPlotKit/data_manager/ps_costing.py +++ b/src/psPlotKit/data_manager/ps_costing.py @@ -52,6 +52,7 @@ """ import difflib +import time import numpy as np @@ -526,17 +527,18 @@ def build(self, load_data=True, error_on_validation_failure=True): after key registration. Set to *False* if data is already loaded. """ + + ts = time.time() self.data_manager._load_registered_data_files() self._discover_keys() self._check_discovery_status() self._register_parameters() self._register_discovered_keys() self._register_validation_keys() - if load_data: self.data_manager.load_data(evaluate_expressions=False) - self._register_zero_sentinel() + self._build_group_expressions() self._build_flow_expressions() self._build_per_group_flow_expressions() @@ -544,16 +546,14 @@ def build(self, load_data=True, error_on_validation_failure=True): self._build_per_group_formula_expressions() self._build_total_formula_expressions() self._build_fraction_expressions() - self.data_manager.evaluate_expressions() - _logger.info( - "Costing build complete - {} group(s), {} formula(e).".format( + "Costing build complete took {} seconds - {} group(s), {} formula(e).".format( + round(time.time() - ts, 2), len(self.costing_groups), len(self.costing_package.formulae), ) ) - self._validate(error_on_failure=error_on_validation_failure) # ------------------------------------------------------------------ @@ -1166,12 +1166,17 @@ def _build_per_group_formula_expressions(self): flow_type in self.costing_package.flow_costs and group_flow_types.get(flow_type) ) + alias_map["{}_flow_cost".format(flow_type)] = ( gft_rk if has_cost else None ) alias_map["aggregate_flow_cost_{}".format(flow_type)] = ( gft_rk if has_cost else None ) + gft_rk = ("costing", group.name, "{}_flow".format(flow_type)) + alias_map["aggregate_{}_flow".format(flow_type)] = ( + gft_rk if has_cost else None + ) for fdef in self.costing_package.formulae: ek = self.data_manager.get_expression_keys() @@ -1199,7 +1204,7 @@ def _build_per_group_formula_expressions(self): self._per_group_formula_keys.setdefault(fdef["return_key"], []).append( group_rk ) - _logger.info("Registered per-group formula '{}'.".format(group_rk)) + _logger.debug("Registered per-group formula '{}'.".format(group_rk)) def _build_total_formula_expressions(self): """Build total (sum across groups) expressions for each formula. diff --git a/src/psPlotKit/data_manager/ps_data_manager.py b/src/psPlotKit/data_manager/ps_data_manager.py index ba43047..a053ccf 100644 --- a/src/psPlotKit/data_manager/ps_data_manager.py +++ b/src/psPlotKit/data_manager/ps_data_manager.py @@ -825,8 +825,15 @@ def sort_idxs(idxs): map_idxs = temp_map_idxs units = np.unique(map_units) if len(units) > 1: - _logger.info("Units are inconsistent, using dimensionless") - units = "dimensionless" + _map_units = [] + for k in range(len(map_units)): + true_sum = np.nansum(map_data[k]) + if true_sum != 0: + _map_units.append(map_units[k]) + units = np.unique(_map_units) + if len(units) > 1: + _logger.info("Units are inconsistent, using dimensionless") + units = "dimensionless" else: units = units[0] @@ -867,9 +874,10 @@ def sort_idxs(idxs): new_keys.append(data_key) new_keys.append(stack_keys) except: - _logger.error( - "Could not stack data for dir {}, key {}".format(uq, data_key) - ) + pass + # _logger.error( + # "Could not stack data for dir {}, key {}".format(uq, data_key) + # ) return unique_dirs, keys_to_process def add_mask(self, directory, indexes, data_shape=None, shape="1D"): diff --git a/src/psPlotKit/data_plotter/fig_generator.py b/src/psPlotKit/data_plotter/fig_generator.py index dd8be22..b7aa4a7 100644 --- a/src/psPlotKit/data_plotter/fig_generator.py +++ b/src/psPlotKit/data_plotter/fig_generator.py @@ -408,10 +408,15 @@ def init_figure( ncols=1, sharex=False, sharey=False, + panel_shape=None, + panel_size=None, twinx=False, twiny=False, grid=None, subplot_adjust=None, + wspace=None, + hspace=None, + constrained_layout=False, projection=None, **kwargs, ): @@ -425,15 +430,28 @@ def init_figure( ncols: Number of subplot columns. sharex: Whether subplots share the x-axis. sharey: Whether subplots share the y-axis. + panel_shape: Optional (rows, cols) tuple overriding nrows/ncols. + panel_size: Optional (width, height) in inches for each panel. + If set, overall figure size becomes + ``(panel_width * ncols, panel_height * nrows)``. twinx: If True, create a twin x-axis. twiny: If True, create a twin y-axis. grid: If set, overrides ncols and enables shared y-axis. subplot_adjust: Horizontal spacing between subplots. + wspace: Width space between subplots. + hspace: Height space between subplots. + constrained_layout: Enable matplotlib constrained layout. projection: Axes projection type (e.g., '3d'). """ + if panel_shape is not None: + nrows, ncols = panel_shape if grid is not None: sharey = True ncols = grid + if panel_size is not None: + panel_width, panel_height = panel_size + width = panel_width * ncols + height = panel_height * nrows self.projection = projection if projection == None: self.mode_3d = False @@ -442,6 +460,7 @@ def init_figure( ncols, sharex=sharex, sharey=sharey, + constrained_layout=constrained_layout, ) else: self.mode_3d = True @@ -452,6 +471,7 @@ def init_figure( sharex=sharex, sharey=sharey, subplot_kw={"projection": projection}, + constrained_layout=constrained_layout, ) self.idx_totals = (nrows, ncols) self.sharex = sharex @@ -470,12 +490,110 @@ def init_figure( if twiny: self.ax = [self.ax[0], self.ax[0].twinx()] self.twiny = True - if subplot_adjust is not None: - self.fig.subplots_adjust(wspace=subplot_adjust) + + # Default subplot spacing policy: + # - shared axes: keep panels close. + # - non-shared axes: leave larger gaps so axis labels do not collide. + auto_wspace = None + auto_hspace = None + if not constrained_layout: + if sharey and ncols > 1 and wspace is None and subplot_adjust is None: + auto_wspace = 0.05 + elif not sharey and ncols > 1 and wspace is None and subplot_adjust is None: + auto_wspace = 0.35 + if sharex and nrows > 1 and hspace is None: + auto_hspace = 0.05 + elif not sharex and nrows > 1 and hspace is None: + auto_hspace = 0.45 + + final_wspace = wspace + if final_wspace is None and subplot_adjust is not None: + final_wspace = subplot_adjust + if final_wspace is None: + final_wspace = auto_wspace + + final_hspace = hspace if hspace is not None else auto_hspace + + if not constrained_layout and ( + final_wspace is not None or final_hspace is not None + ): + self.fig.subplots_adjust( + wspace=( + final_wspace + if final_wspace is not None + else self.fig.subplotpars.wspace + ), + hspace=( + final_hspace + if final_hspace is not None + else self.fig.subplotpars.hspace + ), + ) self.fig.set_dpi(dpi) self.fig.set_size_inches(width, height, forward=True) return self + def init_panel( + self, + panel_shape=(1, 1), + panel_size=(3.25, 3.25), + dpi=150, + sharex=False, + sharey=False, + **kwargs, + ): + """Initialize a subplot panel layout using per-panel dimensions. + + Args: + panel_shape: (rows, cols) panel shape. + panel_size: (width, height) of each panel in inches. + dpi: Figure DPI. + sharex: Whether subplots share x-axis. + sharey: Whether subplots share y-axis. + + Returns: + Self for fluent chaining. + """ + nrows, ncols = panel_shape + panel_width, panel_height = panel_size + return self.init_figure( + nrows=nrows, + ncols=ncols, + width=panel_width * ncols, + height=panel_height * nrows, + dpi=dpi, + sharex=sharex, + sharey=sharey, + panel_shape=panel_shape, + panel_size=panel_size, + **kwargs, + ) + + def _normalize_ax_idx(self, idx): + """Normalize axis selection to the internal axis representation.""" + if hasattr(idx, "plot"): + return idx + if isinstance(idx, tuple): + return idx + if isinstance(idx, (int, np.integer)): + if self.idx_totals[0] > 1 and self.idx_totals[1] > 1: + ncols = self.idx_totals[1] + return (int(idx) // ncols, int(idx) % ncols) + return int(idx) + raise TypeError("Unsupported axis index type: {}".format(type(idx))) + + def _iter_axes(self): + """Yield all axes in row-major order.""" + if self.idx_totals[0] > 1 and self.idx_totals[1] > 1: + for row in self.ax: + for axis in row: + yield axis + elif self.idx_totals[0] == 1 and self.idx_totals[1] == 1: + yield self.ax[0] + else: + for axis in self.ax: + yield axis + def get_color(self, ax, val_update=0): """Get the current color index for the given axis and optionally advance it. @@ -486,12 +604,13 @@ def get_color(self, ax, val_update=0): Returns: Integer color index into the current colormap. """ + norm_ax = self._normalize_ax_idx(ax) if self.idx_totals[0] > 1 and self.idx_totals[1] > 1: - self.current_color_index[ax[0]][ax[1]] += val_update - return int(self.current_color_index[ax[0]][ax[1]]) + self.current_color_index[norm_ax[0]][norm_ax[1]] += val_update + return int(self.current_color_index[norm_ax[0]][norm_ax[1]]) else: - self.current_color_index[ax] += val_update - return int(self.current_color_index[ax]) + self.current_color_index[norm_ax] += val_update + return int(self.current_color_index[norm_ax]) def plot_bar( self, @@ -754,7 +873,7 @@ def plot_line( zorder=zorder, markersize=markersize, ) - for value, marker in marker_dict.items(): + for i, (value, marker) in enumerate(marker_dict.items()): plot_range = np.where(marker_overlay == value)[0] if len(plot_range) > 0: @@ -2096,6 +2215,130 @@ def add_legend( bbox_to_anchor=bbox_to_anchor, ) + def add_shared_legend( + self, + loc="top", + fontsize=9, + ncol=None, + bbox_to_anchor=None, + deduplicate=True, + reverse_legend=False, + ax_indices=None, + **kwargs, + ): + """Add a figure-level legend shared across multiple subplots. + + Args: + loc: Shared legend placement key. One of ``top``, ``bottom``, + ``right``, or any matplotlib legend location string. + fontsize: Legend font size. + ncol: Number of legend columns. Defaults to number of labels. + bbox_to_anchor: Explicit legend anchor override. + deduplicate: If True, remove duplicate labels while preserving + insertion order. + reverse_legend: If True, reverse final legend order. + ax_indices: Optional iterable of axis indices to include. + + Returns: + The created matplotlib Legend instance. + """ + if ax_indices is None: + axes = list(self._iter_axes()) + else: + axes = [self.get_axis(ax_idx) for ax_idx in ax_indices] + + handles = [] + labels = [] + for axis in axes: + axis_handles, axis_labels = axis.get_legend_handles_labels() + handles.extend(axis_handles) + labels.extend(axis_labels) + + if deduplicate: + unique_handles = [] + unique_labels = [] + seen = set() + for handle, label in zip(handles, labels): + if label not in seen and label != "": + seen.add(label) + unique_handles.append(handle) + unique_labels.append(label) + handles, labels = unique_handles, unique_labels + + if reverse_legend: + handles, labels = handles[::-1], labels[::-1] + + if ncol is None: + ncol = max(1, len(labels)) + + loc_map = { + "top": ("upper center", (0.5, 1.02)), + "bottom": ("lower center", (0.5, -0.02)), + "right": ("center left", (1.02, 0.5)), + } + + if loc in loc_map: + mpl_loc, default_anchor = loc_map[loc] + anchor = bbox_to_anchor if bbox_to_anchor is not None else default_anchor + else: + mpl_loc = loc + anchor = bbox_to_anchor + + legend = self.fig.legend( + handles, + labels, + frameon=False, + loc=mpl_loc, + ncol=ncol, + prop={"size": fontsize}, + labelspacing=0.2, + columnspacing=0.8, + handlelength=1, + handleheight=1, + bbox_to_anchor=anchor, + ) + + if loc == "top" and bbox_to_anchor is None: + self._auto_adjust_top_shared_legend(legend) + + return legend + + def _auto_adjust_top_shared_legend(self, legend, gap=0.01, max_top=0.995): + """Place top shared legend close to panels, adjusting top margin if needed.""" + self.fig.canvas.draw() + renderer = self.fig.canvas.get_renderer() + legend_bbox = legend.get_window_extent(renderer=renderer).transformed( + self.fig.transFigure.inverted() + ) + legend_height = legend_bbox.height + + panel_top = max(axis.get_position().y1 for axis in self._iter_axes()) + max_panel_top = max_top - (legend_height + gap) + + if panel_top > max_panel_top: + self.fig.subplots_adjust(top=max_panel_top) + self.fig.canvas.draw() + panel_top = max(axis.get_position().y1 for axis in self._iter_axes()) + + target_anchor_y = panel_top + gap + legend_height + legend.set_bbox_to_anchor( + (0.5, target_anchor_y), transform=self.fig.transFigure + ) + self.fig.canvas.draw() + + # Renderers can slightly shift legend extents; correct once so the + # legend remains above the panel with the requested gap. + corrected_bbox = legend.get_window_extent( + renderer=self.fig.canvas.get_renderer() + ).transformed(self.fig.transFigure.inverted()) + current_panel_top = max(axis.get_position().y1 for axis in self._iter_axes()) + current_gap = corrected_bbox.y0 - current_panel_top + if current_gap < gap: + desired_panel_top = corrected_bbox.y0 - gap + self.fig.subplots_adjust( + top=max(0.1, min(desired_panel_top, self.fig.subplotpars.top)) + ) + def get_axis(self, idx): """Return the matplotlib axes object for the given index. @@ -2105,10 +2348,13 @@ def get_axis(self, idx): Returns: Matplotlib Axes object. """ + norm_idx = self._normalize_ax_idx(idx) + if hasattr(norm_idx, "plot"): + return norm_idx if self.idx_totals[0] > 1 and self.idx_totals[1] > 1: - return self.ax[idx[0], idx[1]] + return self.ax[norm_idx[0], norm_idx[1]] else: - return self.ax[idx] + return self.ax[norm_idx] def remove_ticks(self, ax_idx=0, y_axis=None, x_axis=None): """Hide tick marks and labels for the specified axes. diff --git a/src/psPlotKit/data_plotter/ps_break_down_plotter.py b/src/psPlotKit/data_plotter/ps_break_down_plotter.py index 33ae2da..0b9c4e5 100644 --- a/src/psPlotKit/data_plotter/ps_break_down_plotter.py +++ b/src/psPlotKit/data_plotter/ps_break_down_plotter.py @@ -137,7 +137,7 @@ def check_key_in_dir(self, udir, key): return True else: for di in d: - if isinstance(di, str): + if isinstance(di, (str, int, float)): if key == di: return True else: diff --git a/src/psPlotKit/data_plotter/tests/test_fig_generator.py b/src/psPlotKit/data_plotter/tests/test_fig_generator.py index 5656442..d745ad2 100644 --- a/src/psPlotKit/data_plotter/tests/test_fig_generator.py +++ b/src/psPlotKit/data_plotter/tests/test_fig_generator.py @@ -1041,3 +1041,140 @@ def test_set_axis_auto_ticks_from_data_storage(self): assert ylim[0] == pytest.approx(10.0) assert ylim[1] == pytest.approx(50.0) fig.close() + + +# --------------------------------------------------------------------------- +# Panel and shared legend +# --------------------------------------------------------------------------- + + +class TestPanelAndSharedLegend: + def test_init_panel_sets_expected_shape_and_size(self): + fig = FigureGenerator() + fig.init_panel( + panel_shape=(2, 3), panel_size=(2.0, 1.5), sharex=True, sharey=True + ) + + assert fig.idx_totals == (2, 3) + width, height = fig.fig.get_size_inches() + assert width == pytest.approx(6.0) + assert height == pytest.approx(3.0) + + # Flat and tuple indexing should resolve to the same axis. + assert fig.get_axis(4) is fig.get_axis((1, 1)) + fig.close() + + def test_flat_axis_index_works_on_multiplot_grid(self): + fig = FigureGenerator() + fig.init_figure(nrows=2, ncols=2) + + fig.plot_line([1, 2], [1, 2], label="A", ax_idx=0) + fig.plot_line([1, 2], [2, 3], label="B", ax_idx=3) + + ax00 = fig.get_axis((0, 0)) + ax11 = fig.get_axis((1, 1)) + assert len(ax00.lines) == 1 + assert len(ax11.lines) == 1 + fig.close() + + def test_add_shared_legend_top_deduplicates_labels(self): + fig = FigureGenerator() + fig.init_figure(nrows=1, ncols=2) + + fig.plot_line([1, 2], [1, 2], label="Series A", ax_idx=0) + fig.plot_line([1, 2], [2, 3], label="Series A", ax_idx=1) + fig.plot_line([1, 2], [3, 4], label="Series B", ax_idx=1) + + legend = fig.add_shared_legend(loc="top") + labels = [text.get_text() for text in legend.get_texts()] + + assert labels == ["Series A", "Series B"] + assert len(fig.fig.legends) == 1 + fig.close() + + def test_add_shared_legend_right_location(self): + fig = FigureGenerator() + fig.init_figure(nrows=1, ncols=2) + + fig.plot_line([1, 2], [1, 2], label="Left", ax_idx=0) + fig.plot_line([1, 2], [2, 3], label="Right", ax_idx=1) + + legend = fig.add_shared_legend(loc="right", ncol=1) + assert legend is not None + assert len(legend.get_texts()) == 2 + fig.close() + + def test_sharex_defaults_to_small_vertical_panel_gap(self): + fig = FigureGenerator() + fig.init_figure(nrows=3, ncols=1, sharex=True) + + ax0 = fig.get_axis(0) + ax1 = fig.get_axis(1) + vertical_gap = ax0.get_position().y0 - ax1.get_position().y1 + + assert vertical_gap < 0.08 + fig.close() + + def test_sharey_defaults_to_small_horizontal_panel_gap(self): + fig = FigureGenerator() + fig.init_figure(nrows=1, ncols=3, sharey=True) + + ax0 = fig.get_axis(0) + ax1 = fig.get_axis(1) + horizontal_gap = ax1.get_position().x0 - ax0.get_position().x1 + + assert horizontal_gap < 0.08 + fig.close() + + def test_add_shared_legend_top_auto_positions_close_to_panel(self): + fig = FigureGenerator() + fig.init_figure(nrows=1, ncols=2) + + # Many labels in one column force a tall legend. + for i in range(8): + fig.plot_line([1, 2], [i, i + 1], label=f"Series {i}", ax_idx=0) + + legend = fig.add_shared_legend(loc="top", ncol=1) + fig.fig.canvas.draw() + + renderer = fig.fig.canvas.get_renderer() + legend_bbox = legend.get_window_extent(renderer=renderer).transformed( + fig.fig.transFigure.inverted() + ) + panel_top = max(axis.get_position().y1 for axis in fig._iter_axes()) + gap = legend_bbox.y0 - panel_top + + assert gap >= 0 + assert gap < 0.04 + assert legend_bbox.y1 <= 1.0 + fig.close() + + def test_nonshared_vertical_labels_do_not_overlap_adjacent_panel(self): + fig = FigureGenerator() + fig.init_figure(nrows=2, ncols=1, sharex=False) + + fig.get_axis(0).set_xlabel("Very Long X Label For Top Panel") + fig.get_axis(1).set_xlabel("Bottom X Label") + fig.fig.canvas.draw() + + renderer = fig.fig.canvas.get_renderer() + top_xlabel_bbox = fig.get_axis(0).xaxis.label.get_window_extent(renderer) + bottom_axes_bbox = fig.get_axis(1).get_window_extent(renderer) + + assert not top_xlabel_bbox.overlaps(bottom_axes_bbox) + fig.close() + + def test_nonshared_horizontal_labels_do_not_overlap_adjacent_panel(self): + fig = FigureGenerator() + fig.init_figure(nrows=1, ncols=2, sharey=False) + + fig.get_axis(0).set_ylabel("Very Long Y Label For Left Panel") + fig.get_axis(1).set_ylabel("Right Y Label") + fig.fig.canvas.draw() + + renderer = fig.fig.canvas.get_renderer() + left_ylabel_bbox = fig.get_axis(0).yaxis.label.get_window_extent(renderer) + right_axes_bbox = fig.get_axis(1).get_window_extent(renderer) + + assert not left_ylabel_bbox.overlaps(right_axes_bbox) + fig.close()