flowMC 0.3.2__tar.gz → 0.3.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flowMC
-Version: 0.3.2
+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.2 → flowmc-0.3.4}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = flowMC
-version = 0.3.2
+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.2 → 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.2 → 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.2 → flowmc-0.3.4}/src/flowMC/strategy/optimization.py

@@ -10,23 +10,31 @@ from flowMC.strategy.base import Strategy
 
 
 class optimization_Adam(Strategy):
-
     """
     Optimize a set of chains using Adam optimization.
     Note that if the posterior can go to infinity, this optimization scheme is likely to return NaNs.
-    
+
+    Args:
+        n_steps: int = 100
+            Number of optimization steps.
+        learning_rate: float = 1e-2
+            Learning rate for the optimization.
+        noise_level: float = 10
+            Variance of the noise added to the gradients.
     """
 
     n_steps: int = 100
     learning_rate: float = 1e-2
     noise_level: float = 10
+    bounds: Float[Array, "n_dim 2"] = jnp.array([[-jnp.inf, jnp.inf]])
 
     @property
     def __name__(self):
         return "AdamOptimization"
-    
+
     def __init__(
         self,
+        bounds: Float[Array, "n_dim 2"] = jnp.array([[-jnp.inf, jnp.inf]]),
         **kwargs,
     ):
         class_keys = list(self.__class__.__annotations__.keys())
@@ -39,6 +47,8 @@ class optimization_Adam(Strategy):
             optax.adam(learning_rate=self.learning_rate),
         )
 
+        self.bounds = bounds
+
     def __call__(
         self,
         rng_key: PRNGKeyArray,
@@ -61,6 +71,7 @@ class optimization_Adam(Strategy):
             grad = grad_fn(params) * (1 + jax.random.normal(subkey) * self.noise_level)
             updates, opt_state = self.solver.update(grad, opt_state, params)
             params = optax.apply_updates(params, updates)
+            params = optax.projections.projection_box(params, self.bounds[:, 0], self.bounds[:, 1])
             return (key, params, opt_state), None
 
         def _single_optimize(
@@ -91,11 +102,77 @@ class optimization_Adam(Strategy):
         summary["final_positions"] = optimized_positions
         summary["final_log_prob"] = local_sampler.logpdf_vmap(optimized_positions, data)
 
-        if jnp.isinf(summary['final_log_prob']).any() or jnp.isnan(summary['final_log_prob']).any():
+        if (
+            jnp.isinf(summary["final_log_prob"]).any()
+            or jnp.isnan(summary["final_log_prob"]).any()
+        ):
             print("Warning: Optimization accessed infinite or NaN log-probabilities.")
-            
+
         return rng_key, optimized_positions, local_sampler, global_sampler, summary
 
+    def optimize(
+        self,
+        rng_key: PRNGKeyArray,
+        objective: Callable,
+        initial_position: Float[Array, " n_chain n_dim"],
+    ):
+        """
+        Standalone optimization function that takes an objective function and returns the optimized positions.
+
+        Args:
+            rng_key: PRNGKeyArray
+                Random key for the optimization.
+            objective: Callable
+                Objective function to optimize.
+            initial_position: Float[Array, " n_chain n_dim"]
+                Initial positions for the optimization.
+        """
+        grad_fn = jax.jit(jax.grad(objective))
+
+        def _kernel(carry, data):
+            key, params, opt_state = carry
+
+            key, subkey = jax.random.split(key)
+            grad = grad_fn(params) * (1 + jax.random.normal(subkey) * self.noise_level)
+            updates, opt_state = self.solver.update(grad, opt_state, params)
+            params = optax.apply_updates(params, updates)
+            return (key, params, opt_state), None
+
+        def _single_optimize(
+            key: PRNGKeyArray,
+            initial_position: Float[Array, " n_dim"],
+        ) -> Float[Array, " n_dim"]:
+
+            opt_state = self.solver.init(initial_position)
+
+            (key, params, opt_state), _ = jax.lax.scan(
+                _kernel,
+                (key, initial_position, opt_state),
+                jnp.arange(self.n_steps),
+            )
+
+            return params  # type: ignore
+
+        print("Using Adam optimization")
+        rng_key, subkey = jax.random.split(rng_key)
+        keys = jax.random.split(subkey, initial_position.shape[0])
+        optimized_positions = jax.vmap(_single_optimize, in_axes=(0, 0))(
+            keys, initial_position
+        )
+
+        summary = {}
+        summary["initial_positions"] = initial_position
+        summary["initial_log_prob"] = jax.jit(jax.vmap(objective))(initial_position)
+        summary["final_positions"] = optimized_positions
+        summary["final_log_prob"] = jax.jit(jax.vmap(objective))(optimized_positions)
+
+        if (
+            jnp.isinf(summary["final_log_prob"]).any()
+            or jnp.isnan(summary["final_log_prob"]).any()
+        ):
+            print("Warning: Optimization accessed infinite or NaN log-probabilities.")
+
+        return rng_key, optimized_positions, summary
 
 class Evosax_CMA_ES(Strategy):
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flowMC
-Version: 0.3.2
+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