flowMC 0.3.3__tar.gz → 0.3.4__tar.gz

{flowmc-0.3.3/src/flowMC.egg-info → flowmc-0.3.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flowMC
-Version: 0.3.3
+Version: 0.3.4
 Summary: Normalizing flow exhanced sampler in jax
 Home-page: https://github.com/kazewong/flowMC
 Author: Kaze Wong, Marylou Gabrié, Dan Foreman-Mackey

{flowmc-0.3.3 → flowmc-0.3.4}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = flowMC
-version = 0.3.3
+version = 0.3.4
 author = Kaze Wong, Marylou Gabrié, Dan Foreman-Mackey
 author_email = kazewong.physics@gmail.com
 url = https://github.com/kazewong/flowMC

{flowmc-0.3.3 → flowmc-0.3.4}/src/flowMC/Sampler.py

@@ -271,7 +271,7 @@ class Sampler:
         Args:
             path (str): Path to save the normalizing flow.
         """
-        self.nf_model.load_model(path)
+        self.nf_model = self.nf_model.load_model(path)
 
     def reset(self):
         """
@@ -388,3 +388,46 @@ class Sampler:
         """
         with open(path, "wb") as f:
             pickle.dump(self.summary, f)
+
+    def print_summary(self) -> None:
+        """
+        Print summary to the screen about log probabilities and local/global acceptance rates.
+        """
+        train_summary = self.get_sampler_state(training=True)
+        production_summary = self.get_sampler_state(training=False)
+
+        training_log_prob = train_summary["log_prob"]
+        training_local_acceptance = train_summary["local_accs"]
+        training_global_acceptance = train_summary["global_accs"]
+        training_loss = train_summary["loss_vals"]
+
+        production_log_prob = production_summary["log_prob"]
+        production_local_acceptance = production_summary["local_accs"]
+        production_global_acceptance = production_summary["global_accs"]
+
+        print("Training summary")
+        print("=" * 10)
+        print(
+            f"Log probability: {training_log_prob.mean():.3f} +/- {training_log_prob.std():.3f}"
+        )
+        print(
+            f"Local acceptance: {training_local_acceptance.mean():.3f} +/- {training_local_acceptance.std():.3f}"
+        )
+        print(
+            f"Global acceptance: {training_global_acceptance.mean():.3f} +/- {training_global_acceptance.std():.3f}"
+        )
+        print(
+            f"Max loss: {training_loss.max():.3f}, Min loss: {training_loss.min():.3f}"
+        )
+
+        print("Production summary")
+        print("=" * 10)
+        print(
+            f"Log probability: {production_log_prob.mean():.3f} +/- {production_log_prob.std():.3f}"
+        )
+        print(
+            f"Local acceptance: {production_local_acceptance.mean():.3f} +/- {production_local_acceptance.std():.3f}"
+        )
+        print(
+            f"Global acceptance: {production_global_acceptance.mean():.3f} +/- {production_global_acceptance.std():.3f}"
+        )

{flowmc-0.3.3 → flowmc-0.3.4}/src/flowMC/nfmodel/base.py

@@ -78,7 +78,7 @@ class NFModel(eqx.Module):
         eqx.tree_serialise_leaves(path + ".eqx", self)
 
     def load_model(self, path: str):
-        self = eqx.tree_deserialise_leaves(path + ".eqx", self)
+        return eqx.tree_deserialise_leaves(path + ".eqx", self)
 
     @eqx.filter_value_and_grad
     def loss_fn(self, x):

{flowmc-0.3.3 → flowmc-0.3.4/src/flowMC.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flowMC
-Version: 0.3.3
+Version: 0.3.4
 Summary: Normalizing flow exhanced sampler in jax
 Home-page: https://github.com/kazewong/flowMC
 Author: Kaze Wong, Marylou Gabrié, Dan Foreman-Mackey