Code snippet for colormaps by linear interpolation and 2D heatmaps in dex-lang

heatmap.dx

import png

' ## Colormaps
  Here we model colormaps as sequences of RGB colors which are then
  piecewise linearly interpolated in the interval [0,1] upon evaluation.

RGB = (Fin 3)=>Float

cmap_viridis = [[253.0, 231.0,  37.0]
               ,[ 94.0, 201.0,  98.0]
               ,[ 33.0, 145.0, 140.0] 
               ,[ 59.0,  82.0, 139.0]
               ,[ 68.0,   1.0,  84.0]] / 255.0

def piecewise_linear {n v} [VSpace v] (points: n=>v) (x: Float) : v =
  x = max x 0
  x = min x 1
  xn = x * (n_to_f (size n) - 1.0)
  lo = floor xn
  hi = ceil  xn
  p1 = points.(f_to_n lo @ _)
  p2 = points.(f_to_n hi @ _)
  (1 - (xn - lo)) .* p1 + (xn - lo) .* p2

-- util to inspect a cmap visually
def plot_cmap {n} (cmap: n=>RGB) : (Fin 10)=>(Fin 100)=>RGB =
  for i:(Fin 10). for j:(Fin 100).
    piecewise_linear cmap (n_to_f (ordinal j) / 100.0)

:html imshow $ plot_cmap cmap_viridis

cmap_jax = [[157.0,  34.0, 178.0]
           ,[106.0,  20.0, 155.0]
           ,[ 37.0,  85.0, 199.0]
           ,[ 94.0, 152.0, 246.0]
           ,[ 33.0, 167.0, 155.0]] / 255.0

:html imshow $ plot_cmap cmap_jax

-- to reverse a cmap we just reverse the array of control points
:html imshow $ plot_cmap (reverse cmap_jax)

' ## 2D heatmap

def heatmap2d {n m c} 
    (ymin: Float)
    (ymax: Float)
    (ygrid: n=>m=>Float)
    (cmap: c=>RGB)
    : (n=>m=>RGB) =
  img = for i j.
    y = (ygrid.i.j - ymin) / (ymax - ymin)
    piecewise_linear cmap y
  img

' ## 2D discretization of a function on a grid

def discretize2d {n m} 
    (f: Float -> Float -> Float)
    (xs1: n=>Float)
    (xs2: m=>Float)
    : n=>m=>Float =
  fgrid = for i j. f xs1.i xs2.j
  fgrid

' ## Example: Himmelblau's function

def himmelblau (x1: Float) (x2: Float) : Float =
  sq (sq x1 + x2 - 11.0) + sq (x1 + sq x2 - 7.0)

example_fn = himmelblau
ymin = -30.0
ymax = 200.0
xs1 = for i:(Fin 100). n_to_f (ordinal i) / 10.0 - 5.0
xs2 = for i:(Fin 100). n_to_f (ordinal i) / 10.0 - 5.0

ygrid = discretize2d example_fn xs1 xs2
:html imshow $ heatmap2d ymin ymax ygrid cmap_viridis

def hcat {l n m a} (mat1: l=>n=>a) (mat2: l=>m=>a) : l=>(n|m)=>a =
  for i j. case j of
    Left  j1 -> mat1.i.j1
    Right j2 -> mat2.i.j2

img_viridis  = heatmap2d ymin ymax ygrid cmap_viridis          
img_viridisr = heatmap2d ymin ymax ygrid (reverse cmap_viridis)
img_jax      = heatmap2d ymin ymax ygrid cmap_jax              
img_jaxr     = heatmap2d ymin ymax ygrid (reverse cmap_jax)    

:html imshow $ hcat img_viridis $ hcat img_viridisr $ hcat img_jax img_jaxr