We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Cannot create Spin hilbert with odd sites and S=1
Reported by @fabienalet 

```python
>>> import netket as nk
>>> nk.hilbert.Spin(1.0, N=3, total_sz=1.0)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py", line 88, in __init__
    _check_total_sz(total_sz, N)
  File "/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py", line 40, in _check_total_sz
    raise Exception(
Exception: Cannot fix the total magnetization: Nspins + totalSz must be even.
```

the cuprit is the second check in here: (size is number of spins)
```python
def _check_total_sz(total_sz, size):
    if total_sz is None:
        return

    m = round(2 * total_sz)
    if np.abs(m) > size:
        raise Exception(
            "Cannot fix the total magnetization: 2|M| cannot " "exceed Nspins."
        )

    if (size + m) % 2 != 0:
        raise Exception(
            "Cannot fix the total magnetization: Nspins + " "totalSz must be even."
        )
```

The second check works only for Spins 1/2.
Can we generalise it to always work?
I'm too tired to think how to now.

Cannot create Spin hilbert with odd sites and S=1
Reported by @fabienalet 

```python
>>> import netket as nk
>>> nk.hilbert.Spin(1.0, N=3, total_sz=1.0)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py", line 88, in __init__
    _check_total_sz(total_sz, N)
  File "/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py", line 40, in _check_total_sz
    raise Exception(
Exception: Cannot fix the total magnetization: Nspins + totalSz must be even.
```

the cuprit is the second check in here: (size is number of spins)
```python
def _check_total_sz(total_sz, size):
    if total_sz is None:
        return

    m = round(2 * total_sz)
    if np.abs(m) > size:
        raise Exception(
            "Cannot fix the total magnetization: 2|M| cannot " "exceed Nspins."
        )

    if (size + m) % 2 != 0:
        raise Exception(
            "Cannot fix the total magnetization: Nspins + " "totalSz must be even."
        )
```

The second check works only for Spins 1/2.
Can we generalise it to always work?
I'm too tired to think how to now.

Cannot create Spin hilbert with odd sites and S=1
Reported by @fabienalet 

```python
>>> import netket as nk
>>> nk.hilbert.Spin(1.0, N=3, total_sz=1.0)
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py", line 88, in __init__
    _check_total_sz(total_sz, N)
  File "/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py", line 40, in _check_total_sz
    raise Exception(
Exception: Cannot fix the total magnetization: Nspins + totalSz must be even.
```

the cuprit is the second check in here: (size is number of spins)
```python
def _check_total_sz(total_sz, size):
    if total_sz is None:
        return

    m = round(2 * total_sz)
    if np.abs(m) > size:
        raise Exception(
            "Cannot fix the total magnetization: 2|M| cannot " "exceed Nspins."
        )

    if (size + m) % 2 != 0:
        raise Exception(
            "Cannot fix the total magnetization: Nspins + " "totalSz must be even."
        )
```

The second check works only for Spins 1/2.
Can we generalise it to always work?
I'm too tired to think how to now.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `netket/hilbert/spin.py`
Content:
```
1 # Copyright 2021 The NetKet Authors - All rights reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #    http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 from fractions import Fraction
16 from typing import Optional, List, Union, Iterable
17 
18 import jax
19 from jax import numpy as jnp
20 import numpy as np
21 from netket.graph import AbstractGraph
22 from numba import jit
23 
24 
25 from .custom_hilbert import CustomHilbert
26 from ._deprecations import graph_to_N_depwarn
27 
28 
29 def _check_total_sz(total_sz, size):
30     if total_sz is None:
31         return
32 
33     m = round(2 * total_sz)
34     if np.abs(m) > size:
35         raise Exception(
36             "Cannot fix the total magnetization: 2|M| cannot " "exceed Nspins."
37         )
38 
39     if (size + m) % 2 != 0:
40         raise Exception(
41             "Cannot fix the total magnetization: Nspins + " "totalSz must be even."
42         )
43 
44 
45 @jit(nopython=True)
46 def _sum_constraint(x, total_sz):
47     return np.sum(x, axis=1) == round(2 * total_sz)
48 
49 
50 class Spin(CustomHilbert):
51     r"""Hilbert space obtained as tensor product of local spin states."""
52 
53     def __init__(
54         self,
55         s: float,
56         N: int = 1,
57         total_sz: Optional[float] = None,
58         graph: Optional[AbstractGraph] = None,
59     ):
60         r"""Hilbert space obtained as tensor product of local spin states.
61 
62         Args:
63            s: Spin at each site. Must be integer or half-integer.
64            N: Number of sites (default=1)
65            total_sz: If given, constrains the total spin of system to a particular value.
66            graph: (deprecated) a graph from which to extract the number of sites.
67 
68         Examples:
69            Simple spin hilbert space.
70 
71            >>> from netket.hilbert import Spin
72            >>> g = Hypercube(length=10,n_dim=2,pbc=True)
73            >>> hi = Spin(s=0.5, N=4)
74            >>> print(hi.size)
75            4
76         """
77         N = graph_to_N_depwarn(N=N, graph=graph)
78 
79         local_size = round(2 * s + 1)
80         local_states = np.empty(local_size)
81 
82         assert int(2 * s + 1) == local_size
83 
84         for i in range(local_size):
85             local_states[i] = -round(2 * s) + 2 * i
86         local_states = local_states.tolist()
87 
88         _check_total_sz(total_sz, N)
89         if total_sz is not None:
90 
91             def constraints(x):
92                 return _sum_constraint(x, total_sz)
93 
94         else:
95             constraints = None
96 
97         self._total_sz = total_sz if total_sz is None else int(total_sz)
98         self._s = s
99         self._local_size = local_size
100 
101         super().__init__(local_states, N, constraints)
102 
103     def __pow__(self, n):
104         if self._total_sz is None:
105             total_sz = None
106         else:
107             total_sz = total_sz * n
108 
109         return Spin(self._s, self.size * n, total_sz=total_sz)
110 
111     def _mul_sametype_(self, other):
112         assert type(self) == type(other)
113         if self._s == other._s:
114             if self._total_sz is None and other._total_sz is None:
115                 return Spin(s=self._s, N=self.size + other.size)
116 
117         return NotImplemented
118 
119     def ptrace(self, sites: Union[int, List]) -> Optional["Spin"]:
120         if isinstance(sites, int):
121             sites = [sites]
122 
123         for site in sites:
124             if site < 0 or site >= self.size:
125                 raise ValueError(
126                     f"Site {site} not in this hilbert space of site {self.size}"
127                 )
128 
129         if self._total_sz is not None:
130             raise TypeError(
131                 "Cannot take the partial trace with a total magnetization constraint."
132             )
133 
134         Nsites = len(sites)
135 
136         if self.size - Nsites == 0:
137             return None
138         else:
139             return Spin(s=self._s, N=self.size - Nsites)
140 
141     def __repr__(self):
142         total_sz = (
143             ", total_sz={}".format(self._total_sz) if self._total_sz is not None else ""
144         )
145         return "Spin(s={}{}, N={})".format(Fraction(self._s), total_sz, self._size)
146 
147     @property
148     def _attrs(self):
149         return (self.size, self._s, self._constraint_fn)
150 
151     def _random_state_legacy(self, size=None, *, out=None, rgen=None):
152         if isinstance(size, int):
153             size = (size,)
154         shape = (*size, self._size) if size is not None else (self._size,)
155 
156         if out is None:
157             out = np.empty(shape=shape)
158 
159         if rgen is None:
160             rgen = np.random.default_rng()
161 
162         if self._total_sz is None:
163             out[:] = rgen.choice(self.local_states, size=shape)
164         else:
165             # TODO: this can most likely be done more efficiently
166             if size is not None:
167                 out_r = out.reshape(-1, self._size)
168                 for b in range(out_r.shape[0]):
169                     self._random_state_with_constraint_legacy(out_r[b], rgen)
170             else:
171                 self._random_state_with_constraint_legacy(out, rgen)
172 
173         return out
174 
175     def _random_state_with_constraint_legacy(self, out, rgen):
176         sites = list(range(self.size))
177         out.fill(-round(2 * self._s))
178         ss = self.size
179 
180         for i in range(round(self._s * self.size) + self._total_sz):
181             s = rgen.integers(0, ss, size=())
182 
183             out[sites[s]] += 2
184 
185             if out[sites[s]] > round(2 * self._s - 1):
186                 sites.pop(s)
187                 ss -= 1
188
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/netket/hilbert/spin.py b/netket/hilbert/spin.py
--- a/netket/hilbert/spin.py
+++ b/netket/hilbert/spin.py
@@ -26,20 +26,31 @@
 from ._deprecations import graph_to_N_depwarn
 
 
-def _check_total_sz(total_sz, size):
+def _check_total_sz(total_sz, S, size):
     if total_sz is None:
         return
 
+    local_size = 2 * S + 1
+
     m = round(2 * total_sz)
-    if np.abs(m) > size:
-        raise Exception(
+    if np.abs(m) > size * (2 * S):
+        raise ValueError(
             "Cannot fix the total magnetization: 2|M| cannot " "exceed Nspins."
         )
 
-    if (size + m) % 2 != 0:
-        raise Exception(
-            "Cannot fix the total magnetization: Nspins + " "totalSz must be even."
-        )
+    # If half-integer spins (1/2, 3/2)
+    if local_size % 2 == 0:
+        # Check that the total magnetization is odd if odd spins or even if even # of spins
+        if (size + m) % 2 != 0:
+            raise ValueError(
+                "Cannot fix the total magnetization: Nspins + 2*totalSz must be even."
+            )
+    # else if full-integer (S=1,2)
+    else:
+        if m % 2 != 0:
+            raise ValueError(
+                "Cannot fix the total magnetization to a half-integer number"
+            )
 
 
 @jit(nopython=True)
@@ -85,7 +96,7 @@
             local_states[i] = -round(2 * s) + 2 * i
         local_states = local_states.tolist()
 
-        _check_total_sz(total_sz, N)
+        _check_total_sz(total_sz, s, N)
         if total_sz is not None:
 
             def constraints(x):

{"golden_diff": "diff --git a/netket/hilbert/spin.py b/netket/hilbert/spin.py\n--- a/netket/hilbert/spin.py\n+++ b/netket/hilbert/spin.py\n@@ -26,20 +26,31 @@\n from ._deprecations import graph_to_N_depwarn\n \n \n-def _check_total_sz(total_sz, size):\n+def _check_total_sz(total_sz, S, size):\n     if total_sz is None:\n         return\n \n+    local_size = 2 * S + 1\n+\n     m = round(2 * total_sz)\n-    if np.abs(m) > size:\n-        raise Exception(\n+    if np.abs(m) > size * (2 * S):\n+        raise ValueError(\n             \"Cannot fix the total magnetization: 2|M| cannot \" \"exceed Nspins.\"\n         )\n \n-    if (size + m) % 2 != 0:\n-        raise Exception(\n-            \"Cannot fix the total magnetization: Nspins + \" \"totalSz must be even.\"\n-        )\n+    # If half-integer spins (1/2, 3/2)\n+    if local_size % 2 == 0:\n+        # Check that the total magnetization is odd if odd spins or even if even # of spins\n+        if (size + m) % 2 != 0:\n+            raise ValueError(\n+                \"Cannot fix the total magnetization: Nspins + 2*totalSz must be even.\"\n+            )\n+    # else if full-integer (S=1,2)\n+    else:\n+        if m % 2 != 0:\n+            raise ValueError(\n+                \"Cannot fix the total magnetization to a half-integer number\"\n+            )\n \n \n @jit(nopython=True)\n@@ -85,7 +96,7 @@\n             local_states[i] = -round(2 * s) + 2 * i\n         local_states = local_states.tolist()\n \n-        _check_total_sz(total_sz, N)\n+        _check_total_sz(total_sz, s, N)\n         if total_sz is not None:\n \n             def constraints(x):\n", "issue": "Cannot create Spin hilbert with odd sites and S=1\nReported by @fabienalet \r\n\r\n```python\r\n>>> import netket as nk\r\n>>> nk.hilbert.Spin(1.0, N=3, total_sz=1.0)\r\nTraceback (most recent call last):\r\n  File \"<stdin>\", line 1, in <module>\r\n  File \"/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py\", line 88, in __init__\r\n    _check_total_sz(total_sz, N)\r\n  File \"/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py\", line 40, in _check_total_sz\r\n    raise Exception(\r\nException: Cannot fix the total magnetization: Nspins + totalSz must be even.\r\n```\r\n\r\nthe cuprit is the second check in here: (size is number of spins)\r\n```python\r\ndef _check_total_sz(total_sz, size):\r\n    if total_sz is None:\r\n        return\r\n\r\n    m = round(2 * total_sz)\r\n    if np.abs(m) > size:\r\n        raise Exception(\r\n            \"Cannot fix the total magnetization: 2|M| cannot \" \"exceed Nspins.\"\r\n        )\r\n\r\n    if (size + m) % 2 != 0:\r\n        raise Exception(\r\n            \"Cannot fix the total magnetization: Nspins + \" \"totalSz must be even.\"\r\n        )\r\n```\r\n\r\nThe second check works only for Spins 1/2.\r\nCan we generalise it to always work?\r\nI'm too tired to think how to now.\r\n\nCannot create Spin hilbert with odd sites and S=1\nReported by @fabienalet \r\n\r\n```python\r\n>>> import netket as nk\r\n>>> nk.hilbert.Spin(1.0, N=3, total_sz=1.0)\r\nTraceback (most recent call last):\r\n  File \"<stdin>\", line 1, in <module>\r\n  File \"/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py\", line 88, in __init__\r\n    _check_total_sz(total_sz, N)\r\n  File \"/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py\", line 40, in _check_total_sz\r\n    raise Exception(\r\nException: Cannot fix the total magnetization: Nspins + totalSz must be even.\r\n```\r\n\r\nthe cuprit is the second check in here: (size is number of spins)\r\n```python\r\ndef _check_total_sz(total_sz, size):\r\n    if total_sz is None:\r\n        return\r\n\r\n    m = round(2 * total_sz)\r\n    if np.abs(m) > size:\r\n        raise Exception(\r\n            \"Cannot fix the total magnetization: 2|M| cannot \" \"exceed Nspins.\"\r\n        )\r\n\r\n    if (size + m) % 2 != 0:\r\n        raise Exception(\r\n            \"Cannot fix the total magnetization: Nspins + \" \"totalSz must be even.\"\r\n        )\r\n```\r\n\r\nThe second check works only for Spins 1/2.\r\nCan we generalise it to always work?\r\nI'm too tired to think how to now.\r\n\nCannot create Spin hilbert with odd sites and S=1\nReported by @fabienalet \r\n\r\n```python\r\n>>> import netket as nk\r\n>>> nk.hilbert.Spin(1.0, N=3, total_sz=1.0)\r\nTraceback (most recent call last):\r\n  File \"<stdin>\", line 1, in <module>\r\n  File \"/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py\", line 88, in __init__\r\n    _check_total_sz(total_sz, N)\r\n  File \"/home/filippovicentini/Dropbox/Ricerca/Codes/Python/netket/netket/hilbert/spin.py\", line 40, in _check_total_sz\r\n    raise Exception(\r\nException: Cannot fix the total magnetization: Nspins + totalSz must be even.\r\n```\r\n\r\nthe cuprit is the second check in here: (size is number of spins)\r\n```python\r\ndef _check_total_sz(total_sz, size):\r\n    if total_sz is None:\r\n        return\r\n\r\n    m = round(2 * total_sz)\r\n    if np.abs(m) > size:\r\n        raise Exception(\r\n            \"Cannot fix the total magnetization: 2|M| cannot \" \"exceed Nspins.\"\r\n        )\r\n\r\n    if (size + m) % 2 != 0:\r\n        raise Exception(\r\n            \"Cannot fix the total magnetization: Nspins + \" \"totalSz must be even.\"\r\n        )\r\n```\r\n\r\nThe second check works only for Spins 1/2.\r\nCan we generalise it to always work?\r\nI'm too tired to think how to now.\r\n\n", "before_files": [{"content": "# Copyright 2021 The NetKet Authors - All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom fractions import Fraction\nfrom typing import Optional, List, Union, Iterable\n\nimport jax\nfrom jax import numpy as jnp\nimport numpy as np\nfrom netket.graph import AbstractGraph\nfrom numba import jit\n\n\nfrom .custom_hilbert import CustomHilbert\nfrom ._deprecations import graph_to_N_depwarn\n\n\ndef _check_total_sz(total_sz, size):\n    if total_sz is None:\n        return\n\n    m = round(2 * total_sz)\n    if np.abs(m) > size:\n        raise Exception(\n            \"Cannot fix the total magnetization: 2|M| cannot \" \"exceed Nspins.\"\n        )\n\n    if (size + m) % 2 != 0:\n        raise Exception(\n            \"Cannot fix the total magnetization: Nspins + \" \"totalSz must be even.\"\n        )\n\n\n@jit(nopython=True)\ndef _sum_constraint(x, total_sz):\n    return np.sum(x, axis=1) == round(2 * total_sz)\n\n\nclass Spin(CustomHilbert):\n    r\"\"\"Hilbert space obtained as tensor product of local spin states.\"\"\"\n\n    def __init__(\n        self,\n        s: float,\n        N: int = 1,\n        total_sz: Optional[float] = None,\n        graph: Optional[AbstractGraph] = None,\n    ):\n        r\"\"\"Hilbert space obtained as tensor product of local spin states.\n\n        Args:\n           s: Spin at each site. Must be integer or half-integer.\n           N: Number of sites (default=1)\n           total_sz: If given, constrains the total spin of system to a particular value.\n           graph: (deprecated) a graph from which to extract the number of sites.\n\n        Examples:\n           Simple spin hilbert space.\n\n           >>> from netket.hilbert import Spin\n           >>> g = Hypercube(length=10,n_dim=2,pbc=True)\n           >>> hi = Spin(s=0.5, N=4)\n           >>> print(hi.size)\n           4\n        \"\"\"\n        N = graph_to_N_depwarn(N=N, graph=graph)\n\n        local_size = round(2 * s + 1)\n        local_states = np.empty(local_size)\n\n        assert int(2 * s + 1) == local_size\n\n        for i in range(local_size):\n            local_states[i] = -round(2 * s) + 2 * i\n        local_states = local_states.tolist()\n\n        _check_total_sz(total_sz, N)\n        if total_sz is not None:\n\n            def constraints(x):\n                return _sum_constraint(x, total_sz)\n\n        else:\n            constraints = None\n\n        self._total_sz = total_sz if total_sz is None else int(total_sz)\n        self._s = s\n        self._local_size = local_size\n\n        super().__init__(local_states, N, constraints)\n\n    def __pow__(self, n):\n        if self._total_sz is None:\n            total_sz = None\n        else:\n            total_sz = total_sz * n\n\n        return Spin(self._s, self.size * n, total_sz=total_sz)\n\n    def _mul_sametype_(self, other):\n        assert type(self) == type(other)\n        if self._s == other._s:\n            if self._total_sz is None and other._total_sz is None:\n                return Spin(s=self._s, N=self.size + other.size)\n\n        return NotImplemented\n\n    def ptrace(self, sites: Union[int, List]) -> Optional[\"Spin\"]:\n        if isinstance(sites, int):\n            sites = [sites]\n\n        for site in sites:\n            if site < 0 or site >= self.size:\n                raise ValueError(\n                    f\"Site {site} not in this hilbert space of site {self.size}\"\n                )\n\n        if self._total_sz is not None:\n            raise TypeError(\n                \"Cannot take the partial trace with a total magnetization constraint.\"\n            )\n\n        Nsites = len(sites)\n\n        if self.size - Nsites == 0:\n            return None\n        else:\n            return Spin(s=self._s, N=self.size - Nsites)\n\n    def __repr__(self):\n        total_sz = (\n            \", total_sz={}\".format(self._total_sz) if self._total_sz is not None else \"\"\n        )\n        return \"Spin(s={}{}, N={})\".format(Fraction(self._s), total_sz, self._size)\n\n    @property\n    def _attrs(self):\n        return (self.size, self._s, self._constraint_fn)\n\n    def _random_state_legacy(self, size=None, *, out=None, rgen=None):\n        if isinstance(size, int):\n            size = (size,)\n        shape = (*size, self._size) if size is not None else (self._size,)\n\n        if out is None:\n            out = np.empty(shape=shape)\n\n        if rgen is None:\n            rgen = np.random.default_rng()\n\n        if self._total_sz is None:\n            out[:] = rgen.choice(self.local_states, size=shape)\n        else:\n            # TODO: this can most likely be done more efficiently\n            if size is not None:\n                out_r = out.reshape(-1, self._size)\n                for b in range(out_r.shape[0]):\n                    self._random_state_with_constraint_legacy(out_r[b], rgen)\n            else:\n                self._random_state_with_constraint_legacy(out, rgen)\n\n        return out\n\n    def _random_state_with_constraint_legacy(self, out, rgen):\n        sites = list(range(self.size))\n        out.fill(-round(2 * self._s))\n        ss = self.size\n\n        for i in range(round(self._s * self.size) + self._total_sz):\n            s = rgen.integers(0, ss, size=())\n\n            out[sites[s]] += 2\n\n            if out[sites[s]] > round(2 * self._s - 1):\n                sites.pop(s)\n                ss -= 1\n", "path": "netket/hilbert/spin.py"}], "after_files": [{"content": "# Copyright 2021 The NetKet Authors - All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#    http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom fractions import Fraction\nfrom typing import Optional, List, Union, Iterable\n\nimport jax\nfrom jax import numpy as jnp\nimport numpy as np\nfrom netket.graph import AbstractGraph\nfrom numba import jit\n\n\nfrom .custom_hilbert import CustomHilbert\nfrom ._deprecations import graph_to_N_depwarn\n\n\ndef _check_total_sz(total_sz, S, size):\n    if total_sz is None:\n        return\n\n    local_size = 2 * S + 1\n\n    m = round(2 * total_sz)\n    if np.abs(m) > size * (2 * S):\n        raise ValueError(\n            \"Cannot fix the total magnetization: 2|M| cannot \" \"exceed Nspins.\"\n        )\n\n    # If half-integer spins (1/2, 3/2)\n    if local_size % 2 == 0:\n        # Check that the total magnetization is odd if odd spins or even if even # of spins\n        if (size + m) % 2 != 0:\n            raise ValueError(\n                \"Cannot fix the total magnetization: Nspins + 2*totalSz must be even.\"\n            )\n    # else if full-integer (S=1,2)\n    else:\n        if m % 2 != 0:\n            raise ValueError(\n                \"Cannot fix the total magnetization to a half-integer number\"\n            )\n\n\n@jit(nopython=True)\ndef _sum_constraint(x, total_sz):\n    return np.sum(x, axis=1) == round(2 * total_sz)\n\n\nclass Spin(CustomHilbert):\n    r\"\"\"Hilbert space obtained as tensor product of local spin states.\"\"\"\n\n    def __init__(\n        self,\n        s: float,\n        N: int = 1,\n        total_sz: Optional[float] = None,\n        graph: Optional[AbstractGraph] = None,\n    ):\n        r\"\"\"Hilbert space obtained as tensor product of local spin states.\n\n        Args:\n           s: Spin at each site. Must be integer or half-integer.\n           N: Number of sites (default=1)\n           total_sz: If given, constrains the total spin of system to a particular value.\n           graph: (deprecated) a graph from which to extract the number of sites.\n\n        Examples:\n           Simple spin hilbert space.\n\n           >>> from netket.hilbert import Spin\n           >>> g = Hypercube(length=10,n_dim=2,pbc=True)\n           >>> hi = Spin(s=0.5, N=4)\n           >>> print(hi.size)\n           4\n        \"\"\"\n        N = graph_to_N_depwarn(N=N, graph=graph)\n\n        local_size = round(2 * s + 1)\n        local_states = np.empty(local_size)\n\n        assert int(2 * s + 1) == local_size\n\n        for i in range(local_size):\n            local_states[i] = -round(2 * s) + 2 * i\n        local_states = local_states.tolist()\n\n        _check_total_sz(total_sz, s, N)\n        if total_sz is not None:\n\n            def constraints(x):\n                return _sum_constraint(x, total_sz)\n\n        else:\n            constraints = None\n\n        self._total_sz = total_sz if total_sz is None else int(total_sz)\n        self._s = s\n        self._local_size = local_size\n\n        super().__init__(local_states, N, constraints)\n\n    def __pow__(self, n):\n        if self._total_sz is None:\n            total_sz = None\n        else:\n            total_sz = total_sz * n\n\n        return Spin(self._s, self.size * n, total_sz=total_sz)\n\n    def _mul_sametype_(self, other):\n        assert type(self) == type(other)\n        if self._s == other._s:\n            if self._total_sz is None and other._total_sz is None:\n                return Spin(s=self._s, N=self.size + other.size)\n\n        return NotImplemented\n\n    def ptrace(self, sites: Union[int, List]) -> Optional[\"Spin\"]:\n        if isinstance(sites, int):\n            sites = [sites]\n\n        for site in sites:\n            if site < 0 or site >= self.size:\n                raise ValueError(\n                    f\"Site {site} not in this hilbert space of site {self.size}\"\n                )\n\n        if self._total_sz is not None:\n            raise TypeError(\n                \"Cannot take the partial trace with a total magnetization constraint.\"\n            )\n\n        Nsites = len(sites)\n\n        if self.size - Nsites == 0:\n            return None\n        else:\n            return Spin(s=self._s, N=self.size - Nsites)\n\n    def __repr__(self):\n        total_sz = (\n            \", total_sz={}\".format(self._total_sz) if self._total_sz is not None else \"\"\n        )\n        return \"Spin(s={}{}, N={})\".format(Fraction(self._s), total_sz, self._size)\n\n    @property\n    def _attrs(self):\n        return (self.size, self._s, self._constraint_fn)\n\n    def _random_state_legacy(self, size=None, *, out=None, rgen=None):\n        if isinstance(size, int):\n            size = (size,)\n        shape = (*size, self._size) if size is not None else (self._size,)\n\n        if out is None:\n            out = np.empty(shape=shape)\n\n        if rgen is None:\n            rgen = np.random.default_rng()\n\n        if self._total_sz is None:\n            out[:] = rgen.choice(self.local_states, size=shape)\n        else:\n            # TODO: this can most likely be done more efficiently\n            if size is not None:\n                out_r = out.reshape(-1, self._size)\n                for b in range(out_r.shape[0]):\n                    self._random_state_with_constraint_legacy(out_r[b], rgen)\n            else:\n                self._random_state_with_constraint_legacy(out, rgen)\n\n        return out\n\n    def _random_state_with_constraint_legacy(self, out, rgen):\n        sites = list(range(self.size))\n        out.fill(-round(2 * self._s))\n        ss = self.size\n\n        for i in range(round(self._s * self.size) + self._total_sz):\n            s = rgen.integers(0, ss, size=())\n\n            out[sites[s]] += 2\n\n            if out[sites[s]] > round(2 * self._s - 1):\n                sites.pop(s)\n                ss -= 1\n", "path": "netket/hilbert/spin.py"}]}