NetKet 3.15.1__tar.gz → 3.16__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (330) hide show
  1. {netket-3.15.1 → netket-3.16}/.gitignore +1 -0
  2. {netket-3.15.1 → netket-3.16}/PKG-INFO +10 -9
  3. {netket-3.15.1 → netket-3.16}/netket/_version.py +2 -2
  4. {netket-3.15.1 → netket-3.16}/netket/errors.py +36 -1
  5. {netket-3.15.1 → netket-3.16}/netket/experimental/driver/tdvp.py +0 -9
  6. {netket-3.15.1 → netket-3.16}/netket/experimental/driver/tdvp_common.py +5 -52
  7. {netket-3.15.1 → netket-3.16}/netket/experimental/driver/tdvp_schmitt.py +10 -12
  8. {netket-3.15.1 → netket-3.16}/netket/experimental/dynamics/_integrator.py +3 -7
  9. {netket-3.15.1 → netket-3.16}/netket/experimental/dynamics/_integrator_params.py +2 -3
  10. {netket-3.15.1 → netket-3.16}/netket/experimental/dynamics/_rk/_solver.py +1 -1
  11. {netket-3.15.1 → netket-3.16}/netket/experimental/dynamics/_solver.py +1 -3
  12. {netket-3.15.1 → netket-3.16}/netket/experimental/dynamics/_utils.py +1 -26
  13. {netket-3.15.1 → netket-3.16}/netket/experimental/hilbert/__init__.py +16 -1
  14. netket-3.16/netket/experimental/models/__init__.py +38 -0
  15. netket-3.16/netket/experimental/operator/__init__.py +47 -0
  16. netket-3.16/netket/experimental/operator/fermion.py +44 -0
  17. {netket-3.15.1 → netket-3.16}/netket/experimental/operator/pyscf.py +2 -2
  18. {netket-3.15.1 → netket-3.16}/netket/experimental/sampler/__init__.py +8 -2
  19. netket-3.16/netket/experimental/sampler/rules/__init__.py +17 -0
  20. {netket-3.15.1 → netket-3.16}/netket/graph/common_lattices.py +31 -2
  21. {netket-3.15.1 → netket-3.16}/netket/graph/lattice.py +2 -1
  22. {netket-3.15.1 → netket-3.16}/netket/hilbert/__init__.py +1 -0
  23. {netket-3.15.1 → netket-3.16}/netket/hilbert/abstract_hilbert.py +4 -5
  24. {netket-3.15.1 → netket-3.16}/netket/hilbert/constraint/base.py +1 -1
  25. {netket-3.15.1 → netket-3.16}/netket/hilbert/custom_hilbert.py +2 -2
  26. {netket-3.15.1 → netket-3.16}/netket/hilbert/doubled_hilbert.py +19 -2
  27. {netket-3.15.1 → netket-3.16}/netket/hilbert/fock.py +0 -1
  28. {netket-3.15.1 → netket-3.16}/netket/hilbert/homogeneous.py +5 -7
  29. {netket-3.15.1 → netket-3.16}/netket/hilbert/index/constrained_generic.py +1 -1
  30. {netket-3.15.1 → netket-3.16}/netket/hilbert/random/base.py +3 -3
  31. {netket-3.15.1 → netket-3.16}/netket/hilbert/random/qubit.py +3 -0
  32. {netket-3.15.1 → netket-3.16}/netket/hilbert/spin.py +8 -57
  33. {netket-3.15.1/netket/experimental → netket-3.16/netket}/hilbert/spin_orbital_fermions.py +1 -1
  34. {netket-3.15.1 → netket-3.16}/netket/hilbert/tensor_hilbert_discrete.py +1 -1
  35. {netket-3.15.1 → netket-3.16}/netket/jax/_utils_dtype.py +43 -1
  36. {netket-3.15.1 → netket-3.16}/netket/jax/_utils_random.py +1 -1
  37. {netket-3.15.1 → netket-3.16}/netket/jax/_vjp_chunked.py +4 -4
  38. {netket-3.15.1 → netket-3.16}/netket/jax/sharding.py +1 -1
  39. {netket-3.15.1 → netket-3.16}/netket/logging/tensorboard.py +2 -2
  40. {netket-3.15.1 → netket-3.16}/netket/models/__init__.py +2 -0
  41. {netket-3.15.1 → netket-3.16}/netket/models/slater.py +1 -1
  42. {netket-3.15.1 → netket-3.16}/netket/nn/blocks/mlp.py +1 -1
  43. {netket-3.15.1 → netket-3.16}/netket/operator/__init__.py +3 -1
  44. {netket-3.15.1 → netket-3.16}/netket/operator/_bose_hubbard/jax.py +15 -3
  45. {netket-3.15.1 → netket-3.16}/netket/operator/_discrete_operator_jax.py +17 -4
  46. netket-3.16/netket/operator/_fermion2nd/__init__.py +19 -0
  47. netket-3.15.1/netket/experimental/operator/_fermion_operator_2nd_base.py → netket-3.16/netket/operator/_fermion2nd/base.py +6 -6
  48. netket-3.15.1/netket/experimental/operator/_fermion_operator_2nd_jax.py → netket-3.16/netket/operator/_fermion2nd/jax.py +6 -5
  49. netket-3.15.1/netket/experimental/operator/_fermion_operator_2nd_numba.py → netket-3.16/netket/operator/_fermion2nd/numba.py +7 -7
  50. {netket-3.15.1 → netket-3.16}/netket/operator/_local_operator/convert.py +2 -3
  51. {netket-3.15.1 → netket-3.16}/netket/operator/_pauli_strings/base.py +1 -1
  52. {netket-3.15.1/netket/experimental → netket-3.16/netket}/operator/fermion.py +14 -8
  53. {netket-3.15.1 → netket-3.16}/netket/optimizer/solver/solvers.py +1 -6
  54. {netket-3.15.1 → netket-3.16}/netket/optimizer/sr.py +2 -2
  55. {netket-3.15.1 → netket-3.16}/netket/sampler/__init__.py +2 -1
  56. {netket-3.15.1 → netket-3.16}/netket/sampler/autoreg.py +1 -1
  57. {netket-3.15.1 → netket-3.16}/netket/sampler/base.py +13 -2
  58. {netket-3.15.1 → netket-3.16}/netket/sampler/exact.py +1 -1
  59. {netket-3.15.1 → netket-3.16}/netket/sampler/metropolis.py +67 -8
  60. {netket-3.15.1 → netket-3.16}/netket/sampler/metropolis_numpy.py +9 -0
  61. netket-3.15.1/netket/sampler/metropolis_pt.py → netket-3.16/netket/sampler/parallel_tempering.py +7 -9
  62. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/__init__.py +1 -0
  63. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/exchange.py +1 -1
  64. {netket-3.15.1/netket/experimental → netket-3.16/netket}/sampler/rules/fermion_2nd.py +8 -8
  65. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/hamiltonian.py +6 -6
  66. {netket-3.15.1 → netket-3.16}/netket/utils/_dependencies_check.py +6 -6
  67. {netket-3.15.1 → netket-3.16}/netket/utils/config_flags.py +3 -0
  68. {netket-3.15.1 → netket-3.16}/netket/utils/history/accum.py +2 -1
  69. {netket-3.15.1 → netket-3.16}/netket/utils/history/history.py +19 -3
  70. {netket-3.15.1 → netket-3.16}/netket/utils/history/history_dict.py +6 -5
  71. {netket-3.15.1 → netket-3.16}/netket/utils/model_frameworks/__init__.py +3 -1
  72. {netket-3.15.1 → netket-3.16}/netket/utils/model_frameworks/base.py +15 -2
  73. {netket-3.15.1 → netket-3.16}/netket/utils/model_frameworks/equinox.py +14 -4
  74. netket-3.16/netket/utils/model_frameworks/nnx.py +132 -0
  75. netket-3.16/netket/utils/model_frameworks/nnx_wrapped.py +55 -0
  76. {netket-3.15.1 → netket-3.16}/netket/utils/mpi/mpi.py +1 -1
  77. {netket-3.15.1 → netket-3.16}/netket/utils/static_range.py +18 -7
  78. {netket-3.15.1 → netket-3.16}/netket/utils/struct/pytree.py +40 -33
  79. {netket-3.15.1 → netket-3.16}/netket/vqs/base.py +3 -0
  80. {netket-3.15.1 → netket-3.16}/netket/vqs/full_summ/state.py +16 -0
  81. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_state/state.py +56 -8
  82. {netket-3.15.1 → netket-3.16}/pyproject.toml +38 -50
  83. netket-3.15.1/netket/experimental/models/__init__.py +0 -18
  84. netket-3.15.1/netket/experimental/operator/__init__.py +0 -26
  85. netket-3.15.1/netket/experimental/sampler/metropolis.py +0 -40
  86. netket-3.15.1/netket/experimental/sampler/rules/__init__.py +0 -1
  87. {netket-3.15.1 → netket-3.16}/LICENSE +0 -0
  88. {netket-3.15.1 → netket-3.16}/README.md +0 -0
  89. {netket-3.15.1 → netket-3.16}/hatch.toml +0 -0
  90. {netket-3.15.1 → netket-3.16}/netket/__init__.py +0 -0
  91. {netket-3.15.1 → netket-3.16}/netket/callbacks/__init__.py +0 -0
  92. {netket-3.15.1 → netket-3.16}/netket/callbacks/convergence_stopping.py +0 -0
  93. {netket-3.15.1 → netket-3.16}/netket/callbacks/early_stopping.py +0 -0
  94. {netket-3.15.1 → netket-3.16}/netket/callbacks/invalid_loss_stopping.py +0 -0
  95. {netket-3.15.1 → netket-3.16}/netket/callbacks/timeout.py +0 -0
  96. {netket-3.15.1 → netket-3.16}/netket/driver/__init__.py +0 -0
  97. {netket-3.15.1 → netket-3.16}/netket/driver/abstract_variational_driver.py +0 -0
  98. {netket-3.15.1 → netket-3.16}/netket/driver/steady_state.py +0 -0
  99. {netket-3.15.1 → netket-3.16}/netket/driver/vmc.py +0 -0
  100. {netket-3.15.1 → netket-3.16}/netket/driver/vmc_common.py +0 -0
  101. {netket-3.15.1 → netket-3.16}/netket/exact.py +0 -0
  102. {netket-3.15.1 → netket-3.16}/netket/experimental/__init__.py +0 -0
  103. {netket-3.15.1 → netket-3.16}/netket/experimental/driver/__init__.py +0 -0
  104. {netket-3.15.1 → netket-3.16}/netket/experimental/driver/vmc_srt.py +0 -0
  105. {netket-3.15.1 → netket-3.16}/netket/experimental/dynamics/__init__.py +0 -0
  106. {netket-3.15.1 → netket-3.16}/netket/experimental/dynamics/_integrator_state.py +0 -0
  107. {netket-3.15.1 → netket-3.16}/netket/experimental/dynamics/_rk/__init__.py +0 -0
  108. {netket-3.15.1 → netket-3.16}/netket/experimental/dynamics/_rk/_tableau.py +0 -0
  109. {netket-3.15.1 → netket-3.16}/netket/experimental/logging/__init__.py +0 -0
  110. {netket-3.15.1 → netket-3.16}/netket/experimental/logging/hdf5_log.py +0 -0
  111. {netket-3.15.1 → netket-3.16}/netket/experimental/models/fast_rnn.py +0 -0
  112. {netket-3.15.1 → netket-3.16}/netket/experimental/models/rnn.py +0 -0
  113. {netket-3.15.1 → netket-3.16}/netket/experimental/nn/__init__.py +0 -0
  114. {netket-3.15.1 → netket-3.16}/netket/experimental/nn/rnn/__init__.py +0 -0
  115. {netket-3.15.1 → netket-3.16}/netket/experimental/nn/rnn/cells.py +0 -0
  116. {netket-3.15.1 → netket-3.16}/netket/experimental/nn/rnn/layers.py +0 -0
  117. {netket-3.15.1 → netket-3.16}/netket/experimental/nn/rnn/layers_fast.py +0 -0
  118. {netket-3.15.1 → netket-3.16}/netket/experimental/nn/rnn/ordering.py +0 -0
  119. {netket-3.15.1 → netket-3.16}/netket/experimental/observable/__init__.py +0 -0
  120. {netket-3.15.1 → netket-3.16}/netket/experimental/observable/renyi2/S2_operator.py +0 -0
  121. {netket-3.15.1 → netket-3.16}/netket/experimental/observable/renyi2/__init__.py +0 -0
  122. {netket-3.15.1 → netket-3.16}/netket/experimental/observable/renyi2/exact.py +0 -0
  123. {netket-3.15.1 → netket-3.16}/netket/experimental/observable/renyi2/expect.py +0 -0
  124. {netket-3.15.1 → netket-3.16}/netket/experimental/observable/variance/__init__.py +0 -0
  125. {netket-3.15.1 → netket-3.16}/netket/experimental/observable/variance/exact.py +0 -0
  126. {netket-3.15.1 → netket-3.16}/netket/experimental/observable/variance/expect.py +0 -0
  127. {netket-3.15.1 → netket-3.16}/netket/experimental/observable/variance/variance_operator.py +0 -0
  128. {netket-3.15.1 → netket-3.16}/netket/experimental/qsr/__init__.py +0 -0
  129. {netket-3.15.1 → netket-3.16}/netket/experimental/qsr/dataset.py +0 -0
  130. {netket-3.15.1 → netket-3.16}/netket/experimental/qsr/driver.py +0 -0
  131. {netket-3.15.1 → netket-3.16}/netket/experimental/qsr/logic_helpers.py +0 -0
  132. {netket-3.15.1 → netket-3.16}/netket/experimental/sampler/metropolis_pmap.py +0 -0
  133. {netket-3.15.1 → netket-3.16}/netket/experimental/vqs/__init__.py +0 -0
  134. {netket-3.15.1 → netket-3.16}/netket/experimental/vqs/io.py +0 -0
  135. {netket-3.15.1 → netket-3.16}/netket/graph/__init__.py +0 -0
  136. {netket-3.15.1 → netket-3.16}/netket/graph/_lattice_draw.py +0 -0
  137. {netket-3.15.1 → netket-3.16}/netket/graph/_lattice_edge_logic.py +0 -0
  138. {netket-3.15.1 → netket-3.16}/netket/graph/abstract_graph.py +0 -0
  139. {netket-3.15.1 → netket-3.16}/netket/graph/graph.py +0 -0
  140. {netket-3.15.1 → netket-3.16}/netket/graph/space_group.py +0 -0
  141. {netket-3.15.1 → netket-3.16}/netket/hilbert/constraint/__init__.py +0 -0
  142. {netket-3.15.1 → netket-3.16}/netket/hilbert/constraint/multiple.py +0 -0
  143. {netket-3.15.1 → netket-3.16}/netket/hilbert/constraint/sum.py +0 -0
  144. {netket-3.15.1 → netket-3.16}/netket/hilbert/constraint/sum_partition.py +0 -0
  145. {netket-3.15.1 → netket-3.16}/netket/hilbert/continuous_hilbert.py +0 -0
  146. {netket-3.15.1 → netket-3.16}/netket/hilbert/discrete_hilbert.py +0 -0
  147. {netket-3.15.1 → netket-3.16}/netket/hilbert/index/__init__.py +0 -0
  148. {netket-3.15.1 → netket-3.16}/netket/hilbert/index/base.py +0 -0
  149. {netket-3.15.1 → netket-3.16}/netket/hilbert/index/constrained_sum.py +0 -0
  150. {netket-3.15.1 → netket-3.16}/netket/hilbert/index/constrained_sum_partitions.py +0 -0
  151. {netket-3.15.1 → netket-3.16}/netket/hilbert/index/unconstrained.py +0 -0
  152. {netket-3.15.1 → netket-3.16}/netket/hilbert/index/uniform_tensor.py +0 -0
  153. {netket-3.15.1 → netket-3.16}/netket/hilbert/particle.py +0 -0
  154. {netket-3.15.1 → netket-3.16}/netket/hilbert/qubit.py +0 -0
  155. {netket-3.15.1 → netket-3.16}/netket/hilbert/random/__init__.py +0 -0
  156. {netket-3.15.1 → netket-3.16}/netket/hilbert/random/custom.py +0 -0
  157. {netket-3.15.1 → netket-3.16}/netket/hilbert/random/doubled.py +0 -0
  158. {netket-3.15.1 → netket-3.16}/netket/hilbert/random/fock.py +0 -0
  159. {netket-3.15.1 → netket-3.16}/netket/hilbert/random/homogeneous.py +0 -0
  160. {netket-3.15.1 → netket-3.16}/netket/hilbert/random/particle.py +0 -0
  161. {netket-3.15.1 → netket-3.16}/netket/hilbert/random/tensor_hilbert.py +0 -0
  162. {netket-3.15.1 → netket-3.16}/netket/hilbert/tensor_hilbert.py +0 -0
  163. {netket-3.15.1 → netket-3.16}/netket/jax/__init__.py +0 -0
  164. {netket-3.15.1 → netket-3.16}/netket/jax/_chunk_utils.py +0 -0
  165. {netket-3.15.1 → netket-3.16}/netket/jax/_expect.py +0 -0
  166. {netket-3.15.1 → netket-3.16}/netket/jax/_grad.py +0 -0
  167. {netket-3.15.1 → netket-3.16}/netket/jax/_jacobian/__init__.py +0 -0
  168. {netket-3.15.1 → netket-3.16}/netket/jax/_jacobian/default_mode.py +0 -0
  169. {netket-3.15.1 → netket-3.16}/netket/jax/_jacobian/jacobian_dense.py +0 -0
  170. {netket-3.15.1 → netket-3.16}/netket/jax/_jacobian/jacobian_pytree.py +0 -0
  171. {netket-3.15.1 → netket-3.16}/netket/jax/_jacobian/logic.py +0 -0
  172. {netket-3.15.1 → netket-3.16}/netket/jax/_math.py +0 -0
  173. {netket-3.15.1 → netket-3.16}/netket/jax/_scanmap.py +0 -0
  174. {netket-3.15.1 → netket-3.16}/netket/jax/_sort.py +0 -0
  175. {netket-3.15.1 → netket-3.16}/netket/jax/_utils_tree.py +0 -0
  176. {netket-3.15.1 → netket-3.16}/netket/jax/_vjp.py +0 -0
  177. {netket-3.15.1 → netket-3.16}/netket/jax/_vmap_chunked.py +0 -0
  178. {netket-3.15.1 → netket-3.16}/netket/logging/__init__.py +0 -0
  179. {netket-3.15.1 → netket-3.16}/netket/logging/base.py +0 -0
  180. {netket-3.15.1 → netket-3.16}/netket/logging/json_log.py +0 -0
  181. {netket-3.15.1 → netket-3.16}/netket/logging/runtime_log.py +0 -0
  182. {netket-3.15.1 → netket-3.16}/netket/logging/state_log.py +0 -0
  183. {netket-3.15.1 → netket-3.16}/netket/models/README.md +0 -0
  184. {netket-3.15.1 → netket-3.16}/netket/models/autoreg.py +0 -0
  185. {netket-3.15.1 → netket-3.16}/netket/models/deepset.py +0 -0
  186. {netket-3.15.1 → netket-3.16}/netket/models/equivariant.py +0 -0
  187. {netket-3.15.1 → netket-3.16}/netket/models/fast_autoreg.py +0 -0
  188. {netket-3.15.1 → netket-3.16}/netket/models/full_space.py +0 -0
  189. {netket-3.15.1 → netket-3.16}/netket/models/gaussian.py +0 -0
  190. {netket-3.15.1 → netket-3.16}/netket/models/jastrow.py +0 -0
  191. {netket-3.15.1 → netket-3.16}/netket/models/mlp.py +0 -0
  192. {netket-3.15.1 → netket-3.16}/netket/models/ndm.py +0 -0
  193. {netket-3.15.1 → netket-3.16}/netket/models/rbm.py +0 -0
  194. {netket-3.15.1 → netket-3.16}/netket/models/tensor_networks/__init__.py +0 -0
  195. {netket-3.15.1 → netket-3.16}/netket/models/tensor_networks/mpdo.py +0 -0
  196. {netket-3.15.1 → netket-3.16}/netket/models/tensor_networks/mps.py +0 -0
  197. {netket-3.15.1 → netket-3.16}/netket/models/utils.py +0 -0
  198. {netket-3.15.1 → netket-3.16}/netket/nn/__init__.py +0 -0
  199. {netket-3.15.1 → netket-3.16}/netket/nn/activation.py +0 -0
  200. {netket-3.15.1 → netket-3.16}/netket/nn/blocks/__init__.py +0 -0
  201. {netket-3.15.1 → netket-3.16}/netket/nn/blocks/deepset.py +0 -0
  202. {netket-3.15.1 → netket-3.16}/netket/nn/blocks/symmetry_sum.py +0 -0
  203. {netket-3.15.1 → netket-3.16}/netket/nn/fast_masked_linear.py +0 -0
  204. {netket-3.15.1 → netket-3.16}/netket/nn/masked_linear.py +0 -0
  205. {netket-3.15.1 → netket-3.16}/netket/nn/symmetric_linear.py +0 -0
  206. {netket-3.15.1 → netket-3.16}/netket/nn/utils.py +0 -0
  207. {netket-3.15.1 → netket-3.16}/netket/operator/_abstract_observable.py +0 -0
  208. {netket-3.15.1 → netket-3.16}/netket/operator/_abstract_operator.py +0 -0
  209. {netket-3.15.1 → netket-3.16}/netket/operator/_abstract_super_operator.py +0 -0
  210. {netket-3.15.1 → netket-3.16}/netket/operator/_bose_hubbard/__init__.py +0 -0
  211. {netket-3.15.1 → netket-3.16}/netket/operator/_bose_hubbard/base.py +0 -0
  212. {netket-3.15.1 → netket-3.16}/netket/operator/_bose_hubbard/numba.py +0 -0
  213. {netket-3.15.1 → netket-3.16}/netket/operator/_continuous_operator.py +0 -0
  214. {netket-3.15.1 → netket-3.16}/netket/operator/_discrete_operator.py +0 -0
  215. /netket-3.15.1/netket/experimental/operator/_fermion_operator_2nd_utils.py → /netket-3.16/netket/operator/_fermion2nd/utils.py +0 -0
  216. {netket-3.15.1 → netket-3.16}/netket/operator/_graph_operator.py +0 -0
  217. {netket-3.15.1 → netket-3.16}/netket/operator/_hamiltonian.py +0 -0
  218. {netket-3.15.1 → netket-3.16}/netket/operator/_heisenberg.py +0 -0
  219. {netket-3.15.1 → netket-3.16}/netket/operator/_ising/__init__.py +0 -0
  220. {netket-3.15.1 → netket-3.16}/netket/operator/_ising/base.py +0 -0
  221. {netket-3.15.1 → netket-3.16}/netket/operator/_ising/jax.py +0 -0
  222. {netket-3.15.1 → netket-3.16}/netket/operator/_ising/numba.py +0 -0
  223. {netket-3.15.1 → netket-3.16}/netket/operator/_kinetic.py +0 -0
  224. {netket-3.15.1 → netket-3.16}/netket/operator/_lazy.py +0 -0
  225. {netket-3.15.1 → netket-3.16}/netket/operator/_local_liouvillian.py +0 -0
  226. {netket-3.15.1 → netket-3.16}/netket/operator/_local_operator/__init__.py +0 -0
  227. {netket-3.15.1 → netket-3.16}/netket/operator/_local_operator/base.py +0 -0
  228. {netket-3.15.1 → netket-3.16}/netket/operator/_local_operator/compile_helpers.py +0 -0
  229. {netket-3.15.1 → netket-3.16}/netket/operator/_local_operator/helpers.py +0 -0
  230. {netket-3.15.1 → netket-3.16}/netket/operator/_local_operator/jax.py +0 -0
  231. {netket-3.15.1 → netket-3.16}/netket/operator/_local_operator/numba.py +0 -0
  232. {netket-3.15.1 → netket-3.16}/netket/operator/_pauli_strings/__init__.py +0 -0
  233. {netket-3.15.1 → netket-3.16}/netket/operator/_pauli_strings/jax.py +0 -0
  234. {netket-3.15.1 → netket-3.16}/netket/operator/_pauli_strings/numba.py +0 -0
  235. {netket-3.15.1 → netket-3.16}/netket/operator/_potential.py +0 -0
  236. {netket-3.15.1 → netket-3.16}/netket/operator/_sumoperators.py +0 -0
  237. {netket-3.15.1 → netket-3.16}/netket/operator/boson.py +0 -0
  238. {netket-3.15.1 → netket-3.16}/netket/operator/spin.py +0 -0
  239. {netket-3.15.1 → netket-3.16}/netket/optimizer/__init__.py +0 -0
  240. {netket-3.15.1 → netket-3.16}/netket/optimizer/linear_operator.py +0 -0
  241. {netket-3.15.1 → netket-3.16}/netket/optimizer/preconditioner.py +0 -0
  242. {netket-3.15.1 → netket-3.16}/netket/optimizer/qgt/__init__.py +0 -0
  243. {netket-3.15.1 → netket-3.16}/netket/optimizer/qgt/common.py +0 -0
  244. {netket-3.15.1 → netket-3.16}/netket/optimizer/qgt/default.py +0 -0
  245. {netket-3.15.1 → netket-3.16}/netket/optimizer/qgt/qgt_jacobian.py +0 -0
  246. {netket-3.15.1 → netket-3.16}/netket/optimizer/qgt/qgt_jacobian_common.py +0 -0
  247. {netket-3.15.1 → netket-3.16}/netket/optimizer/qgt/qgt_jacobian_dense.py +0 -0
  248. {netket-3.15.1 → netket-3.16}/netket/optimizer/qgt/qgt_jacobian_pytree.py +0 -0
  249. {netket-3.15.1 → netket-3.16}/netket/optimizer/qgt/qgt_onthefly.py +0 -0
  250. {netket-3.15.1 → netket-3.16}/netket/optimizer/qgt/qgt_onthefly_logic.py +0 -0
  251. {netket-3.15.1 → netket-3.16}/netket/optimizer/solver/__init__.py +0 -0
  252. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/base.py +0 -0
  253. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/continuous_gaussian.py +0 -0
  254. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/custom_numpy.py +0 -0
  255. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/fixed.py +0 -0
  256. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/hamiltonian_numpy.py +0 -0
  257. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/langevin.py +0 -0
  258. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/local.py +0 -0
  259. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/local_numpy.py +0 -0
  260. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/multiple.py +0 -0
  261. {netket-3.15.1 → netket-3.16}/netket/sampler/rules/tensor.py +0 -0
  262. {netket-3.15.1 → netket-3.16}/netket/stats/__init__.py +0 -0
  263. {netket-3.15.1 → netket-3.16}/netket/stats/_autocorr.py +0 -0
  264. {netket-3.15.1 → netket-3.16}/netket/stats/mc_stats.py +0 -0
  265. {netket-3.15.1 → netket-3.16}/netket/stats/mc_stats_old.py +0 -0
  266. {netket-3.15.1 → netket-3.16}/netket/stats/mpi_stats.py +0 -0
  267. {netket-3.15.1 → netket-3.16}/netket/tools/__init__.py +0 -0
  268. {netket-3.15.1 → netket-3.16}/netket/tools/_common.py +0 -0
  269. {netket-3.15.1 → netket-3.16}/netket/tools/_cpu_info.py +0 -0
  270. {netket-3.15.1 → netket-3.16}/netket/tools/_mpi_info.py +0 -0
  271. {netket-3.15.1 → netket-3.16}/netket/tools/check_mpi.py +0 -0
  272. {netket-3.15.1 → netket-3.16}/netket/tools/info.py +0 -0
  273. {netket-3.15.1 → netket-3.16}/netket/utils/__init__.py +0 -0
  274. {netket-3.15.1 → netket-3.16}/netket/utils/api_utils.py +0 -0
  275. {netket-3.15.1 → netket-3.16}/netket/utils/array.py +0 -0
  276. {netket-3.15.1 → netket-3.16}/netket/utils/deprecation.py +0 -0
  277. {netket-3.15.1 → netket-3.16}/netket/utils/dispatch.py +0 -0
  278. {netket-3.15.1 → netket-3.16}/netket/utils/display/__init__.py +0 -0
  279. {netket-3.15.1 → netket-3.16}/netket/utils/display/repr.py +0 -0
  280. {netket-3.15.1 → netket-3.16}/netket/utils/float.py +0 -0
  281. {netket-3.15.1 → netket-3.16}/netket/utils/group/__init__.py +0 -0
  282. {netket-3.15.1 → netket-3.16}/netket/utils/group/_group.py +0 -0
  283. {netket-3.15.1 → netket-3.16}/netket/utils/group/_permutation_group.py +0 -0
  284. {netket-3.15.1 → netket-3.16}/netket/utils/group/_point_group.py +0 -0
  285. {netket-3.15.1 → netket-3.16}/netket/utils/group/_semigroup.py +0 -0
  286. {netket-3.15.1 → netket-3.16}/netket/utils/group/axial.py +0 -0
  287. {netket-3.15.1 → netket-3.16}/netket/utils/group/cubic.py +0 -0
  288. {netket-3.15.1 → netket-3.16}/netket/utils/group/icosa.py +0 -0
  289. {netket-3.15.1 → netket-3.16}/netket/utils/group/planar.py +0 -0
  290. {netket-3.15.1 → netket-3.16}/netket/utils/history/__init__.py +0 -0
  291. {netket-3.15.1 → netket-3.16}/netket/utils/holomorphic.py +0 -0
  292. {netket-3.15.1 → netket-3.16}/netket/utils/jax.py +0 -0
  293. {netket-3.15.1 → netket-3.16}/netket/utils/model_frameworks/flax.py +0 -0
  294. {netket-3.15.1 → netket-3.16}/netket/utils/model_frameworks/haiku.py +0 -0
  295. {netket-3.15.1 → netket-3.16}/netket/utils/model_frameworks/jax.py +0 -0
  296. {netket-3.15.1 → netket-3.16}/netket/utils/moduletools.py +0 -0
  297. {netket-3.15.1 → netket-3.16}/netket/utils/mpi/__init__.py +0 -0
  298. {netket-3.15.1 → netket-3.16}/netket/utils/mpi/gpu_autorank_util.py +0 -0
  299. {netket-3.15.1 → netket-3.16}/netket/utils/mpi/primitives.py +0 -0
  300. {netket-3.15.1 → netket-3.16}/netket/utils/numbers.py +0 -0
  301. {netket-3.15.1 → netket-3.16}/netket/utils/optional_deps.py +0 -0
  302. {netket-3.15.1 → netket-3.16}/netket/utils/partial.py +0 -0
  303. {netket-3.15.1 → netket-3.16}/netket/utils/seed.py +0 -0
  304. {netket-3.15.1 → netket-3.16}/netket/utils/static_number.py +0 -0
  305. {netket-3.15.1 → netket-3.16}/netket/utils/struct/__init__.py +0 -0
  306. {netket-3.15.1 → netket-3.16}/netket/utils/struct/dataclass.py +0 -0
  307. {netket-3.15.1 → netket-3.16}/netket/utils/struct/fields.py +0 -0
  308. {netket-3.15.1 → netket-3.16}/netket/utils/struct/pytree_serialization_sharding.py +0 -0
  309. {netket-3.15.1 → netket-3.16}/netket/utils/struct/utils.py +0 -0
  310. {netket-3.15.1 → netket-3.16}/netket/utils/summation.py +0 -0
  311. {netket-3.15.1 → netket-3.16}/netket/utils/timing.py +0 -0
  312. {netket-3.15.1 → netket-3.16}/netket/utils/types.py +0 -0
  313. {netket-3.15.1 → netket-3.16}/netket/utils/version_check.py +0 -0
  314. {netket-3.15.1 → netket-3.16}/netket/vqs/__init__.py +0 -0
  315. {netket-3.15.1 → netket-3.16}/netket/vqs/full_summ/__init__.py +0 -0
  316. {netket-3.15.1 → netket-3.16}/netket/vqs/full_summ/expect.py +0 -0
  317. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/__init__.py +0 -0
  318. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/common.py +0 -0
  319. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/kernels.py +0 -0
  320. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_mixed_state/__init__.py +0 -0
  321. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_mixed_state/expect.py +0 -0
  322. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_mixed_state/expect_chunked.py +0 -0
  323. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_mixed_state/state.py +0 -0
  324. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_state/__init__.py +0 -0
  325. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_state/expect.py +0 -0
  326. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_state/expect_chunked.py +0 -0
  327. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_state/expect_forces.py +0 -0
  328. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_state/expect_forces_chunked.py +0 -0
  329. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_state/expect_grad.py +0 -0
  330. {netket-3.15.1 → netket-3.16}/netket/vqs/mc/mc_state/expect_grad_chunked.py +0 -0
@@ -21,6 +21,7 @@ pip-wheel-metadata
21
21
  .virtual_documents
22
22
  .ruff_cache
23
23
  oldest_requirements.txt
24
+ uv.lock
24
25
 
25
26
  # Packages #
26
27
  ############
@@ -1,6 +1,6 @@
1
- Metadata-Version: 2.3
1
+ Metadata-Version: 2.4
2
2
  Name: NetKet
3
- Version: 3.15.1
3
+ Version: 3.16
4
4
  Summary: Netket : Machine Learning toolbox for many-body quantum systems.
5
5
  Project-URL: homepage, https://www.netket.org
6
6
  Project-URL: documentation, https://netket.readthedocs.io/en/latest/#
@@ -8,6 +8,7 @@ Project-URL: repository, https://github.com/netket/netket
8
8
  Project-URL: changelog, https://netket.readthedocs.io/en/latest/docs/changelog.html
9
9
  Author: Giuseppe Carleo, Filippo Vicentini, The NetKet authors
10
10
  License: Apache 2.0
11
+ License-File: LICENSE
11
12
  Keywords: Jax,Machine Learning,Monte Carlo,Neural Networks,Neural Quantum States,Quantum
12
13
  Classifier: Development Status :: 5 - Production/Stable
13
14
  Classifier: Intended Audience :: Science/Research
@@ -18,19 +19,19 @@ Classifier: Programming Language :: Python :: 3
18
19
  Classifier: Topic :: Scientific/Engineering :: Physics
19
20
  Requires-Python: >=3.10
20
21
  Requires-Dist: equinox>=0.11.3
21
- Requires-Dist: flax>=0.7.1
22
+ Requires-Dist: flax>=0.10.2
22
23
  Requires-Dist: igraph>=0.10.0
23
- Requires-Dist: jax>=0.4.16
24
+ Requires-Dist: jax>=0.4.38
24
25
  Requires-Dist: numba>=0.59
25
26
  Requires-Dist: numpy<3.0,>=1.20
26
- Requires-Dist: optax<0.3,>=0.1.3
27
+ Requires-Dist: optax<0.3,>=0.2.2
27
28
  Requires-Dist: orjson<4,>=3.4
28
29
  Requires-Dist: plum-dispatch>=2.4
29
- Requires-Dist: rich>=10.0
30
+ Requires-Dist: rich>=12.0
30
31
  Requires-Dist: scipy<2,>=1.5.3
31
32
  Requires-Dist: tqdm<5,>=4.60
32
33
  Provides-Extra: dev
33
- Requires-Dist: black==24.10.0; extra == 'dev'
34
+ Requires-Dist: black==24.1.0; extra == 'dev'
34
35
  Requires-Dist: build>=1; extra == 'dev'
35
36
  Requires-Dist: coverage>=5; extra == 'dev'
36
37
  Requires-Dist: matplotlib>=3; extra == 'dev'
@@ -41,12 +42,12 @@ Requires-Dist: pytest-json-report>=1.3; extra == 'dev'
41
42
  Requires-Dist: pytest-xdist[psutil]>=2; extra == 'dev'
42
43
  Requires-Dist: pytest>=6; extra == 'dev'
43
44
  Requires-Dist: qutip>=5; extra == 'dev'
44
- Requires-Dist: ruff==0.7.3; extra == 'dev'
45
+ Requires-Dist: ruff==0.9.4; extra == 'dev'
45
46
  Requires-Dist: wheel>=0.38.1; extra == 'dev'
46
47
  Provides-Extra: docs
47
48
  Requires-Dist: myst-nb<1.2,>=1.0; extra == 'docs'
48
49
  Requires-Dist: myst-parser<4.1.0,>=2.0.0; extra == 'docs'
49
- Requires-Dist: sphinx-autodoc-typehints<3.0,>=1.22; extra == 'docs'
50
+ Requires-Dist: sphinx-autodoc-typehints<4.0,>=1.22; extra == 'docs'
50
51
  Requires-Dist: sphinx-book-theme~=1.1.0; extra == 'docs'
51
52
  Requires-Dist: sphinx-design>=0.4.0; extra == 'docs'
52
53
  Requires-Dist: sphinx<8.2,>=7.0; extra == 'docs'
@@ -12,5 +12,5 @@ __version__: str
12
12
  __version_tuple__: VERSION_TUPLE
13
13
  version_tuple: VERSION_TUPLE
14
14
 
15
- __version__ = version = '3.15.1'
16
- __version_tuple__ = version_tuple = (3, 15, 1)
15
+ __version__ = version = '3.16'
16
+ __version_tuple__ = version_tuple = (3, 16)
@@ -657,7 +657,7 @@ class SymmModuleInvalidInputShape(NetketError):
657
657
 
658
658
  def __init__(self, name, x):
659
659
  super().__init__(
660
- """
660
+ f"""
661
661
  Input to DenseSymmFFT has {x.ndim =} but 3 are required.
662
662
 
663
663
  The input format is (B,C,L), aka (batches, features/channels, length). If
@@ -1128,3 +1128,38 @@ def concrete_or_error(force, value, error_class, *args, **kwargs):
1128
1128
  )
1129
1129
  except ConcretizationTypeError as err:
1130
1130
  raise error_class(*args, **kwargs) from err
1131
+
1132
+
1133
+ #################################################
1134
+ # Graph errors #
1135
+ #################################################
1136
+
1137
+
1138
+ class InitializePeriodicLatticeOnSmallLatticeWarning(NetketWarning):
1139
+ """
1140
+ Warning thrown when attempting to create a periodic lattice on a lattice with less than two sites in one direction.
1141
+
1142
+ In a periodic lattice with two sites (a, b) in one direction,
1143
+ the expected behavior is to have an edge connecting both a->b and b->a
1144
+
1145
+ However, as the lattice uses an undirected graph to represent the lattice, it does not support this behavior.
1146
+ Hence, only one edge a->b is created, which makes the behavior equivalent to an open boundary condition in this direction.
1147
+
1148
+ This may cause unexpected behavior if you intend to loop over the edges to create a Hamiltonian, for example.
1149
+
1150
+ To avoid this warning, consider either using a lattice with more than two sites in the direction you want to be periodic,
1151
+ or define the graph using :class:`~netket.graph.Graph` by adding the edges manually.
1152
+ """
1153
+
1154
+ def __init__(self, extent, dimension):
1155
+ super().__init__(
1156
+ f"""
1157
+ You are attempting to define a lattice with length {extent} in dimension {dimension} using periodic boundary condition.
1158
+
1159
+ Lattice with less than two sites in one direction does not support periodic boundary condition.
1160
+ The behavior of the lattice is equivalent to an open boundary condition in this direction.
1161
+
1162
+ To avoid this warning, consider either using a lattice with more than two sites in the direction you want to be periodic,
1163
+ or define the graph using :class:`~netket.graph.Graph` by adding the edges manually.
1164
+ """
1165
+ )
@@ -38,15 +38,6 @@ class TDVP(TDVPBaseDriver):
38
38
  Variational time evolution based on the time-dependent variational principle which,
39
39
  when used with Monte Carlo sampling via :class:`netket.vqs.MCState`, is the time-dependent VMC
40
40
  (t-VMC) method.
41
-
42
- .. note::
43
- This TDVP Driver uses the time-integrators from the `nkx.dynamics` module, which are
44
- automatically executed under a `jax.jit` context.
45
-
46
- When running computations on GPU, this can lead to infinite hangs or extremely long
47
- compilation times. In those cases, you might try setting the configuration variable
48
- `nk.config.netket_experimental_disable_ode_jit = True` to mitigate those issues.
49
-
50
41
  """
51
42
 
52
43
  def __init__(
@@ -12,20 +12,17 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from collections.abc import Callable
16
- from collections.abc import Sequence
15
+ from collections.abc import Callable, Sequence
16
+ from functools import partial
17
17
  import warnings
18
18
 
19
- import jax
20
19
  import jax.numpy as jnp
21
20
  import numpy as np
22
21
  from tqdm.auto import tqdm
23
22
 
24
23
  import netket as nk
25
- from netket import config
26
24
  from netket.driver import AbstractVariationalDriver
27
25
  from netket.driver.abstract_variational_driver import _to_iterable
28
- from netket.jax import HashablePartial
29
26
  from netket.logging.json_log import JsonLog
30
27
  from netket.operator import AbstractOperator
31
28
  from netket.utils import mpi, timing
@@ -45,15 +42,6 @@ class TDVPBaseDriver(AbstractVariationalDriver):
45
42
  Variational time evolution based on the time-dependent variational principle which,
46
43
  when used with Monte Carlo sampling via :class:`netket.vqs.MCState`, is the time-dependent VMC
47
44
  (t-VMC) method.
48
-
49
- .. note::
50
- This TDVP Driver uses the time-integrators from the `nkx.dynamics` module, which are
51
- automatically executed under a `jax.jit` context.
52
-
53
- When running computations on GPU, this can lead to infinite hangs or extremely long
54
- compilation times. In those cases, you might try setting the configuration variable
55
- :py:`nk.config.netket_experimental_disable_ode_jit = True` to mitigate those issues.
56
-
57
45
  """
58
46
 
59
47
  def __init__(
@@ -107,8 +95,6 @@ class TDVPBaseDriver(AbstractVariationalDriver):
107
95
  self._last_qgt = None
108
96
  self._integrator = None
109
97
 
110
- self._odefun = HashablePartial(odefun_host_callback, self.state, self)
111
-
112
98
  self.error_norm = error_norm
113
99
 
114
100
  if integrator is not None:
@@ -149,7 +135,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
149
135
  t0 = self.t
150
136
 
151
137
  self._integrator = Integrator(
152
- self._odefun,
138
+ partial(odefun, self.state, self),
153
139
  new_ode_solver,
154
140
  t0,
155
141
  self.state.parameters,
@@ -186,18 +172,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
186
172
  elif error_norm == "maximum":
187
173
  self._error_norm = maximum_norm
188
174
  elif error_norm == "qgt":
189
- if config.netket_experimental_disable_ode_jit:
190
- self._error_norm = HashablePartial(qgt_norm, self)
191
- else:
192
- w = self.state.parameters
193
- norm_dtype = nk.jax.dtype_real(nk.jax.tree_dot(w, w))
194
- # QGT norm is called via host callback since it accesses the driver
195
- # TODO: make this also an hashablepartial on self to reduce recompilation
196
- self._error_norm = lambda x: jax.pure_callback(
197
- HashablePartial(qgt_norm, self),
198
- jax.ShapeDtypeStruct((), norm_dtype),
199
- x,
200
- )
175
+ self._error_norm = partial(qgt_norm, self)
201
176
  else:
202
177
  raise ValueError(
203
178
  "error_norm must be a callable or one of 'euclidean', 'qgt', 'maximum',"
@@ -479,7 +454,7 @@ class TDVPBaseDriver(AbstractVariationalDriver):
479
454
  t = self.t
480
455
  if w is None:
481
456
  w = self.state.parameters
482
- return self._odefun(t, w)
457
+ return odefun(self.state, self, t, w)
483
458
 
484
459
 
485
460
  def qgt_norm(driver: TDVPBaseDriver, x: PyTree):
@@ -495,25 +470,3 @@ def qgt_norm(driver: TDVPBaseDriver, x: PyTree):
495
470
  def odefun(state, driver, t, w, **kwargs):
496
471
  # pylint: disable=unused-argument
497
472
  raise NotImplementedError(f"odefun not implemented for {type(state)}")
498
-
499
-
500
- def odefun_host_callback(state, driver, *args, **kwargs):
501
- """
502
- Calls odefun through a host callback in order to make the rest of the
503
- ODE inteagrator jit-able.
504
- """
505
- if config.netket_experimental_disable_ode_jit:
506
- return odefun(state, driver, *args, **kwargs)
507
-
508
- result_shape = jax.tree_util.tree_map(
509
- lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype),
510
- state.parameters,
511
- )
512
-
513
- return jax.pure_callback(
514
- lambda args_and_kw: odefun(state, driver, *args_and_kw[0], **args_and_kw[1]),
515
- result_shape,
516
- # pack args and kwargs together, since host_callback passes a single argument:
517
- (args, kwargs),
518
- )
519
- return odefun(state, driver, *args, **kwargs)
@@ -88,16 +88,6 @@ class TDVPSchmitt(TDVPBaseDriver):
88
88
  :math:`\rho_k=V_{k,k'}^{\dagger}F_{k'}` (see
89
89
  `[arXiv:1912.08828] <https://arxiv.org/pdf/1912.08828.pdf>`_ for details).
90
90
 
91
-
92
- .. note::
93
-
94
- This TDVP Driver uses the time-integrators from the `nkx.dynamics` module, which are
95
- automatically executed under a `jax.jit` context.
96
-
97
- When running computations on GPU, this can lead to infinite hangs or extremely long
98
- compilation times. In those cases, you might try setting the configuration variable
99
- `nk.config.netket_experimental_disable_ode_jit = True` to mitigate those issues.
100
-
101
91
  """
102
92
 
103
93
  def __init__(
@@ -225,8 +215,16 @@ def _impl(parameters, n_samples, E_loc, S, rhs_coeff, rcond, rcond_smooth, snr_a
225
215
  QEdata = Q.conj() * ΔE_loc
226
216
  rho = V.conj().T @ F
227
217
 
228
- # Compute the SNR according to Eq. 21
229
- snr = jnp.abs(rho) * jnp.sqrt(n_samples) / jnp.sqrt(stats.var(QEdata, axis=0))
218
+ # Compute the SNR according to Eq. 21 but taking care of where sigma_k is zero
219
+ sigma_k = jnp.maximum(jnp.sqrt(stats.var(QEdata, axis=0)), rcond)
220
+ # Here we are hardcoding the case where rho==0 and sigma_k==0 to have infinite snr.
221
+ # This is an arbitrary choice, but avoids generating NaNs in the snr calculation.
222
+ # See netket#1959 and #1960 for more details.
223
+ snr = jnp.where(
224
+ sigma_k == 0,
225
+ jnp.inf,
226
+ jnp.abs(rho) * jnp.sqrt(n_samples) / sigma_k,
227
+ )
230
228
 
231
229
  # Discard eigenvalues below numerical precision
232
230
  ev_inv = jnp.where(jnp.abs(ev / ev[-1]) > rcond, 1.0 / ev, 0.0)
@@ -12,8 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Callable, Optional
16
- from functools import partial
15
+ from collections.abc import Callable
17
16
 
18
17
  import jax
19
18
  import jax.numpy as jnp
@@ -27,7 +26,6 @@ from ._integrator_state import IntegratorState, IntegratorFlags
27
26
  from ._integrator_params import IntegratorParameters
28
27
  from ._solver import AbstractSolver
29
28
  from ._utils import (
30
- maybe_jax_jit,
31
29
  scaled_error,
32
30
  propose_time_step,
33
31
  set_flag_jax,
@@ -217,12 +215,11 @@ class Integrator(struct.Pytree, mutable=True):
217
215
  )
218
216
 
219
217
  @staticmethod
220
- @partial(maybe_jax_jit, static_argnames=["f"])
221
218
  def _step_fixed(
222
219
  solver: AbstractSolver,
223
220
  f: Callable,
224
221
  state: IntegratorState,
225
- max_dt: Optional[float],
222
+ max_dt: float | None,
226
223
  parameters: IntegratorParameters,
227
224
  ) -> IntegratorState:
228
225
  r"""
@@ -263,12 +260,11 @@ class Integrator(struct.Pytree, mutable=True):
263
260
  )
264
261
 
265
262
  @staticmethod
266
- @partial(maybe_jax_jit, static_argnames=["f", "norm_fn"])
267
263
  def _step_adaptive(
268
264
  solver: AbstractSolver,
269
265
  f: Callable,
270
266
  state: IntegratorState,
271
- max_dt: Optional[float],
267
+ max_dt: float | None,
272
268
  parameters: IntegratorParameters,
273
269
  norm_fn: Callable,
274
270
  ) -> IntegratorState:
@@ -12,7 +12,6 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Optional
16
15
 
17
16
  from netket.utils import struct
18
17
 
@@ -29,7 +28,7 @@ class IntegratorParameters(struct.Pytree):
29
28
  rtol: float
30
29
  """The tolerance for the relative error on the solution."""
31
30
 
32
- dt_limits: Optional[LimitsDType] = struct.field(pytree_node=False)
31
+ dt_limits: LimitsDType | None = struct.field(pytree_node=False)
33
32
  """The extremal accepted values for the time-step size `dt`."""
34
33
 
35
34
  def __init__(
@@ -37,7 +36,7 @@ class IntegratorParameters(struct.Pytree):
37
36
  dt: float,
38
37
  atol: float = 0.0,
39
38
  rtol: float = 1e-7,
40
- dt_limits: Optional[LimitsDType] = None,
39
+ dt_limits: LimitsDType | None = None,
41
40
  ):
42
41
  r"""
43
42
  Args:
@@ -12,7 +12,7 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from typing import Callable
15
+ from collections.abc import Callable
16
16
 
17
17
  import jax
18
18
  import jax.numpy as jnp
@@ -114,9 +114,7 @@ class AbstractSolver(Pytree):
114
114
  )
115
115
 
116
116
  def __repr__(self) -> str:
117
- return "{}(dt={}, adaptive={}, integrator_parameters={})".format(
118
- self.__class__.__name__, self.dt, self.adaptive, self.integrator_params
119
- )
117
+ return f"{self.__class__.__name__}(dt={self.dt}, adaptive={self.adaptive}, integrator_parameters={self.integrator_params})"
120
118
 
121
119
  @property
122
120
  def is_explicit(self) -> bool:
@@ -12,14 +12,12 @@
12
12
  # See the License for the specific language governing permissions and
13
13
  # limitations under the License.
14
14
 
15
- from functools import wraps
16
15
 
17
16
  import jax
18
17
  import jax.numpy as jnp
19
18
 
20
19
  import netket.jax as nkjax
21
- from netket import config
22
- from netket.utils.types import Array, PyTree, Callable
20
+ from netket.utils.types import Array, PyTree
23
21
 
24
22
 
25
23
  LimitsDType = tuple[float | None, float | None]
@@ -53,29 +51,6 @@ def propose_time_step(
53
51
  )
54
52
 
55
53
 
56
- def maybe_jax_jit(fun: Callable, *jit_args, **jit_kwargs) -> Callable:
57
- """
58
- Only jit if `config.netket_experimental_disable_ode_jit` is False.
59
-
60
- This is used to disable jitting when this config is set. The switch is
61
- performed at runtime so that the flag can be changed as desired.
62
- """
63
-
64
- # jit the function only once:
65
- jitted_fun = jax.jit(fun, *jit_args, **jit_kwargs)
66
-
67
- @wraps(fun)
68
- def _maybe_jitted_fun(*args, **kwargs):
69
- if config.netket_experimental_disable_ode_jit:
70
- with jax.spmd_mode("allow_all"):
71
- res = fun(*args, **kwargs)
72
- return res
73
- else:
74
- return jitted_fun(*args, **kwargs)
75
-
76
- return _maybe_jitted_fun
77
-
78
-
79
54
  def set_flag_jax(condition, flags, flag):
80
55
  """
81
56
  If `condition` is true, `flags` is updated by setting `flag` to 1.
@@ -15,7 +15,22 @@
15
15
 
16
16
  __all__ = ["SpinOrbitalFermions"]
17
17
 
18
- from .spin_orbital_fermions import SpinOrbitalFermions
18
+ from netket.hilbert import SpinOrbitalFermions as _deprecated_SpinOrbitalFermions
19
+
20
+
21
+ _deprecations = {
22
+ # May 2024
23
+ "SpinOrbitalFermions": (
24
+ "netket.experimental.hilbert.SpinOrbitalFermions is deprecated: use "
25
+ "netket.hilbert.SpinOrbitalFermions (netket >= 3.12)",
26
+ _deprecated_SpinOrbitalFermions,
27
+ ),
28
+ }
29
+
30
+ from netket.utils.deprecation import deprecation_getattr as _deprecation_getattr
31
+
32
+ __getattr__ = _deprecation_getattr(__name__, _deprecations)
33
+ del _deprecation_getattr
19
34
 
20
35
  from netket.utils import _hide_submodules
21
36
 
@@ -0,0 +1,38 @@
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 .rnn import RNN, LSTMNet, GRUNet1D
16
+ from .fast_rnn import FastRNN, FastLSTMNet, FastGRUNet1D
17
+
18
+ from netket.models import Slater2nd as _deprecated_Slater2nd
19
+ from netket.models import MultiSlater2nd as _deprecated_MultiSlater2nd
20
+
21
+ _deprecations = {
22
+ # May 2024
23
+ "Slater2nd": (
24
+ "netket.experimental.models.Slater2nd is deprecated: use "
25
+ "netket.models.Slater2nd (netket >= 3.12)",
26
+ _deprecated_Slater2nd,
27
+ ),
28
+ "MultiSlater2nd": (
29
+ "netket.experimental.models.MultiSlater2nd is deprecated: use "
30
+ "netket.models.MultiSlater2nd (netket >= 3.12)",
31
+ _deprecated_MultiSlater2nd,
32
+ ),
33
+ }
34
+
35
+ from netket.utils.deprecation import deprecation_getattr as _deprecation_getattr
36
+
37
+ __getattr__ = _deprecation_getattr(__name__, _deprecations)
38
+ del _deprecation_getattr
@@ -0,0 +1,47 @@
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
+
16
+ __all__ = ["FermionOperator2nd", "FermionOperator2ndJax"]
17
+
18
+ from .pyscf import from_pyscf_molecule
19
+
20
+
21
+ from . import fermion
22
+
23
+ from netket.operator import FermionOperator2nd as _deprecated_FermionOperator2nd
24
+ from netket.operator import FermionOperator2ndJax as _deprecated_FermionOperator2ndJax
25
+
26
+ _deprecations = {
27
+ # June 2024, NetKet 3.13
28
+ "FermionOperator2nd": (
29
+ "netket.experimental.operator.FermionOperator2nd is deprecated: use "
30
+ "netket.operator.FermionOperator2nd (netket >= 3.13)",
31
+ _deprecated_FermionOperator2nd,
32
+ ),
33
+ "FermionOperator2ndJax": (
34
+ "netket.experimental.operator.FermionOperator2ndJax is deprecated: use "
35
+ "netket.operator.FermionOperator2ndJax (netket >= 3.13)",
36
+ _deprecated_FermionOperator2ndJax,
37
+ ),
38
+ }
39
+
40
+
41
+ from netket.utils import _auto_export
42
+ from netket.utils.deprecation import deprecation_getattr as _deprecation_getattr
43
+
44
+ __getattr__ = _deprecation_getattr(__name__, _deprecations)
45
+ _auto_export(__name__)
46
+
47
+ del _deprecation_getattr, _auto_export
@@ -0,0 +1,44 @@
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
+
16
+ """
17
+ This file only contains deprecated bindings. To be removed.
18
+ """
19
+
20
+ from netket.operator import fermion as _deprecated_fermion
21
+
22
+ _deprecations = {
23
+ # June 2024, NetKet 3.13
24
+ f"{_name}": (
25
+ f"netket.experimental.operator.fermion.{_name} is deprecated: use "
26
+ f"netket.operator.fermion.{_name} (netket >= 3.13)",
27
+ getattr(_deprecated_fermion, _name),
28
+ )
29
+ for _name in [
30
+ "destroy",
31
+ "create",
32
+ "number",
33
+ "identity",
34
+ "zero",
35
+ ]
36
+ }
37
+
38
+ from netket.utils import _auto_export # noqa: E402
39
+ from netket.utils.deprecation import ( # noqa: E402
40
+ deprecation_getattr as _deprecation_getattr,
41
+ )
42
+
43
+ __getattr__ = _deprecation_getattr(__name__, _deprecations)
44
+ _auto_export(__name__)
@@ -19,9 +19,9 @@ import jax
19
19
  import jax.numpy as jnp
20
20
 
21
21
  from netket.operator import DiscreteOperator
22
- from netket.experimental.hilbert import SpinOrbitalFermions
22
+ from netket.hilbert import SpinOrbitalFermions
23
23
  from netket.utils.optional_deps import import_optional_dependency
24
- from ._fermion_operator_2nd_numba import FermionOperator2nd
24
+ from netket.operator import FermionOperator2nd
25
25
 
26
26
 
27
27
  def compute_pyscf_integrals(mol, mo_coeff):
@@ -16,14 +16,13 @@ from . import rules
16
16
 
17
17
  from .metropolis_pmap import MetropolisSamplerPmap
18
18
 
19
- from .metropolis import MetropolisParticleExchange
20
-
21
19
 
22
20
  # Deprecated bindings from stabilisaation
23
21
  from netket.sampler import (
24
22
  ParallelTemperingSampler as _deprecated_MetropolisPtSampler,
25
23
  ParallelTemperingLocal as _deprecated_MetropolisLocalPtSampler,
26
24
  ParallelTemperingExchange as _deprecated_MetropolisExchangePtSampler,
25
+ MetropolisFermionHop as _deprecated_MetropolisParticleExchange,
27
26
  )
28
27
 
29
28
  _deprecations = {
@@ -43,8 +42,15 @@ _deprecations = {
43
42
  "netket.sampler.ParallelTemperingExchange (netket >= 3.12)",
44
43
  _deprecated_MetropolisExchangePtSampler,
45
44
  ),
45
+ # June 2024, NetKet 3.13
46
+ "MetropolisParticleExchange": (
47
+ "netket.experimental.sampler.MetropolisParticleExchange is deprecated: use "
48
+ "netket.sampler.MetropolisFermionHop (netket >= 3.13)",
49
+ _deprecated_MetropolisParticleExchange,
50
+ ),
46
51
  }
47
52
 
53
+
48
54
  from netket.utils.deprecation import deprecation_getattr as _deprecation_getattr
49
55
  from netket.utils import _hide_submodules
50
56
 
@@ -0,0 +1,17 @@
1
+ from netket.sampler.rules import (
2
+ FermionHopRule as _deprecated_FermionHopRule,
3
+ )
4
+
5
+ _deprecations = {
6
+ # June 2024, NetKet 3.13
7
+ "ParticleExchangeRule": (
8
+ "netket.experimental.sampler.rules.ParticleExchangeRule is deprecated: use "
9
+ "netket.sampler.rules.FermionHopRule (netket >= 3.13)",
10
+ _deprecated_FermionHopRule,
11
+ ),
12
+ }
13
+
14
+ from netket.utils.deprecation import deprecation_getattr as _deprecation_getattr
15
+
16
+ __getattr__ = _deprecation_getattr(__name__, _deprecations)
17
+ del _deprecation_getattr