chore: 🔧 enforce type checking with ty (fixes #16)

prek.toml

@@ -40,3 +40,10 @@ hooks = [
 repo = "https://github.com/astral-sh/uv-pre-commit"
 rev = "0.8.14"
 hooks = [{ id = "uv-lock" }]
+
+
+[[repos]]
+# Keep in sync with the ty pin in pyproject.toml.
+repo = "https://github.com/astral-sh/ty-pre-commit"
+rev = "v0.0.82"
+hooks = [{ id = "ty" }]

pyproject.toml

@@ -71,6 +71,7 @@ dev = [
   "prek>=0.2.6",
   "ruff==0.16.5",
   "cyclopts>=3.20",
+  "ty==0.0.82",
   { include-group = "test" },
 ]
 interactive = [
@@ -269,25 +270,8 @@ convention = "google"
 required-imports = ["from __future__ import annotations"]
 
 
-[tool.pyright]
-pythonVersion = "3.11"
+[tool.ty.environment]
+python-version = "3.11"
 
-include = ["src/src_method/**/*.py"]
-
-exclude = ["test_*.py", "docs/**", "**/__pycache__"]
-
-typeCheckingMode = "basic"
-useLibraryCodeForTypes = true
-
-reportGeneralTypeIssues = "information"
-reportMissingImports = true
-reportMissingTypeStubs = "warning"
-reportImportCycles = "information"
-reportPrivateUsage = "information"
-reportPrivateImportUsage = "information"
-reportUnnecessaryIsInstance = "information"
-reportUnnecessaryCast = "information"
-reportUnnecessaryComparison = "information"
-reportUnnecessaryContains = "information"
-reportUnnecessaryTypeIgnoreComment = "information"
-reportMatchNotExhaustive = "information"
+[tool.ty.src]
+include = ["src/src_method"]

src/src_method/apply.py

@@ -35,7 +35,7 @@
     from collections.abc import Sequence
     from types import ModuleType
 
-    from numpy.typing import NDArray
+    from numpy.typing import DTypeLike, NDArray
 
 # Set up logger
 setup_logging()
@@ -61,7 +61,7 @@ def apply(
     chi_out: int,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike = np.float64,
     seed: int | None = None,
     device: str = "cpu",
 ) -> list[NDArray]:
@@ -146,7 +146,7 @@ def _src_mpo_mps(
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike = np.float64,
 ) -> list[NDArray]:
     """Computes the compressed product |η> ≈ H|ψ> using the SRC method.
 
@@ -233,7 +233,7 @@ def _src_mpo_mps(
         eta[j] = Q_trunc.reshape(phys_dim, chi_right, rank).transpose(2, 1, 0)
         S = contract(
             "acb,debf,ghf,ceh->adg",
-            eta[j].conj(),
+            eta[j].conj(),  # ty: ignore[unresolved-attribute]
             mpo_arrs[j],
             mps_arrs[j],
             S,
@@ -248,7 +248,7 @@ def _src_mpo_mps(
     logger.debug(LOG_TIME, t_rtl=tms * 1e-9)
     logger.info("SRC MPO-MPS complete.")
 
-    return [to_numpy(t) for t in eta]
+    return [to_numpy(t) for t in eta]  # ty: ignore[invalid-argument-type]
 
 
 def _src_mpo_mpo(
@@ -259,7 +259,7 @@ def _src_mpo_mpo(
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike = np.float64,
 ) -> list[NDArray]:
     """Computes the compressed product H_new ≈ H1 @ H2 using the SRC method.
 
@@ -354,7 +354,7 @@ def _src_mpo_mpo(
         )
         S = contract(
             "abcd,hecg,ifgd,bef->ahi",
-            eta[j].conj(),
+            eta[j].conj(),  # ty: ignore[unresolved-attribute]
             mpo_left_arrs[j],
             mpo_right_arrs[j],
             S,
@@ -369,4 +369,4 @@ def _src_mpo_mpo(
     logger.debug(LOG_TIME, t_rtl=tms * 1e-9)
     logger.info("SRC MPO-MPO complete.")
 
-    return [to_numpy(t) for t in eta]
+    return [to_numpy(t) for t in eta]  # ty: ignore[invalid-argument-type]

src/src_method/compress.py

@@ -35,7 +35,7 @@
     from collections.abc import Sequence
     from types import ModuleType
 
-    from numpy.typing import NDArray
+    from numpy.typing import DTypeLike, NDArray
 
 # Set up logger
 setup_logging()
@@ -60,7 +60,7 @@ def compress(
     chi_out: int,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike = np.float64,
     seed: int | None = None,
     device: str = "cpu",
 ) -> list[NDArray]:
@@ -129,7 +129,7 @@ def _src_mpo(
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike = np.float64,
 ) -> list[NDArray]:
     """Compress an MPO using the SRC method.
 
@@ -212,7 +212,7 @@ def _src_mpo(
         eta[j] = Q_trunc.reshape(chi_right, phys_up, phys_down, rank).transpose(
             3, 0, 1, 2
         )
-        S = contract("abcd,fecd,eb->fa", eta[j].conj(), mpo_arrs[j], S)
+        S = contract("abcd,fecd,eb->fa", eta[j].conj(), mpo_arrs[j], S)  # ty: ignore[unresolved-attribute]
         C[j - 1] = None
         chi_right = rank
 
@@ -223,7 +223,7 @@ def _src_mpo(
     logger.debug(LOG_TIME, t_rtl=tms * 1e-9)
     logger.info("SRC MPO complete.")
 
-    return [to_numpy(t) for t in eta]
+    return [to_numpy(t) for t in eta]  # ty: ignore[invalid-argument-type]
 
 
 def _src_mps(
@@ -233,7 +233,7 @@ def _src_mps(
     xp: ModuleType,
     *,
     cutoff: float = 0.0,
-    dtype: type = np.float64,
+    dtype: DTypeLike = np.float64,
 ) -> list[NDArray]:
     """Compress an MPS using the SRC method.
 
@@ -306,7 +306,7 @@ def _src_mps(
         Q_trunc = truncated_qr(M, cutoff, xp)
         rank = Q_trunc.shape[1]
         eta[j] = Q_trunc.reshape(phys_dim, chi_right, rank).transpose(2, 1, 0)
-        S = contract("acb,deb,ec->da", eta[j].conj(), mps_arrs[j], S)
+        S = contract("acb,deb,ec->da", eta[j].conj(), mps_arrs[j], S)  # ty: ignore[unresolved-attribute]
         C[j - 1] = None
         chi_right = rank
 
@@ -317,4 +317,4 @@ def _src_mps(
     logger.debug(LOG_TIME, t_rtl=tms * 1e-9)
     logger.info("SRC MPS complete.")
 
-    return [to_numpy(t) for t in eta]
+    return [to_numpy(t) for t in eta]  # ty: ignore[invalid-argument-type]

src/src_method/utils/_backend.py

@@ -34,7 +34,7 @@ def get_xp(device: str) -> ModuleType:
     if device == "cpu":
         return np
     if device == "gpu":
-        import cupy  # noqa: PLC0415  (lazy: optional dependency)
+        import cupy  # noqa: PLC0415  (lazy: optional dependency) # ty: ignore[unresolved-import]
 
         return cupy
     msg = f"Unknown device {device!r}; expected 'cpu' or 'gpu'."

src/src_method/utils/linalg.py

@@ -12,7 +12,11 @@
     from numpy.typing import NDArray
 
 
-def truncated_qr(matrix: NDArray, cutoff: float, xp: ModuleType = np) -> NDArray:
+def truncated_qr(
+    matrix: NDArray[np.floating | np.complexfloating],
+    cutoff: float,
+    xp: ModuleType = np,
+) -> NDArray[np.floating | np.complexfloating]:
     """QR with SVD-based rank truncation, returning only the isometry.
 
     Decomposes ``matrix = Q @ R``, then truncates via SVD on R,

src/src_method/utils/logging_config.py

@@ -25,9 +25,8 @@ def setup_logging(default_level: str = "INFO") -> None:
     """
     # Configure Python's standard logging
     level_name = os.environ.get("LOG_LEVEL_SRC", default_level)
-    level = logging.getLevelName(level_name.upper())
     logging.basicConfig(
-        level=level,
+        level=level_name.upper(),
         format="%(message)s",
         stream=sys.stdout,
     )