Mypy fix/exact state

.github/workflows/build-test

@@ -61,5 +61,6 @@ fi
 
 if [[ "${MYPY}" = "mypy" ]]
 then
+    python -m pip install --pre -r ${GITHUB_WORKSPACE}/.github/workflows/mypy/requirements.txt
     ${GITHUB_WORKSPACE}/mypy-check ${GITHUB_WORKSPACE}
 fi

.github/workflows/mypy/requirements.txt

@@ -0,0 +1,3 @@
+scipy
+ipython
+pytket

docs/api.rst

@@ -2,5 +2,5 @@ API documentation
 -----------------
 
 .. toctree::
-    modules/fullTN.rst
+    modules/general_state.rst
     modules/structured_state.rst

docs/modules/general_state.rst

@@ -0,0 +1,32 @@
+Full tensor network (general state) contraction
+===============================================
+
+.. automodule:: pytket.extensions.cutensornet.general_state
+
+cuQuantum `contract` API interface
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: pytket.extensions.cutensornet.general_state.TensorNetwork
+
+.. autoclass:: pytket.extensions.cutensornet.general_state.PauliOperatorTensorNetwork
+
+.. autoclass:: pytket.extensions.cutensornet.general_state.ExpectationValueTensorNetwork
+
+.. autofunction:: pytket.extensions.cutensornet.general_state.measure_qubits_state
+
+.. autofunction:: pytket.extensions.cutensornet.general_state.tk_to_tensor_network
+
+cuQuantum `high-level` API interface
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: pytket.extensions.cutensornet.general_state.GeneralState
+
+.. autoclass:: pytket.extensions.cutensornet.general_state.GeneralOperator
+
+.. autoclass:: pytket.extensions.cutensornet.general_state.GeneralExpectationValue
+
+Pytket backend
+~~~~~~~~~~~~~~
+
+.. automodule:: pytket.extensions.cutensornet
+    :members: CuTensorNetBackend

pytket/extensions/cutensornet/__init__.py

@@ -13,17 +13,8 @@
 # limitations under the License.
 """Module for conversion from tket primitives to cuQuantum primitives."""
 
-# _metadata.py is copied to the folder after installation.
-from ._metadata import __extension_version__, __extension_name__  # type: ignore
-
 from .backends import CuTensorNetBackend
+from .general import CuTensorNetHandle
 
-from .tensor_network_convert import (
-    TensorNetwork,
-    PauliOperatorTensorNetwork,
-    ExpectationValueTensorNetwork,
-    tk_to_tensor_network,
-    measure_qubits_state,
-)
-
-from .utils import circuit_statevector_postselect
+# _metadata.py is copied to the folder after installation.
+from ._metadata import __extension_version__, __extension_name__  # type: ignore

pytket/extensions/cutensornet/backends/cutensornet_backend.py

@@ -30,7 +30,7 @@
 from pytket.backends.backend import KwargTypes, Backend, BackendResult
 from pytket.backends.backendinfo import BackendInfo
 from pytket.backends.resulthandle import _ResultIdTuple
-from pytket.extensions.cutensornet.tensor_network_convert import (
+from pytket.extensions.cutensornet.general_state import (
     TensorNetwork,
     ExpectationValueTensorNetwork,
     tk_to_tensor_network,

pytket/extensions/cutensornet/general.py

@@ -11,9 +11,76 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+from __future__ import annotations  # type: ignore
+import warnings
 import logging
 from logging import Logger
 
+from typing import Any, Optional
+
+try:
+    import cupy as cp  # type: ignore
+except ImportError:
+    warnings.warn("local settings failed to import cupy", ImportWarning)
+try:
+    import cuquantum.cutensornet as cutn  # type: ignore
+except ImportError:
+    warnings.warn("local settings failed to import cutensornet", ImportWarning)
+
+
+class CuTensorNetHandle:
+    """Initialise the cuTensorNet library with automatic workspace memory
+    management.
+
+    Note:
+        Always use as ``with CuTensorNetHandle() as libhandle:`` so that cuTensorNet
+        handles are automatically destroyed at the end of execution.
+
+    Attributes:
+        handle (int): The cuTensorNet library handle created by this initialisation.
+        device_id (int): The ID of the device (GPU) where cuTensorNet is initialised.
+            If not provided, defaults to ``cp.cuda.Device()``.
+    """
+
+    def __init__(self, device_id: Optional[int] = None):
+        self._is_destroyed = False
+
+        # Make sure CuPy uses the specified device
+        dev = cp.cuda.Device(device_id)
+        dev.use()
+
+        self.dev = dev
+        self.device_id = dev.id
+
+        self.handle = cutn.create()
+
+    def destroy(self) -> None:
+        """Destroys the memory handle, releasing memory.
+
+        Only call this method if you are initialising a ``CuTensorNetHandle`` outside
+        a ``with CuTensorNetHandle() as libhandle`` statement.
+        """
+        cutn.destroy(self.handle)
+        self._is_destroyed = True
+
+    def __enter__(self) -> CuTensorNetHandle:
+        return self
+
+    def __exit__(self, exc_type: Any, exc_value: Any, exc_tb: Any) -> None:
+        self.destroy()
+
+    def print_device_properties(self, logger: Logger) -> None:
+        """Prints local GPU properties."""
+        device_props = cp.cuda.runtime.getDeviceProperties(self.dev.id)
+        logger.debug("===== device info ======")
+        logger.debug("GPU-name:", device_props["name"].decode())
+        logger.debug("GPU-clock:", device_props["clockRate"])
+        logger.debug("GPU-memoryClock:", device_props["memoryClockRate"])
+        logger.debug("GPU-nSM:", device_props["multiProcessorCount"])
+        logger.debug("GPU-major:", device_props["major"])
+        logger.debug("GPU-minor:", device_props["minor"])
+        logger.debug("========================")
+
 
 def set_logger(
     logger_name: str,

pytket/extensions/cutensornet/general_state/__init__.py

@@ -0,0 +1,26 @@
+# Copyright 2019-2024 Quantinuum
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Module for conversion from tket primitives to cuQuantum primitives."""
+
+from .utils import circuit_statevector_postselect
+
+from .tensor_network_convert import (
+    TensorNetwork,
+    PauliOperatorTensorNetwork,
+    ExpectationValueTensorNetwork,
+    tk_to_tensor_network,
+    measure_qubits_state,
+)
+
+from .tensor_network_state import GeneralState, GeneralOperator, GeneralExpectationValue

pytket/extensions/cutensornet/tensor_network_convert.py → pytket/extensions/cutensornet/general_state/tensor_network_convert.py

@@ -22,10 +22,9 @@
 from networkx.classes.reportviews import OutMultiEdgeView, OutMultiEdgeDataView  # type: ignore
 import numpy as np
 from numpy.typing import NDArray
-from pytket import Qubit  # type: ignore
-from pytket.utils import Graph
-from pytket.pauli import QubitPauliString  # type: ignore
-from pytket.circuit import Circuit, Qubit  # type: ignore
+from pytket.utils import Graph  # type: ignore
+from pytket.pauli import QubitPauliString
+from pytket.circuit import Circuit, Qubit
 from pytket.utils import permute_rows_cols_in_unitary
 from pytket.extensions.cutensornet.general import set_logger
 
@@ -53,18 +52,18 @@ def __init__(
         # self._circuit.replace_implicit_wire_swaps()
         self._qubit_names_ilo = [str(q) for q in self._circuit.qubits]
         self._logger.debug(f"ILO-ordered qubit names: {self._qubit_names_ilo}")
-        self._graph = Graph(self._circuit)
+        self._graph = Graph(self._circuit)  # type: ignore
         qname_to_q = {
             qname: q for qname, q in zip(self._qubit_names_ilo, self._circuit.qubits)
         }
         self._output_index_to_qubit = {
-            oi: qname_to_q[qname] for oi, qname in self._graph.output_names.items()
+            oi: qname_to_q[qname] for oi, qname in self._graph.output_names.items()  # type: ignore
         }
         self._logger.debug(
             f"NX output index to (possibly re-labeled) qubit objects map: "
             f"{self._output_index_to_qubit}"
         )
-        self._network = self._graph.as_nx()
+        self._network = self._graph.as_nx()  # type: ignore
         self._node_tensors = self._assign_node_tensors(adj=adj)
         self._node_tensor_indices, self.sticky_indices = self._get_tn_indices(
             self._network, adj=adj
@@ -264,14 +263,14 @@ def _get_tn_indices(
         nodes_out = self._output_nodes
         # Re-order outward edges indices according to ILO
         edges_out = [
-            edge for edge in net.edges() if edge[1] in self._graph.output_names
+            edge for edge in net.edges() if edge[1] in self._graph.output_names  # type: ignore
         ]
         eids = [
             record[0][0] for key, record in edge_indices.items() if key in edges_out
         ]
         eids_sorted = sorted(eids, key=abs)
-        qnames_graph_ordered = [qname for qname in self._graph.output_names.values()]
-        oids_graph_ordered = [oid for oid in self._graph.output_names.keys()]
+        qnames_graph_ordered = [qname for qname in self._graph.output_names.values()]  # type: ignore
+        oids_graph_ordered = [oid for oid in self._graph.output_names.keys()]  # type: ignore
         eids_qubit_ordered = [
             eids_sorted[qnames_graph_ordered.index(q)] for q in self._qubit_names_ilo
         ]  # Order eid's in the same way as qnames_graph_ordered as compared to ILO