From 52c59aa68943ba1c32e0cb51703b9d6551f6ea7f Mon Sep 17 00:00:00 2001 From: Anton Hoyer <156099087+hoyer-a@users.noreply.github.com> Date: Tue, 21 Oct 2025 17:18:35 +0200 Subject: [PATCH 01/10] add utils functions for plots --- spharpy/plot/__init__.py | 5 +++ spharpy/plot/_utils.py | 82 +++++++++++++++++++++++++++++++++++++++ tests/test_plot__utils.py | 45 +++++++++++++++++++++ 3 files changed, 132 insertions(+) create mode 100644 spharpy/plot/_utils.py create mode 100644 tests/test_plot__utils.py diff --git a/spharpy/plot/__init__.py b/spharpy/plot/__init__.py index 87207a45..7ed9b921 100644 --- a/spharpy/plot/__init__.py +++ b/spharpy/plot/__init__.py @@ -12,6 +12,10 @@ MidpointNormalize, ) +from ._utils import ( + _prepare_plot, +) + from .cmap import phase_twilight @@ -26,4 +30,5 @@ 'contour_map', 'MidpointNormalize', 'phase_twilight', + '_prepare_plot' ] diff --git a/spharpy/plot/_utils.py b/spharpy/plot/_utils.py new file mode 100644 index 00000000..f53c09a2 --- /dev/null +++ b/spharpy/plot/_utils.py @@ -0,0 +1,82 @@ +"""Private utility functions for plot module.""" + +import numpy as np +import matplotlib.pyplot as plt + + +def _add_colorbar(colorbar, fig, ax, mappable, label): + """ + Add colorbar to plot + + Parameters + ---------- + + colorbar : bool + Flag indicating if a colobar should be added to the plot + fig : matplotlib.figure.Figure + Figure to plot on. + ax : list[matplotlib.axes.Axes] + Either a list of two axes objects or a list with one axis and None + object. + mappable : matplotlib.cm.ScalarMappable + The matplotlib.cm.ScalarMappable described by the colorbar. + label : string + colorbar label + + Returns + ------- + cb : matplotlib colorbar object + """ + if colorbar: + if ax[1] is None: + cb = fig.colorbar(mappable, ax=ax[0]) + else: + cb = fig.colorbar(mappable, ax=ax[1]) + cb.set_label(label) + else: + cb = None + + return cb + + +def _prepare_plot(ax, projection=None): + """ + Returns a figure to plot on. + + Parameters + ---------- + ax : matplotlib.axes.Axes or list, tuple or ndarray of maplotlib.axes.Axes + Axes to plot on. The default is None in which case the axes are + obtained from the current figure. A new figure is created if it does + not exist. + projection : str, optional + Type of projection for the axes. This is only applied if new axes are + created. Default is ``None`` (2D projection). + + Returns + ------- + fig : matplotlib.figure.Figure + Returns the active figure. + ax : maptlotlib.axes.Axes + Returns the current axes. + """ + if ax is None: + # get current figure or create a new one + fig = plt.gcf() + if fig.axes: + ax = plt.gca() + else: + ax = plt.axes(projection=projection) + + else: + # get figure from axis + # ax can be array or tuple of two ax objects + # first axis for the plot, second axis for colorbar placement + if isinstance(ax, np.ndarray): + fig = ax.flatten()[0].figure + elif isinstance(ax, (list, tuple)): + fig = ax[0].figure + else: + fig = ax.figure + + return fig, ax diff --git a/tests/test_plot__utils.py b/tests/test_plot__utils.py new file mode 100644 index 00000000..b5a5dabd --- /dev/null +++ b/tests/test_plot__utils.py @@ -0,0 +1,45 @@ +import pytest +import spharpy as sp +import matplotlib.pyplot as plt +import matplotlib as mpl + + +@pytest.mark.parametrize( + ("input_ax", "output_type", "projection"), + [(None, plt.Axes, '3d'), (plt.gca(), plt.Axes, None), + ([plt.gca(), plt.gca()], list, None)] +) +def test_prepare_plot(input_ax, output_type, projection): + """ + Test output of :py:func:`~spharpy.plot._utils._prepare_plot`. + """ + fig, ax = sp.plot._prepare_plot(input_ax, projection) + + assert isinstance(fig, plt.Figure) + assert isinstance(ax, output_type) + + if isinstance(ax, list): + assert all(isinstance(ax_, plt.Axes) for ax_ in input_ax) + + if ax is None: + assert ax.name == projection + + +@pytest.mark.parametrize( + ("colorbar", "ax", "return_type"), + [(True, [plt.axes(), None], mpl.colorbar.Colorbar), + (True, [plt.axes(), plt.axes()], mpl.colorbar.Colorbar), + (False, [plt.axes(), None], None)] +) +def test_add_colorbar(colorbar, ax, return_type): + norm = mpl.colors.Normalize(vmin=0, vmax=10) + cmap = plt.get_cmap('viridis') + mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap) + + fig = plt.gcf() + cb = sp.plot._utils._add_colorbar(colorbar, fig, ax, mappable, None) + + if return_type is None: + assert cb is None + else: + assert isinstance(cb, return_type) From 6232cef992a04b2ffd7be51ef34ba7d9c7467db9 Mon Sep 17 00:00:00 2001 From: Anton Hoyer <156099087+hoyer-a@users.noreply.github.com> Date: Tue, 21 Oct 2025 17:23:40 +0200 Subject: [PATCH 02/10] fix ruff and remove private function from init --- spharpy/plot/__init__.py | 6 +----- spharpy/plot/_utils.py | 3 +-- tests/test_plot__utils.py | 6 +++--- 3 files changed, 5 insertions(+), 10 deletions(-) diff --git a/spharpy/plot/__init__.py b/spharpy/plot/__init__.py index 7ed9b921..fedd45f7 100644 --- a/spharpy/plot/__init__.py +++ b/spharpy/plot/__init__.py @@ -12,10 +12,6 @@ MidpointNormalize, ) -from ._utils import ( - _prepare_plot, -) - from .cmap import phase_twilight @@ -30,5 +26,5 @@ 'contour_map', 'MidpointNormalize', 'phase_twilight', - '_prepare_plot' + ] diff --git a/spharpy/plot/_utils.py b/spharpy/plot/_utils.py index f53c09a2..4a61fc78 100644 --- a/spharpy/plot/_utils.py +++ b/spharpy/plot/_utils.py @@ -6,11 +6,10 @@ def _add_colorbar(colorbar, fig, ax, mappable, label): """ - Add colorbar to plot + Add colorbar to plot. Parameters ---------- - colorbar : bool Flag indicating if a colobar should be added to the plot fig : matplotlib.figure.Figure diff --git a/tests/test_plot__utils.py b/tests/test_plot__utils.py index b5a5dabd..aad94eaf 100644 --- a/tests/test_plot__utils.py +++ b/tests/test_plot__utils.py @@ -7,13 +7,13 @@ @pytest.mark.parametrize( ("input_ax", "output_type", "projection"), [(None, plt.Axes, '3d'), (plt.gca(), plt.Axes, None), - ([plt.gca(), plt.gca()], list, None)] + ([plt.gca(), plt.gca()], list, None)], ) def test_prepare_plot(input_ax, output_type, projection): """ Test output of :py:func:`~spharpy.plot._utils._prepare_plot`. """ - fig, ax = sp.plot._prepare_plot(input_ax, projection) + fig, ax = sp.plot._utils._prepare_plot(input_ax, projection) assert isinstance(fig, plt.Figure) assert isinstance(ax, output_type) @@ -29,7 +29,7 @@ def test_prepare_plot(input_ax, output_type, projection): ("colorbar", "ax", "return_type"), [(True, [plt.axes(), None], mpl.colorbar.Colorbar), (True, [plt.axes(), plt.axes()], mpl.colorbar.Colorbar), - (False, [plt.axes(), None], None)] + (False, [plt.axes(), None], None)], ) def test_add_colorbar(colorbar, ax, return_type): norm = mpl.colors.Normalize(vmin=0, vmax=10) From 4280c4fbcf500c036b3e4ad345d7e146e071af5b Mon Sep 17 00:00:00 2001 From: Anton Hoyer <156099087+hoyer-a@users.noreply.github.com> Date: Tue, 21 Oct 2025 17:29:43 +0200 Subject: [PATCH 03/10] remove blank line --- spharpy/plot/__init__.py | 1 - 1 file changed, 1 deletion(-) diff --git a/spharpy/plot/__init__.py b/spharpy/plot/__init__.py index fedd45f7..87207a45 100644 --- a/spharpy/plot/__init__.py +++ b/spharpy/plot/__init__.py @@ -26,5 +26,4 @@ 'contour_map', 'MidpointNormalize', 'phase_twilight', - ] From d547680e7a8cb9fece88bf62c3166c4d95def88e Mon Sep 17 00:00:00 2001 From: Anton Hoyer <156099087+hoyer-a@users.noreply.github.com> Date: Tue, 21 Oct 2025 17:31:03 +0200 Subject: [PATCH 04/10] add docstring to test --- tests/test_plot__utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_plot__utils.py b/tests/test_plot__utils.py index aad94eaf..8914fa54 100644 --- a/tests/test_plot__utils.py +++ b/tests/test_plot__utils.py @@ -32,6 +32,7 @@ def test_prepare_plot(input_ax, output_type, projection): (False, [plt.axes(), None], None)], ) def test_add_colorbar(colorbar, ax, return_type): + """Test return type of :py:func:`~spharpy.plot._utils._add_colorbar`""" norm = mpl.colors.Normalize(vmin=0, vmax=10) cmap = plt.get_cmap('viridis') mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap) From f81f5c89e4101f10d96b3543f4d0b77cfac9a920 Mon Sep 17 00:00:00 2001 From: Anton Hoyer <156099087+hoyer-a@users.noreply.github.com> Date: Tue, 21 Oct 2025 17:32:42 +0200 Subject: [PATCH 05/10] specify return type --- spharpy/plot/_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/spharpy/plot/_utils.py b/spharpy/plot/_utils.py index 4a61fc78..8f9ab88b 100644 --- a/spharpy/plot/_utils.py +++ b/spharpy/plot/_utils.py @@ -24,7 +24,8 @@ def _add_colorbar(colorbar, fig, ax, mappable, label): Returns ------- - cb : matplotlib colorbar object + cb : matplotlib.colorbar.Colorbar + Returns matplotlib colorbar object. """ if colorbar: if ax[1] is None: From a16c3b43731eb9df894194fa2d089e440b220d87 Mon Sep 17 00:00:00 2001 From: Anton Hoyer <156099087+hoyer-a@users.noreply.github.com> Date: Tue, 21 Oct 2025 21:20:17 +0200 Subject: [PATCH 06/10] fix tests --- tests/test_plot__utils.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_plot__utils.py b/tests/test_plot__utils.py index 8914fa54..f5b961db 100644 --- a/tests/test_plot__utils.py +++ b/tests/test_plot__utils.py @@ -2,6 +2,7 @@ import spharpy as sp import matplotlib.pyplot as plt import matplotlib as mpl +from spharpy.plot._utils import _prepare_plot, _add_colorbar @pytest.mark.parametrize( @@ -13,7 +14,7 @@ def test_prepare_plot(input_ax, output_type, projection): """ Test output of :py:func:`~spharpy.plot._utils._prepare_plot`. """ - fig, ax = sp.plot._utils._prepare_plot(input_ax, projection) + fig, ax = _prepare_plot(input_ax, projection) assert isinstance(fig, plt.Figure) assert isinstance(ax, output_type) @@ -38,7 +39,7 @@ def test_add_colorbar(colorbar, ax, return_type): mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap) fig = plt.gcf() - cb = sp.plot._utils._add_colorbar(colorbar, fig, ax, mappable, None) + cb = _add_colorbar(colorbar, fig, ax, mappable, None) if return_type is None: assert cb is None From 73e0798ab5cfba4bc15aa58979958a301bda2291 Mon Sep 17 00:00:00 2001 From: Anton Hoyer <156099087+hoyer-a@users.noreply.github.com> Date: Tue, 21 Oct 2025 21:26:35 +0200 Subject: [PATCH 07/10] fix ruff --- tests/test_plot__utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/test_plot__utils.py b/tests/test_plot__utils.py index f5b961db..e1557ed3 100644 --- a/tests/test_plot__utils.py +++ b/tests/test_plot__utils.py @@ -1,5 +1,4 @@ import pytest -import spharpy as sp import matplotlib.pyplot as plt import matplotlib as mpl from spharpy.plot._utils import _prepare_plot, _add_colorbar @@ -33,7 +32,7 @@ def test_prepare_plot(input_ax, output_type, projection): (False, [plt.axes(), None], None)], ) def test_add_colorbar(colorbar, ax, return_type): - """Test return type of :py:func:`~spharpy.plot._utils._add_colorbar`""" + """Test return type of :py:func:`~spharpy.plot._utils._add_colorbar`.""" norm = mpl.colors.Normalize(vmin=0, vmax=10) cmap = plt.get_cmap('viridis') mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap) From 1e2f95fca5453e93c3313cf2eb1af19480ac3b98 Mon Sep 17 00:00:00 2001 From: Anton Hoyer <156099087+hoyer-a@users.noreply.github.com> Date: Tue, 21 Oct 2025 21:42:02 +0200 Subject: [PATCH 08/10] revise tests --- tests/test_plot__utils.py | 34 ++++++++++++++++++++++++---------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/tests/test_plot__utils.py b/tests/test_plot__utils.py index e1557ed3..f4f4c829 100644 --- a/tests/test_plot__utils.py +++ b/tests/test_plot__utils.py @@ -5,34 +5,48 @@ @pytest.mark.parametrize( - ("input_ax", "output_type", "projection"), - [(None, plt.Axes, '3d'), (plt.gca(), plt.Axes, None), - ([plt.gca(), plt.gca()], list, None)], + ("ax_case", "output_type", "projection"), + [("none", plt.Axes, "3d"), ("single", plt.Axes, None), + ("two", list, None)], ) -def test_prepare_plot(input_ax, output_type, projection): +def test_prepare_plot(ax_case, output_type, projection): """ Test output of :py:func:`~spharpy.plot._utils._prepare_plot`. """ + if ax_case == "none": + input_ax = None + elif ax_case == "single": + _, input_ax = plt.subplots() + else: + _, axs = plt.subplots(1, 2) + input_ax = [axs[0], axs[1]] + fig, ax = _prepare_plot(input_ax, projection) assert isinstance(fig, plt.Figure) assert isinstance(ax, output_type) if isinstance(ax, list): - assert all(isinstance(ax_, plt.Axes) for ax_ in input_ax) + assert all(isinstance(ax_, plt.Axes) for ax_ in ax) if ax is None: assert ax.name == projection @pytest.mark.parametrize( - ("colorbar", "ax", "return_type"), - [(True, [plt.axes(), None], mpl.colorbar.Colorbar), - (True, [plt.axes(), plt.axes()], mpl.colorbar.Colorbar), - (False, [plt.axes(), None], None)], + ("colorbar", "ax_case", "return_type"), + [(True, "single", mpl.colorbar.Colorbar), + (True, "two", mpl.colorbar.Colorbar), (False, "single", None)], ) -def test_add_colorbar(colorbar, ax, return_type): +def test_add_colorbar(colorbar, ax_case, return_type): """Test return type of :py:func:`~spharpy.plot._utils._add_colorbar`.""" + if ax_case == "single": + fig, ax0 = plt.subplots() + ax = [ax0, None] + else: + fig, axs = plt.subplots(1, 2) + ax = [axs[0], axs[1]] + norm = mpl.colors.Normalize(vmin=0, vmax=10) cmap = plt.get_cmap('viridis') mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap) From 8325add9079cb13c4cba4f267eb9665bd1111b48 Mon Sep 17 00:00:00 2001 From: Anton Hoyer <156099087+hoyer-a@users.noreply.github.com> Date: Tue, 21 Oct 2025 21:45:40 +0200 Subject: [PATCH 09/10] extend assertion --- tests/test_plot__utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_plot__utils.py b/tests/test_plot__utils.py index f4f4c829..542bd832 100644 --- a/tests/test_plot__utils.py +++ b/tests/test_plot__utils.py @@ -29,7 +29,7 @@ def test_prepare_plot(ax_case, output_type, projection): if isinstance(ax, list): assert all(isinstance(ax_, plt.Axes) for ax_ in ax) - if ax is None: + if ax is None and projection is not None: assert ax.name == projection From a020f540fd46338db05ea2cfa9338981c66c7f17 Mon Sep 17 00:00:00 2001 From: Anton Hoyer <156099087+hoyer-a@users.noreply.github.com> Date: Wed, 22 Oct 2025 10:09:23 +0200 Subject: [PATCH 10/10] apply review --- spharpy/plot/_utils.py | 43 +++++---------------------------------- tests/test_plot__utils.py | 30 +-------------------------- 2 files changed, 6 insertions(+), 67 deletions(-) diff --git a/spharpy/plot/_utils.py b/spharpy/plot/_utils.py index 8f9ab88b..750857d1 100644 --- a/spharpy/plot/_utils.py +++ b/spharpy/plot/_utils.py @@ -4,42 +4,7 @@ import matplotlib.pyplot as plt -def _add_colorbar(colorbar, fig, ax, mappable, label): - """ - Add colorbar to plot. - - Parameters - ---------- - colorbar : bool - Flag indicating if a colobar should be added to the plot - fig : matplotlib.figure.Figure - Figure to plot on. - ax : list[matplotlib.axes.Axes] - Either a list of two axes objects or a list with one axis and None - object. - mappable : matplotlib.cm.ScalarMappable - The matplotlib.cm.ScalarMappable described by the colorbar. - label : string - colorbar label - - Returns - ------- - cb : matplotlib.colorbar.Colorbar - Returns matplotlib colorbar object. - """ - if colorbar: - if ax[1] is None: - cb = fig.colorbar(mappable, ax=ax[0]) - else: - cb = fig.colorbar(mappable, ax=ax[1]) - cb.set_label(label) - else: - cb = None - - return cb - - -def _prepare_plot(ax, projection=None): +def _prepare_plot(ax=None, projection=None): """ Returns a figure to plot on. @@ -51,7 +16,9 @@ def _prepare_plot(ax, projection=None): not exist. projection : str, optional Type of projection for the axes. This is only applied if new axes are - created. Default is ``None`` (2D projection). + created. Default is ``None`` (2D projection). See + `matplotlib.projections `_ + for more information on projections. Returns ------- @@ -59,7 +26,7 @@ def _prepare_plot(ax, projection=None): Returns the active figure. ax : maptlotlib.axes.Axes Returns the current axes. - """ + """ # noqa: E501 if ax is None: # get current figure or create a new one fig = plt.gcf() diff --git a/tests/test_plot__utils.py b/tests/test_plot__utils.py index 542bd832..9d39b8a3 100644 --- a/tests/test_plot__utils.py +++ b/tests/test_plot__utils.py @@ -1,7 +1,6 @@ import pytest import matplotlib.pyplot as plt -import matplotlib as mpl -from spharpy.plot._utils import _prepare_plot, _add_colorbar +from spharpy.plot._utils import _prepare_plot @pytest.mark.parametrize( @@ -31,30 +30,3 @@ def test_prepare_plot(ax_case, output_type, projection): if ax is None and projection is not None: assert ax.name == projection - - -@pytest.mark.parametrize( - ("colorbar", "ax_case", "return_type"), - [(True, "single", mpl.colorbar.Colorbar), - (True, "two", mpl.colorbar.Colorbar), (False, "single", None)], -) -def test_add_colorbar(colorbar, ax_case, return_type): - """Test return type of :py:func:`~spharpy.plot._utils._add_colorbar`.""" - if ax_case == "single": - fig, ax0 = plt.subplots() - ax = [ax0, None] - else: - fig, axs = plt.subplots(1, 2) - ax = [axs[0], axs[1]] - - norm = mpl.colors.Normalize(vmin=0, vmax=10) - cmap = plt.get_cmap('viridis') - mappable = mpl.cm.ScalarMappable(norm=norm, cmap=cmap) - - fig = plt.gcf() - cb = _add_colorbar(colorbar, fig, ax, mappable, None) - - if return_type is None: - assert cb is None - else: - assert isinstance(cb, return_type)