add more docs

Author: refraction-ray

docs/source/advance.rst

@@ -220,6 +220,135 @@ Please refer to :py:meth:`tensorcircuit.templates.measurements.sparse_expectatio
 For different representations to evaluate Hamiltonian expectation in tensorcircuit, please refer to :doc:`tutorials/tfim_vqe_diffreph`.
 
 
+Hamiltonian Matrix Building
+----------------------------
+
+TensorCircuit-NG provides multiple ways to build Hamiltonian matrices, especially for sparse Hamiltonians constructed from Pauli strings. This is crucial for quantum many-body physics simulations and variational quantum algorithms.
+
+**Pauli String Based Construction:**
+
+The most flexible way to build Hamiltonians is through Pauli strings:
+
+.. code-block:: python
+
+    import tensorcircuit as tc
+    
+    # Define Pauli strings and their weights
+    # Each Pauli string is represented by a list of integers:
+    # 0: Identity, 1: X, 2: Y, 3: Z
+    pauli_strings = [
+        [1, 1, 0],  # X₁X₂I₃
+        [3, 3, 0],  # Z₁Z₂I₃
+        [0, 0, 1],  # I₁I₂X₃
+    ]
+    weights = [0.5, 1.0, -0.2]
+    
+    # Build sparse Hamiltonian
+    h_sparse = tc.quantum.PauliStringSum2COO(pauli_strings, weights)
+    
+    # Or dense Hamiltonian if preferred
+    h_dense = tc.quantum.PauliStringSum2Dense(pauli_strings, weights)
+
+
+**High-Level Hamiltonian Construction:**
+
+For common Hamiltonians like Heisenberg model:
+
+.. code-block:: python
+
+    # Create a 1D chain with 10 sites
+    g = tc.templates.graphs.Line1D(10, pbc=True)  # periodic boundary condition
+    
+    # XXZ model
+    h = tc.quantum.heisenberg_hamiltonian(
+        g,
+        hxx=1.0,  # XX coupling
+        hyy=1.0,  # YY coupling
+        hzz=1.2,  # ZZ coupling
+        hx=0.5,   # X field
+        sparse=True
+    )
+
+
+**Advanced Usage:**
+
+1. Converting between xyz and Pauli string representations:
+
+.. code-block:: python
+
+    # Convert Pauli string to xyz format
+    xyz_dict = tc.quantum.ps2xyz([1, 2, 2, 0])  # X₁Y₂Y₃I₄
+    print(xyz_dict)  # {'x': [0], 'y': [1, 2], 'z': []}
+    
+    # Convert back to Pauli string
+    ps = tc.quantum.xyz2ps(xyz_dict, n=4)
+    print(ps)  # [1, 2, 2, 0]
+
+
+2. Working with MPO format:
+
+TensorCircuit-NG supports conversion from different MPO (Matrix Product Operator) formats, particularly from TensorNetwork and Quimb libraries. This is useful when you want to leverage existing MPO implementations or convert between different frameworks.
+
+**TensorNetwork MPO:**
+
+For TensorNetwork MPOs, you can convert predefined models like the Transverse Field Ising (TFI) model:
+
+.. code-block:: python
+
+    import tensorcircuit as tc
+    import tensornetwork as tn
+    
+    # Create TFI Hamiltonian MPO from TensorNetwork
+    nwires = 6
+    Jx = np.array([1.0] * (nwires - 1))  # XX coupling strength
+    Bz = np.array([-1.0] * nwires)       # Transverse field strength
+    
+    # Create TensorNetwork MPO
+    tn_mpo = tn.matrixproductstates.mpo.FiniteTFI(
+        Jx, Bz, 
+        dtype=np.complex64
+    )
+    
+    # Convert to TensorCircuit format
+    tc_mpo = tc.quantum.tn2qop(tn_mpo)
+    
+    # Get dense matrix representation
+    h_matrix = tc_mpo.eval_matrix()
+
+Note: TensorNetwork MPO currently only supports open boundary conditions.
+
+**Quimb MPO:**
+
+Quimb provides more flexible MPO construction options:
+
+.. code-block:: python
+
+    import tensorcircuit as tc
+    import quimb.tensor as qtn
+    
+    # Create Ising Hamiltonian MPO using Quimb
+    nwires = 6
+    J = 4.0    # ZZ coupling
+    h = 2.0    # X field
+    qb_mpo = qtn.MPO_ham_ising(
+        nwires, 
+        J, h,
+        cyclic=True  # Periodic boundary conditions
+    )
+    
+    # Convert to TensorCircuit format
+    tc_mpo = tc.quantum.quimb2qop(qb_mpo)
+    
+    # Custom Hamiltonian construction
+    builder = qtn.SpinHam1D()
+    builder += 1.0, "Y"  # Add Y term with strength 1.0
+    builder += 0.5, "X"  # Add X term with strength 0.5
+    H = builder.build_mpo(3)  # Build for 3 sites
+    
+    # Convert to TensorCircuit MPO
+    h_tc = tc.quantum.quimb2qop(H)
+
+
 Fermion Gaussian State Simulator
 --------------------------------
 

docs/source/quickstart.rst

@@ -224,21 +224,104 @@ The related API design in TensorCircuit-NG closely follows the functional progra
 
 **AD Support:**
 
-Gradients, vjps, jvps, natural gradients, Jacobians, and Hessians.
-AD is the base for all modern machine learning libraries.
+Automatic Differentiation (AD) is crucial for quantum circuit optimization. TensorCircuit-NG supports various differentiation operations:
 
+* Gradients: First-order derivatives
+* Vector-Jacobian products (vjps): Efficient backward-mode differentiation
+* Jacobian-vector products (jvps): Efficient forward-mode differentiation
+* Natural gradients: Geometry-aware optimization
+* Jacobians: Full derivative matrices
+* Hessians: Second-order derivatives
+
+Example of gradient computation:
+
+.. code-block:: python
+
+    import tensorcircuit as tc
+    K = tc.set_backend("tensorflow")
+
+    def circuit(params):
+        c = tc.Circuit(2)
+        c.rx(0, theta=params[0])
+        c.ry(1, theta=params[1])
+        c.cnot(0, 1)
+        return K.real(c.expectation([tc.gates.z(), [0]]))
+
+    # Get value and gradient
+    params = K.ones([2])
+    value, grad = K.value_and_grad(circuit)(params)
+    print("Value:", value)
+    print("Gradient:", grad)
+
+    # Compute Hessian
+    hess = K.hessian(circuit)(params)
+    print("Hessian:", hess)
 
 **JIT Support:**
 
-Parameterized quantum circuits can run in a blink. Always use jit if the circuit will get evaluations multiple times, it can greatly boost the simulation with two or three order time reduction. But also be cautious, users need to be familiar with jit, otherwise, the jitted function may return unexpected results or recompile on every hit (wasting lots of time).
-To learn more about the jit mechanism, one can refer to documentation or blogs on ``tf.function`` or ``jax.jit``, though these two still have subtle differences.
+Just-In-Time (JIT) compilation significantly accelerates quantum circuit simulation by optimizing the computation graph. Key points:
+
+* Use JIT for functions that will be called multiple times
+* JIT compilation has some overhead, so it's most beneficial for repeated executions
+* Ensure input shapes and types are consistent to avoid recompilation
+* The input and output of the functions are all tensors, except static inputs.
+
+Example of JIT acceleration:
+
+.. code-block:: python
+
+    import time
+    
+    # Define a quantum circuit function
+    def noisy_circuit(key):
+        c = tc.Circuit(5)
+        for i in range(5):
+            c.h(i)
+            c.depolarizing(i, px=0.01, py=0.01, pz=0.01, status=key[i])
+        return c.expectation_ps(z=[0])
+
+    # Compare performance with and without JIT
+    start = time.time()
+    for _ in range(100):
+        noisy_circuit(K.ones([5]))
+    print("Without JIT:", time.time() - start)
 
+    jitted_circuit = K.jit(noisy_circuit)
+    start = time.time()
+    for _ in range(100):
+        jitted_circuit(K.ones([5]))
+    print("With JIT:", time.time() - start)
 
 **VMAP Support:**
 
-Inputs, parameters, measurements, circuit structures, and Monte Carlo noise can all be evaluated in parallel.
-To learn more about vmap mechanism, one can refer to documentation or blogs on ``tf.vectorized_map`` or ``jax.vmap``.
-One can also refer to `tutorial <https://tensorcircuit-ng.readthedocs.io/en/latest/whitepaper/6-3-vmap.html>`_ for more details on the vmap usage in TensorCircuit-NG.
+Vectorized mapping (vmap) enables parallel evaluation across multiple inputs or parameters:
+
+* Batch processing of quantum circuit input wavefunctions
+* Batch processing quantum circuit structure
+* Parallel parameter optimization
+* Efficient Monte Carlo sampling for noise simulation
+* Vectorized measurement operations
+
+Example of vmap for parallel circuit evaluation:
+
+.. code-block:: python
+
+    # Define a parameterized circuit
+    def param_circuit(params):
+        c = tc.Circuit(2)
+        c.rx(0, theta=params[0])
+        c.ry(1, theta=params[1])
+        return K.real(c.expectation([tc.gates.z(), [0]]))
+
+    # Create batch of parameters
+    batch_params = K.ones([10, 2])
+
+    # Vectorize the circuit evaluation
+    vmap_circuit = K.vmap(param_circuit)
+    results = vmap_circuit(batch_params)
+    
+
+For more advanced usage patterns and detailed examples of vmap, refer to our `vmap tutorial <https://tensorcircuit-ng.readthedocs.io/en/latest/whitepaper/6-3-vmap.html>`_.
 
 
 Backend Agnosticism
@@ -424,15 +507,14 @@ and the other part is implemented in `TensorCircuit package <modules.html#module
     'vvag',
     'zeros']
 
-​
 
 Switch the Dtype
 --------------------
 
 TensorCircuit-NG supports simulation using 32/64 bit precession. The default dtype is 32-bit as "complex64".
 Change this by ``tc.set_dtype("complex128")``.
 
-``tc.dtypestr`` always returns the current dtype string: either "complex64" or "complex128".
+``tc.dtypestr`` always returns the current dtype string: either "complex64" or "complex128". Accordingly, ``tc.rdtypestr`` always returns the current real dtype string: either "float32" or "float64".
 
 
 Setup the Contractor
@@ -769,6 +851,79 @@ We also provider wrapper of quantum function for keras layer as :py:meth:`tensor
         l = layer(v)
     grad = tape.gradient(l, layer.trainable_variables)
 
+**JAX interfaces:**
+
+TensorCircuit-NG also newly introduces JAX interface to seamlessly integrate with JAX's ecosystem. 
+This allows you to use JAX's powerful features like automatic differentiation, JIT compilation, and vectorization with quantum circuits or functions running on any backend.
+
+Basic usage with JAX interface:
+
+.. code-block:: python
+
+    import tensorcircuit as tc
+    import jax
+    import jax.numpy as jnp
+
+    # Set non-jax backend
+    tc.set_backend("tensorflow")
+
+    def circuit(params):
+        c = tc.Circuit(2)
+        c.rx(0, theta=params[0])
+        c.ry(1, theta=params[1])
+        c.cnot(0, 1)
+        return tc.backend.real(c.expectation_ps(z=[1]))
+
+    # Wrap the circuit with JAX interface
+    jax_circuit = tc.interfaces.jax_interface(circuit, jit=True)
+
+    # Now you can use JAX features
+    params = jnp.ones(2)
+    value, grad = jax.value_and_grad(jax_circuit)(params)
+    print("Value:", value)
+    print("Gradient:", grad)
+
+Some advanced features:
+
+1. DLPack support for efficient tensor conversion:
+
+.. code-block:: python
+
+    # Enable DLPack for zero-copy tensor conversion
+    jax_circuit = tc.interfaces.jax_interface(circuit, 
+                                            jit=True, 
+                                            enable_dlpack=True)
+
+2. Explicit output shape specification for better performance:
+
+.. code-block:: python
+
+    # Specify output shape and dtype
+    jax_circuit = tc.interfaces.jax_interface(circuit,
+                                            jit=True,
+                                            output_shape=(1,),
+                                            output_dtype=jnp.float32)
+
+3. Multiple outputs support:
+
+.. code-block:: python
+
+    def multi_output_circuit(params):
+    c = tc.Circuit(2)
+    c.rx(0, theta=params[0])
+    c.ry(1, theta=params[1])
+    z0 = c.expectation([tc.gates.z(), [0]])
+    z1 = c.expectation([tc.gates.z(), [1]])
+    return tc.backend.real(z0), tc.backend.real(z1)
+
+    jax_circuit = tc.interfaces.jax_interface(multi_output_circuit,
+                                            jit=True,
+                                            output_shape=[[], []],
+                                            output_dtype=[jnp.float32, jnp.float32])
+    # Now you can use JAX features
+    params = jnp.ones(2)
+    value, grad = jax.value_and_grad(tc.utils.append(jax_circuit, sum))(params)
+
 
 
 **Scipy Interface to Utilize Scipy Optimizers:**
@@ -827,3 +982,4 @@ See :py:meth:`tensorcircuit.templates.measurements.heisenberg_measurements`
 
 .. figure:: statics/bell_pair_block.png
     :scale: 50%
+

docs/source/sharpbits.rst

@@ -191,6 +191,9 @@ If the device is not consistent, one can move the tensor between devices by ``tc
 AD Consistency
 ---------------------
 
+Gradients in terms of complex dtypes
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
 TF and JAX backend manage the differentiation rules differently for complex-valued function (actually up to a complex conjuagte). See issue discussion `tensorflow issue <https://github.com/tensorflow/tensorflow/issues/3348>`_.
 
 In TensorCircuit-NG, currently we make the difference in AD transparent, namely, when switching the backend, the AD behavior and result for complex valued function can be different and determined by the nature behavior of the corresponding backend framework.
@@ -222,4 +225,19 @@ Also see the code below for a reference:
     # tf.Tensor([0.90929717 0.90929717], shape=(2,), dtype=float32)
     # jax backend
     # [0.90929747-0.9228759j 0.90929747-0.9228759j]
-    # [0.90929747 0.90929747]
+    # [0.90929747 0.90929747]
+
+
+    VMAP outside grad-like function on tensorflow backend
+    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+    Vmap (vectorized map) outside a grad-like function may cause incorrected results on TensorFlow backends due to a long existing `bug <https://github.com/tensorflow/tensorflow/issues/52148>`_ in TensorFlow codebase. So better always stick to the first-vmap-then-differentiated paradigm.
+
+    Grad over vmap function
+    ~~~~~~~~~~~~~~~~~~~~~~~~~
+
+    A related issue is the different behavior for `K.grad(K.vmap(f))` on different backends. For tensorflow backend, the function to be differentiated has a scalar output which is the sum of all outputs.
+
+    However, for Jax backend, the function simply raise error as only scalar output function can be differentiated, no implicit sum of the vectorized ``f`` is assumed. For non-scalar output, one should use `jacrev` or `jacfwd` to get the gradient information.
+
+    Specifically, `K.grad(K.vmap(f))` on TensorFlow backend is equilvalent to `K.grad(K.append(K.vamp(f), K.sum))` on Jax backend.