one function in datasets submodule with better docs

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Unreleased
 
 - `base` and `utils` modules are now tested directly.
+- Better docs of `datasets` submodule and the three functions replaced in only one.
 
 
 ## v0.3.1 (2024-01-29)

galpynostatic/datasets/__init__.py

@@ -48,16 +48,26 @@
 # ============================================================================
 
 
-def load_planar():
-    """Galvanostatic planar map for a cut-off potential of 150 mV."""
-    return pd.read_csv(PATH / "planar.csv")
-
-
-def load_cylindrical():
-    """Galvanostatic cylindrical map for a cut-off potential of 150 mV."""
-    return pd.read_csv(PATH / "cylindrical.csv")
-
-
-def load_spherical():
-    """Galvanostatic spherical map for a cut-off potential of 150 mV."""
-    return pd.read_csv(PATH / "spherical.csv")
+def load_dataset(geometry="spherical"):
+    """Galvanostatic map for a cut-off potential of 150 mV.
+
+    Parameters
+    ----------
+    geometry : str, default="spherical"
+        The geometry of the electrode. It can be `"spherical"`, `"cylindrical"`
+        or `"planar"`.
+
+    Returns
+    -------
+    pandas.DataFrame
+        The geometry dataset as a ``pandas.DataFrame``.
+
+    Raises
+    ------
+    ValueError
+        If the geometry is not `"spherical"`, `"cylindrical"` or `"planar"`.
+    """
+    if geometry not in ("spherical", "cylindrical", "planar"):
+        raise ValueError(f"{geometry} is not a valid geometry.")
+
+    return pd.read_csv(PATH / f"{geometry}.csv")

galpynostatic/model.py

@@ -28,7 +28,7 @@
 from sklearn.utils import validation as _skl_validation
 
 from .base import MapSpline
-from .datasets import load_cylindrical, load_planar, load_spherical
+from .datasets import load_dataset
 from .plot import GalvanostaticPlotter
 from .utils import logell, logxi
 
@@ -158,14 +158,8 @@ def __init__(
     def _validate_geometry(self):
         """Validate geometry (when dataset is a string)."""
         if isinstance(self.dataset, str):
-            load_geometry = {
-                "planar": load_planar,
-                "cylindrical": load_cylindrical,
-                "spherical": load_spherical,
-            }
-
-            if self.dataset in load_geometry:
-                self.dataset = load_geometry[self.dataset]()
+            if self.dataset in ("spherical", "cylindrical", "planar"):
+                self.dataset = load_dataset(geometry=self.dataset)
             else:
                 raise ValueError(f"{self.dataset} is not a valid geometry.")
 

galpynostatic/plot.py

@@ -66,6 +66,11 @@ def render_map(self, ax=None, clb=True, clb_label="maximum SOC"):
 
         clb_label : str, default="maximum SOC"
             The label for the colorbar.
+
+        Returns
+        -------
+        ax : matplotlib.axes.Axes
+            The current axes.
         """
         ax = plt.gca() if ax is None else ax
 

tests/conftest.py

@@ -32,19 +32,9 @@ def data_path():
     )
 
 
-@pytest.fixture()
-def planar():
-    return galpynostatic.datasets.load_planar()
-
-
-@pytest.fixture()
-def cylindrical():
-    return galpynostatic.datasets.load_cylindrical()
-
-
 @pytest.fixture()
 def spherical():
-    return galpynostatic.datasets.load_spherical()
+    return galpynostatic.datasets.load_dataset()
 
 
 @pytest.fixture()

tests/test_datasets.py

@@ -17,6 +17,8 @@
 
 import pandas as pd
 
+import pytest
+
 
 # =============================================================================
 # TESTS
@@ -25,7 +27,7 @@
 
 def test_load_planar():
     """Test the planar dataset."""
-    pla = galpynostatic.datasets.load_planar()
+    pla = galpynostatic.datasets.load_dataset(geometry="planar")
 
     assert isinstance(pla, pd.DataFrame)
 
@@ -44,7 +46,7 @@ def test_load_planar():
 
 def test_load_cylindrical():
     """Test the cylindrical dataset."""
-    cyl = galpynostatic.datasets.load_cylindrical()
+    cyl = galpynostatic.datasets.load_dataset(geometry="cylindrical")
 
     assert isinstance(cyl, pd.DataFrame)
 
@@ -63,7 +65,7 @@ def test_load_cylindrical():
 
 def test_load_spherical():
     """Test the spherical dataset."""
-    sph = galpynostatic.datasets.load_spherical()
+    sph = galpynostatic.datasets.load_dataset()
 
     assert isinstance(sph, pd.DataFrame)
 
@@ -78,3 +80,9 @@ def test_load_spherical():
     np.testing.assert_almost_equal(sph.xmax.min(), 0.0, 6)
     np.testing.assert_almost_equal(sph.xmax.max(), 0.99706, 6)
     np.testing.assert_almost_equal(sph.xmax.mean(), 0.518575, 6)
+
+
+def test_raise():
+    """Test the raise of the ValueError."""
+    with pytest.raises(ValueError):
+        galpynostatic.datasets.load_dataset(geometry="plane")