An example to reproduce a bug(maybe) of mkl

mwe.py

import TAT

Tensor = TAT.No.D.Tensor

L1 = 2
L2 = 4
D = 5


def construct_tensor(l1, l2):
    names = ["P0", "P1"]
    edges = [2, 2]
    if l1 != 0:
        names.append("U")
        edges.append(D)
    if l2 != 0:
        names.append("L")
        edges.append(D)
    if l1 != L1 - 1:
        names.append("D")
        edges.append(D)
    if l2 != L2 - 1:
        names.append("R")
        edges.append(D)
    result = Tensor(names, edges).zero()
    rank = result.rank
    block = result.blocks[result.names]
    block[tuple(0 for i in range(rank))] = 1
    block[tuple(1 if i <= 1 else 0 for i in range(rank))] = 1
    return result


lattice = [[construct_tensor(l1, l2) for l2 in range(L2)] for l1 in range(L1)]
environment_h = [[None for l2 in range(L2 - 1)] for l1 in range(L1)]
environment_v = [[None for l2 in range(L2)] for l1 in range(L1 - 1)]

h_terms = [('H', 0, 0), ('H', 0, 2), ('H', 1, 0), ('H', 1, 2), ('H', 0, 1), ('H', 1, 1), ('V', 0, 0), ('V', 0, 3), ('V', 0, 1), ('V', 0, 2)]

SS = Tensor("I0,I1,O0,O1".split(","), [2, 2, 2, 2]).zero()
SS_block = SS.blocks[SS.names].reshape([4, 4])
SS_block[0, 0] = +1 / 4
SS_block[1, 1] = -1 / 4
SS_block[2, 2] = -1 / 4
SS_block[3, 3] = +1 / 4
SS_block[1, 2] = +1 / 2
SS_block[2, 1] = +1 / 2
U = (-0.01 * SS).exponential({("I0", "O0"), ("I1", "O1")}, step=10)


def energy():
    for l1 in range(L1):
        for l2 in range(L2):
            rename_map = {f"P{orbit}": f"P_{l1}_{l2}_{orbit}" for orbit in [0, 1]}
            if l1 != L1 - 1:
                rename_map["D"] = f"D_{l2}"
            this = lattice[l1][l2].edge_rename(rename_map)
            this = try_multiple(this, l1, l2, "L", division=True)
            this = try_multiple(this, l1, l2, "U", division=True)
            if l1 == l2 == 0:
                vector = this
            else:
                contract_pair = set()
                if l2 != 0:
                    contract_pair.add(("R", "L"))
                if l1 != 0:
                    contract_pair.add((f"D_{l2}", "U"))
                vector = vector.contract(this, contract_pair)

    energy = 0
    for h_or_v, l1, l2 in h_terms:
        coord = [(l1, l2), (l1 + 1 if h_or_v == "V" else 0, l2 + 1 if h_or_v == "H" else 0)]
        Hrho = vector.contract(SS.edge_rename({f"O{t}": f"P_{l1}_{l2}_{0}" for t, [l1, l2] in enumerate(coord)}), {(f"P_{l1}_{l2}_{0}", f"I{t}") for t, [l1, l2] in enumerate(coord)})
        trHrho = (Hrho.trace({(f"P_{l1}_{l2}_{0}", f"P_{l1}_{l2}_{1}") for l1 in range(L1) for l2 in range(L2)}))
        energy += float(trHrho)

    return energy / float(vector.trace({(f"P_{l1}_{l2}_{0}", f"P_{l1}_{l2}_{1}") for l1 in range(L1) for l2 in range(L2)})) / (L1 * L2)


def environment(l1, l2, direction, value=None):
    if value is not None:
        if direction == "R":
            environment_h[l1][l2] = value
        elif direction == "L":
            l2 -= 1
            environment_h[l1][l2] = value
        elif direction == "D":
            environment_v[l1][l2] = value
        elif direction == "U":
            l1 -= 1
            environment_v[l1][l2] = value
    else:
        if direction == "R":
            if 0 <= l1 < L1 and 0 <= l2 < L2 - 1:
                return environment_h[l1][l2]
        elif direction == "L":
            l2 -= 1
            if 0 <= l1 < L1 and 0 <= l2 < L2 - 1:
                return environment_h[l1][l2]
        elif direction == "D":
            if 0 <= l1 < L1 - 1 and 0 <= l2 < L2:
                return environment_v[l1][l2]
        elif direction == "U":
            l1 -= 1
            if 0 <= l1 < L1 - 1 and 0 <= l2 < L2:
                return environment_v[l1][l2]
        return None


def try_multiple(tensor, i, j, direction, *, division=False, square_root=False):
    environment_tensor = environment(i, j, direction)
    if environment_tensor is not None:
        if division:
            environment_tensor = environment_tensor.reciprocal()
        if square_root:
            identity = environment_tensor.same_shape().identity({tuple(environment_tensor.names)})
            delta = environment_tensor.sqrt()
            # Delivery former identity and former environment tensor to four direction.
            if direction in ("D", "R"):
                environment_tensor = identity * delta
            else:
                environment_tensor = environment_tensor * delta.reciprocal()
        if direction == "L":
            tensor = tensor.contract(environment_tensor, {("L", "R")})
        if direction == "R":
            tensor = tensor.contract(environment_tensor, {("R", "L")})
        if direction == "U":
            tensor = tensor.contract(environment_tensor, {("U", "D")})
        if direction == "D":
            tensor = tensor.contract(environment_tensor, {("D", "U")})
    return tensor


def single_term_horizontal(i, j):
    left = lattice[i][j]
    right = lattice[i][j + 1]
    right = try_multiple(right, i, j + 1, "L", division=True)
    left_q, left_r = left.qr("r", {"P0", "R"}, "R", "L")
    right_q, right_r = right.qr("r", {"P0", "L"}, "L", "R")
    u, s, v = (
        left_r  #
        .edge_rename({"P0": "P0"})  #
        .contract(right_r.edge_rename({"P0": "P1"}), {("R", "L")})  #
        .contract(U, {("P0", "I0"), ("P1", "I1")})  #
        .svd({"O0", "L"}, "R", "L", "L", "R", D))
    s /= s.norm_2()
    environment(i, j, "R", s)
    u = try_multiple(u, i, j, "R")
    u = u.contract(left_q, {("L", "R")}).edge_rename({"O0": "P0"})
    lattice[i][j] = u
    v = try_multiple(v, i, j + 1, "L")
    v = v.contract(right_q, {("R", "L")}).edge_rename({"O1": "P0"})
    lattice[i][j + 1] = v


def single_term_vertical(i, j):
    up = lattice[i][j]
    down = lattice[i + 1][j]
    down = try_multiple(down, i + 1, j, "U", division=True)
    up_q, up_r = up.qr("r", {"P0", "D"}, "D", "U")
    down_q, down_r = down.qr("r", {"P0", "U"}, "U", "D")
    u, s, v = (
        up_r  #
        .edge_rename({"P0": "P0"})  #
        .contract(down_r.edge_rename({"P0": "P1"}), {("D", "U")})  #
        .contract(U, {("P0", "I0"), ("P1", "I1")})  #
        .svd({"O0", "U"}, "D", "U", "U", "D", D))
    s /= s.norm_2()
    environment(i, j, "D", s)
    u = try_multiple(u, i, j, "D")
    u = u.contract(up_q, {("U", "D")}).edge_rename({"O0": "P0"})
    lattice[i][j] = u
    v = try_multiple(v, i + 1, j, "U")
    v = v.contract(down_q, {("D", "U")}).edge_rename({"O1": "P0"})
    lattice[i + 1][j] = v


for _ in range(3):
    for h_or_v, l1, l2 in h_terms:
        if h_or_v == "H":
            single_term_horizontal(l1, l2)
        else:
            single_term_vertical(l1, l2)
        print(energy())
    print()