Merge pull request #162 from DaveCasson/develop

Addition of arguments to ensemble.py and simulation.py for configuration writing. Minor addition file_manager.py

Author: arbennett

pysumma/ensemble.py

@@ -92,8 +92,8 @@ def _generate_coords(self):
                 decision_dims[k] = v
             for k, v in conf.get('file_manager', {}).items():
                 manager_dims[k] = v
-            #for k, v in conf.get('parameters', {}).items():
-            #    parameter_dims[k] = v
+            for k, v in conf.get('parameters', {}).items():
+                parameter_dims[k] = v
             for k, v in conf.get('trial_parameters', {}).items():
                 parameter_dims[k] = v
         return {'decisions': decision_dims,
@@ -133,13 +133,15 @@ def merge_output(self):
 
     def open_output(self):
         """
-        Open all of the output datasets from the ensembe and
+        Open all of the output datasets from the ensemble and
         return as a dictionary of datasets
         """
         return {n: s.output for n, s in self.simulations.items()}
 
 
-    def start(self, run_option: str='local', prerun_cmds: list=None):
+    def start(self, run_option: str='local',  run_suffix: str='',
+              prerun_cmds: list=None, freq_restart: str=None,
+              write_config: bool=True,write_manager: bool=False):
         """
         Start running the ensemble members.
 
@@ -154,9 +156,13 @@ def start(self, run_option: str='local', prerun_cmds: list=None):
             # Sleep calls are to ensure writeout happens
             config = self.configuration[n]
             self.submissions.append(self._client.submit(
-                _submit, s, n, run_option, prerun_cmds, config))
+                _submit, s, n, run_option,
+                run_suffix, prerun_cmds, freq_restart,
+                write_config,write_manager,config))
 
-    def run(self, run_option: str='local', prerun_cmds=None, monitor: bool=True):
+    def run(self, run_option: str='local', run_suffix: str='',
+            prerun_cmds: list=None, freq_restart: str=None, write_config: bool=True,
+            write_manager: bool=False, monitor: bool=True):
         """
         Run the ensemble
 
@@ -169,7 +175,8 @@ def run(self, run_option: str='local', prerun_cmds=None, monitor: bool=True):
         monitor:
             Whether to halt operation until runs are complete
         """
-        self.start(run_option, prerun_cmds)
+        self.start(run_option, run_suffix, prerun_cmds,
+                   freq_restart, write_config, write_manager)
         if monitor:
             return self.monitor()
         else:
@@ -191,11 +198,15 @@ def map(self, fun, args, include_sims=True, monitor=True):
 
     def monitor(self):
         """
-        Halt computation until submitted simulations are complete
+        Halt computation until submitted simulations are complete.
+        Update submission name from config.
+
         """
         simulations = self._client.gather(self.submissions)
-        for s in simulations:
-            self.simulations[s.run_suffix] = s
+        names = self.simulations.keys()
+
+        for s,name in zip(simulations,names):
+            self.simulations[name] = s
 
     def summary(self):
         """
@@ -211,8 +222,9 @@ def summary(self):
                 other.append(n)
         return {'Success': success, 'Error': error, 'Other': other}
 
-    def rerun_failed(self, run_option: str='local', prerun_cmds=None,
-                     monitor: bool=True):
+    def rerun_failed(self, run_option: str='local', run_suffix: str='',
+              prerun_cmds: list=None, freq_restart: str=None, write_config: bool=True,
+              write_manager: bool=False, monitor: bool=True):
         """
         Try to re-run failed simulations.
 
@@ -232,18 +244,27 @@ def rerun_failed(self, run_option: str='local', prerun_cmds=None,
             s = self.simulations[n]
             s.reset()
             self.submissions.append(self._client.submit(
-                _submit, s, n, run_option, prerun_cmds, config))
+                _submit, s, n, run_option, run_suffix,
+                prerun_cmds, freq_restart,write_config,
+                write_manager, config))
+
         if monitor:
             return self.monitor()
         else:
             return True
 
+def _submit(s: Simulation, name: str, run_option: str,  run_suffix: str,
+            prerun_cmds: list, freq_restart: str,write_config: bool,
+            write_manager: bool,config: Dict, **kwargs):
+    """
+    Run simulation objects
+    """
 
-def _submit(s: Simulation, name: str, run_option: str, prerun_cmds: List,
-            config: Dict, **kwargs):
     s.initialize()
     s.apply_config(config)
-    s.run(run_option, run_suffix=name, prerun_cmds=prerun_cmds, freq_restart='e')
+
+    s.run(run_option, run_suffix=run_suffix, freq_restart=freq_restart,
+          write_config=write_config,write_manager=write_manager)
     s.process = None
     return s
 

pysumma/file_manager.py

@@ -108,7 +108,10 @@ def trial_params(self):
 
     @property
     def initial_conditions(self):
-        p1 = self.get_value('settingsPath')
+        if self.get_value('statePath'):
+            p1 = self.get_value('statePath')
+        else:
+            p1 = self.get_value('settingsPath')
         p2 = self.get_value('initConditionFile')
         with xr.open_dataset(p1 + p2) as self._init_cond:
         	self._init_cond.load()

pysumma/simulation.py

@@ -114,6 +114,8 @@ def apply_config(self, config: dict):
         """
         if 'file_manager' in config:
             self.manager_path = Path(os.path.abspath(config['file_manager']))
+        for k, v in config.get('manager', {}).items():
+            self.manager.set_option(k, v)
         for k, v in config.get('decisions', {}).items():
             self.decisions.set_option(k, v)
         for k, v in config.get('parameters', {}).items():
@@ -302,7 +304,8 @@ def _run_docker(self, run_suffix, processes=1,
 
     def start(self, run_option='local',  run_suffix='pysumma_run', processes=1,
               prerun_cmds=[], startGRU=None, countGRU=None, iHRU=None,
-              freq_restart=None, progress=None, **kwargs):
+              freq_restart=None, write_config=True,write_manager=False,
+              progress=None, **kwargs):
         """
         Run a SUMMA simulation without halting. Most likely this should
         not be used. Use the ``run`` method for most common use cases.
@@ -311,7 +314,10 @@ def start(self, run_option='local',  run_suffix='pysumma_run', processes=1,
         if not prerun_cmds:
             prerun_cmds = []
         self.run_suffix = run_suffix
-        self._write_configuration(name=run_suffix)
+        if write_config:
+            self._write_configuration(name=run_suffix)
+        if write_manager:
+            self._write_file_manager()
         if run_option == 'local':
             self._run_local(run_suffix, processes, prerun_cmds,
                             startGRU, countGRU, iHRU, freq_restart, progress)
@@ -324,7 +330,8 @@ def start(self, run_option='local',  run_suffix='pysumma_run', processes=1,
 
     def run(self, run_option='local',  run_suffix='pysumma_run', processes=1,
             prerun_cmds=None, startGRU=None, countGRU=None, iHRU=None,
-            freq_restart=None, progress=None, **kwargs):
+            freq_restart=None, write_config=True,write_manager=False,
+            progress=None, **kwargs):
         """
         Run a SUMMA simulation and halt until completion or error.
 
@@ -360,7 +367,8 @@ def run(self, run_option='local',  run_suffix='pysumma_run', processes=1,
             hourly output.
         """
         self.start(run_option, run_suffix, processes, prerun_cmds,
-                   startGRU, countGRU, iHRU, freq_restart, progress, **kwargs)
+                   startGRU, countGRU, iHRU, freq_restart,write_config,
+                   write_manager, progress, **kwargs)
         self.monitor()
 
     def monitor(self):
@@ -452,6 +460,9 @@ def _write_configuration(self, name=''):
         with open(settings_path / self.manager['vegTableFile'].value, 'w+') as f:
             f.writelines(self.vegparm)
 
+    def _write_file_manager(self, name=''):
+        self.manager.write(path=self.config_path.parent)
+
     def get_forcing_data_list(self) -> List[xr.Dataset]:
         return self.force_file_list.open_forcing_data()