add scipy interface

Author: refraction-ray

Diff for: CHANGELOG.md

@@ -12,6 +12,8 @@
 
 - add `state_centric` decorator in `tc.templates.blocks` to transform circuit-to-circuit funtion to state-to-state function
 
+- add `interfaces.scipy_optimize_interface` to transform quantum function into `scipy.optimize.minimize` campatible form
+
 ### Fixed
 
 - avoid error on watch non `tf.Tensor` in tensorflow backend grad method

Diff for: tensorcircuit/interfaces.py

@@ -2,12 +2,12 @@
 interfaces bridging different backends
 """
 
-from typing import Any, Callable, Tuple
+from typing import Any, Callable, Tuple, Optional
 
 import numpy as np
 import torch
 
-from .cons import backend
+from .cons import backend, dtypestr
 from .backends import get_backend  # type: ignore
 
 Tensor = Any
@@ -17,16 +17,10 @@
 
 
 def tensor_to_numpy(t: Tensor) -> Array:
-    from jax import numpy as jnp
-    import tensorflow as tf
-
-    if isinstance(t, torch.Tensor):
-        return t.numpy()
-    if isinstance(t, tf.Tensor) or isinstance(t, tf.Variable):
+    try:
         return t.numpy()
-    if isinstance(t, jnp.ndarray):
+    except AttributeError:
         return np.array(t)
-    return t
 
 
 def general_args_to_numpy(args: Any, same_pytree: bool = True) -> Any:
@@ -133,3 +127,27 @@ def backward(ctx: Any, *grad_y: Any) -> Any:
 
     # currently, memory transparent dlpack in these ML framework has broken support on complex dtypes
     return Fun.apply  # type: ignore
+
+
+def scipy_optimize_interface(
+    fun: Callable[..., Any], shape: Optional[Tuple[int, ...]] = None, jit: bool = True
+) -> Callable[..., Any]:
+    vag = backend.value_and_grad(fun, argnums=0)
+    if jit:
+        vag = backend.jit(vag)
+
+    def scipy_vag(*args: Any, **kws: Any) -> Tuple[Tensor, Tensor]:
+        scipy_args = numpy_args_to_backend(args, dtype=dtypestr)
+        if shape is not None:
+            scipy_args = list(scipy_args)
+            scipy_args[0] = backend.reshape(scipy_args[0], shape)
+            scipy_args = tuple(scipy_args)
+        vs, gs = vag(*scipy_args, **kws)
+        scipy_vs = general_args_to_numpy(vs)
+        gs = backend.reshape(gs, [-1])
+        scipy_gs = general_args_to_numpy(gs)
+        scipy_vs = scipy_vs.astype(np.float64)
+        scipy_gs = scipy_gs.astype(np.float64)
+        return scipy_vs, scipy_gs
+
+    return scipy_vag

Diff for: tests/test_interfaces.py

@@ -2,6 +2,7 @@
 import sys
 import pytest
 from pytest_lazyfixture import lazy_fixture as lf
+from scipy import optimize
 
 thisfile = os.path.abspath(__file__)
 modulepath = os.path.dirname(os.path.dirname(thisfile))
@@ -98,3 +99,29 @@ def f3(x):
     l3.backward()
     pg = param3.grad
     np.testing.assert_allclose(pg, 2 * np.ones([2]).astype(np.complex64), atol=1e-5)
+
+
+@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
+def test_scipy_interface(backend):
+    n = 3
+
+    def f(param):
+        c = tc.Circuit(n)
+        for i in range(n):
+            c.rx(i, theta=param[0, i])
+            c.rz(i, theta=param[1, i])
+        loss = c.expectation(
+            [
+                tc.gates.y(),
+                [
+                    0,
+                ],
+            ]
+        )
+        return tc.backend.real(loss)
+
+    f_scipy = interfaces.scipy_optimize_interface(f, shape=[2, n])
+    r = optimize.minimize(f_scipy, np.zeros([2 * n]), method="L-BFGS-B", jac=True)
+    # L-BFGS-B may has issue with float32
+    # see: https://github.com/scipy/scipy/issues/5832
+    np.testing.assert_allclose(r["fun"], -1.0, atol=1e-5)