Merge pull request GraphingLib#634 from yalap13/add-norm-param-heatmap

Added norm param to Heatmap

Author: yalap13

docs/release_notes/upcoming_changes/634.new_feature.rst

@@ -0,0 +1,3 @@
+`norm` parameter for `Heatmap`
+==============================
+Added a `norm` parameter for `Heatmap` as exists for `imshow`.

graphinglib/data_plotting_2d.py

@@ -6,7 +6,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from matplotlib.colors import Colormap
+from matplotlib.colors import Colormap, Normalize
 from matplotlib.image import imread
 from numpy.typing import ArrayLike
 from scipy.interpolate import griddata
@@ -67,6 +67,8 @@ class Heatmap(Plottable2D):
 
             For other interpolation methods, refer to
             `Interpolations for imshow <https://matplotlib.org/stable/gallery/images_contours_and_fields/interpolation_methods.html>`_.
+    norm : str or Normalize, optional
+        Normalization of the colormap. Default is ``None``.
     """
 
     def __init__(
@@ -81,6 +83,7 @@ def __init__(
         aspect_ratio: str | float = "default",
         origin_position: str = "default",
         interpolation: str = "none",
+        norm: Optional[str | Normalize] = None,
     ) -> None:
         """
         The class implements heatmaps.
@@ -119,6 +122,8 @@ def __init__(
 
                 For other interpolation methods, refer to
                 `Interpolations for imshow <https://matplotlib.org/stable/gallery/images_contours_and_fields/interpolation_methods.html>`_.
+        norm : str or Normalize, optional
+            Normalization of the colormap. Default is ``None``.
         """
         self._image = image
         self._x_axis_range = x_axis_range
@@ -130,6 +135,7 @@ def __init__(
         self._aspect_ratio = aspect_ratio
         self._origin_position = origin_position
         self._interpolation = interpolation
+        self._norm = norm
 
         self._color_bar_params: dict = {}
 
@@ -153,6 +159,7 @@ def from_function(
         origin_position: str = "default",
         interpolation: str = "none",
         number_of_points: tuple[int, int] = (50, 50),
+        norm: Optional[str | Normalize] = None,
     ) -> Self:
         """
         Creates a heatmap from a function.
@@ -194,6 +201,8 @@ def from_function(
         number_of_points : tuple[int, int]
             Number of points in the x and y coordinates.
             Defaults to ``(50, 50)``.
+        norm : str or Normalize, optional
+            Normalization of the colormap. Default is ``None``.
 
         Returns
         -------
@@ -214,6 +223,7 @@ def from_function(
             aspect_ratio,
             origin_position,
             interpolation,
+            norm,
         )
 
     @classmethod
@@ -233,6 +243,7 @@ def from_points(
         origin_position: str = "default",
         interpolation: str = "none",
         number_of_points: tuple[int, int] = (50, 50),
+        norm: Optional[str | Normalize] = None,
     ):
         """
         Creates a heatmap by interpolating unevenly distributed data points on a grid.
@@ -278,6 +289,8 @@ def from_points(
         number_of_points : tuple[int, int]
             Number of points in the x and y coordinates.
             Defaults to ``(50, 50)``.
+        norm : str or Normalize, optional
+            Normalization of the colormap. Default is ``None``.
 
         Returns
         -------
@@ -304,6 +317,7 @@ def from_points(
             aspect_ratio,
             origin_position,
             interpolation,
+            norm,
         )
 
     @property
@@ -445,6 +459,7 @@ def _plot_element(self, axes: plt.Axes, z_order: int, **kwargs) -> None:
             "origin": self._origin_position,
             "interpolation": self._interpolation,
             "extent": self._xy_range,
+            "norm": self._norm,
         }
 
         params = {k: v for k, v in params.items() if v != "default"}