Repository: cirKITers/masKIT

Issue:
Remove iris leftover
During creation of `datasets` submodule the loading of iris dataset was integrated there. But the old method does still exist in [maskit/iris.py](https://github.com/cirKITers/masKIT/blob/342378fdfc8affec5e30230628eb85cb5980ec04/maskit/iris.py) and should therefore be removed so that nobody could accidentally use this implementation.

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/55?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) Report
> Merging [#55](https://codecov.io/gh/cirKITers/masKIT/pull/55?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) (6eaf9ef) into [main](https://codecov.io/gh/cirKITers/masKIT/commit/342378fdfc8affec5e30230628eb85cb5980ec04?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) (342378f) will **increase** coverage by `1.61%`.
> The diff coverage is `9.67%`.

[![Impacted file tree graph](https://codecov.io/gh/cirKITers/masKIT/pull/55/graphs/tree.svg?width=650&height=150&src=pr&token=8Or2uHjmFE&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers)](https://codecov.io/gh/cirKITers/masKIT/pull/55?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers)

```diff
@@            Coverage Diff             @@
##             main      #55      +/-   ##
==========================================
+ Coverage   81.21%   82.83%   +1.61%     
==========================================
  Files          24       23       -1     
  Lines        1363     1340      -23     
  Branches      155      157       +2     
==========================================
+ Hits         1107     1110       +3     
+ Misses        247      221      -26     
  Partials        9        9              
```


| [Impacted Files](https://codecov.io/gh/cirKITers/masKIT/pull/55?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers) | Coverage Δ | |
|---|---|---|
| [main.py](https://codecov.io/gh/cirKITers/masKIT/pull/55/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFpbi5weQ==) | `0.00% <0.00%> (ø)` | |
| [maskit/circuits.py](https://codecov.io/gh/cirKITers/masKIT/pull/55/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers#diff-bWFza2l0L2NpcmN1aXRzLnB5) | `80.43% <66.66%> (+3.16%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/55?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/cirKITers/masKIT/pull/55?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers). Last update [342378f...6eaf9ef](https://codecov.io/gh/cirKITers/masKIT/pull/55?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=cirKITers).

pep8speaks: Hello @eileen-kuehn! Thanks for updating this PR. We checked the lines you've touched for [PEP 8](https://www.python.org/dev/peps/pep-0008) issues, and found:





* In the file [`maskit/circuits.py`](https://github.com/cirKITers/masKIT/blob/05012d871fe29b8d84ab74df14e85fec97879410/maskit/circuits.py):

> [Line 34:41](https://github.com/cirKITers/masKIT/blob/05012d871fe29b8d84ab74df14e85fec97879410/maskit/circuits.py#L34): [E231](https://duckduckgo.com/?q=pep8%20E231) missing whitespace after ','
















eileen-kuehn: My dear @maxfischer2781 you might re-consider approving this PR now 🧐 

Patch:
diff --git a/main.py b/main.py
index 1ae4861..c5d91d4 100644
--- a/main.py
+++ b/main.py
@@ -1,4 +1,6 @@
-from typing import List, Optional
+from maskit.datasets.utils import DataSet
+from typing import Dict, List, Optional, Tuple, Type
+from dataclasses import dataclass, field
 
 import random
 import pennylane as qml
@@ -6,18 +8,20 @@ from pennylane import numpy as np
 
 from maskit._masks import (
     DropoutMask,
+    Mask,
     PerturbationAxis as Axis,
     PerturbationMode as Mode,
 )
 from maskit._masked_circuits import MaskedCircuit
 from maskit.datasets import load_data
-from maskit.utils import cross_entropy, check_params
-from maskit.circuits import variational_circuit, basis_circuit
+from maskit.utils import check_params
+from maskit.circuits import cost, cost_basis, variational_circuit, basis_circuit
 from maskit.log_results import log_results
 from maskit.optimizers import ExtendedOptimizers
 from maskit.ensembles import (
     AdaptiveEnsemble,
     Ensemble,
+    EnsembleResult,
 )
 
 
@@ -28,37 +32,13 @@ def get_device(sim_local: bool, wires: int, shots: Optional[int] = None):
     return dev
 
 
-def cost(
-    circuit,
-    params,
-    rotations: List,
-    masked_circuit: MaskedCircuit,
-):
-    return 1 - circuit(params, rotations, masked_circuit)[0]
-
-
-def cost_basis(
-    circuit,
-    params,
-    data,
-    target,
-    rotations: List,
-    masked_circuit: MaskedCircuit,
-    wires: int,
-    wires_to_measure: List[int],
-):
-    prediction = circuit(
-        params, data, rotations, masked_circuit, wires, wires_to_measure
-    )
-    return cross_entropy(predictions=prediction, targets=target)
-
-
 def init_parameters(
     layers: int,
     current_layers: int,
     wires: int,
     default_value: Optional[float],
     dynamic_parameters: bool = True,
+    mask_type: Type[Mask] = DropoutMask,
 ) -> MaskedCircuit:
     params_uniform = np.random.uniform(
         low=-np.pi, high=np.pi, size=(current_layers, wires)
@@ -72,76 +52,103 @@ def init_parameters(
         default_value=default_value,
         entangling_mask=DropoutMask(shape=(layers, wires - 1)),
         dynamic_parameters=dynamic_parameters,
+        mask_type=mask_type,
     )
-    mc.mask(Axis.LAYERS)[current_layers:] = True
+    mc.mask(Axis.LAYERS, mask_type=mask_type)[current_layers:] = True
     return mc
 
 
+@dataclass
+class LoggingData:
+    interval: int = 1
+    costs: Dict = field(default_factory=dict)
+    branch_selection: Dict = field(default_factory=dict)
+    branch_cost: Dict = field(default_factory=lambda: {"netto": {}, "brutto": {}})
+    branch_cost_step: Dict = field(default_factory=lambda: {"netto": {}, "brutto": {}})
+    active_count: Dict = field(default_factory=dict)
+    cost_values: List = field(default_factory=list)
+    active_count_values: List = field(default_factory=list)
+
+    def log_result(self, result: EnsembleResult, step):
+        self.cost_values.append(result.cost)
+        self.active_count_values.append(result.active)
+        if result.ensemble:
+            self.branch_selection[step] = result.branch_name
+            self.branch_cost["brutto"][step] = result.brutto
+            self.branch_cost["netto"][step] = result.netto
+            self.branch_cost_step["brutto"][step] = result.brutto_steps
+            self.branch_cost_step["netto"][step] = result.netto_steps
+        if step % self.interval == 0:  # perform logging
+            self._persist_average_cost(step)
+            self._persist_average_dropout(step)
+
+    def _persist_average_cost(self, step):
+        self.costs[step] = np.average(self.cost_values)
+        self.cost_values.clear()
+
+    def _persist_average_dropout(self, step):
+        self.active_count[step] = np.average(self.active_count_values)
+        self.active_count_values.clear()
+
+
 def train(
-    train_params,
-    train_data: Optional[np.ndarray] = None,
-    train_target: Optional[np.ndarray] = None,
+    wires: int = 1,
+    wires_to_measure: Tuple[int, ...] = (0,),
+    interpret: Tuple[int, ...] = (0,),
+    layers: int = 1,
+    starting_layers: Optional[int] = None,
+    steps: int = 1000,
+    sim_local: bool = True,
+    shots: Optional[int] = None,
+    step_size: Optional[float] = None,
+    optimizer: ExtendedOptimizers = ExtendedOptimizers.GD,
+    default_value: Optional[float] = None,
+    log_interval: int = 5,
+    ensemble_type: Type[Ensemble] = Ensemble,
+    ensemble_kwargs: Optional[Dict] = None,
+    data: Optional[np.ndarray] = None,
+    target: Optional[np.ndarray] = None,
+    mask_type: Type[Mask] = DropoutMask,
 ):
-    logging_costs = {}
-    logging_branch_selection = {}
-    logging_branch_cost = {"netto": {}, "brutto": {}}
-    logging_branch_cost_step = {"netto": {}, "brutto": {}}
-    logging_dropout_count = {}
-    logging_cost_values = []
-    logging_dropout_count_values = []
-
-    np.random.seed(train_params["seed"])
-    random.seed(train_params["seed"])
+    log_data = LoggingData(interval=log_interval)
 
     # set up circuit, training, dataset
-    wires = train_params["wires"]
-    wires_to_measure = train_params["wires_to_measure"]
-    layers = train_params["layers"]
-    steps = train_params.get("steps", 1000)
-    dev = get_device(
-        train_params.get("sim_local", True),
-        wires=wires,
-        shots=train_params.get("shots", None),
-    )
-    step_size = train_params.get("step_size", None)
-    if step_size:
-        opt = train_params["optimizer"].value(step_size)
-    else:
-        opt = train_params["optimizer"].value()
-    dropout_ensemble = train_params.get("ensemble_type", Ensemble)(
-        **train_params.get("ensemble_kwargs", {"dropout": None})
+    dev = get_device(sim_local=sim_local, wires=wires, shots=shots)
+    opt = optimizer.value(step_size) if step_size else optimizer.value()
+    dropout_ensemble = ensemble_type(
+        **(ensemble_kwargs if ensemble_kwargs is not None else {"dropout": None})
     )
 
     rotation_choices = [0, 1, 2]
     rotations = [np.random.choice(rotation_choices) for _ in range(layers * wires)]
 
-    current_layers = train_params.get("starting_layers", layers)
-    default_value = train_params.get("default_value", None)
+    current_layers = starting_layers if starting_layers else layers
 
-    if train_params["dataset"] == "simple":
-        circuit = qml.QNode(variational_circuit, dev)
+    if data is not None and target is not None:
+        circuit = qml.QNode(basis_circuit, dev)
 
         def cost_fn(params, masked_circuit=None):
-            return cost(
+            return cost_basis(
                 circuit,
                 params,
+                current_data,
+                current_target,
                 rotations,
                 masked_circuit,
+                wires,
+                wires_to_measure,
+                interpret,
             )
 
-    elif train_params["dataset"] in ["iris", "mnist", "circles"]:
-        circuit = qml.QNode(basis_circuit, dev)
+    else:  # learning without encoding any data
+        circuit = qml.QNode(variational_circuit, dev)
 
         def cost_fn(params, masked_circuit=None):
-            return cost_basis(
+            return cost(
                 circuit,
                 params,
-                data,
-                target,
                 rotations,
                 masked_circuit,
-                wires,
-                wires_to_measure,
             )
 
     # set up parameters
@@ -150,38 +157,22 @@ def train(
         current_layers,
         wires,
         default_value,
-        dynamic_parameters=False
-        if train_params["optimizer"] == ExtendedOptimizers.ADAM
-        else True,
+        dynamic_parameters=False if optimizer == ExtendedOptimizers.ADAM else True,
+        mask_type=mask_type,
     )
 
     # -----------------------------
     # ======= TRAINING LOOP =======
     # -----------------------------
     for step in range(steps):
-        if train_data is not None and train_target is not None:
-            data = train_data[step % len(train_data)]
-            target = train_target[step % len(train_target)]
+        if data is not None and target is not None:
+            current_data = data[step % len(data)]
+            current_target = target[step % len(target)]
 
         # TODO: add logging for adaptive ensembles
         result = dropout_ensemble.step(masked_circuit, opt, cost_fn, ensemble_steps=1)
         masked_circuit = result.branch
-        if result.ensemble:
-            logging_branch_selection[step] = result.branch_name
-            logging_branch_cost["brutto"][step] = result.brutto
-            logging_branch_cost["netto"][step] = result.netto
-            logging_branch_cost_step["brutto"][step] = result.brutto_steps
-            logging_branch_cost_step["netto"][step] = result.netto_steps
-        logging_cost_values.append(result.cost)
-        logging_dropout_count_values.append(
-            np.sum(masked_circuit.full_mask(DropoutMask))
-        )
-        if step % train_params["log_interval"] == 0:
-            # perform logging
-            logging_costs[step] = np.average(logging_cost_values)
-            logging_dropout_count[step] = np.average(logging_dropout_count_values)
-            logging_cost_values.clear()
-            logging_dropout_count_values.clear()
+        log_data.log_result(result, step)
 
         if __debug__:
             print(
@@ -191,88 +182,87 @@ def train(
 
     if __debug__:
         print(masked_circuit.parameters)
-        print(masked_circuit.full_mask(DropoutMask))
+        print(masked_circuit.full_mask(mask_type))
 
     return {
-        "costs": logging_costs,
+        "costs": log_data.costs,
         "final_cost": result.cost,
-        "dropouts": logging_dropout_count,
-        "branch_selections": logging_branch_selection,
-        "branch_costs": logging_branch_cost,
-        "branch_step_costs": logging_branch_cost_step,
+        "maximum_active": masked_circuit.parameters.size,
+        "active": log_data.active_count,
+        "branch_selections": log_data.branch_selection,
+        "branch_costs": log_data.branch_cost,
+        "branch_step_costs": log_data.branch_cost_step,
         "final_layers": current_layers,
         "params": masked_circuit.parameters.unwrap(),
-        "dropout_mask": masked_circuit.full_mask(DropoutMask).unwrap(),
-        "__wire_mask": masked_circuit.mask(Axis.WIRES).mask,
-        "__layer_mask": masked_circuit.mask(Axis.LAYERS).mask,
-        "__parameter_mask": masked_circuit.mask(Axis.PARAMETERS).mask,
+        "mask": masked_circuit.full_mask(mask_type).unwrap(),
         "__rotations": rotations,
+        "__masked_circuit": masked_circuit,
     }
 
 
 def test(
-    train_params,
-    params,
-    wire_mask,
-    layer_mask,
-    parameter_mask,
-    layers: int,
+    masked_circuit: MaskedCircuit,
     rotations: List,
-    test_data: Optional[np.ndarray] = None,
-    test_target: Optional[np.ndarray] = None,
+    wires_to_measure: Tuple[int, ...],
+    interpret: Tuple[int, ...],
+    shots: Optional[int],
+    sim_local: bool,
+    data: Optional[np.ndarray] = None,
+    target: Optional[np.ndarray] = None,
 ):
-    if test_data is None or test_target is None:
+    if data is None or target is None:
         pass
-    elif test_data is not None and test_target is not None:
-        wires = train_params["wires"]
-        wires_to_measure = train_params["wires_to_measure"]
-        dev = get_device(
-            train_params["sim_local"],
-            wires=wires,
-            shots=train_params.get("shots", None),
-        )
+    elif data is not None and target is not None:
+        dev = get_device(sim_local, wires=masked_circuit.wires, shots=shots)
         circuit = qml.QNode(basis_circuit, dev)
         correct = 0
-        N = len(test_data)
+        N = len(data)
         costs = []
-        masked_circuit = MaskedCircuit.full_circuit(
-            parameters=params,
-            layers=layers,
-            wires=wires,
-            wire_mask=wire_mask,
-            layer_mask=layer_mask,
-            parameter_mask=parameter_mask,
-        )
-        for _step, (data, target) in enumerate(zip(test_data, test_target)):
+        for current_data, current_target in zip(data, target):
             output = circuit(
                 masked_circuit.differentiable_parameters,
-                data,
+                current_data,
                 rotations,
                 masked_circuit,
-                wires,
+                masked_circuit.wires,
                 wires_to_measure,
             )
-            c = cost_basis(
-                circuit,
-                masked_circuit.differentiable_parameters,
-                data,
-                target,
-                rotations,
-                masked_circuit,
-                wires,
-                wires_to_measure,
+            costs.append(
+                cost_basis(
+                    circuit,
+                    masked_circuit.differentiable_parameters,
+                    current_data,
+                    current_target,
+                    rotations,
+                    masked_circuit,
+                    masked_circuit.wires,
+                    wires_to_measure,
+                    interpret,
+                )
             )
-            costs.append(c)
-            same = np.argmax(target) == np.argmax(output)
+            selected_output = output[
+                interpret,
+            ]
+            same = np.argmax(current_target) == np.argmax(selected_output)
             if same:
                 correct += 1
             if __debug__:
-                print(f"Label: {target} Output: {output} Correct: {same}")
+                print(
+                    f"Label: {current_target} Output: {selected_output} "
+                    f"({output}) Correct: {same}"
+                )
+
+        accuracy = correct / N
+        loss = np.average(costs)
         if __debug__:
             print(
-                f"Accuracy = {correct} / {N} = {correct/N} \n",
-                f"Avg Cost: {np.average(costs)}",
+                f"Accuracy = {correct} / {N} = {accuracy} \n",
+                f"Avg Cost: {loss}",
             )
+        return {
+            "accuracy": accuracy,
+            "loss": loss,
+        }
 
 
 if __name__ == "__main__":
@@ -295,6 +285,7 @@ if __name__ == "__main__":
                             "amount": 1,
                             "mode": Mode.SET,
                             "axis": Axis.PARAMETERS,
+                            "mask": DropoutMask,
                         },
                     },
                 ],
@@ -305,6 +296,7 @@ if __name__ == "__main__":
                             "amount": 0.05,
                             "mode": Mode.RESET,
                             "axis": Axis.PARAMETERS,
+                            "mask": DropoutMask,
                         }
                     },
                 ],
@@ -318,10 +310,28 @@ if __name__ == "__main__":
         "logging": True,
         "seed": 1337,
         "log_interval": 5,
+        "mask_type": DropoutMask,
     }
     check_params(train_params)
-    if train_params.get("logging", True):
-        train = log_results(train)
+    logging_activated = train_params.pop("logging", True)
+    if logging_activated:
+        train = log_results(
+            train,
+            exclude=(
+                "data",
+                "target",
+            ),
+        )
+        test = log_results(
+            test,
+            exclude=(
+                "data",
+                "target",
+            ),
+        )
+    seed = train_params.pop("seed", 1337)
+    np.random.seed(seed)
+    random.seed(seed)
 
     data_params = {
         "wires": train_params["wires"],
@@ -329,20 +339,22 @@ if __name__ == "__main__":
         "train_size": 120,
         "test_size": 100,
         "shuffle": True,
+        "target_length": len(train_params.get("interpret", (0,))),
     }
-    data = load_data(train_params.get("dataset", "simple"), **data_params)
-    result = train(
-        train_params, train_data=data.train_data, train_target=data.train_target
-    )
-    if train_params["testing"]:
+    try:
+        data = load_data(train_params.pop("dataset"), **data_params)
+    except ValueError:
+        data = DataSet(None, None, None, None)
+    testing = train_params.pop("testing", False)
+    result = train(**train_params, data=data.train_data, target=data.train_target)
+    if testing:
         test(
-            train_params,
-            result["params"],
-            result["__wire_mask"],
-            result["__layer_mask"],
-            result["__parameter_mask"],
-            result["final_layers"],
-            result["__rotations"],
-            test_data=data.test_data,
-            test_target=data.test_target,
+            masked_circuit=result["__masked_circuit"],
+            rotations=result["__rotations"],
+            wires_to_measure=train_params.get("wires_to_measure", (0,)),
+            interpret=train_params.get("interpret", (0,)),
+            shots=train_params.get("shots", None),
+            sim_local=train_params.get("sim_local", True),
+            data=data.test_data,
+            target=data.test_target,
         )
diff --git a/maskit/_masked_circuits.py b/maskit/_masked_circuits.py
index b48d942..521059b 100644
--- a/maskit/_masked_circuits.py
+++ b/maskit/_masked_circuits.py
@@ -32,7 +32,14 @@ class MaskedCircuit(object):
         which masks are supported, defaults to `None`
     """
 
-    __slots__ = ("parameters", "default_value", "_dynamic_parameters", "masks")
+    __slots__ = (
+        "parameters",
+        "default_value",
+        "_dynamic_parameters",
+        "masks",
+        "wires",
+        "layers",
+    )
 
     def __init__(
         self,
@@ -49,6 +56,8 @@ class MaskedCircuit(object):
         assert (
             wires == parameters.shape[1]
         ), "Second dimension of parameters shape must be equal to number of wires"
+        self.wires = wires
+        self.layers = layers
         self.parameters = parameters
         self.default_value = default_value
         self._dynamic_parameters = dynamic_parameters
@@ -92,7 +101,20 @@ class MaskedCircuit(object):
         :param axis: The axis the mask is performed on
         :param mask_type: The type of `Mask` to return
         """
-        return self.masks[axis][mask_type]
+        try:
+            return self.masks[axis][mask_type]
+        except KeyError as err:
+            mask_types = list(self.masks[axis].keys())
+            if len(mask_types) > 0:
+                return mask_type(shape=self.masks[axis][mask_types[0]].shape)
+            if axis == Axis.WIRES:
+                return mask_type(shape=(self.wires,))
+            elif axis == Axis.LAYERS:
+                return mask_type(shape=(self.layers,))
+            elif axis == Axis.PARAMETERS:
+                return mask_type(shape=self.parameters.shape)
+            else:
+                raise ValueError from err
 
     def _accumulated_mask(self, for_differentiable=True) -> np.ndarray:
         result = np.zeros(
@@ -257,6 +279,8 @@ class MaskedCircuit(object):
     def copy(self: Self) -> Self:
         """Returns a copy of the current MaskedCircuit."""
         clone = object.__new__(type(self))
+        clone.wires = self.wires
+        clone.layers = self.layers
         clone.parameters = self.parameters.copy()
         clone.default_value = self.default_value
         clone._dynamic_parameters = self._dynamic_parameters
@@ -377,6 +401,7 @@ class MaskedCircuit(object):
         wire_mask: Optional[Mask] = None,
         layer_mask: Optional[Mask] = None,
         parameter_mask: Optional[Mask] = None,
+        mask_type: Type[Mask] = DropoutMask,
     ):
         """Helper to create a full circuit with `DropoutMasks`."""
         initializable_masks = [axis for axis in Axis if axis != Axis.ENTANGLING]
@@ -391,7 +416,7 @@ class MaskedCircuit(object):
             parameters=parameters,
             layers=layers,
             wires=wires,
-            masks=[(axis, DropoutMask) for axis in initializable_masks],
+            masks=[(axis, mask_type) for axis in initializable_masks],
             dynamic_parameters=dynamic_parameters,
             default_value=default_value,
         )
diff --git a/maskit/circuits.py b/maskit/circuits.py
index 3df04ff..e5d7996 100644
--- a/maskit/circuits.py
+++ b/maskit/circuits.py
@@ -1,12 +1,46 @@
-from maskit._masks import PerturbationAxis as Axis, DropoutMask
-from maskit._masked_circuits import MaskedCircuit
 import pennylane as qml
 from pennylane import numpy as np
+from typing import List, Tuple
+
+from maskit._masks import PerturbationAxis as Axis, DropoutMask
+from maskit._masked_circuits import MaskedCircuit
+from maskit.utils import cross_entropy
+
+
+def cost(
+    circuit,
+    params,
+    rotations: List,
+    masked_circuit: MaskedCircuit,
+):
+    return 1 - circuit(params, rotations, masked_circuit)[0]
+
+
+def cost_basis(
+    circuit,
+    params,
+    data,
+    target,
+    rotations: List,
+    masked_circuit: MaskedCircuit,
+    wires: int,
+    wires_to_measure: Tuple[int, ...],
+    interpret: Tuple[int, ...],
+):
+    prediction = circuit(
+        params, data, rotations, masked_circuit, wires, wires_to_measure
+    )
+    return cross_entropy(
+        predictions=prediction[
+            interpret,
+        ],
+        targets=target,
+    )
 
 
 def basic_variational_circuit(params, rotations, masked_circuit: MaskedCircuit):
     full_parameters = masked_circuit.expanded_parameters(params)
-    wires = len(masked_circuit.mask(Axis.WIRES))
+    wires = masked_circuit.wires
     dropout_mask = masked_circuit.full_mask(DropoutMask)
     for wire, _is_masked in enumerate(masked_circuit.mask(Axis.WIRES)):
         qml.RY(np.pi / 4, wires=wire)
@@ -26,14 +60,14 @@ def basic_variational_circuit(params, rotations, masked_circuit: MaskedCircuit):
 
         for wire in range(0, wires - 1, 2):
             if (
-                masked_circuit.mask(Axis.ENTANGLING) is not None
+                Axis.ENTANGLING in masked_circuit.masks
                 and masked_circuit.mask(Axis.ENTANGLING)[layer, wire]
             ):
                 continue
             qml.CZ(wires=[wire, wire + 1])
         for wire in range(1, wires - 1, 2):
             if (
-                masked_circuit.mask(Axis.ENTANGLING) is not None
+                Axis.ENTANGLING in masked_circuit.masks
                 and masked_circuit.mask(Axis.ENTANGLING)[layer, wire]
             ):
                 continue
@@ -52,15 +86,6 @@ def variational_circuit(params, rotations, masked_circuit):
     )
 
 
-def iris_circuit(params, data, rotations, masked_circuit):
-    masked_circuit = masked_circuit.unwrap()
-    qml.templates.embeddings.AngleEmbedding(features=data, wires=range(4), rotation="X")
-    basic_variational_circuit(
-        params=params, rotations=rotations, masked_circuit=masked_circuit
-    )
-    return qml.probs(wires=[0, 1])
-
-
 def basis_circuit(params, data, rotations, masked_circuit, wires, wires_to_measure):
     masked_circuit = masked_circuit.unwrap()
     qml.templates.embeddings.AngleEmbedding(
diff --git a/maskit/datasets/__init__.py b/maskit/datasets/__init__.py
index 582e800..e2daa10 100644
--- a/maskit/datasets/__init__.py
+++ b/maskit/datasets/__init__.py
@@ -1,9 +1,9 @@
-from typing import Tuple
+from typing import Tuple, Optional
 
 from maskit.datasets.circles import circles
 from maskit.datasets.iris import iris
 from maskit.datasets.mnist import mnist
-from maskit.datasets.utils import DataSet
+from maskit.datasets.utils import DataSet, pad_data
 
 
 def load_data(
@@ -13,6 +13,7 @@ def load_data(
     shuffle: bool = True,
     classes: Tuple[int, ...] = (6, 9),
     wires: int = 4,
+    target_length: Optional[int] = None,
 ) -> DataSet:
     """
     Returns the data for the requested ``dataset``.
@@ -24,16 +25,30 @@ def load_data(
     :param shuffle: if the dataset should be shuffled
     :param classes: which numbers of the mnist dataset should be included
     :param wires: number of wires in the circuit
+    :param target_length: Normalised length of target arrays
 
     :raises ValueError: Raised if a not supported dataset is requested
     """
-    if dataset == "simple":
-        return DataSet(None, None, None, None)
-    elif dataset == "iris":
-        return iris(train_size, test_size, shuffle)
+    if dataset == "iris":
+        result = iris(train_size, test_size, shuffle)
     elif dataset == "mnist":
-        return mnist(wires, classes, train_size, test_size, shuffle)
+        result = mnist(wires, classes, train_size, test_size, shuffle)
     elif dataset == "circles":
-        return circles(train_size, test_size, shuffle)
+        result = circles(train_size, test_size, shuffle)
     else:
         raise ValueError(f"Unknown dataset: {dataset}")
+    if target_length is None:
+        target_length = 2 ** wires
+    difference = target_length - result.train_target.shape[1]
+    assert difference >= 0, (
+        f"Target length ({target_length}) must support at least "
+        f"{result.train_target.shape[1]} classes"
+    )
+    if difference > 0:
+        # extend train and test target arrays
+        new_train_target = pad_data(result.train_target, 1, difference)
+        new_test_target = pad_data(result.test_target, 1, difference)
+        result = DataSet(
+            result.train_data, new_train_target, result.test_data, new_test_target
+        )
+    return result
diff --git a/maskit/datasets/iris.py b/maskit/datasets/iris.py
index 50c1570..686421f 100644
--- a/maskit/datasets/iris.py
+++ b/maskit/datasets/iris.py
@@ -24,7 +24,7 @@ def iris(train_size=100, test_size=50, shuffle=True) -> DataSet:
     x_train, y_train = np.split(train, [4], axis=1)
     x_test, y_test = np.split(test, [4], axis=1)
 
-    y_train = one_hot(y_train, 4)
-    y_test = one_hot(y_test, 4)
+    y_train = one_hot(y_train, 3)
+    y_test = one_hot(y_test, 3)
 
     return DataSet(x_train, y_train, x_test, y_test)
diff --git a/maskit/datasets/mnist.py b/maskit/datasets/mnist.py
index 7c0b901..3defacd 100644
--- a/maskit/datasets/mnist.py
+++ b/maskit/datasets/mnist.py
@@ -1,7 +1,6 @@
 from typing import List
 import tensorflow as tf
 import collections
-import math
 from sklearn.decomposition import PCA
 from pennylane import numpy as np
 from sklearn.preprocessing import minmax_scale
@@ -40,11 +39,7 @@ def convert_to_binary(x: np.ndarray) -> np.ndarray:
 
 def convert_label(y: int, classes: List[int]) -> List[float]:
     assert y in classes
-    # Measuring n qubits gets 2^n results to compare against this vector
-    num_classes = nearest_power_of_two(len(classes))
-    a = [0.0 for i in range(num_classes)]
-    a[classes.index(y)] = 1.0
-    return a
+    return [1.0 if the_class == y else 0.0 for the_class in classes]
 
 
 def apply_PCA(wires: int, x_train: np.ndarray):
@@ -53,10 +48,6 @@ def apply_PCA(wires: int, x_train: np.ndarray):
     return pca
 
 
-def nearest_power_of_two(x: int) -> int:
-    return 2 ** (math.ceil(math.log(x, 2)))
-
-
 def mnist(
     wires=4, classes=(6, 9), train_size=100, test_size=50, shuffle=True
 ) -> DataSet:
diff --git a/maskit/datasets/utils.py b/maskit/datasets/utils.py
index e5e74db..e6660ac 100644
--- a/maskit/datasets/utils.py
+++ b/maskit/datasets/utils.py
@@ -6,6 +6,20 @@ def one_hot(a: np.ndarray, num_classes: int) -> np.ndarray:
     return np.squeeze(np.eye(num_classes)[a.astype(int).reshape(-1)])
 
 
+def pad_data(data: np.ndarray, axis: int, padding: int) -> np.ndarray:
+    """
+    Function pads 0 data to the end of the given axis.
+
+    :param data: The data to pad 0 data to
+    :param axis: Axis to pad on
+    :param padding: How many elements to pad
+    """
+    size = data.shape[0]
+    if size == 0:
+        return data.reshape((size, data.shape[axis] + padding))
+    return np.append(data, [[0] * padding for _ in range(size)], axis)
+
+
 class DataSet(NamedTuple):
     train_data: Optional[np.ndarray]
     train_target: Optional[np.ndarray]
diff --git a/maskit/ensembles.py b/maskit/ensembles.py
index 2262237..7d0b30a 100644
--- a/maskit/ensembles.py
+++ b/maskit/ensembles.py
@@ -10,6 +10,8 @@ class EnsembleResult(NamedTuple):
     branch: MaskedCircuit
     #: name of branch as configured
     branch_name: str
+    #: number of active gates of selected branch
+    active: int
     #: cost of selected branch
     cost: float
     #: gradient of selected branch
@@ -73,6 +75,7 @@ class Ensemble(object):
             return EnsembleResult(
                 branch=masked_circuit,
                 branch_name="center",
+                active=basic_active_gates,
                 cost=objective_fn(
                     masked_circuit.differentiable_parameters,
                     masked_circuit=masked_circuit,
@@ -109,6 +112,7 @@ class Ensemble(object):
         return EnsembleResult(
             branch=selected_branch,
             branch_name=branch_name,
+            active=selected_branch.active(),
             cost=branch_costs[minimum_index].unwrap(),
             gradient=branch_gradients[minimum_index],
             brutto_steps=1 + len(branches) * ensemble_steps,
diff --git a/maskit/iris.py b/maskit/iris.py
deleted file mode 100644
index aee1938..0000000
--- a/maskit/iris.py
+++ /dev/null
@@ -1,32 +0,0 @@
-from pennylane import numpy as np
-from sklearn import datasets
-
-np.random.seed(1337)
-
-
-def one_hot(a, num_classes):
-    return np.squeeze(np.eye(num_classes)[a.astype(int).reshape(-1)])
-
-
-def load_iris(train_size=120):
-    train_size = min(train_size, 150)
-    data, target = datasets.load_iris(return_X_y=True)
-    target = target.reshape((150, 1))
-    dataset = np.concatenate((data, target), axis=1)
-
-    np.random.shuffle(dataset)
-
-    train, test = dataset[:train_size, :], dataset[train_size:, :]
-
-    x_train, y_train = np.split(train, [4], axis=1)
-    x_test, y_test = np.split(test, [4], axis=1)
-
-    y_train = one_hot(y_train, 4)
-    y_test = one_hot(y_test, 4)
-
-    return x_train, y_train, x_test, y_test
-
-
-if __name__ == "__main__":
-    x_train, y_train, x_test, y_test = load_iris()
-    print(y_train)
diff --git a/maskit/log_results.py b/maskit/log_results.py
index 17e8567..ba1733d 100644
--- a/maskit/log_results.py
+++ b/maskit/log_results.py
@@ -1,4 +1,4 @@
-from typing import List, Dict, TypeVar, Union, Any, Callable
+from typing import List, Dict, Tuple, TypeVar, Union, Any, Callable
 import os
 import time
 import json
@@ -23,7 +23,7 @@ def serialize(o):
     # raise TypeError(f"Cannot serialize {o} to JSON")
 
 
-def log_results(executor: CJ) -> CJ:
+def log_results(executor: CJ, exclude: Tuple[str, ...]) -> CJ:
     @wraps(executor)
     def wrapper(*args, **kwargs):
         wallclock, cpuclock = time.perf_counter, time.process_time
@@ -37,7 +37,11 @@ def log_results(executor: CJ) -> CJ:
                     "call": [
                         f"{executor.__module__}.{executor.__qualname__}",
                         args,
-                        kwargs,
+                        {
+                            key: value
+                            for key, value in kwargs.items()
+                            if key not in exclude
+                        },
                     ],
                     "result": result
                     if type(result) is not dict
diff --git a/maskit/utils.py b/maskit/utils.py
index 806b715..310cd2e 100644
--- a/maskit/utils.py
+++ b/maskit/utils.py
@@ -6,6 +6,10 @@ from maskit.optimizers import ExtendedOptimizers
 def check_params(train_params):
     assert train_params["dataset"] in ["simple", "iris", "mnist", "circles"]
     assert isinstance(train_params["optimizer"], ExtendedOptimizers)
+    if "interpret" not in train_params:
+        train_params["interpret"] = tuple(
+            range(2 ** len(train_params.get("wires_to_measure", [0])))
+        )
 
 
 def cross_entropy(


Test patch:
diff --git a/tests/test_load_data.py b/tests/test_load_data.py
index 9f8e522..c534e68 100644
--- a/tests/test_load_data.py
+++ b/tests/test_load_data.py
@@ -3,10 +3,12 @@ from maskit.datasets import load_data
 
 class TestLoadData:
     def test_iris_basic(self):
+        num_classes = 3
         data = load_data(
             "iris",
             train_size=120,
             test_size=30,
+            target_length=num_classes,  # num of classes
         )
 
         assert len(data.train_data) == 120
@@ -15,9 +17,9 @@ class TestLoadData:
         assert len(data.test_target) == 30
 
         assert len(data.train_data[0]) == 4
-        assert len(data.train_target[0]) == 4
+        assert len(data.train_target[0]) == num_classes
         assert len(data.test_data[0]) == 4
-        assert len(data.test_target[0]) == 4
+        assert len(data.test_target[0]) == num_classes
 
     def test_iris_empty_test(self):
         data = load_data(
@@ -29,13 +31,11 @@ class TestLoadData:
         assert len(data.train_target) == 150
         assert len(data.test_data) == 0
         assert len(data.test_target) == 0
+        assert data.train_target.shape[1] == 16
+        assert data.test_target.shape[1] == 16
 
     def test_circles_basic(self):
-        data = load_data(
-            "circles",
-            train_size=150,
-            test_size=50,
-        )
+        data = load_data("circles", train_size=150, test_size=50, target_length=2)
 
         assert len(data.train_data) == 150
         assert len(data.train_target) == 150
@@ -48,8 +48,14 @@ class TestLoadData:
         assert len(data.test_target[0]) == 2
 
     def test_mnist_basic(self):
+        classes = [6, 7, 8]
         data = load_data(
-            "mnist", wires=5, classes=[6, 7, 8], train_size=150, test_size=50
+            "mnist",
+            wires=5,
+            classes=classes,
+            train_size=150,
+            test_size=50,
+            target_length=len(classes),
         )
 
         assert len(data.train_data) == 150
@@ -58,6 +64,6 @@ class TestLoadData:
         assert len(data.test_target) == 50
 
         assert len(data.train_data[0]) == 5
-        assert len(data.train_target[0]) == 4
+        assert len(data.train_target[0]) == len(classes)
         assert len(data.test_data[0]) == 5
-        assert len(data.test_target[0]) == 4
+        assert len(data.test_target[0]) == len(classes)