''' PRMS-Python: Powerful, sane tools for manipulating PRMS input data to create new scenarios or parameterizations for sensitivity analysis, scenario modeling, or whatever other uses this might have. The fundamental process in scenario development is to modify some "base" starting data to create some "scenario" data. No matter what data we're using, once it's ready, we run a PRMS "simulation" on that data. This module presents a Simulation and Scenario class, where each tracks relevant provenance information and input files to facilitate better data management techniques to streamline later analyses. ''' import datetime import io import itertools import numpy as np from collections import OrderedDict def modify_params(params_in, params_out, param_mods=None): ''' Given a parameter file in and a dictionary of param_mods, write modified parameters to params_out. Example: Below we modify the monthly jh_coef by increasing it 10% for every month. >>> params_in = 'models/lbdc/params' >>> params_out = 'scenarios/jh_coef_1.1/params' >>> scale_10pct = lambda x: x * 1.1 >>> modify_params(params_in, params_out, {'jh_coef': scale_10pct}) So param_mods is a dictionary of with keys being parameter names and values a function that operates on a single value. Currently we only accept functions that operate on single values without reference to any other parameters. The function will be applied to every cell, month, or cascade routing rule for which the parameter is defined. Arguments: params_in (str): location on disk of the base parameter file params_out (str): location on disk where the modified parameters will be written param_mods (dict): param name-keyed, param modification function-valued Returns: None ''' p_in = Parameters(params_in) for k in param_mods: p_in[k] = param_mods[k](p_in[k]) p_in.write(params_out) class Parameters(object): ''' Disk-based representation of a PRMS parameter file. For the sake of memory efficiency, we only load parameters from ``base_file`` that get modifified through item assignment, for example >>> p = Parameters('example_params') >>> p['jh_coef'] = p['jh_coef']*1.1 >>> p.write('example_modified_params') will read parameter information from the params file to check that ``jh_coef`` is present in the parameter file, read the lines corresponding to ``jh_coef`` data and assign the new value as requested. Internally, a reference is kept to only modified parameter data, so when ``p.write(modified_params_file)`` is called, mostly this will copy from ``base_file`` to ``modified_params_file``. ''' def __init__(self, base_file): self.base_file = base_file self.base_file_reader = open(base_file) self.dimensions, self.base_params = self.__read_base(base_file) self.param_arrays = dict() def write(self, out_name): with open(self.base_file, 'r') as base_file: with open(out_name, 'w') as out_file: # write metadata out_file.write('File Auto-generated by PRMS-Python\n') out_file.write(datetime.datetime.now().isoformat() + '\n') # # write dimensions out_file.write('** Dimensions **\n') # write parameters; pre-sorted by data start line on read name_is_next = False params_start = False write_params_lines = False for l in base_file: if not params_start and l.strip() == '** Parameters **': out_file.write('** Parameters **\n') params_start = True elif l.strip() == '####': name_is_next = True elif name_is_next: name = l.strip().split()[0] if name not in self.param_arrays: out_file.write('####\n') out_file.write(name + '\n') name_is_next = False write_params_lines = True else: write_params_lines = False name_is_next = False elif write_params_lines: out_file.write(l.strip() + '\n') # write all parameters that had been accessed and/or modified for param, new_arr in self.param_arrays.items(): out_file.write('####\n') param_info = [el for el in self.base_params if el['name'] == param].pop() out_file.write(str(param_info['name']) + '\n') out_file.write(str(param_info['ndims']) + '\n') for dimname in param_info['dimnames']: out_file.write(dimname + '\n') out_file.write(str(param_info['length']) + '\n') out_file.write(str(param_info['vartype']) + '\n') out_file.writelines([str(a) + '\n' for a in new_arr.flatten()]) def __read_base(self, base_file): "Read base file returning 2-tuple of dimension and params dict" params_startline, dimensions = self.__make_dimensions_dict(base_file) base_params = self.__make_parameter_dict(base_file, params_startline) return (dimensions, base_params) def __make_dimensions_dict(self, base_file): """ Extract dimensions and each dimension length. Run before __make_parameter_dict. """ ret = OrderedDict() dim_name = '' dim_len = 0 # finished = False found_dim_start = False # while not finished: for idx, l in enumerate(self.base_file_reader): if l.strip() == '** Dimensions **': # start of dimensions found_dim_start = True elif '#' in l: # comments pass elif l.strip() == '** Parameters **': # start of parameters dimlines = idx # finished = True break elif found_dim_start: if dim_name == '': dim_name = l.strip() else: dim_len = int(l) ret.update({dim_name: dim_len}) dim_name = '' return (dimlines, ret) def __make_parameter_dict(self, base_file, params_startline=0): ret = [] name = '' ndims = 0 dimnames = [] length = 0 vartype = '' dimnames_read = 0 data_startline = 0 for idx, l in enumerate(self.base_file_reader): if '#' in l: # we have a comment; the next lines will be new # parameter metadata. No data for the first time through, so # we don't want to append an metadata blob with empty values if name: ret.append( dict( name=name, ndims=ndims, dimnames=dimnames, length=length, vartype=vartype, data_startline=data_startline ) ) name = '' ndims = 0 dimnames = [] length = 0 vartype = '' dimnames_read = 0 elif not name: name = l.strip().split()[0] # in case old format with integer after name elif not ndims: ndims = int(l.strip()) elif not (dimnames_read == ndims): dimnames.append(l.strip()) dimnames_read += 1 elif not length: length = int(l.strip()) elif not vartype: vartype = l.strip() # advance one from current position and account for starting # to count from zero data_startline = params_startline + idx + 2 # need to append one more time since iteration will have stopped after # last line ret.append( dict( name=name, ndims=ndims, dimnames=dimnames, length=length, vartype=vartype, data_startline=data_startline ) ) return ret def __getitem__(self, key): """ Look up a parameter by its name. Raises: KeyError if parameter name is not valid """ def load_parameter_array(param_metadata): startline = param_metadata['data_startline'] endline = startline + param_metadata['length'] + 1 param_slice = itertools.islice( io.open(self.base_file, 'rb'), startline, endline ) arr = np.genfromtxt(param_slice) if param_metadata['ndims'] > 1: dimsizes = [ self.dimensions[d] for d in param_metadata['dimnames'] ] dimsizes.reverse() arr = arr.reshape(dimsizes) return arr if key in self.param_arrays: return self.param_arrays[key] else: try: param_metadata = [ el for el in self.base_params if el['name'] == key ].pop() except IndexError: raise KeyError(key) arr = load_parameter_array(param_metadata) # cache the value for future access (but maybe shouldn't?) self.param_arrays.update({key: arr}) return arr def __setitem__(self, key, value): if key in self.param_arrays: cur_arr = self.param_arrays[key] if not value.shape == cur_arr.shape: raise ValueError('New array does not match existing') self.param_arrays[key] = value