<s> import calendar <EOL> from datetime import timedelta , datetime <EOL> from time import strptime <EOL> from glob import glob <EOL> import re <EOL> import string <EOL> import os <EOL> from ocw . dataset import Dataset <EOL> import ocw . utils as utils <EOL> import netCDF4 <EOL> import h5py <EOL> import numpy <EOL> import numpy . ma as ma <EOL> LAT_NAMES = [ '<STR_LIT:y>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> LON_NAMES = [ '<STR_LIT:x>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> TIME_NAMES = [ '<STR_LIT:time>' , '<STR_LIT>' , '<STR_LIT:date>' , '<STR_LIT>' , '<STR_LIT>' ] <EOL> def _get_netcdf_variable_name ( valid_var_names , netcdf , netcdf_var ) : <EOL> '''<STR_LIT>''' <EOL> dimensions = netcdf . variables [ netcdf_var ] . dimensions <EOL> dims_lower = [ dim . encode ( ) . lower ( ) for dim in dimensions ] <EOL> intersect = set ( valid_var_names ) . intersection ( dims_lower ) <EOL> if len ( intersect ) == <NUM_LIT:1> : <EOL> index = dims_lower . index ( intersect . pop ( ) ) <EOL> dimension_name = dimensions [ index ] . encode ( ) <EOL> possible_vars = [ ] <EOL> for var in netcdf . variables . keys ( ) : <EOL> var_dimensions = netcdf . variables [ var ] . dimensions <EOL> if len ( var_dimensions ) != <NUM_LIT:1> : <EOL> continue <EOL> if var_dimensions [ <NUM_LIT:0> ] . encode ( ) == dimension_name : <EOL> possible_vars . append ( var ) <EOL> if len ( possible_vars ) == <NUM_LIT:1> : <EOL> return possible_vars [ <NUM_LIT:0> ] <EOL> variables = netcdf . variables . keys ( ) <EOL> vars_lower = [ var . encode ( ) . lower ( ) for var in variables ] <EOL> intersect = set ( valid_var_names ) . intersection ( vars_lower ) <EOL> if len ( intersect ) == <NUM_LIT:1> : <EOL> index = vars_lower . index ( intersect . pop ( ) ) <EOL> return variables [ index ] <EOL> error = ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> ) <EOL> raise ValueError ( error ) <EOL> def load_WRF_2d_files ( file_path = None , <EOL> filename_pattern = None , <EOL> filelist = None , <EOL> variable_name = '<STR_LIT>' , <EOL> name = '<STR_LIT>' ) : <EOL> '''<STR_LIT>''' <EOL> if not filelist : <EOL> WRF_files = [ ] <EOL> for pattern in filename_pattern : <EOL> WRF_files . extend ( glob ( file_path + pattern ) ) <EOL> else : <EOL> WRF_files = [ line . rstrip ( '<STR_LIT:\n>' ) for line in open ( filelist ) ] <EOL> WRF_files . sort ( ) <EOL> file_object_first = netCDF4 . Dataset ( WRF_files [ <NUM_LIT:0> ] ) <EOL> lats = file_object_first . variables [ '<STR_LIT>' ] [ <NUM_LIT:0> , : ] <EOL> lons = file_object_first . variables [ '<STR_LIT>' ] [ <NUM_LIT:0> , : ] <EOL> times = [ ] <EOL> nfile = len ( WRF_files ) <EOL> for ifile , file in enumerate ( WRF_files ) : <EOL> print '<STR_LIT>' + str ( ifile + <NUM_LIT:1> ) + '<STR_LIT:/>' + str ( nfile ) , file <EOL> file_object = netCDF4 . Dataset ( file ) <EOL> time_struct_parsed = strptime ( file [ - <NUM_LIT> : ] , "<STR_LIT>" ) <EOL> for ihour in numpy . arange ( <NUM_LIT> ) : <EOL> times . append ( datetime ( * time_struct_parsed [ : <NUM_LIT:6> ] ) + timedelta ( hours = ihour ) ) <EOL> values0 = file_object . variables [ variable_name ] [ : ] <EOL> if ifile == <NUM_LIT:0> : <EOL> values = values0 <EOL> variable_unit = file_object . variables [ variable_name ] . units <EOL> else : <EOL> values = numpy . concatenate ( ( values , values0 ) ) <EOL> file_object . close ( ) <EOL> times = numpy . array ( times ) <EOL> return Dataset ( lats , lons , times , values , variable_name , units = variable_unit , name = name ) <EOL> def load_file ( file_path , <EOL> variable_name , <EOL> variable_unit = None , <EOL> elevation_index = <NUM_LIT:0> , <EOL> name = '<STR_LIT>' , <EOL> lat_name = None , <EOL> lon_name = None , <EOL> time_name = None ) : <EOL> '''<STR_LIT>''' <EOL> try : <EOL> netcdf = netCDF4 . Dataset ( file_path , mode = '<STR_LIT:r>' ) <EOL> except RuntimeError : <EOL> err = "<STR_LIT>" <EOL> raise ValueError ( err ) <EOL> except : <EOL> err = ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> ) <EOL> raise ValueError ( err ) <EOL> if not lat_name : <EOL> lat_name = _get_netcdf_variable_name ( LAT_NAMES , netcdf , variable_name ) <EOL> if not lon_name : <EOL> lon_name = _get_netcdf_variable_name ( LON_NAMES , netcdf , variable_name ) <EOL> if not time_name : <EOL> time_name = _get_netcdf_variable_name ( TIME_NAMES , netcdf , variable_name ) <EOL> lats = netcdf . variables [ lat_name ] [ : ] <EOL> lons = netcdf . variables [ lon_name ] [ : ] <EOL> time_raw_values = netcdf . variables [ time_name ] [ : ] <EOL> times = utils . decode_time_values ( netcdf , time_name ) <EOL> times = numpy . array ( times ) <EOL> values = ma . array ( netcdf . variables [ variable_name ] [ : ] ) <EOL> variable_unit = netcdf . variables [ variable_name ] . units <EOL> if len ( values . shape ) == <NUM_LIT:4> : <EOL> dims = netcdf . variables [ variable_name ] . dimensions <EOL> dimension_names = [ dim_name . encode ( ) for dim_name in dims ] <EOL> lat_lon_time_var_names = [ lat_name , lon_name , time_name ] <EOL> elev_names = set ( dimension_names ) - set ( lat_lon_time_var_names ) <EOL> level_index = <NUM_LIT:1> <EOL> if level_index == <NUM_LIT:0> : <EOL> values = values [ elevation_index , : , : , : ] <EOL> elif level_index == <NUM_LIT:1> : <EOL> values = values [ : , elevation_index , : , : ] <EOL> elif level_index == <NUM_LIT:2> : <EOL> values = values [ : , : , elevation_index , : ] <EOL> else : <EOL> values = values [ : , : , : , elevation_index ] <EOL> origin = { <EOL> '<STR_LIT:source>' : '<STR_LIT>' , <EOL> '<STR_LIT:path>' : file_path , <EOL> '<STR_LIT>' : lat_name , <EOL> '<STR_LIT>' : lon_name , <EOL> '<STR_LIT>' : time_name <EOL> } <EOL> if elevation_index != <NUM_LIT:0> : origin [ '<STR_LIT>' ] = elevation_index <EOL> return Dataset ( lats , lons , times , values , variable = variable_name , <EOL> units = variable_unit , name = name , origin = origin ) <EOL> def load_multiple_files ( file_path , <EOL> variable_name , <EOL> dataset_name = '<STR_LIT>' , <EOL> variable_unit = None , <EOL> lat_name = None , <EOL> lon_name = None , <EOL> time_name = None ) : <EOL> '''<STR_LIT>''' <EOL> data_filenames = [ ] <EOL> data_filenames . extend ( glob ( file_path ) ) <EOL> data_filenames . sort ( ) <EOL> ndata = len ( data_filenames ) <EOL> if ndata == <NUM_LIT:1> : <EOL> data_name = [ dataset_name ] <EOL> else : <EOL> data_name = [ ] <EOL> data_filenames_reversed = [ ] <EOL> for element in data_filenames : <EOL> data_filenames_reversed . append ( element [ : : - <NUM_LIT:1> ] ) <EOL> prefix = os . path . commonprefix ( data_filenames ) <EOL> postfix = os . path . commonprefix ( data_filenames_reversed ) [ : : - <NUM_LIT:1> ] <EOL> for element in data_filenames : <EOL> data_name . append ( element . replace ( prefix , '<STR_LIT>' ) . replace ( postfix , '<STR_LIT>' ) ) <EOL> datasets = [ ] <EOL> for ifile , filename in enumerate ( data_filenames ) : <EOL> datasets . append ( load_file ( filename , variable_name , variable_unit , name = data_name [ ifile ] , <EOL> lat_name = lat_name , lon_name = lon_name , time_name = time_name ) ) <EOL> return datasets , data_name <EOL> def load_WRF_2d_files_RAIN ( file_path = None , <EOL> filename_pattern = None , <EOL> filelist = None , <EOL> name = '<STR_LIT>' ) : <EOL> '''<STR_LIT>''' <EOL> if not filelist : <EOL> WRF_files = [ ] <EOL> for pattern in filename_pattern : <EOL> WRF_files . extend ( glob ( file_path + pattern ) ) <EOL> WRF_files . sort ( ) <EOL> else : <EOL> WRF_files = [ line . rstrip ( '<STR_LIT:\n>' ) for line in open ( filelist ) ] <EOL> file_object_first = netCDF4 . Dataset ( WRF_files [ <NUM_LIT:0> ] ) <EOL> lats = file_object_first . variables [ '<STR_LIT>' ] [ <NUM_LIT:0> , : ] <EOL> lons = file_object_first . variables [ '<STR_LIT>' ] [ <NUM_LIT:0> , : ] <EOL> times = [ ] <EOL> nfile = len ( WRF_files ) <EOL> for ifile , file in enumerate ( WRF_files ) : <EOL> print '<STR_LIT>' + str ( ifile + <NUM_LIT:1> ) + '<STR_LIT:/>' + str ( nfile ) , file <EOL> file_object = netCDF4 . Dataset ( file ) <EOL> time_struct_parsed = strptime ( file [ - <NUM_LIT> : ] , "<STR_LIT>" ) <EOL> for ihour in range ( <NUM_LIT> ) : <EOL> times . append ( datetime ( * time_struct_parsed [ : <NUM_LIT:6> ] ) + timedelta ( hours = ihour ) ) <EOL> if ifile == <NUM_LIT:0> : <EOL> values0 = file_object . variables [ '<STR_LIT>' ] [ : ] + file_object . variables [ '<STR_LIT>' ] [ : ] <EOL> else : <EOL> values0 = numpy . concatenate ( ( values0 , file_object . variables [ '<STR_LIT>' ] [ : ] + file_object . variables [ '<STR_LIT>' ] [ : ] ) ) <EOL> file_object . close ( ) <EOL> times = numpy . array ( times ) <EOL> years = numpy . array ( [ d . year for d in times ] ) <EOL> ncycle = numpy . unique ( years ) . size <EOL> print '<STR_LIT>' , ncycle <EOL> nt , ny , nx = values0 . shape <EOL> values = numpy . zeros ( [ nt - ncycle * <NUM_LIT> , ny , nx ] ) <EOL> times2 = [ ] <EOL> nt2 = nt / ncycle <EOL> nt3 = nt2 - <NUM_LIT> <EOL> t_index = <NUM_LIT:0> <EOL> for icycle in numpy . arange ( ncycle ) : <EOL> for it in numpy . arange ( nt3 ) + <NUM_LIT> : <EOL> values [ t_index , : ] = values0 [ icycle * nt2 + it , : ] - values0 [ icycle * nt2 + it - <NUM_LIT:1> , : ] <EOL> times2 . append ( times [ icycle * nt2 + it ] ) <EOL> t_index = t_index + <NUM_LIT:1> <EOL> variable_name = '<STR_LIT>' <EOL> variable_unit = '<STR_LIT>' <EOL> times2 = numpy . array ( times2 ) <EOL> return Dataset ( lats , lons , times2 , values , variable_name , units = variable_unit , name = name ) <EOL> def load_dataset_from_multiple_netcdf_files ( variable_name , <EOL> lat_name = None , lon_name = None , time_name = None , <EOL> name = '<STR_LIT>' , file_list = None , file_path = None , filename_pattern = None , <EOL> mask_file = None , mask_variable = None , mask_value = <NUM_LIT:0> ) : <EOL> '''<STR_LIT>''' <EOL> nc_files = [ ] <EOL> if not file_list : <EOL> for pattern in filename_pattern : <EOL> nc_files . extend ( glob ( file_path + pattern ) ) <EOL> else : <EOL> nc_files = [ line . rstrip ( '<STR_LIT:\n>' ) for line in open ( file_list ) ] <EOL> nc_files . sort ( ) <EOL> dataset0 = load_file ( nc_files [ <NUM_LIT:0> ] , variable_name = variable_name , lat_name = lat_name , lon_name = lon_name , time_name = time_name ) <EOL> if dataset0 . lons . ndim == <NUM_LIT:1> and dataset0 . lats . ndim == <NUM_LIT:1> : <EOL> lons , lats = numpy . meshgrid ( dataset0 . lons , dataset0 . lats ) <EOL> elif dataset0 . lons . ndim == <NUM_LIT:2> and dataset0 . lats . ndim == <NUM_LIT:2> : <EOL> lons = dataset0 . lons <EOL> lats = dataset0 . lats <EOL> if mask_file : <EOL> mask_dataset = load_file ( mask_file , mask_variable ) <EOL> y_index , x_index = numpy . where ( mask_dataset . values == mask_value ) <EOL> times = [ ] <EOL> nfile = len ( nc_files ) <EOL> for ifile , file in enumerate ( nc_files ) : <EOL> print '<STR_LIT>' + str ( ifile + <NUM_LIT:1> ) + '<STR_LIT:/>' + str ( nfile ) , file <EOL> file_object0 = load_file ( file , variable_name ) <EOL> values0 = file_object0 . values <EOL> times . extend ( file_object0 . times ) <EOL> if mask_file : <EOL> values0 = values0 [ : , y_index , x_index ] <EOL> if ifile == <NUM_LIT:0> : <EOL> data_values = values0 <EOL> else : <EOL> data_values = numpy . concatenate ( ( data_values , values0 ) ) <EOL> times = numpy . array ( times ) <EOL> return Dataset ( lats , lons , times , data_values , variable_name , name = name ) <EOL> def load_NLDAS_forcingA_files ( file_path = None , <EOL> filename_pattern = None , <EOL> filelist = None , <EOL> variable_name = '<STR_LIT>' , <EOL> name = '<STR_LIT>' ) : <EOL> '''<STR_LIT>''' <EOL> if not filelist : <EOL> NLDAS_files = [ ] <EOL> for pattern in filename_pattern : <EOL> NLDAS_files . extend ( glob ( file_path + pattern ) ) <EOL> else : <EOL> NLDAS_files = [ line . rstrip ( '<STR_LIT:\n>' ) for line in open ( filelist ) ] <EOL> NLDAS_files . sort ( ) <EOL> file_object_first = netCDF4 . Dataset ( NLDAS_files [ <NUM_LIT:0> ] ) <EOL> lats = file_object_first . variables [ '<STR_LIT>' ] [ : ] <EOL> lons = file_object_first . variables [ '<STR_LIT>' ] [ : ] <EOL> lons , lats = numpy . meshgrid ( lons , lats ) <EOL> times = [ ] <EOL> nfile = len ( NLDAS_files ) <EOL> for ifile , file in enumerate ( NLDAS_files ) : <EOL> print '<STR_LIT>' + str ( ifile + <NUM_LIT:1> ) + '<STR_LIT:/>' + str ( nfile ) , file <EOL> file_object = netCDF4 . Dataset ( file ) <EOL> time_struct_parsed = strptime ( file [ - <NUM_LIT:20> : - <NUM_LIT:7> ] , "<STR_LIT>" ) <EOL> times . append ( datetime ( * time_struct_parsed [ : <NUM_LIT:6> ] ) ) <EOL> values0 = file_object . variables [ variable_name ] [ : ] <EOL> values0 = numpy . expand_dims ( values0 , axis = <NUM_LIT:0> ) <EOL> if ifile == <NUM_LIT:0> : <EOL> values = values0 <EOL> variable_unit = file_object . variables [ variable_name ] . units <EOL> else : <EOL> values = numpy . concatenate ( ( values , values0 ) ) <EOL> file_object . close ( ) <EOL> times = numpy . array ( times ) <EOL> return Dataset ( lats , lons , times , values , variable_name , units = variable_unit , name = name ) <EOL> def load_GPM_IMERG_files ( file_path = None , <EOL> filename_pattern = None , <EOL> filelist = None , <EOL> variable_name = '<STR_LIT>' , <EOL> name = '<STR_LIT>' ) : <EOL> '''<STR_LIT>''' <EOL> if not filelist : <EOL> GPM_files = [ ] <EOL> for pattern in filename_pattern : <EOL> GPM_files . extend ( glob ( file_path + pattern ) ) <EOL> else : <EOL> GPM_files = [ line . rstrip ( '<STR_LIT:\n>' ) for line in open ( filelist ) ] <EOL> GPM_files . sort ( ) <EOL> file_object_first = h5py . File ( GPM_files [ <NUM_LIT:0> ] ) <EOL> lats = file_object_first [ '<STR_LIT>' ] [ '<STR_LIT>' ] [ : ] <EOL> lons = file_object_first [ '<STR_LIT>' ] [ '<STR_LIT>' ] [ : ] <EOL> lons , lats = numpy . meshgrid ( lons , lats ) <EOL> variable_unit = "<STR_LIT>" <EOL> times = [ ] <EOL> nfile = len ( GPM_files ) <EOL> for ifile , file in enumerate ( GPM_files ) : <EOL> print '<STR_LIT>' + str ( ifile + <NUM_LIT:1> ) + '<STR_LIT:/>' + str ( nfile ) , file <EOL> file_object = h5py . File ( file ) <EOL> time_struct_parsed = strptime ( file [ - <NUM_LIT> : - <NUM_LIT> ] , "<STR_LIT>" ) <EOL> times . append ( datetime ( * time_struct_parsed [ : <NUM_LIT:6> ] ) ) <EOL> values0 = numpy . transpose ( ma . masked_less ( file_object [ '<STR_LIT>' ] [ variable_name ] [ : ] , <NUM_LIT:0.> ) ) <EOL> values0 = numpy . expand_dims ( values0 , axis = <NUM_LIT:0> ) <EOL> if ifile == <NUM_LIT:0> : <EOL> values = values0 <EOL> else : <EOL> values = numpy . concatenate ( ( values , values0 ) ) <EOL> file_object . close ( ) <EOL> times = numpy . array ( times ) <EOL> return Dataset ( lats , lons , times , values , variable_name , units = variable_unit , name = name ) </s>
