qadence 1.10.2__py3-none-any.whl → 1.11.0__py3-none-any.whl

qadence/ml_tools/config.py

@@ -20,6 +20,7 @@ from qadence.types import (
     ReuploadScaling,
     Strategy,
 )
+from torch import dtype
 
 logger = getLogger(__file__)
 
@@ -116,10 +117,9 @@ class TrainConfig:
     """The log folder for saving checkpoints and tensorboard logs.
 
     This stores the path where all logs and checkpoints are being saved
-    for this training session. `log_folder` takes precedence over `root_folder` and
-    `create_subfolder_per_run` arguments. If the user specifies a log_folder,
-    all checkpoints will be saved in this folder and `root_folder` argument
-    will not be used.
+    for this training session. `log_folder` takes precedence over `root_folder`,
+    but it is ignored if `create_subfolders_per_run=True` (in which case, subfolders
+    will be spawned in the root folder).
     """
 
     checkpoint_best_only: bool = False
@@ -195,7 +195,7 @@ class TrainConfig:
     plots that are logged or saved at specified intervals.
     """
 
-    _subfolders: list = field(default_factory=list)
+    _subfolders: list[str] = field(default_factory=list)
     """List of subfolders used for logging different runs using the same config inside the.
 
     root folder.
@@ -203,6 +203,67 @@ class TrainConfig:
     Each subfolder is of structure `<id>_<timestamp>_<PID>`.
     """
 
+    nprocs: int = 1
+    """
+    The number of processes to use for training when spawning subprocesses.
+
+    For effective parallel processing, set this to a value greater than 1.
+    - In case of Multi-GPU or Multi-Node-Multi-GPU setups, nprocs should be equal to
+    the total number of GPUs across all nodes (world size), or total number of GPU to be used.
+
+    If nprocs > 1, multiple processes will be spawned for training. The training framework will launch
+    additional processes (e.g., for distributed or parallel training).
+    - For CPU setup, this will launch a true parallel processes
+    - For GPU setup, this will launch a distributed training routine.
+    This uses the DistributedDataParallel framework from PyTorch.
+    """
+
+    compute_setup: str = "cpu"
+    """
+    Compute device setup; options are "auto", "gpu", or "cpu".
+
+    - "auto": Automatically uses GPU if available; otherwise, falls back to CPU.
+    - "gpu": Forces GPU usage, raising an error if no CUDA device is available.
+    - "cpu": Forces the use of CPU regardless of GPU availability.
+    """
+
+    backend: str = "gloo"
+    """
+    Backend used for distributed training communication.
+
+    The default is "gloo". Other options may include "nccl" - which is optimized for GPU-based training or "mpi",
+    depending on your system and requirements.
+    It should be one of the backends supported by `torch.distributed`. For further details, please look at
+    [torch backends](https://pytorch.org/docs/stable/distributed.html#torch.distributed.Backend)
+    """
+
+    log_setup: str = "cpu"
+    """
+    Logging device setup; options are "auto" or "cpu".
+
+    - "auto": Uses the same device for logging as for computation.
+    - "cpu": Forces logging to occur on the CPU. This can be useful to avoid potential conflicts with GPU processes.
+    """
+
+    dtype: dtype | None = None
+    """
+    Data type (precision) for computations.
+
+    Both model parameters, and dataset will be of the provided precision.
+
+    If not specified or None, the default torch precision (usually torch.float32) is used.
+    If provided dtype is torch.complex128, model parameters will be torch.complex128, and data parameters will be torch.float64
+    """
+
+    all_reduce_metrics: bool = False
+    """
+    Whether to aggregate metrics (e.g., loss, accuracy) across processes.
+
+    When True, metrics from different training processes are averaged to provide a consolidated metrics.
+    Note: Since aggregation requires synchronization/all_reduce operation, this can increase the
+     computation time significantly.
+    """
+
 
 @dataclass
 class FeatureMapConfig:

qadence/ml_tools/constructors.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from typing import Callable
 import numpy as np
 from sympy import Basic
 
@@ -24,7 +25,6 @@ from qadence.constructors.iia import iia
 from qadence.measurements import Measurements
 from qadence.noise import NoiseHandler
 from qadence.operations import CNOT, RX, RY, I, N, Z
-from qadence.parameters import Parameter
 from qadence.register import Register
 from qadence.types import (
     AnsatzType,
@@ -33,10 +33,10 @@ from qadence.types import (
     InputDiffMode,
     Interaction,
     MultivariateStrategy,
-    ObservableTransform,
     ReuploadScaling,
     Strategy,
     TParameter,
+    TArray,
 )
 
 from .config import AnsatzConfig, FeatureMapConfig
@@ -706,35 +706,6 @@ def _interleave_ansatz_in_fm(
     return chain(*full_fm)
 
 
-def load_observable_transformations(config: ObservableConfig) -> tuple[Parameter, Parameter]:
-    """
-    Get the observable shifting and scaling factors.
-
-    Args:
-        config (ObservableConfig): Observable configuration.
-
-    Returns:
-        tuple[Parameter, Parameter]: The observable shifting and scaling factors.
-    """
-    shift = config.shift
-    scale = config.scale
-    if config.trainable_transform is not None:
-        shift = Parameter(name=shift, trainable=config.trainable_transform)
-        scale = Parameter(name=scale, trainable=config.trainable_transform)
-    else:
-        shift = Parameter(shift)
-        scale = Parameter(scale)
-    return scale, shift
-
-
-ObservableTransformMap = {
-    ObservableTransform.RANGE: lambda detuning, scale, shift: (
-        (shift, shift - scale) if detuning is N else (0.5 * (shift - scale), 0.5 * (scale + shift))
-    ),
-    ObservableTransform.SCALE: lambda _, scale, shift: (scale, shift),
-}
-
-
 def _global_identity(register: int | Register) -> KronBlock:
     """Create a global identity block."""
     return kron(
@@ -742,7 +713,7 @@ def _global_identity(register: int | Register) -> KronBlock:
     )
 
 
-def observable_from_config(
+def create_observable(
     register: int | Register,
     config: ObservableConfig,
 ) -> AbstractBlock:
@@ -756,35 +727,11 @@ def observable_from_config(
     Returns:
         AbstractBlock: The observable block.
     """
-    scale, shift = load_observable_transformations(config)
-    return create_observable(register, config.detuning, scale, shift, config.transformation_type)
-
-
-def create_observable(
-    register: int | Register,
-    detuning: TDetuning = Z,
-    scale: TParameter | None = None,
-    shift: TParameter | None = None,
-    transformation_type: ObservableTransform = ObservableTransform.NONE,  # type: ignore[assignment]
-) -> AbstractBlock:
-    """
-    Create an observable block.
-
-    Args:
-        register (int | Register): Number of qubits or a register object.
-        detuning: The type of detuning.
-        scale: A parameter for the scale.
-        shift: A parameter for the shift.
-
-    Returns:
-        AbstractBlock: The observable block.
-    """
-    if transformation_type == ObservableTransform.RANGE:
-        scale, shift = ObservableTransformMap[transformation_type](detuning, scale, shift)  # type: ignore[index]
-    shifting_term: AbstractBlock = shift * _global_identity(register)  # type: ignore[operator]
-    detuning_hamiltonian: AbstractBlock = scale * hamiltonian_factory(  # type: ignore[operator]
+    shifting_term: AbstractBlock = config.shift * _global_identity(register)  # type: ignore[operator]
+    detuning_hamiltonian: AbstractBlock = config.scale * hamiltonian_factory(  # type: ignore[operator]
         register=register,
-        detuning=detuning,
+        interaction=config.interaction,
+        detuning=config.detuning,
     )
     return add(shifting_term, detuning_hamiltonian)
 
@@ -844,9 +791,9 @@ def build_qnn_from_configs(
     circ = QuantumCircuit(register, *blocks)
 
     observable: AbstractBlock | list[AbstractBlock] = (
-        [observable_from_config(register=register, config=cfg) for cfg in observable_config]
+        [create_observable(register=register, config=cfg) for cfg in observable_config]
         if isinstance(observable_config, list)
-        else observable_from_config(register=register, config=observable_config)
+        else create_observable(register=register, config=observable_config)
     )
 
     ufa = QNN(

qadence/ml_tools/data.py

@@ -34,6 +34,10 @@ class OptimizeResult:
     """Metrics that can be saved during training."""
     extra: dict = field(default_factory=lambda: dict())
     """Extra dict for saving anything else to be used in callbacks."""
+    rank: int = 0
+    """Rank of the process for which this result was generated."""
+    device: str | None = "cpu"
+    """Device on which this result for calculated."""
 
 
 @dataclass

qadence/ml_tools/information/__init__.py

@@ -0,0 +1,3 @@
+from __future__ import annotations
+
+from .information_content import InformationContent

qadence/ml_tools/information/information_content.py

@@ -0,0 +1,339 @@
+from __future__ import annotations
+
+import functools
+from logging import getLogger
+from math import log, sqrt
+from statistics import NormalDist
+from typing import Any, Callable
+
+import torch
+from torch import nn
+from torch.func import functional_call  # type: ignore
+
+logger = getLogger("ml_tools")
+
+
+class InformationContent:
+    def __init__(
+        self,
+        model: nn.Module,
+        loss_fn: Callable,
+        xs: Any,
+        epsilons: torch.Tensor,
+        variation_multiple: int = 20,
+    ) -> None:
+        """Information Landscape class.
+
+        This class handles the study of loss landscape from information theoretic
+        perspective and provides methods to get bounds on the norm of the
+        gradient from the Information Content of the loss landscape.
+
+        Args:
+            model: The quantum or classical model to analyze.
+            loss_fn: Loss function that takes model output and calculates loss
+            xs: Input data to evaluate the model on
+            epsilons: The thresholds to use for discretization of the finite derivatives
+            variation_multiple: The number of sets of variational parameters to generate per each
+                variational parameter. The number of variational parameters required for the
+                statistical analysis scales linearly with the amount of them present in the
+                model. This is that linear factor.
+
+        Notes:
+            This class provides flexibility in terms of what the model, the loss function,
+            and the xs are. The only requirement is that the loss_fn takes the model and xs as
+            arguments and returns the loss, and another dictionary of other metrics.
+
+            Thus, assumed structure:
+                loss_fn(model, xs) -> (loss, metrics, ...)
+
+            Example: A Classifier
+                ```python
+                model = nn.Linear(10, 1)
+
+                def loss_fn(
+                    model: nn.Module,
+                    xs: tuple[torch.Tensor, torch.Tensor]
+                ) -> tuple[torch.Tensor, dict[str, float]:
+                    criterion = nn.MSELoss()
+                    inputs, labels = xs
+                    outputs = model(inputs)
+                    loss = criterion(outputs, labels)
+                    metrics = {"loss": loss.item()}
+                    return loss, metrics
+
+                xs = (torch.randn(10, 10), torch.randn(10, 1))
+
+                info_landscape = InfoLandscape(model, loss_fn, xs)
+                ```
+                In this example, the model is a linear classifier, and the `xs` include both the
+                inputs and the target labels. The logic for calculation of the loss from this lies
+                entirely within the `loss_fn` function. This can then further be used to obtain the
+                bounds on the average norm of the gradient of the loss function.
+
+            Example: A Physics Informed Neural Network
+                ```python
+                class PhysicsInformedNN(nn.Module):
+                    // <Initialization Logic>
+
+                    def forward(self, xs: dict[str, torch.Tensor]):
+                        return {
+                            "pde_residual": pde_residual(xs["pde"]),
+                            "boundary_condition": bc_term(xs["bc"]),
+                        }
+
+                def loss_fn(
+                    model: PhysicsInformedNN,
+                    xs: dict[str, torch.Tensor]
+                ) -> tuple[torch.Tensor, dict[str, float]:
+                    pde_residual, bc_term = model(xs)
+                    loss = torch.mean(torch.sum(pde_residual**2, dim=1), dim=0)
+                        + torch.mean(torch.sum(bc_term**2, dim=1), dim=0)
+
+                    return loss, {"pde_residual": pde_residual, "bc_term": bc_term}
+
+                xs = {
+                    "pde": torch.linspace(0, 1, 10),
+                    "bc": torch.tensor([0.0]),
+                }
+
+                info_landscape = InfoLandscape(model, loss_fn, xs)
+                ```
+
+                In this example, the model is a Physics Informed Neural Network, and the `xs`
+                are the inputs to the different residual components of the model. The logic
+                for calculation of the residuals lies within the PhysicsInformedNN class, and
+                the loss function is defined to calculate the loss that is to be optimized
+                from these residuals. This can then further be used to obtain the
+                bounds on the average norm of the gradient of the loss function.
+
+            The first value that the `loss_fn` returns is the loss value that is being optimized.
+            The function is also expected to return other value(s), often the metrics that are
+            used to calculate the loss. These values are ignored for the purpose of this class.
+        """
+        self.model = model
+        self.loss_fn = loss_fn
+        self.xs = xs
+        self.epsilons = epsilons
+        self.device = next(model.parameters()).device
+
+        self.param_shapes = {}
+        self.total_params = 0
+
+        for name, param in model.named_parameters():
+            self.param_shapes[name] = param.shape
+            self.total_params += param.numel()
+        self.n_variations = variation_multiple * self.total_params
+        self.all_variations = torch.empty(
+            (self.n_variations, self.total_params), device=self.device
+        ).uniform_(0, 2 * torch.pi)
+
+    def reshape_param_variations(self) -> dict[str, torch.Tensor]:
+        """Reshape variations of the model's variational parameters.
+
+        Returns:
+            Dictionary of parameter tensors, each with shape [n_variations, *param_shape]
+        """
+        param_variations = {}
+        start_idx = 0
+
+        for name, shape in self.param_shapes.items():
+            param_size = torch.prod(torch.tensor(shape)).item()
+            param_variations[name] = self.all_variations[
+                :, start_idx : start_idx + param_size
+            ].view(self.n_variations, *shape)
+            start_idx += param_size
+
+        return param_variations
+
+    def batched_loss(self) -> torch.Tensor:
+        """Calculate loss for all parameter variations in a batched manner.
+
+        Returns: Tensor of loss values for each parameter variation
+        """
+        param_variations = self.reshape_param_variations()
+        losses = torch.zeros(self.n_variations, device=self.device)
+
+        for i in range(self.n_variations):
+            params = {name: param[i] for name, param in param_variations.items()}
+            current_model = lambda x: functional_call(self.model, params, (x,))
+            losses[i] = self.loss_fn(current_model, self.xs)[0]
+
+        return losses
+
+    def randomized_finite_der(self) -> torch.Tensor:
+        """
+        Calculate normalized finite difference of loss on doing random walk in the parameter space.
+
+        This serves as a proxy for the derivative of the loss with respect to parameters.
+
+        Returns:
+            Tensor containing normalized finite differences (approximate directional derivatives)
+            between consecutive points in the random walk. Shape: [n_variations - 1]
+        """
+        losses = self.batched_loss()
+
+        return (losses[1:] - losses[:-1]) / (
+            torch.norm(self.all_variations[1:] - self.all_variations[:-1], dim=1) + 1e-8
+        )
+
+    def discretize_derivatives(self) -> torch.Tensor:
+        """
+        Convert finite derivatives into discrete values.
+
+        Returns:
+            Tensor containing discretized derivatives with shape [n_epsilons, n_variations-2]
+            Each row contains {-1, 0, 1} values for that epsilon
+        """
+        derivatives = self.randomized_finite_der()
+
+        derivatives = derivatives.unsqueeze(0)
+        epsilons = self.epsilons.unsqueeze(1)
+
+        discretized = torch.zeros((len(epsilons), len(derivatives[0])), device=self.device)
+        discretized[derivatives > epsilons] = 1
+        discretized[derivatives < -epsilons] = -1
+
+        return discretized
+
+    def calculate_transition_probabilities_batch(self) -> torch.Tensor:
+        """
+        Calculate transition probabilities for multiple epsilon values.
+
+        Returns:
+            Tensor of shape [n_epsilons, 6] containing probabilities for each transition type
+            Columns order: [+1to0, +1to-1, 0to+1, 0to-1, -1to0, -1to+1]
+        """
+        discretized = self.discretize_derivatives()
+
+        current = discretized[:, :-1]
+        next_val = discretized[:, 1:]
+
+        transitions = torch.stack(
+            [
+                ((current == 1) & (next_val == 0)).sum(dim=1),
+                ((current == 1) & (next_val == -1)).sum(dim=1),
+                ((current == 0) & (next_val == 1)).sum(dim=1),
+                ((current == 0) & (next_val == -1)).sum(dim=1),
+                ((current == -1) & (next_val == 0)).sum(dim=1),
+                ((current == -1) & (next_val == 1)).sum(dim=1),
+            ],
+            dim=1,
+        ).float()
+
+        total_transitions = current.size(1)
+        probabilities = transitions / total_transitions
+
+        return probabilities
+
+    @functools.cached_property
+    def calculate_IC(self) -> torch.Tensor:
+        """
+        Calculate Information Content for multiple epsilon values.
+
+        Returns: Tensor of IC values for each epsilon [n_epsilons]
+        """
+        probs = self.calculate_transition_probabilities_batch()
+
+        mask = probs > 1e-4
+
+        ic_terms = torch.where(mask, -probs * torch.log(probs), torch.zeros_like(probs))
+        ic_values = ic_terms.sum(dim=1) / torch.log(torch.tensor(6.0))
+
+        return ic_values
+
+    def max_IC(self) -> tuple[float, float]:
+        """
+        Get the maximum Information Content and its corresponding epsilon.
+
+        Returns: Tuple of (maximum IC value, optimal epsilon)
+        """
+        max_ic, max_idx = torch.max(self.calculate_IC, dim=0)
+        max_epsilon = self.epsilons[max_idx]
+        return max_ic.item(), max_epsilon.item()
+
+    def sensitivity_IC(self, eta: float) -> float:
+        """
+        Find the minimum value of epsilon such that the information content is less than eta.
+
+        Args:
+            eta: Threshold value, the sensitivity IC.
+
+        Returns: The epsilon value that gives IC that is less than the sensitivity IC.
+        """
+        ic_values = self.calculate_IC
+        mask = ic_values < eta
+        epsilons = self.epsilons[mask]
+        return float(epsilons.min().item())
+
+    @staticmethod
+    @functools.lru_cache
+    def q_value(H_value: float) -> float:
+        """
+        Compute the q value.
+
+        q is the solution to the equation:
+        H(x) = 4h(x) + 2h(1/2 - 2x)
+
+        It is the value of the probability of 4 of the 6 transitions such that
+        the IC is the same as the IC of our system.
+
+        This quantity is useful in calculating the bounds on the norms of the gradients.
+
+        Args:
+            H_value (float): The information content.
+
+        Returns:
+            float: The q value
+        """
+
+        x = torch.linspace(0.001, 0.16667, 10000)
+
+        H = -4 * x * torch.log(x) / torch.log(torch.tensor(6)) - 2 * (0.5 - 2 * x) * torch.log(
+            0.5 - 2 * x
+        ) / torch.log(torch.tensor(6))
+        err = torch.abs(H - H_value)
+        idx = torch.argmin(err)
+        return float(x[idx].item())
+
+    def get_grad_norm_bounds_max_IC(self) -> tuple[float, float]:
+        """
+        Compute the bounds on the average norm of the gradient.
+
+        Returns:
+            tuple[Tensor, Tensor]: The lower and upper bounds.
+        """
+        max_IC, epsilon_m = self.max_IC()
+        lower_bound = (
+            epsilon_m
+            * sqrt(self.total_params)
+            / (NormalDist().inv_cdf(1 - 2 * self.q_value(max_IC)))
+        )
+        upper_bound = (
+            epsilon_m
+            * sqrt(self.total_params)
+            / (NormalDist().inv_cdf(0.5 * (1 + 2 * self.q_value(max_IC))))
+        )
+
+        if max_IC < log(2, 6):
+            logger.warning(
+                "Warning: The maximum IC is less than the required value. The bounds may be"
+                + " inaccurate."
+            )
+
+        return lower_bound, upper_bound
+
+    def get_grad_norm_bounds_sensitivity_IC(self, eta: float) -> float:
+        """
+        Compute the bounds on the average norm of the gradient.
+
+        Args:
+            eta (float): The sensitivity IC.
+
+        Returns:
+            Tensor: The lower bound.
+        """
+        epsilon_sensitivity = self.sensitivity_IC(eta)
+        upper_bound = (
+            epsilon_sensitivity * sqrt(self.total_params) / (NormalDist().inv_cdf(1 - 3 * eta / 2))
+        )
+        return upper_bound

qadence/ml_tools/models.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from collections import Counter
+from collections import Counter, OrderedDict
 from logging import getLogger
 from typing import Any, Callable
 
@@ -19,6 +19,7 @@ from qadence.model import QuantumModel
 from qadence.noise import NoiseHandler
 from qadence.register import Register
 from qadence.types import BackendName, DiffMode, Endianness, InputDiffMode, ParamDictType
+from qadence.utils import block_to_mathematical_expression
 
 logger = getLogger(__name__)
 
@@ -208,6 +209,8 @@ class QNN(QuantumModel):
         else:
             raise ValueError(f"Unkown forward diff mode: {self.input_diff_mode}")
 
+        self._model_configs: dict = dict()
+
     @classmethod
     def from_configs(
         cls,
@@ -255,7 +258,7 @@ class QNN(QuantumModel):
         from qadence.constructors import ObservableConfig
         from qadence.operations import Z
         from qadence.types import (
-            AnsatzType, BackendName, BasisSet, ObservableTransform, ReuploadScaling, Strategy
+            AnsatzType, BackendName, BasisSet, ReuploadScaling, Strategy
         )
 
         register = 4
@@ -263,7 +266,6 @@ class QNN(QuantumModel):
             detuning=Z,
             scale=5.0,
             shift=0.0,
-            transformation_type=ObservableTransform.SCALE,
             trainable_transform=None,
         )
         fm_config = FeatureMapConfig(
@@ -293,7 +295,7 @@ class QNN(QuantumModel):
         """
         from .constructors import build_qnn_from_configs
 
-        return build_qnn_from_configs(
+        qnn = build_qnn_from_configs(
             register=register,
             observable_config=obs_config,
             fm_config=fm_config,
@@ -305,6 +307,69 @@ class QNN(QuantumModel):
             configuration=configuration,
             input_diff_mode=input_diff_mode,
         )
+        qnn._model_configs = {
+            "register": register,
+            "observable_config": obs_config,
+            "fm_config": fm_config,
+            "ansatz_config": ansatz_config,
+        }
+        return qnn
+
+    def __str__(self) -> str | Any:
+        """Return a string representation of a QNN.
+
+        When creating a QNN from a set of configurations,
+        we print the configurations used. Otherwise, we use the default printing.
+
+        Returns:
+            str | Any: A string representation of a QNN.
+
+        Example:
+        ```python exec="on" source="material-block" result="json"
+        from qadence import QNN
+        from qadence.constructors.hamiltonians import Interaction
+        from qadence.ml_tools.config import AnsatzConfig, FeatureMapConfig
+        from qadence.ml_tools.constructors import (
+            ObservableConfig,
+        )
+        from qadence.operations import Z
+        from qadence.types import BackendName
+
+        backend = BackendName.PYQTORCH
+        fm_config = FeatureMapConfig(num_features=1)
+        ansatz_config = AnsatzConfig()
+        observable_config = ObservableConfig(detuning=Z, interaction=Interaction.ZZ, scale=2)
+
+        qnn = QNN.from_configs(
+            register=2,
+            obs_config=observable_config,
+            fm_config=fm_config,
+            ansatz_config=ansatz_config,
+            backend=backend,
+        )
+        print(qnn) # markdown-exec: hide
+        ```
+        """
+        if bool(self._model_configs):
+            configs_str = "\n".join(
+                (
+                    k + " = " + str(self._model_configs[k])
+                    for k in sorted(self._model_configs.keys())
+                    if k != "observable_config"
+                )
+            )
+            observable_str = ""
+            if self._observable:
+                observable_str = (
+                    "observable_config = [\n"
+                    + "\n".join(
+                        (block_to_mathematical_expression(obs.original) for obs in self._observable)
+                    )
+                    + "\n]"
+                )
+            return f"{type(self).__name__}(\n{configs_str}\n{observable_str}\n)"
+
+        return super().__str__()
 
     def forward(
         self,