refactor norm

src/qailo/__init__.py

@@ -2,7 +2,7 @@
 from . import operator as op
 from . import state_vector as sv
 from ._version import version
-from .dispatch import apply, apply_seq, num_qubits, probability, vector
+from .dispatch import apply, apply_seq, norm, num_qubits, probability, vector
 
 __all__ = [
     alg,
@@ -15,6 +15,7 @@
     version,
     apply,
     apply_seq,
+    norm,
     num_qubits,
     probability,
     vector,

src/qailo/dispatch.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from . import mps
 from . import operator as op
 from . import state_vector as sv
@@ -23,6 +25,13 @@ def apply_seq(v, seq):
     return v
 
 
+def norm(v):
+    if sv.is_state_vector(v):
+        return np.linalg.norm(v)
+    elif mps.is_mps(v):
+        return v._norm()
+
+
 def num_qubits(v):
     if sv.is_state_vector(v):
         return sv.num_qubits(v)

src/qailo/mps/__init__.py

@@ -1,7 +1,6 @@
 from .apply import apply, apply_seq
 from .mps_c import canonical_mps
 from .mps_p import projector_mps
-from .norm import norm
 from .product_state import one, product_state, tensor_decomposition, zero
 from .projector import projector
 from .state_vector import state_vector
@@ -12,7 +11,6 @@
     apply_seq,
     canonical_mps,
     projector_mps,
-    norm,
     one,
     product_state,
     tensor_decomposition,

src/qailo/mps/mps_c.py

@@ -34,6 +34,13 @@ def __init__(self, tensors, nkeep=None):
     def _num_qubits(self):
         return len(self.tensors)
 
+    def _norm(self):
+        A = np.identity(1)
+        for t in range(self._num_qubits()):
+            A = np.einsum("ij,jkl->ikl", A, self._tensor(t))
+            A = np.einsum("ijk,ijl->kl", A, self._tensor(t).conj())
+        return np.sqrt(np.trace(A))
+
     def _state_vector(self):
         n = self._num_qubits()
         v = self._tensor(0)

src/qailo/mps/mps_p.py

@@ -38,6 +38,13 @@ def __init__(self, tensors, nkeep=None):
     def _num_qubits(self):
         return len(self.tensors)
 
+    def _norm(self):
+        A = np.identity(1)
+        for t in range(self._num_qubits()):
+            A = np.einsum("ij,jkl->ikl", A, self._tensor(t))
+            A = np.einsum("ijk,ijl->kl", A, self._tensor(t).conj())
+        return np.sqrt(np.trace(A))
+
     def _state_vector(self):
         n = self._num_qubits()
         v = self._tensor(0)

test/mps/test_canonical.py

@@ -16,40 +16,40 @@ def test_canonical():
     tensors.append(np.random.random((d, 2, 1)))
     for mps in [q.mps.canonical_mps, q.mps.projector_mps]:
         m = mps(tensors)
-        norm = q.mps.norm(m)
-        assert q.mps.norm(m) == approx(norm)
+        norm = q.norm(m)
+        assert q.norm(m) == approx(norm)
         assert q.mps.is_canonical(m)
 
         for _ in range(n):
             p = np.random.randint(n)
             m._canonicalize(p)
-            assert q.mps.norm(m) == approx(norm)
+            assert q.norm(m) == approx(norm)
             assert q.mps.is_canonical(m)
             assert q.mps.is_canonical(m)
 
         for _ in range(n):
             p = np.random.randint(n - 1)
             m._canonicalize(p, p + 1)
-            assert q.mps.norm(m) == approx(norm)
+            assert q.norm(m) == approx(norm)
             assert q.mps.is_canonical(m)
 
     v = np.random.random(2**n).reshape((2,) * n + (1,))
     v /= np.linalg.norm(v)
     tensors = q.mps.tensor_decomposition(v, maxdim)
     for mps in [q.mps.canonical_mps, q.mps.projector_mps]:
         m = mps(tensors)
-        norm = q.mps.norm(m)
+        norm = q.norm(m)
 
         for _ in range(n):
             p = np.random.randint(n)
             m._canonicalize(p)
-            assert q.mps.norm(m) == approx(norm)
+            assert q.norm(m) == approx(norm)
             assert q.mps.is_canonical(m)
 
         for _ in range(n):
             p = np.random.randint(n - 1)
             m._canonicalize(p, p + 1)
-            assert q.mps.norm(m) == approx(norm)
+            assert q.norm(m) == approx(norm)
             assert q.mps.is_canonical(m)
 
 

test/mps/test_move.py

@@ -19,7 +19,7 @@ def test_swap():
     for mps in [q.mps.canonical_mps, q.mps.projector_mps]:
         m = mps(tensors)
         q.mps.is_canonical(m)
-        norm = q.mps.norm(m)
+        norm = q.norm(m)
         v = q.sv.vector(q.mps.state_vector(m))
         for _ in range(64):
             s = np.random.randint(n - 1)
@@ -28,7 +28,7 @@ def test_swap():
             _swap_tensors(m, s)
             print(q.sv.vector(q.mps.state_vector(m)))
             q.mps.is_canonical(m)
-            assert q.mps.norm(m) == approx(norm)
+            assert q.norm(m) == approx(norm)
 
         vn = q.sv.vector(q.mps.state_vector(m))
         assert len(v) == len(vn)
@@ -51,7 +51,7 @@ def test_move():
     for mps in [q.mps.canonical_mps, q.mps.projector_mps]:
         m = mps(tensors)
         q.mps.is_canonical(m)
-        norm = q.mps.norm(m)
+        norm = q.norm(m)
         v = q.sv.vector(q.mps.state_vector(m))
 
         for _ in range(16):
@@ -62,7 +62,7 @@ def test_move():
             _move_qubit(m, p, s)
             print(q.sv.vector(q.mps.state_vector(m)))
             q.mps.is_canonical(m)
-            assert q.mps.norm(m) == approx(norm)
+            assert q.norm(m) == approx(norm)
 
         vn = q.sv.vector(q.mps.state_vector(m))
         assert len(v) == len(vn)