Interaction matrix calculation (#735)

* adding method to calculate rydberg interaction.

* adding test.

* updating docs to be more clear.

* making loop break.

* removing skip on vacant sites.

* removing old print.

* updating comment to match test.

* bump minor version

* adding scikit-optimize to dev deps.

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bloqade"
-version = "0.8.0"
+version = "0.9.0"
 description = "Neutral atom software development kit"
 authors = [
     {name = "QuEra Computing Inc.", email = "[email protected]"},
@@ -96,6 +96,7 @@ dev = [
     "pytest-recording>=0.12.2",
     "vcrpy==4.4.0",
     "pyinstrument>=4.5.3",
+    "scikit-optimize>=0.9.0",
 ]
 
 [tool.pdm.scripts]

src/bloqade/emulate/sparse_operator.py

@@ -149,7 +149,6 @@ def _matvec_imp(col_indices, row_indices, scale, input, output):
                 )
 
         elif isinstance(self.col_indices, slice):
-            print(self.col_indices, self.row_indices)
 
             def _matvec_imp(col_indices, row_indices, scale, input, output):
                 return _index_mapping_col_sliced(

src/bloqade/ir/location/base.py

@@ -10,6 +10,8 @@
 from enum import Enum
 import plotext as pltxt
 import sys
+import numpy as np
+from numpy.typing import NDArray
 from bloqade.visualization import get_atom_arrangement_figure
 from bloqade.visualization import display_ir
 
@@ -112,6 +114,35 @@ def figure(self, fig_kwargs=None, **assignments):
     def show(self, **assignments) -> None:
         display_ir(self, assignments)
 
+    def rydberg_interaction(self, **assignments) -> NDArray:
+        """calculate the Rydberg interaction matrix.
+
+        Args:
+            **assignments: the values to assign to the variables in the register.
+
+        Returns:
+            NDArray: the Rydberg interaction matrix in the lower triangular form.
+
+        """
+
+        from bloqade.constants import RB_C6
+
+        # calculate the Interaction matrix
+        V_ij = np.zeros((self.n_sites, self.n_sites))
+        for i, site_i in enumerate(self.enumerate()):
+            pos_i = np.array([float(ele(**assignments)) for ele in site_i.position])
+
+            for j, site_j in enumerate(self.enumerate()):
+                if j >= i:
+                    break  # enforce lower triangular form
+
+                pos_j = np.array([float(ele(**assignments)) for ele in site_j.position])
+                r_ij = np.linalg.norm(pos_i - pos_j)
+
+                V_ij[i, j] = RB_C6 / r_ij**6
+
+        return V_ij
+
     @property
     def n_atoms(self) -> int:
         """number of atoms (filled sites) in the register."""

tests/test_lattice.py

@@ -1,8 +1,41 @@
 import bloqade.ir as ir
 from bloqade.ir.location import ListOfLocations, AtomArrangement
 from bloqade.ir.location import Square
+from bloqade.constants import RB_C6
 from bloqade import cast
 import pytest
+import numpy as np
+
+
+def test_rydberg_interactions():
+    geometry = ListOfLocations([(0, 0), (1, 0), (0, 1), (1, 1)]).scale(5.0)
+
+    V_ij = geometry.rydberg_interaction()
+
+    # 2  3
+    #  \/
+    #  /\
+    # 0  1
+
+    d_01 = 5.0
+    d_02 = 5.0
+    d_03 = np.sqrt(2) * 5.0
+    d_12 = np.sqrt(2) * 5.0
+    d_13 = 5.0
+    d_23 = 5.0
+
+    V_01 = RB_C6 / d_01**6
+    V_02 = RB_C6 / d_02**6
+    V_03 = RB_C6 / d_03**6
+    V_12 = RB_C6 / d_12**6
+    V_13 = RB_C6 / d_13**6
+    V_23 = RB_C6 / d_23**6
+
+    V_ij_expected = np.array(
+        [[0, 0, 0, 0], [V_01, 0, 0, 0], [V_02, V_12, 0, 0], [V_03, V_13, V_23, 0]]
+    )
+
+    assert np.allclose(V_ij, V_ij_expected)
 
 
 def test_listOfLocations():