Skip to content

Commit

Permalink
TN: add psi.normalize_simple
Browse files Browse the repository at this point in the history
  • Loading branch information
jcmgray committed Jul 22, 2024
1 parent a3950bc commit aeca738
Show file tree
Hide file tree
Showing 2 changed files with 43 additions and 1 deletion.
26 changes: 25 additions & 1 deletion quimb/tensor/tensor_arbgeom.py
Original file line number Diff line number Diff line change
Expand Up @@ -631,6 +631,30 @@ def __iadd__(self, other):
def __isub__(self, other):
return tensor_network_ag_sum(self, other, negate=True, inplace=True)

def normalize_simple(self, gauges, **contract_opts):
"""Normalize this network using simple local gauges. After calling
this, any tree-like sub network gauged with ``gauges`` will have
2-norm 1. Inplace operation on both the tensor network and ``gauges``.
Parameters
----------
gauges : dict[str, array_like]
The gauges to normalize with.
"""
# normalize gauges
for ix, g in gauges.items():
gauges[ix] = g / do("linalg.norm", g)

# normalize sites
for site in self.sites:
tn_site = self.select(site)
tn_site_gauged = tn_site.copy()
tn_site_gauged.gauge_simple_insert(gauges)
lnorm = (tn_site_gauged.H | tn_site_gauged).contract(
all, **contract_opts
) ** 0.5
tn_site /= lnorm


def gauge_product_boundary_vector(
tn,
Expand Down Expand Up @@ -894,7 +918,7 @@ def gate_with_op_lazy(self, A, transpose=False, inplace=False, **kwargs):
which_A="upper" if transpose else "lower",
contract=False,
inplace=inplace,
**kwargs
**kwargs,
)

gate_with_op_lazy_ = functools.partialmethod(
Expand Down
18 changes: 18 additions & 0 deletions tests/test_tensor/test_tensor_arbgeom.py
Original file line number Diff line number Diff line change
Expand Up @@ -116,3 +116,21 @@ def test_gate_sandwich_with_op():
y = A.to_dense() @ B.to_dense() @ A.to_dense().conj().T
B.gate_sandwich_with_op_lazy_(A)
assert_allclose(B.to_dense(), y)


def test_normalize_simple():
psi = qtn.PEPS.rand(3, 3, 2, dtype=complex)
gauges = {}
psi.gauge_all_simple_(100, 5e-6, gauges=gauges)
psi.normalize_simple(gauges)

for where in [
[(0, 0)],
[(1, 1), (1, 2)],
[(0, 0), (0, 1), (0, 2), (1, 2), (2, 2), (2, 1)],
]:
tags = [psi.site_tag(w) for w in where]
k = psi.select_any(tags, virtual=False)
k.gauge_simple_insert(gauges)

assert k.H @ k == pytest.approx(1.0)

0 comments on commit aeca738

Please sign in to comment.