optimize mpo construction

Author: jjren

renormalizer/mps/mpo.py

@@ -271,12 +271,12 @@ def __init__(self, model: Model = None, terms: Union[Op, List[Op]] = None, offse
         if len(terms) == 0:
             raise ValueError("Terms all have factor 0.")
 
-        table, factor = _terms_to_table(model, terms, -self.offset)
+        table, primary_ops, factor = _terms_to_table(model, terms, -self.offset)
 
         self.dtype = factor.dtype
 
         self.symbolic_mpo, self.qn, self.qntot, self.qnidx, self.symbolic_out_ops_list, self.primary_ops \
-            = construct_symbolic_mpo(table, factor, algo=algo)
+            = construct_symbolic_mpo(table, primary_ops, factor, algo=algo)
         # from renormalizer.mps.symbolic_mpo import _format_symbolic_mpo
         # print(_format_symbolic_mpo(mpo_symbol))
         self.model = model

renormalizer/mps/symbolic_mpo.py

@@ -19,7 +19,7 @@
 OpTuple = namedtuple("OpTuple", ["symbol", "qn", "factor"])
 
 
-def construct_symbolic_mpo(table, factor, algo="Hopcroft-Karp"):
+def construct_symbolic_mpo(table, primary_ops, factor, algo="Hopcroft-Karp"):
     r"""
     A General Compact (Symbolic) MPO Construction Routine
 
@@ -102,17 +102,18 @@ def construct_symbolic_mpo(table, factor, algo="Hopcroft-Karp"):
     The local mpo is the transformation matrix between 0'',1'' to 0'''
     """
 
-    qn_size = len(table[0][0].qn)
+    qn_size = len(primary_ops[0].qn)
+
     # Simplest case. Cut to the chase
-    if len(table) == 1:
+    if table.shape[0] == 1:
         # The first layer: number of sites. The middle array: in and out virtual bond
         # the 4th layer: operator sums
         mpo: List[np.ndarray[List[Op]]] = []
         mpoqn = [np.zeros((1, qn_size), dtype=int)]
-        primary_ops = list(set(table[0]))
         op2idx = dict(zip(primary_ops, range(len(primary_ops))))
         out_ops_list: List[List[OpTuple]] = [[OpTuple([0], qn=0, factor=1)]]
-        for op in table[0]:
+        for idx in table[0]:
+            op = primary_ops[idx]
             mo = np.full((1, 1), None)
             mo[0][0] = [op]
             mpo.append(mo)
@@ -130,14 +131,10 @@ def construct_symbolic_mpo(table, factor, algo="Hopcroft-Karp"):
         return mpo, mpoqn, qntot, qnidx, out_ops_list, primary_ops
 
     logger.debug(f"symbolic mpo algorithm: {algo}")
-    logger.debug(f"Input operator terms: {len(table)}")
-
-    table, factor, primary_ops = _transform_table(table, factor)
 
     # add the first and last column for convenience
     ta = np.zeros((table.shape[0], 1), dtype=np.uint16)
     table = np.concatenate((ta, table, ta), axis=1)
-    logger.debug(f"After combination of the same terms: {table.shape[0]}")
 
     # 0 represents the identity symbol. Identity might not present
     # in `primary_ops` but the algorithm still works.
@@ -364,22 +361,38 @@ def _terms_to_table(model: Model, terms: List[Op], const: float):
 
     table = []
     factor_list = []
-    
+
+    primary_ops_eachsite = []
+    primary_ops = []
+
+    index = 0
+
     dummy_table_entry = []
     for b in model.basis:
         if b.multi_dof:
             dof = b.dof[0]
         else:
             dof = b.dof
         op = Op.identity(dof, qn_size=model.qn_size)
-        dummy_table_entry.append(op)
+
+        primary_ops_eachsite.append({op:index})
+        primary_ops.append(op)
+        dummy_table_entry.append(index)
+        index += 1
+
     for op in terms:
         elem_ops, factor = op.split_elementary(model.dof_to_siteidx)
         table_entry = dummy_table_entry.copy()
+        
         for elem_op in elem_ops:
             # it is ensured in `elem_op` every symbol is on the same site
             site_idx = model.dof_to_siteidx[elem_op.dofs[0]]
-            table_entry[site_idx] = elem_op
+            if elem_op not in primary_ops_eachsite[site_idx].keys():
+                primary_ops_eachsite[site_idx][elem_op] = index
+                primary_ops.append(elem_op)
+                index += 1
+            table_entry[site_idx] = primary_ops_eachsite[site_idx][elem_op]
+
         table.append(table_entry)
         factor_list.append(factor)
 
@@ -389,10 +402,19 @@ def _terms_to_table(model: Model, terms: List[Op], const: float):
         factor_list.append(const)
         table.append(table_entry)
 
-    factor_list = np.array(factor_list)
+    factor = np.array(factor_list)
     logger.debug(f"# of operator terms: {len(table)}")
+    
+    # use np.uint32, np.uint16 to save memory
+    max_uint16 = np.iinfo(np.uint16).max
+    assert len(primary_ops) < max_uint16
+    
+    table = np.array(table, dtype=np.uint16)
+    logger.debug(f"Input operator terms: {table.shape[0]}")
+    table, factor = _deduplicate_table(table, factor)
+    logger.debug(f"After combination of the same terms: {table.shape[0]}")
 
-    return table, factor_list
+    return table, primary_ops, factor
 
 
 def _deduplicate_table(table, factor):
@@ -416,48 +438,6 @@ def _deduplicate_table(table, factor):
     return new_table, factor
 
 
-def _transform_table(table, factor):
-    """Transforms the table to integer table and combine duplicate terms."""
-
-    # use np.uint32, np.uint16 to save memory
-    max_uint16 = np.iinfo(np.uint16).max
-
-    # translate the symbolic operator table to an easy to manipulate numpy array
-    table = np.array(table)
-    # unique operators with DoF names taken into consideration
-    # The inclusion of DoF names is necessary for multi-dof basis.
-    unique_op = OrderedDict.fromkeys(table.ravel())
-    # Convert Set(table.ravel()) to List will change the Op order in list, OrderedDict made reproducible
-    unique_op = list(unique_op.keys())
-
-    # check the index of different operators could be represented with np.uint16
-    assert len(unique_op) < max_uint16
-
-    # Construct mapping from easy-to-manipulate integer to actual Op
-    primary_ops = list(unique_op)
-
-    op2idx = dict(zip(unique_op, range(len(unique_op))))
-    new_table = np.vectorize(op2idx.get)(table).astype(np.uint16)
-
-    del unique_op
-
-    if __debug__:
-        qn_size = len(table[0][0].qn)
-        qn_table = np.array([[x.qn for x in ta] for ta in table])
-        factor_table = np.array([[x.factor for x in ta] for ta in table])
-        for idx in range(len(primary_ops)):
-            coord = np.nonzero(new_table == idx)
-            # check that op with the same symbol has the same factor and qn
-            for j in range(qn_size):
-                assert np.unique(qn_table[:, :, j][coord]).size == 1
-            assert np.all(factor_table[coord] == factor_table[coord][0])
-
-        del factor_table, qn_table
-
-    table, factor = _deduplicate_table(new_table, factor)
-
-    return table, factor, primary_ops
-
 
 # translate the numbers into symbolic Matrix Operator
 def compose_symbolic_mo(in_ops, out_ops, primary_ops):

renormalizer/tn/symbolic_ttno.py

@@ -6,7 +6,7 @@
 from renormalizer import Op, Model
 from renormalizer.model.basis import BasisSet
 from renormalizer.tn.treebase import BasisTree
-from renormalizer.mps.symbolic_mpo import _terms_to_table, _transform_table, _construct_symbolic_mpo_one_site, OpTuple
+from renormalizer.mps.symbolic_mpo import _terms_to_table, _construct_symbolic_mpo_one_site, OpTuple
 
 
 logger = logging.getLogger(__name__)
@@ -57,8 +57,7 @@ def construct_symbolic_ttno(tn: BasisTree, terms: List[Op], const: float = 0, al
     basis = list(chain(*[n.basis_sets for n in nodes]))
     model = Model(basis, [])
     qn_size = model.qn_size
-    table, factor = _terms_to_table(model, terms, const)
-    table, factor, primary_ops = _transform_table(table, factor)
+    table, primary_ops, factor = _terms_to_table(model, terms, const)
 
     dummy_in_ops = [[OpTuple([0], qn=np.zeros(qn_size, dtype=int), factor=1)]]
     out_ops: List[List[OpTuple]]