<s> """<STR_LIT>""" <EOL> from __future__ import division <EOL> import numpy as np <EOL> import matplotlib . pyplot as plt <EOL> import matplotlib . colors as colors <EOL> import os <EOL> import shutil <EOL> from . . exceptions import MissingOptionalDependency <EOL> try : <EOL> from osgeo import gdal <EOL> IMPORT_FLAG = True <EOL> except ImportError : <EOL> IMPORT_FLAG = False <EOL> def write_grid_geotiff ( grid , filename , field , rgb = False , level = None , <EOL> cmap = '<STR_LIT>' , vmin = <NUM_LIT:0> , vmax = <NUM_LIT> , color_levels = None , <EOL> warp = False , sld = False ) : <EOL> """<STR_LIT>""" <EOL> if not IMPORT_FLAG : <EOL> raise MissingOptionalDependency ( <EOL> '<STR_LIT>' ) <EOL> if field not in grid . fields . keys ( ) : <EOL> raise KeyError ( '<STR_LIT>' , field , '<STR_LIT>' ) <EOL> if '<STR_LIT:.>' not in filename : <EOL> name = filename <EOL> end = '<STR_LIT>' <EOL> ofile = name + "<STR_LIT:.>" + end <EOL> else : <EOL> ofile = filename <EOL> nz , ny , nx = grid . fields [ field ] [ '<STR_LIT:data>' ] . shape <EOL> dist = max ( grid . x [ '<STR_LIT:data>' ] ) <EOL> rangestep = grid . x [ '<STR_LIT:data>' ] [ <NUM_LIT:1> ] - grid . x [ '<STR_LIT:data>' ] [ <NUM_LIT:2> ] <EOL> lat = grid . origin_latitude [ '<STR_LIT:data>' ] [ <NUM_LIT:0> ] <EOL> lon = grid . origin_longitude [ '<STR_LIT:data>' ] [ <NUM_LIT:0> ] <EOL> filled = np . ma . filled ( grid . fields [ field ] [ '<STR_LIT:data>' ] , fill_value = - <NUM_LIT> ) <EOL> if level is None : <EOL> data = np . amax ( filled , <NUM_LIT:0> ) <EOL> else : <EOL> data = filled [ level ] <EOL> data = data . astype ( float ) <EOL> data [ data == - <NUM_LIT> ] = np . nan <EOL> iproj = '<STR_LIT>' + '<STR_LIT>' + '<STR_LIT>' + '<STR_LIT>' + '<STR_LIT>' + '<STR_LIT>' + str ( lat ) + '<STR_LIT>' + '<STR_LIT>' + str ( lon ) + '<STR_LIT>' + '<STR_LIT>' + '<STR_LIT>' + '<STR_LIT>' <EOL> out_driver = gdal . GetDriverByName ( "<STR_LIT>" ) <EOL> if not rgb : <EOL> dst_options = [ '<STR_LIT>' , '<STR_LIT>' ] <EOL> dst_ds = out_driver . Create ( <EOL> ofile , data . shape [ <NUM_LIT:1> ] , data . shape [ <NUM_LIT:0> ] , <NUM_LIT:1> , <EOL> gdal . GDT_Float32 , dst_options ) <EOL> else : <EOL> rarr , garr , barr = _get_rgb_values ( <EOL> data , vmin , vmax , color_levels , cmap ) <EOL> dst_ds = out_driver . Create ( ofile , data . shape [ <NUM_LIT:1> ] , <EOL> data . shape [ <NUM_LIT:0> ] , <NUM_LIT:3> , gdal . GDT_Byte ) <EOL> dst_ds . SetGeoTransform ( [ - dist , - rangestep , <NUM_LIT:0> , dist , <NUM_LIT:0> , rangestep ] ) <EOL> dst_ds . SetProjection ( iproj ) <EOL> if not rgb : <EOL> dst_ds . GetRasterBand ( <NUM_LIT:1> ) . WriteArray ( data [ : : - <NUM_LIT:1> , : ] ) <EOL> else : <EOL> dst_ds . GetRasterBand ( <NUM_LIT:1> ) . WriteArray ( rarr [ : : - <NUM_LIT:1> , : ] ) <EOL> dst_ds . GetRasterBand ( <NUM_LIT:2> ) . WriteArray ( garr [ : : - <NUM_LIT:1> , : ] ) <EOL> dst_ds . GetRasterBand ( <NUM_LIT:3> ) . WriteArray ( barr [ : : - <NUM_LIT:1> , : ] ) <EOL> dst_ds . FlushCache ( ) <EOL> dst_ds = None <EOL> if sld : <EOL> _create_sld ( cmap , vmin , vmax , ofile , color_levels ) <EOL> if warp : <EOL> os . system ( '<STR_LIT>' + <EOL> '<STR_LIT>' + ofile + '<STR_LIT:U+0020>' + <EOL> ofile + '<STR_LIT>' ) <EOL> shutil . move ( ofile + '<STR_LIT>' , ofile ) <EOL> def _get_rgb_values ( data , vmin , vmax , color_levels , cmap ) : <EOL> """<STR_LIT>""" <EOL> frac = ( data - vmin ) / np . float ( vmax - vmin ) <EOL> if color_levels is None : <EOL> color_levels = <NUM_LIT:255> <EOL> index = ( frac * color_levels ) . ravel ( ) <EOL> index [ index < <NUM_LIT:0> ] = <NUM_LIT:0> <EOL> index [ index > <NUM_LIT:255> ] = <NUM_LIT:255> <EOL> rarr = [ ] <EOL> garr = [ ] <EOL> barr = [ ] <EOL> cmap = plt . cm . get_cmap ( cmap ) <EOL> for val in index : <EOL> if not np . isnan ( val ) : <EOL> ind = np . int ( np . round ( val ) ) <EOL> r , g , b , t = cmap ( ind ) <EOL> rarr . append ( np . int ( np . round ( r * <NUM_LIT:255> ) ) ) <EOL> garr . append ( np . int ( np . round ( g * <NUM_LIT:255> ) ) ) <EOL> barr . append ( np . int ( np . round ( b * <NUM_LIT:255> ) ) ) <EOL> else : <EOL> rarr . append ( np . nan ) <EOL> garr . append ( np . nan ) <EOL> barr . append ( np . nan ) <EOL> rarr = np . reshape ( rarr , data . shape ) <EOL> garr = np . reshape ( garr , data . shape ) <EOL> barr = np . reshape ( barr , data . shape ) <EOL> return rarr , garr , barr <EOL> def _create_sld ( cmap , vmin , vmax , filename , color_levels = None ) : <EOL> """<STR_LIT>""" <EOL> cmap = plt . cm . get_cmap ( cmap ) <EOL> if color_levels is None : <EOL> color_levels = <NUM_LIT:255> <EOL> name , end = filename . split ( '<STR_LIT:.>' ) <EOL> ofile = name + '<STR_LIT>' <EOL> fileobj = open ( ofile , '<STR_LIT:w>' ) <EOL> header = """<STR_LIT>""" + str ( name ) + """<STR_LIT>""" <EOL> fileobj . write ( header ) <EOL> for i in np . arange ( color_levels + <NUM_LIT:1> ) : <EOL> val = i * ( vmax - vmin ) / ( color_levels ) + vmin <EOL> rgbt = cmap ( i ) <EOL> if i == <NUM_LIT:0> : <EOL> op = <NUM_LIT:0.0> <EOL> else : <EOL> op = <NUM_LIT:1.0> <EOL> hexval = colors . rgb2hex ( rgbt [ <NUM_LIT:0> : <NUM_LIT:3> ] ) <EOL> wstr = '<STR_LIT>' + str ( hexval ) . upper ( ) + '<STR_LIT>' + str ( op ) + '<STR_LIT>' + str ( val ) + '<STR_LIT>' <EOL> fileobj . write ( wstr ) <EOL> footer = """<STR_LIT>""" <EOL> fileobj . write ( footer ) <EOL> fileobj . close ( ) </s>
