Repository: htjb/margarine (language: python)

Issue:
Tqdm and small compile
A small quality of life request in a progress bar for training. 

A second small ( but potentially breaking change?) in wrapping the train step in a jit wrapper. Naively gives a factor 2 improvement in run time on CPU and GPU as far as I can see, but may need more advanced checking on other hardware.

More broadly I tried putting the training on a GPU and saw little speed up, I started digging through the code and I think there are quite a lot of conversions to numpy and non-TF operations. So at the moment there isn't much gain from GPU usage here, but the JIT compilation still looks advantageous.
Some updates to improve speed of sampling
Decorating the relevant functions inside a `MAF.__call__` with a @tf.function(jit_compile=True) decorator to speed up the sampling. Seen marginal improvements on some test examples.
Runtime
Curious whether the code can be sped up by using the `tf.function` tool provided in tensorflow to wrap the `train_step` function in `maf.py`. `tf.function` converts functions in to graphs rather than executing them using eager execution.
Tqdm and small compile
A small quality of life request in a progress bar for training. 

A second small ( but potentially breaking change?) in wrapping the train step in a jit wrapper. Naively gives a factor 2 improvement in run time on CPU and GPU as far as I can see, but may need more advanced checking on other hardware.

More broadly I tried putting the training on a GPU and saw little speed up, I started digging through the code and I think there are quite a lot of conversions to numpy and non-TF operations. So at the moment there isn't much gain from GPU usage here, but the JIT compilation still looks advantageous.
MPI Clustering
Need to get MPI training set up for the clustering so that the cluster MAFs (or KDEs #29) are trained in parallel rather than one after the other.


Discussion / hints:
Ahh thanks David! Both of these things looks super helpful. I did play around with jit compile a bit as I had seen it lead to improvements with some other work I had been doing but just hadn't got around to implementing it here. Should help solve #20. 

This shouldn't be a breaking change but we will have to bump the version number up to 0.6.0.

@yallup need to add tqdm to the requirements.txt for the tests I think.
Hmm yes the numpy to tf conversions probably aren't helping with runtime. I think speed up from GPUs is limited for small models as well because there is an overhead to running on GPU.
Hi @yallup, have you had a chance to look at this? Be great to get these features merged in. Thanks!
For some reason, the clustering is causing a 

```bash
ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information.
```
error with the `@tf.function(jit_compile=True)`. Happens on the `test_maf_clustering()` test which is the only one that has a `clustering=True` in the MAF(). If I turn this of the test passes.
I think it's because tensorflow is trying to create the function multiple times... in the for loop

```python
for i in range(len(self.theta)):
        self.maf[i] = self._training(self.theta[i],
                                     self.sample_weights[i],
                                     self.maf[i], self.theta_min[i],
                                     self.theta_max[i])
```

which I can validate by running

```python
self.maf[0] = self._training(self.theta[0],
                                 self.sample_weights[0],
                                 self.maf[0], self.theta_min[0],
                                 self.theta_max[0])
self.maf[1] = self._training(self.theta[1],
                                 self.sample_weights[1],
                                 self.maf[1], self.theta_min[1],
                                 self.theta_max[1])
```
which gives the same error when it gets to self.maf[1].
@yallup I've been trying to resolve this issue but couldn't do it. I'm keen to take advantage of the tqdm addition though so I have removed the `@tf.function(jit_compile=True)` decorator on the training step for now.

I will come back to this when I look at #31 again. I think I has the same error with the clustering and @tf.function decorator there as well. If the tests pass could you please squash and merge. Thanks!
## [Codecov](https://app.codecov.io/gh/htjb/margarine/pull/26?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins) Report
Patch coverage: **`100.00%`** and project coverage change: **`+0.02%`** :tada:
> Comparison is base [(`ae63d2a`)](https://app.codecov.io/gh/htjb/margarine/commit/ae63d2ab66daa949edbf5f3320d86b8cebacf78e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins) 88.93% compared to head [(`b34b8f9`)](https://app.codecov.io/gh/htjb/margarine/pull/26?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins) 88.95%.

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master      #26      +/-   ##
==========================================
+ Coverage   88.93%   88.95%   +0.02%     
==========================================
  Files           4        4              
  Lines         470      471       +1     
==========================================
+ Hits          418      419       +1     
  Misses         52       52              
```


| [Files Changed](https://app.codecov.io/gh/htjb/margarine/pull/26?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins) | Coverage Δ | |
|---|---|---|
| [margarine/maf.py](https://app.codecov.io/gh/htjb/margarine/pull/26?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins#diff-bWFyZ2FyaW5lL21hZi5weQ==) | `95.97% <100.00%> (+0.01%)` | :arrow_up: |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/htjb/margarine/pull/26?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins).

Totally agree on putting a pin in the speedup stuff, it was just a little thing I was trying. I got a bit caught up trying out some other bits such as reducing translations between tf and numpy operations as that might be a bigger speedup in the long run, one to put a pin in and revisit I think!
No worries! Thanks for having a look at it. I started the process of rewriting the MAF class line by line this evening to try and eliminate some of the back and forth between numpy/tensorflow and to add in some tf.function decorators where possible. Seems to be going well and hopefully that will help with run time issues. I'll put a PR in at some point.
Reviewing runtime (#20) in more detail with the intention of making a more complete PR.

Ahh thanks David! Both of these things looks super helpful. I did play around with jit compile a bit as I had seen it lead to improvements with some other work I had been doing but just hadn't got around to implementing it here. Should help solve #20. 

This shouldn't be a breaking change but we will have to bump the version number up to 0.6.0.

@yallup need to add tqdm to the requirements.txt for the tests I think.
Hmm yes the numpy to tf conversions probably aren't helping with runtime. I think speed up from GPUs is limited for small models as well because there is an overhead to running on GPU.
Hi @yallup, have you had a chance to look at this? Be great to get these features merged in. Thanks!
For some reason, the clustering is causing a 

```bash
ValueError: tf.function only supports singleton tf.Variables created on the first call. Make sure the tf.Variable is only created once or created outside tf.function. See https://www.tensorflow.org/guide/function#creating_tfvariables for more information.
```
error with the `@tf.function(jit_compile=True)`. Happens on the `test_maf_clustering()` test which is the only one that has a `clustering=True` in the MAF(). If I turn this of the test passes.
I think it's because tensorflow is trying to create the function multiple times... in the for loop

```python
for i in range(len(self.theta)):
        self.maf[i] = self._training(self.theta[i],
                                     self.sample_weights[i],
                                     self.maf[i], self.theta_min[i],
                                     self.theta_max[i])
```

which I can validate by running

```python
self.maf[0] = self._training(self.theta[0],
                                 self.sample_weights[0],
                                 self.maf[0], self.theta_min[0],
                                 self.theta_max[0])
self.maf[1] = self._training(self.theta[1],
                                 self.sample_weights[1],
                                 self.maf[1], self.theta_min[1],
                                 self.theta_max[1])
```
which gives the same error when it gets to self.maf[1].
@yallup I've been trying to resolve this issue but couldn't do it. I'm keen to take advantage of the tqdm addition though so I have removed the `@tf.function(jit_compile=True)` decorator on the training step for now.

I will come back to this when I look at #31 again. I think I has the same error with the clustering and @tf.function decorator there as well. If the tests pass could you please squash and merge. Thanks!
## [Codecov](https://app.codecov.io/gh/htjb/margarine/pull/26?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins) Report
Patch coverage: **`100.00%`** and project coverage change: **`+0.02%`** :tada:
> Comparison is base [(`ae63d2a`)](https://app.codecov.io/gh/htjb/margarine/commit/ae63d2ab66daa949edbf5f3320d86b8cebacf78e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins) 88.93% compared to head [(`b34b8f9`)](https://app.codecov.io/gh/htjb/margarine/pull/26?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins) 88.95%.

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master      #26      +/-   ##
==========================================
+ Coverage   88.93%   88.95%   +0.02%     
==========================================
  Files           4        4              
  Lines         470      471       +1     
==========================================
+ Hits          418      419       +1     
  Misses         52       52              
```


| [Files Changed](https://app.codecov.io/gh/htjb/margarine/pull/26?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins) | Coverage Δ | |
|---|---|---|
| [margarine/maf.py](https://app.codecov.io/gh/htjb/margarine/pull/26?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins#diff-bWFyZ2FyaW5lL21hZi5weQ==) | `95.97% <100.00%> (+0.01%)` | :arrow_up: |


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/htjb/margarine/pull/26?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=Harry+Bevins).

Totally agree on putting a pin in the speedup stuff, it was just a little thing I was trying. I got a bit caught up trying out some other bits such as reducing translations between tf and numpy operations as that might be a bigger speedup in the long run, one to put a pin in and revisit I think!
No worries! Thanks for having a look at it. I started the process of rewriting the MAF class line by line this evening to try and eliminate some of the back and forth between numpy/tensorflow and to add in some tf.function decorators where possible. Seems to be going well and hopefully that will help with run time issues. I'll put a PR in at some point.


Fix patch:
diff --git a/README.rst b/README.rst
index 742fb94..6629571 100755
--- a/README.rst
+++ b/README.rst
@@ -7,7 +7,7 @@ Introduction
 
 :margarine: Marginal Bayesian Statistics
 :Authors: Harry T.J. Bevins
-:Version: 0.6.0
+:Version: 1.0.0
 :Homepage:  https://github.com/htjb/margarine
 :Documentation: https://margarine.readthedocs.io/
 
diff --git a/docs/source/functions.rst b/docs/source/functions.rst
index 2c1519b..44fbf48 100644
--- a/docs/source/functions.rst
+++ b/docs/source/functions.rst
@@ -8,22 +8,23 @@ margarine.maf.MAF()
 ~~~~~~~~~~~~~~~~~~~~
 
 .. automodule:: margarine.maf.MAF()
-    :members: _train_step, train, __call__, sample, save, load
+    :members: gen_mades, train, _train_step, _test_step, _training, __call__, sample, log_prob, log_like, save, load
 
 margarine.kde.KDE()
 ~~~~~~~~~~~~~~~~~~~~
 
 .. automodule:: margarine.kde.KDE()
-    :members: generate_kde, __call__, sample, save, load
+    :members: generate_kde, __call__, sample, log_prob, log_like, save, load
 
-margarine.marginal_stats.maf_calculations()
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+margarine.clustered.clusterMAF()
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. automodule:: margarine.marginal_stats.maf_calculations
-    :members: _calc_logL, klDiv, bayesian_dimensionality
+.. automodule:: margarine.clustered.clusterMAF()
+    :members: train, __call__, sample, log_prob, log_like, save, load
 
-margarine.marginal_stats.kde_calculations()
+margarine.marginal_stats.calculate()
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-.. automodule:: margarine.marginal_stats.kde_calculations
-    :members: _calc_logL, klDiv, bayesian_dimensionality
+.. automodule:: margarine.marginal_stats.calculate
+    :members: statistics
+
diff --git a/margarine/clustered.py b/margarine/clustered.py
new file mode 100644
index 0000000..ac5f7a9
--- /dev/null
+++ b/margarine/clustered.py
@@ -0,0 +1,463 @@
+from scipy.special import logsumexp
+from sklearn.cluster import KMeans
+from tensorflow import keras
+from margarine.maf import MAF
+import anesthetic
+import numpy as np
+import tensorflow as tf
+import warnings
+import pickle
+
+
+class clusterMAF():
+
+    r"""
+
+    This class is used to train, load and call a piecewise normalising flow
+    built from a set of masked autoregressive flows. The class
+    is essentially a wrapper around the MAF class with some additional
+    clustering functionality. This class has all the same
+    functionality as the MAF class and can be used in the same way.
+
+    **Parameters:**
+
+        theta: **numpy array or anesthetic.samples**
+            | The samples from the probability distribution that we require the
+                MAF to learn. This can either be a numpy array or an anesthetic
+                NestedSamples or MCMCSamples object.
+
+    **kwargs:**
+
+        weights: **numpy array / default=np.ones(len(theta))**
+            | The weights associated with the samples above. If an anesthetic
+                NestedSamples or MCMCSamples object is passed the code
+                draws the weights from this.
+
+        number_networks: **int / default = 6**
+            | The bijector is built by chaining a series of
+                autoregressive neural
+                networks together and this parameter is used to determine
+                how many networks there are in the chain.
+
+        learning_rate: **float / default = 1e-3**
+            | The learning rate determines the 'step size' of the optimization
+                algorithm used to train the MAF. Its value can effect the
+                quality of emulation.
+
+        hidden_layers: **list / default = [50, 50]**
+            | The number of layers and number of nodes in each hidden layer for
+                each neural network. The default is two hidden layers with
+                50 nodes each and each network in the chain has the same hidden
+                layer structure.
+
+        cluster_labels: **list / default = None**
+            | If clustering has been performed externally to margarine you can
+                provide a list of labels for the samples theta. The labels
+                should be integers from 0 to k corresponding to the cluster
+                that each sample is in.
+
+        cluster_number: **int / default = None**
+            | If clustering has been performed externally to margarine you
+                need to provide the number of clusters, k, alongside the
+                cluster labels.
+
+        parameters: **list of strings**
+            | A list of the relevant parameters to train on. Only needed
+                if theta is an anestehetic samples object. If not provided,
+                all parameters will be used.
+
+    """
+
+    def __init__(self, theta, **kwargs):
+
+        self.number_networks = kwargs.pop('number_networks', 6)
+        self.learning_rate = kwargs.pop('learning_rate', 1e-3)
+        self.hidden_layers = kwargs.pop('hidden_layers', [50, 50])
+        self.cluster_labels = kwargs.pop('cluster_labels', None)
+        self.cluster_number = kwargs.pop('cluster_number', None)
+        self.parameters = kwargs.pop('parameters', None)
+
+        if isinstance(theta, 
+                      (anesthetic.samples.NestedSamples, 
+                       anesthetic.samples.MCMCSamples)):
+            weights = theta.get_weights()
+            if self.parameters:
+                theta = theta[self.parameters].values
+            else:
+                if isinstance(theta, anesthetic.samples.NestedSamples):
+                    self.parameters = theta.columns[:-3].values
+                    theta = theta[theta.columns[:-3]].values
+                else:
+                    self.parameters = theta.columns[:-1].values
+                    theta = theta[theta.columns[:-1]].values
+        else:
+            weights = kwargs.pop('weights', np.ones(len(theta)))
+
+        self.theta = theta
+        self.sample_weights = weights
+
+        # needed for marginal stats
+        self.n = np.sum(self.sample_weights)**2 / \
+            np.sum(self.sample_weights**2)
+        theta_max = np.max(self.theta, axis=0)
+        theta_min = np.min(self.theta, axis=0)
+        a = ((self.n-2)*theta_max-theta_min)/(self.n-3)
+        b = ((self.n-2)*theta_min-theta_max)/(self.n-3)
+        self.theta_min = b
+        self.theta_max = a
+
+        if type(self.number_networks) is not int:
+            raise TypeError("'number_networks' must be an integer.")
+        if not isinstance(self.learning_rate,
+                          (int, float,
+                           keras.optimizers.schedules.LearningRateSchedule)):
+            raise TypeError("'learning_rate', " +
+                            "must be an integer, float or keras scheduler.")
+        if type(self.hidden_layers) is not list:
+            raise TypeError("'hidden_layers' must be a list of integers.")
+        else:
+            for i in range(len(self.hidden_layers)):
+                if type(self.hidden_layers[i]) is not int:
+                    raise TypeError(
+                        "One or more valus in 'hidden_layers'" +
+                        "is not an integer.")
+
+        mask = np.isfinite(theta).all(axis=-1)
+        self.theta = self.theta[mask]
+        self.sample_weights = self.sample_weights[mask]
+
+        if self.cluster_number is not None:
+            if self.cluster_labels is None:
+                raise ValueError("'cluster_labels' should be provided if " +
+                                 "'cluster_number' is specified.")
+        else:
+            if self.cluster_labels is not None:
+                raise ValueError("'cluster_number' should be provided if " +
+                                 "'cluster_labels' is specified.")
+
+        if self.cluster_number is not None:
+            if self.cluster_number % 1 != 0:
+                raise TypeError("'cluster_number' must be a whole number.")
+        if self.cluster_labels is not None:
+            if not isinstance(self.cluster_labels, (np.ndarray, list)):
+                raise TypeError("'cluster_labels' must be an array " +
+                                "or a list.")
+
+        if self.cluster_number is None:
+            # this code currently performs the clustering with
+            # kmeans but we could have anything here.
+            from sklearn.metrics import silhouette_score
+            ks = np.arange(2, 21)
+            losses = []
+            for k in ks:
+                kmeans = KMeans(k, random_state=0)
+                labels = kmeans.fit(self.theta).predict(self.theta)
+                losses.append(-silhouette_score(self.theta, labels))
+            losses = np.array(losses)
+            minimum_index = np.argmin(losses)
+            self.cluster_number = ks[minimum_index]
+
+            kmeans = KMeans(self.cluster_number, random_state=0)
+            self.cluster_labels = kmeans.fit(self.theta).predict(self.theta)
+
+        if self.cluster_number == 20:
+            warnings.warn("The number of clusters is 20. " +
+                          "This is the maximum " +
+                          "number of clusters that can be used. If you " +
+                          "require more clusters, please specify the " +
+                          "'cluster_number' kwarg. " +
+                          "margarine will continue " +
+                          "with 20 clusters.")
+
+        if np.array(list(self.cluster_labels)).dtype == 'float':
+            # convert cluster labels to integers
+            self.cluster_labels = self.cluster_labels.astype(int)
+        # count the number of times a cluster label appears in cluster_labels
+        self.cluster_count = np.bincount(self.cluster_labels)
+        # While loop to make sure clusters are not too small
+        while self.cluster_count.min() < 100:
+            warnings.warn("One or more clusters are too small " +
+                          "(n_cluster < 100). " +
+                          "Reducing the number of clusters by 1.")
+            minimum_index -= 1
+            self.cluster_number = ks[minimum_index]
+            kmeans = KMeans(self.cluster_number, random_state=0)
+            self.cluster_labels = kmeans.fit(self.theta).predict(self.theta)
+            self.cluster_count = np.bincount(self.cluster_labels)
+            if self.cluster_number == 2:
+                # break if two clusters
+                warnings.warn("The number of clusters is 2. This is the " +
+                              "minimum number of clusters that can be used. " +
+                              "Some clusters may be too small and the " +
+                              "train/test split may fail." +
+                              "Try running without clusting. ")
+                break
+
+        split_theta = []
+        split_sample_weights = []
+        for i in range(self.cluster_number):
+            split_theta.append(self.theta[self.cluster_labels == i])
+            split_sample_weights.append(
+                self.sample_weights[self.cluster_labels == i])
+
+        self.split_theta = split_theta
+        self.split_sample_weights = split_sample_weights
+
+        self.flow = []
+        for i in range(len(split_theta)):
+            self.flow.append(MAF(split_theta[i], 
+                                 weights=split_sample_weights[i],
+                                 number_networks=self.number_networks,
+                                 learning_rate=self.learning_rate,
+                                 hidden_layers=self.hidden_layers))
+
+    def train(self, epochs=100, early_stop=False, loss_type='sum'):
+
+        r"""
+        This function is called to train the clusterMAF once it has been
+        initialised. It calls the `train()` function for each flow.
+
+        .. code:: python
+
+            from margarine.clustered import clusterMAF
+
+            bij = clusterMAF(theta, weights=weights)
+            bij.train()
+
+        **Kwargs:**
+
+            epochs: **int / default = 100**
+                | The number of iterations to train the neural networks for.
+
+            early_stop: **boolean / default = False**
+                | Determines whether or not to implement an early stopping
+                    algorithm or
+                    train for the set number of epochs. If set to True then the
+                    algorithm will stop training when test loss has not
+                    improved for 2% of the requested epochs. At this point
+                    margarine will roll back to the best model and return this
+                    to the user.
+
+            loss_type: **string / default = 'sum'**
+                | Determines whether to use the sum or mean of the weighted
+                    log probabilities to calculate the loss function.
+
+        """
+
+        for i in range(len(self.flow)):
+            self.flow[i].train(epochs=epochs, early_stop=early_stop,
+                               loss_type=loss_type)
+
+    def log_prob(self, params):
+
+        """
+        Function to caluclate the log-probability for a given clusterMAF and
+        set of parameters.
+
+        While each density estimator has its own built in log probability
+        function, a correction has to be applied for the transformation of
+        variables that is used to improve accuracy when learning and we
+        have to sum probabilities over the series of flows. The
+        correction and the sum are implemented here.
+
+        **Parameters:**
+
+            params: **numpy array**
+                | The set of samples for which to calculate the log
+                    probability.
+
+        """
+
+        # currently working with numpy not tensorflow
+        # nan repalacement with 0 difficult in tensorflow
+        logprob = []
+        for flow in self.flow:
+            probs = flow.log_prob(params).numpy()
+            for j in range(len(probs)):
+                if np.isnan(probs[j]):
+                    probs[j] = np.log(1e-300)
+            logprob.append(probs)
+        logprob = np.array(logprob)
+        print(logprob.shape)
+        logprob = logsumexp(logprob, axis=0)
+
+        return logprob
+
+    def log_like(self, params, logevidence, prior_de=None):
+
+        r"""
+        This function should return the log-likelihood for a given set of
+        parameters.
+
+        It requires the logevidence from the original nested sampling run
+        in order to do this and in the case that the prior is non-uniform
+        a trained prior density estimator should be provided.
+
+        **Parameters:**
+
+            params: **numpy array**
+                | The set of samples for which to calculate the log
+                    probability.
+
+            logevidence: **float**
+                | Should be the log-evidence from the full nested sampling
+                    run with nuisance parameters.
+
+            prior_de: **margarine.maf.MAF / default=None**
+                | If the prior is non-uniform then a trained prior density
+                    estimator should be provided. Otherwise the prior
+                    is assumed to be uniform and the prior probability
+                    is calculated analytically from the minimum and maximum
+                    values of the parameters.
+
+        """
+
+        if prior_de is None:
+            warnings.warn('Assuming prior is uniform!')
+
+            n = (np.sum(self.sample_weights)**2) / \
+                (np.sum(self.sample_weights**2))
+
+            theta_max = np.max(self.theta, axis=0)
+            theta_min = np.min(self.theta, axis=0)
+            a = ((n-2)*theta_max-theta_min)/(n-3)
+            b = ((n-2)*theta_min-theta_max)/(n-3)
+            prior_logprob = np.log(
+                np.prod([1/(a - b)
+                         for i in range(len(b))]))
+        else:
+            prior_logprob = prior_de.log_prob(params).numpy()
+
+        posterior_logprob = self.log_prob(params)
+
+        loglike = posterior_logprob + prior_logprob - logevidence
+
+        return loglike
+
+    @tf.function(jit_compile=True)
+    def __call__(self, u):
+
+        r"""
+
+        This function is used when calling the clusterMAF class to transform
+        samples from the unit hypercube to samples on the clusterMAF.
+
+        **Parameters:**
+
+            u: **numpy array**
+                | Samples on the uniform hypercube.
+
+        """
+
+        len_thetas = [len(self.split_theta[i])
+                      for i in range(len(self.split_theta))]
+        probabilities = [len_thetas[i]/np.sum(len_thetas)
+                         for i in range(len(self.split_theta))]
+        options = np.arange(0, self.cluster_number)
+        choice = np.random.choice(options,
+                                  p=probabilities, size=len(u))
+
+        totals = [len(choice[choice == options[i]])
+                  for i in range(len(options))]
+        totals = np.hstack([0, np.cumsum(totals)])
+
+        values = []
+        for i in range(len(options)):
+            x = self.flow[i](u[totals[i]:totals[i+1]])
+            values.append(x)
+
+        x = tf.concat(values, axis=0)
+        return x
+
+    def sample(self, length=1000):
+
+        r"""
+
+        This function is used to generate samples on the clusterMAF via the
+        clusterMAF __call__ function.
+
+        **Kwargs:**
+
+            length: **int / default=1000**
+                | This should be an integer and is used to determine how many
+                    samples are generated when calling the clusterMAF.
+
+        """
+
+        u = np.random.uniform(size=(length, self.theta.shape[-1]))
+        return self(u)
+
+    def save(self, filename):
+
+        r"""
+
+        This function can be used to save an instance of
+        a trained clusterMAF as
+        a pickled class so that it can be loaded and used in differnt scripts.
+
+        **Parameters:**
+
+            filename: **string**
+                | Path in which to save the pickled MAF.
+
+        """
+
+        nn_weights = {}
+        for i in range(self.cluster_number):
+            made = self.flow[i].mades
+            w = [m.get_weights() for m in made]
+            nn_weights[i] = w
+
+        with open(filename, 'wb') as f:
+            pickle.dump([self.theta,
+                        nn_weights,
+                        self.sample_weights,
+                        self.number_networks,
+                        self.hidden_layers,
+                        self.learning_rate,
+                        self.cluster_number,
+                        self.cluster_labels], f)
+
+    @classmethod
+    def load(cls, filename):
+
+        r"""
+
+        This function can be used to load a saved MAF. For example
+
+        .. code:: python
+
+            from margarine.clustered import clusterMAF
+
+            file = 'path/to/pickled/MAF.pkl'
+            bij = clusterMAF.load(file)
+
+        **Parameters:**
+
+            filename: **string**
+                | Path to the saved MAF.
+
+        """
+
+        with open(filename, 'rb') as f:
+            data = pickle.load(f)
+            theta, nn_weights, \
+                sample_weights, \
+                number_networks, \
+                hidden_layers, \
+                learning_rate, \
+                cluster_number, \
+                labels = data
+
+        maf = cls(theta, weights=sample_weights,
+                  number_networks=number_networks,
+                  hidden_layers=hidden_layers,
+                  learning_rate=learning_rate,
+                  cluster_number=cluster_number,
+                  cluster_labels=labels)
+        maf(np.random.uniform(size=(1, theta.shape[-1])))
+        for j in range(len(maf.flow)):
+            for made, nnw in zip(maf.flow[j].mades, nn_weights[j]):
+                made.set_weights(nnw)
+
+        return maf
diff --git a/margarine/kde.py b/margarine/kde.py
index 33dbd9e..e2a2465 100755
--- a/margarine/kde.py
+++ b/margarine/kde.py
@@ -2,10 +2,12 @@
 from scipy.stats import gaussian_kde, norm
 from margarine.processing import _forward_transform, _inverse_transform
 from scipy.optimize import root_scalar
+import tensorflow as tf
 import pickle
 import warnings
 from tensorflow_probability import bijectors as tfb
 import margarine
+import anesthetic
 
 
 class KDE(object):
@@ -17,15 +19,18 @@ class KDE(object):
 
     **Parameters:**
 
-        theta: **numpy array**
+        theta: **numpy array or anesthetic.samples**
             | The samples from the probability distribution that we require the
-                bijector to learn.
-
-        weights: **numpy array**
-            | The weights associated with the samples above.
+                MAF to learn. This can either be a numpy array or an anesthetic
+                NestedSamples or MCMCSamples object.
 
     **kwargs:**
 
+        weights: **numpy array / default=np.ones(len(theta))**
+            | The weights associated with the samples above. If an anesthetic
+                NestedSamples or MCMCSamples object is passed the code
+                draws the weights from this.
+
         bw_method: **str, scalar or callable**
             | The bandwidth for the KDE.
 
@@ -36,6 +41,10 @@ class KDE(object):
         theta_min: **numpy array**
             | As above but the true lower limits of the priors.
 
+        parameters: **list of strings**
+            | A list of the relevant parameters to train on. Only needed
+                if theta is an anestehetic samples object. If not provided,
+                all parameters will be used.
     **Attributes:**
 
     A list of some key attributes accessible to the user.
@@ -75,7 +84,25 @@ class KDE(object):
 
     """
 
-    def __init__(self, theta, weights, **kwargs):
+    def __init__(self, theta, **kwargs):
+
+        self.parameters = kwargs.pop('parameters', None)
+
+        if isinstance(theta, 
+                      (anesthetic.samples.NestedSamples, 
+                       anesthetic.samples.MCMCSamples)):
+            weights = theta.get_weights()
+            if self.parameters:
+                theta = theta[self.parameters].values
+            else:
+                if isinstance(theta, anesthetic.samples.NestedSamples):
+                    self.parameters = theta.columns[:-3].values
+                    theta = theta[theta.columns[:-3]].values
+                else:
+                    self.parameters = theta.columns[:-1].values
+                    theta = theta[theta.columns[:-1]].values
+        else:
+            weights = kwargs.pop('weights', np.ones(len(theta)))
 
         self.theta = theta
         self.sample_weights = weights
@@ -96,8 +123,11 @@ def generate_kde(self):
         Function noramlises the input data into a standard normal parameter
         space and then generates a weighted KDE.
         """
+        theta = tf.convert_to_tensor(self.theta, dtype=tf.float32)
+        theta_min = tf.convert_to_tensor(self.theta_min, dtype=tf.float32)
+        theta_max = tf.convert_to_tensor(self.theta_max, dtype=tf.float32)
 
-        phi = _forward_transform(self.theta, self.theta_min, self.theta_max)
+        phi = _forward_transform(theta, theta_min, theta_max).numpy()
         mask = np.isfinite(phi).all(axis=-1)
         phi = phi[mask, :]
         weights_phi = self.sample_weights[mask]
@@ -145,7 +175,10 @@ def __call__(self, u):
                     bracket=(mu[:, i].min()*2, mu[:, i].max()*2),
                     method='bisect').root
             transformed_samples.append(
-                _inverse_transform(y, self.theta_min, self.theta_max))
+                _inverse_transform(
+                    tf.convert_to_tensor(y, dtype=tf.float32),
+                    tf.convert_to_tensor(self.theta_min, dtype=tf.float32),
+                    tf.convert_to_tensor(self.theta_max, dtype=tf.float32)))
         transformed_samples = np.array(transformed_samples)
         return transformed_samples
 
@@ -167,7 +200,10 @@ def sample(self, length=1000):
 
         """
         x = self.kde.resample(length).T
-        return _inverse_transform(x, self.theta_min, self.theta_max)
+        return _inverse_transform(
+            tf.convert_to_tensor(x, dtype='float32'),
+            tf.convert_to_tensor(self.theta_min, dtype=tf.float32),
+            tf.convert_to_tensor(self.theta_max, dtype=tf.float32))
 
     def log_prob(self, params):
 
@@ -191,7 +227,9 @@ def log_prob(self, params):
         maxs = self.theta_max.astype(np.float32)
 
         transformed_x = _forward_transform(
-            params, mins, maxs)
+            tf.convert_to_tensor(params, dtype=tf.float32),
+            tf.convert_to_tensor(mins, dtype=tf.float32),
+            tf.convert_to_tensor(maxs, dtype=tf.float32)).numpy()
 
         transform_chain = tfb.Chain([
             tfb.Invert(tfb.NormalCDF()),
@@ -243,7 +281,7 @@ def log_like(self, params, logevidence, prior=None, prior_weights=None):
         if prior is None:
             warnings.warn('Assuming prior is uniform!')
             prior_logprob = np.log(np.prod(
-                                    [1/(self.theta_max[i] - self.theta_min[i])
+                                   [1/(self.theta_max[i] - self.theta_min[i])
                                     for i in range(len(self.theta_min))]))
         else:
             self.prior = margarine.kde.KDE(prior, prior_weights)
@@ -280,7 +318,7 @@ def load(cls, filename):
 
         .. code:: python
 
-            from bayesstats.kde import KDE
+            from margarine.kde import KDE
 
             file = 'path/to/pickled/bijector.pkl'
             KDE_class = KDE.load(file)
@@ -295,7 +333,7 @@ def load(cls, filename):
         with open(filename, 'rb') as f:
             theta, sample_weights, kde = pickle.load(f)
 
-        kde_class = cls(theta, sample_weights)
+        kde_class = cls(theta, weights=sample_weights)
         kde_class.kde = kde
         kde_class(np.random.uniform(0, 1, size=(2, theta.shape[-1])))
 
diff --git a/margarine/maf.py b/margarine/maf.py
index a9e0ad8..ed2e279 100755
--- a/margarine/maf.py
+++ b/margarine/maf.py
@@ -1,18 +1,16 @@
-import numpy as np
 import tensorflow as tf
 from tensorflow import keras
 from tensorflow_probability import (bijectors as tfb, distributions as tfd)
-from margarine.processing import _forward_transform, _inverse_transform
-from sklearn.cluster import KMeans
-import pickle
+from margarine.processing import (_forward_transform, _inverse_transform,
+                                  pure_tf_train_test_split)
+import numpy as np
+import tqdm
 import warnings
-import margarine
-from scipy.special import logsumexp
-from sklearn.model_selection import train_test_split
-import tqdm 
+import pickle
+import anesthetic
 
 
-class MAF(object):
+class MAF():
 
     r"""
 
@@ -21,15 +19,18 @@ class MAF(object):
 
     **Parameters:**
 
-        theta: **numpy array**
+        theta: **numpy array or anesthetic.samples**
             | The samples from the probability distribution that we require the
-                MAF to learn.
-
-        weights: **numpy array**
-            | The weights associated with the samples above.
+                MAF to learn. This can either be a numpy array or an anesthetic
+                NestedSamples or MCMCSamples object.
 
     **kwargs:**
 
+        weights: **numpy array / default=np.ones(len(theta))**
+            | The weights associated with the samples above. If an anesthetic
+                NestedSamples or MCMCSamples object is passed the code
+                draws the weights from this.
+
         number_networks: **int / default = 6**
             | The bijector is built by chaining a series of
                 autoregressive neural
@@ -53,54 +54,16 @@ class MAF(object):
 
         theta_min: **numpy array**
             | As above but the true lower limits of the priors.
-
-        clustering: **Bool / default = False**
-            | Whether or not to perform clustering as in
-                https://arxiv.org/abs/2305.02930.
-
-        cluster_labels: **list / default = None**
-            | If clustering has been performed externally to margarine you can
-                provide a list of labels for the samples theta. The labels
-                should be integers from 0 to k corresponding to the cluster
-                that each sample is in. Clustering is turned on if cluster
-                labels are supplied.
-
-        cluster_number: **int / default = None**
-            | If clustering has been performed externally to margarine you
-                need to provide the number of clusters, k, alongside the
-                cluster labels.
-
-
+            
+        parameters: **list of strings**
+            | A list of the relevant parameters to train on. Only needed
+                if theta is an anestehetic samples object. If not provided,
+                all parameters will be used.
+            
     **Attributes:**
 
     A list of some key attributes accessible to the user.
 
-        maf: **Instance of tfd.TransformedDistribution**
-            | By loading a trained instance of this class and accessing the
-                ``maf`` (masked autoregressive flow) attribute, which is an
-                instance of
-                ``tensorflow_probability.distributions.TransformedDistribution``,
-                the used can sample the trained MAF. e.g.
-
-                .. code:: python
-
-                    from ...maf import MAF
-
-                    file = '/trained_maf.pkl'
-                    bij = MAF.load(file)
-
-                    samples = bij.maf.sample(1000)
-
-                It can also be used to calculate log probabilities via
-
-                .. code:: python
-
-                    log_prob = bij.log_prob(samples)
-
-                For more information on the attributes associated with
-                ``tensorflow_probability.distributions.TransformedDistribution``
-                see the tensorflow documentation.
-
         theta_max: **numpy array**
             | The true upper limits of the priors used to generate the
                 samples that we want the MAF to learn. If theta_max is not
@@ -117,38 +80,47 @@ class MAF(object):
 
     """
 
-    def __init__(self, theta, weights, **kwargs):
+    def __init__(self, theta, **kwargs):
         self.number_networks = kwargs.pop('number_networks', 6)
         self.learning_rate = kwargs.pop('learning_rate', 1e-3)
         self.hidden_layers = kwargs.pop('hidden_layers', [50, 50])
-        self.clustering = kwargs.pop('clustering', False)
-        self.cluster_labels = kwargs.pop('cluster_labels', None)
-        self.cluster_number = kwargs.pop('cluster_number', None)
-
-        if self.cluster_number is not None:
-            if self.cluster_labels is None:
-                raise ValueError("'cluster_labels' should be provided if " +
-                                 "'cluster_number' is specified.")
+        self.parameters = kwargs.pop('parameters', None)
+
+        if isinstance(theta, 
+                      (anesthetic.samples.NestedSamples, 
+                       anesthetic.samples.MCMCSamples)):
+            weights = theta.get_weights()
+            if self.parameters:
+                theta = theta[self.parameters].values
+            else:
+                if isinstance(theta, anesthetic.samples.NestedSamples):
+                    self.parameters = theta.columns[:-3].values
+                    theta = theta[theta.columns[:-3]].values
+                else:
+                    self.parameters = theta.columns[:-1].values
+                    theta = theta[theta.columns[:-1]].values
         else:
-            if self.cluster_labels is not None:
-                raise ValueError("'cluster_number' should be provided if " +
-                                 "'cluster_labels' is specified.")
+            weights = kwargs.pop('weights', np.ones(len(theta)))
 
-        if self.cluster_labels is not None and self.cluster_number is not None:
-            self.clustering = True
+        self.theta = tf.convert_to_tensor(theta, dtype=tf.float32)
+        self.sample_weights = tf.convert_to_tensor(weights, dtype=tf.float32)
 
-        self.n = (np.sum(weights)**2)/(np.sum(weights**2))
-        self.sample_weights = weights
+        mask = np.isfinite(theta).all(axis=-1)
+        self.theta = tf.boolean_mask(self.theta, mask, axis=0)
+        self.sample_weights = tf.boolean_mask(
+                                              self.sample_weights,
+                                              mask, axis=0)
 
-        theta_max = np.max(theta, axis=0)
-        theta_min = np.min(theta, axis=0)
+        self.n = tf.math.reduce_sum(self.sample_weights)**2 / \
+            tf.math.reduce_sum(self.sample_weights**2)
+
+        theta_max = tf.math.reduce_max(self.theta, axis=0)
+        theta_min = tf.math.reduce_min(self.theta, axis=0)
         a = ((self.n-2)*theta_max-theta_min)/(self.n-3)
         b = ((self.n-2)*theta_min-theta_max)/(self.n-3)
         self.theta_min = kwargs.pop('theta_min', b)
         self.theta_max = kwargs.pop('theta_max', a)
 
-        self.theta = theta
-
         if type(self.number_networks) is not int:
             raise TypeError("'number_networks' must be an integer.")
         if not isinstance(self.learning_rate,
@@ -168,18 +140,7 @@ def __init__(self, theta, weights, **kwargs):
         self.optimizer = tf.keras.optimizers.legacy.Adam(
                 learning_rate=self.learning_rate)
 
-        if self.clustering is False:
-            self.gen_mades()
-            self.cluster_number = None
-        if self.clustering is True:
-            if self.cluster_number is not None:
-                if self.cluster_number%1!=0:
-                    raise TypeError("'cluster_number' must be a whole number.")
-            if self.cluster_labels is not None:
-                if not isinstance(self.cluster_labels, (np.ndarray, list)):
-                    raise TypeError("'cluster_labels' must be an array " +
-                                    "or a list.")
-            self.clustering_call()
+        self.gen_mades()
 
     def gen_mades(self):
 
@@ -201,114 +162,18 @@ def gen_mades(self):
 
         return self.bij, self.maf
 
-    def clustering_call(self):
-
-        """Generating a piecewise masked autoregressive
-        flow with clustering."""
-
-        if self.cluster_number is None:
-            from sklearn.metrics import silhouette_score
-            ks = np.arange(2, 21)
-            losses = []
-            for k in ks:
-                kmeans = KMeans(k, random_state=0)
-                labels = kmeans.fit(self.theta).predict(self.theta)
-                losses.append(-silhouette_score(self.theta, labels))
-            losses = np.array(losses)
-            minimum_index = np.argmin(losses)
-            self.cluster_number = ks[minimum_index]
-
-            kmeans = KMeans(self.cluster_number, random_state=0)
-            self.cluster_labels = kmeans.fit(self.theta).predict(self.theta)
-        
-        if self.cluster_number == 20:
-            warnings.warn("The number of clusters is 20. This is the maximum "+
-                            "number of clusters that can be used. If you " +
-                            "require more clusters, please specify the " +
-                            "'cluster_number' kwarg. margarine will continue "+
-                            "with 20 clusters.")
-        
-        if np.array(list(self.cluster_labels)).dtype == 'float':
-            # convert cluster labels to integers
-            self.cluster_labels = self.cluster_labels.astype(int)
-        # count the number of times a cluster label appears in cluster_labels
-        self.cluster_count = np.bincount(self.cluster_labels)
-        # While loop to make sure clusters are not too small
-        while self.cluster_count.min() < 100:
-            warnings.warn("One or more clusters are too small " +
-                          "(n_cluster < 100). " +
-                          "Reducing the number of clusters by 1.")
-            minimum_index -= 1
-            self.cluster_number = ks[minimum_index]
-            kmeans = KMeans(self.cluster_number, random_state=0)
-            self.cluster_labels = kmeans.fit(self.theta).predict(self.theta)
-            self.cluster_count = np.bincount(self.cluster_labels)
-            if self.cluster_number == 2:
-                # break if two clusters
-                warnings.warn("The number of clusters is 2. This is the " +
-                            "minimum number of clusters that can be used. " +
-                            "Some clusters may be too small and the " +
-                            "train/test split may fail." +
-                            "Try running without clusting. ")
-                break
-
-        self.n, split_theta, self.new_theta_max = [], [], []
-        self.new_theta_min, split_sample_weights = [], []
-        self.bij, self.maf, self.mades = [], [], []
-        for i in range(self.cluster_number):
-            theta = self.theta[self.cluster_labels == i]
-            split_sample_weights.append(
-                self.sample_weights[self.cluster_labels == i])
-
-            self.n.append(
-                (np.sum(self.sample_weights[self.cluster_labels == i])**2) /
-                (np.sum(self.sample_weights[self.cluster_labels == i]**2)))
-
-            if not any(isinstance(tm, list) for tm in self.theta_max):
-                new_theta_max = np.max(theta, axis=0)
-                new_theta_min = np.min(theta, axis=0)
-                a = ((self.n[-1]-2)*new_theta_max-new_theta_min)/(self.n[-1]-3)
-                b = ((self.n[-1]-2)*new_theta_min-new_theta_max)/(self.n[-1]-3)
-                self.new_theta_min.append(b)
-                self.new_theta_max.append(a)
-
-            split_theta.append(theta)
-
-            self.mades.append(
-                [tfb.AutoregressiveNetwork(
-                            params=2,
-                            hidden_units=self.hidden_layers, activation='tanh',
-                            input_order='random')
-                    for _ in range(self.number_networks)])
-
-            self.bij.append(tfb.Chain([
-                tfb.MaskedAutoregressiveFlow(made)
-                for made in self.mades[-1]]))
-
-            self.base = tfd.Blockwise(
-                [tfd.Normal(loc=0, scale=1)
-                    for _ in range(self.theta.shape[-1])])
-
-            self.maf.append(tfd.TransformedDistribution(
-                self.base, bijector=self.bij[-1]))
-        self.theta = split_theta
-        self.sample_weights = split_sample_weights
-        if self.new_theta_max != []:
-            self.theta_max = self.new_theta_max
-            self.theta_min = self.new_theta_min
-
     def train(self, epochs=100, early_stop=False, loss_type='sum'):
+
         r"""
 
         This function is called to train the MAF once it has been
-        initialised. It calls `_training()` for each MAF whether there
-        is only one or whether clustering is turned on. For example
+        initialised. It calls the `_training()` function.
 
         .. code:: python
 
-            from ...maf import MAF
+            from margarine.maf import MAF
 
-            bij = MAF(theta, weights)
+            bij = MAF(theta, weights=weights)
             bij.train()
 
         **Kwargs:**
@@ -325,7 +190,13 @@ def train(self, epochs=100, early_stop=False, loss_type='sum'):
                     margarine will roll back to the best model and return this
                     to the user.
 
+            loss_type: **string / default = 'sum'**
+                | Determines whether to use the sum or mean of the weighted
+                    log probabilities to calculate the loss function.
+
+
         """
+
         if type(epochs) is not int:
             raise TypeError("'epochs' is not an integer.")
         if type(early_stop) is not bool:
@@ -335,35 +206,20 @@ def train(self, epochs=100, early_stop=False, loss_type='sum'):
         self.early_stop = early_stop
         self.loss_type = loss_type
 
-        if self.cluster_number is not None:
-            for i in range(len(self.theta)):
-                self.maf[i] = self._training(self.theta[i],
-                                             self.sample_weights[i],
-                                             self.maf[i], self.theta_min[i],
-                                             self.theta_max[i])
-        else:
-            self.maf = self._training(self.theta,
-                                      self.sample_weights, self.maf,
-                                      self.theta_min, self.theta_max)
+        self.maf = self._training(self.theta,
+                                  self.sample_weights, self.maf,
+                                  self.theta_min, self.theta_max)
 
     def _training(self, theta, sample_weights, maf,
                   theta_min, theta_max):
 
-        """Function to perform the training of each MAF."""
+        """Training the masked autoregressive flow."""
 
         phi = _forward_transform(theta, theta_min, theta_max)
-
-        mask = np.isfinite(phi).all(axis=-1)
-        phi = phi[mask, :]
-        weights_phi = sample_weights[mask]
-        weights_phi /= weights_phi.sum()
-
-        phi = phi.astype('float32')
-        self.phi = phi.copy()
-        weights_phi = weights_phi.astype('float32')
+        weights_phi = sample_weights/tf.reduce_sum(sample_weights)
 
         phi_train, phi_test, weights_phi_train, weights_phi_test = \
-            train_test_split(phi, weights_phi, test_size=0.2)
+            pure_tf_train_test_split(phi, weights_phi, test_size=0.2)
 
         self.loss_history = []
         self.test_loss_history = []
@@ -371,17 +227,12 @@ def _training(self, theta, sample_weights, maf,
         for i in tqdm.tqdm(range(self.epochs)):
             loss = self._train_step(phi_train,
                                     weights_phi_train,
-                                    self.loss_type, maf).numpy()
+                                    self.loss_type, maf)
             self.loss_history.append(loss)
 
-            if self.loss_type == 'sum':
-                loss_test = -tf.reduce_sum(
-                    weights_phi_test*maf.log_prob(phi_test))
-            elif self.loss_type == 'mean':
-                loss_test = -tf.reduce_mean(
-                    weights_phi_test*maf.log_prob(phi_test))
-
-            self.test_loss_history.append(loss_test)
+            self.test_loss_history.append(self._test_step(phi_test,
+                                          weights_phi_test,
+                                          self.loss_type, maf))
 
             if self.early_stop:
                 c += 1
@@ -401,8 +252,22 @@ def _training(self, theta, sample_weights, maf,
                               '. Minimum at epoch = ' + str(minimum_epoch))
                         return minimum_model
         return maf
-    
-    #@tf.function(jit_compile=True)
+
+    @tf.function(jit_compile=True)
+    def _test_step(self, x, w, loss_type, maf):
+
+        r"""
+        This function is used to calculate the test loss value at each epoch
+        for early stopping.
+        """
+
+        if loss_type == 'sum':
+            loss = -tf.reduce_sum(w*maf.log_prob(x))
+        elif loss_type == 'mean':
+            loss = -tf.reduce_mean(w*maf.log_prob(x))
+        return loss
+
+    @tf.function(jit_compile=True)
     def _train_step(self, x, w, loss_type, maf):
 
         r"""
@@ -410,7 +275,7 @@ def _train_step(self, x, w, loss_type, maf):
         adjust the weights and biases of the neural networks via the
         optimizer algorithm.
         """
-        
+
         with tf.GradientTape() as tape:
             if loss_type == 'sum':
                 loss = -tf.reduce_sum(w*maf.log_prob(x))
@@ -422,8 +287,8 @@ def _train_step(self, x, w, loss_type, maf):
                 maf.trainable_variables))
         return loss
 
+    @tf.function(jit_compile=True)
     def __call__(self, u):
-
         r"""
 
         This function is used when calling the MAF class to transform
@@ -435,35 +300,15 @@ def __call__(self, u):
                 | Samples on the uniform hypercube.
 
         """
+        u = tf.cast(u, dtype=tf.float32)
 
-        if self.cluster_number is not None:
-            len_thetas = [len(self.theta[i]) for i in range(len(self.theta))]
-            probabilities = [len_thetas[i]/np.sum(len_thetas)
-                             for i in range(len(self.theta))]
-            options = np.arange(0, self.cluster_number)
-            choice = np.random.choice(options,
-                                      p=probabilities, size=len(u))
-
-            totals = [len(choice[choice == options[i]])
-                      for i in range(len(options))]
-            totals = np.hstack([0, np.cumsum(totals)])
-
-            values = []
-            for i in range(len(options)):
-                x = _forward_transform(u[totals[i]:totals[i+1]])
-                x = self.bij[i](x.astype(np.float32)).numpy()
-                values.append(_inverse_transform(x, self.theta_min[i],
-                                                 self.theta_max[i]))
-
-            x = np.concatenate(values)
-        else:
-            x = _forward_transform(u)
-            x = self.bij(x.astype(np.float32)).numpy()
-            x = _inverse_transform(x, self.theta_min, self.theta_max)
+        x = _forward_transform(u)
+        x = self.bij(x)
+        x = _inverse_transform(x, self.theta_min, self.theta_max)
 
-        mask = np.isfinite(x).all(axis=-1)
-        return x[mask, ...]
+        return x
 
+    @tf.function(jit_compile=True)
     def sample(self, length=1000):
 
         r"""
@@ -481,12 +326,10 @@ def sample(self, length=1000):
         if type(length) is not int:
             raise TypeError("'length' must be an integer.")
 
-        if self.cluster_number is not None:
-            u = np.random.uniform(0, 1, size=(length, self.theta[0].shape[-1]))
-        else:
-            u = np.random.uniform(0, 1, size=(length, self.theta.shape[-1]))
+        u = tf.random.uniform((length, len(self.theta_min)))
         return self(u)
 
+    @tf.function(jit_compile=True)
     def log_prob(self, params):
 
         """
@@ -512,7 +355,7 @@ def calc_log_prob(mins, maxs, maf):
 
             def norm_jac(y):
                 return transform_chain.inverse_log_det_jacobian(
-                    y, event_ndims=0).numpy()
+                    y, event_ndims=0)
 
             transformed_x = _forward_transform(params, mins, maxs)
 
@@ -521,30 +364,15 @@ def norm_jac(y):
                 tfb.Scale(1/(maxs - mins)), tfb.Shift(-mins)])
 
             correction = norm_jac(transformed_x)
-            logprob = (maf.log_prob(transformed_x).numpy() -
-                       np.sum(correction, axis=-1)).astype(np.float64)
+            logprob = (maf.log_prob(transformed_x) -
+                       tf.reduce_sum(correction, axis=-1))
             return logprob
 
-        if self.clustering is True:
-            logprob = []
-            for i in range(len(self.theta_min)):
-                mins = self.theta_min[i].astype(np.float32)
-                maxs = self.theta_max[i].astype(np.float32)
-                probs = calc_log_prob(mins, maxs, self.maf[i])
-                for j in range(len(probs)):
-                    if np.isnan(probs[j]):
-                        probs[j] = np.log(1e-300)
-                logprob.append(probs)
-            logprob = np.array(logprob)
-            logprob = logsumexp(logprob, axis=0)
-        else:
-            mins = self.theta_min.astype(np.float32)
-            maxs = self.theta_max.astype(np.float32)
-            logprob = calc_log_prob(mins, maxs, self.maf)
+        logprob = calc_log_prob(self.theta_min, self.theta_max, self.maf)
 
         return logprob
 
-    def log_like(self, params, logevidence, prior=None, prior_weights=None):
+    def log_like(self, params, logevidence, prior_de=None):
 
         r"""
         This function should return the log-likelihood for a given set of
@@ -552,7 +380,7 @@ def log_like(self, params, logevidence, prior=None, prior_weights=None):
 
         It requires the logevidence from the original nested sampling run
         in order to do this and in the case that the prior is non-uniform
-        prior samples should be provided.
+        a trained prior density estimator should be provided.
 
         **Parameters:**
 
@@ -564,28 +392,25 @@ def log_like(self, params, logevidence, prior=None, prior_weights=None):
                 | Should be the log-evidence from the full nested sampling
                     run with nuisance parameters.
 
-            prior: **numpy array/default=None**
-                | An array of prior samples corresponding to the prior. Default
-                    assumption is that the prior is uniform which is
-                    required if you want to combine likelihoods from different
-                    experiments/data sets. In this case samples and prior
-                    samples should be reweighted prior to any training.
-
-            prior_weights: **numpy array/default=None**
-                | Weights to go with the prior samples.
+            prior_de: **margarine.maf.MAF / default=None**
+                | If the prior is non-uniform then a trained prior density
+                    estimator should be provided. Otherwise the prior
+                    is assumed to be uniform and the prior probability
+                    is calculated analytically from the minimum and maximum
+                    values of the parameters.
 
         """
 
-        if prior is None:
+        if prior_de is None:
             warnings.warn('Assuming prior is uniform!')
-            prior_logprob = np.log(
-                np.prod([1/(self.theta_max[i] - self.theta_min[i])
-                         for i in range(len(self.theta_min))]))
+            prior_logprob = tf.math.log(
+                                        tf.math.reduce_prod(
+                                            [1/(self.theta_max[i] -
+                                                self.theta_min[i])
+                                                for i in range(
+                                                    len(self.theta_min))]))
         else:
-            self.prior = margarine.maf.MAF(prior, prior_weights)
-            self.prior.train()
-            prior_logprob_func = self.prior.log_prob
-            prior_logprob = prior_logprob_func(params)
+            prior_logprob = self.prior_de.log_prob(params)
 
         posterior_logprob = self.log_prob(params)
 
@@ -596,8 +421,7 @@ def log_like(self, params, logevidence, prior=None, prior_weights=None):
     def save(self, filename):
         r"""
 
-        This function can be used to save an instance of a trained MAF
-        or piecewise MAF as
+        This function can be used to save an instance of a trained MAF as
         a pickled class so that it can be loaded and used in differnt scripts.
 
         **Parameters:**
@@ -606,35 +430,17 @@ def save(self, filename):
                 | Path in which to save the pickled MAF.
 
         """
-        if self.cluster_number is None:
-            nn_weights = [made.get_weights() for made in self.mades]
-        else:
-            nn_weights = {}
-            for i in range(self.cluster_number):
-                made = self.mades[i]
-                w = [m.get_weights() for m in made]
-                nn_weights[i] = w
+
+        nn_weights = [made.get_weights() for made in self.mades]
         with open(filename, 'wb') as f:
-            if self.clustering is True:
-                pickle.dump([self.theta,
-                            nn_weights,
-                            self.sample_weights,
-                            self.number_networks,
-                            self.hidden_layers,
-                            self.learning_rate,
-                            self.theta_min,
-                            self.theta_max,
-                            self.cluster_number,
-                            self.cluster_labels], f)
-            else:
-                pickle.dump([self.theta,
-                            nn_weights,
-                            self.sample_weights,
-                            self.number_networks,
-                            self.hidden_layers,
-                            self.learning_rate,
-                            self.theta_min,
-                            self.theta_max], f)
+            pickle.dump([self.theta,
+                        nn_weights,
+                        self.sample_weights,
+                        self.number_networks,
+                        self.hidden_layers,
+                        self.learning_rate,
+                        self.theta_min,
+                        self.theta_max], f)
 
     @classmethod
     def load(cls, filename):
@@ -644,7 +450,7 @@ def load(cls, filename):
 
         .. code:: python
 
-            from ...maf import MAF
+            from margarine.maf import MAF
 
             file = 'path/to/pickled/MAF.pkl'
             bij = MAF.load(file)
@@ -658,49 +464,19 @@ def load(cls, filename):
 
         with open(filename, 'rb') as f:
             data = pickle.load(f)
-            try:
-                theta, nn_weights, \
-                    sample_weights, \
-                    number_networks, \
-                    hidden_layers, \
-                    learning_rate, theta_min, theta_max = data
-                cluster_number = None
-            except Exception:
-                theta, nn_weights, \
-                    sample_weights, \
-                    number_networks, \
-                    hidden_layers, \
-                    learning_rate, theta_min, theta_max, \
-                    cluster_number, \
-                    labels = data
-
-        if cluster_number is None:
-            bijector = cls(
-                theta, sample_weights, number_networks=number_networks,
-                learning_rate=learning_rate, hidden_layers=hidden_layers,
-                theta_min=theta_min, theta_max=theta_max)
-            bijector(
-                np.random.uniform(0, 1, size=(len(theta), theta.shape[-1])))
-            for made, nn_weights in zip(bijector.mades, nn_weights):
-                made.set_weights(nn_weights)
-        else:
-            # have to sort labels because theta is reordered when
-            # it gets concatenated
-            labels.sort()
-            theta = np.concatenate(theta)
-            sample_weights = np.concatenate(sample_weights)
-            bijector = cls(
-                theta, sample_weights, number_networks=number_networks,
-                learning_rate=learning_rate, hidden_layers=hidden_layers,
-                clustering=True,
-                cluster_number=cluster_number, cluster_labels=labels,
-                theta_min=theta_min, theta_max=theta_max)
-            bijector(
-                np.random.uniform(0, 1, size=(len(theta), theta.shape[-1])))
-            for i in range(cluster_number):
-                mades = bijector.mades[i]
-                weights = nn_weights[i]
-                for m, w in zip(mades, weights):
-                    m.set_weights(w)
+            theta, nn_weights, \
+                sample_weights, \
+                number_networks, \
+                hidden_layers, \
+                learning_rate, theta_min, theta_max = data
+
+        bijector = cls(
+            theta, weights=sample_weights, number_networks=number_networks,
+            learning_rate=learning_rate, hidden_layers=hidden_layers,
+            theta_min=theta_min, theta_max=theta_max)
+        bijector(
+            np.random.uniform(0, 1, size=(len(theta), theta.shape[-1])))
+        for made, nn_weights in zip(bijector.mades, nn_weights):
+            made.set_weights(nn_weights)
 
         return bijector
diff --git a/margarine/marginal_stats.py b/margarine/marginal_stats.py
index 7597e51..57d87d1 100755
--- a/margarine/marginal_stats.py
+++ b/margarine/marginal_stats.py
@@ -1,10 +1,7 @@
 import numpy as np
 from tensorflow_probability import distributions as tfd
-import margarine
-import warnings
 import pandas as pd
 from margarine.maf import MAF
-from margarine.kde import KDE
 
 
 class calculate(object):
@@ -18,13 +15,14 @@ class calculate(object):
     **Paramesters:**
 
         de: **instance of MAF class or KDE class**
-            | This should be a loaded and trained instance of a MAF or KDE.
-                Bijectors can be loaded like so
+            | This should be a loaded and trained instance of a MAF, clusterMAF
+                or KDE.Bijectors can be loaded like so
 
                 .. code:: python
 
                     from margarine.maf import MAF
                     from margarine.kde import KDE
+                    from margarine.clustered import clusterMAF
 
                     file = '/trained_maf.pkl'
                     maf = MAF.load(file)
@@ -32,6 +30,9 @@ class calculate(object):
                     file = '/trained_kde.pkl'
                     kde = KDE.load(file)
 
+                    file = '/trained_clustered_maf.pkl'
+                    clustered_maf = clusterMAF.load(file)
+
         samples: **numpy array**
             | This should be the output of the bijector when called to generate
                 a set of samples from the replicated probability
@@ -45,14 +46,10 @@ class calculate(object):
 
     **Kwargs:**
 
-        prior_samples: **numpy array / default=None**
-            | Can be provided if the prior is non-uniform and will be
-                used to generate a prior density estimator to calcualte
-                prior log-probabilities. If not provided the prior is
-                assumed to be uniform.
-
-        prior_weights: **numpy array / default=None**
-            | Weights associated with the above prior samples.
+        prior_de: **instance of MAF class, clusterMAF class or KDE class**
+            | This should be a loaded and trained instance of a MAF, clusterMAF
+                or KDE for the prior.
+                If not provided, a uniform prior will be used.
 
     """
 
@@ -60,46 +57,36 @@ def __init__(self, de, **kwargs):
 
         self.de = de
 
-        try:
-            self.theta = np.concatenate(self.de.theta)
-            self.theta_weights = np.concatenate(self.de.sample_weights)
-            self.theta_min = np.min(self.de.theta_min, axis=1)
-            self.theta_max = np.max(self.de.theta_max, axis=1)
-        except ValueError:
-            self.theta = self.de.theta
-            self.theta_weights = self.de.sample_weights
-            self.theta_min = self.de.theta_min
-            self.theta_max = self.de.theta_max
-    
+        self.theta = self.de.theta
+        self.theta_weights = self.de.sample_weights
+        self.theta_min = self.de.theta_min
+        self.theta_max = self.de.theta_max
+
+        if isinstance(self.de, MAF):
+            self.theta_weights = self.theta_weights.numpy()
+            self.theta_min = self.theta_min.numpy()
+            self.theta_max = self.theta_max.numpy()
+
         self.samples = self.de.sample(len(self.theta))
 
         self.prior_de = kwargs.pop('prior_de', None)
-        self.prior_samples = kwargs.pop('prior_samples', None)
-        self.prior_weights = kwargs.pop('prior_weights', None)
-
-        if self.prior_samples is None:
-            warnings.warn('If prior samples are not provided the prior is ' +
-                          'assumed to be uniform and the' +
-                          'posterior samples are ' +
-                          'are assumed to be from the same uniform space.')
-        else:
-            if self.prior_weights is None:
-                warnings.warn('No prior weights have been provided. ' +
-                              'Assuming there are none.')
 
     def statistics(self):
 
+        """
+        Calculate marginal bayesian KL divergence and dimensionality with
+        approximate errors.
+        """
+
         def mask_arr(arr):
             return arr[np.isfinite(arr)], np.isfinite(arr)
 
-        if isinstance(self.de, margarine.kde.KDE):
-            logprob_func = self.de.log_prob
-            logprob = logprob_func(self.samples)
-            theta_logprob = logprob_func(self.theta)
-        elif isinstance(self.de, margarine.maf.MAF):
-            logprob_func = self.de.log_prob
-            logprob = logprob_func(self.samples)
-            theta_logprob = logprob_func(self.theta)
+        logprob = self.de.log_prob(self.samples)
+        theta_logprob = self.de.log_prob(self.theta)
+
+        if isinstance(self.de, MAF):
+            logprob = logprob.numpy()
+            theta_logprob = theta_logprob.numpy()
 
         args = np.argsort(theta_logprob)
         self.theta_weights = self.theta_weights[args]
@@ -113,15 +100,15 @@ def mask_arr(arr):
 
         mid_point = np.log((np.exp(logprob) + np.exp(theta_logprob))/2)
 
-        if self.prior_samples is None:
+        if self.prior_de is None:
             self.base = tfd.Blockwise(
                 [tfd.Uniform(self.theta_min[i], self.theta_max[i])
                  for i in range(self.samples.shape[-1])])
             prior = self.base.sample(len(self.theta)).numpy()
 
             theta_base_logprob = self.base.log_prob(self.theta).numpy()
-
             base_logprob = self.base.log_prob(prior).numpy()
+
             prior_wde = [np.sum(self.theta_weights)/len(base_logprob)] * \
                 len(base_logprob)
 
@@ -132,35 +119,14 @@ def mask_arr(arr):
 
             theta_base_logprob = theta_base_logprob[args]
         elif self.prior_de is not None:
-            if isinstance(self.prior_de, margarine.kde.KDE):
-                self.base = self.prior_de.copy()
-                base_logprob_func = self.base.log_prob
-                de_prior_samples = self.base.sample(len(self.theta))
-                theta_base_logprob = base_logprob_func(self.prior_samples)
-                base_logprob = base_logprob_func(self.theta)
-            elif isinstance(self.prior_de, margarine.maf.MAF):
-                self.base = self.prior_de.copy()
-                base_logprob_func = self.base.log_prob
-                de_prior_samples = self.base.sample(len(self.theta))
-                theta_base_logprob = base_logprob_func(self.prior_samples)
-                base_logprob = base_logprob_func(self.theta)
-        else:
-            if isinstance(self.de, margarine.kde.KDE):
-                self.base = KDE(
-                    self.prior_samples, self.prior_weights)
-                self.base.generate_kde()
-                base_logprob_func = self.base.log_prob
-                de_prior_samples = self.base.sample(len(self.theta))
-                theta_base_logprob = base_logprob_func(self.prior_samples)
-                base_logprob = base_logprob_func(de_prior_samples)
-            elif isinstance(self.de, margarine.maf.MAF):
-                self.base = MAF(
-                    self.prior_samples, self.prior_weights)
-                self.base.train(epochs=250)
-                base_logprob_func = self.base.log_prob
-                de_prior_samples = self.base.sample(len(self.theta))
-                theta_base_logprob = base_logprob_func(self.prior_samples)
-                base_logprob = base_logprob_func(de_prior_samples)
+            self.base = self.prior_de.copy()
+            de_prior_samples = self.base.sample(len(self.theta))
+            theta_base_logprob = self.base.log_prob(de_prior_samples)
+            base_logprob = self.base.log_prob(self.theta)
+
+            if isinstance(self.prior_de, MAF):
+                theta_base_logprob = theta_base_logprob.numpy()
+                base_logprob = base_logprob.numpy()
 
             base_logprob = base_logprob[deargs]
             base_logprob = np.interp(
@@ -168,7 +134,7 @@ def mask_arr(arr):
 
             base_args = np.argsort(theta_base_logprob)
             theta_base_logprob = theta_base_logprob[base_args]
-            prior_weights = np.cumsum(self.prior_weights[base_args])
+            prior_weights = np.cumsum(self.prior_de.sample_weights[base_args])
 
             prior_weights = (np.cumsum(self.theta_weights).max() -
                              np.cumsum(self.theta_weights).min()) * \
diff --git a/margarine/processing.py b/margarine/processing.py
index f69c485..627456b 100755
--- a/margarine/processing.py
+++ b/margarine/processing.py
@@ -1,6 +1,9 @@
 from tensorflow_probability import distributions as tfd
+import tensorflow as tf
+import random
 
 
+@tf.function(jit_compile=True)
 def _forward_transform(x, min=0, max=1):
     r"""
 
@@ -24,11 +27,12 @@ def _forward_transform(x, min=0, max=1):
                 values...)
 
     """
-    x = tfd.Uniform(min, max).cdf(x).numpy()
-    x = tfd.Normal(0, 1).quantile(x).numpy()
+    x = tfd.Uniform(min, max).cdf(x)
+    x = tfd.Normal(0, 1).quantile(x)
     return x
 
 
+@tf.function(jit_compile=True)
 def _inverse_transform(x, min, max):
     r"""
 
@@ -51,6 +55,39 @@ def _inverse_transform(x, min, max):
                 values...)
 
     """
-    x = tfd.Normal(0, 1).cdf(x).numpy()
-    x = tfd.Uniform(min, max).quantile(x).numpy()
+    x = tfd.Normal(0, 1).cdf(x)
+    x = tfd.Uniform(min, max).quantile(x)
     return x
+
+
+def pure_tf_train_test_split(a, b, test_size=0.2):
+
+    """
+    Splitting data into training and testing sets. Function is equivalent
+    to sklearn.model_selection.train_test_split but a and b
+    are tensorflow tensors.
+
+    **parameters:**
+
+        a: **array**
+            | Samples to be split.
+
+        b: **array**
+            | Weights to be split.
+
+        test_size: **float**
+            | Fraction of data to be used for testing.
+    """
+
+    idx = random.sample(range(len(a)), int(len(a)*test_size))
+
+    a_train = tf.gather(a,
+                        tf.convert_to_tensor(
+                            list(set(range(len(a))) - set(idx))))
+    b_train = tf.gather(b,
+                        tf.convert_to_tensor(
+                            list(set(range(len(b))) - set(idx))))
+    a_test = tf.gather(a, tf.convert_to_tensor(idx))
+    b_test = tf.gather(b, tf.convert_to_tensor(idx))
+
+    return a_train, a_test, b_train, b_test
diff --git a/notebook/Tutorial.ipynb b/notebook/Tutorial.ipynb
index 3022a5d..5bf7ee4 100644
--- a/notebook/Tutorial.ipynb
+++ b/notebook/Tutorial.ipynb
@@ -81,7 +81,7 @@
    "source": [
     "## Masked Autoregressive Flows\n",
     "\n",
-    "Firstly we will look at training a Masked Autoregressive Flow or MAF with `margarine`. To train the MAF we first need to initalise the class with the samples and corresponding weights."
+    "Firstly we will look at training a Masked Autoregressive Flow or MAF with `margarine`. To train the MAF we first need to initalise the class with the samples and corresponding weights. For all of the different density estimators in `margarine` the weights have to be provided by the kwarg `weights` else they are assumed to be equal unless `theta` is an instance of anesthetic.samples."
    ]
   },
   {
@@ -96,7 +96,7 @@
     "\n",
     "from margarine.maf import MAF\n",
     "\n",
-    "bij = MAF(data, weights)\n",
+    "bij = MAF(data, weights=weights)\n",
     "bij.train(10000, early_stop=True)"
    ]
   },
@@ -187,7 +187,7 @@
    "outputs": [],
    "source": [
     "from margarine.kde import KDE\n",
-    "kde = KDE(data, weights)\n",
+    "kde = KDE(data, weights=weights)\n",
     "kde.generate_kde()\n",
     "x = kde.sample(5000)\n",
     "\n",
@@ -228,7 +228,7 @@
    "source": [
     "# Clustering with margarine\n",
     "\n",
-    "The below example demonstrates how to perform clustering with `margarine` which can improve the accuracy of the emulation."
+    "The below example demonstrates how to perform clustering with `margarine` which can improve the accuracy of the emulation. The clustering class has all the same functionality as the MAF class."
    ]
   },
   {
@@ -254,7 +254,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "flow = MAF(data, np.ones(len(data)), clustering=True)\n",
+    "from margarine.clustered import clusterMAF\n",
+    "flow = clusterMAF(data)\n",
     "flow.train(10000, early_stop=True)"
    ]
   },
diff --git a/requirements.txt b/requirements.txt
index cebb997..188edce 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,5 +1,6 @@
 numpy
-tensorflow
+tensorflow-macos; sys_platform == 'darwin'
+tensorflow; sys_platform != 'darwin'
 tensorflow_probability
 scipy
 pandas
diff --git a/setup.py b/setup.py
index e8c77ac..712a480 100755
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,7 @@ def readme(short=False):
 
 setup(
     name='margarine',
-    version='0.6.0',
+    version='1.0.0',
     description='margarine: Posterior Sampling and Marginal Bayesian Statistics',
     long_description=readme(),
     author='Harry T. J. Bevins',


Test patch:
diff --git a/tests/test_clusters.py b/tests/test_clusters.py
index ed96842..8bea60e 100644
--- a/tests/test_clusters.py
+++ b/tests/test_clusters.py
@@ -1,6 +1,6 @@
 import numpy as np
 from anesthetic import MCMCSamples
-from margarine.maf import MAF
+from margarine.clustered import clusterMAF
 from margarine.marginal_stats import calculate
 import pytest
 from numpy.testing import assert_equal, assert_allclose
@@ -30,16 +30,16 @@ def d_g(logL, weights):
 
 def test_maf_clustering():
 
-    bij = MAF(theta, weights, clustering=True)
+    bij = clusterMAF(theta, weights=weights)
     bij.train(10000, early_stop=True)
     file = 'saved_maf_cluster.pkl'
     bij.save(file)
 
-    loaded_bijector = MAF.load('saved_maf_cluster.pkl')
-    for i in range(len(bij.mades)):
-        for j in range(len(bij.mades[i])):
-            assert_equal(bij.mades[i][j].get_weights(),
-                loaded_bijector.mades[i][j].get_weights())
+    loaded_bijector = clusterMAF.load('saved_maf_cluster.pkl')
+    for f in range(len(bij.flow)):
+        for i in range(len(bij.flow[f].mades)):
+            assert_equal(bij.flow[f].mades[i].get_weights(),
+                loaded_bijector.flow[f].mades[i].get_weights())
 
     def check_stats(i):
         if i ==0:
@@ -63,15 +63,13 @@ def check_stats(i):
 def test_maf_cluster_kwargs():
 
     with pytest.raises(ValueError):
-        bij = MAF(theta, weights, cluster_number=3)
+        bij = clusterMAF(theta, weights=weights, cluster_number=3)
     with pytest.raises(ValueError):
-        bij = MAF(theta, weights, cluster_labels=np.ones(len(theta)))
+        bij = clusterMAF(theta, weightes=weights, 
+                         cluster_labels=np.ones(len(theta)))
     
     labels = np.ones(len(theta))
     labels[:len(theta)//2] = 0
-    bij = MAF(theta, weights, cluster_number=2, 
-              cluster_labels=labels)
-    assert_equal(bij.clustering, True)
 
 def test_cluster_size():
     # testing the while loop for cluster size
@@ -80,5 +78,4 @@ def draw(ndims):
                 np.zeros(ndims), np.eye(ndims), ndims*10)
 
     samples = draw(3)
-    flow = MAF(samples, np.ones(len(samples)), clustering=True)
-    return 0
\ No newline at end of file
+    flow = clusterMAF(samples)
\ No newline at end of file
diff --git a/tests/test_stats.py b/tests/test_stats.py
index 38f4e4f..27f1181 100755
--- a/tests/test_stats.py
+++ b/tests/test_stats.py
@@ -4,6 +4,7 @@
 from margarine.marginal_stats import calculate
 import pytest
 from margarine.kde import KDE
+from margarine.clustered import clusterMAF
 from numpy.testing import assert_equal, assert_allclose
 from scipy.stats import ks_2samp
 
@@ -29,7 +30,7 @@ def d_g(logL, weights):
 
 logL = likelihood(theta)
 
-samples = MCMCSamples(data=theta, logL=logL)
+mcmc_samples = MCMCSamples(data=theta, logL=logL)
 samples_kl = D_KL(logL, weights)
 samples_d = d_g(logL, weights)
 names = [i for i in range(theta.shape[-1])]
@@ -43,13 +44,13 @@ def check_stats(i):
             value = samples_d
         assert_allclose(stats['Value'][i], value, rtol=1, atol=1)
 
-    bij = MAF(theta, weights)
+    bij = MAF(theta, weights=weights)
     bij.train(10000, early_stop=True)
 
     stats = calculate(bij).statistics()
     [check_stats(i) for i in range(2)]
 
-    equal_weight_theta = samples.compress(50)[names].values
+    equal_weight_theta = mcmc_samples.compress(50)[names].values
     x = bij.sample(len(equal_weight_theta))
 
     res = [ks_2samp(equal_weight_theta[:, i], x[:, i]) 
@@ -58,7 +59,7 @@ def check_stats(i):
     for i in range(len(p_values)):
         assert(round(p_values[i], 2) >0.05)
 
-    estL = bij.log_like(equal_weight_theta, 0.0)
+    estL = bij.log_like(equal_weight_theta.astype(np.float32), 0.0)
     check_like = 0
     for i in range(len(logL)):
         if np.isclose(np.exp(logL[i]), np.exp(estL[i]), rtol=1, atol=1):
@@ -68,35 +69,35 @@ def check_stats(i):
 def test_maf_kwargs():
 
     with pytest.raises(TypeError):
-        bij = MAF(theta, weights)
+        bij = MAF(theta, weights=weights)
         bij.sample(4.5)
     with pytest.raises(TypeError):
-        MAF(theta, weights, number_networks=4.4)
+        MAF(theta, weights=weights, number_networks=4.4)
     with pytest.raises(TypeError):
-        MAF(theta, weights, learning_rate='foobar')
+        MAF(theta, weights=weights, learning_rate='foobar')
     with pytest.raises(TypeError):
-        MAF(theta, weights, hidden_layers='foobar')
+        MAF(theta, weights=weights, hidden_layers='foobar')
     with pytest.raises(TypeError):
-        MAF(theta, weights, hidden_layers=[4.5, 50])
+        MAF(theta, weights=weights, hidden_layers=[4.5, 50])
     with pytest.raises(TypeError):
-        MAF(theta, weights)
+        MAF(theta, weights=weights)
         bij.train(epochs=4.5)
     with pytest.raises(TypeError):
-        MAF(theta, weights)
+        MAF(theta, weights=weights)
         bij.train(epochs=100, early_stop='foo')
     with pytest.raises(TypeError):
-        MAF(theta, weights)
+        MAF(theta, weights=weights)
         bij.train(epochs=100, clustering=5)
     with pytest.raises(TypeError):
-        MAF(theta, weights)
+        MAF(theta, weights=weights)
         bij.train(epochs=100, cluster_numeber='foo')
     with pytest.raises(TypeError):
-        MAF(theta, weights)
+        MAF(theta, weights=weights)
         bij.train(epochs=100, cluster_labels=5)
 
 def test_maf_save_load():
 
-    bij = MAF(theta, weights)
+    bij = MAF(theta, weights=weights)
     bij.train(100)
     file = 'saved_maf.pkl'
     bij.save(file)
@@ -114,13 +115,13 @@ def check_stats(i):
             value = samples_d
         assert_allclose(stats['Value'][i], value, rtol=1, atol=1)
 
-    kde = KDE(theta, weights)
+    kde = KDE(theta, weights=weights)
     kde.generate_kde()
 
     stats = calculate(kde).statistics()
     [check_stats(i) for i in range(2)]
 
-    equal_weight_theta = samples.compress(50)[names].values
+    equal_weight_theta = mcmc_samples.compress(50)[names].values
     x = kde.sample(len(equal_weight_theta))
 
     res = [ks_2samp(equal_weight_theta[:, i], x[:, i]) 
@@ -138,9 +139,22 @@ def check_stats(i):
 
 def test_kde_save_load():
 
-    kde = KDE(theta, weights)
+    kde = KDE(theta, weights=weights)
     kde.generate_kde()
     file = 'saved_maf.pkl'
     kde.save(file)
     loaded_kde = KDE.load(file)
     assert_equal(kde.kde.covariance, loaded_kde.kde.covariance)
+
+def test_anesthetic():
+
+    kde = KDE(mcmc_samples, parameters=names)
+    maf = MAF(mcmc_samples, parameters=names)
+    cmaf = clusterMAF(mcmc_samples)
+
+    assert_equal(kde.parameters, names)
+    assert_equal(maf.parameters, names)
+
+    # not providing parametes here but deriving them from the 
+    # anesthetic object columns
+    assert(np.all(cmaf.parameters == np.array(names)))