From 72525e325cbe3c3179429c5c41bf244b054a9097 Mon Sep 17 00:00:00 2001
From: Shailesh Shrestha <102210755+shailesh-stha@users.noreply.github.com>
Date: Fri, 29 Dec 2023 05:10:56 +0100
Subject: [PATCH] Modify add_velocity to support custom color list (#654)

* Enhanced LeafMap's add_velocity in leafmap.py to support custom color visualization by including a 'color_scale' parameter, allowing users to specify a list of colors for velocity vectors on the map.

* Update color_scale in add_velocity

---------

Co-authored-by: Qiusheng Wu <giswqs@gmail.com>
---
 examples/notebooks/52_netcdf.ipynb | 21 ++++++++++++++++++++-
 leafmap/leafmap.py                 | 22 ++++++++++++++++++++++
 2 files changed, 42 insertions(+), 1 deletion(-)

diff --git a/examples/notebooks/52_netcdf.ipynb b/examples/notebooks/52_netcdf.ipynb
index 29ddec1398..8afcb226af 100644
--- a/examples/notebooks/52_netcdf.ipynb
+++ b/examples/notebooks/52_netcdf.ipynb
@@ -172,7 +172,14 @@
             "source": [
                 "m = leafmap.Map(layers_control=True)\n",
                 "m.add_basemap('CartoDB.DarkMatter')\n",
-                "m.add_velocity(filename, zonal_speed='u_wind', meridional_speed='v_wind')\n",
+                "m.add_velocity(filename,\n",
+                "               zonal_speed='u_wind',\n",
+                "               meridional_speed='v_wind',\n",
+                "               color_scale = [\"rgb(0,0,150)\",\n",
+                "                              \"rgb(0,150,0)\",\n",
+                "                              \"rgb(255,255,0)\",\n",
+                "                              \"rgb(255,165,0)\",\n",
+                "                              \"rgb(150,0,0)\"])\n",
                 "m"
             ]
         },
@@ -189,6 +196,18 @@
             "display_name": "Python 3",
             "language": "python",
             "name": "python3"
+        },
+        "language_info": {
+            "codemirror_mode": {
+                "name": "ipython",
+                "version": 3
+            },
+            "file_extension": ".py",
+            "mimetype": "text/x-python",
+            "name": "python",
+            "nbconvert_exporter": "python",
+            "pygments_lexer": "ipython3",
+            "version": "3.8.18"
         }
     },
     "nbformat": 4,
diff --git a/leafmap/leafmap.py b/leafmap/leafmap.py
index fcd4bc5080..6fabe96e74 100644
--- a/leafmap/leafmap.py
+++ b/leafmap/leafmap.py
@@ -3861,6 +3861,7 @@ def add_velocity(
         max_velocity=20,
         display_options={},
         name="Velocity",
+        color_scale=None
     ):
         """Add a velocity layer to the map.
 
@@ -3877,6 +3878,7 @@ def add_velocity(
             max_velocity (int, optional): The maximum velocity to display. Defaults to 20.
             display_options (dict, optional): The display options for the velocity layer. Defaults to {}. See https://bit.ly/3uf8t6w.
             name (str, optional): Layer name to use . Defaults to 'Velocity'.
+            color_scale (list, optional): List of RGB color values for the velocity vector color scale. Defaults to []. See https://bit.ly/3uf8t6w.
 
         Raises:
             ImportError: If the xarray package is not installed.
@@ -3911,6 +3913,25 @@ def add_velocity(
         if level_dimension in coords:
             ds = ds.isel(drop=True, **params)
 
+        if color_scale is None:
+            color_scale=[
+                "rgb(36,104, 180)",
+                "rgb(60,157, 194)",
+                "rgb(128,205,193)",
+                "rgb(151,218,168)",
+                "rgb(198,231,181)",
+                "rgb(238,247,217)",
+                "rgb(255,238,159)",
+                "rgb(252,217,125)",
+                "rgb(255,182,100)",
+                "rgb(252,150,75)",
+                "rgb(250,112,52)",
+                "rgb(245,64,32)",
+                "rgb(237,45,28)",
+                "rgb(220,24,32)",
+                "rgb(180,0,35)"
+            ]
+        
         wind = Velocity(
             data=ds,
             zonal_speed=zonal_speed,
@@ -3921,6 +3942,7 @@ def add_velocity(
             max_velocity=max_velocity,
             display_options=display_options,
             name=name,
+            color_scale=color_scale,
         )
         self.add(wind)