Take best params (#17)

* Return best parameter set after training
* Remove flake-nb

.gitignore

@@ -153,3 +153,6 @@ cython_debug/
 
 # vscode
 .vscode/
+
+
+slurm*

.pre-commit-config.yaml

@@ -60,17 +60,6 @@ repos:
  args: ["--ignore-missing-imports"]
  files: "(sbijax|examples)"
 
-- repo: https://github.com/nbQA-dev/nbQA
- rev: 1.6.3
- hooks:
- - id: nbqa-black
- - id: nbqa-pyupgrade
- args: [--py39-plus]
- - id: nbqa-isort
- args: ['--profile=black']
- - id: nbqa-flake8
- args: ['--ignore=E501,E203,E302,E402,E731,W503']
-
 - repo: https://github.com/jorisroovers/gitlint
  rev: v0.18.0
  hooks:

sbijax/__init__.py

@@ -2,7 +2,7 @@
 sbijax: Simulation-based inference in JAX
 """
 
-__version__ = "0.1.2"
+__version__ = "0.1.3"
 
 
 from sbijax.abc.rejection_abc import RejectionABC

sbijax/snl.py

@@ -113,6 +113,7 @@ def loss_fn(params):
 
  losses = np.zeros([n_iter, 2])
  early_stop = EarlyStopping(1e-3, n_early_stopping_patience)
+ best_params, best_loss = None, np.inf
  logging.info("training model")
  for i in range(n_iter):
  train_loss = 0.0
@@ -129,9 +130,12 @@ def loss_fn(params):
  if early_stop.should_stop:
  logging.info("early stopping criterion found")
  break
+ if validation_loss < best_loss:
+ best_loss = validation_loss
+ best_params = params.copy()
 
  losses = jnp.vstack(losses)[:i, :]
- return params, losses
+ return best_params, losses
 
  def _validation_loss(self, params, val_iter):
  @jax.jit

sbijax/snp.py

@@ -9,11 +9,10 @@
 from jax import scipy as jsp
 
 from sbijax._sne_base import SNE
-
-# pylint: disable=too-many-arguments,unused-argument
 from sbijax.nn.early_stopping import EarlyStopping
 
 
+# pylint: disable=too-many-arguments,unused-argument
 class SNP(SNE):
  """
  Sequential neural posterior estimation
@@ -133,6 +132,7 @@ def step(params, rng, state, **batch):
 
  losses = np.zeros([n_iter, 2])
  early_stop = EarlyStopping(1e-3, n_early_stopping_patience)
+ best_params, best_loss = None, np.inf
  logging.info("training model")
  for i in range(n_iter):
  train_loss = 0.0
@@ -156,10 +156,13 @@ def step(params, rng, state, **batch):
  if early_stop.should_stop:
  logging.info("early stopping criterion found")
  break
+ if validation_loss < best_loss:
+ best_loss = validation_loss
+ best_params = params.copy()
 
  self.n_round += 1
  losses = jnp.vstack(losses)[:i, :]
- return params, losses
+ return best_params, losses
 
  def _init_params(self, rng_key, **init_data):
  params = self.model.init(