NetKet 3.22.0.dev1__py3-none-any.whl → 3.22.0.dev3__py3-none-any.whl

netket/_src/driver/abstract_variational_driver.py

@@ -229,6 +229,21 @@ class AbstractDriver(struct.Pytree, mutable=True):
         self.reset_step()
         # self._step_count = 0
 
+    def _accept_step(self) -> bool:
+        """
+        Called after :meth:`compute_loss_and_update` to decide whether the
+        current step should be accepted.
+
+        Return ``False`` to reject the step and retry (``_step_attempt`` will
+        be incremented before the next attempt). Subclasses can override this
+        to implement driver-level acceptance criteria such as adaptive
+        step-size control: mutate driver state (e.g. ``self.dt``) as needed
+        before returning ``False``.
+
+        The default implementation always accepts.
+        """
+        return True
+
     def reset_step(self, hard: bool = False):
         """
         Resets the state of the driver at the beginning of a new step.
@@ -442,7 +457,8 @@ class AbstractDriver(struct.Pytree, mutable=True):
                         reject_step = callbacks.on_compute_update_end(
                             self.step_count, step_log_data, self
                         )
-                        if reject_step:
+                        driver_reject_step = not self._accept_step()
+                        if driver_reject_step or reject_step:
                             self._step_attempt += 1
                             continue
                         else:

netket/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '3.22.0.dev1'
-__version_tuple__ = version_tuple = (3, 22, 0, 'dev1')
+__version__ = version = '3.22.0.dev3'
+__version_tuple__ = version_tuple = (3, 22, 0, 'dev3')
 
 __commit_id__ = commit_id = None

netket/experimental/dynamics/_rk/_tableau.py

@@ -108,24 +108,29 @@ bt_feuler = TableauRKExplicit(
 
 
 bt_midpoint = TableauRKExplicit(
-                order = (2,),
+                order = (2,1),
                 a = jnp.array([[0,   0],
                                [1/2, 0]], dtype=default_dtype),
-                b = jnp.array( [0,   1], dtype=default_dtype),
+                b = jnp.array([[0,   1],   # midpoint, order 2
+                               [1,   0]]),  # Euler, order 1
                 c = jnp.array( [0, 1/2], dtype=default_dtype),
                 name = "Midpoint"
                 )
 
 
 bt_heun = TableauRKExplicit(
-                order = (2,),
+                order = (2,1),
                 a = jnp.array([[0,   0],
                                [1,   0]], dtype=default_dtype),
-                b = jnp.array( [1/2, 1/2], dtype=default_dtype),
+                b = jnp.array([[1/2, 1/2],   # Heun, order 2
+                               [1,   0  ]]),  # Euler, order 1
                 c = jnp.array( [0, 1], dtype=default_dtype),
                 name = "Heun"
                 )
 
+# Adaptive step:
+# Heun Euler https://en.wikipedia.org/wiki/Runge–Kutta_methods
+bt_rk12  = bt_heun
 
 bt_rk4  = TableauRKExplicit(
                 order = (4,),
@@ -138,20 +143,6 @@ bt_rk4  = TableauRKExplicit(
                 name = "RK4"
                 )
 
-
-# Adaptive step:
-# Heun Euler https://en.wikipedia.org/wiki/Runge–Kutta_methods
-bt_rk12  = TableauRKExplicit(
-                order = (2,1),
-                a = jnp.array([[0,   0],
-                               [1,   0]], dtype=default_dtype),
-                b = jnp.array([[1/2, 1/2],
-                               [1,   0]], dtype=default_dtype),
-                c = jnp.array( [0, 1], dtype=default_dtype),
-                name = "RK12"
-                )
-
-
 # Bogacki–Shampine coefficients
 bt_rk23  = TableauRKExplicit(
                 order = (3,2),

{netket-3.22.0.dev1.dist-info → netket-3.22.0.dev3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: NetKet
-Version: 3.22.0.dev1
+Version: 3.22.0.dev3
 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/#
@@ -54,7 +54,7 @@ Requires-Dist: myst-nb<1.5,>=1.0; extra == 'docs'
 Requires-Dist: myst-parser<5.1.0,>=2.0.0; extra == 'docs'
 Requires-Dist: nbconvert>=7.16.6; extra == 'docs'
 Requires-Dist: sphinx-autodoc-typehints<4.0,>=1.22; extra == 'docs'
-Requires-Dist: sphinx-book-theme~=1.1.0; extra == 'docs'
+Requires-Dist: sphinx-book-theme<1.3,>=1.1; extra == 'docs'
 Requires-Dist: sphinx-copybutton>=0.5.2; extra == 'docs'
 Requires-Dist: sphinx-design>=0.4.0; extra == 'docs'
 Requires-Dist: sphinx<9.2.0,>=8.2.3; extra == 'docs'

{netket-3.22.0.dev1.dist-info → netket-3.22.0.dev3.dist-info}/RECORD

@@ -1,6 +1,6 @@
 docs/sphinx_extensions/netket_citations.py,sha256=IkS9Uj_yv12vdKlGQwT_HZHgtr2rlTu7zpRBM6d6iYo,7990
 netket/__init__.py,sha256=bn46oMQ0ahZ3RgzQbx780aTWsGAVx1VO_tvtmGcPxto,1553
-netket/_version.py,sha256=5VwbeC7MNxrikogPRDBsdmr1xcOF3iadJTRlkaJa8hk,719
+netket/_version.py,sha256=6p6qkvhXglYZXbFAvu3a2PworISwuesiQ7mV3qft6Zo,719
 netket/errors.py,sha256=ZVW-_MY3basQU5R-304NYuVW6m3VKe0br569RwoH-bk,60415
 netket/exact.py,sha256=o4vDhiEv0sew3VL-8_NjT7yMlJNXZIlVunytzI4NCBE,8713
 netket/_src/callbacks/auto_chunk_size.py,sha256=P9iggnZE-pz47vj3dRDE8sJRSRDlOEo5uEh1ZonIFyw,9562
@@ -14,7 +14,7 @@ netket/_src/callbacks/save_state.py,sha256=ojt2JEN2xtqyNYSbHlx1tg55kT0E33NBxN1pO
 netket/_src/callbacks/stop_at_time.py,sha256=tEWMX71be4r2CDp1W2v-tSjcaPhe-S1qkJteVxBNokE,1321
 netket/_src/driver/abstract_dynamics_driver.py,sha256=cKyVaf-_GCHMryMKW2X_wQRq70CzP-asOrulvCwic5k,7183
 netket/_src/driver/abstract_optimization_driver.py,sha256=EzEGYH_PYXxPeOdt6udqyWj1idyL0NBFPxF_qHRoH9A,3121
-netket/_src/driver/abstract_variational_driver.py,sha256=MGcjGI_OvwmPDruBw3tDovv4-hgmxrPYBNkBmHcGyuw,21169
+netket/_src/driver/abstract_variational_driver.py,sha256=ChtxMqY-wkGyrwAHjet1GSDnRasGnt5Bn6DU2t7UpoU,21849
 netket/_src/driver/abstract_variational_driver_old.py,sha256=JkMy04pQukwvYMxbnnIU_IzB-PnGKwazL5N7xzg1dMA,16884
 netket/_src/driver/infidelity_sr.py,sha256=Vo119ITrcq-CZi_1AH0B87w6MwycqNxLnrdpjuXbEWg,17842
 netket/_src/driver/vmc_sr.py,sha256=OkpDUwErsRfVRgKmoiw5cIIhCtPVJBABYrda334uLNg,28954
@@ -76,7 +76,7 @@ netket/experimental/dynamics/_solver.py,sha256=-2gaEzWF2qxY9VSVVC8Xm22dQb3wFUYSx
 netket/experimental/dynamics/_utils.py,sha256=j0tq1W0HCGEr8y9iRlypVtVZBuHVU1FeqfnCQvSrFlE,3052
 netket/experimental/dynamics/_rk/__init__.py,sha256=sl2tCIuOqK-pzEaxjCui2v0kSeQQ2XwtH274fFpdWcE,66
 netket/experimental/dynamics/_rk/_solver.py,sha256=-mZyVuYiszTMkKBwzVXt9vlGSAiQHuPa3BTEezTtrA4,7532
-netket/experimental/dynamics/_rk/_tableau.py,sha256=WY1zCzA7KwgYJHBl5H--O0vcduwEwnmAsykcaZ4xVTE,8990
+netket/experimental/dynamics/_rk/_tableau.py,sha256=TOyRICs-d5fCI17IcRsLxTXqIgg5M33zV_3M7aNwcu8,8746
 netket/experimental/geometry/__init__.py,sha256=xuiOXqqbFP8YiQS2XOEa643xZK0_FPyQYwzptARd0xI,133
 netket/experimental/geometry/cell.py,sha256=iaQM4nE-WMCIM2jaDVntaBoecWhBumaK9BrfCcbcLqU,2515
 netket/experimental/hilbert/__init__.py,sha256=A_5xRK8Cz4HsNKiDUDpfTIfnYUHpliCUOAszefUWKSY,834
@@ -405,8 +405,8 @@ netket/vqs/mc/mc_state/expect_forces_chunked.py,sha256=XcG2UvyqNkGy1lMoau2yzscLq
 netket/vqs/mc/mc_state/expect_grad.py,sha256=l_m5y4OQx1aRuklXg8fNQNbPOlHBS3OsHAweBJm51Fw,5260
 netket/vqs/mc/mc_state/expect_grad_chunked.py,sha256=H1NmakACgDKfTArO0tftVdG6dn3UdYAs_RSEL1HFiAQ,2322
 netket/vqs/mc/mc_state/state.py,sha256=A70TlejtZi1P81fGkXEyBRkxI4oHEDYAQV1qth0izOQ,40381
-netket-3.22.0.dev1.dist-info/METADATA,sha256=a86TqtrI6x32_W6KMtUgzDFyqVDxKrhoBb3x0SotP6o,5932
-netket-3.22.0.dev1.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
-netket-3.22.0.dev1.dist-info/entry_points.txt,sha256=mGOesl2TwrfDnkQVvYYb1jvA5TTzGPGchbCFPFpqHao,104
-netket-3.22.0.dev1.dist-info/licenses/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
-netket-3.22.0.dev1.dist-info/RECORD,,
+netket-3.22.0.dev3.dist-info/METADATA,sha256=wbSDOxxAxzu6xH3A86ANWBYCCrB8L6CLLW_X0meXfeA,5935
+netket-3.22.0.dev3.dist-info/WHEEL,sha256=QccIxa26bgl1E6uMy58deGWi-0aeIkkangHcxk2kWfw,87
+netket-3.22.0.dev3.dist-info/entry_points.txt,sha256=mGOesl2TwrfDnkQVvYYb1jvA5TTzGPGchbCFPFpqHao,104
+netket-3.22.0.dev3.dist-info/licenses/LICENSE,sha256=psuoW8kuDP96RQsdhzwOqi6fyWv0ct8CR6Jr7He_P_k,10173
+netket-3.22.0.dev3.dist-info/RECORD,,