Add dedicated FD solver module and public FDSolver API

Author: rezah

src/pepsy/__init__.py

@@ -16,6 +16,7 @@
         boundary_sweeps,
         core,
         fit,
+        ft_solver,
         gate,
         ham,
         gradient_solver,
@@ -103,6 +104,7 @@
     from .optimize_energy import EnergyOptimizer
     from .optimize_mps import MpsOptimizer
     from .optimize_mpo import MpoOptimizer
+    from .gradient_solver import FDSolver
 
 __all__ = [
     "__version__",
@@ -123,6 +125,7 @@
     "register_torch_linalg",
     "reset_default_backends",
     "SweepOptimizer",
+    "FDSolver",
     "EnergyOptimizer",
     "tns_align",
     "measure_obs",
@@ -182,6 +185,7 @@
     "optimize_mps",
     "optimize_mpo",
     "gradient_solver",
+    "ft_solver",
     "ham",
     "boundary_metrics",
     "boundary_states",
@@ -201,6 +205,7 @@ def __getattr__(name):
         "boundary_sweeps",
         "ham",
         "gradient_solver",
+        "ft_solver",
         "optimize_mps",
         "optimize_global",
         "optimize_sweep",
@@ -476,6 +481,11 @@ def __getattr__(name):
 
         return SweepOptimizer
 
+    if name == "FDSolver":
+        from .gradient_solver import FDSolver  # pylint: disable=import-outside-toplevel
+
+        return FDSolver
+
     if name == "EnergyOptimizer":
         from .optimize_energy import EnergyOptimizer  # pylint: disable=import-outside-toplevel
 

src/pepsy/core.py

@@ -773,8 +773,9 @@ def backend_cupy(device=None, dtype=None):
     ----------
     device : int | cupy.cuda.Device | None, optional
         Target CUDA device. If ``None``, use CuPy's current device.
-    dtype : dtype-like | None, optional
-        Target CuPy dtype. If ``None``, infer from input.
+    dtype : dtype-like | torch.dtype | None, optional
+        Target CuPy dtype. If ``None``, infer from input. Torch dtypes are
+        accepted and internally mapped to CuPy-compatible dtypes.
     """
     try:
         import cupy as cp  # pylint: disable=import-outside-toplevel
@@ -790,6 +791,26 @@ def backend_cupy(device=None, dtype=None):
     if isinstance(target_device, int):
         target_device = cp.cuda.Device(target_device)
 
+    if torch is not None and isinstance(dtype, torch.dtype):
+        torch_to_cupy = {
+            torch.complex128: cp.complex128,
+            torch.complex64: cp.complex64,
+            torch.float64: cp.float64,
+            torch.float32: cp.float32,
+            torch.float16: cp.float16,
+            torch.int64: cp.int64,
+            torch.int32: cp.int32,
+            torch.int16: cp.int16,
+            torch.int8: cp.int8,
+            torch.uint8: cp.uint8,
+            torch.bool: cp.bool_,
+        }
+        if dtype not in torch_to_cupy:
+            raise ValueError(
+                f"backend_cupy does not support torch dtype {dtype!r}."
+            )
+        dtype = torch_to_cupy[dtype]
+
     def cast_array(x, device=target_device, dtype=dtype):
         if device is None:
             return cp.asarray(x, dtype=dtype)