16748222285

Committed 05 Aug 2025 11:03AM UTC coverage: 49.518% (-46.5%) from 95.968%

Build Type

github

Committed by ahobeost

Commit Message

Reduce linux job to just test.

Pull Request Pull Request #126: 125 bug step file not read when step used with type 25

Run Details

5 of 6 new or added lines in 2 files covered. (83.33%)

578 existing lines in 11 files now uncovered.

616 of 1244 relevant lines covered (49.52%)

0.99 hits per line

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

25.65

/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

    @staticmethod
    def get_fig_and_ax(kwargs, size):
        if "fig" not in kwargs and "ax" not in kwargs:
            fig, ax = _plt.subplots(
                figsize=size,
                layout="constrained",
            )
            kwargs["ax"] = ax
        else:
            fig = kwargs["fig"]
            ax = kwargs["ax"]
        return fig, ax


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 = self.get_fig_and_ax(kwargs, size)

        plot_kwargs = {
            "stacked": True,
            "legend": use_legend,
            **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 = self.get_fig_and_ax(kwargs, size)

        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 = self.get_fig_and_ax(kwargs, size)

        plot_kwargs = {
            "legend": use_legend,
            **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 = self.get_fig_and_ax(kwargs, size)

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

SPF-OST / pytrnsys_process / 16748222285

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