Added Docstrings (see issue #135)

torchquantum/algorithms/hamiltonian.py

@@ -12,11 +12,52 @@
 __all__ = ["Hamiltonian"]
 
 class Hamiltonian(object):
-    def __init__(self, hamil_info) -> None:
+    """The Hamiltonian.
+
+    This class provides functionality to process and initialize a Hamiltonian.
+
+    Attributes:
+        hamil_info (dict): Information about the Hamiltonian.
+        n_wires (int): Number of wires in the Hamiltonian.
+
+    Methods:
+        __init__(hamil_info): Initializes the Hamiltonian object.
+        process_hamil_info: Process the Hamiltonian information.
+        from_file: Create a Hamiltonian from a file.
+    """
+
+    def __init__(self, hamil_info: dict) -> None:
+        """Initialize the Hamiltonian object.
+
+        Args:
+            hamil_info (dict): Information about the Hamiltonian.
+
+        Returns:
+            None
+
+        Examples:
+            >>> hamil_info = {"hamil_list": [...], "n_wires": 3}
+            >>> hamiltonian = Hamiltonian(hamil_info)
+        """
+
         self.hamil_info = self.process_hamil_info(hamil_info)
         self.n_wires = hamil_info["n_wires"]
 
-    def process_hamil_info(self, hamil_info):
+    def process_hamil_info(self, hamil_info: dict) -> dict:
+        """Process the Hamiltonian information.
+
+        Args:
+            hamil_info (dict): Information about the Hamiltonian.
+
+        Returns:
+            dict: processed Hamiltonian information
+
+        Examples:
+            >>> hamil_info = {"hamil_list": [...], "n_wires": 3}
+
+            >>> processed_info = self.process_hamil_info(hamil_info)
+        """
+
         hamil_list = hamil_info["hamil_list"]
         n_wires = hamil_info["n_wires"]
         all_info = []
@@ -34,7 +75,22 @@ def process_hamil_info(self, hamil_info):
         return hamil_info
 
     @classmethod
-    def from_file(cls, file_path):
+    def from_file(cls, file_path: str) -> 'Hamiltonian':
+        """Create a Hamiltonian object from a file.
+
+        Args:
+            file_path (str): Path to the file containing Hamiltonian information.
+
+        Returns:
+            Hamiltonian: the created Hamiltonian object
+
+        Examples:
+            >>> file_path = "hamiltonian.txt"
+            >>> hamiltonian = Hamiltonian.from_file(file_path)
+
+            >>> hamil = Hamiltonian.from_file("hamil_file.txt")
+        """
+
         hamil_info = parse_hamiltonian_file(file_path)
         return cls(hamil_info)
-
+

torchquantum/algorithms/vqe.py

@@ -14,11 +14,35 @@
 
 
 class VQE(object):
+    """The Variational Quantum Eigensolver (VQE).
+
+    Attributes:
+        hamil: A dictionary containing information about the Hamiltonian.
+        ansatz: An Ansatz object.
+        train_configs: A dictionary containing training configurations.
+        n_wires: An integer number of wires in the Hamiltonian.
+        n_epochs: An integer number of training epochs.
+        n_steps: An integer number of optimization steps per epoch.
+        optimizer_name: Name of the optimizer.
+        scheduler_name: Name of the learning rate scheduler.
+        lr: A float indicating learning rate.
+        device: Device to use for training.
+
+    Methods:
+        __init__: Initialize the VQE object.
+        get_expval: Calculate the expectation value of the Hamiltonian for a given quantum device.
+        get_loss: Calculate the loss function.
+        train: Train the VQE model.
+    """
+
     def __init__(self, hamil, ansatz, train_configs) -> None:
-        """Init function for VQE class
+        """Initialize the VQE object.
+
         Args:
-            hamil (dict): A dictionary containing the information of the hamiltonian
-         ansatz (Ansatz): An Ansatz object
+            hamil (dict): Information about the Hamiltonian
+            ansatz (Ansatz): An Ansatz object representing the variational circuit
+            train_configs (dict): Configuration parameters for training the VQE model
+
         """
         self.hamil = hamil
         self.ansatz = ansatz
@@ -34,6 +58,15 @@ def __init__(self, hamil, ansatz, train_configs) -> None:
         self.ansatz = self.ansatz.to(self.device)
 
     def get_expval(self, qdev):
+        """Calculate the expectation value of the Hamiltonian for a given quantum device.
+
+        Args:
+            qdev (torchquantum.QuantumDevice): Quantum device representing the state of the variational circuit
+
+        Returns:
+            float : expectation value of the Hamiltonian
+        """
+
         hamil_list = self.hamil.hamil_info["hamil_list"]
         expval = 0
         for hamil in hamil_list:
@@ -44,6 +77,12 @@ def get_expval(self, qdev):
         return expval
 
     def get_loss(self):
+        """Calculate the loss function.
+
+        Returns:
+            float: loss value
+        """
+
         qdev = tq.QuantumDevice(
             n_wires=self.n_wires,
             bsz=1,
@@ -54,6 +93,12 @@ def get_loss(self):
         return expval
 
     def train(self):
+        """Train the VQE model.
+
+        Returns:
+            float: final loss value
+        """
+
         optimizer = getattr(torch.optim, self.optimizer_name)(self.ansatz.parameters(), lr=self.lr)
         lr_scheduler = getattr(torch.optim.lr_scheduler, self.scheduler_name)(optimizer, T_max=self.n_epochs)
         loss = None

torchquantum/datasets/layer_regression.py

@@ -6,31 +6,84 @@
 
 
 class LayerRegressionDatasetV:
+    """Layer Regression Dataset with Variational Input"""
+
     def __init__(self):
+        """Initializes a LayerRegressionDatasetV object.
+
+        Example:
+            >>> dataset = LayerRegressionDatasetV()
+        """
+
         self.data = None
         self.n_instance = 10000
         self.input = np.random.rand(5).astype(np.float) * 2 * np.pi
         self.output = np.random.rand(5).astype(np.float) * 2 - 1
 
     def __getitem__(self, index: int):
+        """Gets an instance from the dataset.
+
+        Args:
+            index (int): The index of the instance.
+
+        Returns:
+            dict: A dictionary containing the input and output.
+
+        Example:
+            >>> dataset = LayerRegressionDatasetV()
+            >>> instance = dataset[0]
+        """
+
         instance = {"input": self.input, "output": self.output}
         return instance
 
     def __len__(self) -> int:
+        """Returns the number of instances in the dataset.
+
+        Returns:
+            int: The number of instances.
+
+        Example:
+            >>> dataset = LayerRegressionDatasetV()
+            >>> length = len(dataset)
+        """
+
         return self.n_instance
 
 
 class LayerRegressionV(Dataset):
+    """Layer Regression Dataset with Variational Input"""
+
     def __init__(
         self,
     ):
+        """Initializes a LayerRegressionV object.
+
+        Example:
+            >>> dataset = LayerRegressionV()
+        """
+
         super().__init__(
             {split: LayerRegressionDatasetV() for split in ["train", "valid", "test"]}
         )
 
 
 class LayerRegressionDataset:
+    """Layer Regression Dataset
+
+    Attributes:
+        data: The dataset.
+        n_instance (int): The number of instances.
+        output: The output data.
+    """
+
     def __init__(self):
+        """Initializes a LayerRegressionDataset object.
+
+        Example:
+            >>> dataset = LayerRegressionDataset()
+        """
+
         self.data = None
         self.n_instance = 10000
         mat = torch.randn((2**5, 2**5), dtype=torch.complex64)
@@ -39,17 +92,48 @@ def __init__(self):
         self.output = mat
 
     def __getitem__(self, index: int):
+        """Gets an instance from the dataset.
+
+        Args:
+            index (int): The index of the instance.
+
+        Returns:
+            dict: A dictionary containing the input and output.
+
+        Example:
+            >>> dataset = LayerRegressionDataset()
+            >>> instance = dataset[0]
+        """
+
         instance = {"input": self.output, "output": self.output}
         return instance
 
     def __len__(self) -> int:
+        """Returns the number of instances in the dataset.
+
+        Returns:
+            int: The number of instances.
+
+        Example:
+            >>> dataset = LayerRegressionDataset()
+            >>> length = len(dataset)
+        """
+
         return self.n_instance
 
 
 class LayerRegression(Dataset):
+    """Layer Regression Dataset"""
+
     def __init__(
         self,
     ):
+        """Initializes a LayerRegression object.
+
+        Example:
+            >>> dataset = LayerRegression()
+        """
+
         super().__init__(
             {split: LayerRegressionDataset() for split in ["train", "valid", "test"]}
         )

torchquantum/datasets/vqe.py

@@ -5,24 +5,85 @@
 
 
 class VQEDataset:
+    """Dataset for VQE.
+
+    Attributes:
+        split (str): Split of the dataset ("train", "valid", "test").
+        steps_per_epoch (int): Number of steps per epoch.
+
+    Methods:
+        __getitem__: Get an item from the dataset
+        __len__: Get the length of the dataset
+
+    Examples:
+        >>> dataset = VQEDataset(split='train', steps_per_epoch=100)
+        >>> instance = dataset[0]
+    """
+
     def __init__(self, split, steps_per_epoch):
+        """Initialize the VQEDataset.
+
+        Args:
+            split (str): Split of the dataset ("train", "valid", "test")
+            steps_per_epoch (int): Number of steps per epoch
+        """
+
         self.split = split
         self.steps_per_epoch = steps_per_epoch
 
     def __getitem__(self, index: int):
+        """Get an instance from the dataset.
+
+        Args:
+            index (int): Index of the item
+
+        Returns:
+            dict: instance containing the input and target
+
+        Examples:
+            >>> instance = dataset[0]
+        """
+
         instance = {"input": -1, "target": -1}
 
         return instance
 
     def __len__(self) -> int:
+        """Get the length of the dataset.
+
+        Returns:
+            int: length of the dataset
+
+        Examples:
+            >>> length = len(dataset)
+        """
+
         if self.split == "train":
             return self.steps_per_epoch
         else:
             return 1
 
 
 class VQE(Dataset):
+    """Dataset for the VQE.
+
+    Attributes:
+        steps_per_epoch (int): Number of steps per epoch
+
+    Methods:
+        __init__: Initialize the VQE dataset
+
+    Examples:
+        >>> dataset = VQE(steps_per_epoch=100)
+    """
+
     def __init__(self, steps_per_epoch):
+        """Initialize the VQE dataset.
+
+        Args:
+            steps_per_epoch (int): Number of steps per epoch
+        """
+
         super().__init__(
             {
                 split: VQEDataset(split=split, steps_per_epoch=steps_per_epoch)