DAY 1 #1

Tried to match 2D matplotlib colormaps to 3D mayavi color lookup tables

discretize_showcase.py

@@ -8,6 +8,7 @@
 from mayavi import mlab
 
 from functions_2D import interpolate_discretized_data, plot_contours, isolate_contour_datapoints
+from helpers.colors import DEMScalarMappable
 
 """
 This script creates 3D graphics to showcase data discretization
@@ -37,30 +38,21 @@
 contour_levels = levels
 
 # Set up Colormaps
-cmap = plt.get_cmap("terrain")
-vmin = -0.25 * maximum * 1.1
-vmax = maximum * 1.1
-norm = Normalize(vmin, vmax)  # Leave extra 10% for interpolation overshoot
-colors = ScalarMappable(norm=norm, cmap=cmap)
-contour_colors = ScalarMappable(
-    norm=BoundaryNorm(
-        contour_levels,
-        contour_levels.size - 1,
-    ),
-    cmap=LinearSegmentedColormap.from_list("", colors.to_rgba(contour_levels[:-1]), N=contour_levels.size - 1),
-)
-discretized_lut = contour_colors.to_rgba(np.linspace(vmin, vmax, 255)) * 255
+vmin = minimum / 1.1 if minimum > 0 else minimum * 1.1  # Leave extra 10% for interpolation overshoot
+vmax = maximum * 1.1  # Leave extra 10% for interpolation overshoot
+colors = DEMScalarMappable(vmin, vmax)
 
 # Discretize Data
-discretized_datagrid = height_levels * np.floor(datagrid / height_levels)
+discretized_datagrid = height_levels * (np.floor(datagrid / height_levels))
 
 mlab.options.offscreen = True
 # 1
 mlab.figure(bgcolor=(1, 1, 1))
 surf = mlab.surf(
     y, x, np.rot90(discretized_datagrid.T), warp_scale=warp_scale, colormap="terrain", vmin=vmin, vmax=vmax
 )
-surf.module_manager.scalar_lut_manager.lut.table = discretized_lut
+# Using colormap for contours is a dirty fix for mayavi using weird lookup table values
+surf.module_manager.scalar_lut_manager.lut.table = colors.segmented_lut(levels, kind="middle")
 mlab.draw()
 mlab.gcf().scene._lift()
 mlab.view(azimuth=azimuth, distance="auto")
@@ -69,7 +61,7 @@
 # 2
 mlab.figure(bgcolor=(1, 1, 1))
 surf = mlab.surf(y, x, np.rot90(datagrid.T), warp_scale=warp_scale, colormap="terrain", vmin=vmin, vmax=vmax)
-surf.module_manager.scalar_lut_manager.lut.table = discretized_lut
+surf.module_manager.scalar_lut_manager.lut.table = colors.segmented_lut(levels, kind="bottom")
 mlab.draw()
 mlab.gcf().scene._lift()
 mlab.view(azimuth=azimuth, distance="auto")
@@ -78,9 +70,10 @@
 # 3
 mlab.figure(bgcolor=(1, 1, 1))
 surf = mlab.surf(y, x, np.rot90(datagrid.T), warp_scale=warp_scale, colormap="terrain", vmin=vmin, vmax=vmax)
+surf.module_manager.scalar_lut_manager.lut.table = colors.lut
 contours = isolate_contour_datapoints(x, y, np.rot90(datagrid, 3), levels=levels, return_separate_contours=True)
 for contour in contours:
-    mlab.plot3d(
+    obj = mlab.plot3d(
         contour[:, 0],
         contour[:, 1],
         contour[:, 2] * warp_scale,
@@ -90,6 +83,7 @@
         vmin=vmin,
         vmax=vmax,
     )
+    obj.module_manager.scalar_lut_manager.lut.table = colors.lut
 """
 Contour alternative:
 
@@ -108,7 +102,6 @@
 mlab.gcf().scene._lift()
 mlab.view(azimuth=azimuth, distance="auto")
 mlab.savefig(f"images/discretize/Accurate.png", magnification=10)
-
 # Use ImageMagick to remove background from images and crop out fully transparent region.
 for image in os.listdir(f"images/discretize/"):
     os.system(f"convert images/discretize/{image} -transparent white -trim +repage images/discretize/{image}")

helpers/colors.py

@@ -0,0 +1,69 @@
+from matplotlib import pyplot as plt
+import numpy as np
+from matplotlib.colors import LinearSegmentedColormap, ListedColormap, TwoSlopeNorm, BoundaryNorm, Normalize
+from matplotlib.cm import ScalarMappable
+
+
+class DEMScalarMappable(ScalarMappable):
+    colors_sea = plt.cm.terrain(np.linspace(0, 0.17, 256))
+    colors_land = plt.cm.terrain(np.linspace(0.25, 1, 256))
+    all_colors = np.vstack((colors_sea, colors_land))
+
+    def __init__(self, vmin, vmax):
+        if vmin > 0:
+            norm = Normalize(vmin, vmax, clip=True)
+            cmap = LinearSegmentedColormap.from_list("terrain_map", self.colors_land)
+        elif vmin == 0:
+            norm = Normalize(vmin, vmax, clip=True)
+            cmap = LinearSegmentedColormap.from_list(
+                "terrain_map", np.vstack((self.colors_sea[-1], self.colors_land[1:]))
+            )
+        elif vmax <= 0:
+            norm = Normalize(vmin, vmax, clip=True)
+            cmap = LinearSegmentedColormap.from_list("terrain_map", self.colors_sea)
+        else:
+            cmap = LinearSegmentedColormap.from_list("terrain_map", self.all_colors)
+            norm = TwoSlopeNorm(vmin=vmin, vcenter=0, vmax=vmax)
+        super().__init__(norm, cmap)
+
+    def segmented(self, bin_bounds, kind):
+        if kind == "bottom":
+            norm = BoundaryNorm(bin_bounds, bin_bounds.size - 1, clip=True)
+            cmap = LinearSegmentedColormap.from_list(
+                "terrain_map_segmented", self.to_rgba(bin_bounds[:-1]), N=bin_bounds.size - 1
+            )
+            return ScalarMappable(norm, cmap)
+        elif kind == "middle":
+            norm = BoundaryNorm(
+                [
+                    1.5 * bin_bounds[0] - 0.5 * bin_bounds[1],
+                    *(bin_bounds[0:-1] + bin_bounds[1:]) / 2,
+                    1.5 * bin_bounds[-1] - 0.5 * bin_bounds[-2],
+                ],
+                bin_bounds.size,
+            )
+            cmap = LinearSegmentedColormap.from_list(
+                "terrain_map_segmented", self.to_rgba(bin_bounds), N=bin_bounds.size
+            )
+            return ScalarMappable(norm, cmap)
+        elif kind == "top":
+            norm = BoundaryNorm(bin_bounds, bin_bounds.size - 1, clip=True)
+            cmap = LinearSegmentedColormap.from_list(
+                "terrain_map_segmented", self.to_rgba(bin_bounds[1:]), N=bin_bounds.size - 1
+            )
+            return ScalarMappable(norm, cmap)
+        else:
+            raise TypeError("Type can only be one of 'bottom', 'middle' or 'top'.")
+
+    @property
+    def lut(self):
+        return self.__get_lut(self)
+
+    def segmented_lut(self, bin_bounds, kind):
+        return self.__get_lut(self.segmented(bin_bounds, kind))
+
+    @staticmethod
+    def __get_lut(scalarmappable: ScalarMappable):
+        return (
+            scalarmappable.to_rgba(np.linspace(scalarmappable.get_clim()[0], scalarmappable.get_clim()[1], 256)) * 255
+        )