Skip to content

Commit 6ad2ba5

Browse files
mt-digitalmt-digital
committed
added Simulation.from_data
1 parent ab5a1c7 commit 6ad2ba5

4 files changed

Lines changed: 1837 additions & 51 deletions

File tree

‎prms_python/scenario.py‎

Lines changed: 87 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,9 @@
77

88
from datetime import datetime
99

10-
from .parameters import modify_params
10+
from .parameters import modify_params, Parameters
11+
from .data import Data
12+
from .util import load_statvar
1113
from .simulation import Simulation
1214

1315

@@ -90,6 +92,43 @@ def __init__(self, base_dir, scenarios_dir, title=None, description=None):
9092
uuid_title_map={})
9193
self.scenarios = []
9294

95+
self.outputs = None
96+
97+
@classmethod
98+
def from_parameters_iter(cls, base_directory, parameters_iter,
99+
title=None, description=None):
100+
'''
101+
Create a ScenarioSeries from a list of parameters and the path to a
102+
directory in which scenarios should be saved and a data and control
103+
file can be found.
104+
'''
105+
series = cls(base_directory, base_directory,
106+
title=title, description=description)
107+
108+
for parameters in parameters_iter:
109+
110+
title = parameters['title'] if 'title' in parameters else None
111+
112+
uu = str(uuid.uuid4())
113+
114+
series.metadata['uuid_title_map'].update({uu: title})
115+
116+
scenario_dir = os.path.join(series.scenarios_dir, uu)
117+
118+
scenario = Scenario(series.base_dir, scenario_dir, title=title)
119+
120+
scenario.build()
121+
122+
series.scenarios.append(scenario)
123+
124+
with open(
125+
os.path.join(series.scenarios_dir, 'series_metadata.json'), 'w'
126+
) as f:
127+
f.write(json.dumps(series.metadata, indent=2))
128+
129+
def __len__(self):
130+
return len(self.scenarios)
131+
93132
def build(self, scenarios_list):
94133
"""
95134
Build the scenarios from a list of scenario definitions in dicitonary
@@ -154,6 +193,11 @@ def run(self, prms_exec='prms', nproc=None):
154193
pool = mp.Pool(processes=nproc)
155194
pool.map(_scenario_runner, self.scenarios)
156195

196+
# self.outputs = [
197+
# ScenarioOutput(uu, os.path.join(os.curdir(), d))
198+
# for uu, d in self.metadata['uuid_title_map'].items()
199+
# ]
200+
157201

158202
# multiprocessing req the function be def'd at root scope so it's picklable
159203
def _scenario_runner(scenario, prms_exec='prms'):
@@ -183,36 +227,38 @@ def __init__(self, base_dir, scenario_dir,
183227

184228
def build(self, param_mod_funs=None):
185229

186-
if not isinstance(param_mod_funs, dict):
187-
raise TypeError('param_mod_funs must be a dictionary')
188-
189-
# create scenario_dir that will be used as Simulation input dir
190-
if os.path.isdir(self.scenario_dir):
191-
shutil.rmtree(self.scenario_dir)
192-
193-
os.makedirs(self.scenario_dir)
194-
shutil.copy(
195-
os.path.join(self.base_dir, 'control'), self.scenario_dir
196-
)
197-
shutil.copy(
198-
os.path.join(self.base_dir, 'data'), self.scenario_dir
199-
)
200-
201-
old_params_path = os.path.join(self.base_dir, 'parameters')
202-
new_params_path = os.path.join(self.scenario_dir, 'parameters')
203-
if not param_mod_funs:
204-
shutil.copy(old_params_path, self.scenario_dir)
205-
else:
206-
modify_params(old_params_path, new_params_path, param_mod_funs)
230+
if isinstance(param_mod_funs, dict):
231+
232+
# create scenario_dir that will be used as Simulation input dir
233+
if os.path.isdir(self.scenario_dir):
234+
shutil.rmtree(self.scenario_dir)
235+
236+
os.makedirs(self.scenario_dir)
237+
shutil.copy(
238+
os.path.join(self.base_dir, 'control'), self.scenario_dir
239+
)
240+
shutil.copy(
241+
os.path.join(self.base_dir, 'data'), self.scenario_dir
242+
)
207243

208-
param_mod_funs_metadata = {
209-
param_name: inspect.getsource(param_mod_fun)
210-
for param_name, param_mod_fun in param_mod_funs.items()
211-
}
244+
old_params_path = os.path.join(self.base_dir, 'parameters')
245+
new_params_path = os.path.join(self.scenario_dir, 'parameters')
246+
if not param_mod_funs:
247+
shutil.copy(old_params_path, self.scenario_dir)
248+
else:
249+
modify_params(old_params_path, new_params_path, param_mod_funs)
212250

213-
self.metadata['mod_funs_dict'] = param_mod_funs_metadata
251+
param_mod_funs_metadata = {
252+
param_name: inspect.getsource(param_mod_fun)
253+
for param_name, param_mod_fun in param_mod_funs.items()
254+
}
214255

215-
self.simulation = Simulation(self.scenario_dir, self.scenario_dir)
256+
self.metadata['mod_funs_dict'] = param_mod_funs_metadata
257+
258+
self.simulation = Simulation(self.scenario_dir, self.scenario_dir)
259+
260+
else:
261+
self.simulation = Simulation(self.scenario_dir)
216262

217263
self.__simulation_ready = True
218264

@@ -250,3 +296,16 @@ def __setitem__(self, key, value):
250296
def write(self, output_path):
251297
with open(output_path, 'w') as f:
252298
f.write(json.dumps(self.metadata_dict))
299+
300+
301+
class ScenarioOutput:
302+
303+
def __init__(self, scenario_uu, scenario_directory, title=None):
304+
opj = os.path.join
305+
self.uuid = scenario_uu
306+
self.scenario_directory = scenario_directory
307+
self.title = title
308+
self.data = Data(opj(scenario_directory, 'data'))
309+
self.parameters = Parameters(opj(scenario_directory, 'parameters'))
310+
self.statvar = load_statvar(opj(scenario_directory, 'statvar.dat'))
311+
self.control = open(opj(scenario_directory, 'control')).read()

‎prms_python/simulation.py‎

Lines changed: 102 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,15 +1,40 @@
11
import glob
2+
import multiprocessing as mp
23
import os
34
import shutil
45
import subprocess
56

7+
from .data import Data
8+
from .parameters import Parameters
9+
10+
11+
class SimulationSeries(object):
12+
'''
13+
Series of simulations all to be run through a common interface
14+
'''
15+
16+
def __init__(self, simulations):
17+
self.series = simulations
18+
19+
def run(self, prms_exec='prms', nproc=None):
20+
21+
if not nproc:
22+
nproc = mp.cpu_count()//2
23+
24+
pool = mp.Pool(processes=nproc)
25+
pool.map(_simulation_runner, self.series)
26+
27+
28+
def _simulation_runner(sim):
29+
sim.run(prms_exec='prms')
30+
631

732
class Simulation(object):
833
"""
934
Simulation class for tracking the inputs and outputs of a single
1035
PRMS simulation.
1136
"""
12-
def __init__(self, input_dir, simulation_dir=None):
37+
def __init__(self, input_dir=None, simulation_dir=None):
1338
"""
1439
Create a new Simulation object from a simulation directory. Check that
1540
all required PRMS inputs (control, parameters, data) exist in the
@@ -28,37 +53,93 @@ def __init__(self, input_dir, simulation_dir=None):
2853
"""
2954
idir = input_dir
3055
self.input_dir = idir
31-
self.control_path = os.path.join(idir, 'control')
32-
self.parameters_path = os.path.join(idir, 'parameters')
33-
self.data_path = os.path.join(idir, 'data')
56+
self.simulation_dir = simulation_dir
57+
if idir is not None:
58+
self.control_path = os.path.join(idir, 'control')
59+
self.parameters_path = os.path.join(idir, 'parameters')
60+
self.data_path = os.path.join(idir, 'data')
3461

35-
if not os.path.exists(self.control_path):
36-
raise RuntimeError('Control file missing from ' + idir)
62+
if not os.path.exists(self.control_path):
63+
raise RuntimeError('Control file missing from ' + idir)
3764

38-
if not os.path.exists(self.parameters_path):
39-
raise RuntimeError('Parameter file missing from ' + idir)
65+
if not os.path.exists(self.parameters_path):
66+
raise RuntimeError('Parameter file missing from ' + idir)
4067

41-
if not os.path.exists(self.data_path):
42-
raise RuntimeError('Data file missing from ' + idir)
68+
if not os.path.exists(self.data_path):
69+
raise RuntimeError('Data file missing from ' + idir)
4370

44-
self.simulation_dir = simulation_dir
45-
if simulation_dir and simulation_dir != input_dir:
71+
if simulation_dir is not None:
72+
self.simulation_dir = simulation_dir
73+
if simulation_dir and simulation_dir != input_dir:
74+
75+
if os.path.exists(simulation_dir):
76+
shutil.rmtree(simulation_dir)
4677

47-
if os.path.exists(simulation_dir):
48-
shutil.rmtree(simulation_dir)
78+
os.mkdir(simulation_dir)
4979

50-
os.mkdir(simulation_dir)
80+
shutil.copy(self.control_path, simulation_dir)
81+
shutil.copy(self.data_path, simulation_dir)
82+
shutil.copy(self.parameters_path, simulation_dir)
5183

52-
shutil.copy(self.control_path, simulation_dir)
53-
shutil.copy(self.data_path, simulation_dir)
54-
shutil.copy(self.parameters_path, simulation_dir)
84+
self.control_path = os.path.join(simulation_dir, 'control')
85+
self.parameters_path = os.path.join(simulation_dir,
86+
'parameters')
87+
self.data_path = os.path.join(simulation_dir, 'data')
5588

56-
self.control_path = os.path.join(simulation_dir, 'control')
57-
self.parameters_path = os.path.join(simulation_dir, 'parameters')
58-
self.data_path = os.path.join(simulation_dir, 'data')
89+
else:
90+
self.control_path = None
91+
self.parameters_path = None
92+
self.data_path = None
93+
self.simulation_dir = None
5994

6095
self.has_run = False
6196

97+
@classmethod
98+
def from_data(cls, data, parameters, control_path, simulation_dir):
99+
'''
100+
Create a Simulation from a Data and Parameter object, plus a path
101+
to the control file, and providing a simulation_dir where the
102+
simulation should be run.
103+
104+
Args:
105+
data (Data): weather station data
106+
parameters (Parameters): simulation parameters
107+
control_path (str): path to control file
108+
simulation_dir (str): path to directory where simulations will be
109+
run and output will be stored. If it exists it will be
110+
overwritten.
111+
112+
Returns:
113+
(Simulation) simulation ready to be run using simulation_dir for
114+
inputs and outputs
115+
'''
116+
117+
if not isinstance(data, Data):
118+
raise TypeError('data must be instance of Data')
119+
120+
if not isinstance(parameters, Parameters):
121+
raise TypeError('parameters must be instance of Parameters')
122+
123+
if os.path.exists(simulation_dir):
124+
shutil.rmtree(simulation_dir)
125+
126+
os.makedirs(simulation_dir)
127+
128+
sim = cls()
129+
sim.simulation_dir = simulation_dir
130+
131+
sd = simulation_dir
132+
opj = os.path.join
133+
134+
data_path = opj(sd, 'data')
135+
data.write(data_path)
136+
params_path = opj(sd, 'parameters')
137+
parameters.write(params_path)
138+
shutil.copy(control_path, opj(sd, 'control'))
139+
140+
return sim
141+
142+
62143
def run(self, prms_exec='prms'):
63144

64145
cwd = os.getcwd()

0 commit comments

Comments
 (0)