port tetragono to support tat v0.4

tetragono/tetragono/abstract_state.py

@@ -217,7 +217,8 @@ def Edge(self):
         type
             The edge type of this abstract state.
         """
-        return self.Tensor.model.Edge
+        # Get the compat edge constructor.
+        return self.Tensor.__self__.symmetry.Edge
 
     @property
     def Symmetry(self):
@@ -229,7 +230,8 @@ def Symmetry(self):
         type
             The symmetry type of this abstract state.
         """
-        return self.Tensor.model.Symmetry
+        # Group all input into a tuple
+        return staticmethod(lambda *args: args[0] if args and isinstance(args[0], tuple) else args)
 
     def _v2_to_v3_rename(self, state):
         """
@@ -328,10 +330,8 @@ def _construct_symmetry(self, value):
         Symmetry
             The result symmetry object.
         """
-        if isinstance(value, self.Symmetry):
-            return value
-        else:
-            return self.Symmetry(value)
+        # Symmetry is nothing but a tuple wrapper
+        return self.Symmetry(value)
 
     @property
     def total_symmetry(self):
@@ -367,7 +367,7 @@ def _total_symmetry_edge(self):
         Edge
             The result virtual edge.
         """
-        return self.Edge([(-self._total_symmetry, 1)], False)
+        return self.Edge([(self._total_symmetry, 1)], False).conjugate()
 
     def _construct_edge(self, value):
         """
@@ -383,10 +383,8 @@ def _construct_edge(self, value):
         Edge
             The result edge object.
         """
-        if isinstance(value, self.Edge):
-            return value
-        else:
-            return self.Edge(value)
+        # Edge accept argument with type Edge. So previous condition is not needed.
+        return self.Edge(value)
 
     def _construct_physics_edge(self, edge):
         """
@@ -444,8 +442,9 @@ def _set_hamiltonian(self, points, tensor):
              The hamiltonian tensor.
         """
         body = len(points)
-        if not isinstance(tensor, self.Tensor):
-            raise TypeError("Wrong hamiltonian type")
+        # Do not check hamiltonian type temporarily
+        #if not isinstance(tensor, self.Tensor):
+        #    raise TypeError("Wrong hamiltonian type")
         if {f"{i}" for i in tensor.names} != {f"{i}{j}" for i in ["I", "O"] for j in range(body)}:
             raise ValueError("Wrong hamiltonian name")
         for i in range(body):

tetragono/tetragono/sampling_lattice/lattice.py

@@ -313,7 +313,10 @@ def _get_shrinker(self, l1l2, configuration):
             # P side is dimension one edge
             # Q side is connected to lattice
             shrinker = self.Tensor(["P", "Q"], [[(symmetry, 1)], edge.conjugate()]).zero_()
-            shrinker[{"Q": (-symmetry, index), "P": (symmetry, 0)}] = 1
+            shrinker[{
+                "Q": (tuple(sym if isinstace(sym, bool) else -sym for sym in symmetry), index),
+                "P": (symmetry, 0)
+            }] = 1
             yield orbit, shrinker
 
     def _shrink_configuration(self, l1l2, configuration):

tetragono/tetragono/sampling_lattice/observer.py

@@ -18,7 +18,8 @@
 
 import os
 import numpy as np
-import PyScalapack
+# PyScalapack support should be dropped, I think...
+#import PyScalapack
 from ..utility import (show, showln, allreduce_lattice_buffer, allreduce_buffer, allreduce_number, bcast_buffer,
                        lattice_update, lattice_prod_sum, lattice_conjugate, mpi_rank, mpi_size, mpi_comm, pickle)
 from ..tensor_element import tensor_element
@@ -274,8 +275,10 @@ def add_observer(self, name, observers):
         if self._start:
             raise RuntimeError("Cannot enable hole after sampling start")
         for positions, observer in observers.items():
-            if not isinstance(observer, self.owner.Tensor):
-                raise TypeError("Wrong observer type")
+            pass
+            # Do not check the type of observer temporarily
+            #if not isinstance(observer, self.owner.Tensor):
+            #    raise TypeError("Wrong observer type")
         self._observer[name] = observers
 
     def add_energy(self):
@@ -398,7 +401,7 @@ def __call__(self, possibility, configuration):
                 # Es should be complex here when calculating gradient
             if name == "energy" and self._enable_gradient:
                 Es = whole_name_value
-                if self.owner.Tensor.is_real:
+                if self.owner[0, 0].is_real:
                     Es = Es.real
                 holes = configuration.holes()  # <psi|s|partial_x psi> / <psi|s|psi>
                 for l1, l2 in self.owner.sites():
@@ -519,7 +522,7 @@ def total_energy(self):
 
     def _total_energy_with_imaginary_part(self):
         name = "energy"
-        if self.owner.Tensor.is_complex:
+        if self.owner[0, 0].is_complex:
             return (self._whole_result_reweight[name] + self._total_imaginary_energy_reweight * 1j) / self._total_weight
         else:
             return self._whole_result_reweight[name] / self._total_weight
@@ -592,17 +595,15 @@ def natural_gradient_by_conjugate_gradient(self, step, error):
         energy = self._total_energy_with_imaginary_part()
         delta = self._delta_to_array(self._Delta) / self._total_weight
 
-        dtype = np.dtype(self.owner.Tensor.dtype)
-        btype = self.owner.Tensor.btype
-
         Delta = []
         Energy = []
         for reweight_s, energy_s, delta_s in self._weights_and_deltas():
             param = (reweight_s / self._total_weight)**(1 / 2)
+            param = param.item()
             Delta.append((self._delta_to_array(delta_s) - delta) * param)
             Energy.append((energy_s.conjugate() - energy) * param)
-        Delta = np.asarray(Delta, dtype=dtype)
-        Energy = np.asarray(Energy, dtype=dtype)
+        Delta = np.asarray(Delta)
+        Energy = np.asarray(Energy)
 
         # A x = b
         # DT r D x = DT r E
@@ -663,8 +664,8 @@ def _delta_to_array(self, delta):
         # Both delta and result array is in bra space
         result = []
         for l1, l2 in self.owner.sites():
-            result.append(delta[l1][l2].transpose(self._Delta[l1][l2].names).storage)
-        result = np.concatenate(result, dtype=self.owner.Tensor.dtype)
+            result.append(delta[l1][l2].transpose(self._Delta[l1][l2].names).copy().storage)
+        result = np.concatenate(result, dtype=result[0].dtype)
         return result
 
     def _array_to_delta(self, array):

tetragono/tetragono/sampling_lattice/sampling.py

@@ -320,12 +320,7 @@ def __call__(self):
                         .transpose(["I", "O"]))
                     hole_edge = hole.edge_by_name("O")
                     # Calculate rho for all the segments of the physics edge of this orbit
-                    rho = []
-                    for seg in hole_edge.segments:
-                        symmetry, _ = seg
-                        block_rho = hole.blocks[[("I", -symmetry), ("O", symmetry)]]
-                        diag_rho = np.diagonal(block_rho)
-                        rho = [*rho, *diag_rho]
+                    rho = hole.data.diagonal()
                     rho = np.array(rho).real
                     rho = np.maximum(rho, 0)  # Sometimes there is some negative value because of numeric error.
                     if np.sum(rho) == 0:

tetragono/tetragono/tensor_element.py

@@ -27,66 +27,45 @@ def tensor_element(tensor):
     return element_pool[tensor_id]
 
 
-def loop_nonzero_block(block, symmetries, rank, names, template):
-    indices = [0 for _ in range(rank)]
-    for i in range(rank):
-        if block.shape[i] == 0:
-            # dims == 0
-            return
-    while True:
-        value = block[tuple(indices)]
-        if value != 0:
-            # yield
-            template[{names[i]: 0 for i in range(rank)}] = value
-            yield [(symmetries[i], indices[i]) for i in range(rank)], template.copy()
-
-        edge_position = rank - 1
-
-        indices[edge_position] += 1
-        while indices[edge_position] == block.shape[edge_position]:
-            if edge_position == 0:
-                return
-            indices[edge_position] = 0
-            edge_position -= 1
-            indices[edge_position] += 1
-
-
 def loop_nonzero_tensor(tensor, names, rank):
-    # see include/TAT/structure/edge.hpp
-    edges = [tensor.edges[i].segments for i in range(rank)]
-    arrow = [tensor.edges[i].arrow for i in range(rank)]
-    Edge = tensor.model.Edge
+    if tensor.data.nelement() == 0:
+        return
     if rank == 0:
         # rank == 0
         yield [], tensor
         return
-    symmetry_indices = [0 for _ in range(rank)]
-    for i in range(rank):
-        if len(edges[i]) == 0:
-            # dims == 0
-            return
-    zero_symmetry = tensor.model.Symmetry()
+    indices = [0 for _ in range(rank)]
     while True:
-        symmetries = [edges[i][symmetry_indices[i]][0] for i in range(rank)]
-        if sum(symmetries, start=zero_symmetry) == zero_symmetry:
-            block = tensor.blocks[[(names[i], symmetries[i]) for i in range(rank)]]
-            template_edges = [Edge([symmetries[i]], arrow[i]) for i in range(rank)]
-            template = type(tensor)(names, template_edges)
-            yield from loop_nonzero_block(block, symmetries, rank, names, template)
+        if tensor.data[tuple(indices)] != 0:
+            element = tensor.__class__(
+                names=tensor.names,
+                edges=tuple(tensor.edges[i].__class__(
+                    fermion=tensor.fermion,
+                    dtypes=tensor.dtypes,
+                    symmetry=tuple(symmetryp[indices[i]].reshape([1]) for symmetry in tensor.edges[i].symmetry),
+                    dimension=1,
+                    arrow=tensor.edges[i].arrow,
+                    parity=tensor.edges[i].parity[indices[i]].reshape([1]),
+                ) for i in range(rank)),
+                fermion=tensor.fermion,
+                dtypes=tensor.dtypes,
+                data=tensor.data[tuple(indices)].reshape([1 for _ in range(rank)]),
+                mask=tensor.mask[tuple(indices)].reshape([1 for _ in range(rank)]),
+            )
+            yield [tensor.edges[i].point_by_index(indices[i]) for i in range(rank)], element
 
         edge_position = rank - 1
-
-        symmetry_indices[edge_position] += 1
-        while symmetry_indices[edge_position] == len(edges[edge_position]):
+        indices[edge_position] += 1
+        while indices[edge_position] == tensor.edges[edge_position].dimension:
             if edge_position == 0:
                 return
-            symmetry_indices[edge_position] = 0
+            indices[edge_position] = 0
             edge_position -= 1
-            symmetry_indices[edge_position] += 1
+            indices[edge_position] += 1
 
 
 def conjugate_symmetry(edge_point):
-    return (-edge_point[0], edge_point[1])
+    return (tuple(sym if isinstance(sym) else -sym for sym in edge_point[0]), edge_point[1])
 
 
 def calculate_element(tensor):