mps optimization and gate API cleanup

Author: rezaquant

__init__.py

@@ -0,0 +1 @@
+from .src.pepsy.optimize_mps import MpsOptimizer

src/pepsy/__init__.py

@@ -44,12 +44,20 @@
     from .gate import (
         apply_2dtn_,
         apply_gates,
-        canonize_mps,
         gate_1d,
         gen_long_range_swap_path,
+        rx,
+        ry,
+        rz,
+        rxx,
+        ryy,
+        rzz,
+        u3,
+        su4,
     )
     from .optimize_global import GlobalOptimizer
     from .optimize_sweep import PEPSSweepOptimizer, SweepResult
+    from .optimize_mps import MpsOptimizer
 
 __all__ = [
     "__version__",
@@ -81,7 +89,14 @@
     "apply_2dtn_",
     "apply_gates",
     "gate_1d",
-    "canonize_mps",
+    "rx",
+    "ry",
+    "rz",
+    "rxx",
+    "ryy",
+    "rzz",
+    "u3",
+    "su4",
     "product_state_peps",
     "optimize_global",
     "optimize_sweep",
@@ -93,6 +108,7 @@
     "core",
     "dmrg_fit",
     "debug",
+    "MpsOptimizer",
 ]
 
 
@@ -105,6 +121,7 @@ def __getattr__(name):
         "debug",
         "gate",
         "gradient_solver",
+        "optimize_mps",
         "optimize_global",
         "optimize_sweep",
         "core",
@@ -150,22 +167,43 @@ def __getattr__(name):
         "apply_2dtn_",
         "apply_gates",
         "gate_1d",
-        "canonize_mps",
+        "rx",
+        "ry",
+        "rz",
+        "rxx",
+        "ryy",
+        "rzz",
+        "u3",
+        "su4",
     ):
         from .gate import (  # pylint: disable=import-outside-toplevel
             apply_2dtn_,
             apply_gates,
-            canonize_mps,
             gate_1d,
             gen_long_range_swap_path,
+            rx,
+            ry,
+            rz,
+            rxx,
+            ryy,
+            rzz,
+            u3,
+            su4,
         )
 
         return {
             "gen_long_range_swap_path": gen_long_range_swap_path,
             "apply_2dtn_": apply_2dtn_,
             "apply_gates": apply_gates,
             "gate_1d": gate_1d,
-            "canonize_mps": canonize_mps,
+            "rx": rx,
+            "ry": ry,
+            "rz": rz,
+            "rxx": rxx,
+            "ryy": ryy,
+            "rzz": rzz,
+            "u3": u3,
+            "su4": su4,
         }[name]
 
     if name in ("tn_applied", "product_state_peps"):
@@ -256,4 +294,9 @@ def __getattr__(name):
             "SweepResult": SweepResult,
         }[name]
 
+    if name == "MpsOptimizer":
+        from .optimize_mps import MpsOptimizer  # pylint: disable=import-outside-toplevel
+
+        return MpsOptimizer
+
     raise AttributeError(f"module 'pepsy' has no attribute {name!r}")

src/pepsy/boundary_norm.py

@@ -29,17 +29,6 @@ class BoundaryContractResult:
     max_separation: int
 
 
-def _validate_tensor_network_tags(p):
-    """Ensure PEPS lattice/site tags are present for shape inference."""
-    validate_tensor_network_tags(p)
-
-
-def _normalize_retag_for_direction(direction, re_tag):
-    """Normalize ``re_tag`` flag for direction-specific calls."""
-    _ = direction
-    return bool(re_tag)
-
-
 def _warn_nonstandard_physical_outer_inds(tn, role):
     """Warn when outer physical indices don't match ``k<int>[,<int>...]`` or ``b<int>[,<int>...]``."""
     bad = [
@@ -103,7 +92,7 @@ def prepare_boundary_inputs(
     if ket is None:
         raise ValueError("Provide ket.")
 
-    _validate_tensor_network_tags(ket)
+    validate_tensor_network_tags(ket)
 
     ket_tagged = ket
     auto_bra = bra is None
@@ -196,7 +185,7 @@ def ContractBoundary(
     if not isinstance(mps_boundaries, dict):
         raise TypeError("mps_boundaries must be a dictionary of boundary states.")
 
-    re_tag = _normalize_retag_for_direction(direction, re_tag)
+    re_tag = bool(re_tag)
     norm_tagged = norm.copy()
 
     comp_bdy = CompBdy(

src/pepsy/boundary_states.py

@@ -113,8 +113,6 @@ def __init__(
         default_array_backend = get_default_array_backend()
         if default_array_backend is not None:
             self.array_backend = default_array_backend
-        # Backward-compatible alias used by older tests/callers.
-        self.to_backend = self.array_backend
 
         self.numpy_backend = make_numpy_array_caster(dtype=dtype_name)
 

src/pepsy/boundary_sweeps.py

@@ -10,10 +10,7 @@
 from .core import build_optimizer, tn_fidelity
 from .dmrg_fit import FIT
 
-__all__ = ["BdyMPS", "CompBdy", "tn_fidelity", "build_optimizer", "opt_"]
-
-# Backward-compatible alias.
-opt_ = build_optimizer
+__all__ = ["BdyMPS", "CompBdy", "tn_fidelity", "build_optimizer"]
 
 
 @dataclass(frozen=True)
@@ -97,11 +94,6 @@ def __init__(
         self.Lx = 1 + max_x  # pylint: disable=invalid-name
         self._update_separation()
 
-    @property
-    def fidelity(self):
-        """Alias for ``self.fidel``."""
-        return self.fidel
-
     def _reset_fidelity_history(self):
         """Reset stored fidelity values for a fresh public call."""
         self.fidel = []

src/pepsy/dmrg_fit.py

@@ -18,14 +18,10 @@
 __all__ = [
     "FIT",
     "build_optimizer",
-    "opt_",
     "tn_fidelity",
     "internal_inds",
 ]
 
-# Backward-compatible alias.
-opt_ = build_optimizer
-
 
 def internal_inds(psi):
     """Return all internal (non-open) indices of ``psi``."""

src/pepsy/gate.py

@@ -15,7 +15,14 @@
 __all__ = [
     "apply_gates",
     "gate_1d",
-    "canonize_mps",
+    "rx",
+    "ry",
+    "rz",
+    "rxx",
+    "ryy",
+    "rzz",
+    "u3",
+    "su4",
 ]
 
 
@@ -63,6 +70,98 @@ def _to_torch(x, dtype=target_dtype, device=target_device):
     return None
 
 
+def rx(theta):
+    """Return a one-qubit RX gate for angle ``theta``.
+
+    Parameters
+    ----------
+    theta : float
+        Rotation angle.
+    """
+    return qtn.circuit.rx_gate_param_gen([theta])
+
+
+def ry(theta):
+    """Return a one-qubit RY gate for angle ``theta``.
+
+    Parameters
+    ----------
+    theta : float
+        Rotation angle.
+    """
+    return qtn.circuit.ry_gate_param_gen([theta])
+
+
+def rz(theta):
+    """Return a one-qubit RZ gate for angle ``theta``.
+
+    Parameters
+    ----------
+    theta : float
+        Rotation angle.
+    """
+    return qtn.circuit.rz_gate_param_gen([theta])
+
+
+def rzz(theta):
+    """Return a two-qubit RZZ gate for angle ``theta``.
+
+    Parameters
+    ----------
+    theta : float
+        Rotation angle.
+    """
+    return qtn.circuit.rzz_param_gen([theta])
+
+
+def rxx(theta):
+    """Return a two-qubit RXX gate for angle ``theta``.
+
+    Parameters
+    ----------
+    theta : float
+        Rotation angle.
+    """
+    return qtn.circuit.rxx_param_gen([theta])
+
+
+def ryy(theta):
+    """Return a two-qubit RYY gate for angle ``theta``.
+
+    Parameters
+    ----------
+    theta : float
+        Rotation angle.
+    """
+    return qtn.circuit.ryy_param_gen([theta])
+
+
+def su4(params):
+    """Return a two-qubit SU(4) gate from 15 parameters.
+
+    Parameters
+    ----------
+    params : sequence
+        Sequence of exactly 15 parameters.
+    """
+    if len(params) != 15:
+        raise ValueError("su4 expects exactly 15 parameters.")
+    return qtn.circuit.su4_gate_param_gen(params)
+
+
+def u3(params):
+    """Return a one-qubit U3 gate from 3 parameters.
+
+    Parameters
+    ----------
+    params : sequence
+        Sequence of exactly 3 parameters.
+    """
+    if len(params) != 3:
+        raise ValueError("u3 expects exactly 3 parameters.")
+    return qtn.circuit.u3_gate_param_gen(params)
+
+
 def gen_long_range_swap_path(  # pylint: disable=too-many-branches,too-many-locals,too-many-statements
     ij_a, ij_b, sequence=None
 ):
@@ -427,15 +526,6 @@ def apply_gates(  # pylint: disable=too-many-arguments,too-many-positional-argum
     return peps
 
 
-def canonize_mps(p, where, cur_orthog):
-    xmin, xmax = sorted(where)
-    p.canonize([xmin, xmax], cur_orthog=cur_orthog, 
-               #info=info_c
-              )
-    # update cur_orthog in place (preserving reference)
-    cur_orthog[:] = [xmin, xmax]
-
-
 def gate_1d(tn, where, G, ind_id="k{}", site_tags="I{}",
             cutoff=1.e-12, contract='split-gate', 
             inplace=False):

src/pepsy/optimize_global.py

@@ -61,6 +61,17 @@ def _merge_opts(base, extra):
             merged.update(dict(extra))
         return merged
 
+    @staticmethod
+    def _normalize_optimizer_name(optimizer):
+        """Normalize common optimizer aliases for ``qtn.TNOptimizer``."""
+        if not isinstance(optimizer, str):
+            return optimizer
+
+        key = optimizer.strip().lower().replace("_", "-")
+        if key in {"lbfgs", "l-bfgs", "lbfgsb", "l-bfgs-b"}:
+            return "L-BFGS-B"
+        return optimizer
+
     @classmethod
     def _pick_known_keys(cls, options, allowed_keys, *, warn_unknown=True):
         incoming = dict(options or {})
@@ -123,7 +134,7 @@ def _norm_peps(  # pylint: disable=too-many-arguments,too-many-positional-argume
             sequence = ["xmin", "xmax", "ymin", "ymax"]
         if (mode == "hyper") and (copt is None):
             warnings.warn(
-                "mode='hyper' requested but copt is None; provide copt_() for stable behavior.",
+                "mode='hyper' requested but copt is None; provide a compressed optimizer via `copt`.",
                 RuntimeWarning,
                 stacklevel=2,
             )
@@ -203,7 +214,7 @@ def _loss_peps(  # pylint: disable=too-many-arguments,too-many-positional-argume
         """Compute overlap-based loss between trainable and target PEPS."""
         if (mode == "hyper") and (copt is None):
             warnings.warn(
-                "mode='hyper' requested but copt is None; provide copt_() for stable behavior.",
+                "mode='hyper' requested but copt is None; provide a compressed optimizer via `copt`.",
                 RuntimeWarning,
                 stacklevel=2,
             )
@@ -395,6 +406,7 @@ def make_tn_optimizer(
     ):
         """Construct a configured :class:`quimb.tensor.TNOptimizer`."""
         merged_loss_kwargs = self._merge_opts(self.loss_kwargs, loss_kwargs)
+        optimizer = self._normalize_optimizer_name(optimizer)
 
         constants = {}
         if self.peps_target is not None: