Add option on some functions to return used ResComp object

Author: lostinnowhere314

rescomp/optimizer/optimizer_controller.py

@@ -34,7 +34,7 @@ def __init__(self, system, map_initial, prediction_type, method, res_ode=None,
                 parallel=False, parallel_profile=None, **res_params):
         """
         Arguments:
-            system (string or template.System): the system to use. If not a template.System object, will attempt to load one using rescomp.optimizer.get_system(system)
+            system (string or rescomp.optimizer.System): the system to use. If not a rescomp.optimizer.System object, will attempt to load one using rescomp.optimizer.get_system(system)
             map_initial (string): initial condition mapping for reservoir computer to use
             prediction_type (string): 'random' or 'continue'; prediction type to use while optimizing.
             method (string): training method; 'standard' or 'augmented'
@@ -141,7 +141,7 @@ def run_tests(self, test_ntrials, lyap_reps=20, parameters=None):
                 results_dict["rand_deriv_fit"].append(rand_df)
         return results_dict
 
-    def generate_orbits(self, n_orbits, parameters=None):
+    def generate_orbits(self, n_orbits, parameters=None, return_rescomp=False):
         """
         Trains a reservoir computer and has it predict, using the given hyperparameters
         If parameters are not specified, uses the optimized hyperparameters.
@@ -158,12 +158,12 @@ def generate_orbits(self, n_orbits, parameters=None):
 
         if self.parallel:
             results, _ = self._run_n_times_parallel(n_orbits, create_orbit,
-                        self.system, self.prediction_type, **self.res_params, **parameters)
+                        self.system, self.prediction_type, **self.res_params, **parameters, return_rescomp=return_rescomp)
             #Collapse into single list of outputs
             results = [item for sublist in results for item in sublist]
         else:
             results = self._run_n_times(n_orbits, create_orbit,
-                        self.system, self.prediction_type, **self.res_params, **parameters)
+                        self.system, self.prediction_type, **self.res_params, **parameters, return_rescomp=return_rescomp)
 
         return results
 

rescomp/optimizer/optimizer_functions.py

@@ -120,7 +120,7 @@ def meanlyap(system, rcomp, pre, r0, ts, pert_size=1e-6, lyap_reps=20):
         lam += rc.lyapunov(ts[:i], pre[:i, :], predelta[:i, :], delta0)
     return lam / lyap_reps
 
-def create_orbit(*args, **kwargs):
+def create_orbit(*args, return_rescomp=False, **kwargs):
     """ 
     Trains a reservoir computer and has it predict, using the given arguments
     
@@ -148,7 +148,10 @@ def create_orbit(*args, **kwargs):
     init_cond = make_initial(pred_type, rcomp, Uts)
     pre = rcomp_prediction(system, rcomp, ts, init_cond)
 
-    return tr, Utr, ts, Uts, pre
+    if return_rescomp:
+        return rcomp, tr, Utr, ts, Uts, pre
+    else:
+        return tr, Utr, ts, Uts, pre
 
 #########################
 ## Rescomp creation