Fix minor problems with QML

Co-authored-by: Florian Kiwit <[email protected]>
Co-authored-by: Maroua Marso <[email protected]>

src/modules/applications/QML/data_handler/DataHandler.py

@@ -106,7 +106,7 @@ def postprocess(self, input_data: dict, config: dict, **kwargs):
 
         if self.generalization_mark is not None:
             self.metrics.add_metric_batch({"KL_best": evaluation["KL_best"]})
-            metrics, _ = self.generalisation(input_data)
+            metrics, _ = self.generalisation()
 
             # Save generalisation metrics
             with open(f"{store_dir_iter}/record_gen_metrics_{kwargs['rep_count']}.pkl", 'wb') as f:
@@ -169,7 +169,7 @@ def data_load(self, gen_mod: dict, config: dict) -> dict:
         """
         pass
 
-    def generalisation(self, solution: list) -> (dict, float):
+    def generalisation(self) -> (dict, float):
         """
         Compute generalisation metrics
 
@@ -179,7 +179,10 @@ def generalisation(self, solution: list) -> (dict, float):
         :rtype: tuple(any, float)
 
         """
-        pass
+        # Compute your metrics here
+        metrics = {}  # Replace with actual metric calculations
+        time_taken = 0.0  # Replace with actual time calculation
+        return metrics, time_taken
 
     @abstractmethod
     def evaluate(self, solution: any) -> (dict, float):

src/modules/applications/QML/data_handler/DiscreteData.py

@@ -158,7 +158,7 @@ def data_load(self, gen_mod: dict, config: Config) -> dict:
 
         return application_config
 
-    def generalisation(self, solution: list) -> (dict, float):
+    def generalisation(self) -> (dict, float):
         """
         Calculate generalization metrics for the generated.
 

src/modules/applications/QML/generative_modeling/GenerativeModeling.py

@@ -50,13 +50,13 @@ def get_requirements() -> list[dict]:
     def get_solution_quality_unit(self) -> str:
         return "minimum KL"
 
-    def get_default_submodule(self, data_option: str) -> Union[ContinuousData, DiscreteData]:
-        if data_option == "Continuous Data":
+    def get_default_submodule(self, option: str) -> Union[ContinuousData, DiscreteData]:
+        if option == "Continuous Data":
             self.data = ContinuousData()
-        elif data_option == "Discrete Data":
+        elif option == "Discrete Data":
             self.data = DiscreteData()
         else:
-            raise NotImplementedError(f"Transformation Option {data_option} not implemented")
+            raise NotImplementedError(f"Transformation Option {option} not implemented")
         return self.data
 
     def get_parameter_options(self) -> dict:

src/modules/applications/QML/libraries/Library.py

@@ -28,10 +28,11 @@ class Library(Core, ABC):
     This class is an abstract base class for mapping a library-agnostic gate sequence to a library such as Qiskit 
     """
 
-    def __init__(self):
+    def __init__(self, name):
         """
         Constructor method
         """
+        self.name = name
         super().__init__()
 
     class Config(TypedDict):

src/modules/applications/QML/libraries/LibraryQiskit.py

@@ -37,7 +37,7 @@ def __init__(self):
         """
         Constructor method
         """
-        super().__init__()
+        super().__init__("LibraryQiskit")
         self.submodule_options = ["QCBM", "Inference"]
 
     @staticmethod

src/modules/applications/QML/training/Inference.py

@@ -112,9 +112,9 @@ def start_training(self, input_data: dict, config: Config, **kwargs: dict) -> di
 
         loss = self.kl_divergence(pmfs.reshape([-1, 1]), self.target)
 
-        input_data["best_parameter"] = parameters
+        input_data["best_parameter"] = parameters.get() if GPU else parameters
         input_data["inference"] = True
         input_data["KL"] = [loss.get() if GPU else loss]
-        input_data["best_sample"] = samples.astype(int)
+        input_data["best_sample"] = samples.astype(int).get() if GPU else samples.astype(int)
 
         return input_data

src/modules/applications/QML/training/QCBM.py

@@ -16,6 +16,7 @@
 import logging
 from cma import CMAEvolutionStrategy
 from tensorboardX import SummaryWriter
+from matplotlib import figure, axes
 import matplotlib.pyplot as plt
 
 from modules.applications.QML.training.Training import *
@@ -39,18 +40,12 @@ def __init__(self):
 
         self.n_states_range: list
         self.target: np.array
-        self.n_shots: int
-        self.execute_circuit: callable
-        self.timing: object
         self.study_generalization: bool
         self.generalization_metrics: dict
         self.writer: SummaryWriter
-        self.image_epochs: list
-        self.tb_images: bool
         self.loss_func: callable
-        self.n_params: int
-        self.x0: np.array
-        self.options: dict
+        self.fig: figure
+        self.ax: axes.Axes
 
     @staticmethod
     def get_requirements() -> list[dict]:
@@ -167,7 +162,7 @@ class Config(TypedDict):
     def get_default_submodule(self, option: str) -> Core:
         raise ValueError("This module has no submodules.")
 
-    def setup_training(self, input_data, config) -> dict:
+    def setup_training(self, input_data, config) -> tuple:
         """
         Method to configure the training setup including CMA-ES and tensorboard.
 
@@ -178,20 +173,10 @@ def setup_training(self, input_data, config) -> dict:
         :return: Updated input_data 
         :rtype: dict
         """
-        size = None
-        input_data['MPI_size'] = size
-        input_data["store_dir_iter"] += f"_{input_data['dataset_name']}_qubits{input_data['n_qubits']}"
+
         logging.info(
             f"Running config: [backend={input_data['backend']}] [n_qubits={input_data['n_qubits']}] "\
-            f"[n_shots={input_data['n_shots']}] [population_size={config['population_size']}]")
-
-        if comm.Get_rank() == 0:
-            self.target = np.asarray(input_data["histogram_train"])
-            self.target[self.target == 0] = 1e-8
-        self.n_states_range = range(2 ** input_data['n_qubits'])
-        self.n_shots = input_data["n_shots"]
-        self.execute_circuit = input_data["execute_circuit"]
-        self.timing = self.Timing()
+            f"[population_size={config['population_size']}]")
 
         self.study_generalization = "generalization_metrics" in list(input_data.keys())
         if self.study_generalization:
@@ -200,9 +185,6 @@ def setup_training(self, input_data, config) -> dict:
             input_data["store_dir_iter"] += f"_alpha{self.generalization_metrics.train_size}_depth{input_data['depth']}"
 
         self.writer = SummaryWriter(input_data["store_dir_iter"])
-        epochs = int(config['max_evaluations'] / config['population_size'])
-        self.image_epochs = range(0, epochs, 1) if int(epochs / 100) < 1 else range(0, epochs, int(epochs / 100))
-        self.tb_images = "2D" in input_data["dataset_name"]
 
         if config['loss'] == "KL":
             self.loss_func = self.kl_divergence
@@ -211,23 +193,22 @@ def setup_training(self, input_data, config) -> dict:
         else:
             raise NotImplementedError("Loss function not implemented")
 
-        self.n_params = len(input_data["circuit"].parameters)
-        self.x0 = (np.random.rand(self.n_params) - 0.5) * np.pi
-
+        n_params = len(input_data["circuit"].parameters)
+        x0 = (np.random.rand(n_params) - 0.5) * np.pi
         if config["pretrained"] != "False":
             parameters = np.load(config["pretrained"])
-            self.x0[:len(parameters)] = parameters
+            x0[:len(parameters)] = parameters
             logging.info(f'Training starting from parameters in path {config["pretrained"]}')
 
-        self.options = {
-            'bounds': [self.n_params * [-np.pi], self.n_params * [np.pi]],
+        options = {
+            'bounds': [n_params * [-np.pi], n_params * [np.pi]],
             'maxfevals': config['max_evaluations'],
             'popsize': config['population_size'],
             'verbose': -3,
             'tolfun': 1e-12
         }
 
-        return input_data
+        return x0, options
 
     def start_training(self, input_data: dict, config: Config, **kwargs: dict) -> (dict, float):
         """
@@ -242,76 +223,54 @@ def start_training(self, input_data: dict, config: Config, **kwargs: dict) -> (d
         :return: Dictionary including the information of previous modules as well as of the training
         :rtype: dict
         """
-        input_data = self.setup_training(input_data, config)
 
-        es = CMAEvolutionStrategy(self.x0.get() if GPU else self.x0, config['sigma'], self.options)
+        size = None # TODO: define correct mpi size
+        input_data['MPI_size'] = size
+        input_data["store_dir_iter"] += f"_{input_data['dataset_name']}_qubits{input_data['n_qubits']}"
+        x0, options = self.setup_training(input_data, config)
+
+        if comm.Get_rank() == 0:
+            self.target = np.asarray(input_data["histogram_train"])
+            self.target[self.target == 0] = 1e-8
+        self.n_states_range = range(2 ** input_data['n_qubits'])
+        execute_circuit = input_data["execute_circuit"]
+        timing = self.Timing()
+
+        es = CMAEvolutionStrategy(x0.get() if GPU else x0, config['sigma'], options)
         for parameter in ["best_parameters", "time_circuit", "time_loss", "KL", "best_sample"]:
             input_data[parameter] = []
 
         best_loss = float("inf")
-        fig, ax = plt.subplots()
+        self.fig, self.ax = plt.subplots()
         while not es.stop():
             solutions = es.ask()
             epoch = es.result[4]
             sigma = es.sigma
 
-            self.timing.start_recording()
-            pmfs_model, samples = self.execute_circuit(solutions)
+            timing.start_recording()
+            pmfs_model, samples = execute_circuit(solutions)
             pmfs_model = np.asarray(pmfs_model)
-            time_circ = self.timing.stop_recording()
+            time_circ = timing.stop_recording()
 
-            self.timing.start_recording()
+            timing.start_recording()
             if comm.Get_rank() == 0:
                 loss_epoch = self.loss_func(pmfs_model.reshape([config['population_size'], -1]), self.target)
             else:
                 loss_epoch = np.empty(config["population_size"])
             comm.Bcast(loss_epoch, root=0)
             comm.Barrier()
 
-            time_loss = self.timing.stop_recording()
+            time_loss = timing.stop_recording()
 
             es.tell(solutions, loss_epoch.get() if GPU else loss_epoch)
 
             if es.result[1] < best_loss:
                 best_loss = es.result[1]
-                index = loss_epoch.argmin()
-                best_pmf = pmfs_model[index] / pmfs_model[index].sum()
-                if self.study_generalization:
-
-                    if samples is None:
-                        counts = self.sample_from_pmf(
-                            n_states_range=self.n_states_range,
-                            pmf=best_pmf.get() if GPU else best_pmf,
-                            n_shots=self.generalization_metrics.n_shots)
-                    else:
-                        counts = samples[int(index)]
-
-                    metrics = self.generalization_metrics.get_metrics(counts.get() if GPU else counts)
-                    for (key, value) in metrics.items():
-                        self.writer.add_scalar(f"metrics/{key}", value, epoch)
-
-                nll = self.nll(best_pmf.reshape([1, -1]), self.target)
-                kl = self.kl_divergence(best_pmf.reshape([1, -1]), self.target)
-                # best_pmf_cpu = best_pmf.reshape([1, -1]) if GPU else best_pmf.reshape([1, -1])
-                self.writer.add_scalar("metrics/NLL", nll.get() if GPU else nll, epoch)
-                self.writer.add_scalar("metrics/KL", kl.get() if GPU else kl, epoch)
-                self.writer.add_scalar("time/circuit", time_circ, epoch)
-                self.writer.add_scalar("time/loss", time_loss, epoch)
-            input_data["best_parameters"].append(es.result[0])
+                best_pmf = self.data_visualization(loss_epoch, pmfs_model, samples, epoch)
 
+            input_data["best_parameters"].append(es.result[0])
             input_data["KL"].append(float(es.result[1]))
 
-            if self.tb_images and epoch in self.image_epochs:
-                ax.clear()
-                ax.imshow(
-                    best_pmf.reshape(int(np.sqrt(best_pmf.size)), int(np.sqrt(best_pmf.size))).get() if GPU
-                    else best_pmf.reshape(int(np.sqrt(best_pmf.size)),
-                                          int(np.sqrt(best_pmf.size))),
-                    cmap='binary',
-                    interpolation='none')
-                ax.set_title(f'Iteration {epoch}')
-                self.writer.add_figure('grid_figure', fig, global_step=epoch)
-
             logging.info(
                 f"[Iteration {es.result[4]}] "
                 f"[{config['loss']}: {es.result[1]:.5f}] "\
@@ -327,6 +286,41 @@ def start_training(self, input_data: dict, config: Config, **kwargs: dict) -> (d
         best_sample = self.sample_from_pmf(self.n_states_range,
                                            best_pmf.get() if GPU else best_pmf,
                                            n_shots=input_data["n_shots"])
-        input_data["best_sample"] = best_sample.get() if GPU else best_sample
+        input_data["best_sample"] = best_sample.get() if GPU else best_sample # pylint: disable=E1101
 
         return input_data
+
+
+    def data_visualization(self, loss_epoch, pmfs_model, samples, epoch):
+        index = loss_epoch.argmin()
+        best_pmf = pmfs_model[index] / pmfs_model[index].sum()
+        if self.study_generalization:
+
+            if samples is None:
+                counts = self.sample_from_pmf(
+                    n_states_range=self.n_states_range,
+                    pmf=best_pmf.get() if GPU else best_pmf,
+                    n_shots=self.generalization_metrics.n_shots)
+            else:
+                counts = samples[int(index)]
+
+            metrics = self.generalization_metrics.get_metrics(counts.get() if GPU else counts)
+            for (key, value) in metrics.items():
+                self.writer.add_scalar(f"metrics/{key}", value, epoch)
+
+        nll = self.nll(best_pmf.reshape([1, -1]), self.target)
+        kl = self.kl_divergence(best_pmf.reshape([1, -1]), self.target)
+        self.writer.add_scalar("metrics/NLL", nll.get() if GPU else nll, epoch)
+        self.writer.add_scalar("metrics/KL", kl.get() if GPU else kl, epoch)
+
+        self.ax.clear()
+        self.ax.imshow(
+            best_pmf.reshape(int(np.sqrt(best_pmf.size)), int(np.sqrt(best_pmf.size))).get() if GPU
+            else best_pmf.reshape(int(np.sqrt(best_pmf.size)),
+                                    int(np.sqrt(best_pmf.size))),
+            cmap='binary',
+            interpolation='none')
+        self.ax.set_title(f'Iteration {epoch}')
+        self.writer.add_figure('grid_figure', self.fig, global_step=epoch)
+
+        return best_pmf