Update Emulator interface. (#943)

* refactor `AnalogGate` interface to accept and return `StateVector` objects. fixing tests.

* adding `tocsr` to Hamiltonian.

* adding API to `HamiltonianData`.

* adding doc strings.

* switching to list as hamiltonian output.

* rename class.

* adding fock state API.

* adding fock state API.

* removing print.

* adding tests.

* adding more tests.

* removing unused code.

* expanding tests.

* rename file.

* fixing bug in `_apply`.

* updating tests.

pyproject.toml

@@ -101,6 +101,7 @@ dev = [
     "scikit-optimize>=0.9.0",
     "matplotlib>=3.8.1",
     "icecream>=2.1.3",
+    "tqdm>=4.66.4",
 ]
 
 [tool.pdm.scripts]

src/bloqade/emulate/ir/emulator.py

@@ -172,15 +172,35 @@ def __post_init__(self):
     def __len__(self):
         return len(self.sites)
 
+    def same(self, other: Any) -> bool:
+        if not isinstance(other, Register):
+            return False
+
+        return (
+            self.atom_type == other.atom_type
+            and self.blockade_radius == other.blockade_radius
+            and set(self.sites) == set(other.sites)
+        )
+
+    # for hashing/comparing I only want to compare based on the effective
+    # fock space that gets generated which is determined by the geometry
+    # only if the blockade radius is non-zero. I overload the __eq__ and
+    # __hash__ methods to achieve this.
+
     def __eq__(self, other: Any):
-        if isinstance(other, Register):
-            return (
-                self.atom_type == other.atom_type
-                and self.blockade_radius == other.blockade_radius
-                and set(self.sites) == set(other.sites)
+        if not isinstance(other, Register):
+            return False
+
+        if self.blockade_radius == Decimal("0") and other.blockade_radius == Decimal(
+            "0"
+        ):
+            # if blockade radius is zero, then the positions are irrelevant
+            # because the fock states generated by the geometry are the same
+            return self.atom_type == other.atom_type and len(self.sites) == len(
+                other.sites
             )
 
-        return False
+        return self.same(other)
 
     def __hash__(self) -> int:
         if self.blockade_radius == Decimal("0"):

src/bloqade/emulate/ir/space.py

@@ -1,12 +1,13 @@
 from dataclasses import dataclass
 from numpy.typing import NDArray
-from beartype.typing import TYPE_CHECKING
+from beartype.typing import TYPE_CHECKING, Any
 import numpy as np
 from enum import Enum
 
 if TYPE_CHECKING:
     from .emulator import Register
     from .atom_type import AtomType
+    from .state_vector import StateVector
 
 MAX_PRINT_SIZE = 30
 
@@ -28,6 +29,12 @@ def __post_init__(self):
 
         assert isinstance(self.program_register, Register)
 
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, Space):
+            return False
+
+        return self.program_register == other.program_register
+
     @classmethod
     def create(cls, register: "Register"):
         sites = register.sites
@@ -193,10 +200,12 @@ def index_to_fock_state(self, index: int) -> str:
                 self.configurations[index], self.n_atoms
             )
 
-    def zero_state(self, dtype=np.float64) -> NDArray:
+    def zero_state(self, dtype=np.float64) -> "StateVector":
+        from .state_vector import StateVector
+
         state = np.zeros(self.size, dtype=dtype)
         state[0] = 1.0
-        return state
+        return StateVector(state, self)
 
     def sample_state_vector(
         self, state_vector: NDArray, n_samples: int, project_hyperfine: bool = True