Add feature to make life eaiser for users of the expectation function classes

Author: JanLuca

peps_ad/expectation/two_sites.py

@@ -6,6 +6,7 @@
 import jax.numpy as jnp
 from jax import jit
 
+from peps_ad import peps_ad_config
 from peps_ad.peps import PEPS_Tensor, PEPS_Unit_Cell
 from peps_ad.contractions import apply_contraction
 from .model import Expectation_Model
@@ -419,6 +420,12 @@ class Two_Sites_Expectation_Value(Expectation_Model):
     spiral_unitary_operator: Optional[jnp.ndarray] = None
 
     def __post_init__(self) -> None:
+        if isinstance(self.horizontal_gates, jnp.ndarray):
+            self.horizontal_gates = (self.horizontal_gates,)
+
+        if isinstance(self.vertical_gates, jnp.ndarray):
+            self.vertical_gates = (self.vertical_gates,)
+
         if (
             len(self.horizontal_gates) > 0
             and len(self.horizontal_gates) > 0

peps_ad/mapping/florett_pentagon.py

@@ -775,6 +775,15 @@ class Florett_Pentagon_Expectation_Value(Expectation_Model):
     spiral_unitary_operator: Optional[jnp.ndarray] = None
 
     def __post_init__(self) -> None:
+        if isinstance(self.black_gates, jnp.ndarray):
+            self.black_gates = (self.black_gates,)
+
+        if isinstance(self.green_gates, jnp.ndarray):
+            self.green_gates = (self.green_gates,)
+
+        if isinstance(self.blue_gates, jnp.ndarray):
+            self.blue_gates = (self.blue_gates,)
+
         if (len(self.green_gates) != len(self.blue_gates)) or (
             len(self.green_gates) != len(self.black_gates)
         ):

peps_ad/mapping/honeycomb.py

@@ -55,6 +55,15 @@ class Honeycomb_Expectation_Value(Expectation_Model):
     spiral_unitary_operator: Optional[jnp.ndarray] = None
 
     def __post_init__(self) -> None:
+        if isinstance(self.x_gates, jnp.ndarray):
+            self.x_gates = (self.x_gates,)
+
+        if isinstance(self.y_gates, jnp.ndarray):
+            self.y_gates = (self.y_gates,)
+
+        if isinstance(self.z_gates, jnp.ndarray):
+            self.z_gates = (self.z_gates,)
+
         if (
             (
                 len(self.x_gates) > 0

peps_ad/mapping/kagome.py

@@ -53,6 +53,12 @@ class Kagome_PESS3_Expectation_Value(Expectation_Model):
     spiral_unitary_operator: Optional[jnp.ndarray] = None
 
     def __post_init__(self) -> None:
+        if isinstance(self.upward_triangle_gates, jnp.ndarray):
+            self.upward_triangle_gates = (self.upward_triangle_gates,)
+
+        if isinstance(self.downward_triangle_gates, jnp.ndarray):
+            self.downward_triangle_gates = (self.downward_triangle_gates,)
+
         if (
             len(self.upward_triangle_gates) > 0
             and len(self.downward_triangle_gates) > 0

peps_ad/mapping/maple_leaf.py

@@ -426,6 +426,15 @@ class Maple_Leaf_Expectation_Value(Expectation_Model):
     spiral_unitary_operator: Optional[jnp.ndarray] = None
 
     def __post_init__(self) -> None:
+        if isinstance(self.green_gates, jnp.ndarray):
+            self.green_gates = (self.green_gates,)
+
+        if isinstance(self.blue_gates, jnp.ndarray):
+            self.blue_gates = (self.blue_gates,)
+
+        if isinstance(self.red_gates, jnp.ndarray):
+            self.red_gates = (self.red_gates,)
+
         if (len(self.green_gates) != len(self.blue_gates)) or (
             len(self.green_gates) != len(self.red_gates)
         ):

peps_ad/mapping/square_kagome.py

@@ -535,6 +535,18 @@ class Square_Kagome_Expectation_Value(Expectation_Model):
     spiral_unitary_operator: Optional[jnp.ndarray] = None
 
     def __post_init__(self) -> None:
+        if isinstance(self.triangle_gates, jnp.ndarray):
+            self.triangle_gates = (self.triangle_gates,)
+
+        if isinstance(self.square_gates, jnp.ndarray):
+            self.square_gates = (self.square_gates,)
+
+        if isinstance(self.plus_gates, jnp.ndarray):
+            self.plus_gates = (self.plus_gates,)
+
+        if isinstance(self.cross_gates, jnp.ndarray):
+            self.cross_gates = (self.cross_gates,)
+
         if len(self.cross_gates) > 0:
             raise NotImplementedError("Cross term calculation is not implemented yet.")
 

peps_ad/mapping/triangular.py

@@ -55,6 +55,10 @@ class Triangular_Expectation_Value(Expectation_Model):
     is_spiral_peps: bool = False
     spiral_unitary_operator: Optional[jnp.ndarray] = None
 
+    def __post_init__(self) -> None:
+        if isinstance(self.nearest_neighbor_gates, jnp.ndarray):
+            self.nearest_neighbor_gates = (self.nearest_neighbor_gates,)
+
     def __call__(
         self,
         peps_tensors: Sequence[jnp.ndarray],