SymbolicHamiltonian refactor

src/qibo/hamiltonians/__init__.py

@@ -1,6 +1,2 @@
-from qibo.hamiltonians.hamiltonians import (
-    Hamiltonian,
-    SymbolicHamiltonian,
-    TrotterHamiltonian,
-)
+from qibo.hamiltonians.hamiltonians import Hamiltonian, SymbolicHamiltonian
 from qibo.hamiltonians.models import TFIM, XXX, XXZ, Heisenberg, MaxCut, X, Y, Z

src/qibo/hamiltonians/models.py

@@ -1,9 +1,10 @@
 from functools import reduce
-from typing import Union
+from typing import Optional, Union
 
 import numpy as np
 
-from qibo.backends import _check_backend, matrices
+from qibo import symbols
+from qibo.backends import Backend, _check_backend, matrices
 from qibo.config import raise_error
 from qibo.hamiltonians.hamiltonians import Hamiltonian, SymbolicHamiltonian
 from qibo.hamiltonians.terms import HamiltonianTerm
@@ -25,7 +26,7 @@ def X(nqubits, dense: bool = True, backend=None):
             in the execution. If ``None``, it uses the current backend.
             Defaults to ``None``.
     """
-    return _OneBodyPauli(nqubits, matrices.X, dense, backend=backend)
+    return _OneBodyPauli(nqubits, symbols.X, dense, backend=backend)
 
 
 def Y(nqubits, dense: bool = True, backend=None):
@@ -43,7 +44,7 @@ def Y(nqubits, dense: bool = True, backend=None):
             in the execution. If ``None``, it uses the current backend.
             Defaults to ``None``.
     """
-    return _OneBodyPauli(nqubits, matrices.Y, dense, backend=backend)
+    return _OneBodyPauli(nqubits, symbols.Y, dense, backend=backend)
 
 
 def Z(nqubits, dense: bool = True, backend=None):
@@ -61,7 +62,7 @@ def Z(nqubits, dense: bool = True, backend=None):
             in the execution. If ``None``, it uses the current backend.
             Defaults to ``None``.
     """
-    return _OneBodyPauli(nqubits, matrices.Z, dense, backend=backend)
+    return _OneBodyPauli(nqubits, symbols.Z, dense, backend=backend)
 
 
 def TFIM(nqubits, h: float = 0.0, dense: bool = True, backend=None):
@@ -80,25 +81,31 @@ def TFIM(nqubits, h: float = 0.0, dense: bool = True, backend=None):
     """
     if nqubits < 2:
         raise_error(ValueError, "Number of qubits must be larger than one.")
+    backend = _check_backend(backend)
     if dense:
         condition = lambda i, j: i in {j % nqubits, (j + 1) % nqubits}
-        ham = -_build_spin_model(nqubits, matrices.Z, condition)
+        ham = -_build_spin_model(nqubits, backend.matrices.Z, condition, backend)
         if h != 0:
             condition = lambda i, j: i == j % nqubits
-            ham -= h * _build_spin_model(nqubits, matrices.X, condition)
+            ham -= h * _build_spin_model(
+                nqubits, backend.matrices.X, condition, backend
+            )
         return Hamiltonian(nqubits, ham, backend=backend)
 
-    matrix = -(
-        _multikron([matrices.Z, matrices.Z]) + h * _multikron([matrices.X, matrices.I])
-    )
-    terms = [HamiltonianTerm(matrix, i, i + 1) for i in range(nqubits - 1)]
-    terms.append(HamiltonianTerm(matrix, nqubits - 1, 0))
-    ham = SymbolicHamiltonian(backend=backend)
-    ham.terms = terms
+    term = lambda q1, q2: symbols.Z(q1, backend=backend) * symbols.Z(
+        q2, backend=backend
+    ) + h * symbols.X(q1, backend=backend)
+    form = -1 * sum(term(i, i + 1) for i in range(nqubits - 1)) - term(nqubits - 1, 0)
+    ham = SymbolicHamiltonian(form=form, nqubits=nqubits, backend=backend)
     return ham
 
 
-def MaxCut(nqubits, dense: bool = True, backend=None):
+def MaxCut(
+    nqubits,
+    dense: bool = True,
+    adj_matrix: Optional[Union[list[list[float]], np.ndarray]] = None,
+    backend: Optional[Backend] = None,
+):
     """Max Cut Hamiltonian.
 
     .. math::
@@ -109,26 +116,29 @@ def MaxCut(nqubits, dense: bool = True, backend=None):
         dense (bool): If ``True`` it creates the Hamiltonian as a
             :class:`qibo.core.hamiltonians.Hamiltonian`, otherwise it creates
             a :class:`qibo.core.hamiltonians.SymbolicHamiltonian`.
+        adj_matrix (list[list[float]] | np.ndarray): Adjecency matrix of the graph. Defaults to a
+            homogeneous fully connected graph with all edges having an equal 1.0 weight.
         backend (:class:`qibo.backends.abstract.Backend`, optional): backend to be used
             in the execution. If ``None``, it uses the current backend.
             Defaults to ``None``.
     """
-    import sympy as sp
+    if adj_matrix is None:
+        adj_matrix = np.ones((nqubits, nqubits))
+    elif len(adj_matrix) != nqubits:
+        raise_error(
+            RuntimeError,
+            f"Expected an adjacency matrix of shape ({nqubits},{nqubits}) for a {nqubits}-qubits system.",
+        )
 
-    Z = sp.symbols(f"Z:{nqubits}")
-    V = sp.symbols(f"V:{nqubits**2}")
-    sham = -sum(
-        V[i * nqubits + j] * (1 - Z[i] * Z[j])
+    form = -sum(
+        adj_matrix[i][j]
+        * (1 - symbols.Z(i, backend=backend) * symbols.Z(j, backend=backend))
         for i in range(nqubits)
         for j in range(nqubits)
     )
-    sham /= 2
+    form /= 2
 
-    v = np.ones(nqubits**2)
-    smap = {s: (i, matrices.Z) for i, s in enumerate(Z)}
-    smap.update({s: (i, v[i]) for i, s in enumerate(V)})
-
-    ham = SymbolicHamiltonian(sham, symbol_map=smap, backend=backend)
+    ham = SymbolicHamiltonian(form, nqubits=nqubits, backend=backend)
     if dense:
         return ham.dense
     return ham
@@ -203,14 +213,16 @@ def Heisenberg(
 
     backend = _check_backend(backend)
 
-    paulis = [matrices.X, matrices.Y, matrices.Z]
+    paulis = (symbols.X, symbols.Y, symbols.Z)
 
     if dense:
         condition = lambda i, j: i in {j % nqubits, (j + 1) % nqubits}
         matrix = np.zeros((2**nqubits, 2**nqubits), dtype=complex)
         matrix = backend.cast(matrix, dtype=matrix.dtype)
         for ind, pauli in enumerate(paulis):
-            double_term = _build_spin_model(nqubits, pauli, condition)
+            double_term = _build_spin_model(
+                nqubits, pauli(0, backend=backend).matrix, condition, backend
+            )
             double_term = backend.cast(double_term, dtype=double_term.dtype)
             matrix = matrix - coupling_constants[ind] * double_term
             matrix = (
@@ -221,31 +233,24 @@ def Heisenberg(
 
         return Hamiltonian(nqubits, matrix, backend=backend)
 
-    hx = _multikron([matrices.X, matrices.X])
-    hy = _multikron([matrices.Y, matrices.Y])
-    hz = _multikron([matrices.Z, matrices.Z])
-
-    matrix = (
-        -coupling_constants[0] * hx
-        - coupling_constants[1] * hy
-        - coupling_constants[2] * hz
-    )
+    def h(symbol):
+        return lambda q1, q2: symbol(q1, backend=backend) * symbol(q2, backend=backend)
 
-    terms = [HamiltonianTerm(matrix, i, i + 1) for i in range(nqubits - 1)]
-    terms.append(HamiltonianTerm(matrix, nqubits - 1, 0))
+    def term(q1, q2):
+        return sum(
+            coeff * h(operator)(q1, q2)
+            for coeff, operator in zip(coupling_constants, paulis)
+        )
 
-    terms.extend(
-        [
-            -field_strength * HamiltonianTerm(pauli, qubit)
-            for qubit in range(nqubits)
-            for field_strength, pauli in zip(external_field_strengths, paulis)
-            if field_strength != 0.0
-        ]
+    form = -1 * sum(term(i, i + 1) for i in range(nqubits - 1)) - term(nqubits - 1, 0)
+    form -= sum(
+        field_strength * pauli(qubit)
+        for qubit in range(nqubits)
+        for field_strength, pauli in zip(external_field_strengths, paulis)
+        if field_strength != 0.0
     )
 
-    ham = SymbolicHamiltonian(backend=backend)
-    ham.terms = terms
-
+    ham = SymbolicHamiltonian(form=form, backend=backend)
     return ham
 
 
@@ -342,7 +347,7 @@ def XXZ(nqubits, delta=0.5, dense: bool = True, backend=None):
     return Heisenberg(nqubits, [-1, -1, -delta], 0, dense=dense, backend=backend)
 
 
-def _multikron(matrix_list):
+def _multikron(matrix_list, backend):
     """Calculates Kronecker product of a list of matrices.
 
     Args:
@@ -351,28 +356,90 @@ def _multikron(matrix_list):
     Returns:
         ndarray: Kronecker product of all matrices in ``matrix_list``.
     """
-    return reduce(np.kron, matrix_list)
+    # TO DO: check whether this scales better on gpu
+    """
+    indices = list(range(2 * len(matrix_list)))
+    even, odd = indices[::2], indices[1::2]
+    lhs = zip(even, odd)
+    rhs = even + odd
+    einsum_args = [item for pair in zip(matrix_list, lhs) for item in pair]
+    dim = 2 ** len(matrix_list)
+    if backend.platform != "tensorflow":
+        h = backend.np.einsum(*einsum_args, rhs)
+    else:
+        h = np.einsum(*einsum_args, rhs)
+    h = backend.np.sum(backend.np.reshape(h, (-1, dim, dim)), axis=0)
+    return h
+    """
+    # reduce appears to be faster especially when matrix_list is long
+    return reduce(backend.np.kron, matrix_list)
 
 
-def _build_spin_model(nqubits, matrix, condition):
+def _build_spin_model(nqubits, matrix, condition, backend):
     """Helper method for building nearest-neighbor spin model Hamiltonians."""
     h = sum(
-        _multikron(matrix if condition(i, j) else matrices.I for j in range(nqubits))
+        reduce(
+            backend.np.kron,
+            (
+                matrix if condition(i, j) else backend.matrices.I()
+                for j in range(nqubits)
+            ),
+        )
         for i in range(nqubits)
     )
+    """
+    indices = list(range(2 * nqubits))
+    even, odd = indices[::2], indices[1::2]
+    lhs = zip(
+        nqubits
+        * [
+            len(indices),
+        ],
+        even,
+        odd,
+    )
+    rhs = (
+        [
+            len(indices),
+        ]
+        + even
+        + odd
+    )
+    columns = [
+        backend.np.reshape(
+            backend.np.concatenate(
+                [
+                    matrix if condition(i, j) else backend.matrices.I()
+                    for i in range(nqubits)
+                ],
+                axis=0,
+            ),
+            (nqubits, 2, 2),
+        )
+        for j in range(nqubits)
+    ]
+    einsum_args = [item for pair in zip(columns, lhs) for item in pair]
+    dim = 2**nqubits
+    if backend.platform == "tensorflow":
+        h = np.einsum(*einsum_args, rhs)
+    else:
+        h = backend.np.einsum(*einsum_args, rhs, optimize=True)
+    h = backend.np.sum(backend.np.reshape(h, (nqubits, dim, dim)), axis=0)
+    """
     return h
 
 
-def _OneBodyPauli(nqubits, matrix, dense: bool = True, backend=None):
-    """Helper method for constracting non-interacting
+def _OneBodyPauli(nqubits, operator, dense: bool = True, backend=None):
+    """Helper method for constructing non-interacting
     :math:`X`, :math:`Y`, and :math:`Z` Hamiltonians."""
+    backend = _check_backend(backend)
     if dense:
         condition = lambda i, j: i == j % nqubits
-        ham = -_build_spin_model(nqubits, matrix, condition)
+        ham = -_build_spin_model(
+            nqubits, operator(0, backend=backend).matrix, condition, backend
+        )
         return Hamiltonian(nqubits, ham, backend=backend)
 
-    matrix = -matrix
-    terms = [HamiltonianTerm(matrix, i) for i in range(nqubits)]
-    ham = SymbolicHamiltonian(backend=backend)
-    ham.terms = terms
+    form = sum([-1 * operator(i, backend=backend) for i in range(nqubits)])
+    ham = SymbolicHamiltonian(form=form, backend=backend)
     return ham