Refine MPS optimizer API and clean up PEPS helpers

Author: rezaquant

.gitignore

@@ -13,3 +13,6 @@ store/
 nohup.out
 
 *.egg-info/
+.eggs/
+build/
+dist/

__init__.py

@@ -36,7 +36,7 @@
         reset_default_backends,
         set_default_array_backend,
         set_default_grad_backend,
-        tn_applied,
+        tns_align,
     )
     from .linalg_registrations import reg_complex_svd_jax, reg_complex_svd_torch
     from .debug import plot_sweep_diagnostics, plot_inner_loss, plot_global_loss_trajectory
@@ -84,7 +84,7 @@
     "plot_sweep_diagnostics",
     "plot_inner_loss",
     "plot_global_loss_trajectory",
-    "tn_applied",
+    "tns_align",
     "gen_long_range_swap_path",
     "apply_2dtn_",
     "apply_gates",
@@ -206,14 +206,14 @@ def __getattr__(name):
             "su4": su4,
         }[name]
 
-    if name in ("tn_applied", "product_state_peps"):
+    if name in ("tns_align", "product_state_peps"):
         from .core import (  # pylint: disable=import-outside-toplevel
             product_state_peps,
-            tn_applied,
+            tns_align,
         )
 
         return {
-            "tn_applied": tn_applied,
+            "tns_align": tns_align,
             "product_state_peps": product_state_peps,
         }[name]
 

src/pepsy/__init__.py

@@ -2,6 +2,8 @@
 
 import importlib.util
 import math
+import os
+import tempfile
 import warnings
 from typing import Any
 
@@ -25,12 +27,74 @@
     "build_optimizer",
     "build_compressed_optimizer",
     "tn_fidelity",
+    "tns_align",
+    "product_state_peps",
 ]
 
 _DEFAULT_ARRAY_BACKEND = None
 _DEFAULT_GRAD_BACKEND = None
 
 
+def _default_cache_root():
+    """Return default cache root for pepsy artifacts."""
+    env_cache = os.environ.get("PEPSY_CACHE_DIR")
+    if env_cache:
+        return env_cache
+
+    try:
+        from platformdirs import user_cache_dir  # pylint: disable=import-outside-toplevel
+
+        return user_cache_dir("pepsy")
+    except Exception:  # pragma: no cover - optional dependency fallback
+        return os.path.join(os.path.expanduser("~"), ".cache", "pepsy")
+
+
+def _resolve_cache_directory(directory, subdir):
+    """Resolve cache directory, honoring global disable and env override."""
+    if directory is not None:
+        return _ensure_cache_directory(directory, warn=True)
+
+    disable_cache = os.environ.get("PEPSY_DISABLE_CACHE", "").strip().lower()
+    if disable_cache in {"1", "true", "yes", "on"}:
+        return None
+
+    default_cache = _ensure_cache_directory(os.path.join(_default_cache_root(), subdir), warn=False)
+    if default_cache is not None:
+        return default_cache
+
+    # Fallback for restricted environments where user-cache is not writable.
+    fallback_cache = _ensure_cache_directory(
+        os.path.join(tempfile.gettempdir(), "pepsy-cache", subdir),
+        warn=False,
+    )
+    if fallback_cache is not None:
+        return fallback_cache
+
+    warnings.warn(
+        "No writable cache directory available. Disabling optimizer cache.",
+        RuntimeWarning,
+        stacklevel=2,
+    )
+    return None
+
+
+def _ensure_cache_directory(path, warn=False):
+    """Create cache directory when possible, else return ``None``."""
+    if path is None:
+        return None
+    try:
+        os.makedirs(path, exist_ok=True)
+        return path
+    except OSError:
+        if warn:
+            warnings.warn(
+                f"Cache directory '{path}' is not writable. Disabling optimizer cache.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+        return None
+
+
 def _validate_backend_callable(name, fn):
     if fn is not None and not callable(fn):
         raise TypeError(f"{name} must be callable or None")
@@ -183,17 +247,26 @@ def build_optimizer(
     max_repeats=2**6,
     parallel=True,
     optlib="cmaes",
-    directory="cash/",
+    directory=None,
     hash_method="b",
 ):
-    """Build and return a reusable cotengra contraction optimizer."""
+    """Build and return a reusable cotengra contraction optimizer.
+
+    Parameters
+    ----------
+    directory : str | None, optional
+        Cache directory for optimizer artifacts. If ``None``, defaults to
+        ``$PEPSY_CACHE_DIR/cotengra`` (or OS user-cache dir fallback). Set
+        environment variable ``PEPSY_DISABLE_CACHE=1`` to force ``None``.
+    """
     selected_optlib = optlib
     if selected_optlib == "cmaes" and importlib.util.find_spec("cmaes") is None:
         warnings.warn(
             "Package 'cmaes' not found. Falling back to optlib='random'.",
             RuntimeWarning,
         )
         selected_optlib = "random"
+    cache_dir = _resolve_cache_directory(directory, "cotengra")
     opt = ctg.ReusableHyperOptimizer(
         minimize=f"combo-{int(alpha)}",
         slicing_opts={"target_size": 2**40},
@@ -204,7 +277,7 @@ def build_optimizer(
         optlib=selected_optlib,
         max_time=max_time,
         hash_method=hash_method,
-        directory=directory,
+        directory=cache_dir,
         progbar=progbar,
         on_trial_error="ignore",
     )
@@ -218,14 +291,23 @@ def build_compressed_optimizer(
     max_repeats=2**8,
     max_time="rate:1e8",
 ):
-    """Build and return a reusable cotengra compressed optimizer."""
+    """Build and return a reusable cotengra compressed optimizer.
+
+    Parameters
+    ----------
+    directory : str | None, optional
+        Cache directory for optimizer artifacts. If ``None``, defaults to
+        ``$PEPSY_CACHE_DIR/cotengra-compressed`` (or OS user-cache dir
+        fallback). Set ``PEPSY_DISABLE_CACHE=1`` to force ``None``.
+    """
+    cache_dir = _resolve_cache_directory(directory, "cotengra-compressed")
     copt = ctg.ReusableHyperCompressedOptimizer(
         chi,
         max_repeats=max_repeats,
         minimize="combo-compressed",
         progbar=progbar,
         max_time=max_time,
-        directory=directory,
+        directory=cache_dir,
     )
     return copt
 
@@ -259,7 +341,7 @@ def add_cycle(peps, bond_dim, cylinder=False):
     return peps
 
 
-def tn_applied(p, pepo):
+def tns_align(p, pepo):
     r"""Apply a PEPO operator to a PEPS ket: :math:`\hat{O}|\psi\rangle`.
 
     The PEPO ``k``-indices contract with the PEPS ``k``-indices on join.

src/pepsy/core.py

@@ -55,15 +55,15 @@ def __init__(
         self,
         tn: qtn.TensorNetwork,
         p: Optional[qtn.TensorNetwork] = None,
-        cutoffs: float = 1.e-12,
+        cutoffs: float = 1e-12,
         backend: Optional[str] = None,
         site_tag_id: str = "I{}",
         opt: str = "auto-hq",
         range_int: Optional[Sequence[int]] = None,
         re_tag: bool = False,
         info: Optional[Dict[str, Any]] = None,
         warning: bool = False,
-        inplace=False,
+        inplace: bool = False,
     ):  # pylint: disable=too-many-arguments,too-many-positional-arguments
 
         if p is None:
@@ -79,19 +79,19 @@ def __init__(
 
         self.tn = tn.copy()
 
-
         if site_tag_id:
+            site_ind_id = getattr(self.p, "site_ind_id", None)
             self.p.view_as_(
                 qtn.MatrixProductState,
                 L=self.L,
                 site_tag_id=site_tag_id,
-                site_ind_id=None,
+                site_ind_id=site_ind_id,
                 cyclic=False,
             )
 
         self.site_tag_id = site_tag_id
 
-        # cotengra path finder
+        # Contraction path optimizer spec.
         self.opt = opt
 
         # cutoffs and underlying backend
@@ -101,9 +101,9 @@ def __init__(
         # warnings being printed or not
         self.warning = warning
 
-        # store cost function results
-        self.loss: List[float] = []
-        self.loss_: List[float] = []
+        # Diagnostics collected during sweeps.
+        self.fidelity_trace: List[float] = []
+        self.local_norm_trace: List[float] = []
         self.info: Dict[str, Any] = info or {}
         self.range_int: List[int] = list(range_int) if range_int is not None else []
         if self.range_int:
@@ -114,7 +114,6 @@ def __init__(
                 raise ValueError("range_int must satisfy start < stop.")
 
 
-        # Is there a better solution?
         # Reindex tensor network with random UUIDs for internal indices
         self.tn.reindex_({idx: qtn.rand_uuid() for idx in self.tn.inner_inds()})
 
@@ -125,7 +124,6 @@ def __init__(
         if re_tag:
             self._re_tag()
 
-
     def visual(
         self,
         figsize=(14, 14),
@@ -215,7 +213,6 @@ def _re_tag(self):
                 )
             self._deep_tag()
 
-
     def run(self, n_iter=6, verbose=False):
         """Run basic left-to-right local fitting sweeps.
 
@@ -224,7 +221,7 @@ def run(self, n_iter=6, verbose=False):
         n_iter : int
             Number of complete sweeps.
         verbose : bool
-            If ``True``, append per-sweep fidelity values to ``self.loss``.
+            If ``True``, append per-sweep fidelity values to ``self.fidelity_trace``.
         """
         if self.p is None:
             raise ValueError("Initial state `p` must be provided.")
@@ -250,15 +247,15 @@ def run(self, n_iter=6, verbose=False):
                 f = f.transpose(*psi[site].inds)
 
                 norm_f = (f.H & f).contract(all) ** 0.5
-                self.loss_.append(complex(norm_f).real)
+                self.local_norm_trace.append(complex(norm_f).real)
 
                 # Update tensor data
                 psi[site].modify(data=f.data)
 
             # Compute fidelity if verbose mode is enabled
             if verbose:
                 fidelity = tn_fidelity(self.tn, psi)
-                self.loss.append(ar.do("real", fidelity))
+                self.fidelity_trace.append(ar.do("real", fidelity))
 
     def _build_env_right(self, psi, env_right):
         """
@@ -286,9 +283,6 @@ def _build_env_right(self, psi, env_right):
                 t |= env_right[site_tag_id.format(i + 1)]
                 env_right[site_tag_id.format(i)] = t.contract(all, optimize=opt)
 
-
-
-
     def _right_range(self, psi, env_right, start, stop):
         """
         Build right environments env_right["I{i}"] for i in 0..L-1.
@@ -355,8 +349,6 @@ def _left_range(self, psi, site, count, env_left):
                 t |= env_left[site_tag_id.format(site - 1)]
                 env_left[site_tag_id.format(site)] = t.contract(all, optimize=opt)
 
-
-
     def _update_env_left(self, psi, site: int, env_left):
         """Update left environment incrementally for current site."""
 
@@ -376,8 +368,11 @@ def _update_env_left(self, psi, site: int, env_left):
             t |= env_left[site_tag_id.format(site - 1)]
             env_left[site_tag_id.format(site)] = t.contract(all, optimize=opt)
 
-
-    def run_eff(self, n_iter=6, verbose=False):  # pylint: disable=too-many-branches,too-many-locals,too-many-statements
+    def run_eff(
+        self,
+        n_iter=6,
+        verbose=False,
+    ):  # pylint: disable=too-many-branches,too-many-locals,too-many-statements
         """Run environment-based fitting sweeps with cached left/right blocks.
 
         This method avoids rebuilding full contractions at each site by
@@ -456,7 +451,7 @@ def run_eff(self, n_iter=6, verbose=False):  # pylint: disable=too-many-branches
                     raise TypeError("Unexpected effective tensor type during run_eff.")
 
                 norm_f = (f.H & f).contract(all) ** 0.5
-                self.loss_.append(complex(norm_f).real)
+                self.local_norm_trace.append(complex(norm_f).real)
 
                 # Contract and normalize
                 # Update tensor data
@@ -465,9 +460,7 @@ def run_eff(self, n_iter=6, verbose=False):  # pylint: disable=too-many-branches
             # Compute fidelity if verbose mode is enabled
             if verbose:
                 fidelity = tn_fidelity(self.tn, psi)
-                self.loss.append(ar.do("real", fidelity))
-
-
+                self.fidelity_trace.append(ar.do("real", fidelity))
 
     def run_gate(
         self, n_iter=6, verbose=False
@@ -525,7 +518,6 @@ def run_gate(
                         tn = env_right[site_tag_id.format(site + 1)]
 
                 if 0 < site < L - 1:
-
                     # Boundary consistency: the left and right indices must match between tn and p
                     if count_ == 0:
                         indx = psi.bond(start - 1, start)
@@ -579,10 +571,7 @@ def run_gate(
                     raise TypeError("Unexpected effective tensor type during run_gate.")
 
                 norm_f = (f.H & f).contract(all) ** 0.5
-
-                # norm_f = ar.do("norm", f.data)
-
-                self.loss_.append(complex(norm_f).real)
+                self.local_norm_trace.append(complex(norm_f).real)
 
                 # Contract and normalize
                 # Update tensor data
@@ -591,8 +580,7 @@ def run_gate(
                 if site < stop:
                     psi.left_canonize_site(site, bra=None)
 
-
             # Compute fidelity if verbose mode is enabled
             if verbose:
                 fidelity = tn_fidelity(self.tn, psi)
-                self.loss.append(ar.do("real", fidelity))
+                self.fidelity_trace.append(ar.do("real", fidelity))