Drawing wind barbs: GDAL python

Wind on a map is usually represented using wind barbs, a graphical way to show both speed and direction. The problem is that most GIS software is not prepared to draw them directly.
In this entry, I'll show how I use to generate PNG files from the meteorological model GFS, the same from the entry showing how to do some raster calculations.

As usual, the code is available at GitHub.

First the file that generates the png:

'''
Draws wind barbs on a png from a GRIB file using PIL
'''

import Image
import ImageDraw
from osgeo import gdal
from osgeo import osr
from osgeo.gdalconst import GA_ReadOnly
from math import floor
from math import ceil
from math import sqrt
from math import atan2
from math import degrees


def find_band_number(dataset, variable, level):
    '''
    Finds the band number inside the GRIB file, given the variable and the level names
    '''
    for i in range(1,dataset.RasterCount + 1):
        band = dataset.GetRasterBand(i)
        metadata = band.GetMetadata()
        band_level = metadata['GRIB_SHORT_NAME']
        band_variable = metadata['GRIB_ELEMENT']
        if (variable == band_variable) and (level == band_level):
            return i
    return None

def draw_barb(size, module, angle, color):
    extra_pixels = 10 #Avoid cutting the barb when rotating
    barb = Image.new('RGBA', (size + extra_pixels, size + extra_pixels))
    barb_draw = ImageDraw.Draw(barb)
    size = float(size) #Force it to float  
    separation = size/6
    module5 = int(round(module/5))

    x,v = divmod(module5,2)
    l,x = divmod(x,5)

    barb_draw.line([(extra_pixels,extra_pixels+size/2),(size+extra_pixels,extra_pixels+size/2)],color)
    pos = 0
    for nl in range(0,l,1): #50 kt triangles
        barb_draw.polygon([(extra_pixels+pos,extra_pixels+size/2),(extra_pixels+pos + extra_pixels+size/8,0),(extra_pixels+pos+size/4,extra_pixels+size/2),(extra_pixels+pos,extra_pixels+size/2)],color)

        pos = pos + size/4 + separation
    for nx in range(0,x,1):
        barb_draw.line([(extra_pixels+pos,extra_pixels+size/2),(extra_pixels+pos,extra_pixels+0)],color)

        pos = pos + separation
    if pos == 0: #advance a little for 5kt barbs
        pos = pos + separation
    if v == 1: # Only 0 or 1 are possible
        barb_draw.line([(extra_pixels+pos,extra_pixels+size/2),(extra_pixels+pos,extra_pixels+size/4)],color)

    barb = barb.rotate(angle, Image.BICUBIC)

    return barb

def draw_wind_barbs(xsize, ysize, out_file, data_file, epsg, geotransform, u_var, v_var, level, separation, barb_size, barb_color):
    out_img = Image.new('RGBA', (xsize, ysize) )

    dataset = gdal.Open(data_file, GA_ReadOnly )
    u_band_id = find_band_number(dataset, u_var, level)
    v_band_id = find_band_number(dataset, v_var, level)

    band_u = dataset.GetRasterBand(u_band_id)
    band_v = dataset.GetRasterBand(v_band_id)

    geotransform_in = dataset.GetGeoTransform()   

    xsize_in = band_u.XSize
    ysize_in = band_u.YSize

    values_u = band_u.ReadAsArray(0, 0, xsize_in, ysize_in)
    values_v = band_v.ReadAsArray(0, 0, xsize_in, ysize_in)
    

    proj_in = osr.SpatialReference()
    proj_in.SetFromUserInput(dataset.GetProjection())
    proj_out = osr.SpatialReference()
    proj_out.ImportFromEPSG(epsg)

    transf = osr.CoordinateTransformation(proj_out,proj_in)

    for px in range(0,xsize + separation,separation):
        x_out = geotransform[0] + px * geotransform[1]
        for py in range(0,ysize + separation,separation):
            y_out = geotransform[3] + py * geotransform[5]
            point = transf.TransformPoint(x_out, y_out,0)
            px_in = (float(point[0]) - geotransform_in[0]) / geotransform_in[1]
            py_in = (float(point[1]) - geotransform_in[3]) / geotransform_in[5]
           
            d1 = 1/sqrt( pow((px_in - floor(px_in)), 2) + pow((py_in - floor(py_in)), 2) )
            d2 = 1/sqrt( pow((ceil(px_in) - px_in), 2) + pow((py_in - floor(py_in)), 2) )
            d3 = 1/sqrt( pow((ceil(px_in) - px_in), 2) + pow((ceil(py_in) - py_in), 2) )
            d4 = 1/sqrt( pow((px_in - floor(px_in)), 2) + pow((ceil(py_in) - py_in), 2) ) 
            d_sum = d1 + d2 + d3 + d4   

            u = (d1*values_u[floor(py_in)][floor(px_in)] + d2*values_u[floor(py_in)][ceil(px_in)] + d3*values_u[ceil(py_in)][ceil(px_in)] + d4*values_u[ceil(py_in)][floor(px_in)]) / d_sum
            v = (d1*values_v[floor(py_in)][floor(px_in)] + d2*values_v[floor(py_in)][ceil(px_in)] + d3*values_v[ceil(py_in)][ceil(px_in)] + d4*values_v[ceil(py_in)][floor(px_in)]) / d_sum

            module = 1.943844494 * sqrt((u*u)+(v*v)) #in knots
           
            angle = degrees(atan2(v,u)) #Angle like in trigonometry, not wind direction where 0 is northern wind
            try:
                barb = draw_barb(barb_size, module, angle, barb_color)
                out_img.paste(barb, (px-(barb_size/2),py-(barb_size/2)), barb)
            except Exception, ex:
                raise Exception("No es pot dibuixar la barba: " + str(ex))
   
    out_img.save( out_file )

    #Creating the pngw worldfile
    fp = open(out_file.replace("png","pngw"),"w")
    print str(geotransform[1])+"\n0\n0\n"+str(geotransform[5])+"\n"+str(geotransform[0])+"\n"+str(geotransform[3])
    fp.write(str(geotransform[1])+"\n0\n0\n"+str(geotransform[5])+"\n"+str(geotransform[0])+"\n"+str(geotransform[3]))
    fp.close()

if __name__ == '__main__':
   
    draw_wind_barbs(600, 400, 'wind.png', 'gfs_3_20130524_0000_000.grb', 4326, [0, 0.1, 0, 70, 0, -0.1], 'UGRD', 'VGRD', '850-ISBL', 20, 20, '#000000')

• The last part of the file (line 122) shows how to call tha main function draw_wind_barbs
• draw_wind_barbs takes the following arguments:
• xsize, ysize: The output image size
• out_file: the output png file. An other file will be created with the extension pngw, as the worldfile.
• data_file: The input file from the model (actually, any GRIB file with the wind in U and V components)
• epsg, geotransform: The output projection parameters.
• u_var, v_var, level: The wind component variable names (south-north is v, east-west is u) as they appear in the GRIB_ELEMENT metadata in the GRIB file. Level as it appears in the GRIB_SHORT_NAME metadata.
• separation, barb_size, color: The separation and size in pixels at the output file (it's possible change the barb density)
• From line 61 to 84, the function creates the output image according to the input file features and the output file as asked in the parameters
• The number of the band corresponding to the u and v components is calculated with the function find_band_number, that looks into the file metadata to match the asked values with some of the existing bands. Lines 17-28
• From 86 to 110, the barbs are drawn. For every n pixels as asked in the separation parameter, a barb is calculated. To do it, is necessary to know the position in the input file (lines 91 to 98), with the projection transformation,
• Then the wind value is calculated from the closest four pixels in the original file, using the inverse of the distance method . lines 94-102
• After the u and v components are known, the module and angle are calculated. Note that the original speeds are in m/s, but the barbs use knots, so the transformation is done here.The angle is not S-N, but W-E, which is fine for drawing, but a bit confusing. The coordinates origins are S and W.
• With the angle and module, the barb can be drawn, using the function draw_barb, and the result pasted into the output file.
• The draw_barb function divides the module by 5, 10 and 50 to know how many ticks the barb has to have, and draws the barb with the PIL polygon and line methods.
• The file is written and the worldfile is created at lines 114-118

To download the data files and draw directly the png, I wrote the following script:

from urllib2 import urlopen
from urllib2 import HTTPError
from sys import exit
from datetime import datetime
from barbs import draw_wind_barbs

def download_file(year,month,day,hour,valid_time):
    filename = "gfs_3_%d%02d%02d_%02d00_%03d.grb"%(year, month, day, hour, valid_time)
    url = "http://nomads.ncdc.noaa.gov/data/gfs-avn-hi/%d%02d/%d%02d%02d/gfs_3_%d%02d%02d_%02d00_%03d.grb"%(year, month, year, month, day, year, month, day, hour, valid_time)
    try: 
        req = urlopen(url)
    except HTTPError:
        print "File not found: May be the date is incorrect?"
        exit()

    f = open(filename, 'w')
    f.write(req.read())
    f.close()
    return filename

if __name__ == "__main__":
    now = datetime.now()
    print now.year
    year = raw_input('Type year ('+str(now.year)+'): ')
    month = raw_input('Type month ('+str(now.month)+'): ')
    day = raw_input('Type day ('+str(now.day)+'): ')
    hour = raw_input('Type hour (0): ')
    valid_time = raw_input('Type valid time (0): ')
    if (year == ''):
        year = now.year
    if (month == ''):
        month = now.month
    if (day == ''):
        day = now.day  
    if (hour == ''):
        hour = 0
    if (valid_time == ''):
        valid_time = 0     
    filename = download_file(int(year),int(month),int(day),int(hour),int(valid_time))
    out_filename = filename.replace("grb","png")

    draw_wind_barbs(600, 400, out_filename, filename, 4326, [0, 0.1, 0, 70, 0, -0.1], 'UGRD', 'VGRD', '850-ISBL', 20, 20, '#000000')

• The main part takes the date wanted by the user and sets it at the current day at 00h as default. 
• After that, calls the download_file function to download the GRIB file from the server.
• Then call the draw_wind_barbs function to draw the map. Try changing the geotransform and image size to see how the output changes.

Finally, to draw the file with the map at the background, I have  opened the png with QGIS (since it has the worldfile, is possible to do it), and added the background from naturalearth before exporting the image as a png.