flowMC 0.3.0__tar.gz → 0.3.2__tar.gz

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

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

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

@@ -1,15 +1,15 @@
 import pickle
-import jax
-import jax.numpy as jnp
-from jaxtyping import Array, Int, Float, PRNGKeyArray
-from tqdm import tqdm
+
 import equinox as eqx
+import jax.numpy as jnp
 import optax
-from flowMC.proposal.NF_proposal import NFProposal
-from flowMC.proposal.base import ProposalBase
+from jaxtyping import Array, Float, Int, PRNGKeyArray
+
 from flowMC.nfmodel.base import NFModel
+from flowMC.proposal.base import ProposalBase
+from flowMC.proposal.NF_proposal import NFProposal
 from flowMC.strategy.base import Strategy
-from flowMC.strategy.global_tuning import GlobalTuning, GlobalSampling
+from flowMC.strategy.global_tuning import GlobalSampling, GlobalTuning
 
 
 class Sampler:
@@ -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.0 → flowmc-0.3.2}/src/flowMC/nfmodel/base.py

@@ -1,13 +1,14 @@
-from abc import abstractmethod
 import copy
+from abc import abstractmethod
+from typing import Optional, overload
+
 import equinox as eqx
-from typing import overload, Optional
-from typing_extensions import Self
-from jaxtyping import Array, PRNGKeyArray, Float
-import optax
-from tqdm import trange, tqdm
-import jax.numpy as jnp
 import jax
+import jax.numpy as jnp
+import optax
+from jaxtyping import Array, Float, PRNGKeyArray
+from tqdm import tqdm, trange
+from typing_extensions import Self
 
 
 class NFModel(eqx.Module):
@@ -102,9 +103,10 @@ class NFModel(eqx.Module):
             model (eqx.Model): Updated model.
             opt_state (optax.OptState): Updated optimizer state.
         """
-        loss, grads = self.loss_fn(x)
+        model = self
+        loss, grads = model.loss_fn(x)
         updates, state = optim.update(grads, state)
-        model = eqx.apply_updates(self, updates)
+        model = eqx.apply_updates(model, updates)
         return loss, model, state
 
     def train_epoch(
@@ -162,16 +164,21 @@ class NFModel(eqx.Module):
             pbar = trange(num_epochs, desc="Training NF", miniters=int(num_epochs / 10))
         else:
             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
             rng, input_rng = jax.random.split(rng)
             # Run an optimization step over a training batch
-            value, model, state = model.train_epoch(input_rng, optim, state, data, batch_size)
+            value, model, state = model.train_epoch(
+                input_rng, optim, state, data, batch_size
+            )
             loss_values = loss_values.at[epoch].set(value)
             if loss_values[epoch] < best_loss:
                 best_model = model
+                best_state = state
                 best_loss = loss_values[epoch]
             if verbose:
                 assert isinstance(pbar, tqdm)
@@ -182,7 +189,7 @@ class NFModel(eqx.Module):
                     if epoch == num_epochs:
                         pbar.set_description(f"Training NF, current loss: {value:.3f}")
 
-        return rng, best_model, state, loss_values
+        return rng, best_model, best_state, loss_values
 
 
 class Bijection(eqx.Module):

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/nfmodel/common.py

@@ -1,17 +1,12 @@
 from typing import Callable, List, Tuple
 
-import jax
-import jax.numpy as jnp
-from jaxtyping import Array
 import equinox as eqx
-
-from flowMC.nfmodel.base import Bijection, Distribution
-
 import jax
 import jax.numpy as jnp
-import equinox as eqx
 from jaxtyping import Array, Float, PRNGKeyArray
 
+from flowMC.nfmodel.base import Bijection, Distribution
+
 
 class MLP(eqx.Module):
     r"""Multilayer perceptron.

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/nfmodel/realNVP.py

@@ -1,13 +1,15 @@
+from functools import partial
 from typing import List, Tuple
+
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 import numpy as np
-import equinox as eqx
-from flowMC.nfmodel.base import NFModel, Distribution
-from flowMC.nfmodel.common import MLP, MaskedCouplingLayer, MLPAffine, Gaussian
-from jaxtyping import Array, Float, PRNGKeyArray
-from functools import partial
 import optax
+from jaxtyping import Array, Float, PRNGKeyArray
+
+from flowMC.nfmodel.base import Distribution, NFModel
+from flowMC.nfmodel.common import MLP, Gaussian, MaskedCouplingLayer, MLPAffine
 
 
 class AffineCoupling(eqx.Module):

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/nfmodel/rqSpline.py

@@ -1,11 +1,13 @@
+from functools import partial
+
+import equinox as eqx
 import jax
 import jax.numpy as jnp
-from jaxtyping import Array, PRNGKeyArray, Float
-import equinox as eqx
+from jaxtyping import Array, Float, PRNGKeyArray
 
-from flowMC.nfmodel.base import NFModel, Bijection, Distribution
-from flowMC.nfmodel.common import MaskedCouplingLayer, ScalarAffine, MLP, Gaussian
-from functools import partial
+from flowMC.nfmodel.base import Bijection, Distribution, NFModel
+from flowMC.nfmodel.common import (MLP, Gaussian, MaskedCouplingLayer,
+                                   ScalarAffine)
 
 
 @partial(jax.vmap, in_axes=(0, None, None))

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/proposal/Gaussian_random_walk.py

@@ -1,9 +1,11 @@
 from typing import Callable
+
 import jax
 import jax.numpy as jnp
+from jaxtyping import Array, Float, Int, PRNGKeyArray, PyTree
 from tqdm import tqdm
+
 from flowMC.proposal.base import ProposalBase
-from jaxtyping import PyTree, Array, Float, Int, PRNGKeyArray
 
 
 class GaussianRandomWalk(ProposalBase):

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/proposal/HMC.py

@@ -1,9 +1,11 @@
 from typing import Callable
+
 import jax
 import jax.numpy as jnp
+from jaxtyping import Array, Float, Int, PRNGKeyArray, PyTree
 from tqdm import tqdm
+
 from flowMC.proposal.base import ProposalBase
-from jaxtyping import PyTree, Array, Float, Int, PRNGKeyArray
 
 
 class HMC(ProposalBase):

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/proposal/MALA.py

@@ -1,11 +1,13 @@
+from functools import partialmethod
 from typing import Callable
+
 import jax
 import jax.numpy as jnp
 from jax.scipy.stats import multivariate_normal
+from jaxtyping import Array, Bool, Float, Int, PRNGKeyArray, PyTree
 from tqdm import tqdm
+
 from flowMC.proposal.base import ProposalBase
-from functools import partialmethod
-from jaxtyping import PyTree, Array, Float, Int, PRNGKeyArray, Bool
 
 
 class MALA(ProposalBase):

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/proposal/NF_proposal.py

@@ -1,12 +1,14 @@
+from math import ceil
+from typing import Callable
+
 import jax
 import jax.numpy as jnp
 from jax import random
+from jaxtyping import Array, Float, Int, PRNGKeyArray, PyTree
 from tqdm import tqdm
+
 from flowMC.nfmodel.base import NFModel
-from typing import Callable
 from flowMC.proposal.base import ProposalBase
-from jaxtyping import Array, Float, Int, PRNGKeyArray, PyTree
-from math import ceil
 
 
 @jax.tree_util.register_pytree_node_class

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/proposal/base.py

@@ -1,8 +1,9 @@
 from abc import abstractmethod
 from typing import Callable
+
 import jax
 import jax.numpy as jnp
-from jaxtyping import PyTree, Array, Float, Int, PRNGKeyArray
+from jaxtyping import Array, Float, Int, PRNGKeyArray, PyTree
 
 
 @jax.tree_util.register_pytree_node_class

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/proposal/flowHMC.py

@@ -1,14 +1,15 @@
+from math import ceil
+from typing import Callable
+
 import jax
 import jax.numpy as jnp
+from jax import random
+from jaxtyping import Array, Float, Int, PRNGKeyArray, PyTree
+from tqdm import tqdm
+
 from flowMC.nfmodel.base import NFModel
-from jaxtyping import Array, PRNGKeyArray, PyTree
-from typing import Callable
 from flowMC.proposal.HMC import HMC
 from flowMC.proposal.NF_proposal import NFProposal
-from jaxtyping import Array, Float, Int, PRNGKeyArray
-from math import ceil
-from jax import random
-from tqdm import tqdm
 
 ###################################
 # This is not in production yet

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/strategy/base.py

@@ -1,7 +1,11 @@
 from abc import abstractmethod
+
+from jaxtyping import Array, Float, PRNGKeyArray, PyTree
+
 from flowMC.proposal.base import ProposalBase
 from flowMC.proposal.NF_proposal import NFProposal
-from jaxtyping import Array, Float, PRNGKeyArray, PyTree
+
+
 class Strategy:
     """
     Base class for strategies, which are basically wrapper blocks that modify the state of the sampler

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/strategy/global_tuning.py

@@ -1,13 +1,14 @@
-from flowMC.proposal.base import ProposalBase
-from flowMC.proposal.NF_proposal import NFProposal
-from flowMC.strategy.base import Strategy
+import equinox as eqx
 import jax
 import jax.numpy as jnp
-from jaxtyping import Array, Float, PRNGKeyArray, PyTree
 import optax
-import equinox as eqx
+from jaxtyping import Array, Float, PRNGKeyArray, PyTree
 from tqdm import tqdm
 
+from flowMC.proposal.base import ProposalBase
+from flowMC.proposal.NF_proposal import NFProposal
+from flowMC.strategy.base import Strategy
+
 
 class GlobalTuning(Strategy):
 
@@ -76,6 +77,7 @@ class GlobalTuning(Strategy):
         summary["local_accs"] = jnp.empty((self.n_chains, 0))
         summary["global_accs"] = jnp.empty((self.n_chains, 0))
         summary["loss_vals"] = jnp.empty((0, self.n_epochs))
+        current_position = initial_position
         for _ in tqdm(
             range(self.n_loop),
             desc="Global Tuning",
@@ -90,7 +92,7 @@ class GlobalTuning(Strategy):
             ) = local_sampler.sample(
                 rng_keys_mcmc,
                 self.n_local_steps,
-                initial_position,
+                current_position,
                 data,
                 verbose=self.verbose,
             )
@@ -112,6 +114,8 @@ class GlobalTuning(Strategy):
                 axis=1,
             )
 
+            current_position = summary["chains"][:, -1]
+
             rng_key, rng_keys_nf = jax.random.split(rng_key)
             positions = summary["chains"][:, :: self.train_thinning]
             chain_size = positions.shape[0] * positions.shape[1]
@@ -142,8 +146,8 @@ class GlobalTuning(Strategy):
 
             (
                 rng_keys_nf,
-                global_sampler.model,
-                self.optim_state,
+                model,
+                optim_state,
                 loss_values,
             ) = global_sampler.model.train(
                 rng_keys_nf,
@@ -154,6 +158,8 @@ class GlobalTuning(Strategy):
                 self.batch_size,
                 self.verbose,
             )
+            global_sampler.model = model
+            self.optim_state = optim_state
             summary["loss_vals"] = jnp.append(
                 summary["loss_vals"],
                 loss_values.reshape(1, -1),
@@ -168,7 +174,7 @@ class GlobalTuning(Strategy):
             ) = global_sampler.sample(
                 rng_keys_nf,
                 self.n_global_steps,
-                positions[:, -1],
+                current_position,
                 data,
                 verbose=self.verbose,
             )
@@ -190,7 +196,9 @@ class GlobalTuning(Strategy):
                 axis=1,
             )
 
-        return rng_key, summary['chains'][:, -1], local_sampler, global_sampler, summary
+            current_position = summary["chains"][:, -1]
+
+        return rng_key, current_position, local_sampler, global_sampler, summary
 
 
 class GlobalSampling(Strategy):
@@ -239,6 +247,7 @@ class GlobalSampling(Strategy):
         summary["local_accs"] = jnp.empty((self.n_chains, 0))
         summary["global_accs"] = jnp.empty((self.n_chains, 0))
 
+        current_position = initial_position
         for _ in tqdm(
             range(self.n_loop),
             desc="Global Sampling",
@@ -253,7 +262,7 @@ class GlobalSampling(Strategy):
             ) = local_sampler.sample(
                 rng_keys_mcmc,
                 self.n_local_steps,
-                initial_position,
+                current_position,
                 data,
                 verbose=self.verbose,
             )
@@ -275,6 +284,8 @@ class GlobalSampling(Strategy):
                 axis=1,
             )
 
+            current_position = summary["chains"][:, -1]
+
             rng_key, rng_keys_nf = jax.random.split(rng_key)
             (
                 rng_keys_nf,
@@ -306,4 +317,6 @@ class GlobalSampling(Strategy):
                 axis=1,
             )
 
+            current_position = summary["chains"][:, -1]
+
         return rng_key, summary['chains'][:, -1], local_sampler, global_sampler, summary

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.0 → flowmc-0.3.2}/src/flowMC/utils/PythonFunctionWrap.py

@@ -1,6 +1,7 @@
 import warnings
 from functools import wraps
-from typing import Any, List, Dict, Callable, Iterable, Tuple, NamedTuple, Union
+from typing import (Any, Callable, Dict, Iterable, List, NamedTuple, Tuple,
+                    Union)
 
 Array = Any
 PyTree = Union[Array, Iterable[Array], Dict[Any, Array], NamedTuple]

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/utils/postprocessing.py

@@ -1,7 +1,8 @@
-import matplotlib.pyplot as plt
 import jax.numpy as jnp
+import matplotlib.pyplot as plt
 # from flowMC.sampler.Sampler import Sampler
-from jaxtyping import Float, Array
+from jaxtyping import Array, Float
+
 
 def plot_summary(sampler: object, training: bool = False, **plotkwargs) -> None:
     """

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: flowMC
-Version: 0.3.0
+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.0 → 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

{flowMC-0.3.0 → flowmc-0.3.2}/src/flowMC/utils/EvolutionaryOptimizer.py

@@ -1,8 +1,8 @@
-from evosax import CMA_ES
 import jax
 import jax.numpy as jnp
-from jaxtyping import PRNGKeyArray
 import tqdm
+from evosax import CMA_ES
+from jaxtyping import PRNGKeyArray
 
 
 class EvolutionaryOptimizer: