flowMC 0.3.1__tar.gz → 0.3.2__tar.gz

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

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

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

{flowMC-0.3.1 → flowmc-0.3.2}/src/flowMC/Sampler.py

@@ -131,7 +131,7 @@ class Sampler:
         )
         self.optim_state = self.optim.init(eqx.filter(self.nf_model, eqx.is_array))
 
-        self.strategies = [
+        default_strategies = [
             GlobalTuning(
                 n_dim=self.n_dim,
                 n_chains=self.n_chains,
@@ -161,7 +161,15 @@ class Sampler:
         if kwargs.get("strategies") is not None:
             kwargs_strategies = kwargs.get("strategies")
             assert isinstance(kwargs_strategies, list)
-            self.strategies = kwargs_strategies
+            self.strategies = []
+            for strategy in kwargs_strategies:
+                if isinstance(strategy, str):
+                    if strategy.lower() == "default":
+                        self.strategies += default_strategies
+                else:
+                    self.strategies.append(strategy)
+        else:
+            self.strategies = default_strategies 
 
         self.summary = {}
 

{flowMC-0.3.1 → flowmc-0.3.2}/src/flowMC/nfmodel/base.py

@@ -166,6 +166,7 @@ class NFModel(eqx.Module):
             pbar = range(num_epochs)
 
         best_model = model = self
+        best_state = state
         best_loss = 1e9
         for epoch in pbar:
             # Use a separate PRNG key to permute image data during shuffling

flowmc-0.3.2/src/flowMC/strategy/optimization.py

@@ -0,0 +1,120 @@
+import jax
+import jax.numpy as jnp
+import optax
+from jaxtyping import Array, Float, PRNGKeyArray, PyTree
+from typing import Callable
+
+from flowMC.proposal.base import ProposalBase
+from flowMC.proposal.NF_proposal import NFProposal
+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.
+    
+    """
+
+    n_steps: int = 100
+    learning_rate: float = 1e-2
+    noise_level: float = 10
+
+    @property
+    def __name__(self):
+        return "AdamOptimization"
+    
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        class_keys = list(self.__class__.__annotations__.keys())
+        for key, value in kwargs.items():
+            if key in class_keys:
+                if not key.startswith("__"):
+                    setattr(self, key, value)
+
+        self.solver = optax.chain(
+            optax.adam(learning_rate=self.learning_rate),
+        )
+
+    def __call__(
+        self,
+        rng_key: PRNGKeyArray,
+        local_sampler: ProposalBase,
+        global_sampler: NFProposal,
+        initial_position: Float[Array, " n_chain n_dim"],
+        data: dict,
+    ) -> tuple[
+        PRNGKeyArray, Float[Array, " n_chain n_dim"], ProposalBase, NFProposal, PyTree
+    ]:
+        def loss_fn(params: Float[Array, " n_dim"]) -> Float:
+            return -local_sampler.logpdf(params, data)
+
+        grad_fn = jax.jit(jax.grad(loss_fn))
+
+        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"] = local_sampler.logpdf_vmap(initial_position, data)
+        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():
+            print("Warning: Optimization accessed infinite or NaN log-probabilities.")
+            
+        return rng_key, optimized_positions, local_sampler, global_sampler, summary
+
+
+class Evosax_CMA_ES(Strategy):
+
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        class_keys = list(self.__class__.__annotations__.keys())
+        for key, value in kwargs.items():
+            if key in class_keys:
+                if not key.startswith("__"):
+                    setattr(self, key, value)
+
+    def __call__(
+        self,
+        rng_key: PRNGKeyArray,
+        local_sampler: ProposalBase,
+        global_sampler: NFProposal,
+        initial_position: Array,
+        data: dict,
+    ) -> tuple[PRNGKeyArray, Array, ProposalBase, NFProposal, PyTree]:
+        raise NotImplementedError

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flowMC
-Version: 0.3.1
+Version: 0.3.2
 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.1 → flowmc-0.3.2}/src/flowMC.egg-info/SOURCES.txt

@@ -25,6 +25,7 @@ src/flowMC/strategy/__init__.py
 src/flowMC/strategy/base.py
 src/flowMC/strategy/global_tuning.py
 src/flowMC/strategy/importance_sampling.py
+src/flowMC/strategy/optimization.py
 src/flowMC/utils/EvolutionaryOptimizer.py
 src/flowMC/utils/PythonFunctionWrap.py
 src/flowMC/utils/__init__.py