add jax function load save

Author: refraction-ray

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ## Unreleased
 
+### Added
+
+- Add jax jitted function load/save utilities in experimental module
+
+- Add `circuit.to_openqasm_file` function for compatibility of qiskit>1
+
 ## 1.0.2
 
 ### Added

docs/source/advance.rst

@@ -5,17 +5,65 @@ Advanced Usage
 MPS Simulator
 ----------------
 
-(Still experimental support)
-
-Very simple, we provide the same set of API for ``MPSCircuit`` as ``Circuit``, 
+Very straightforward to use, we provide the same set of API for ``MPSCircuit`` as ``Circuit``, 
 the only new line is to set the bond dimension for the new simulator.
 
 .. code-block:: python
 
     c = tc.MPSCircuit(n)
     c.set_split_rules({"max_singular_values": 50})
 
-The larger bond dimension we set, the better approximation ratio (of course the more computational cost we pay)
+The larger bond dimension we set, the better approximation ratio (of course the more computational cost we pay).
+
+
+Stacked gates
+----------------
+
+Stacked gates is a simple grammar sugar to make constructing the circuit easily when multiple gate of the same type are applied on the different qubits, namely, the index for gate function can accept list of ints instead of one integer.
+
+.. code-block:: python
+
+    >>> import tensorcircuit as tc
+    >>> c = tc.Circuit(4)
+    >>> c.h(range(3))
+    >>> c.draw()
+         ┌───┐
+    q_0: ┤ H ├
+         ├───┤
+    q_1: ┤ H ├
+         ├───┤
+    q_2: ┤ H ├
+         └───┘
+    q_3: ─────
+
+
+    >>> c = tc.Circuit(4)
+    >>> c.cnot([0, 1], [2, 3])
+    >>> c.draw()
+
+    q_0: ──■───────
+           │
+    q_1: ──┼────■──
+         ┌─┴─┐  │
+    q_2: ┤ X ├──┼──
+         └───┘┌─┴─┐
+    q_3: ─────┤ X ├
+              └───┘
+
+    >>> c = tc.Circuit(4)
+    >>> c.rx(range(4), theta=tc.backend.convert_to_tensor([0.1, 0.2, 0.3, 0.4]))
+    >>> c.draw()
+         ┌─────────┐
+    q_0: ┤ Rx(0.1) ├
+         ├─────────┤
+    q_1: ┤ Rx(0.2) ├
+         ├─────────┤
+    q_2: ┤ Rx(0.3) ├
+         ├─────────┤
+    q_3: ┤ Rx(0.4) ├
+         └─────────┘
+
+
 
 Split Two-qubit Gates
 -------------------------
@@ -45,15 +93,92 @@ The two-qubit gates applied on the circuit can be decomposed via SVD, which may
 
 Note ``max_singular_values`` must be specified to make the whole procedure static and thus jittable.
 
+Analog circuit simulation
+-----------------------------
+
+TensorCircuit-NG support digital-analog hybrid simulation (say cases in Rydberg atom arrays), where the analog part is simulated by the neural differential equation solver given the API to specify a time dependent Hamiltonian.
+The simulation is still differentiable and jittable. Only jax backend is supported for analog simulation as the neural ode engine is built on top of jax. 
+This utility is super helpful for optimizing quantum control or investigating digital-analog hybrid variational quantum schemes.
+We support two modes of analog simulation, where :py:meth:`tensorcircuit.experimentaql.evol_global` evolve the state via a Hamiltonian define on the whole system, and :py:meth:`tensorcircuit.experimentaql.evol_local` evolve the state via a Hamiltonian define on a local subsystem.
+
+.. Note::
+
+    ``evol_global`` use sparse Hamiltonian while ``evol_local`` use dense Hamiltonian.
+
+
+.. code-block:: python
+
+    # in this demo, we build a jittable and differentiable simulation function `hybrid_evol` 
+    # with both digital gates and local/global analog Hamiltonian evolutions
+
+    import optax
+    import tensorcircuit as tc
+    from tensorcircuit.experimental import evol_global, evol_local
+
+    K = tc.set_backend("jax")
+
+
+    def h_fun(t, b):
+        return b * tc.gates.x().tensor
+
+
+    hy = tc.quantum.PauliStringSum2COO([[2, 0]])
+
+
+    def h_fun2(t, b):
+        return b[2] * K.cos(b[0] * t + b[1]) * hy
+
+
+    @K.jit
+    @K.value_and_grad
+    def hybrid_evol(params):
+        c = tc.Circuit(2)
+        c.x([0, 1])
+        c = evol_local(c, [1], h_fun, 1.0, params[0])
+        c.cx(1, 0)
+        c.h(0)
+        c = evol_global(c, h_fun2, 1.0, params[1:])
+        return K.real(c.expectation_ps(z=[0, 1]))
+
+
+    b = K.implicit_randn([4])
+    v, gs = hybrid_evol(b)
+
+
 
 Jitted Function Save/Load
 -----------------------------
 
 To reuse the jitted function, we can save it on the disk via support from the TensorFlow `SavedModel <https://www.tensorflow.org/guide/saved_model>`_. That is to say, only jitted quantum function on the TensorFlow backend can be saved on the disk. 
 
+We wrap the tf-backend `SavedModel` as very easy-to-use function :py:meth:`tensorcircuit.keras.save_func` and :py:meth:`tensorcircuit.keras.load_func`.
+
 For the JAX-backend quantum function, one can first transform them into the tf-backend function via JAX experimental support: `jax2tf <https://github.com/google/jax/tree/main/jax/experimental/jax2tf>`_.
 
-We wrap the tf-backend `SavedModel` as very easy-to-use function :py:meth:`tensorcircuit.keras.save_func` and :py:meth:`tensorcircuit.keras.load_func`.
+**Updates**: jax now also support jitted function save/load via ``export`` module, see `jax documentation <https://jax.readthedocs.io/en/latest/export/export.html>_`.
+
+We wrape the jax function export capability in ``experimental`` module and can be used as follows
+
+.. code-block:: python
+
+    from tensorcircuit import experimental
+
+    K = tc.set_backend("jax")
+
+    @K.jit
+    def f(weights):
+        c = tc.Circuit(3)
+        c.rx(range(3), theta=weights)
+        return K.real(c.expectation_ps(z=[0]))
+
+    print(f(K.ones([3])))
+
+    experimental.jax_jitted_function_save("temp.bin", f, K.ones([3]))
+
+    f_load = tc.experimental.jax_jitted_function_load("temp.bin")
+    f_load(K.ones([3]))
+
+
 
 Parameterized Measurements
 -----------------------------

docs/source/quickstart.rst

@@ -29,10 +29,8 @@ For more details on docker setup, please refer to `docker readme <https://github
 
 - For Windows, due to the lack of support for Jax, we recommend to use docker or WSL, please refer to `TC via windows docker <contribs/development_windows.html>`_ or `TC via WSL <contribs/development_wsl2.html>`_.
 
-- For MacOS, please refer to `TC on Mac <contribs/development_Mac.html>`_.
-
-Overall, the installation of TensorCircuit is simple, since it is purely in Python and hence very portable. 
-As long as the users can take care of the installation of ML frameworks on the corresponding system, TensorCircuit will work as expected.
+Overall, the installation of TensorCircuit-NG is simple, since it is purely in Python and hence very portable. 
+As long as the users can take care of the installation of ML frameworks on the corresponding system, TensorCircuit-NG will work as expected.
 
 To debug the installation issue or report bugs, please check the environment information by ``tc.about()``.
 

tensorcircuit/abstractcircuit.py

@@ -792,6 +792,16 @@ def to_openqasm(self, **kws: Any) -> str:
             qasm_str = dumps(qc)  # type: ignore
         return qasm_str  # type: ignore
 
+    def to_openqasm_file(self, file: str, **kws: Any) -> None:
+        """
+        save the circuit to openqasm file
+
+        :param file: the file path to save the circuit
+        :type file: str
+        """
+        with open(file, "w") as f:
+            f.write(self.to_openqasm(**kws))
+
     @classmethod
     def from_openqasm(
         cls,

tensorcircuit/experimental.py

@@ -350,10 +350,6 @@ def grad_f(*args: Any, **kws: Any) -> Any:
     return grad_f
 
 
-# TODO(@refraction-ray): add SPSA gradient wrapper similar to parameter shift
-# -- using noisyopt package instead
-
-
 def finite_difference_differentiator(
     f: Callable[..., Any],
     argnums: Tuple[int, ...] = (0,),
@@ -456,7 +452,7 @@ def evol_local(
 
     :param c: _description_
     :type c: Circuit
-    :param index: _description_
+    :param index: qubit sites to evolve
     :type index: Sequence[int]
     :param h_fun: h_fun should return a dense Hamiltonian matrix
         with input arguments time and *args
@@ -525,3 +521,48 @@ def f(y: Tensor, t: Tensor, *args: Any) -> Tensor:
     ts = backend.cast(ts, dtype=rdtypestr)
     s1 = odeint(f, s, ts, *args, **solver_kws)
     return type(c)(n, inputs=s1[-1])
+
+
+def jax_jitted_function_save(filename: str, f: Callable[..., Any], *args: Any) -> None:
+    """
+    save a jitted jax function as a file
+
+    :param filename: _description_
+    :type filename: str
+    :param f: the jitted function
+    :type f: Callable[..., Any]
+    :param args: example function arguments for ``f``
+    """
+
+    from jax import export
+
+    f_export = export.export(f)(*args)  # type: ignore
+    barray = f_export.serialize()
+
+    with open(filename, "wb") as file:
+        file.write(barray)
+
+
+jax_func_save = jax_jitted_function_save
+
+
+def jax_jitted_function_load(filename: str) -> Callable[..., Any]:
+    """
+    load a jitted function from file
+
+    :param filename: _description_
+    :type filename: str
+    :return: the loaded function
+    :rtype: _type_
+    """
+    from jax import export
+
+    with open(filename, "rb") as f:
+        barray = f.read()
+
+    f_load = export.deserialize(barray)  # type: ignore
+
+    return f_load.call
+
+
+jax_func_load = jax_jitted_function_load

tensorcircuit/fgs.py

@@ -176,6 +176,14 @@ def get_cmatrix(self, now_i: bool = True, now_j: bool = True) -> Tensor:
             return cmatrix
 
     def get_reduced_cmatrix(self, subsystems_to_trace_out: List[int]) -> Tensor:
+        """
+        get reduced correlation matrix by tracing out subsystems
+
+        :param subsystems_to_trace_out: list of sites to be traced out
+        :type subsystems_to_trace_out: List[int]
+        :return: reduced density matrix
+        :rtype: Tensor
+        """
         m = self.get_cmatrix()
         if subsystems_to_trace_out is None:
             subsystems_to_trace_out = []

tests/test_circuit.py

@@ -1497,7 +1497,7 @@ def test_gate_count():
     # {'x': 1, 'h': 2, 'rx': 1, 'multicontrol': 1, 'toffoli': 3}
 
 
-def test_to_openqasm():
+def test_to_openqasm(tmp_path):
     c = tc.Circuit(3)
     c.H(0)
     c.rz(2, theta=0.2)
@@ -1511,8 +1511,8 @@ def test_to_openqasm():
     c1 = tc.Circuit.from_openqasm(s)
     print(c1.draw())
     np.testing.assert_allclose(c.state(), c1.state())
-    c.to_openqasm(filename="test.qasm")
-    c2 = tc.Circuit.from_openqasm_file("test.qasm")
+    c.to_openqasm_file(os.path.join(tmp_path, "test.qasm"))
+    c2 = tc.Circuit.from_openqasm_file(os.path.join(tmp_path, "test.qasm"))
     np.testing.assert_allclose(c.state(), c2.state())
 
 

tests/test_miscs.py

@@ -254,3 +254,24 @@ def h_square_sparse(t, b):
 def test_energy_baseline():
     print(TFIM1Denergy(10))
     print(Heisenberg1Denergy(10))
+
+
+def test_jax_function_load(jaxb, tmp_path):
+    K = tc.backend
+
+    @K.jit
+    def f(weights):
+        c = tc.Circuit(3)
+        c.rx(range(3), theta=weights)
+        return K.real(c.expectation_ps(z=[0]))
+
+    print(f(K.ones([3])))
+
+    experimental.jax_jitted_function_save(
+        os.path.join(tmp_path, "temp.bin"), f, K.ones([3])
+    )
+
+    f_load = tc.experimental.jax_jitted_function_load(
+        os.path.join(tmp_path, "temp.bin")
+    )
+    np.testing.assert_allclose(f_load(K.ones([3])), 0.5403, atol=1e-4)