Repository: netket/netket

Issue:
Allow `GraphOperator` to act on specified sites or subsystems
Imagine a composite Hilbert space with two subsystems like this:
```python
>>> hi = Spin(1/2, N=n1) * Fock(N=n2)
```
Then we have `n1 + n2` sites `[0, ..., n1 - 1]` corresponds to the first and `[n1, ..., n1 + n2 - 1]` to the second subsystem.

Constructing some Hamiltonian as sum of `LocalOperator`s is easy enough. But `GraphOperator` (and also our built-in physical Hamiltonians like `Ising` and `Heisenberg`) does not currently support that. It acts on the sites of the given graph, but currently it is expected that the graph and site indices match exactly, so there is no direct way to specify a `GraphOperator` acting only on, e.g.,  the second subsystem.

I think the best way to address this is to allow users to specify a mapping of graph site indices to Hilbert sites like this:
```python
>>> g1 = Chain(n1)
>>> g2 = Chain(n2)
>>> h1 = Heisenberg(hi, g1, acting_on_sites=list(range(n1)))
>>> h2 = GraphOperator(hi, g2, acting_on_sites=list(range(n1, n1 + n2)), ...)
```
so that for both operators, the graph node with index `Ig` corresponds the Hilbert space site with index `acting_on_sites[Ig]`. For the common pattern of having a start index and a number of sites, we could also support specifying tuples `(start, n_sites)` like this:
```python
>>> h1 = Heisenberg(hi, g1, acting_on_sites=(0, n1))
>>> h2 = GraphOperator(hi, g2, acting_on_sites=(n1, n2), ...)
```

I think this is fully backwards compatible and makes `GraphOperator` and the like more useful in composite spaces. What do you think?

A more fancy alternative, would be to label subsystems and allow to specify that label, like
```python
>>> hi = Spin(1/2, N=n1, label="spin") * Fock(N=n2, label="fock")
>>> hi.subsystem("spin")  # returns subsystem indices
[0, 1, ..., n1]
>>> h2 = GraphOperator(hi, g2, subsystem="fock", ...)
```
This would be a bit less flexible, though (since only subsystems, not arbitrary indices, could be passed).

EDIT: Improved interface in examples a bit.


Discussion:

Resolving patch:
diff --git a/netket/operator/_graph_operator.py b/netket/operator/_graph_operator.py
--- a/netket/operator/_graph_operator.py
+++ b/netket/operator/_graph_operator.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List, Tuple, Union
+
 from netket.utils.types import DType
 
 from netket.graph import AbstractGraph
@@ -20,6 +22,24 @@
 from ._local_operator import LocalOperator
 
 
+def check_acting_on_subspace(acting_on_subspace, hilbert, graph):
+    """Check `acting_on_subspace` argument used by various operators."""
+    if acting_on_subspace is None:
+        acting_on_subspace = list(range(hilbert.size))
+    elif isinstance(acting_on_subspace, int):
+        start = acting_on_subspace
+        acting_on_subspace = [start + i for i in range(graph.n_nodes)]
+    elif isinstance(acting_on_subspace, list):
+        if len(acting_on_subspace) != graph.n_nodes:
+            raise ValueError(
+                "acting_on_subspace must be a list of length graph.n_nodes"
+            )
+    else:
+        raise TypeError("acting_on_subspace must be a list or single integer")
+
+    return acting_on_subspace
+
+
 class GraphOperator(LocalOperator):
     r"""
     A graph-based quantum operator. In its simplest terms, this is the sum of
@@ -34,6 +54,8 @@ def __init__(
         bond_ops=[],
         bond_ops_colors=[],
         dtype: DType = None,
+        *,
+        acting_on_subspace: Union[List[int], int] = None,
     ):
         r"""
         A graph-based quantum operator. In its simplest terms, this is the sum of
@@ -57,6 +79,12 @@ def __init__(
              specified, the user must give a list of site operators.
          bond_ops_colors: A list of edge colors, specifying the color each
              bond operator acts on. The default is an empty list.
+         dtype: Data type type of the matrix elements.
+         acting_on_subspace: Specifies the mapping between nodes of the graph and
+            Hilbert space sites, so that graph node :code:`i ∈ [0, ..., graph.n_nodes]`,
+            corresponds to :code:`acting_on_subspace[i] ∈ [0, ..., hilbert.n_sites]`.
+            Must be a list of length `graph.n_nodes`. Passing a single integer :code:`start`
+            is equivalent to :code:`[start, ..., start + graph.n_nodes - 1]`.
 
         Examples:
          Constructs a ``GraphOperator`` operator for a 2D system.
@@ -74,12 +102,10 @@ def __init__(
          >>> print(op)
          GraphOperator(dim=20, #acting_on=40 locations, constant=0, dtype=float64, graph=Graph(n_nodes=20, n_edges=20))
         """
-
-        if graph.n_nodes != hilbert.size:
-            raise ValueError(
-                """The number of vertices in the graph ({graph.n_nodes})
-                                must match the hilbert space size ({hilbert.size})"""
-            )
+        acting_on_subspace = check_acting_on_subspace(
+            acting_on_subspace, hilbert, graph
+        )
+        self._acting_on_subspace = acting_on_subspace
 
         # Ensure that at least one of SiteOps and BondOps was initialized
         if len(bond_ops) == 0 and len(site_ops) == 0:
@@ -92,9 +118,10 @@ def __init__(
         # Site operators
         if len(site_ops) > 0:
             for i in range(graph.n_nodes):
-                for j, site_op in enumerate(site_ops):
+                for site_op in site_ops:
+                    i_prime = acting_on_subspace[i]
                     operators.append(site_op)
-                    acting_on.append([i])
+                    acting_on.append([i_prime])
 
         # Bond operators
         if len(bond_ops_colors) > 0:
@@ -107,17 +134,18 @@ def __init__(
             if len(bond_ops) > 0:
                 #  Use edge_colors to populate operators
                 for (u, v, color) in graph.edges(return_color=True):
-                    edge = u, v
+                    u, v = acting_on_subspace[u], acting_on_subspace[v]
                     for c, bond_color in enumerate(bond_ops_colors):
                         if bond_color == color:
                             operators.append(bond_ops[c])
-                            acting_on.append(edge)
+                            acting_on.append([u, v])
         else:
             assert len(bond_ops) == 1
 
-            for edge in graph.edges():
+            for u, v in graph.edges():
+                u, v = acting_on_subspace[u], acting_on_subspace[v]
                 operators.append(bond_ops[0])
-                acting_on.append(edge)
+                acting_on.append([u, v])
 
         super().__init__(hilbert, operators, acting_on, dtype=dtype)
         self._graph = graph
@@ -127,6 +155,14 @@ def graph(self) -> AbstractGraph:
         """The graph on which this Operator is defined"""
         return self._graph
 
+    @property
+    def acting_on_subspace(self):
+        """
+        Mapping between nodes of the graph and Hilbert space sites as given in
+        the constructor.
+        """
+        return self._acting_on_subspace
+
     def __repr__(self):
         ao = self.acting_on
 
diff --git a/netket/operator/_hamiltonian.py b/netket/operator/_hamiltonian.py
--- a/netket/operator/_hamiltonian.py
+++ b/netket/operator/_hamiltonian.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import List, Tuple, Union
 from numba import jit
 
 import numpy as np
@@ -146,7 +147,10 @@ def __init__(
 
         self._h = dtype(h)
         self._J = dtype(J)
-        self._edges = np.asarray(list(graph.edges()), dtype=np.intp)
+        self._edges = np.asarray(
+            [[u, v] for u, v in graph.edges()],
+            dtype=np.intp,
+        )
 
         self._dtype = dtype
 
@@ -271,23 +275,7 @@ def _flattened_kernel(
         for i in range(x.shape[0]):
             mels[diag_ind] = 0.0
             for k in range(edges.shape[0]):
-                mels[diag_ind] += (
-                    J
-                    * x[
-                        i,
-                        edges[
-                            k,
-                            0,
-                        ],
-                    ]
-                    * x[
-                        i,
-                        edges[
-                            k,
-                            1,
-                        ],
-                    ]
-                )
+                mels[diag_ind] += J * x[i, edges[k, 0]] * x[i, edges[k, 1]]
 
             odiag_ind = 1 + diag_ind
 
@@ -366,6 +354,8 @@ def __init__(
         graph: AbstractGraph,
         J: Union[float, Sequence[float]] = 1.0,
         sign_rule=None,
+        *,
+        acting_on_subspace: Union[List[int], int] = None,
     ):
         """
         Constructs an Heisenberg operator given a hilbert space and a graph providing the
@@ -378,12 +368,17 @@ def __init__(
                Can pass a sequence of coupling strengths with coloured graphs:
                edges of colour n will have coupling strength J[n]
             sign_rule: If True, Marshal's sign rule will be used. On a bipartite
-                       lattice, this corresponds to a basis change flipping the Sz direction
-                       at every odd site of the lattice. For non-bipartite lattices, the
-                       sign rule cannot be applied. Defaults to True if the lattice is
-                       bipartite, False otherwise.
-                       If a sequence of coupling strengths is passed, defaults to False
-                       and a matching sequence of sign_rule must be specified to override it
+                lattice, this corresponds to a basis change flipping the Sz direction
+                at every odd site of the lattice. For non-bipartite lattices, the
+                sign rule cannot be applied. Defaults to True if the lattice is
+                bipartite, False otherwise.
+                If a sequence of coupling strengths is passed, defaults to False
+                and a matching sequence of sign_rule must be specified to override it
+            acting_on_subspace: Specifies the mapping between nodes of the graph and
+                Hilbert space sites, so that graph node :code:`i ∈ [0, ..., graph.n_nodes - 1]`,
+                corresponds to :code:`acting_on_subspace[i] ∈ [0, ..., hilbert.n_sites]`.
+                Must be a list of length `graph.n_nodes`. Passing a single integer :code:`start`
+                is equivalent to :code:`[start, ..., start + graph.n_nodes - 1]`.
 
         Examples:
          Constructs a ``Heisenberg`` operator for a 1D system.
@@ -446,7 +441,11 @@ def __init__(
             bond_ops_colors = []
 
         super().__init__(
-            hilbert, graph, bond_ops=bond_ops, bond_ops_colors=bond_ops_colors
+            hilbert,
+            graph,
+            bond_ops=bond_ops,
+            bond_ops_colors=bond_ops_colors,
+            acting_on_subspace=acting_on_subspace,
         )
 
     @property


Test patch:
diff --git a/test/hilbert/test_hilbert.py b/test/hilbert/test_hilbert.py
--- a/test/hilbert/test_hilbert.py
+++ b/test/hilbert/test_hilbert.py
@@ -271,3 +271,15 @@ def test_deprecations():
     with pytest.warns(FutureWarning):
         with pytest.raises(ValueError):
             Spin(s=0.5, graph=g, N=3)
+
+
+def test_composite_hilbert():
+    hi1 = Spin(s=1 / 2, N=8)
+    hi2 = Spin(s=3 / 2, N=8)
+
+    hi = hi1 * hi2
+
+    assert hi.size == hi1.size + hi2.size
+
+    for i in range(hi.size):
+        assert hi.size_at_index(i) == 2 if i < 8 else 4
diff --git a/test/operator/test_operator.py b/test/operator/test_operator.py
--- a/test/operator/test_operator.py
+++ b/test/operator/test_operator.py
@@ -33,6 +33,15 @@
     hi, g, site_ops=[sigmax], bond_ops=[mszsz]
 )
 
+g_sub = nk.graph.Graph(edges=edges[:7])  # edges of first eight sites
+operators["Graph Hamiltonian (on subspace)"] = nk.operator.GraphOperator(
+    hi,
+    g_sub,
+    site_ops=[sigmax],
+    bond_ops=[mszsz],
+    acting_on_subspace=8,
+)
+
 # Graph Hamiltonian with colored edges
 edges_c = [(i, j, i % 2) for i, j in edges]
 g = nk.graph.Graph(edges=edges_c)
@@ -470,3 +479,30 @@ def test_pauli_zero():
     all_states = op.hilbert.all_states()
     _, mels = op.get_conn_padded(all_states)
     assert np.allclose(mels, 0)
+
+
+def test_operator_on_subspace():
+    hi = nk.hilbert.Spin(1 / 2, N=3) * nk.hilbert.Qubit(N=3)
+    g = nk.graph.Chain(3, pbc=False)
+
+    h1 = nk.operator.GraphOperator(hi, g, bond_ops=[mszsz], acting_on_subspace=0)
+    assert h1.acting_on_subspace == list(range(3))
+    assert nk.exact.lanczos_ed(h1)[0] == pytest.approx(-2.0)
+
+    h2 = nk.operator.GraphOperator(hi, g, bond_ops=[mszsz], acting_on_subspace=3)
+    assert h2.acting_on_subspace == list(range(3, 6))
+    assert nk.exact.lanczos_ed(h2)[0] == pytest.approx(-2.0)
+
+    h12 = h1 + h2
+    assert sorted(h12.acting_on) == [[0, 1], [1, 2], [3, 4], [4, 5]]
+    assert nk.exact.lanczos_ed(h12)[0] == pytest.approx(-4.0)
+
+    h3 = nk.operator.GraphOperator(
+        hi, g, bond_ops=[mszsz], acting_on_subspace=[0, 2, 4]
+    )
+    assert h3.acting_on_subspace == [0, 2, 4]
+    assert nk.exact.lanczos_ed(h3)[0] == pytest.approx(-2.0)
+    assert h3.acting_on == [[0, 2], [2, 4]]
+
+    h4 = nk.operator.Heisenberg(hi, g, acting_on_subspace=0)
+    assert h4.acting_on == h1.acting_on