diff --git a/plot2d.py b/plot2d.py
index f5abfe4..5a65600 100644
--- a/plot2d.py
+++ b/plot2d.py
@@ -1,7 +1,7 @@
 import numpy as np
 
 from bokeh.plotting import figure
-from bokeh.models import ColumnDataSource
+from bokeh.models import ColumnDataSource, ColorBar, LinearColorMapper
 
 from buki import Object
 
@@ -15,8 +15,16 @@ class Plot2D(Object):
 
         self.source = source = ColumnDataSource(data=data)
 
-        fig = figure(name=name)
-        fig.image(source=source, x=0, y=0, dw=1, dh=1, palette="Spectral11")
+        cmap = LinearColorMapper(palette="Viridis256")
+        cbar = ColorBar(color_mapper=cmap)
+
+        fig = figure(name=name, match_aspect=True)
+        fig.image(source=source, x=0, y=0, dw=1, dh=1, color_mapper=cmap)
+
+        fig.x_range.range_padding = 0
+        fig.y_range.range_padding = 0
+
+        fig.add_layout(cbar, "right")
         super().__init__(fig)