Fermion Tensor Network Framework

quimb/tensor/block_interface.py

@@ -0,0 +1,58 @@
+import sys
+from pyblock3.algebra.fermion_symmetry import U11, U1, Z2, Z4, Z22
+from pyblock3.algebra.symmetry import BondInfo
+from pyblock3.algebra.fermion import eye, SparseFermionTensor
+from pyblock3.algebra import fermion_setting as setting
+from pyblock3.algebra import fermion_ops
+
+this = sys.modules[__name__]
+this.DEFAULT_SYMMETRY = "U1"
+this.USE_CPP = True
+this.USE_FERMION = True
+symmetry_map = setting.symmetry_map
+
+def set_symmetry(symmetry):
+    symmetry = symmetry.upper()
+    if symmetry not in symmetry_map:
+        raise KeyError("input symmetry %s not supported"%symmetry)
+    this.DEFAULT_SYMMETRY = symmetry
+    setting.set_symmetry(symmetry)
+
+def set_backend(use_cpp):
+    this.USE_CPP = use_cpp
+    setting.set_flat(use_cpp)
+
+def set_fermion(use_fermion):
+    this.USE_FERMION = use_fermion
+    setting.set_fermion(use_fermion)
+
+def set_options(**kwargs):
+    symmetry = kwargs.pop("symmetry", this.DEFAULT_SYMMETRY)
+    use_fermion = kwargs.pop("fermion", this.USE_FERMION)
+    use_cpp = kwargs.pop("use_cpp", this.USE_CPP)
+    set_symmetry(symmetry)
+    set_fermion(use_fermion)
+    set_backend(use_cpp)
+
+def dispatch_settings(*keys):
+    dict = {"symmetry": "DEFAULT_SYMMETRY",
+            "fermion": "USE_FERMION",
+            "use_cpp": "USE_CPP"}
+    _settings = []
+    for ikey in keys:
+        if ikey not in dict:
+            raise KeyError("%s not a valid backend setting"%ikey)
+        _settings.append(getattr(this, dict[ikey]))
+    if len(_settings) == 1:
+        _settings = _settings[0]
+    return _settings
+
+def get_symmetry():
+    return setting.symmetry_map[this.DEFAULT_SYMMETRY]
+
+to_exponential = fermion_ops.get_exponential
+H1 = fermion_ops.H1
+Hubbard = fermion_ops.Hubbard
+onsite_U = fermion_ops.onsite_U
+measure_SZ = fermion_ops.measure_SZ
+ParticleNumber = fermion_ops.ParticleNumber

quimb/tensor/block_tools.py

@@ -0,0 +1,63 @@
+from quimb.tensor import block_interface as bitf
+import numpy as np
+
+def apply(T, func):
+    use_cpp = bitf.dispatch_settings("use_cpp")
+    if use_cpp:
+        new_T = T.copy()
+        new_T.data = func(new_T.data)
+    else:
+        new_T = T.copy()
+        for iblk in new_T:
+            iblk[:] = func(iblk[:])
+    return new_T
+
+def sqrt(T):
+    _sqrt = lambda x : x**.5
+    return apply(T, _sqrt)
+
+def inv_with_smudge(T, cutoff=1e-10, gauge_smudge=1e-6):
+    def _inv_with_smudge(arr):
+        new_arr = np.zeros(arr.shape, dtype=arr.dtype)
+        ind = abs(arr) > cutoff
+        new_arr[ind] = (arr[ind] + gauge_smudge) ** -1
+        return new_arr
+    return apply(T, _inv_with_smudge)
+
+def add_with_smudge(T, cutoff=1e-10, gauge_smudge=1e-6):
+    def _add_with_smudge(arr):
+        ind = abs(arr) > cutoff
+        arr[ind] += gauge_smudge
+        return arr
+    return apply(T, _add_with_smudge)
+
+def get_smudge_balance(T1, T2, ix, smudge):
+    flat = bitf.dispatch_settings("use_cpp")
+    if flat:
+        t1, t2 = T1.data.to_sparse(), T2.data.to_sparse()
+    else:
+        t1, t2 = T1.data, T2.data
+    sign1 = t1.pattern[T1.inds.index(ix)]
+    sign2 = t2.pattern[T2.inds.index(ix)]
+    s1_pattern = {"+":"-+", "-":"+-"}[sign1]
+    s2_pattern = {"-":"-+", "+":"+-"}[sign2]
+
+    inv = (sign1 == sign2)
+    block_cls = t1.blocks[0].__class__
+    block_dict = {}
+    for iblk1 in t1:
+        q0 = iblk1.q_labels[0]
+        block_dict[q0] = np.diag(np.asarray(iblk1)) + smudge
+    for iblk2 in t2:
+        q0 = -iblk2.q_labels[0] if inv else iblk2.q_labels[0]
+        if q0 not in block_dict: continue
+        block_dict[q0] = block_dict[q0] / (np.diag(np.asarray(iblk2)) + smudge)
+
+    s1 = [block_cls(reduced=np.diag(s**-0.25), q_labels=(qlab,)*2) for qlab, s in block_dict.items()]
+    s2 = [block_cls(reduced=np.diag(s** 0.25), q_labels=(qlab,)*2) for qlab, s in block_dict.items()]
+    s1 = t1.__class__(blocks=s1, pattern=s1_pattern)
+    s2 = t2.__class__(blocks=s2, pattern=s2_pattern)
+    if flat:
+        s1 = s1.to_flat()
+        s2 = s2.to_flat()
+    return s1, s2