Python-gdal create geotiff from array with colormapping

Question

I am creating grayscale, single channel geotiff files from numpy arrays with the below module, but I would like to apply a fairly standard colormap to the data. I've tried three channels of R,B and G and have had no success if producing anything that looks correct. I'm thinking my approach, in general is flawed.

How can I go from grayscale to color geotiffs?

def geotiff_output(outfile, mag_grid, lons, lats):
    #*********************Test***************************
    from osgeo import gdal, osr, ogr
    xres = lons[1] - lons[0]
    yres = lats[1] - lats[0]
    ysize = len(lats)
    xsize = len(lons)
    ulx = lons[0] - (xres / 2.)
    uly = lats[-1] - (yres / 2.)
    driver = gdal.GetDriverByName('GTiff')
    ds = driver.Create(sc_settings.dynamic_folder_script + outfile[:-4] + ".tif",
                       xsize, ysize, 1, gdal.GDT_Byte, )

    # this assumes the projection is Geographic lat/lon WGS 84
    srs = osr.SpatialReference()
    srs.ImportFromEPSG(4326)
    ds.SetProjection(srs.ExportToWkt())
    gt = [ulx, xres, 0, uly, 0, yres ]
    ds.SetGeoTransform(gt)
    outband=ds.GetRasterBand(1)
    #mask no data points...
    mag_grid1 = np.ma.masked_where(mag_grid==-99.,mag_grid)
    #assign actual data values for inclusion in metadata
    outband.SetStatistics(np.round(np.min(mag_grid1),decimals=4),
                            np.round(np.max(mag_grid1),decimals=4),
                            np.round(np.mean(mag_grid1),decimals=4),
                            np.round(np.std(mag_grid1),decimals=4))

    #normalize the array to a 0-255 scale
    if np.min(mag_grid1) < 0:
        mag_grid1 = mag_grid1 + np.abs(np.min(mag_grid1))
    #Normalize array to a 0-255 scale for raster channel
    mag_grid2 = ((mag_grid1 - np.min(mag_grid1))/(np.max(mag_grid1) - np.min(mag_grid1)))*256
    outband.WriteArray(mag_grid2)
    ds = None

Answer 1

Following the link in the comment from @kersten I found this https://svn.osgeo.org/gdal/trunk/autotest/gcore/tiff_write.py with lots of interesting snippets. I think what you want to do is this:

...
ct = gdal.ColorTable()
# Some examples
ct.SetColorEntry( 0, (0, 0, 0, 255) )
ct.SetColorEntry( 1, (0, 255, 0, 255) )
ct.SetColorEntry( 2, (255, 0, 0, 255) )
ct.SetColorEntry( 3, (255, 0, 255, 255) )
# Set the color table for your band
outband.SetColorTable(color_table)
...

Also, you may want to use CreateColorRamp