Repository: netket/netket

Issue:
NaN energy during optimisation
Sometimes an experiment might produce nan energies during optimisation. However, the optimisation continues. I think it would be best to stop the optimisation because nothing is being optimised.


Discussion:
This can be implemented as a callback, right? Something like:
```python3
def callback(_, _, driver):
    return np.isfinite(driver.loss_stats.mean)  # or `not np.isnan(...)` for a more fine-grained criterion
```
Yes, I am currently dealing with it like that. However, it seems to me that this behaviour should be default. What can we possibly gain from letting the VMC driver run once it evaluates the energy to nan?
I agree. Technically, the energy could be `NaN` but as long as that is a fluke and the gradient is fine, the optimization could recover. Though this seems unlikely to happen in practice (I don't recall seeing this, at least).

In any case, I'd suggest adding something like a `StopOnInvalidLoss` callback and add it to the list of callbacks by default in our drivers, unless it is disabled by some flag, to implement this functionality.
A couple of days ago I saw the energy going to `NaN` and then recover. Not reproducible, though. I'll just create the stop on invalid loss.

Resolving patch:
diff --git a/netket/callbacks/__init__.py b/netket/callbacks/__init__.py
--- a/netket/callbacks/__init__.py
+++ b/netket/callbacks/__init__.py
@@ -14,6 +14,7 @@
 
 from .earlystopping import EarlyStopping
 from .timeout import Timeout
+from .invalidlossstopping import InvalidLossStopping
 
 from netket.utils import _hide_submodules
 
diff --git a/netket/callbacks/invalidlossstopping.py b/netket/callbacks/invalidlossstopping.py
new file mode 100644
--- /dev/null
+++ b/netket/callbacks/invalidlossstopping.py
@@ -0,0 +1,52 @@
+# 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 dataclasses import dataclass
+from typing import Union
+
+import numpy as np
+
+
+@dataclass
+class InvalidLossStopping:
+    """A simple callback to stop NetKet if there are no more improvements in the training.
+    based on `driver._loss_name`."""
+
+    monitor: str = "mean"
+    """Loss statistic to monitor. Should be one of 'mean', 'variance', 'sigma'."""
+    patience: Union[int, float] = 0
+    """Number of epochs with invalid loss after which training will be stopped."""
+
+    def __post_init__(self):
+        self._last_valid_iter = 0
+
+    def __call__(self, step, log_data, driver):
+        """
+        A boolean function that determines whether or not to stop training.
+
+        Args:
+            step: An integer corresponding to the step (iteration or epoch) in training.
+            log_data: A dictionary containing log data for training.
+            driver: A NetKet variational driver.
+
+        Returns:
+            A boolean. If True, training continues, else, it does not.
+        """
+        loss = np.real(getattr(log_data[driver._loss_name], self.monitor))
+        if not np.isfinite(loss):
+            if step - self._last_valid_iter >= self.patience:
+                return False
+        else:
+            self._last_valid_iter = step
+        return True


Test patch:
diff --git a/test/callbacks/test_callbacks.py b/test/callbacks/test_callbacks.py
--- a/test/callbacks/test_callbacks.py
+++ b/test/callbacks/test_callbacks.py
@@ -73,3 +73,14 @@ def test_earlystopping_doesnt_get_stuck_with_patience():
         if not es(step, {"loss": DummyLogEntry(loss_values[step])}, driver):
             break
     assert es._best_val == 9
+
+
+def test_invalid_loss_stopping():
+    loss_values = [10, 9, 8, 7] + [float("nan")] * 11 + [1] * 4
+    ils = nk.callbacks.InvalidLossStopping(patience=10)
+    driver = DummyDriver()
+    for step in range(len(loss_values)):
+        if not ils(step, {"loss": DummyLogEntry(loss_values[step])}, driver):
+            break
+    assert ils._last_valid_iter == 3
+    assert step == 13