Merge pull request #17 from simsaidan/aidanonefileonefunction

refactor: split core definitions into focused modules

Author: simsaidan

examples/simple_algo_benchmark.py

@@ -9,25 +9,27 @@
 if ROOT not in sys.path:
     sys.path.insert(0, ROOT)
 
-from src.core.utils import (
+from src.utils.constants import (
     zero_state_two_qubits,
     one_state_two_qubits,
     two_state_two_qubits,
     three_state_two_qubits,
-    tensor,
     I,
 )
+from src.utils.tensor import tensor
 
 from src.benchmarking.benchmark import Benchmark
-from src.core.hamiltonian import LocalHamiltonian, Hamiltonian
+from src.core.hamiltonian import Hamiltonian
+from src.core.local_hamiltonian import LocalHamiltonian
 from src.simulators.qutip_simulator import QutipSimulator
 from src.simulators.splitjmatrix_simulator import (
     SplitJMatrixSimulator,
     _cavity_population,
     _population_one,
 )
+from src.core.local_operator import LocalOperator
+from src.core.operator import Operator
 from src.core.state import State
-from src.core.operator_set import LocalOperator, Operator
 
 # Define system parameters
 omega_c = 245000

examples/simple_shadow.py

@@ -10,8 +10,9 @@
 import numpy as np
 
 from src.core.hamiltonian import Hamiltonian
-from src.core.operator_set import Operator, OperatorSet
-from src.core.pauli import PauliString
+from src.core.operator import Operator
+from src.core.operator_set import OperatorSet
+from src.core.pauli_string import PauliString
 from src.simulators.qutip_simulator import QutipSimulator
 from src.core.state import State
 from src.shadow.shadow_hamiltonian import ShadowHamiltonian

src/core/combined_hamiltonian_matrix.py

@@ -0,0 +1,67 @@
+import numpy as np
+
+from src.core.hamiltonian import Hamiltonian
+from src.core.local_hamiltonian import LocalHamiltonian
+
+
+def combined_hamiltonian_matrix(
+    hamiltonians: list[Hamiltonian],
+    num_qubits: int,
+) -> np.ndarray:
+    """
+    Sum Hamiltonian terms on the full ``num_qubits``-site tensor space.
+
+    ``LocalHamiltonian`` terms are embedded with identities on the remaining
+    sites; full-domain ``Hamiltonian`` matrices are added as-is. All terms must
+    match a common Hilbert-space dimension (``local_dim ** num_qubits`` for
+    locals, or the matrix size of bare terms).
+
+    Parameter hamiltonians: The list of Hamiltonians to sum.
+    Precondition: hamiltonians is a list of Hamiltonian objects.
+
+    Parameter num_qubits: The number of qubits in the full system.
+    Precondition: num_qubits is a positive integer.
+    """
+    assert isinstance(hamiltonians, list), "hamiltonians must be a list"
+    assert all(
+        isinstance(h, Hamiltonian) for h in hamiltonians
+    ), "all hamiltonians must be Hamiltonian objects"
+    assert isinstance(num_qubits, int), "num_qubits must be an integer"
+    assert num_qubits > 0, "num_qubits must be a positive integer"
+    assert len(hamiltonians) > 0, "hamiltonians must be a non-empty list"
+
+    local_dims = {h.local_dim for h in hamiltonians if isinstance(h, LocalHamiltonian)}
+    if len(local_dims) > 1:
+        raise ValueError(
+            f"mixed local_dim in LocalHamiltonian terms is not supported: {sorted(local_dims)}"
+        )
+
+    target_dim: int | None = None
+    for h in hamiltonians:
+        if isinstance(h, LocalHamiltonian):
+            d = h.local_dim**num_qubits
+        else:
+            d = int(np.asarray(h.matrix).shape[0])
+        if target_dim is None:
+            target_dim = d
+        elif d != target_dim:
+            raise ValueError(
+                f"Hamiltonian dimensions disagree: got {d} vs {target_dim}"
+            )
+
+    H_tot = np.zeros((target_dim, target_dim), dtype=np.complex128)
+    for h in hamiltonians:
+        if isinstance(h, LocalHamiltonian):
+            ld = h.local_dim
+            lo, hi = min(h.sites), max(h.sites)
+            n_before = lo
+            n_after = num_qubits - 1 - hi
+            term = np.asarray(h.matrix, dtype=np.complex128)
+            if n_before > 0:
+                term = np.kron(np.eye(ld**n_before, dtype=np.complex128), term)
+            if n_after > 0:
+                term = np.kron(term, np.eye(ld**n_after, dtype=np.complex128))
+            H_tot += term
+        else:
+            H_tot += np.asarray(h.matrix, dtype=np.complex128)
+    return H_tot

src/core/hamiltonian.py

@@ -1,119 +1,62 @@
 import numpy as np
 
-from src.core.operator_set import LocalOperator, Operator
+from src.core.operator import Operator
+from src.utils.hermitian import hermitian
 
 
 class Hamiltonian(Operator):
+    """
+    Represents a quantum Hamiltonian.
+    The Hamiltonian must be Hermitian to be valid.
+    """
+
     def __init__(self, matrix: np.ndarray):
+        """Initializes a Hamiltonian object which represents a quantum
+        Hamiltonian. The Hamtonian must be Hermitian to be valid.
+
+        Parameter matrix: The matrix representation of the Hamiltonian.
+        Precondition: matrix is a numpy array.
+        """
+        assert isinstance(matrix, np.ndarray), "matrix must be a numpy array"
+        assert hermitian(matrix), "matrix must be Hermitian"
+        assert (
+            matrix.ndim == 2 and matrix.shape[0] == matrix.shape[1]
+        ), "matrix must be a square 2D array"
         super().__init__(matrix)
 
     def __str__(self):
+        """Returns a string representation of the Hamiltonian."""
         return f"Hamiltonian(matrix={self.matrix})"
 
     def __repr__(self):
+        """Returns a string representation of the Hamiltonian."""
         return f"Hamiltonian(matrix={self.matrix})"
 
     def __eq__(self, other):
+        """Returns whether the Hamiltonian is equal to another Hamiltonian."""
         return np.allclose(self.matrix, other.matrix)
 
     def __ne__(self, other):
+        """Returns whether the Hamiltonian is not equal to another Hamiltonian."""
         return not np.allclose(self.matrix, other.matrix)
 
     def to_hamiltonian_set(self):
+        """Returns a HamiltonianSet containing the Hamiltonian."""
+        from src.core.hamiltonian_set import HamiltonianSet
+
         return HamiltonianSet([self])
 
     def to_local_hamiltonian(self, local_dim: int = 2):
-        lo = self.to_local_operator(local_dim)
-        return LocalHamiltonian(lo.matrix, lo.sites, lo.local_dim)
-
-
-class LocalHamiltonian(LocalOperator, Hamiltonian):
-    """
-    Local Hamiltonian acting non-trivially on the provided `sites`.
-
-    ``sites`` must be distinct consecutive integers; see `LocalOperator`.
-    """
-
-    def __init__(
-        self,
-        matrix: np.ndarray,
-        sites: list[int],
-        local_dim: int = 2,
-    ):
-        super().__init__(matrix, sites, local_dim)
-
-    def __str__(self):
-        return (
-            "LocalHamiltonian("
-            f"sites={self.sites}, "
-            f"local_dim={self.local_dim}, "
-            f"shape={self.matrix.shape}"
-            ")"
-        )
+        """Returns a LocalHamiltonian representing the Hamiltonian acting on
+        the given sites.
 
-    def __repr__(self):
-        return (
-            "LocalHamiltonian("
-            f"matrix={self.matrix!r}, "
-            f"sites={self.sites!r}, "
-            f"local_dim={self.local_dim}"
-            ")"
-        )
-
-
-class HamiltonianSet:
-    def __init__(self, hamiltonians: list[Hamiltonian]):
-        self.hamiltonians = hamiltonians
-        self.hamiltonian_count = len(hamiltonians)
-
-    def __str__(self):
-        return f"HamiltonianSet(hamiltonian_count={self.hamiltonian_count})"
-
-    def __repr__(self):
-        return f"HamiltonianSet(hamiltonians={self.hamiltonians!r})"
+        Parameter local_dim: The local dimension of the Hamiltonian.
+        Precondition: local_dim is a positive integer.
+        """
+        assert isinstance(local_dim, int), "local_dim must be an integer"
+        assert local_dim > 0, "local_dim must be a positive integer"
 
+        from src.core.local_hamiltonian import LocalHamiltonian
 
-def combined_hamiltonian_matrix(
-    hamiltonians: list[Hamiltonian],
-    num_qubits: int,
-) -> np.ndarray:
-    """
-    Sum Hamiltonian terms on the full ``num_qubits``-site tensor space.
-
-    ``LocalHamiltonian`` terms are embedded with identities on the remaining
-    sites; full-domain ``Hamiltonian`` matrices are added as-is. All terms must
-    match a common Hilbert-space dimension (``local_dim ** num_qubits`` for
-    locals, or the matrix size of bare terms).
-    """
-    if not hamiltonians:
-        raise ValueError("hamiltonians must be a non-empty list")
-
-    target_dim: int | None = None
-    for h in hamiltonians:
-        if isinstance(h, LocalHamiltonian):
-            d = h.local_dim**num_qubits
-        else:
-            d = int(np.asarray(h.matrix).shape[0])
-        if target_dim is None:
-            target_dim = d
-        elif d != target_dim:
-            raise ValueError(
-                f"Hamiltonian dimensions disagree: got {d} vs {target_dim}"
-            )
-
-    H_tot = np.zeros((target_dim, target_dim), dtype=np.complex128)
-    for h in hamiltonians:
-        if isinstance(h, LocalHamiltonian):
-            ld = h.local_dim
-            lo, hi = min(h.sites), max(h.sites)
-            n_before = lo
-            n_after = num_qubits - 1 - hi
-            term = np.asarray(h.matrix, dtype=np.complex128)
-            if n_before > 0:
-                term = np.kron(np.eye(ld**n_before, dtype=np.complex128), term)
-            if n_after > 0:
-                term = np.kron(term, np.eye(ld**n_after, dtype=np.complex128))
-            H_tot += term
-        else:
-            H_tot += np.asarray(h.matrix, dtype=np.complex128)
-    return H_tot
+        lo = self.to_local_operator(local_dim)
+        return LocalHamiltonian(lo.matrix, lo.sites, lo.local_dim)

src/core/hamiltonian_set.py

@@ -0,0 +1,43 @@
+from src.core.hamiltonian import Hamiltonian
+
+
+class HamiltonianSet:
+    """Represents a set of quantum Hamiltonians."""
+
+    def __init__(self, hamiltonians: list[Hamiltonian]):
+        """Initializes a HamiltonianSet object which represents a set of quantum
+        Hamiltonians.
+
+        Parameter hamiltonians: The list of Hamiltonians to initialize the HamiltonianSet with.
+        Precondition: hamiltonians is a list of Hamiltonian objects.
+        """
+        assert isinstance(hamiltonians, list), "hamiltonians must be a list"
+        assert all(
+            isinstance(h, Hamiltonian) for h in hamiltonians
+        ), "all hamiltonians must be Hamiltonian objects"
+        self.hamiltonians = hamiltonians
+        self.hamiltonian_count = len(hamiltonians)
+
+    def __str__(self):
+        """Returns a string representation of the HamiltonianSet."""
+        return f"HamiltonianSet(hamiltonian_count={self.hamiltonian_count})"
+
+    def __repr__(self):
+        """Returns a string representation of the HamiltonianSet."""
+        return f"HamiltonianSet(hamiltonians={self.hamiltonians!r})"
+
+    def __iter__(self):
+        """Returns an iterator over the Hamiltonians in the HamiltonianSet."""
+        return iter(self.hamiltonians)
+
+    def __len__(self):
+        """Returns the number of Hamiltonians in the HamiltonianSet."""
+        return self.hamiltonian_count
+
+    def __getitem__(self, index: int):
+        """Returns the Hamiltonian at the given index."""
+        return self.hamiltonians[index]
+
+    def __contains__(self, item: Hamiltonian):
+        """Returns whether the HamiltonianSet contains the given Hamiltonian."""
+        return item in self.hamiltonians