Custom colour maps

README.md

@@ -81,6 +81,11 @@ Set a date range:
 python geo_heatmap.py --min-date 2017-01-02 --max-date 2018-12-30 Records.json
 ```
 
+Change the color map:
+```shell
+python geo_heatmap.py -c Set1_09 Records.json
+```
+
 Advanced heatmap settings:
 
 ```shell
@@ -123,6 +128,9 @@ optional arguments:
                         The minimum opacity of the heatmap. (default: 0.2)
   -mz MAX_ZOOM, --max-zoom MAX_ZOOM
                         The maximum zoom of the heatmap. (default: 4)
+  -c COLOR_MAP, --color-map COLOR_MAP
+                        The color map to use, from:
+                        branca (<name>), matplotlib (mpl.<name>), cmcrameri (cmc.<name>)
 ```
 
 ### 6. Review the Results

geo_heatmap.py

@@ -208,6 +208,7 @@ def generateMap(self, settings):
         blur = settings["blur"]
         min_opacity = settings["min_opacity"]
         max_zoom = settings["max_zoom"]
+        color_map = settings["color_map"]
 
         map_data = [(coords[0], coords[1], magnitude)
                     for coords, magnitude in self.coordinates.items()]
@@ -219,12 +220,14 @@ def generateMap(self, settings):
                        attr="<a href=https://github.com/luka1199/geo-heatmap>geo-heatmap</a>")
 
         # Generate heat map
+        gradient = getColormapAsGradient(color_map) if color_map else None
         heatmap = HeatMap(map_data,
                           max_val=self.max_magnitude,
                           min_opacity=min_opacity,
                           radius=radius,
                           blur=blur,
-                          max_zoom=max_zoom)
+                          max_zoom=max_zoom,
+                          gradient=gradient)
 
         m.add_child(heatmap)
         return m
@@ -307,7 +310,9 @@ def run(self, data_files, output_file, date_range, stream_data, settings):
                         help="The minimum opacity of the heatmap. (default: %(default)s)", default=0.2)
     parser.add_argument("-mz", "--max-zoom", dest="max_zoom", type=int, required=False,
                         help="The maximum zoom of the heatmap. (default: %(default)s)", default=4)
-
+    parser.add_argument("-c", "--color-map", dest="color_map", type=str, required=False,
+                        help="The color map to use, from:\n" \
+                        "branca (<name>), matplotlib (mpl.<name>), cmcrameri (cmc.<name>).", default=None)
 
     args = parser.parse_args()
     data_file = args.files
@@ -320,7 +325,8 @@ def run(self, data_files, output_file, date_range, stream_data, settings):
         "radius": args.radius,
         "blur": args.blur,
         "min_opacity": args.min_opacity,
-        "max_zoom": args.max_zoom
+        "max_zoom": args.max_zoom,
+        "color_map": args.color_map,
     }
 
     generator = Generator()

utils.py

@@ -1,4 +1,5 @@
 from datetime import datetime, timezone
+import numpy as np
 import webbrowser
 
 
@@ -63,3 +64,48 @@ def dateInRange(date, date_range) -> bool:
     date = datetime.strptime(date, "%Y-%m-%d")
     return (min_date is None or min_date <= date) and \
         (max_date is None or max_date >= date)
+
+
+def getColormapAsGradient(name: str, n_steps: int=20) -> dict[int, str]:
+    """Returns the desired colormap in gradient format.
+
+    Colormaps are retrieved from branca.colormap.linear by default.
+    Names starting with mpl. are retrieved from matplotlib (optional import).
+    Names starting with cmc. are retrieved from cmcrameri (optional import).
+
+    Arguments:
+        name {str} -- Colormap name.
+        n_steps {int} -- Number of quantization steps.
+
+    Returns:
+        dict -- Keys are a range from 0 to 1, values are hex strings.
+    """
+    # Define steps
+    steps = np.linspace(0, 1, n_steps)
+
+    # Default behavior: Get from branca
+    if "." not in name or name.lower().startswith("branca."):
+        from branca.colormap import linear as cm
+        cmap = getattr(cm, name)
+        gradient = {i: cmap(i) for i in steps}
+
+    # Get from defined library
+    else:
+        lib, name = name.split(".")
+        lib = lib.lower()
+
+        # Matplotlib and derivatives
+        if lib in ("mpl", "cmc"):
+            from matplotlib.colors import to_hex
+            if lib == "mpl":
+                from matplotlib import colormaps as cm
+                cmap = cm.get_cmap(name)
+            elif lib == "cmc":
+                from cmcrameri import cm
+                cmap = getattr(cm, name)
+            gradient = {i: to_hex(cmap(i)) for i in steps}
+
+        else:
+            raise ValueError(f"Unknown colormap library '{lib}'.")
+
+    return gradient