Cleanup of main and optimizer initializations

c3/experiment.py

@@ -80,9 +80,9 @@ def set_created_by(self, config):
 
         self.created_by = config
 
-    def quick_setup(self, filepath: str) -> None:
+    def load_quick_setup(self, filepath: str) -> None:
         """
-        Load a quick setup file and create all necessary components.
+        Load a quick setup file.
 
         Parameters
         ----------
@@ -92,7 +92,18 @@ def quick_setup(self, filepath: str) -> None:
         """
         with open(filepath, "r") as cfg_file:
             cfg = hjson.loads(cfg_file.read())
+        self.quick_setup(cfg)
 
+    def quick_setup(self, cfg) -> None:
+        """
+        Load a quick setup cfg and create all necessary components.
+
+        Parameters
+        ----------
+        cfg : Dict
+            Configuration options
+
+        """
         model = Model()
         model.read_config(cfg["model"])
         gen = Generator()

c3/main.py

@@ -5,79 +5,82 @@
 import os
 import hjson
 import argparse
-import c3.utils.parsers as parsers
 import c3.utils.tf_utils as tf_utils
 import tensorflow as tf
 from c3.parametermap import ParameterMap
 from c3.experiment import Experiment
 from c3.system.model import Model
 from c3.generator.generator import Generator
 
-logging.getLogger("tensorflow").disabled = True
+from c3.optimizers.c1 import C1
+from c3.optimizers.c2 import C2
+from c3.optimizers.c3 import C3
+from c3.optimizers.sensitivity import SET
 
-# flake8: noqa: C901
-if __name__ == "__main__":
-
-    os.nice(5)  # keep responsiveness when we overcommit memory
 
-    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("master_config")
-    args = parser.parse_args()
-
-    opt_config = args.master_config
-    with open(opt_config, "r") as cfg_file:
-        try:
-            cfg = hjson.loads(cfg_file.read())
-        except hjson.decoder.HjsonDecodeError:
-            raise Exception(f"Config {opt_config} is invalid.")
+logging.getLogger("tensorflow").disabled = True
 
-    optim_type = cfg["optim_type"]
+# flake8: noqa: C901
+def run_cfg(cfg, opt_config_filename, debug=False):
+    """Execute an optimization problem described in the cfg file.
+
+    Parameters
+    ----------
+    cfg : Dict[str, Union[str, int, float]]
+        Configuration file containing optimization options and information needed to completely
+        setup the system and optimization problem.
+    debug : bool, optional
+        Skip running the actual optimization, by default False
+
+    """
+    optim_type = cfg.pop("optim_type")
+    optim_lib = {
+        "C1": C1,
+        "C2": C2,
+        "C3": C3,
+        "C3_confirm": C3,
+        "confirm": C3,
+        "SET": SET,
+    }
+    if not optim_type in optim_lib:
+        raise Exception("C3:ERROR:Unknown optimization type specified.")
 
     tf_utils.tf_setup()
     with tf.device("/CPU:0"):
         model = None
         gen = None
+        exp = None
         if "model" in cfg:
             model = Model()
-            model.read_config(cfg["model"])
+            model.read_config(cfg.pop("model"))
         if "generator" in cfg:
             gen = Generator()
-            gen.read_config(cfg["generator"])
+            gen.read_config(cfg.pop("generator"))
         if "instructions" in cfg:
             pmap = ParameterMap(model=model, generator=gen)
-            pmap.read_config(cfg["instructions"])
+            pmap.read_config(cfg.pop("instructions"))
             exp = Experiment(pmap)
         if "exp_cfg" in cfg:
             exp = Experiment()
-            exp.read_config(cfg["exp_cfg"])
-        else:
+            exp.read_config(cfg.pop("exp_cfg"))
+        if exp is None:
             print("C3:STATUS: No instructions specified. Performing quick setup.")
             exp = Experiment()
-            exp.quick_setup(opt_config)
-
-        if optim_type == "C1":
-            opt = parsers.create_c1_opt(opt_config, exp)
-            if cfg.pop("include_model", False):
-                opt.include_model()
-        elif optim_type == "C2":
-            eval_func = cfg["eval_func"]
-            opt = parsers.create_c2_opt(opt_config, eval_func)
-        elif optim_type == "C3" or optim_type == "C3_confirm":
-            print("C3:STATUS: creating c3 opt ...")
-            opt = parsers.create_c3_opt(opt_config)
-        elif optim_type == "SET":
-            print("C3:STATUS: creating set obj")
-            opt = parsers.create_sensitivity(opt_config)
-        elif optim_type == "confirm":
-            print("C3:STATUS: creating c3 opt ...")
-            opt = parsers.create_c3_opt(opt_config)
-            opt.inverse = True
-        else:
-            raise Exception("C3:ERROR:Unknown optimization type specified.")
+            exp.quick_setup(cfg)
+
+        exp.set_opt_gates(cfg.pop("opt_gates", None))
+        if "gateset_opt_map" in cfg:
+            exp.pmap.set_opt_map(
+                [[tuple(par) for par in pset] for pset in cfg.pop("gateset_opt_map")]
+            )
+        if "exp_opt_map" in cfg:
+            exp.pmap.set_opt_map(
+                [[tuple(par) for par in pset] for pset in cfg.pop("exp_opt_map")]
+            )
+
+        opt = optim_lib[optim_type](**cfg, pmap=exp.pmap)
         opt.set_exp(exp)
-        opt.set_created_by(opt_config)
+        opt.set_created_by(opt_config_filename)
 
         if "initial_point" in cfg:
             initial_points = cfg["initial_point"]
@@ -94,18 +97,12 @@
                         "C3:STATUS:Loading initial point from : "
                         f"{os.path.abspath(init_point)}"
                     )
-                    init_dir = os.path.basename(
-                        os.path.normpath(os.path.dirname(init_point))
-                    )
                 except FileNotFoundError as fnfe:
                     raise Exception(
                         f"C3:ERROR:No initial point found at "
                         f"{os.path.abspath(init_point)}. "
                     ) from fnfe
 
-        if "real_params" in cfg:
-            real_params = cfg["real_params"]
-
         if optim_type == "C1":
             if "adjust_exp" in cfg:
                 try:
@@ -121,23 +118,24 @@
                         f"{os.path.abspath(adjust_exp)} "
                         "Continuing with default."
                     ) from fnfe
-            opt.optimize_controls()
 
-        elif optim_type == "C2":
-            opt.optimize_controls()
+        if not debug:
+            opt.run()
 
-        elif optim_type == "C3" or optim_type == "confirm":
-            opt.read_data(cfg["datafile"])
-            opt.learn_model()
 
-        elif optim_type == "C3_confirm":
-            opt.read_data(cfg["datafile"])
-            opt.log_setup()
-            opt.confirm()
+if __name__ == "__main__":
+    os.nice(5)  # keep responsiveness when we overcommit memory
 
-        elif optim_type == "SET":
-            opt.read_data(cfg["datafile"])
-            opt.log_setup()
+    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
 
-            print("sensitivity test ...")
-            opt.sensitivity()
+    parser = argparse.ArgumentParser()
+    parser.add_argument("master_config")
+    args = parser.parse_args()
+
+    opt_config = args.master_config
+    with open(opt_config, "r") as cfg_file:
+        try:
+            cfg = hjson.load(cfg_file)
+        except hjson.decoder.HjsonDecodeError:
+            raise Exception(f"Config {opt_config} is invalid.")
+    run_cfg(cfg, opt_config)

c3/optimizers/c1.py

@@ -4,9 +4,14 @@
 import shutil
 import time
 import tensorflow as tf
+from typing import Callable, List
+
 from c3.optimizers.optimizer import Optimizer
 from c3.utils.utils import log_setup
 
+from c3.libraries.algorithms import algorithms
+from c3.libraries.fidelities import fidelities
+
 
 class C1(Optimizer):
     """
@@ -39,43 +44,79 @@ class C1(Optimizer):
 
     def __init__(
         self,
-        dir_path,
         fid_func,
         fid_subspace,
         pmap,
-        callback_fids=[],
+        dir_path=None,
+        callback_fids=None,
         algorithm=None,
         store_unitaries=False,
         options={},
         run_name=None,
         interactive=True,
+        include_model=False,
         logger=None,
         fid_func_kwargs={},
     ) -> None:
+        if type(algorithm) is str:
+            algorithm = algorithms[algorithm]
         super().__init__(
             pmap=pmap,
             algorithm=algorithm,
             store_unitaries=store_unitaries,
             logger=logger,
         )
-        self.fid_func = fid_func
+        self.set_fid_func(fid_func)
+        self.callback_fids: List[Callable] = []
+        if callback_fids:
+            self.set_callback_fids(callback_fids)
         self.fid_subspace = fid_subspace
-        self.callback_fids = callback_fids
         self.options = options
         self.__dir_path = dir_path
         self.__run_name = run_name
         self.interactive = interactive
+        self.update_model = include_model
         self.fid_func_kwargs = fid_func_kwargs
+        self.run = (
+            self.optimize_controls
+        )  # Alias the legacy name for the method running the
+        # optimization
+
+    def set_fid_func(self, fid_func) -> None:
+        if type(fid_func) is str:
+            if self.pmap.model.lindbladian:
+                fid = "lindbladian_" + fid_func
+            else:
+                fid = fid_func
+            try:
+                self.fid_func = fidelities[fid]
+            except KeyError:
+                raise Exception(f"C3:ERROR:Unkown goal function: {fid} ")
+            print(f"C3:STATUS:Found {fid} in libraries.")
+        else:
+            self.fid_func = fid_func
+
+    def set_callback_fids(self, callback_fids) -> None:
+        if self.pmap.model.lindbladian:
+            cb_fids = ["lindbladian_" + f for f in callback_fids]
+        else:
+            cb_fids = callback_fids
+        for cb_fid in cb_fids:
+            try:
+                cb_fid_func = fidelities[cb_fid]
+            except KeyError:
+                raise Exception(f"C3:ERROR:Unkown goal function: {cb_fid}")
+            print(f"C3:STATUS:Found {cb_fid} in libraries.")
+            self.callback_fids.append(cb_fid_func)
 
     def log_setup(self) -> None:
         """
         Create the folders to store data.
         """
-        dir_path = os.path.abspath(self.__dir_path)
         run_name = self.__run_name
         if run_name is None:
             run_name = "c1_" + self.fid_func.__name__ + "_" + self.algorithm.__name__
-        self.logdir = log_setup(dir_path, run_name)
+        self.logdir = log_setup(self.__dir_path, run_name)
         self.logname = "open_loop.log"
         if isinstance(self.exp.created_by, str):
             shutil.copy2(self.exp.created_by, self.logdir)

c3/optimizers/c2.py

@@ -31,22 +31,23 @@ class C2(Optimizer):
 
     def __init__(
         self,
-        dir_path,
         eval_func,
         pmap,
         algorithm,
+        dir_path=None,
+        exp_type=None,
         exp_right=None,
         options={},
         run_name=None,
     ):
         super().__init__(pmap=pmap, algorithm=algorithm)
-        self.eval_func = eval_func
+        self.set_eval_func(eval_func, exp_type)
         self.options = options
         self.exp_right = exp_right
         self.__dir_path = dir_path
         self.__run_name = run_name
 
-    def set_eval_func(self, eval_func):
+    def set_eval_func(self, eval_func, exp_type):
         """
         Setter for the eval function.
 
@@ -56,6 +57,7 @@ def set_eval_func(self, eval_func):
             Function to be evaluated
 
         """
+        # TODO: Implement shell for experiment communication
         self.eval_func = eval_func
 
     def log_setup(self) -> None:
@@ -70,11 +72,10 @@ def log_setup(self) -> None:
             User specified name for the run
 
         """
-        dir_path = os.path.abspath(self.__dir_path)
         run_name = self.__run_name
         if run_name is None:
             run_name = self.eval_func.__name__ + self.algorithm.__name__
-        self.logdir = log_setup(dir_path, run_name)
+        self.logdir = log_setup(self.__dir_path, run_name)
         self.logname = "calibration.log"
 
         # We create a copy of the source code of the evaluation function in the log