ssb-sgis 1.0.1__py3-none-any.whl → 1.0.3__py3-none-any.whl

sgis/maps/legend.py

@@ -4,18 +4,21 @@ The Legend class is best accessed through the 'legend' attribute of the Thematic
 class.
 
 """
+
+import itertools
 import warnings
+from typing import Any
 
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
 import pandas as pd
+from geopandas import GeoDataFrame
 from matplotlib.lines import Line2D
 from pandas import Series
 
 from ..geopandas_tools.bounds import points_in_bounds
 
-
 # the geopandas._explore raises a deprication warning. Ignoring for now.
 warnings.filterwarnings(
     action="ignore", category=matplotlib.MatplotlibDeprecationWarning
@@ -23,6 +26,82 @@ warnings.filterwarnings(
 pd.options.mode.chained_assignment = None
 
 
+LEGEND_KWARGS = {
+    "title",
+    "size",
+    "position",
+    "fontsize",
+    "title_fontsize",
+    "markersize",
+    "framealpha",
+    "edgecolor",
+    "kwargs",
+    "labelspacing",
+    "title_color",
+    "width",
+    "height",
+    "labels",
+    "pretty_labels",
+    "thousand_sep",
+    "decimal_mark",
+    "label_sep",
+    "label_suffix",
+    "rounding",
+    "facecolor",
+    "labelcolor",
+}
+
+LOWERCASE_WORDS = {
+    "a",
+    "an",
+    "and",
+    "as",
+    "at",
+    "but",
+    "by",
+    "for",
+    "in",
+    "nor",
+    "of",
+    "on",
+    "or",
+    "the",
+    "up",
+}
+
+
+def prettify_label(label: str) -> str:
+    """Replace underscores with spaces and capitalize words that are all lowecase."""
+    return " ".join(
+        word.title() if word.islower() and word not in LOWERCASE_WORDS else word
+        for word in label.replace("_", " ").split()
+    )
+
+
+def prettify_number(x: int | float, rounding: int) -> int:
+    rounding = int(float(f"1e+{abs(rounding)}"))
+    rounded_down = int(x // rounding * rounding)
+    rounded_up = rounded_down + rounding
+    diff_up = abs(x - rounded_up)
+    diff_down = abs(x - rounded_down)
+    if diff_up < diff_down:
+        return rounded_up
+    else:
+        return rounded_down
+
+
+def prettify_bins(bins: list[int | float], rounding: int) -> list[int]:
+    return [
+        (
+            prettify_number(x, rounding)
+            if i != len(bins) - 1
+            else int(x)
+            # else prettify_number(x, rounding) + abs(rounding)
+        )
+        for i, x in enumerate(bins)
+    ]
+
+
 class Legend:
     """Holds the general attributes of the legend in the ThematicMap class.
 
@@ -33,78 +112,65 @@ class Legend:
     If a numeric column is used, additional attributes can be found in the
     ContinousLegend class.
 
-    Attributes:
-        title: Legend title. Defaults to the column name if used in the
-            ThematicMap class.
-        position: The legend's x and y position in the plot, specified as a tuple of
-            x and y position between 0 and 1. E.g. position=(0.8, 0.2) for a position
-            in the bottom right corner, (0.2, 0.8) for the upper left corner.
-        fontsize: Text size of the legend labels. Defaults to the size of
-            the ThematicMap class.
-        title_fontsize: Text size of the legend title. Defaults to the
-            size * 1.2 of the ThematicMap class.
-        markersize: Size of the color circles in the legend. Defaults to the size of
-            the ThematicMap class.
-        framealpha: Transparency of the legend background.
-        edgecolor: Color of the legend border. Defaults to #0f0f0f (almost black).
-        kwargs: Stores additional keyword arguments taken by the matplotlib legend
-            method. Specify this as e.g. m.legend.kwargs["labelcolor"] = "red", where
-            'm' is the name of the ThematicMap instance. See here:
-            https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.legend.html
-
-    Examples
-    --------
-    Create ten random points with a numeric column from 0 to 9.
+    Examples:
+    ---------
+    Create ten points with a numeric column from 0 to 9.
 
     >>> import sgis as sg
-    >>> points = sg.random_points(10)
+    >>> points = sg.to_gdf(
+    ...     [
+    ...         (0, 1),
+    ...         (1, 0),
+    ...         (1, 1),
+    ...         (0, 0),
+    ...         (0.5, 0.5),
+    ...         (0.5, 0.25),
+    ...         (0.25, 0.25),
+    ...         (0.75, 0.75),
+    ...         (0.25, 0.75),
+    ...         (0.75, 0.25),
+    ...     ]
+    ... )
     >>> points["number"] = range(10)
     >>> points
-                    geometry  number
-    0  POINT (0.59780 0.50425)       0
-    1  POINT (0.07019 0.26167)       1
-    2  POINT (0.56475 0.15422)       2
-    3  POINT (0.87293 0.60316)       3
-    4  POINT (0.47373 0.20040)       4
-    5  POINT (0.98661 0.15614)       5
-    6  POINT (0.30951 0.77057)       6
-    7  POINT (0.47802 0.52824)       7
-    8  POINT (0.12215 0.96588)       8
-    9  POINT (0.02938 0.93467)       9
+                      geometry  number
+    0  POINT (0.00000 1.00000)       0
+    1  POINT (1.00000 0.00000)       1
+    2  POINT (1.00000 1.00000)       2
+    3  POINT (0.00000 0.00000)       3
+    4  POINT (0.50000 0.50000)       4
+    5  POINT (0.50000 0.25000)       5
+    6  POINT (0.25000 0.25000)       6
+    7  POINT (0.75000 0.75000)       7
+    8  POINT (0.25000 0.75000)       8
+    9  POINT (0.75000 0.25000)       9
 
     Creating the ThematicMap instance will also create the legend. Since we
-    specify a numeric column, a ContinousLegend instance is created.
-
-    >>> m = sg.ThematicMap(points, column="number")
+    pass a numeric column, a ContinousLegend is created.
+
+    >>> m = sg.ThematicMap(
+    ...     points,
+    ...     column="number"
+    ...     legend_kwargs=dict(
+    ...         title="Meters",
+    ...         label_sep="to",
+    ...         label_suffix="num",
+    ...         rounding=2,
+    ...         position = (0.35, 0.28),
+    ...         title_fontsize=11,
+    ...         fontsize=9,
+    ...         markersize=7.5,
+    ...     ),
+    ... )
+    >>> m.plot()
     >>> m.legend
     <sgis.maps.legend.ContinousLegend object at 0x00000222206738D0>
-
-    Changing the attributes that apply to both numeric and categorical columns.
-
-    >>> m.legend.title = "Meters"
-    >>> m.legend.title_fontsize = 11
-    >>> m.legend.fontsize = 9
-    >>> m.legend.markersize = 7.5
-    >>> m.legend.position = (0.35, 0.28)
-    >>> m.plot()
-
-    Additional matplotlib keyword arguments can be specified as kwargs.
-
-    >>> m.legend.kwargs["labelcolor"] = "red"
-
-    Since we are using a numeric column, the legend is of type ContinousLegend.
-    We can therefore also access the attributes that only apply to numeric columns.
-
-    >>> m.label_sep = "to"
-    >>> m.label_suffix = "num"
-    >>> m.rounding = 2
-    >>> m.plot()
-
     """
 
     def __init__(
         self,
         title: str | None = None,
+        pretty_labels: bool = True,
         labels: list[str] | None = None,
         position: tuple[float] | None = None,
         markersize: int | None = None,
@@ -113,7 +179,32 @@ class Legend:
         framealpha: float = 1.0,
         edgecolor: str = "#0f0f0f",
         **kwargs,
-    ):
+    ) -> None:
+        """Initialiser.
+
+        Args:
+            title: Legend title. Defaults to the column name if used in the
+                ThematicMap class.
+            pretty_labels: If True, words will be capitalized and underscores turned to spaces.
+                If continous values, numbers will be rounded.
+            labels: Labels of the categories.
+            position: The legend's x and y position in the plot, specified as a tuple of
+                x and y position between 0 and 1. E.g. position=(0.8, 0.2) for a position
+                in the bottom right corner, (0.2, 0.8) for the upper left corner.
+            fontsize: Text size of the legend labels. Defaults to the size of
+                the ThematicMap class.
+            title_fontsize: Text size of the legend title. Defaults to the
+                size * 1.2 of the ThematicMap class.
+            markersize: Size of the color circles in the legend. Defaults to the size of
+                the ThematicMap class.
+            framealpha: Transparency of the legend background.
+            edgecolor: Color of the legend border. Defaults to #0f0f0f (almost black).
+            kwargs: Stores additional keyword arguments taken by the matplotlib legend
+                method. Specify this as e.g. m.legend.kwargs["labelcolor"] = "red", where
+                'm' is the name of the ThematicMap instance. See here:
+                https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.legend.html
+
+        """
         self.title = title
 
         if "size" in kwargs:
@@ -124,6 +215,7 @@ class Legend:
             self._fontsize = fontsize
             self._markersize = markersize
 
+        self.pretty_labels = pretty_labels
         self.framealpha = framealpha
         self.edgecolor = edgecolor
         self.width = kwargs.pop("width", 0.1)
@@ -133,12 +225,25 @@ class Legend:
 
         self.labels = labels
         self._position = position
-        self.kwargs = kwargs
         self._position_has_been_set = True if position else False
 
-    def _get_legend_sizes(self, size, kwargs):
-        """Adjust fontsize and markersize to size kwarg."""
+        self.kwargs = {}
+        for key, value in kwargs.items():
+            if key not in LEGEND_KWARGS:
+                self.kwargs[key] = value
+            else:
+                try:
+                    setattr(self, key, value)
+                except Exception:
+                    setattr(self, f"_{key}", value)
 
+    @property
+    def valid_keywords(self) -> set[str]:
+        """List all valid keywords for the class initialiser."""
+        return LEGEND_KWARGS
+
+    def _get_legend_sizes(self, size: int | float, kwargs: dict) -> None:
+        """Adjust fontsize and markersize to size kwarg."""
         if "title_fontsize" in kwargs:
             self._title_fontsize = kwargs["title_fontsize"]
             self._title_fontsize_has_been_set = True
@@ -157,7 +262,9 @@ class Legend:
         else:
             self._markersize = size
 
-    def _prepare_categorical_legend(self, categories_colors: dict, nan_label: str):
+    def _prepare_categorical_legend(
+        self, categories_colors: dict, nan_label: str
+    ) -> None:
         for attr in self.__dict__.keys():
             if attr in self.kwargs:
                 self[attr] = self.kwargs.pop(attr)
@@ -178,6 +285,8 @@ class Legend:
 
         self._patches, self._categories = [], []
         for category, color in categories_colors.items():
+            if self.pretty_labels:
+                category = prettify_label(category)
             if category == nan_label:
                 self._categories.append(nan_label)
             else:
@@ -195,7 +304,9 @@ class Legend:
                 )
             )
 
-    def _actually_add_legend(self, ax):
+    def _actually_add_legend(self, ax: matplotlib.axes.Axes) -> matplotlib.axes.Axes:
+        if self.pretty_labels:
+            self.title = prettify_label(self.title)
         legend = ax.legend(
             self._patches,
             self._categories,
@@ -215,7 +326,9 @@ class Legend:
 
         return ax
 
-    def _get_best_legend_position(self, gdf, k: int):
+    def _get_best_legend_position(
+        self, gdf: GeoDataFrame, k: int
+    ) -> tuple[float, float]:
         minx, miny, maxx, maxy = gdf.total_bounds
         diffx = maxx - minx
         diffy = maxy - miny
@@ -243,51 +356,58 @@ class Legend:
 
         return bestx_01, besty_01
 
-    def __getitem__(self, item):
+    def __getitem__(self, item: str) -> Any:
+        """Get attribute with square brackets."""
         return getattr(self, item)
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key: str, value: Any) -> None:
+        """Set attribute with square brackets."""
         setattr(self, key, value)
 
-    def get(self, key, default=None):
+    def get(self, key: Any, default: Any = None) -> Any:
+        """Get value of an attribute of the Legend."""
         try:
             return self[key]
         except (KeyError, ValueError, IndexError, AttributeError):
             return default
 
     @property
-    def position(self):
+    def position(self) -> tuple[float, float]:
+        """Legend position in x, y."""
         return self._position
 
     @position.setter
-    def position(self, new_value: bool):
+    def position(self, new_value: tuple[float, float]) -> None:
         self._position = new_value
         self._position_has_been_set = True
 
     @property
-    def title_fontsize(self):
+    def title_fontsize(self) -> int:
+        """Legend title fontsize."""
         return self._title_fontsize
 
     @title_fontsize.setter
-    def title_fontsize(self, new_value: bool):
+    def title_fontsize(self, new_value: int) -> None:
         self._title_fontsize = new_value
         self._title_fontsize_has_been_set = True
 
     @property
-    def fontsize(self):
+    def fontsize(self) -> int:
+        """Legend fontsize."""
         return self._fontsize
 
     @fontsize.setter
-    def fontsize(self, new_value: bool):
+    def fontsize(self, new_value: int) -> None:
         self._fontsize = new_value
         self._fontsize_has_been_set = True
 
     @property
-    def markersize(self):
+    def markersize(self) -> int:
+        """Legend markersize."""
         return self._markersize
 
     @markersize.setter
-    def markersize(self, new_value: bool):
+    def markersize(self, new_value: int) -> None:
         self._markersize = new_value
         self._markersize_has_been_set = True
 
@@ -311,18 +431,14 @@ class ContinousLegend(Legend):
             Defaults to None.
         label_sep: Text to put in between the two numbers in each color group in
             the legend. Defaults to '-'.
-        rounding: Number of decimals in the labels. By default, the rounding
-            depends on the column's maximum value and standard deviation.
-            OBS: The bins will not be rounded, meaning the labels might be wrong
-            if not bins are set manually.
         thousand_sep: Separator between each thousand for large numbers. Defaults to
             None, meaning no separator.
         decimal_mark: Text to use as decimal point. Defaults to None, meaning '.' (dot)
             unless 'thousand_sep' is '.'. In this case, ',' (comma) will be used as
             decimal mark.
 
-    Examples
-    --------
+    Examples:
+    ---------
     Create ten random points with a numeric column from 0 to 9.
 
     >>> import sgis as sg
@@ -379,14 +495,43 @@ class ContinousLegend(Legend):
     def __init__(
         self,
         labels: list[str] | None = None,
-        pretty_labels: bool = False,
+        pretty_labels: bool = True,
         label_suffix: str | None = None,
         label_sep: str = "-",
         rounding: int | None = None,
         thousand_sep: str | None = None,
         decimal_mark: str | None = None,
         **kwargs,
-    ):
+    ) -> None:
+        """Initialiser.
+
+        Args:
+            labels: To manually set labels. If set, all other labeling attributes are
+                ignored. Should be given as a list of strings with the same length as
+                the number of color groups.
+            pretty_labels: If False (default), the minimum and maximum values of each
+                color group will be used as legend labels. If True, the labels will end
+                with the maximum value, but start at 1 + the maximum value of the previous
+                group. The labels will be correct but inaccurate.
+            label_suffix: The text to put after each number in the legend labels.
+                Defaults to None.
+            label_sep: Text to put in between the two numbers in each color group in
+                the legend. Defaults to '-'.
+            rounding: Number of decimals in the labels. By default, the rounding
+                depends on the column's maximum value and standard deviation.
+                OBS: The bins will not be rounded, meaning the labels might be wrong
+                if not bins are set manually.
+            thousand_sep: Separator between each thousand for large numbers. Defaults to
+                None, meaning no separator.
+            decimal_mark: Text to use as decimal point. Defaults to None, meaning '.' (dot)
+                unless 'thousand_sep' is '.'. In this case, ',' (comma) will be used as
+                decimal mark.
+            kwargs: Stores additional keyword arguments taken by the matplotlib legend
+                method. Specify this as e.g. m.legend.kwargs["labelcolor"] = "red", where
+                'm' is the name of the ThematicMap instance. See here:
+                https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.legend.html
+
+        """
         super().__init__(**kwargs)
 
         self.pretty_labels = pretty_labels
@@ -396,7 +541,7 @@ class ContinousLegend(Legend):
         self.label_sep = label_sep
         self.label_suffix = "" if not label_suffix else label_suffix
         self._rounding = rounding
-        self._rounding_has_been_set = True if rounding else False
+        # self._rounding_has_been_set = True if rounding else False
 
     def _get_rounding(self, array: Series | np.ndarray) -> int:
         def isinteger(x):
@@ -423,13 +568,15 @@ class ContinousLegend(Legend):
         return int(abs(np.log10(std_))) + 1
 
     @staticmethod
-    def _set_rounding(bins, rounding: int | float):
-        if rounding == 0:
-            return [int(round(bin, 0)) for bin in bins]
+    def _set_rounding(bins, rounding: int | float) -> list[int | float]:
+        if not rounding:
+            return [int(round(bin_, 0)) for bin_ in bins]
+        elif rounding <= 0:
+            return [int(round(bin_, rounding)) for bin_ in bins]
         else:
-            return [round(bin, rounding) for bin in bins]
+            return [round(bin_, rounding) for bin_ in bins]
 
-    def _remove_max_legend_value(self):
+    def _remove_max_legend_value(self) -> None:
         if not self._legend:
             raise ValueError("Cannot modify legend before it is created.")
 
@@ -439,11 +586,11 @@ class ContinousLegend(Legend):
         colors: list[str],
         nan_label: str,
         bin_values: dict,
-    ):
+    ) -> None:
         # TODO: clean up this messy method
 
-        for attr in self.__dict__.keys():
-            if attr in self.kwargs:
+        for attr in self.kwargs:
+            if attr in self.__dict__:
                 self[attr] = self.kwargs.pop(attr)
 
         self._patches, self._categories = [], []
@@ -466,9 +613,10 @@ class ContinousLegend(Legend):
             if len(self.labels) != len(colors):
                 raise ValueError(
                     "Label list must be same length as the number of groups. "
-                    f"Got k={len(colors)} and labels={len(colors)}."
+                    f"Got k={len(colors)} and labels={len(self.labels)}."
                     f"labels: {', '.join(self.labels)}"
                     f"colors: {', '.join(colors)}"
+                    f"bins: {bins}"
                 )
             self._categories = self.labels
 
@@ -491,72 +639,69 @@ class ContinousLegend(Legend):
                     )
 
         else:
-            for i, (cat1, cat2) in enumerate(zip(bins[:-1], bins[1:], strict=True)):
+            for i, (cat1, cat2) in enumerate(itertools.pairwise(bins)):
                 if nan_label in str(cat1) or nan_label in str(cat2):
                     self._categories.append(nan_label)
                     continue
 
                 min_ = np.min(bin_values[i])
                 max_ = np.max(bin_values[i])
-                min_rounded = self._set_rounding([min_], self._rounding)[0]
-                max_rounded = self._set_rounding([max_], self._rounding)[0]
+
                 if self.pretty_labels:
-                    if i != 0 and self._rounding == 0:
-                        cat1 = int(cat1 + 1)
-                    elif i != 0:
-                        cat1 = cat1 + float(f"1e-{self._rounding}")
+                    if i == 0:
+                        cat1 = int(min_) if (self.rounding or 0) <= 0 else min_
+
+                    is_last = i == len(bins) - 2
+                    if is_last:
+                        cat2 = int(max_) if (self.rounding or 0) <= 0 else max_
+
+                    if (self.rounding or 0) <= 0:
+                        cat1 = int(cat1)
+                        cat2 = int(cat2 - 1) if not is_last else int(cat2)
+                    elif (self.rounding or 0) > 0:
+                        cat1 = round(cat1, self._rounding)
+                        cat2 = round(
+                            cat2 - float(f"1e-{self._rounding}"), self._rounding
+                        )
+                    else:
+                        cat1 = round(cat1, self._rounding)
+                        cat2 = round(cat2, self._rounding)
 
                     cat1 = self._format_number(cat1)
                     cat2 = self._format_number(cat2)
 
                     if min_ == max_:
-                        label = self._two_value_label(cat1, cat2)
+                        label = self._get_two_value_label(cat1, cat2)
                         self._categories.append(label)
                         continue
 
-                    label = self._two_value_label(cat1, cat2)
+                    label = self._get_two_value_label(cat1, cat2)
                     self._categories.append(label)
 
-                elif min_ == max_:
+                    continue
+
+                min_rounded = self._set_rounding([min_], self._rounding)[0]
+                max_rounded = self._set_rounding([max_], self._rounding)[0]
+                if min_ == max_:
                     min_rounded = self._format_number(min_rounded)
-                    label = self._one_value_label(min_rounded)
+                    label = self._get_one_value_label(min_rounded)
                     self._categories.append(label)
                 else:
                     min_rounded = self._format_number(min_rounded)
                     max_rounded = self._format_number(max_rounded)
-                    label = self._two_value_label(min_rounded, max_rounded)
+                    label = self._get_two_value_label(min_rounded, max_rounded)
                     self._categories.append(label)
 
-    def _actually_add_legend(self, ax):
-        legend = ax.legend(
-            self._patches,
-            self._categories,
-            fontsize=self._fontsize,
-            title=self.title,
-            title_fontsize=self._title_fontsize,
-            bbox_to_anchor=self._position + (self.width, self.height),
-            fancybox=False,
-            framealpha=self.framealpha,
-            edgecolor=self.edgecolor,
-            labelspacing=self.labelspacing,
-            **self.kwargs,
-        )
-
-        if self.title_color:
-            plt.setp(legend.get_title(), color=self.title_color)
-
-        return ax
-
-    def _two_value_label(self, value1, value2):
+    def _get_two_value_label(self, value1: int | float, value2: int | float) -> str:
         return (
             f"{value1} {self.label_suffix} {self.label_sep} "
             f"{value2} {self.label_suffix}"
         )
 
-    def _one_value_label(self, value1):
+    def _get_one_value_label(self, value1: int | float) -> str:
         return f"{value1} {self.label_suffix}"
 
-    def _format_number(self, number):
+    def _format_number(self, number: int | float) -> int | float:
         if not self.thousand_sep and not self.decimal_mark:
             return number
 
@@ -581,10 +726,16 @@ class ContinousLegend(Legend):
         return number
 
     @property
-    def rounding(self):
+    def rounding(self) -> int:
+        """Number of decimals in the labels.
+
+        By default, the rounding
+        depends on the column's maximum value and standard deviation.
+        OBS: The bins will not be rounded, meaning the labels might be wrong
+        if not bins are set manually.
+        """
         return self._rounding
 
     @rounding.setter
-    def rounding(self, new_value: bool):
+    def rounding(self, new_value: int) -> None:
         self._rounding = new_value
-        self._rounding_has_been_set = True