15485406069

Committed 06 Jun 2025 07:39AM UTC coverage: 95.847% (+0.5%) from 95.307%

Build Type

push

github

Committed by ahobeost

Commit Message

Ci adjustments.

Run Details

10 of 10 new or added lines in 4 files covered. (100.0%)

10 existing lines in 2 files now uncovered.

1177 of 1228 relevant lines covered (95.85%)

1.91 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

99.45

/pytrnsys_process/plot/plotters.py

import typing as _tp
from abc import abstractmethod
from dataclasses import dataclass

import matplotlib.pyplot as _plt
import numpy as _np
import pandas as _pd

from pytrnsys_process import config as conf

# TODO: provide A4 and half A4 plots to test sizes in latex # pylint: disable=fixme
# TODO: provide height as input for plot?  # pylint: disable=fixme
# TODO: deal with legends (curve names, fonts, colors, linestyles) # pylint: disable=fixme
# TODO: clean up old stuff by refactoring # pylint: disable=fixme
# TODO: make issue for docstrings of plotting # pylint: disable=fixme
# TODO: Add colormap support # pylint: disable=fixme


# TODO find a better place for this to live in # pylint : disable=fixme
plot_settings = conf.global_settings.plot
"Settings shared by all plots"


class ChartBase:
    cmap: str | None = None

    def plot(
        self,
        df: _pd.DataFrame,
        columns: list[str],
        **kwargs,
    ) -> tuple[_plt.Figure, _plt.Axes]:
        fig, ax = self._do_plot(df, columns, **kwargs)
        return fig, ax

    @abstractmethod
    def _do_plot(
        self,
        df: _pd.DataFrame,
        columns: list[str],
        use_legend: bool = True,
        size: tuple[float, float] = conf.PlotSizes.A4.value,
        **kwargs: _tp.Any,
    ) -> tuple[_plt.Figure, _plt.Axes]:
        """Implement actual plotting logic in subclasses"""

    def check_for_cmap(self, kwargs, plot_kwargs):
        if "cmap" not in kwargs and "colormap" not in kwargs:
            plot_kwargs["cmap"] = self.cmap
        return plot_kwargs

    def get_cmap(self, kwargs) -> str | None:
        if not kwargs:
            return self.cmap

        if "cmap" not in kwargs and "colormap" not in kwargs:
            return self.cmap

        if "cmap" in kwargs:
            return kwargs["cmap"]

        if "colormap" in kwargs:
            return kwargs["colormap"]

        raise ValueError  # pragma: no cover


class StackedBarChart(ChartBase):
    cmap: str | None = "inferno_r"

    def _do_plot(
        self,
        df: _pd.DataFrame,
        columns: list[str],
        use_legend: bool = True,
        size: tuple[float, float] = conf.PlotSizes.A4.value,
        **kwargs: _tp.Any,
    ) -> tuple[_plt.Figure, _plt.Axes]:
        fig, ax = _plt.subplots(
            figsize=size,
            layout="constrained",
        )
        plot_kwargs = {
            "stacked": True,
            "legend": use_legend,
            "ax": ax,
            **kwargs,
        }
        self.check_for_cmap(kwargs, plot_kwargs)
        ax = df[columns].plot.bar(**plot_kwargs)
        ax.set_xticklabels(
            _pd.to_datetime(df.index).strftime(plot_settings.date_format)
        )

        return fig, ax


class BarChart(ChartBase):
    cmap = None

    def _do_plot(
        self,
        df: _pd.DataFrame,
        columns: list[str],
        use_legend: bool = True,
        size: tuple[float, float] = conf.PlotSizes.A4.value,
        **kwargs: _tp.Any,
    ) -> tuple[_plt.Figure, _plt.Axes]:
        # TODO: deal with colors  # pylint: disable=fixme
        fig, ax = _plt.subplots(
            figsize=size,
            layout="constrained",
        )
        x = _np.arange(len(df.index))
        width = 0.8 / len(columns)

        cmap = self.get_cmap(kwargs)
        if cmap:
            cm = _plt.cm.get_cmap(cmap)
            colors = cm(_np.linspace(0, 1, len(columns)))
        else:
            colors = [None] * len(columns)

        for i, col in enumerate(columns):
            ax.bar(x + i * width, df[col], width, label=col, color=colors[i])

        if use_legend:
            ax.legend()

        ax.set_xticks(x + width * (len(columns) - 1) / 2)
        ax.set_xticklabels(
            _pd.to_datetime(df.index).strftime(plot_settings.date_format)
        )
        ax.tick_params(axis="x", labelrotation=90)
        return fig, ax


class LinePlot(ChartBase):
    cmap: str | None = None

    def _do_plot(
        self,
        df: _pd.DataFrame,
        columns: list[str],
        use_legend: bool = True,
        size: tuple[float, float] = conf.PlotSizes.A4.value,
        **kwargs: _tp.Any,
    ) -> tuple[_plt.Figure, _plt.Axes]:
        fig, ax = _plt.subplots(
            figsize=size,
            layout="constrained",
        )
        plot_kwargs = {
            "legend": use_legend,
            "ax": ax,
            **kwargs,
        }
        self.check_for_cmap(kwargs, plot_kwargs)

        df[columns].plot.line(**plot_kwargs)
        return fig, ax


@dataclass()
class Histogram(ChartBase):
    bins: int = 50

    def _do_plot(
        self,
        df: _pd.DataFrame,
        columns: list[str],
        use_legend: bool = True,
        size: tuple[float, float] = conf.PlotSizes.A4.value,
        **kwargs: _tp.Any,
    ) -> tuple[_plt.Figure, _plt.Axes]:
        fig, ax = _plt.subplots(
            figsize=size,
            layout="constrained",
        )
        plot_kwargs = {
            "legend": use_legend,
            "ax": ax,
            "bins": self.bins,
            **kwargs,
        }
        self.check_for_cmap(kwargs, plot_kwargs)
        df[columns].plot.hist(**plot_kwargs)
        return fig, ax


def _validate_inputs(
    current_class,
    columns: list[str],
) -> None:
    if len(columns) != 2:
        raise ValueError(
            f"{current_class.__name__} requires exactly 2 columns (x and y)"
        )


class ScatterPlot(ChartBase):
    cmap = "Paired"  # This is ignored when no categorical groupings are used.

    def _do_plot(
        self,
        df: _pd.DataFrame,
        columns: list[str],
        use_legend: bool = True,
        size: tuple[float, float] = conf.PlotSizes.A4.value,
        **kwargs: _tp.Any,
    ) -> tuple[_plt.Figure, _plt.Axes]:
        _validate_inputs(self, columns)
        x_column, y_column = columns

        fig, ax = _plt.subplots(
            figsize=size,
            layout="constrained",
        )
        df.plot.scatter(x=x_column, y=y_column, ax=ax, **kwargs)
        return fig, ax


class ScalarComparePlot(ChartBase):
    """Handles comparative scatter plots with dual grouping by color and markers."""

    cmap = "Paired"  # This is ignored when no categorical groupings are used.

    # pylint: disable=too-many-arguments,too-many-locals, too-many-positional-arguments
    def _do_plot(  # type: ignore[override]
        self,
        df: _pd.DataFrame,
        columns: list[str],
        use_legend: bool = True,
        size: tuple[float, float] = conf.PlotSizes.A4.value,
        group_by_color: str | None = None,
        group_by_marker: str | None = None,
        line_kwargs: dict[str, _tp.Any] | None = None,
        scatter_kwargs: dict[str, _tp.Any] | None = None,
    ) -> tuple[_plt.Figure, _plt.Axes]:

        _validate_inputs(self, columns)
        x_column, y_column = columns

        # ===========================================
        # The following simplifies the code later on,
        # while being compatible with linting.
        if not line_kwargs:
            line_kwargs = {}
        if not scatter_kwargs:
            scatter_kwargs = {}
        # ===========================================

        if group_by_color and group_by_marker:
            # See: https://stackoverflow.com/questions/4700614/
            # how-to-put-the-legend-outside-the-plot
            # This is required to place the legend in a dedicated subplot
            fig, (ax, lax) = _plt.subplots(
                layout="constrained",
                figsize=size,
                ncols=2,
                gridspec_kw={"width_ratios": [4, 1]},
            )
            secondary_axis_used = True
        else:
            secondary_axis_used = False
            fig, ax = _plt.subplots(
                layout="constrained",
                figsize=size,
            )
            lax = ax

        df_grouped, group_values = self._prepare_grouping(
            df, group_by_color, group_by_marker
        )
        cmap = self.get_cmap(line_kwargs)
        color_map, marker_map = self._create_style_mappings(
            *group_values, cmap=cmap
        )

        self._plot_groups(
            df_grouped,
            x_column,
            y_column,
            color_map,
            marker_map,
            ax,
            line_kwargs,
            scatter_kwargs,
        )

        use_color_legend = False
        if group_by_color:
            use_color_legend = True

        if use_legend:
            self._create_legends(
                lax,
                color_map,
                marker_map,
                group_by_color,
                group_by_marker,
                use_color_legend=use_color_legend,
                secondary_axis_used=secondary_axis_used,
            )

        return fig, ax

    @staticmethod
    def _prepare_grouping(
        df: _pd.DataFrame,
        by_color: str | None,
        by_marker: str | None,
    ) -> tuple[
        _pd.core.groupby.generic.DataFrameGroupBy, tuple[list[str], list[str]]
    ]:
        group_by = []
        if by_color:
            group_by.append(by_color)
        if by_marker:
            group_by.append(by_marker)

        df_grouped = df.groupby(group_by)

        color_values = sorted(df[by_color].unique()) if by_color else []
        marker_values = sorted(df[by_marker].unique()) if by_marker else []

        return df_grouped, (color_values, marker_values)

    @staticmethod
    def _create_style_mappings(
        color_values: list[str],
        marker_values: list[str],
        cmap: str | None,
    ) -> tuple[dict[str, _tp.Any], dict[str, str]]:
        if color_values:
            cm = _plt.get_cmap(cmap, len(color_values))
            color_map = {val: cm(i) for i, val in enumerate(color_values)}
        else:
            cm = _plt.get_cmap(cmap, len(marker_values))
            color_map = {val: cm(i) for i, val in enumerate(marker_values)}
        if marker_values:
            marker_map = dict(zip(marker_values, plot_settings.markers))
        else:
            marker_map = {}

        return color_map, marker_map

    # pylint: disable=too-many-arguments
    @staticmethod
    def _plot_groups(
        df_grouped: _pd.core.groupby.generic.DataFrameGroupBy,
        x_column: str,
        y_column: str,
        color_map: dict[str, _tp.Any],
        marker_map: dict[str, str] | str,
        ax: _plt.Axes,
        line_kwargs: dict[str, _tp.Any],
        scatter_kwargs: dict[str, _tp.Any],
    ) -> None:
        ax.set_xlabel(x_column, fontsize=plot_settings.label_font_size)
        ax.set_ylabel(y_column, fontsize=plot_settings.label_font_size)
        for val, group in df_grouped:
            sorted_group = group.sort_values(x_column)
            x = sorted_group[x_column]
            y = sorted_group[y_column]

            plot_args = {"color": "black"}
            if color_map:
                plot_args["color"] = color_map[val[0]]

            for key, value in line_kwargs.items():
                if key not in ["cmap", "colormap"]:
                    plot_args[key] = value

            scatter_args = {"marker": "None", "color": "black", "alpha": 0.5}
            if marker_map:
                scatter_args["marker"] = marker_map[val[-1]]

            for key, value in scatter_kwargs.items():
                if key in ["marker"] and marker_map:
                    continue

                scatter_args[key] = value

            ax.plot(x, y, **plot_args)  # type: ignore
            ax.scatter(x, y, **scatter_args)  # type: ignore

    def _create_legends(
        self,
        lax: _plt.Axes,
        color_map: dict[str, _tp.Any],
        marker_map: dict[str, str],
        color_legend_title: str | None,
        marker_legend_title: str | None,
        use_color_legend: bool,
        secondary_axis_used: bool,
    ) -> None:

        if secondary_axis_used:
            # Secondary axis should be turned off.
            # Primary axis should stay the same.
            lax.axis("off")

        if use_color_legend:
            self._create_color_legend(
                lax,
                color_map,
                color_legend_title,
                bool(marker_map),
                secondary_axis_used,
            )
        if marker_map:
            self._create_marker_legend(
                lax,
                marker_map,
                marker_legend_title,
                bool(use_color_legend),
                secondary_axis_used,
            )

    @staticmethod
    def _create_color_legend(
        lax: _plt.Axes,
        color_map: dict[str, _tp.Any],
        color_legend_title: str | None,
        has_markers: bool,
        secondary_axis_used: bool,
    ) -> None:
        color_handles = [
            _plt.Line2D([], [], color=color, linestyle="-", label=label)
            for label, color in color_map.items()
        ]

        if secondary_axis_used:
            loc = "upper left"
            alignment = "left"
        else:
            loc = "best"
            alignment = "center"

        legend = lax.legend(
            handles=color_handles,
            title=color_legend_title,
            bbox_to_anchor=(0, 0, 1, 1),
            loc=loc,
            alignment=alignment,
            fontsize=plot_settings.legend_font_size,
            borderaxespad=0,
        )

        if has_markers:
            lax.add_artist(legend)

    @staticmethod
    def _create_marker_legend(
        lax: _plt.Axes,
        marker_map: dict[str, str],
        marker_legend_title: str | None,
        has_colors: bool,
        secondary_axis_used: bool,
    ) -> None:
        marker_position = 0.7 if has_colors else 1
        marker_handles = [
            _plt.Line2D(
                [],
                [],
                color="black",
                marker=marker,
                linestyle="None",
                label=label,
            )
            for label, marker in marker_map.items()
            if label is not None
        ]

        if secondary_axis_used:
            loc = "upper left"
            alignment = "left"
        else:
            loc = "best"
            alignment = "center"

        lax.legend(
            handles=marker_handles,
            title=marker_legend_title,
            bbox_to_anchor=(0, 0, 1, marker_position),
            loc=loc,
            alignment=alignment,
            fontsize=plot_settings.legend_font_size,
            borderaxespad=0,
        )

1	import typing as _tp	2✔
2	from abc import abstractmethod	2✔
3	from dataclasses import dataclass	2✔
4
5	import matplotlib.pyplot as _plt	2✔
6	import numpy as _np	2✔
7	import pandas as _pd	2✔
8
9	from pytrnsys_process import config as conf	2✔
10
11	# TODO: provide A4 and half A4 plots to test sizes in latex # pylint: disable=fixme
12	# TODO: provide height as input for plot? # pylint: disable=fixme
13	# TODO: deal with legends (curve names, fonts, colors, linestyles) # pylint: disable=fixme
14	# TODO: clean up old stuff by refactoring # pylint: disable=fixme
15	# TODO: make issue for docstrings of plotting # pylint: disable=fixme
16	# TODO: Add colormap support # pylint: disable=fixme
17
18
19	# TODO find a better place for this to live in # pylint : disable=fixme
20	plot_settings = conf.global_settings.plot	2✔
21	"Settings shared by all plots"	2✔
22
23
24	class ChartBase:	2✔
25	cmap: str \| None = None	2✔
26
27	def plot(	2✔
28	self,
29	df: _pd.DataFrame,
30	columns: list[str],
31	**kwargs,
32	) -> tuple[_plt.Figure, _plt.Axes]:
33	fig, ax = self._do_plot(df, columns, **kwargs)	2✔
34	return fig, ax	2✔
35
36	@abstractmethod	2✔
37	def _do_plot(	2✔
38	self,
39	df: _pd.DataFrame,
40	columns: list[str],
41	use_legend: bool = True,
42	size: tuple[float, float] = conf.PlotSizes.A4.value,
43	**kwargs: _tp.Any,
44	) -> tuple[_plt.Figure, _plt.Axes]:
45	"""Implement actual plotting logic in subclasses"""
46
47	def check_for_cmap(self, kwargs, plot_kwargs):	2✔
48	if "cmap" not in kwargs and "colormap" not in kwargs:	2✔
49	plot_kwargs["cmap"] = self.cmap	2✔
50	return plot_kwargs	2✔
51
52	def get_cmap(self, kwargs) -> str \| None:	2✔
53	if not kwargs:	2✔
54	return self.cmap	2✔
55
56	if "cmap" not in kwargs and "colormap" not in kwargs:	2✔
57	return self.cmap	2✔
58
59	if "cmap" in kwargs:	2✔
60	return kwargs["cmap"]	2✔
61
62	if "colormap" in kwargs:	2✔
63	return kwargs["colormap"]	2✔
64
65	raise ValueError # pragma: no cover
66
67
68	class StackedBarChart(ChartBase):	2✔
69	cmap: str \| None = "inferno_r"	2✔
70
71	def _do_plot(	2✔
72	self,
73	df: _pd.DataFrame,
74	columns: list[str],
75	use_legend: bool = True,
76	size: tuple[float, float] = conf.PlotSizes.A4.value,
77	**kwargs: _tp.Any,
78	) -> tuple[_plt.Figure, _plt.Axes]:
79	fig, ax = _plt.subplots(	2✔
80	figsize=size,
81	layout="constrained",
82	)
83	plot_kwargs = {	2✔
84	"stacked": True,
85	"legend": use_legend,
86	"ax": ax,
87	**kwargs,
88	}
89	self.check_for_cmap(kwargs, plot_kwargs)	2✔
90	ax = df[columns].plot.bar(**plot_kwargs)	2✔
91	ax.set_xticklabels(	2✔
92	_pd.to_datetime(df.index).strftime(plot_settings.date_format)
93	)
94
95	return fig, ax	2✔
96
97
98	class BarChart(ChartBase):	2✔
99	cmap = None	2✔
100
101	def _do_plot(	2✔
102	self,
103	df: _pd.DataFrame,
104	columns: list[str],
105	use_legend: bool = True,
106	size: tuple[float, float] = conf.PlotSizes.A4.value,
107	**kwargs: _tp.Any,
108	) -> tuple[_plt.Figure, _plt.Axes]:
109	# TODO: deal with colors # pylint: disable=fixme
110	fig, ax = _plt.subplots(	2✔
111	figsize=size,
112	layout="constrained",
113	)
114	x = _np.arange(len(df.index))	2✔
115	width = 0.8 / len(columns)	2✔
116
117	cmap = self.get_cmap(kwargs)	2✔
118	if cmap:	2✔
119	cm = _plt.cm.get_cmap(cmap)	2✔
120	colors = cm(_np.linspace(0, 1, len(columns)))	2✔
121	else:
122	colors = [None] * len(columns)	2✔
123
124	for i, col in enumerate(columns):	2✔
125	ax.bar(x + i * width, df[col], width, label=col, color=colors[i])	2✔
126
127	if use_legend:	2✔
128	ax.legend()	2✔
129
130	ax.set_xticks(x + width * (len(columns) - 1) / 2)	2✔
131	ax.set_xticklabels(	2✔
132	_pd.to_datetime(df.index).strftime(plot_settings.date_format)
133	)
134	ax.tick_params(axis="x", labelrotation=90)	2✔
135	return fig, ax	2✔
136
137
138	class LinePlot(ChartBase):	2✔
139	cmap: str \| None = None	2✔
140
141	def _do_plot(	2✔
142	self,
143	df: _pd.DataFrame,
144	columns: list[str],
145	use_legend: bool = True,
146	size: tuple[float, float] = conf.PlotSizes.A4.value,
147	**kwargs: _tp.Any,
148	) -> tuple[_plt.Figure, _plt.Axes]:
149	fig, ax = _plt.subplots(	2✔
150	figsize=size,
151	layout="constrained",
152	)
153	plot_kwargs = {	2✔
154	"legend": use_legend,
155	"ax": ax,
156	**kwargs,
157	}
158	self.check_for_cmap(kwargs, plot_kwargs)	2✔
159
160	df[columns].plot.line(**plot_kwargs)	2✔
161	return fig, ax	2✔
162
163
164	@dataclass()	2✔
165	class Histogram(ChartBase):	2✔
166	bins: int = 50	2✔
167
168	def _do_plot(	2✔
169	self,
170	df: _pd.DataFrame,
171	columns: list[str],
172	use_legend: bool = True,
173	size: tuple[float, float] = conf.PlotSizes.A4.value,
174	**kwargs: _tp.Any,
175	) -> tuple[_plt.Figure, _plt.Axes]:
176	fig, ax = _plt.subplots(	2✔
177	figsize=size,
178	layout="constrained",
179	)
180	plot_kwargs = {	2✔
181	"legend": use_legend,
182	"ax": ax,
183	"bins": self.bins,
184	**kwargs,
185	}
186	self.check_for_cmap(kwargs, plot_kwargs)	2✔
187	df[columns].plot.hist(**plot_kwargs)	2✔
188	return fig, ax	2✔
189
190
191	def _validate_inputs(	2✔
192	current_class,
193	columns: list[str],
194	) -> None:
195	if len(columns) != 2:	2✔
UNCOV 196	raise ValueError(	×
197	f"{current_class.__name__} requires exactly 2 columns (x and y)"
198	)
199
200
201	class ScatterPlot(ChartBase):	2✔
202	cmap = "Paired" # This is ignored when no categorical groupings are used.	2✔
203
204	def _do_plot(	2✔
205	self,
206	df: _pd.DataFrame,
207	columns: list[str],
208	use_legend: bool = True,
209	size: tuple[float, float] = conf.PlotSizes.A4.value,
210	**kwargs: _tp.Any,
211	) -> tuple[_plt.Figure, _plt.Axes]:
212	_validate_inputs(self, columns)	2✔
213	x_column, y_column = columns	2✔
214
215	fig, ax = _plt.subplots(	2✔
216	figsize=size,
217	layout="constrained",
218	)
219	df.plot.scatter(x=x_column, y=y_column, ax=ax, **kwargs)	2✔
220	return fig, ax	2✔
221
222
223	class ScalarComparePlot(ChartBase):	2✔
224	"""Handles comparative scatter plots with dual grouping by color and markers."""
225
226	cmap = "Paired" # This is ignored when no categorical groupings are used.	2✔
227
228	# pylint: disable=too-many-arguments,too-many-locals, too-many-positional-arguments
229	def _do_plot( # type: ignore[override]	2✔
230	self,
231	df: _pd.DataFrame,
232	columns: list[str],
233	use_legend: bool = True,
234	size: tuple[float, float] = conf.PlotSizes.A4.value,
235	group_by_color: str \| None = None,
236	group_by_marker: str \| None = None,
237	line_kwargs: dict[str, _tp.Any] \| None = None,
238	scatter_kwargs: dict[str, _tp.Any] \| None = None,
239	) -> tuple[_plt.Figure, _plt.Axes]:
240
241	_validate_inputs(self, columns)	2✔
242	x_column, y_column = columns	2✔
243
244	# ===========================================
245	# The following simplifies the code later on,
246	# while being compatible with linting.
247	if not line_kwargs:	2✔
248	line_kwargs = {}	2✔
249	if not scatter_kwargs:	2✔
250	scatter_kwargs = {}	2✔
251	# ===========================================
252
253	if group_by_color and group_by_marker:	2✔
254	# See: https://stackoverflow.com/questions/4700614/
255	# how-to-put-the-legend-outside-the-plot
256	# This is required to place the legend in a dedicated subplot
257	fig, (ax, lax) = _plt.subplots(	2✔
258	layout="constrained",
259	figsize=size,
260	ncols=2,
261	gridspec_kw={"width_ratios": [4, 1]},
262	)
263	secondary_axis_used = True	2✔
264	else:
265	secondary_axis_used = False	2✔
266	fig, ax = _plt.subplots(	2✔
267	layout="constrained",
268	figsize=size,
269	)
270	lax = ax	2✔
271
272	df_grouped, group_values = self._prepare_grouping(	2✔
273	df, group_by_color, group_by_marker
274	)
275	cmap = self.get_cmap(line_kwargs)	2✔
276	color_map, marker_map = self._create_style_mappings(	2✔
277	*group_values, cmap=cmap
278	)
279
280	self._plot_groups(	2✔
281	df_grouped,
282	x_column,
283	y_column,
284	color_map,
285	marker_map,
286	ax,
287	line_kwargs,
288	scatter_kwargs,
289	)
290
291	use_color_legend = False	2✔
292	if group_by_color:	2✔
293	use_color_legend = True	2✔
294
295	if use_legend:	2✔
296	self._create_legends(	2✔
297	lax,
298	color_map,
299	marker_map,
300	group_by_color,
301	group_by_marker,
302	use_color_legend=use_color_legend,
303	secondary_axis_used=secondary_axis_used,
304	)
305
306	return fig, ax	2✔
307
308	@staticmethod	2✔
309	def _prepare_grouping(	2✔
310	df: _pd.DataFrame,
311	by_color: str \| None,
312	by_marker: str \| None,
313	) -> tuple[
314	_pd.core.groupby.generic.DataFrameGroupBy, tuple[list[str], list[str]]
315	]:
316	group_by = []	2✔
317	if by_color:	2✔
318	group_by.append(by_color)	2✔
319	if by_marker:	2✔
320	group_by.append(by_marker)	2✔
321
322	df_grouped = df.groupby(group_by)	2✔
323
324	color_values = sorted(df[by_color].unique()) if by_color else []	2✔
325	marker_values = sorted(df[by_marker].unique()) if by_marker else []	2✔
326
327	return df_grouped, (color_values, marker_values)	2✔
328
329	@staticmethod	2✔
330	def _create_style_mappings(	2✔
331	color_values: list[str],
332	marker_values: list[str],
333	cmap: str \| None,
334	) -> tuple[dict[str, _tp.Any], dict[str, str]]:
335	if color_values:	2✔
336	cm = _plt.get_cmap(cmap, len(color_values))	2✔
337	color_map = {val: cm(i) for i, val in enumerate(color_values)}	2✔
338	else:
339	cm = _plt.get_cmap(cmap, len(marker_values))	2✔
340	color_map = {val: cm(i) for i, val in enumerate(marker_values)}	2✔
341	if marker_values:	2✔
342	marker_map = dict(zip(marker_values, plot_settings.markers))	2✔
343	else:
344	marker_map = {}	2✔
345
346	return color_map, marker_map	2✔
347
348	# pylint: disable=too-many-arguments
349	@staticmethod	2✔
350	def _plot_groups(	2✔
351	df_grouped: _pd.core.groupby.generic.DataFrameGroupBy,
352	x_column: str,
353	y_column: str,
354	color_map: dict[str, _tp.Any],
355	marker_map: dict[str, str] \| str,
356	ax: _plt.Axes,
357	line_kwargs: dict[str, _tp.Any],
358	scatter_kwargs: dict[str, _tp.Any],
359	) -> None:
360	ax.set_xlabel(x_column, fontsize=plot_settings.label_font_size)	2✔
361	ax.set_ylabel(y_column, fontsize=plot_settings.label_font_size)	2✔
362	for val, group in df_grouped:	2✔
363	sorted_group = group.sort_values(x_column)	2✔
364	x = sorted_group[x_column]	2✔
365	y = sorted_group[y_column]	2✔
366
367	plot_args = {"color": "black"}	2✔
368	if color_map:	2✔
369	plot_args["color"] = color_map[val[0]]	2✔
370
371	for key, value in line_kwargs.items():	2✔
372	if key not in ["cmap", "colormap"]:	2✔
373	plot_args[key] = value	2✔
374
375	scatter_args = {"marker": "None", "color": "black", "alpha": 0.5}	2✔
376	if marker_map:	2✔
377	scatter_args["marker"] = marker_map[val[-1]]	2✔
378
379	for key, value in scatter_kwargs.items():	2✔
380	if key in ["marker"] and marker_map:	2✔
381	continue	2✔
382
383	scatter_args[key] = value	2✔
384
385	ax.plot(x, y, **plot_args) # type: ignore	2✔
386	ax.scatter(x, y, **scatter_args) # type: ignore	2✔
387
388	def _create_legends(	2✔
389	self,
390	lax: _plt.Axes,
391	color_map: dict[str, _tp.Any],
392	marker_map: dict[str, str],
393	color_legend_title: str \| None,
394	marker_legend_title: str \| None,
395	use_color_legend: bool,
396	secondary_axis_used: bool,
397	) -> None:
398
399	if secondary_axis_used:	2✔
400	# Secondary axis should be turned off.
401	# Primary axis should stay the same.
402	lax.axis("off")	2✔
403
404	if use_color_legend:	2✔
405	self._create_color_legend(	2✔
406	lax,
407	color_map,
408	color_legend_title,
409	bool(marker_map),
410	secondary_axis_used,
411	)
412	if marker_map:	2✔
413	self._create_marker_legend(	2✔
414	lax,
415	marker_map,
416	marker_legend_title,
417	bool(use_color_legend),
418	secondary_axis_used,
419	)
420
421	@staticmethod	2✔
422	def _create_color_legend(	2✔
423	lax: _plt.Axes,
424	color_map: dict[str, _tp.Any],
425	color_legend_title: str \| None,
426	has_markers: bool,
427	secondary_axis_used: bool,
428	) -> None:
429	color_handles = [	2✔
430	_plt.Line2D([], [], color=color, linestyle="-", label=label)
431	for label, color in color_map.items()
432	]
433
434	if secondary_axis_used:	2✔
435	loc = "upper left"	2✔
436	alignment = "left"	2✔
437	else:
438	loc = "best"	2✔
439	alignment = "center"	2✔
440
441	legend = lax.legend(	2✔
442	handles=color_handles,
443	title=color_legend_title,
444	bbox_to_anchor=(0, 0, 1, 1),
445	loc=loc,
446	alignment=alignment,
447	fontsize=plot_settings.legend_font_size,
448	borderaxespad=0,
449	)
450
451	if has_markers:	2✔
452	lax.add_artist(legend)	2✔
453
454	@staticmethod	2✔
455	def _create_marker_legend(	2✔
456	lax: _plt.Axes,
457	marker_map: dict[str, str],
458	marker_legend_title: str \| None,
459	has_colors: bool,
460	secondary_axis_used: bool,
461	) -> None:
462	marker_position = 0.7 if has_colors else 1	2✔
463	marker_handles = [	2✔
464	_plt.Line2D(
465	[],
466	[],
467	color="black",
468	marker=marker,
469	linestyle="None",
470	label=label,
471	)
472	for label, marker in marker_map.items()
473	if label is not None
474	]
475
476	if secondary_axis_used:	2✔
477	loc = "upper left"	2✔
478	alignment = "left"	2✔
479	else:
480	loc = "best"	2✔
481	alignment = "center"	2✔
482
483	lax.legend(	2✔
484	handles=marker_handles,
485	title=marker_legend_title,
486	bbox_to_anchor=(0, 0, 1, marker_position),
487	loc=loc,
488	alignment=alignment,
489	fontsize=plot_settings.legend_font_size,
490	borderaxespad=0,
491	)

SPF-OST / pytrnsys_process / 15485406069

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous