NetKet 3.12.4__tar.gz → 3.13.0__tar.gz

{netket-3.12.4 → netket-3.13.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: NetKet
-Version: 3.12.4
+Version: 3.13.0
 Summary: Netket : Machine Learning toolbox for many-body quantum systems.
 Project-URL: homepage, https://www.netket.org
 Project-URL: documentation, https://netket.readthedocs.io/en/latest/#
@@ -18,7 +18,6 @@ Classifier: Operating System :: Unix
 Classifier: Programming Language :: Python :: 3
 Classifier: Topic :: Scientific/Engineering :: Physics
 Requires-Python: >=3.9
-Requires-Dist: beartype>=0.16.2
 Requires-Dist: equinox<0.12,>=0.11.3
 Requires-Dist: flax<0.9,>=0.7.1
 Requires-Dist: igraph<0.12.0,>=0.10.0
@@ -26,18 +25,17 @@ Requires-Dist: jax<0.5,>=0.4.16
 Requires-Dist: jaxlib<0.5,>=0.4.16
 Requires-Dist: numba4jax<0.1,>=0.0.10.post1
 Requires-Dist: numba<0.61,>=0.52
-Requires-Dist: numpy~=1.20
+Requires-Dist: numpy<3.0,>=1.20
 Requires-Dist: optax<0.3,>=0.1.3
 Requires-Dist: orjson<4,>=3.4
+Requires-Dist: plum-dispatch<3,>=2.4
 Requires-Dist: rich>=10.0
 Requires-Dist: scipy<2,>=1.5.3
 Requires-Dist: tqdm<5,>=4.60
-Requires-Dist: typing-extensions; python_version <= '3.10'
 Provides-Extra: dev
 Requires-Dist: black==23.9.1; extra == 'dev'
 Requires-Dist: build; extra == 'dev'
 Requires-Dist: coverage>=5; extra == 'dev'
-Requires-Dist: ipython; extra == 'dev'
 Requires-Dist: matplotlib>=3; extra == 'dev'
 Requires-Dist: networkx<4,>=2.4; extra == 'dev'
 Requires-Dist: pre-commit>=2.7; extra == 'dev'

{netket-3.12.4 → netket-3.13.0}/netket/_version.py

@@ -12,5 +12,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '3.12.4'
-__version_tuple__ = version_tuple = (3, 12, 4)
+__version__ = version = '3.13.0'
+__version_tuple__ = version_tuple = (3, 13, 0)

{netket-3.12.4 → netket-3.13.0}/netket/driver/abstract_variational_driver.py

@@ -161,6 +161,11 @@ class AbstractVariationalDriver(abc.ABC):
                 **modified in-place** adding new keys.
             step: the current step number.
         """
+        # Always log the acceptance.
+        if hasattr(self.state, "sampler_state"):
+            acceptance = getattr(self.state.sampler_state, "acceptance", None)
+            if acceptance is not None:
+                log_dict["acceptance"] = acceptance
 
     def reset(self):
         """

{netket-3.12.4 → netket-3.13.0}/netket/experimental/driver/tdvp_common.py

@@ -401,6 +401,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
         return loggers
 
     def _log_additional_data(self, log_dict, step):
+        super()._log_additional_data(log_dict, step)
         log_dict["t"] = self.t
 
     @property

{netket-3.12.4 → netket-3.13.0}/netket/experimental/nn/rnn/layers_fast.py

@@ -75,8 +75,7 @@ class FastRNNLayer(RNNLayer):
         hidden = self._extract_hidden(outputs, index, prev_neighbors)
         cell_mem, hidden = self.cell(inputs, cell_mem, hidden)
 
-        initializing = self.is_mutable_collection("params")
-        if not initializing:
+        if not self.is_initializing():
             _cell_mem.value = cell_mem
             _outputs.value = outputs.at[:, index, :].set(hidden)
 

{netket-3.12.4 → netket-3.13.0}/netket/experimental/observable/__init__.py

@@ -15,3 +15,5 @@
 from netket.operator._abstract_observable import AbstractObservable
 
 from .renyi2 import Renyi2EntanglementEntropy
+
+from .variance import VarianceObservable

netket-3.13.0/netket/experimental/observable/variance/__init__.py

@@ -0,0 +1,18 @@
+# Copyright 2022 The NetKet Authors - All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .variance_operator import VarianceObservable
+
+from . import expect
+from . import exact

netket-3.13.0/netket/experimental/observable/variance/exact.py

@@ -0,0 +1,94 @@
+# Copyright 2022 The NetKet Authors - All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax.numpy as jnp
+import jax
+from functools import partial
+
+from netket.vqs import FullSumState, expect, expect_and_grad
+import netket.jax as nkjax
+from netket.utils import mpi
+from netket.stats import Stats
+
+from .variance_operator import VarianceObservable
+
+
+@expect.dispatch
+def expect(vstate: FullSumState, variance_operator: VarianceObservable):
+    if variance_operator.hilbert != vstate.hilbert:
+        raise TypeError("Hilbert spaces should match")
+
+    operator_mtrx = variance_operator.operator.to_dense()
+    operator_squared_mtrx = variance_operator.operator_squared.to_dense()
+
+    return expect_and_grad_inner_fs(
+        vstate._apply_fun,
+        vstate.parameters,
+        vstate.model_state,
+        vstate._all_states,
+        operator_mtrx,
+        operator_squared_mtrx,
+        return_grad=False,
+    )
+
+
+@expect_and_grad.dispatch
+def expect_and_grad(
+    vstate: FullSumState,
+    variance_operator: VarianceObservable,
+    *,
+    mutable,
+):
+    if variance_operator.hilbert != vstate.hilbert:
+        raise TypeError("Hilbert spaces should match")
+
+    operator_mtrx = variance_operator.operator.to_dense()
+    operator_squared_mtrx = variance_operator.operator_squared.to_dense()
+
+    return expect_and_grad_inner_fs(
+        vstate._apply_fun,
+        vstate.parameters,
+        vstate.model_state,
+        vstate._all_states,
+        operator_mtrx,
+        operator_squared_mtrx,
+        return_grad=True,
+    )
+
+
+@partial(jax.jit, static_argnames=("afun", "return_grad"))
+def expect_and_grad_inner_fs(
+    afun, params, model_state, sigma, operator_mtrx, operator_squared_mtrx, return_grad
+):
+    def expect_kernel_var(params):
+        W = {"params": params, **model_state}
+
+        state = jnp.exp(afun(W, sigma))
+        state = state / jnp.linalg.norm(state)
+
+        O_mean = state.conj() @ (operator_mtrx @ state)
+        O2_mean = state.conj() @ (operator_squared_mtrx @ state)
+
+        return O2_mean - O_mean**2
+
+    if not return_grad:
+        var = expect_kernel_var(params)
+        return Stats(mean=var, error_of_mean=0.0, variance=0.0)
+
+    var, var_vjp_fun = nkjax.vjp(expect_kernel_var, params, conjugate=True)
+
+    var_grad = var_vjp_fun(jnp.ones_like(var))[0]
+    var_grad = jax.tree_map(lambda x: mpi.mpi_mean_jax(x)[0], var_grad)
+
+    return Stats(mean=var, error_of_mean=0.0, variance=0.0), var_grad

netket-3.13.0/netket/experimental/observable/variance/expect.py

@@ -0,0 +1,141 @@
+# Copyright 2022 The NetKet Authors - All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import jax.numpy as jnp
+import jax
+from functools import partial
+
+from netket.vqs import (
+    MCState,
+    expect,
+    expect_and_grad,
+    get_local_kernel,
+    get_local_kernel_arguments,
+)
+import netket.jax as nkjax
+from netket.utils import mpi
+
+from .variance_operator import VarianceObservable
+
+
+@expect.dispatch
+def expect(vstate: MCState, variance_operator: VarianceObservable, chunk_size: None):
+    if variance_operator.hilbert != vstate.hilbert:
+        raise TypeError("Hilbert spaces should match")
+
+    local_kernel = get_local_kernel(vstate, variance_operator.operator)
+    local_kernel2 = get_local_kernel(vstate, variance_operator.operator_squared)
+
+    sigma, args = get_local_kernel_arguments(vstate, variance_operator.operator)
+    sigma, args2 = get_local_kernel_arguments(
+        vstate, variance_operator.operator_squared
+    )
+
+    return expect_and_grad_inner_mc(
+        vstate._apply_fun,
+        vstate.parameters,
+        vstate.model_state,
+        local_kernel,
+        local_kernel2,
+        sigma,
+        args,
+        args2,
+        return_grad=False,
+    )
+
+
+@expect_and_grad.dispatch
+def expect_and_grad(
+    vstate: MCState,
+    variance_operator: VarianceObservable,
+    chunk_size: None,
+    *,
+    mutable,
+):
+    if variance_operator.hilbert != vstate.hilbert:
+        raise TypeError("Hilbert spaces should match")
+
+    local_kernel = get_local_kernel(vstate, variance_operator.operator)
+    local_kernel2 = get_local_kernel(vstate, variance_operator.operator_squared)
+
+    sigma, args = get_local_kernel_arguments(vstate, variance_operator.operator)
+    sigma, args2 = get_local_kernel_arguments(
+        vstate, variance_operator.operator_squared
+    )
+
+    return expect_and_grad_inner_mc(
+        vstate._apply_fun,
+        vstate.parameters,
+        vstate.model_state,
+        local_kernel,
+        local_kernel2,
+        sigma,
+        args,
+        args2,
+        return_grad=True,
+    )
+
+
+@partial(
+    jax.jit, static_argnames=("afun", "local_kernel", "local_kernel2", "return_grad")
+)
+def expect_and_grad_inner_mc(
+    afun,
+    params,
+    model_state,
+    local_kernel,
+    local_kernel2,
+    sigma,
+    args,
+    args2,
+    return_grad,
+):
+    n_chains = sigma.shape[0]
+
+    σ = sigma.reshape(-1, sigma.shape[-1])
+
+    def expect_kernel_var(params):
+        log_pdf = lambda params, σ: 2 * afun({"params": params, **model_state}, σ).real
+
+        def local_kernel_var(params, σ):
+            W = {"params": params, **model_state}
+
+            O_mean = nkjax.expect(
+                log_pdf,
+                lambda params, σ: local_kernel(
+                    afun, {"params": params, **model_state}, σ, args
+                ),
+                params,
+                σ,
+                n_chains=n_chains,
+            )[0]
+
+            O2_loc = local_kernel2(afun, W, σ, args2)
+
+            return (O2_loc - O_mean**2).real
+
+        return nkjax.expect(log_pdf, local_kernel_var, params, σ, n_chains=n_chains)
+
+    if not return_grad:
+        var, var_stats = expect_kernel_var(params)
+        return var_stats
+
+    var, var_vjp_fun, var_stats = nkjax.vjp(
+        expect_kernel_var, params, has_aux=True, conjugate=True
+    )
+
+    var_grad = var_vjp_fun(jnp.ones_like(var))[0]
+    var_grad = jax.tree_map(lambda x: mpi.mpi_sum_jax(x)[0], var_grad)
+
+    return var_stats, var_grad

netket-3.13.0/netket/experimental/observable/variance/variance_operator.py

@@ -0,0 +1,84 @@
+# Copyright 2022 The NetKet Authors - All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import netket as nk
+from netket.operator import AbstractOperator
+
+from netket.operator._abstract_observable import AbstractObservable
+
+
+class VarianceObservable(AbstractObservable):
+    r"""
+    Observable computing the variance of a quantum operator :math:`O`.
+    """
+
+    def __init__(self, operator: AbstractOperator, use_Oloc2: bool = False):
+        r"""
+        Constructs the observable computing the variance of an arbitrary quantum operator :math:`O` as:
+
+        .. math::
+
+            \text{Var} = \frac{\langle \Psi | O^2 | \Psi \rangle}{\langle \Psi | \Psi \rangle} - \bigg( \frac{\langle \Psi | O | \Psi \rangle}{\langle \Psi | \Psi \rangle}\bigg)^2
+
+        It can compute the first term using either the estimator of the squared operator :math:`O^2` (more precise but less efficient, since it requires
+        the connected configurations and the matrix elements of :math:`O^2`):
+
+        .. math::
+
+            \text{Var} = \mathbb{E}_{\sigma \sim |\Psi(\sigma)|^2}\bigg[\frac{\langle \sigma | O^2 | \Psi \rangle}{\langle \sigma | \Psi \rangle}\bigg] - \bigg(\mathbb{E}_{\sigma \sim |\Psi(\sigma)|^2}\bigg[\frac{\langle \sigma | O | \Psi \rangle}{\langle \sigma | \Psi \rangle}\bigg]\bigg)^2.
+
+        or using the square modulus of the estimator of :math:`O` (more noisy but more efficient):
+
+        .. math::
+
+            \text{Var} = \mathbb{E}_{\sigma \sim |\Psi(\sigma)|^2}\bigg[\bigg(\frac{\langle \sigma | O | \Psi \rangle}{\langle \sigma | \Psi \rangle} - \mathbb{E}_{\sigma \sim |\Psi(\sigma)|^2}\bigg[\frac{\langle \sigma | O | \Psi \rangle}{\langle \sigma | \Psi \rangle}\bigg]\bigg)^2\bigg]
+
+        This VariationalOperator wraps an operator such that the gradient will not be computed
+        with respect to the expectation value, but with respect to the variance.
+
+        Args:
+            operator: The operator for which the variance is to be computed.
+            use_Oloc2: If True, uses the local estimator of the squared operator `O^2` for variance computation.
+                        If False, uses only the operator `O` for variance computation (defaults to False).
+
+        Returns:
+            Observable computing the variance of `operator`.
+        """
+        super().__init__(operator.hilbert)
+        self._operator = operator
+
+        if use_Oloc2:
+            self._operator_squared = nk.operator.Squared(operator)
+        else:
+            self._operator_squared = operator @ operator
+
+    @property
+    def operator(self) -> AbstractOperator:
+        """
+        The operator for which the variance is to be computed.
+        """
+        return self._operator
+
+    @property
+    def operator_squared(self) -> AbstractOperator:
+        """
+        The squared of the operator for which the variance is to be computed.
+        Depending on the flag `use_Oloc2`, this can be the operator using the local
+        estimator of `O^2` (True), or the one using the square modulus of the
+        local estimator of `O` (False).
+        """
+        return self._operator_squared
+
+    def __repr__(self):
+        return f"VarianceObservable(op={self.operator})"

netket-3.13.0/netket/experimental/sampler/__init__.py

@@ -0,0 +1,54 @@
+# Copyright 2020, 2021 The NetKet Authors - All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from . import rules
+
+from .metropolis_pmap import MetropolisSamplerPmap
+
+from .metropolis import MetropolisParticleExchange
+
+
+# Deprecated bindings from stabilisaation
+from netket.sampler import (
+    ParallelTemperingSampler as _deprecated_MetropolisPtSampler,
+    ParallelTemperingLocal as _deprecated_MetropolisLocalPtSampler,
+    ParallelTemperingExchange as _deprecated_MetropolisExchangePtSampler,
+)
+
+_deprecations = {
+    # May 2024, NetKet 3.12
+    "MetropolisPtSampler": (
+        "netket.experimental.sampler.MetropolisPtSampler is deprecated: use "
+        "netket.sampler.ParallelTemperingSampler (netket >= 3.12)",
+        _deprecated_MetropolisPtSampler,
+    ),
+    "MetropolisLocalPt": (
+        "netket.experimental.sampler.MetropolisLocalPt is deprecated: use "
+        "netket.sampler.ParallelTemperingLocal (netket >= 3.12)",
+        _deprecated_MetropolisLocalPtSampler,
+    ),
+    "MetropolisExchangePt": (
+        "netket.experimental.sampler.MetropolisExchangePt is deprecated: use "
+        "netket.sampler.ParallelTemperingExchange (netket >= 3.12)",
+        _deprecated_MetropolisExchangePtSampler,
+    ),
+}
+
+from netket.utils.deprecation import deprecation_getattr as _deprecation_getattr
+from netket.utils import _hide_submodules
+
+__getattr__ = _deprecation_getattr(__name__, _deprecations)
+_hide_submodules(__name__)
+
+del _deprecation_getattr

{netket-3.12.4 → netket-3.13.0}/netket/experimental/sampler/rules/fermion_2nd.py

@@ -1,7 +1,3 @@
-from functools import partial
-
-import jax
-import jax.numpy as jnp
 from typing import Optional
 import numpy as np
 
@@ -9,7 +5,6 @@ from netket.sampler.rules import ExchangeRule
 from netket.graph import AbstractGraph
 from netket.graph import disjoint_union
 from netket.experimental.hilbert import SpinOrbitalFermions
-from netket.jax.sharding import sharding_decorator
 
 
 class ParticleExchangeRule(ExchangeRule):
@@ -72,52 +67,5 @@ class ParticleExchangeRule(ExchangeRule):
                 )
         super().__init__(clusters=clusters, graph=graph, d_max=d_max)
 
-    def transition(rule, sampler, machine, parameters, state, key, σ):
-        n_chains = σ.shape[0]
-
-        # compute a mask for the clusters that can be hopped
-        hoppable_clusters = _compute_hoppable_clusters_mask(rule.clusters, σ)
-
-        keys = jnp.asarray(jax.random.split(key, n_chains))
-
-        # we use shard_map to avoid the all-gather coming from the batched jnp.take / indexing
-        @partial(sharding_decorator, sharded_args_tree=(True, True, True))
-        @jax.vmap
-        def _update_samples(key, σ, hoppable_clusters):
-            # pick a random cluster, taking into account the mask
-            n_conn = hoppable_clusters.sum(axis=-1)
-            cluster = jax.random.choice(
-                key,
-                a=jnp.arange(rule.clusters.shape[0]),
-                p=hoppable_clusters,
-                replace=True,
-            )
-
-            # sites to be exchanged
-            si = rule.clusters[cluster, 0]
-            sj = rule.clusters[cluster, 1]
-
-            σp = σ.at[si].set(σ[sj])
-            σp = σp.at[sj].set(σ[si])
-
-            # compute the number of connected sites
-            hoppable_clusters_proposed = _compute_hoppable_clusters_mask(
-                rule.clusters, σp
-            )
-            n_conn_proposed = hoppable_clusters_proposed.sum(axis=-1)
-            log_prob_corr = jnp.log(n_conn) - jnp.log(n_conn_proposed)
-            return σp, log_prob_corr
-
-        return _update_samples(keys, σ, hoppable_clusters)
-
     def __repr__(self):
         return f"ParticleExchangeRule(# of clusters: {len(self.clusters)})"
-
-
-@jax.jit
-def _compute_hoppable_clusters_mask(clusters, σ):
-    # mask the clusters to include only feasible moves (occ -> unocc, or the inverse)
-    hoppable_clusters_mask = ~jnp.isclose(
-        σ[..., clusters[:, 0]], σ[..., clusters[:, 1]]
-    )
-    return hoppable_clusters_mask

{netket-3.12.4 → netket-3.13.0}/netket/graph/common_lattices.py

@@ -508,7 +508,7 @@ def KitaevHoneycomb(
     Example:
         Construct a Kitaev honeycomb lattice with 3 × 3 unit cells:
 
-        >>> from netket.graph import Honeycomb
+        >>> from netket.graph import KitaevHoneycomb
         >>> g = KitaevHoneycomb(extent=[3, 3])
         >>> print(g.n_nodes)
         18

{netket-3.12.4 → netket-3.13.0}/netket/graph/graph.py

@@ -261,8 +261,24 @@ def DoubledGraph(graph: AbstractGraph) -> Graph:
     return Graph(n_nodes=dnodes, edges=dedges)
 
 
-def disjoint_union(graph_1: Graph, graph_2: Graph) -> Graph:
+def disjoint_union(*graphs: tuple[Graph, ...]) -> Graph:
     """
-    Returns the disjoint union of two graphs.
+    Returns the ordered disjoint union of the input graphs.
+
+    The ordering of the nodes is the same as the input order, so if
+    the I-th input graph has :math:`n_i` nodes, the first `[0, n_1[` nodes
+    will correspond to the first graph, the nodes between :math:`[n_1, n_1+n2[`
+    will correspond to the second graph, and so on and so forth.
+
+    Args:
+       *graphs: A variable number of NetKet graphs (or lattices).
+
+    Output:
+        A NetKet Graph object.
     """
-    return Graph.from_igraph(igraph.disjoint_union([graph_1._igraph, graph_2._igraph]))
+    if len(graphs) == 0:
+        raise ValueError("Must provide at least one graph")
+    elif len(graphs) == 1:
+        return graphs[0]
+    else:
+        return Graph.from_igraph(igraph.disjoint_union([g._igraph for g in graphs]))

{netket-3.12.4 → netket-3.13.0}/netket/hilbert/discrete_hilbert.py

@@ -24,7 +24,7 @@ import jax.numpy as jnp
 
 from equinox import error_if
 
-from netket.utils.types import Array
+from netket.utils.types import Array, DType
 
 from .abstract_hilbert import AbstractHilbert
 from .index import is_indexable
@@ -35,6 +35,17 @@ class DiscreteHilbert(AbstractHilbert):
 
     This class defines the common interface that can be used to
     interact with hilbert spaces on lattices.
+
+    The local degrees of freedom are discrete and numerable, therefore they can always be
+    converted to and from integers using the methods `states_to_local_indices` and
+    `local_indices_to_states`.
+    This can be used to simplify the implementation of operators that might act on the
+    hilbert space, to avoid reimplementing the logic for different values of the local
+    degrees of freedom.
+
+    If the Hilbert space is small enough, individual states can be converted to and from
+    integers labelling all the basis states. This is done using the methods `numbers_to_states`
+    and `states_to_numbers`.
     """
 
     def __init__(self, shape: tuple[int, ...]):
@@ -236,6 +247,30 @@ class DiscreteHilbert(AbstractHilbert):
             f"implemented for Hilbert space {self} of type {type(self)}"
         )
 
+    def local_indices_to_states(self, x: Array, dtype: DType = None):
+        r"""
+        Converts a tensor of integers to the corresponding local_values in
+        this hilbert space.
+
+        Equivalent to
+
+        .. code::py
+
+            hilbert.local_states[x]
+
+        The input last dimension must match the size of this Hilbert space.
+        This function can be jax-jitted.
+
+        Args:
+            x: a tensor with integer dtype and whose last dimension matches
+                the size of this Hilbert space.
+
+        Returns:
+            a tensor with the same shape as the input, and values corresponding
+            to the local_state indexed by the input tensor `x`.
+        """
+        raise NotImplementedError()
+
     @property
     def is_indexable(self) -> bool:
         """Whether the space can be indexed with an integer"""

{netket-3.12.4 → netket-3.13.0}/netket/hilbert/doubled_hilbert.py

@@ -15,9 +15,11 @@
 import jax.numpy as jnp
 
 from netket.utils.dispatch import parametric
+from netket.utils.types import Array, DType
 
 from .abstract_hilbert import AbstractHilbert
 from .discrete_hilbert import DiscreteHilbert
+from .index import is_indexable
 
 
 @parametric
@@ -63,6 +65,12 @@ class DoubledHilbert(DiscreteHilbert):
     def is_finite(self):
         return self.physical.is_finite
 
+    @property
+    def is_indexable(self) -> bool:
+        """Whether the space can be indexed with an integer"""
+        n = self.physical.n_states
+        return self.physical.is_indexable and is_indexable([n, n])
+
     @property
     def local_size(self):
         return self.physical.local_size
@@ -143,6 +151,9 @@ class DoubledHilbert(DiscreteHilbert):
     def states_to_local_indices(self, x):
         return self.physical.states_to_local_indices(x)
 
+    def local_indices_to_states(self, x: Array, dtype: DType = None):
+        return self.physical.local_indices_to_states(x, dtype=dtype)
+
     def __repr__(self):
         return f"DoubledHilbert({self.physical})"