''' optimizer.py -- Optimization routines for PRMS parameters and data. ''' from __future__ import print_function import pandas as pd import numpy as np import matplotlib.pyplot as plt import datetime as dt import os, sys, json, re import multiprocessing as mp from copy import copy from copy import deepcopy from numpy import log10 from .data import Data from .parameters import Parameters from .simulation import Simulation, SimulationSeries from .util import load_statvar, nash_sutcliffe, percent_bias, rmse OPJ = os.path.join class Optimizer: ''' Container for a PRMS parameter optimization routine consisting of stages similar to what is described in Hay, et al, 2006 (ftp://brrftp.cr.usgs.gov/pub/mows/software/luca_s/jawraHay.pdf). Example: >>> from prms_python import Data, Optimizer, Parameters >>> params = Parameters('path/to/parameters') >>> data = Data('path/to/data') >>> control = 'path/to/control' >>> work_directory = 'path/to/create/simulations' >>> optr = Optimizer(params, data, control, work_directory, \ title='the title', description='desc') >>> measured = 'path/to/measured/csv' >>> statvar_name = 'basin_cfs' # or any other valid statvar >>> params_to_resample = ['dday_intcp', 'dday_slope'] # list >>> optr.monte_carlo(measured, params_to_resample, statvar_name) ''' #dic for min/max of parameter allowable ranges, add more when needed param_ranges = {'dday_intcp': (-60.0, 10.0), 'dday_slope': (0.2, 0.9),\ 'jh_coef': (0.005, 0.06), 'pt_alpha': (1.0, 2.0), \ 'potet_coef_hru_mo': (1.0, 1.6), 'tmax_index': \ (-10.0, 110.0)\ } #changed potet max def __init__(self, parameters, data, control_file, working_dir, title, description=None): if isinstance(parameters, Parameters): self.parameters = parameters else: raise TypeError('parameters must be instance of Parameters') if isinstance(data, Data): self.data = data else: raise TypeError('data must be instance of Data') input_dir = '{}'.format(os.sep).join(control_file.split(os.sep)[:-1]) if not os.path.isfile(OPJ(input_dir, 'statvar.dat')): print('You have no statvar.dat file in your model directory') print('Running PRMS on original data in {} for later comparison'\ .format(input_dir)) sim = Simulation(input_dir) sim.run() if not os.path.isdir(working_dir): os.mkdir(working_dir) self.input_dir = input_dir self.control_file = control_file self.working_dir = working_dir self.title = title self.description = description self.measured_arb = None # for arbitrary output methods self.statvar_name = None self.arb_outputs = [] def monte_carlo(self, reference_path, param_names, statvar_name, \ stage, n_sims=10, method='uniform',\ noise_factor=0.1, nproc=None): ''' Optimize the monthly dday_intcp and dday_slope parameters (two key parameters in the ddsolrad module in PRMS) by one of multiple methods (in development): Monte Carlo default method Args: reference_path (str): path to measured data for optimization param_names (list): list of parameter names to resample statvar_name (str): name of statisical variable output name for optimization stage (str): custom name of optimization stage e.g. 'ddsolrad' Kwargs: n_sims (int): number of simulations to conduct parameter optimization/uncertaitnty analysis. method (str): resampling method for parameters (normal or uniform) noise_factor (float): scales the variance of noise to add to parameter values when adding normal rv (method='normal') nproc (int): number of processors available to run PRMS simulations ''' if '_' in stage: raise ValueError('stage name cannot contain an underscore') # assign the optimization object a copy of measured srad for plots self.measured_arb = pd.Series.from_csv( reference_path, parse_dates=True ) # statistical variable output name self.statvar_name = statvar_name start_time = dt.datetime.now() start_time = start_time.replace(second=0, microsecond=0) # resample params for all simulations- potential place to serialize params = [] for name in param_names: # create list of lists of resampled params tmp = [] for idx in range(n_sims): tmp.append(resample_param(self.parameters, name, how=method,\ noise_factor=noise_factor)) params.append(list(tmp)) # SimulationSeries comprised of each resampled param set series = SimulationSeries( Simulation.from_data( self.data, _mod_params(self.parameters,\ [params[n][i] for n in range(len(params))],\ param_names), self.control_file, OPJ( self.working_dir, # name of sim: first param and mean value '{0}:{1:.6f}'.format(param_names[0], np.mean(params[0][i])) ) ) for i in range(n_sims) ) # run if not nproc: nproc = mp.cpu_count() // 2 outputs = list(series.run(nproc=nproc).outputs_iter()) self.arb_outputs.extend(outputs) # for current instance- add outputs end_time = dt.datetime.now() end_time = end_time.replace(second=0, microsecond=0) # json metadata for Monte Carlo method meta = { 'params_adjusted' : param_names, 'statvar_name' : self.statvar_name, 'optimization_title' : self.title, 'optimization_description' : self.description, 'start_time' : str(start_time), 'end_time' : str(end_time), 'measured' : reference_path, 'method' : 'Monte Carlo', 'noise_factor' : noise_factor, 'resample': method, 'sim_dirs' : [], 'stage': '{}'.format(stage), 'original_params' : self.parameters.base_file, 'nproc': nproc, 'n_sims' : n_sims } for output in outputs: meta['sim_dirs'].append(output['simulation_dir']) json_outfile = OPJ(self.working_dir, _create_metafile_name(\ self.working_dir, self.title, stage)) with open(json_outfile, 'w') as outf: json.dump(meta, outf, sort_keys = True, indent = 4,\ ensure_ascii = False) print('{0}\nOutput information sent to {1}\n'.\ format('-' * 80, json_outfile)) def plot_optimization(self, freq='daily', method='time_series',\ plot_vars='both', plot_1to1=True, return_fig=False,\ n_plots=4): """ Basic plotting of current optimization results with limited options. Plots measured, original simluated, and optimization simulated variabes Not recommended for plotting results when n_sims is very large, instead use options from an OptimizationResult object, or employ a user-defined method using the result data. Kwargs: freq (str): frequency of time series plots, value can be 'daily' or 'monthly' for solar radiation method (str): 'time_series' for time series sub plot of each simulation alongside measured radiation. Another choice is 'correlation' which plots each measured daily solar radiation value versus the corresponding simulated variable as subplots one for each simulation in the optimization. With coefficients of determiniationi i.e. square of pearson correlation coef. plot_vars (str): what to plot alongside simulated srad: 'meas': plot simulated along with measured swrad 'orig': plot simulated along with the original simulated swrad 'both': plot simulated, with original simulation and measured plot_1to1 (bool): if True plot one to one line on correlation scatter plot, otherwise exclude. return_fig (bool): flag whether to return matplotlib figure Returns: f (matplotlib.figure.Figure): If kwarg return_fig=True, then return copy of the figure that is generated to the user. """ if not self.arb_outputs: raise ValueError('You have not run any optimizations') var_name = self.statvar_name X = self.measured_arb idx = X.index.intersection(self.arb_outputs[0]['statvar']\ ['{}'.format(var_name)].index) X = X[idx] orig = load_statvar(OPJ(self.input_dir, 'statvar.dat'))['{}'\ .format(var_name)][idx] meas = self.measured_arb[idx] sims = [out['statvar']['{}'.format(var_name)][idx] for \ out in self.arb_outputs] simdirs = [out['simulation_dir'].split(os.sep)[-1].\ replace('_', ' ') for out in self.arb_outputs] var_name = '{}'.format(self.statvar_name) n = len(self.arb_outputs) # number of simulations to plot # user defined number of subplots from first n_plots results if (n > n_plots): n = n_plots # styles for each plot ms = 4 # markersize for all points orig_sty = dict(linestyle='none',markersize=ms,\ markerfacecolor='none', marker='s',\ markeredgecolor='royalblue', color='royalblue') meas_sty = dict(linestyle='none',markersize=ms+1,\ markerfacecolor='none', marker='1',\ markeredgecolor='k', color='k') sim_sty = dict(linestyle='none',markersize=ms,\ markerfacecolor='none', marker='o',\ markeredgecolor='r', color='r') ## number of subplots and rows (two plots per row) nrow = n//2 # round down if odd n ncol = 2 odd_n = False if n/2. - nrow == 0.5: nrow+=1 # odd number need extra row odd_n = True ######## ## Start plots depnding on key word arguments ######## if freq == 'daily' and method == 'time_series': fig, ax = plt.subplots(n, sharex=True, sharey=True,\ figsize=(12,n*3.5)) axs = ax.ravel() for i,sim in enumerate(sims[:n]): if plot_vars in ('meas', 'both'): axs[i].plot(meas, label='Measured', **meas_sty) if plot_vars in ('orig', 'both'): axs[i].plot(orig, label='Original sim.', **orig_sty) axs[i].plot(sim, **sim_sty) axs[i].set_ylabel('sim: {}'.format(simdirs[i]), fontsize=10) if i == 0: axs[i].legend(markerscale=5, loc='best') fig.subplots_adjust(hspace=0) fig.autofmt_xdate() #monthly means elif freq == 'monthly' and method == 'time_series': # compute monthly means meas = meas.groupby(meas.index.month).mean() orig = orig.groupby(orig.index.month).mean() # change line styles for monthly plots to lines not points for d in (orig_sty, meas_sty, sim_sty): d['linestyle'] = '-' d['marker'] = None fig, ax = plt.subplots(nrows=nrow, ncols=ncol, figsize=(12,n*3.5)) axs = ax.ravel() for i,sim in enumerate(sims[:n]): if plot_vars in ('meas', 'both'): axs[i].plot(meas, label='Measured', **meas_sty) if plot_vars in ('orig', 'both'): axs[i].plot(orig, label='Original sim.', **orig_sty) sim = sim.groupby(sim.index.month).mean() axs[i].plot(sim, **sim_sty) axs[i].set_ylabel('sim: {}\nmean'.format(simdirs[i]),\ fontsize=10) axs[i].set_xlim(0.5,12.5) if i == 0: axs[i].legend(markerscale=5, loc='best') if odd_n: # empty subplot if odd number of simulations fig.delaxes(axs[n]) fig.text(0.5, 0.1, 'month') #x-y scatter elif method == 'correlation': ## figure fig, ax = plt.subplots(nrows=nrow, ncols=ncol, figsize=(12,n*3)) axs = ax.ravel() ## subplot dimensions meas_min = min(X) meas_max = max(X) for i, sim in enumerate(sims[:n]): Y = sim sim_max = max(Y) sim_min = min(Y) m = max(meas_max,sim_max) if plot_1to1: axs[i].plot([0, m], [0, m], 'k--', lw=2) ## one to one line axs[i].set_xlim(meas_min,meas_max) axs[i].set_ylim(sim_min, sim_max) axs[i].plot(X, Y, **sim_sty) axs[i].set_ylabel('sim: {}'.format(simdirs[i])) axs[i].set_xlabel('Measured {}'.format(var_name)) axs[i].text(0.05, 0.95,r'$R^2 = {0:.2f}$'.format(\ X.corr(Y)**2), fontsize=16,\ ha='left', va='center', transform=axs[i].transAxes) if odd_n: # empty subplot if odd number of simulations fig.delaxes(axs[n]) if return_fig: return fig def _create_metafile_name(out_dir, opt_title, stage): """ Search through output directory where simulations are conducted look for all metadata simulation json files and find out if the current simulation is a replicate. Then use that information to build the correct file name for the output json file. The series are typically run in parallel that is why this step has to be done after running multiple simulations from an optimization stage. Args: out_dir (str): path to directory with model results, i.e. location where simulation series outputs and optimization json files are located, aka Optimizer.working_dir opt_title (str): optimization instance title for file search stage (str): stage of optimization, e.g. 'swrad', 'pet' Returns: name (str): file name for the current optimization simulation series metadata json file. E.g 'dry_creek_swrad_opt.json', or if this is the second time you have run an optimization titled 'dry_creek' the next json file will be returned as 'dry_creek_swrad_opt1.json' and so on with integer increments """ meta_re = re.compile(r'^{}_{}_opt(\d*)\.json'.format(opt_title, stage)) reps = [] for f in os.listdir(out_dir): if meta_re.match(f): nrep = meta_re.match(f).group(1) if nrep == '': reps.append(0) else: reps.append(nrep) if not reps: name = '{}_{}_opt.json'.format(opt_title, stage) else: # this is the nth optimization done under the same title n = max(map(int, reps)) + 1 name = '{}_{}_opt{}.json'.format(opt_title, stage, n) return name def resample_param(params, param_name, how='uniform', noise_factor=0.1): """ Resample PRMS parameter by shifting all values by a constant that is taken from a uniform distribution, where the range of the uniform values is equal to the difference between the min(max) of the parameter set and the min(max) of the allowable range from PRMS. For parameters that have array length <= 366 add noise to each parameter element by adding a RV from a normal distribution with mean 0, sigma = param allowable range / 10. Args: params (parameters.Parameters): parameter object param_name (str): name of PRMS parameter to resample Kwargs: how (str): distribution to resample parameters from in the case that each parameter element can be resampled (len <=366) Currently works for uniform and normal distributions. noise_factor (float): factor to multiply parameter range by, use the result as the standard deviation for the normal rand. variable used to add element wise noise. i.e. higher noise facter will result in higher noise added to each param element. Must be > 0. Returns: ret (numpy.ndarry): ndarray of param after uniform random mean shift or element-wise noise addition (normal r.v.) """ p_min, p_max = Optimizer.param_ranges.get(param_name,(-1,-1)) # create dictionary of parameter basic info (not values) param_dic = {param['name']: param for param in params.base_params} if not param_dic.get(param_name): raise KeyError('{} is not a valid parameter'.format(param_name)) if p_min == p_max == -1: raise ValueError("""{} has not been added to the dictionary of parameters to resample, add it's allowable min and max value to the param_ranges dictionary in the resample function in Optimizer.py""".format(param_name)) dim_case = None nhru = params.dimensions['nhru'] ndims = param_dic.get(param_name)['ndims'] dimnames = param_dic.get(param_name)['dimnames'] length = param_dic.get(param_name)['length'] param = params[param_name] # could expand list and check parameter name also e.g. cascade_flg # is a parameter that should not be changed dims_to_not_change = set(['ncascade','ncascdgw','nreach',\ 'nsegment']) if (len(set.intersection(dims_to_not_change, set(dimnames))) > 0): raise ValueError("""{} should not be resampled as it relates to the location of cascade flow parameters.""".format(param_name)) # use param info to get dimension info- e.g. if multidimensional if (ndims == 1 and length <= 366): dim_case = 'resample_each_value' # for smaller param dimensions elif (ndims == 1 and length > 366): dim_case = 'resample_all_values_once' # covers nssr, ngw, etc. elif (ndims == 2 and dimnames[1] == 'nmonths' and \ nhru == params.dimensions[dimnames[0]]): dim_case = 'nhru_nmonths' elif not dim_case: raise ValueError('The {} parameter should not be resampled'.\ format(param_name)) # #testing purposes # print('name: ', param_name) # print('max_val: ', p_max) # print('min_val: ', p_min) # print('ndims: ', ndims) # print('dimnames: ', dimnames) # print('length: ', length) # print('resample_method: ', dim_case) low_bnd = p_min - np.min(param) # lowest param value minus allowable min up_bnd = p_max - np.max(param) s = (p_max - p_min) * noise_factor # variance noise, default: range*(1/10) #do resampling differently based on param dimensions if dim_case == 'resample_all_values_once': if how == 'uniform': #uniform RV for shifting all values once shifted_param = np.random.uniform(low=low_bnd, high=up_bnd) + param ret = shifted_param elif how == 'normal': while True: tmp = np.random.normal(0, s, size=param.shape) + param if np.max(tmp) <= p_max and np.min(tmp) >= p_min: ret = tmp break elif dim_case == 'resample_each_value': ret = copy(param) if how == 'uniform': for i, el in enumerate(param): low_bnd = p_min - np.min(el) # a,b for uniform RV to add up_bnd = p_max - np.max(el) ret[i] = el + np.random.uniform(low=low_bnd, high=up_bnd) elif how == 'normal': for i, el in enumerate(param): while True: low_bnd = p_min - np.min(el) up_bnd = p_max - np.max(el) tmp = el + np.random.normal(0, s) if np.max(tmp) <= p_max and np.min(tmp) >= p_min: ret[i] = tmp break # nhru by nmonth dimensional params elif dim_case == 'nhru_nmonths': ret = copy(param) if how == 'uniform': rvs = [np.random.uniform(low=low_bnd, high=up_bnd)\ for i in range(12)] for month in range(12): ret[month] += rvs[month] elif how == 'normal': for i, el in enumerate(param): while True: low_bnd = p_min - np.min(el) up_bnd = p_max - np.max(el) tmp = el + np.random.normal(0, s) if np.max(tmp) <= p_max and np.min(tmp) >= p_min: ret[i] = tmp break return ret def _mod_params(parameters, params, param_names): ret = copy(parameters) for idx, param in enumerate(params): ret[param_names[idx]] = param return ret class OptimizationResult: def __init__(self, working_dir, stage): self.working_dir = working_dir self.stage = stage self.metadata_json_paths = self._get_optr_jsons(working_dir, stage) self.statvar_name = self.get_statvar_name(stage) self.measured = self.get_measured(stage) self.input_dir = self._get_input_dir(stage) self.input_params = self._get_input_params(stage) # if there are more than one input param for given stage if len(self.input_params) > 1: print( """Warning: there were more than one initial parameter sets used for the optimization for stage: {}. Make sure to compare the the correct input params with their corresponding output sims.""".format(stage)) print('\nThis optimization stage used the\ following input parameter files:\n{}'.format('\n'.join(self.input_params))) def _get_optr_jsons(self, work_dir, stage): """ Retrieve locations of optimization output jsons which contain metadata needed to understand optimization results. Create dictionary of each optimization with stage as key and lists of corresponding json file paths as values. Arguments: work_dir (str): path to directory with model results, i.e. location where simulation series outputs and optimization json files are located, aka Optimizer.working_dir stage (str): the stage ('ddsolrad', 'jhpet', 'flow', etc.) of the optimization in which to gather the jsons Returns: ret (dict): dictionary of stage (keys) and lists of json file paths for that stage (values). """ ret = {} optr_metafile_re = re.compile(r'^.*_{}_opt(\d*)\.json'.\ format(stage)) ret[stage] = [OPJ(work_dir, f) for f in\ os.listdir(work_dir) if\ optr_metafile_re.match(f)] return ret def _get_input_dir(self, stage): # retrieves the input directory from the first json file of given stage json_file = self.metadata_json_paths[stage][0] with open(json_file) as jf: meta_dic = json.load(jf) return '{}'.format(os.sep).join(meta_dic['original_params'].\ split(os.sep)[:-1]) def _get_input_params(self, stage): json_files = self.metadata_json_paths[stage] param_paths = [] for json_file in json_files: with open(json_file) as jf: meta_dic = json.load(jf) param_paths.append(meta_dic['original_params']) return list(set(param_paths)) def get_sim_dirs(self, stage): jsons = self.metadata_json_paths[stage] json_files = [] sim_dirs = [] for inf in jsons: with open(inf) as json_file: json_files.append(json.load(json_file)) for json_file in json_files: sim_dirs.extend(json_file['sim_dirs']) # list of all simulation directory paths for stage return sim_dirs def get_measured(self, stage): # only need to open one json file to get this information if not self.metadata_json_paths.get(stage): return # no optimization json files exist for given stage first_json = self.metadata_json_paths[stage][0] with open(first_json) as json_file: json_data = json.load(json_file) measured_series = pd.Series.from_csv(json_data.get('measured'),\ parse_dates=True) return measured_series def get_statvar_name(self, stage): # only need to open one json file to get this information try: first_json = self.metadata_json_paths[stage][0] except: raise ValueError("""No optimizatin has been run for stage: {}""".format(stage)) with open(first_json) as json_file: json_data = json.load(json_file) var_name = json_data.get('statvar_name') return var_name def result_table(self, freq='daily', top_n=5, latex=False): ##TODO: add stats for freq options monthly, annual (means or sum) sim_dirs = self.get_sim_dirs(self.stage) if top_n >= len(sim_dirs): top_n = len(sim_dirs) sim_names = [path.split(os.sep)[-1] for path in sim_dirs] meas_var = self.get_measured(self.stage) statvar_name = self.get_statvar_name(self.stage) result_df = pd.DataFrame(columns=\ ['NSE','RMSE','PBIAS','COEF_DET','ABS(PBIAS)']) for i, sim in enumerate(sim_dirs): sim_out = load_statvar(OPJ(sim, 'outputs', 'statvar.dat'))\ ['{}'.format(statvar_name)] idx = meas_var.index.intersection(sim_out.index) meas_var = copy(meas_var[idx]) sim_out = sim_out[idx] if freq == 'daily': result_df.loc[sim_names[i]] = [\ nash_sutcliffe(meas_var, sim_out),\ rmse(meas_var, sim_out),\ percent_bias(meas_var,sim_out),\ meas_var.corr(sim_out)**2,\ np.abs(percent_bias(meas_var, sim_out)) ] elif freq == 'monthly': meas_mo = meas_var.groupby(meas_var.index.month).mean() sim_out = sim_out.groupby(sim_out.index.month).mean() result_df.loc[sim_names[i]] = [\ nash_sutcliffe(meas_mo, sim_out),\ rmse(meas_mo, sim_out),\ percent_bias(meas_mo, sim_out),\ meas_mo.corr(sim_out)**2,\ np.abs(percent_bias(meas_mo, sim_out)) ] sorted_result = result_df.sort_values(by=['NSE','RMSE','ABS(PBIAS)',\ 'COEF_DET'], ascending=[False,True,True,False]) sorted_result.columns.name = '{} parameters'.format(self.stage) sorted_result = sorted_result[['NSE','RMSE','PBIAS','COEF_DET']] if latex: return sorted_result[:top_n].to_latex(escape=False) else: return sorted_result[:top_n] def get_top_ranked_sims(self, sorted_df): ## use result table to make dic with best param and statvar paths # index of table is the simulation directory names ret = { 'dir_name' : [], 'param_path' : [], 'statvar_path' : [] } for i,el in enumerate(sorted_df.index): ret['dir_name'].append(el) ret['param_path'].append(OPJ(self.working_dir,el,'inputs',\ 'parameters')) ret['statvar_path'].append(OPJ(self.working_dir,el,'outputs',\ 'statvar.dat')) return ret