Skip to content

Commit 4a4b4fc

Browse files
committed
generic Optimizer.plot_optimization method for srad and pet
1 parent 4837ac4 commit 4a4b4fc

1 file changed

Lines changed: 72 additions & 42 deletions

File tree

prms_python/optimizer.py

Lines changed: 72 additions & 42 deletions
Original file line numberDiff line numberDiff line change
@@ -181,9 +181,9 @@ def pet(self, reference_pet_path, station_nhru, n_sims=10, method='',\
181181
Args:
182182
reference_pet_path (str): path to measured pet data
183183
station_nhru (int or str): hru index in PRMS that is geographically
184-
near the measured/estimated pet location. TODO: if value
185-
if 'basin' then you wish to compare to area-weighted
186-
simulated pet (`basin_potet_1` in PRMS).
184+
near the measured/estimated pet location. If value
185+
is 'basin' then you wish to compare to area-weighted
186+
basin-wide simulated pet (`basin_potet_1` in PRMS).
187187
Kwargs:
188188
method (str): XXX not yet implemented- (default-Monte Carlo)
189189
n_sims (int): number of simulations to conduct
@@ -212,7 +212,7 @@ def pet(self, reference_pet_path, station_nhru, n_sims=10, method='',\
212212
self.control_file,
213213
OPJ(
214214
self.working_dir,
215-
'jh_coef:{0:.3f}'.format(np.mean(jh_coefs[i]))
215+
'jh_coef:{0:.5f}'.format(np.mean(jh_coefs[i]))
216216
)
217217
)
218218
for i in range(n_sims)
@@ -249,18 +249,19 @@ def pet(self, reference_pet_path, station_nhru, n_sims=10, method='',\
249249

250250
print('{0}\nOutput information sent to {1}\n'.format('-' * 80, json_outfile))
251251

252-
#TODO: make single plot function that takes stage as input (swrad, pet, flow...)
253-
def plot_srad_optimization(self, freq='daily', method='time_series',\
254-
plot_vars='both', return_fig=False):
252+
def plot_optimization(self, stage, freq='daily', method='time_series',\
253+
plot_vars='both', plot_1to1=True, return_fig=False):
255254
"""
256-
Basic plotting of current srad optimization results with limited options.
257-
Plots measured, original simluated, and optimization simulated swrad at
258-
the correspinding HRU either as time series (all three) or scatter (measured
259-
versus simulated). Not recommended for plotting results when n_sims,
260-
instead use plotting options from an OptimizationResult object, or employ
261-
a user-defined method using the result data if necessary.
255+
Basic plotting of current optimization results with limited options.
256+
Plots measured, original simluated, and optimization simulated variabes
257+
either swrad, pet, or streamflow (TODO) depending on stage at the
258+
corresponding HRU or basin-wide scale either as time series (all three)
259+
or scatter (measured versus simulated). Not recommended for plotting
260+
results when n_sims, instead use options from an OptimizationResult
261+
object, or employ a user-defined method using the result data if necessary.
262262
263263
Kwargs:
264+
stage (str): stage of optimization to plot (swrad,pet,flow)
264265
freq (str): frequency of time series plots, value can be 'daily'
265266
or 'monthly' for solar radiation
266267
method (str): 'time_series' for time series sub plot of each
@@ -273,22 +274,55 @@ def plot_srad_optimization(self, freq='daily', method='time_series',\
273274
'meas': plot simulated along with measured swrad
274275
'orig': plot simulated along with the original simulated swrad
275276
'both': plot simulated, with original simulation and measured
277+
plot_1to1 (bool): if True plot one to one line on correlation
278+
scatter plot, otherwise exclude.
276279
return_fig (bool): flag whether to return matplotlib figure
277280
Returns:
278281
f (matplotlib.figure.Figure): If kwarg return_fig=True, then return
279282
copy of the figure that is generated to the user.
280283
"""
281-
if not self.srad_outputs:
282-
raise ValueError('You have not run any srad optimizations')
283-
284-
# indices that measured and simulated swrad share (the intersection)
285-
X = self.measured_srad
286-
idx = X.index.intersection(self.srad_outputs[0]['statvar']\
287-
['swrad_{}'.format(self.srad_hru)].index)
288-
X = X[idx]
289-
orig = load_statvar(OPJ(self.input_dir, 'statvar.dat'))['swrad_{}'\
290-
.format(self.srad_hru)][idx]
291-
meas = self.measured_srad[idx]
284+
#use optimization stage to set plot parameters and get appropriate data
285+
if (stage=='swrad'):
286+
if not self.srad_outputs:
287+
raise ValueError('You have not run any srad optimizations')
288+
if self.srad_hru == 'basin':
289+
pet_srad_name = 'basin_swrad_1'
290+
else:
291+
srad_var_name = 'swrad_{}'.format(self.srad_hru)
292+
#indices that measured and simulated share (intersection)
293+
X = self.measured_srad
294+
idx = X.index.intersection(self.srad_outputs[0]['statvar']\
295+
['{}_{}'.format(stage, self.srad_hru)].index)
296+
X = X[idx]
297+
orig = load_statvar(OPJ(self.input_dir, 'statvar.dat'))['{}'\
298+
.format(srad_var_name)][idx]
299+
meas = self.measured_srad[idx]
300+
sims = [out['statvar']['{}'.format(srad_var_name)][idx] for \
301+
out in self.srad_outputs]
302+
simdirs = [out['simulation_dir'].split(os.sep)[-1].replace('_', ' ')\
303+
for out in self.srad_outputs]
304+
var_name = 'shortwave radiation'
305+
n = len(self.srad_outputs) # number of simulations to plot
306+
elif (stage=='pet'):
307+
if not self.pet_outputs:
308+
raise ValueError('You have not run any pet optimizations')
309+
if self.pet_hru == 'basin':
310+
pet_var_name = 'basin_potet_1'
311+
else:
312+
pet_var_name = 'potet_{}'.format(self.pet_hru)
313+
X = self.measured_pet
314+
idx = X.index.intersection(self.pet_outputs[0]['statvar']\
315+
['{}'.format(pet_var_name)].index)
316+
X = X[idx]
317+
orig = load_statvar(OPJ(self.input_dir, 'statvar.dat'))['{}'\
318+
.format(pet_var_name)][idx]
319+
meas = self.measured_pet[idx]
320+
sims = [out['statvar']['{}'.format(pet_var_name)][idx] for \
321+
out in self.pet_outputs]
322+
simdirs = [out['simulation_dir'].split(os.sep)[-1].replace('_', ' ')\
323+
for out in self.pet_outputs]
324+
var_name = 'potential ET'
325+
n = len(self.pet_outputs) # number of simulations to plot
292326
# styles for each plot
293327
ms = 4 # markersize for all points
294328
orig_sty = dict(linestyle='none',markersize=ms,\
@@ -300,7 +334,6 @@ def plot_srad_optimization(self, freq='daily', method='time_series',\
300334
sim_sty = dict(linestyle='none',markersize=ms,\
301335
markerfacecolor='none', marker='o',\
302336
markeredgecolor='r', color='r')
303-
n = len(self.srad_outputs) # number of simulations to plot
304337
## number of subplots and rows (two plots per row)
305338
nrow = n//2 # round down if odd n
306339
ncol = 2
@@ -315,19 +348,17 @@ def plot_srad_optimization(self, freq='daily', method='time_series',\
315348
fig, ax = plt.subplots(n, sharex=True, sharey=True,\
316349
figsize=(12,n*3.5))
317350
axs = ax.ravel()
318-
for i,out in enumerate(self.srad_outputs):
351+
for i,sim in enumerate(sims):
319352
if plot_vars in ('meas', 'both'):
320353
axs[i].plot(meas, label='Measured', **meas_sty)
321354
if plot_vars in ('orig', 'both'):
322355
axs[i].plot(orig, label='Original sim.', **orig_sty)
323-
axs[i].plot(out['statvar']['swrad_{}'.format(self.srad_hru)]\
324-
[idx], **sim_sty)
325-
axs[i].set_ylabel('sim: {}'.format(out['simulation_dir'].\
326-
split(os.sep)[-1].replace('_', ' ')), fontsize=10)
356+
axs[i].plot(sim, **sim_sty)
357+
axs[i].set_ylabel('sim: {}'.format(simdirs[i]), fontsize=10)
327358
if i == 0: axs[i].legend(markerscale=5, loc='best')
328359
fig.subplots_adjust(hspace=0)
329360
fig.autofmt_xdate()
330-
361+
#monthly means
331362
elif freq == 'monthly' and method == 'time_series':
332363
# compute monthly means
333364
meas = meas.groupby(meas.index.month).mean()
@@ -339,22 +370,21 @@ def plot_srad_optimization(self, freq='daily', method='time_series',\
339370

340371
fig, ax = plt.subplots(nrows=nrow, ncols=ncol, figsize=(12,n*3.5))
341372
axs = ax.ravel()
342-
for i,out in enumerate(self.srad_outputs):
373+
for i,sim in enumerate(sims):
343374
if plot_vars in ('meas', 'both'):
344375
axs[i].plot(meas, label='Measured', **meas_sty)
345376
if plot_vars in ('orig', 'both'):
346377
axs[i].plot(orig, label='Original sim.', **orig_sty)
347-
sim = out['statvar']['swrad_{}'.format(self.srad_hru)][idx]
348378
sim = sim.groupby(sim.index.month).mean()
349379
axs[i].plot(sim, **sim_sty)
350-
axs[i].set_ylabel('sim: {}\nmean swrad'.format(out['simulation_dir'].\
351-
split(os.sep)[-1].replace('_', ' ')), fontsize=10)
380+
axs[i].set_ylabel('sim: {}\nmean {}'.format(simdirs[i], stage),\
381+
fontsize=10)
352382
axs[i].set_xlim(0.5,12.5)
353383
if i == 0: axs[i].legend(markerscale=5, loc='best')
354384
if odd_n: # empty subplot if odd number of simulations
355385
fig.delaxes(axs[n])
356386
fig.text(0.5, 0.1, 'month')
357-
387+
#x-y scatter
358388
elif method == 'correlation':
359389
## figure
360390
fig, ax = plt.subplots(nrows=nrow, ncols=ncol, figsize=(12,n*3))
@@ -363,18 +393,18 @@ def plot_srad_optimization(self, freq='daily', method='time_series',\
363393
meas_min = min(X)
364394
meas_max = max(X)
365395

366-
for i, out in enumerate(self.srad_outputs):
367-
Y = out['statvar']['swrad_{}'.format(self.srad_hru)][idx]
396+
for i, sim in enumerate(sims):
397+
Y = sim
368398
sim_max = max(Y)
369399
sim_min = min(Y)
370400
m = max(meas_max,sim_max)
371-
axs[i].plot([0, m], [0, m], 'k--', lw=2) ## one to one line
401+
if plot_1to1:
402+
axs[i].plot([0, m], [0, m], 'k--', lw=2) ## one to one line
372403
axs[i].set_xlim(meas_min,meas_max)
373404
axs[i].set_ylim(sim_min, sim_max)
374405
axs[i].plot(X, Y, **sim_sty)
375-
axs[i].set_ylabel('sim: {}'.format(out['simulation_dir']\
376-
.split(os.sep)[-1].replace('_', ' ')))
377-
axs[i].set_xlabel('Measured shortwave radiation')
406+
axs[i].set_ylabel('sim: {}'.format(simdirs[i]))
407+
axs[i].set_xlabel('Measured {}'.format(var_name))
378408
axs[i].text(0.05, 0.95,r'$R^2 = {0:.2f}$'.format(\
379409
X.corr(Y)**2), fontsize=16,\
380410
ha='left', va='center', transform=axs[i].transAxes)

0 commit comments

Comments
 (0)