Skip to content

API Reference

margarine.base.baseflow

Base density estimator for margarine package.

Defines a base class for density estimators with common interface methods including:

  • train
  • sample
  • __call__
  • log_prob
  • log_like
  • save
  • load

BaseDensityEstimator

Bases: ABC

Base class for density estimators in the margarine package.

Source code in margarine/base/baseflow.py 
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
class BaseDensityEstimator(ABC):
    """Base class for density estimators in the margarine package."""

    @abstractmethod
    def train(self) -> None:
        """Train the density estimator on the provided data."""
        raise NotImplementedError("Train method must be implemented.")

    @abstractmethod
    def sample(self, key: jnp.ndarray, num_samples: int) -> jnp.ndarray:
        """Generate samples from the density estimator.

        Args:
            key: JAX random key for sampling.
            num_samples: Number of samples to generate.

        Returns:
            jnp.ndarray: Generated samples as a JAX array.
        """
        u = jax.random.uniform(key, shape=(num_samples, self.theta.shape[1]))
        raise self(u)

    def __call__(self, u: jnp.ndarray) -> jnp.ndarray:
        """Evaluate the density estimator at given points.

        Args:
            u: Samples from the unit hypercube.

        Returns:
            jnp.ndarray: samples from the density estimator.
        """
        raise NotImplementedError("Call method must be implemented.")

    @abstractmethod
    def log_prob(self, x: jnp.ndarray) -> jnp.ndarray:
        """Compute the log-probability of given samples.

        Args:
            x: Samples for which to compute the log-probability.

        Returns:
            jnp.ndarray: Log-probabilities of the samples.
        """
        raise NotImplementedError("log_prob method must be implemented.")

    @abstractmethod
    def log_like(
        self,
        x: jnp.ndarray,
        logevidence: float,
        prior_density: jnp.ndarray,
    ) -> jnp.ndarray:
        """Compute the log-likelihood of given samples.

        Args:
            x: Samples for which to compute the log-likelihood.
            logevidence: Log-evidence value.
            prior_density: Prior density estimator or densities.

        Returns:
            jnp.ndarray: Log-likelihoods of the samples.
        """
        return NotImplementedError("log_like method must be implemented.")

    def save(self, filepath: str) -> None:
        """Save the density estimator to a file.

        Args:
            filepath: Path to the file where the estimator will be saved.
        """
        raise NotImplementedError("save method must be implemented.")

    @classmethod
    def load(cls, filepath: str) -> "BaseDensityEstimator":
        """Load a density estimator from a file.

        Args:
            filepath: Path to the file from which to load the estimator.

        Returns:
            BaseDensityEstimator: Loaded density estimator instance.
        """
        raise NotImplementedError("load method must be implemented.")

__call__(u)

Evaluate the density estimator at given points.

Parameters:

Name Type Description Default
u ndarray

Samples from the unit hypercube.

 required

Returns:

Type Description
ndarray

jnp.ndarray: samples from the density estimator.
Source code in margarine/base/baseflow.py 
44
45
46
47
48
49
50
51
52
53
def __call__(self, u: jnp.ndarray) -> jnp.ndarray:
    """Evaluate the density estimator at given points.

    Args:
        u: Samples from the unit hypercube.

    Returns:
        jnp.ndarray: samples from the density estimator.
    """
    raise NotImplementedError("Call method must be implemented.")

load(filepath) classmethod

Load a density estimator from a file.

Parameters:

Name Type Description Default
filepath str

Path to the file from which to load the estimator.

 required

Returns:

Name Type Description
BaseDensityEstimator BaseDensityEstimator

Loaded density estimator instance.
Source code in margarine/base/baseflow.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
@classmethod
def load(cls, filepath: str) -> "BaseDensityEstimator":
    """Load a density estimator from a file.

    Args:
        filepath: Path to the file from which to load the estimator.

    Returns:
        BaseDensityEstimator: Loaded density estimator instance.
    """
    raise NotImplementedError("load method must be implemented.")

log_like(x, logevidence, prior_density) abstractmethod

Compute the log-likelihood of given samples.

Parameters:

Name Type Description Default
x ndarray

Samples for which to compute the log-likelihood.

 required
logevidence float

Log-evidence value.

 required
prior_density ndarray

Prior density estimator or densities.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Log-likelihoods of the samples.
Source code in margarine/base/baseflow.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
@abstractmethod
def log_like(
    self,
    x: jnp.ndarray,
    logevidence: float,
    prior_density: jnp.ndarray,
) -> jnp.ndarray:
    """Compute the log-likelihood of given samples.

    Args:
        x: Samples for which to compute the log-likelihood.
        logevidence: Log-evidence value.
        prior_density: Prior density estimator or densities.

    Returns:
        jnp.ndarray: Log-likelihoods of the samples.
    """
    return NotImplementedError("log_like method must be implemented.")

log_prob(x) abstractmethod

Compute the log-probability of given samples.

Parameters:

Name Type Description Default
x ndarray

Samples for which to compute the log-probability.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Log-probabilities of the samples.
Source code in margarine/base/baseflow.py 
55
56
57
58
59
60
61
62
63
64
65
@abstractmethod
def log_prob(self, x: jnp.ndarray) -> jnp.ndarray:
    """Compute the log-probability of given samples.

    Args:
        x: Samples for which to compute the log-probability.

    Returns:
        jnp.ndarray: Log-probabilities of the samples.
    """
    raise NotImplementedError("log_prob method must be implemented.")

sample(key, num_samples) abstractmethod

Generate samples from the density estimator.

Parameters:

Name Type Description Default
key ndarray

JAX random key for sampling.

 required
num_samples int

Number of samples to generate.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Generated samples as a JAX array.
Source code in margarine/base/baseflow.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
@abstractmethod
def sample(self, key: jnp.ndarray, num_samples: int) -> jnp.ndarray:
    """Generate samples from the density estimator.

    Args:
        key: JAX random key for sampling.
        num_samples: Number of samples to generate.

    Returns:
        jnp.ndarray: Generated samples as a JAX array.
    """
    u = jax.random.uniform(key, shape=(num_samples, self.theta.shape[1]))
    raise self(u)

save(filepath)

Save the density estimator to a file.

Parameters:

Name Type Description Default
filepath str

Path to the file where the estimator will be saved.

 required
Source code in margarine/base/baseflow.py 
86
87
88
89
90
91
92
def save(self, filepath: str) -> None:
    """Save the density estimator to a file.

    Args:
        filepath: Path to the file where the estimator will be saved.
    """
    raise NotImplementedError("save method must be implemented.")

train() abstractmethod

Train the density estimator on the provided data.

Source code in margarine/base/baseflow.py 
25
26
27
28
@abstractmethod
def train(self) -> None:
    """Train the density estimator on the provided data."""
    raise NotImplementedError("Train method must be implemented.")

margarine.estimators.kde

KDE implementation using JAX.

KDE

Bases: BaseDensityEstimator

Kernel Density Estimator (KDE) using JAX.

Source code in margarine/estimators/kde.py 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
class KDE(BaseDensityEstimator):
    """Kernel Density Estimator (KDE) using JAX."""

    def __init__(
        self,
        theta: jnp.ndarray,
        weights: jnp.ndarray | None = None,
        theta_ranges: jnp.ndarray | None = None,
        bandwidth: float | str = "silverman",
    ) -> None:
        """Initialize the KDE.

        Args:
            theta: Parameters of the density estimator.
            weights: Optional weights for the parameters.
            theta_ranges: Optional ranges for the parameters.
            bandwidth: Bandwidth for the KDE.
        """
        self.theta = theta
        self.weights = weights
        self.theta_ranges = theta_ranges
        self.bandwidth = bandwidth

        if self.weights is None:
            self.weights = jnp.ones(len(self.theta))

        if theta_ranges is None:
            self.theta_ranges = approximate_bounds(self.theta, self.weights)

    def train(self) -> stats.gaussian_kde:
        """Generates a weighted KDE."""
        phi = forward_transform(
            self.theta, self.theta_ranges[0], self.theta_ranges[1]
        )
        weights = self.weights / jnp.sum(self.weights)
        self.kde = stats.gaussian_kde(
            phi.T, weights=weights, bw_method=self.bandwidth
        )
        return self.kde

    def sample(self, key: jnp.ndarray, num_samples: int) -> jnp.ndarray:
        """Sample from the KDE.

        Args:
            key: JAX random key for sampling.
            num_samples: Number of samples to draw.

        Returns:
            jnp.ndarray: Samples drawn from the KDE.
        """
        x = self.kde.resample(key, (num_samples,)).T
        x = inverse_transform(x, self.theta_ranges[0], self.theta_ranges[1])
        return x

    def __call__(self, u: jnp.ndarray) -> jnp.ndarray:
        r"""Transform samples from the unit hypercube to samples on the KDE.

        Uses the Rosenblatt transformation
        (conditional inverse transform sampling)
        to map uniform samples to a Gaussian mixture KDE distribution.

        In detail, for a d-dimensional KDE, each output
        dimension is computed sequentially:

            x_1 = F_{X_1}^{-1}(u_1)
            x_2 = F_{X_2 | X_1}^{-1}(u_2 | x_1)
            ...
            x_d = F_{X_d | X_{1:d-1}}^{-1}(u_d | x_1, ..., x_{d-1})

        where \(u_i \sim \text{Uniform}[0,1]\), \(F^{-1}\)
        denotes the inverse CDF, and
        the conditional CDFs are computed
        from the Gaussian mixture KDE. Since these
        inverse CDFs do not have a closed form, we
        solve for each x_i using a
        root-finding algorithm (e.g., Newton-Raphson).

        Args:
            u (jnp.ndarray): Samples from the unit hypercube [0,1]^d
                Shape: (n_samples, n_dims)

        Returns:
            jnp.ndarray: The transformed samples following the
                KDE distribution.
                Shape: (n_samples, n_dims)
        """
        # generate useful parameters for __call__ function to transform
        # hypercube into samples on the KDE.
        S = self.kde.covariance
        mu = self.kde.dataset.T
        steps, s = [], []
        for i in range(mu.shape[-1]):
            if i == 0:
                step = jnp.array([])
                s_i = jnp.sqrt(S[0, 0])
            else:
                step = jnp.linalg.solve(S[:i, :i].T, S[i, :i]).T
                s_i = jnp.sqrt(S[i, i] - step @ S[:i, i])
            steps.append(step)  # conditional means
            s.append(s_i)  # conditional stddevs

        @jax.jit
        def newton_root(
            m: jnp.ndarray,
            s_i: float,
            target: float,
            init: float,
            weights: jnp.ndarray,
            maxiter: int = 20,
        ) -> float:
            """Solve for x in the conditional CDF equation with Newton-Raphson.

            Finds the scalar root of:

                sum_k w_k * Φ((x - m_k) / s_i) - target = 0

            where Φ is the standard normal CDF, m is the
            conditional mean vector
            for the current dimension, s_i is the
            conditional standard deviation,
            weights are the mixture weights of the KDE, and target ∈ [0,1] is
            the uniform sample to invert.

            Args:
                m (jnp.ndarray): Conditional means of
                    the mixture components, shape (n_kernels,).
                s_i (float): Conditional standard deviation for this dimension.
                target (float): Uniform sample in [0,1] to map to the KDE.
                init (float): Initial guess for the Newton-Raphson iteration.
                weights (jnp.ndarray): Mixture weights, shape (n_kernels,).
                maxiter (int, optional): Maximum number of iterations.
                    Default is 20.

            Returns:
                float: The root x such that the conditional CDF equals target.
            """

            def f(x: float) -> float:
                """Function for which we want to find the root."""
                return (
                    tfd.Normal(0, 1).cdf((x - m) / s_i) * weights
                ).sum() - target

            def f_prime(x: float) -> float:
                """Derivative of the function f."""
                pdf = tfd.Normal(0, 1).prob((x - m) / s_i) / s_i
                return (pdf * weights).sum()

            x = init
            for _ in range(maxiter):
                x = x - f(x) / f_prime(x)
            return x

        @jax.jit
        def transform_one(x: jnp.ndarray) -> jnp.ndarray:
            """Transform a single sample from unit hypercube to KDE."""
            y = jnp.zeros_like(x)
            for i in range(len(x)):
                m_i = (
                    mu[:, i] + steps[i] @ (y[:i] - mu[:, :i]).T
                    if i > 0
                    else mu[:, i]
                )
                init = mu[:, i].mean()
                y = y.at[i].set(
                    newton_root(m_i, s[i], x[i], init, self.kde.weights)
                )
            return inverse_transform(
                y, self.theta_ranges[0], self.theta_ranges[1]
            )

        transformed_samples = jax.vmap(transform_one)(u)
        return transformed_samples

    def log_prob(self, x: jnp.ndarray) -> jnp.ndarray:
        """Compute the log-probability of given samples.

        While the 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. The
        correction is implemented here.

        Args:
            x: Samples for which to compute the log-probability.

        Returns:
            jnp.ndarray: Log-probabilities of the samples.
        """
        transformed_x = forward_transform(
            x, self.theta_ranges[0], self.theta_ranges[1]
        )

        transform_chain = tfb.Chain(
            [
                tfb.Invert(tfb.NormalCDF()),
                tfb.Scale(1 / (self.theta_ranges[1] - self.theta_ranges[0])),
                tfb.Shift(-self.theta_ranges[0]),
            ]
        )

        def norm_jac(y: jnp.ndarray) -> jnp.ndarray:
            """Calculate the normalising jacobian for the transformation."""
            return transform_chain.inverse_log_det_jacobian(y, event_ndims=0)

        correction = -norm_jac(transformed_x).sum(axis=1)

        log_probs = jnp.log(self.kde.pdf(transformed_x.T)) + correction
        return log_probs

    def log_like(
        self,
        x: jnp.ndarray,
        logevidence: float,
        prior_density: jnp.ndarray | BaseDensityEstimator,
    ) -> jnp.ndarray:
        """Compute the marginal log-likelihood of given samples.

        Args:
            x: Samples for which to compute the log-likelihood.
            logevidence: Log-evidence term.
            prior_density: Prior density or density estimator.

        Returns:
            jnp.ndarray: Log-likelihoods of the samples.
        """
        if isinstance(prior_density, BaseDensityEstimator):
            prior_density = prior_density.log_prob(x)

        return self.log_prob(x) + logevidence - prior_density

    def save(self, filename: str) -> None:
        """Save the KDE to a file.

        Args:
            filename (str): Path to the file where the KDE will be saved.
        """
        path = Path(filename).resolve()
        if path.exists():
            shutil.rmtree(path)

        os.makedirs(path)

        with open(f"{path}/kde.pkl", "wb") as f:
            pickle.dump(self.kde, f)

        config = {
            "bandwidth": self.bandwidth,
            "theta_ranges": self.theta_ranges,
        }

        with open(f"{path}/config.yaml", "w") as f:
            yaml.dump(config, f)

        metadata = {
            "theta": self.theta,
            "weights": self.weights,
            "margarine_version": _version.__version__,
        }

        with open(f"{path}/metadata.yaml", "w") as f:
            yaml.dump(metadata, f)

        with ZipFile(filename + ".marg", "w") as z:
            for subpath in path.rglob("*"):
                if subpath.is_file():
                    z.write(subpath, arcname=subpath.relative_to(path))

        shutil.rmtree(path)

    @classmethod
    def load(cls, filename: str) -> BaseDensityEstimator | None:
        """Load a KDE from a file.

        Args:
            filename (str): Path to the file from which to load the KDE.

        Returns:
            KDE: Loaded KDE instance.
        """
        zip_path = Path(f"{filename}.marg")
        path = Path(filename + ".tmp").resolve()
        with ZipFile(zip_path) as z:
            # Extract all files to a folder
            z.extractall(path)

        with open(f"{path}/metadata.yaml") as f:
            metadata = yaml.unsafe_load(f)

        version = metadata.get("margarine_version", None)
        if version is None:
            print(
                "Warning: The KDE was saved with a version of margarine ",
                " < 2.0.0. In order to load it you will need to downgrade ",
                "margarine to a version < 2.0.0. e.g. ",
                "pip install margarine<2.0.0",
            )
            return
        if version != _version.__version__:
            print(
                f"Warning: The KDE was saved with margarine version "
                f"{version}, but you are loading it with version "
                f"{_version.__version__}. This may lead to "
                f"incompatibilities."
            )

        with open(f"{path}/config.yaml") as f:
            config = yaml.unsafe_load(f)

        instance = cls(
            theta=jnp.array([]),  # Placeholder, will be overwritten
            weights=jnp.array([]),  # Placeholder, will be overwritten
            bandwidth=config["bandwidth"],
            theta_ranges=config["theta_ranges"],
        )

        instance.theta = jnp.array(metadata["theta"])
        instance.weights = jnp.array(metadata["weights"])

        with open(f"{path}/kde.pkl", "rb") as f:
            kde = pickle.load(f)

        instance.kde = kde

        shutil.rmtree(path)
        return instance

__call__(u)

Transform samples from the unit hypercube to samples on the KDE.

Uses the Rosenblatt transformation (conditional inverse transform sampling) to map uniform samples to a Gaussian mixture KDE distribution.

In detail, for a d-dimensional KDE, each output dimension is computed sequentially:

x_1 = F_{X_1}^{-1}(u_1)
x_2 = F_{X_2 | X_1}^{-1}(u_2 | x_1)
...
x_d = F_{X_d | X_{1:d-1}}^{-1}(u_d | x_1, ..., x_{d-1})

where (u_i \sim \text{Uniform}[0,1]), (F^{-1}) denotes the inverse CDF, and the conditional CDFs are computed from the Gaussian mixture KDE. Since these inverse CDFs do not have a closed form, we solve for each x_i using a root-finding algorithm (e.g., Newton-Raphson).

Parameters:

Name Type Description Default
u ndarray

Samples from the unit hypercube [0,1]^d Shape: (n_samples, n_dims)

 required

Returns:

Type Description
ndarray

jnp.ndarray: The transformed samples following the KDE distribution. Shape: (n_samples, n_dims)
Source code in margarine/estimators/kde.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
def __call__(self, u: jnp.ndarray) -> jnp.ndarray:
    r"""Transform samples from the unit hypercube to samples on the KDE.

    Uses the Rosenblatt transformation
    (conditional inverse transform sampling)
    to map uniform samples to a Gaussian mixture KDE distribution.

    In detail, for a d-dimensional KDE, each output
    dimension is computed sequentially:

        x_1 = F_{X_1}^{-1}(u_1)
        x_2 = F_{X_2 | X_1}^{-1}(u_2 | x_1)
        ...
        x_d = F_{X_d | X_{1:d-1}}^{-1}(u_d | x_1, ..., x_{d-1})

    where \(u_i \sim \text{Uniform}[0,1]\), \(F^{-1}\)
    denotes the inverse CDF, and
    the conditional CDFs are computed
    from the Gaussian mixture KDE. Since these
    inverse CDFs do not have a closed form, we
    solve for each x_i using a
    root-finding algorithm (e.g., Newton-Raphson).

    Args:
        u (jnp.ndarray): Samples from the unit hypercube [0,1]^d
            Shape: (n_samples, n_dims)

    Returns:
        jnp.ndarray: The transformed samples following the
            KDE distribution.
            Shape: (n_samples, n_dims)
    """
    # generate useful parameters for __call__ function to transform
    # hypercube into samples on the KDE.
    S = self.kde.covariance
    mu = self.kde.dataset.T
    steps, s = [], []
    for i in range(mu.shape[-1]):
        if i == 0:
            step = jnp.array([])
            s_i = jnp.sqrt(S[0, 0])
        else:
            step = jnp.linalg.solve(S[:i, :i].T, S[i, :i]).T
            s_i = jnp.sqrt(S[i, i] - step @ S[:i, i])
        steps.append(step)  # conditional means
        s.append(s_i)  # conditional stddevs

    @jax.jit
    def newton_root(
        m: jnp.ndarray,
        s_i: float,
        target: float,
        init: float,
        weights: jnp.ndarray,
        maxiter: int = 20,
    ) -> float:
        """Solve for x in the conditional CDF equation with Newton-Raphson.

        Finds the scalar root of:

            sum_k w_k * Φ((x - m_k) / s_i) - target = 0

        where Φ is the standard normal CDF, m is the
        conditional mean vector
        for the current dimension, s_i is the
        conditional standard deviation,
        weights are the mixture weights of the KDE, and target ∈ [0,1] is
        the uniform sample to invert.

        Args:
            m (jnp.ndarray): Conditional means of
                the mixture components, shape (n_kernels,).
            s_i (float): Conditional standard deviation for this dimension.
            target (float): Uniform sample in [0,1] to map to the KDE.
            init (float): Initial guess for the Newton-Raphson iteration.
            weights (jnp.ndarray): Mixture weights, shape (n_kernels,).
            maxiter (int, optional): Maximum number of iterations.
                Default is 20.

        Returns:
            float: The root x such that the conditional CDF equals target.
        """

        def f(x: float) -> float:
            """Function for which we want to find the root."""
            return (
                tfd.Normal(0, 1).cdf((x - m) / s_i) * weights
            ).sum() - target

        def f_prime(x: float) -> float:
            """Derivative of the function f."""
            pdf = tfd.Normal(0, 1).prob((x - m) / s_i) / s_i
            return (pdf * weights).sum()

        x = init
        for _ in range(maxiter):
            x = x - f(x) / f_prime(x)
        return x

    @jax.jit
    def transform_one(x: jnp.ndarray) -> jnp.ndarray:
        """Transform a single sample from unit hypercube to KDE."""
        y = jnp.zeros_like(x)
        for i in range(len(x)):
            m_i = (
                mu[:, i] + steps[i] @ (y[:i] - mu[:, :i]).T
                if i > 0
                else mu[:, i]
            )
            init = mu[:, i].mean()
            y = y.at[i].set(
                newton_root(m_i, s[i], x[i], init, self.kde.weights)
            )
        return inverse_transform(
            y, self.theta_ranges[0], self.theta_ranges[1]
        )

    transformed_samples = jax.vmap(transform_one)(u)
    return transformed_samples

__init__(theta, weights=None, theta_ranges=None, bandwidth='silverman')

Initialize the KDE.

Parameters:

Name Type Description Default
theta ndarray

Parameters of the density estimator.

 required
weights ndarray | None

Optional weights for the parameters.

 None
theta_ranges ndarray | None

Optional ranges for the parameters.

 None
bandwidth float | str

Bandwidth for the KDE.

 'silverman'
Source code in margarine/estimators/kde.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
def __init__(
    self,
    theta: jnp.ndarray,
    weights: jnp.ndarray | None = None,
    theta_ranges: jnp.ndarray | None = None,
    bandwidth: float | str = "silverman",
) -> None:
    """Initialize the KDE.

    Args:
        theta: Parameters of the density estimator.
        weights: Optional weights for the parameters.
        theta_ranges: Optional ranges for the parameters.
        bandwidth: Bandwidth for the KDE.
    """
    self.theta = theta
    self.weights = weights
    self.theta_ranges = theta_ranges
    self.bandwidth = bandwidth

    if self.weights is None:
        self.weights = jnp.ones(len(self.theta))

    if theta_ranges is None:
        self.theta_ranges = approximate_bounds(self.theta, self.weights)

load(filename) classmethod

Load a KDE from a file.

Parameters:

Name Type Description Default
filename str

Path to the file from which to load the KDE.

 required

Returns:

Name Type Description
KDE BaseDensityEstimator | None

Loaded KDE instance.
Source code in margarine/estimators/kde.py 
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
@classmethod
def load(cls, filename: str) -> BaseDensityEstimator | None:
    """Load a KDE from a file.

    Args:
        filename (str): Path to the file from which to load the KDE.

    Returns:
        KDE: Loaded KDE instance.
    """
    zip_path = Path(f"{filename}.marg")
    path = Path(filename + ".tmp").resolve()
    with ZipFile(zip_path) as z:
        # Extract all files to a folder
        z.extractall(path)

    with open(f"{path}/metadata.yaml") as f:
        metadata = yaml.unsafe_load(f)

    version = metadata.get("margarine_version", None)
    if version is None:
        print(
            "Warning: The KDE was saved with a version of margarine ",
            " < 2.0.0. In order to load it you will need to downgrade ",
            "margarine to a version < 2.0.0. e.g. ",
            "pip install margarine<2.0.0",
        )
        return
    if version != _version.__version__:
        print(
            f"Warning: The KDE was saved with margarine version "
            f"{version}, but you are loading it with version "
            f"{_version.__version__}. This may lead to "
            f"incompatibilities."
        )

    with open(f"{path}/config.yaml") as f:
        config = yaml.unsafe_load(f)

    instance = cls(
        theta=jnp.array([]),  # Placeholder, will be overwritten
        weights=jnp.array([]),  # Placeholder, will be overwritten
        bandwidth=config["bandwidth"],
        theta_ranges=config["theta_ranges"],
    )

    instance.theta = jnp.array(metadata["theta"])
    instance.weights = jnp.array(metadata["weights"])

    with open(f"{path}/kde.pkl", "rb") as f:
        kde = pickle.load(f)

    instance.kde = kde

    shutil.rmtree(path)
    return instance

log_like(x, logevidence, prior_density)

Compute the marginal log-likelihood of given samples.

Parameters:

Name Type Description Default
x ndarray

Samples for which to compute the log-likelihood.

 required
logevidence float

Log-evidence term.

 required
prior_density ndarray | BaseDensityEstimator

Prior density or density estimator.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Log-likelihoods of the samples.
Source code in margarine/estimators/kde.py 
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
def log_like(
    self,
    x: jnp.ndarray,
    logevidence: float,
    prior_density: jnp.ndarray | BaseDensityEstimator,
) -> jnp.ndarray:
    """Compute the marginal log-likelihood of given samples.

    Args:
        x: Samples for which to compute the log-likelihood.
        logevidence: Log-evidence term.
        prior_density: Prior density or density estimator.

    Returns:
        jnp.ndarray: Log-likelihoods of the samples.
    """
    if isinstance(prior_density, BaseDensityEstimator):
        prior_density = prior_density.log_prob(x)

    return self.log_prob(x) + logevidence - prior_density

log_prob(x)

Compute the log-probability of given samples.

While the 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. The correction is implemented here.

Parameters:

Name Type Description Default
x ndarray

Samples for which to compute the log-probability.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Log-probabilities of the samples.
Source code in margarine/estimators/kde.py 
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
def log_prob(self, x: jnp.ndarray) -> jnp.ndarray:
    """Compute the log-probability of given samples.

    While the 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. The
    correction is implemented here.

    Args:
        x: Samples for which to compute the log-probability.

    Returns:
        jnp.ndarray: Log-probabilities of the samples.
    """
    transformed_x = forward_transform(
        x, self.theta_ranges[0], self.theta_ranges[1]
    )

    transform_chain = tfb.Chain(
        [
            tfb.Invert(tfb.NormalCDF()),
            tfb.Scale(1 / (self.theta_ranges[1] - self.theta_ranges[0])),
            tfb.Shift(-self.theta_ranges[0]),
        ]
    )

    def norm_jac(y: jnp.ndarray) -> jnp.ndarray:
        """Calculate the normalising jacobian for the transformation."""
        return transform_chain.inverse_log_det_jacobian(y, event_ndims=0)

    correction = -norm_jac(transformed_x).sum(axis=1)

    log_probs = jnp.log(self.kde.pdf(transformed_x.T)) + correction
    return log_probs

sample(key, num_samples)

Sample from the KDE.

Parameters:

Name Type Description Default
key ndarray

JAX random key for sampling.

 required
num_samples int

Number of samples to draw.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Samples drawn from the KDE.
Source code in margarine/estimators/kde.py 
67
68
69
70
71
72
73
74
75
76
77
78
79
def sample(self, key: jnp.ndarray, num_samples: int) -> jnp.ndarray:
    """Sample from the KDE.

    Args:
        key: JAX random key for sampling.
        num_samples: Number of samples to draw.

    Returns:
        jnp.ndarray: Samples drawn from the KDE.
    """
    x = self.kde.resample(key, (num_samples,)).T
    x = inverse_transform(x, self.theta_ranges[0], self.theta_ranges[1])
    return x

save(filename)

Save the KDE to a file.

Parameters:

Name Type Description Default
filename str

Path to the file where the KDE will be saved.

 required
Source code in margarine/estimators/kde.py 
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
def save(self, filename: str) -> None:
    """Save the KDE to a file.

    Args:
        filename (str): Path to the file where the KDE will be saved.
    """
    path = Path(filename).resolve()
    if path.exists():
        shutil.rmtree(path)

    os.makedirs(path)

    with open(f"{path}/kde.pkl", "wb") as f:
        pickle.dump(self.kde, f)

    config = {
        "bandwidth": self.bandwidth,
        "theta_ranges": self.theta_ranges,
    }

    with open(f"{path}/config.yaml", "w") as f:
        yaml.dump(config, f)

    metadata = {
        "theta": self.theta,
        "weights": self.weights,
        "margarine_version": _version.__version__,
    }

    with open(f"{path}/metadata.yaml", "w") as f:
        yaml.dump(metadata, f)

    with ZipFile(filename + ".marg", "w") as z:
        for subpath in path.rglob("*"):
            if subpath.is_file():
                z.write(subpath, arcname=subpath.relative_to(path))

    shutil.rmtree(path)

train()

Generates a weighted KDE.

Source code in margarine/estimators/kde.py 
56
57
58
59
60
61
62
63
64
65
def train(self) -> stats.gaussian_kde:
    """Generates a weighted KDE."""
    phi = forward_transform(
        self.theta, self.theta_ranges[0], self.theta_ranges[1]
    )
    weights = self.weights / jnp.sum(self.weights)
    self.kde = stats.gaussian_kde(
        phi.T, weights=weights, bw_method=self.bandwidth
    )
    return self.kde

margarine.estimators.nice

Implementation of the NICE estimator.

NICE

Bases: BaseDensityEstimator, Module

Implementation of the NICE architecture for density estimation.

Details in https://arxiv.org/abs/1410.8516.

Source code in margarine/estimators/nice.py 
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
class NICE(BaseDensityEstimator, nnx.Module):
    """Implementation of the NICE architecture for density estimation.

    Details in https://arxiv.org/abs/1410.8516.
    """

    def __init__(
        self,
        theta: jnp.ndarray,
        weights: jnp.ndarray | None = None,
        theta_ranges: jnp.ndarray | None = None,
        in_size: int = 2,
        hidden_size: int = 128,
        num_layers: int = 2,
        num_coupling_layers: int = 4,
        nnx_rngs: dict | None = None,
    ) -> None:
        """Initialize the NICE estimator.

        Args:
            theta: Parameters of the density estimator.
            weights: Optional weights for the parameters.
            theta_ranges: Optional ranges for the parameters.
            in_size: Input size.
            hidden_size: Sizes of hidden layers.
            num_layers: Number of layers in each coupling network.
            num_coupling_layers: Number of coupling layers.
            nnx_rngs: Optional RNGs for Flax.
        """
        super().__init__()
        self.theta = nnx.data(theta)
        self.weights = nnx.data(weights)
        self.theta_ranges = nnx.data(theta_ranges)

        if self.weights is None:
            self.weights = jnp.ones(len(self.theta))
        if theta_ranges is None:
            self.theta_ranges = approximate_bounds(self.theta, self.weights)

        if nnx_rngs is None:
            nnx_rngs = nnx.Rngs(0)

        self.in_size = in_size
        self.hidden_size = hidden_size
        self.nlayers = num_layers
        self.num_coupling_layers = num_coupling_layers

        self.pass_size = in_size // 2
        self.net_in_size = in_size - self.pass_size

        layers = nnx.List()
        layers.append(
            nnx.Linear(self.net_in_size, self.hidden_size, rngs=nnx_rngs)
        )
        layers.append(lambda x: jax.nn.gelu(x))
        for _ in range(self.nlayers):
            layers.append(
                nnx.Linear(self.hidden_size, self.hidden_size, rngs=nnx_rngs)
            )
            layers.append(lambda x: jax.nn.gelu(x))

        layers.append(
            nnx.Linear(self.hidden_size, self.pass_size, rngs=nnx_rngs)
        )
        self.mlp = nnx.List(
            [nnx.Sequential(*layers) for _ in range(self.num_coupling_layers)]
        )

        self.S = nnx.Param(jnp.zeros(in_size))

    @jax.jit
    def forward(self, x: jnp.ndarray) -> jnp.ndarray:
        """NICE forward pass.

        This is the forward pass of the NICE coupling layer from
        samples in the target space to the base distribution space.

        Args:
            x (jnp.ndarray): Input samples.
        """
        mlp_fns = tuple((lambda xb, mlp=mlp: mlp(xb)) for mlp in self.mlp)

        def body(i: int, x: jnp.ndarray) -> jnp.ndarray:
            """Body function for the coupling layers.

            Args:
                i: Forward counter.
                x: Input samples.

            Returns:
                Transformed samples.
            """
            x = jax.lax.cond(
                i % 2 == 0,
                lambda x: jnp.flip(x, axis=-1),
                lambda x: x,
                x,
            )
            xa, xb = x[:, : self.pass_size], x[:, self.pass_size :]
            m = jax.lax.switch(i, mlp_fns, xb)
            x = jnp.concatenate([xa + m, xb], axis=1)
            return x

        x = jax.lax.fori_loop(0, self.num_coupling_layers, body, x)

        x = jnp.exp(self.S) * x

        if self.num_coupling_layers % 2 == 1:
            x = jnp.flip(x, axis=-1)

        return x

    @jax.jit
    def inverse(self, x: jnp.ndarray) -> jnp.ndarray:
        """NICE inverse pass.

        Args:
            x: Input samples.
        """
        if self.num_coupling_layers % 2 == 1:
            x = jnp.flip(x, axis=-1)

        x = x / jnp.exp(self.S)  # Undo scaling first

        mlp_fns = tuple((lambda xb, mlp=mlp: mlp(xb)) for mlp in self.mlp)

        def body(k: int, x: jnp.ndarray) -> jnp.ndarray:
            """Body function for the inverse coupling layers.

            Args:
                k: Forward counter.
                x: Input samples.

            Returns:
                Inverted samples.
            """
            # convert upward counter to downward counter
            i = self.num_coupling_layers - 1 - k

            xa, xb = x[:, : self.pass_size], x[:, self.pass_size :]
            m = jax.lax.switch(i, mlp_fns, xb)
            x = jnp.concatenate([xa - m, xb], axis=1)
            x = jax.lax.cond(
                i % 2 == 0,
                lambda x: jnp.flip(x, axis=-1),
                lambda x: x,
                x,
            )
            return x

        return jax.lax.fori_loop(0, self.num_coupling_layers, body, x)

    def train(
        self,
        key: jnp.ndarray,
        learning_rate: float = 1e-4,
        epochs: int = 1000,
        patience: int = 50,
        batch_size: int = 256,
    ) -> None:
        """Train the NICE model.

        Args:
            key: JAX random key for data splitting.
            learning_rate: Learning rate for the optimizer.
            epochs: Number of training epochs.
            patience: Patience for early stopping.
            batch_size: Batch size for training.
        """

        @nnx.jit
        def loss_function(
            model: "NICE",
            targets: jnp.ndarray,
            weights: jnp.ndarray,
        ) -> jnp.ndarray:
            """Loss function for training NICE model.

            Args:
                model (NICE): NICE model.
                targets (jnp.ndarray): Samples from the target distribution
                    that have been passed through the forward_transform.
                weights (jnp.ndarray): Weights for the samples.

            Returns:
                jnp.ndarray: Computed loss.
            """
            return -jnp.mean(weights * model.log_prob_under_NICE(targets))

        @nnx.jit
        def train_step(
            model: NICE,
            optimizer: nnx.Optimizer,
            train_phi: jnp.ndarray,
            train_weights: jnp.ndarray,
        ) -> None:
            """Single training step for NICE model.

            Args:
                model (NICE): NICE model.
                optimizer (nnx.Optimizer): Optimizer for
                    updating model parameters.
                train_phi (jnp.ndarray): Training samples.
                train_weights (jnp.ndarray): Weights for the training samples.
            """
            grad = nnx.grad(loss_function)(model, train_phi, train_weights)
            optimizer.update(model, grad)

        phi = forward_transform(
            self.theta, self.theta_ranges[0], self.theta_ranges[1]
        )
        weights = self.weights / jnp.sum(self.weights)

        key, subkey = jax.random.split(key)
        # need to split the data into training and validation sets
        (
            self.train_phi,
            self.test_phi,
            self.train_weights,
            self.test_weights,
        ) = train_test_split(
            phi,
            weights,
            subkey,
            test_size=0.2,
        )
        key, subkey = jax.random.split(key)
        self.test_phi, self.val_phi, self.test_weights, self.val_weights = (
            train_test_split(
                self.test_phi, self.test_weights, subkey, test_size=0.5
            )
        )

        tx = optax.adam(learning_rate)
        optimizer = nnx.Optimizer(self, tx, wrt=nnx.Param)

        pbar = tqdm.tqdm(range(epochs), desc="Training")

        best_loss = jnp.inf
        best_model = nnx.state(self, nnx.Param)
        c = 0

        data_size = len(self.train_phi)

        tl, vl = [], []
        for _ in pbar:
            data_permutations = jax.random.permutation(
                subkey, jnp.arange(data_size)
            )
            accumulated_train_loss = 0.0
            for i in range(0, len(self.train_phi), batch_size):
                batch_indices = data_permutations[i : i + batch_size]
                batch_phi = self.train_phi[batch_indices]
                batch_weights = self.train_weights[batch_indices]
                loss = loss_function(self, batch_phi, batch_weights)
                train_step(self, optimizer, batch_phi, batch_weights)
                accumulated_train_loss += loss * len(batch_phi)
            loss = accumulated_train_loss / len(self.train_phi)
            tl.append(loss)

            val_loss = loss_function(self, self.val_phi, self.val_weights)
            vl.append(val_loss)

            if val_loss < best_loss:
                best_loss = val_loss
                best_model = nnx.state(self, nnx.Param)
                best_epoch = _
                c = 0
            else:
                c += 1
                if c >= patience:
                    print(
                        f"Early stopping at epoch {_}, best "
                        + f"epoch was {best_epoch} with val loss {best_loss}"
                    )
                    break

            pbar.set_postfix(
                loss=f"{loss:.3e}",
                val_loss=f"{val_loss:.3e}",
                best_loss=f"{best_loss:.3e}",
            )

        self.train_loss = jnp.array(tl)
        self.val_loss = jnp.array(vl)

        if best_model is not None:
            nnx.update(self, best_model)

    def sample(self, key: jnp.ndarray, num_samples: int) -> jnp.ndarray:
        """Sample from the trained NICE model.

        Args:
            key: JAX random key.
            num_samples: Number of samples to draw.

        Returns:
            Samples drawn from the NICE model.
        """
        u = jax.random.uniform(key, shape=(num_samples, self.theta.shape[1]))

        samples = self(u)
        return samples

    def __call__(self, u: jnp.ndarray) -> jnp.ndarray:
        """Transform samples from the unit hypercube.

        Args:
            u (jnp.ndarray): Samples from the unit hypercube.

        Returns:
            Log probabilities of the input data.
        """
        # from unit hypercube to standard normal space
        x = forward_transform(u, 0, 1)

        # from standard normal to target space
        x = self.inverse(x)

        # from gaussianized target space to original target space
        x = inverse_transform(x, self.theta_ranges[0], self.theta_ranges[1])
        return x

    @jax.jit
    def log_prob_under_NICE(self, x: jnp.ndarray) -> jnp.ndarray:
        """Compute the log probability under the NICE model.

        Args:
            x (jnp.ndarray): Input data.

        Returns:
            Log probabilities of the input data.
        """
        # calculate the actual log prob under the NICE model
        # assuming a standard normal base distribution
        z = self.forward(x)
        log_pz = -0.5 * jnp.sum(z**2 + jnp.log(2 * jnp.pi), axis=1)
        log_det_J = jnp.sum(self.S)
        log_prob_under_nice = log_pz + log_det_J
        return log_prob_under_nice

    @jax.jit
    def log_prob(self, x: jnp.ndarray) -> jnp.ndarray:
        """Compute the log probability of the input data.

        Args:
            x (jnp.ndarray): Input data.

        Returns:
            Log probabilities of the input data.
        """
        transformed_x = forward_transform(
            x, self.theta_ranges[0], self.theta_ranges[1]
        )
        transform_chain = tfb.Chain(
            [
                tfb.Invert(tfb.NormalCDF()),
                tfb.Scale(1 / (self.theta_ranges[1] - self.theta_ranges[0])),
                tfb.Shift(-self.theta_ranges[0]),
            ]
        )

        def norm_jac(y: jnp.ndarray) -> jnp.ndarray:
            """Calculate the normalising jacobian for the transformation."""
            return transform_chain.inverse_log_det_jacobian(y, event_ndims=0)

        correction = -norm_jac(transformed_x).sum(axis=1)

        return self.log_prob_under_NICE(transformed_x) + correction

    def log_like(
        self,
        x: jnp.ndarray,
        logevidence: float,
        prior_density: jnp.ndarray | BaseDensityEstimator,
    ) -> jnp.ndarray:
        """Compute the marginal log-likelihood of given samples.

        Args:
            x: Samples for which to compute the log-likelihood.
            logevidence: Log-evidence term.
            prior_density: Prior density or density estimator.

        Returns:
            Log likelihoods of the input data.
        """
        if isinstance(prior_density, BaseDensityEstimator):
            prior_density = prior_density.log_prob(x)

        return self.log_prob(x) + logevidence - prior_density

    def save(self, filename: str) -> None:
        """Save the trained NICE model to a file.

        Args:
            filename: Path to the file where the model will be saved.
        """
        path = Path(filename).resolve()
        if path.exists():
            shutil.rmtree(path)

        state = nnx.state(self)

        checkpointer = orbax.PyTreeCheckpointer()
        checkpointer.save(f"{path}/state", state)

        config = {
            "in_size": self.in_size,
            "hidden_size": self.hidden_size,
            "num_layers": self.nlayers,
            "num_coupling_layers": self.num_coupling_layers,
            "theta_ranges": self.theta_ranges,
        }
        with open(f"{path}/config.yaml", "w") as f:
            yaml.dump(config, f)

        metadata = {
            "theta": self.theta,
            "weights": self.weights,
            "train_loss": self.train_loss,
            "val_loss": self.val_loss,
            "train_phi": self.train_phi,
            "test_phi": self.test_phi,
            "train_weights": self.train_weights,
            "test_weights": self.test_weights,
            "val_phi": self.val_phi,
            "val_weights": self.val_weights,
            "margarine_version": _version.__version__,
        }

        with open(f"{path}/metadata.yaml", "w") as f:
            yaml.dump(metadata, f)

        with ZipFile(filename + ".marg", "w") as z:
            for subpath in path.rglob("*"):
                if subpath.is_file():
                    z.write(subpath, arcname=subpath.relative_to(path))

        shutil.rmtree(path)

    @classmethod
    def load(cls, filename: str) -> BaseDensityEstimator | None:
        """Load a trained NICE model from a file.

        Args:
            filename (str): Path to the file from which the
                model will be loaded.

        Returns:
            Loaded NICE model.
        """
        zip_path = Path(f"{filename}.marg")
        path = Path(filename + ".tmp").resolve()
        with ZipFile(zip_path) as z:
            # Extract all files to a folder
            z.extractall(path)

        with open(f"{path}/metadata.yaml") as f:
            metadata = yaml.unsafe_load(f)

        version = metadata.get("margarine_version", None)
        if version is None:
            print(
                "Warning: The KDE was saved with a version of margarine ",
                " < 2.0.0. In order to load it you will need to downgrade ",
                "margarine to a version < 2.0.0. e.g. ",
                "pip install margarine<2.0.0",
            )
            return
        if version != _version.__version__:
            print(
                f"Warning: The KDE was saved with margarine version "
                f"{version}, but you are loading it with version "
                f"{_version.__version__}. This may lead to "
                f"incompatibilities."
            )

        with open(f"{path}/config.yaml") as f:
            config = yaml.unsafe_load(f)

        instance = cls(
            theta=jnp.zeros((1, 2)),
            in_size=config["in_size"],
            hidden_size=config["hidden_size"],
            num_layers=config["num_layers"],
            num_coupling_layers=config["num_coupling_layers"],
            theta_ranges=config["theta_ranges"],
        )

        abstract_state = nnx.state(instance)
        checkpointer = orbax.PyTreeCheckpointer()
        state = checkpointer.restore(
            f"{path}/state", item=abstract_state, partial_restore=True
        )
        nnx.update(instance, state)

        instance.theta = metadata["theta"]
        instance.weights = metadata["weights"]
        instance.train_loss = metadata["train_loss"]
        instance.val_loss = metadata["val_loss"]
        instance.train_phi = metadata["train_phi"]
        instance.test_phi = metadata["test_phi"]
        instance.train_weights = metadata["train_weights"]
        instance.test_weights = metadata["test_weights"]
        instance.val_phi = metadata["val_phi"]
        instance.val_weights = metadata["val_weights"]

        shutil.rmtree(path)
        return instance

__call__(u)

Transform samples from the unit hypercube.

Parameters:

Name Type Description Default
u ndarray

Samples from the unit hypercube.

 required

Returns:

Type Description
ndarray

Log probabilities of the input data.
Source code in margarine/estimators/nice.py 
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
def __call__(self, u: jnp.ndarray) -> jnp.ndarray:
    """Transform samples from the unit hypercube.

    Args:
        u (jnp.ndarray): Samples from the unit hypercube.

    Returns:
        Log probabilities of the input data.
    """
    # from unit hypercube to standard normal space
    x = forward_transform(u, 0, 1)

    # from standard normal to target space
    x = self.inverse(x)

    # from gaussianized target space to original target space
    x = inverse_transform(x, self.theta_ranges[0], self.theta_ranges[1])
    return x

__init__(theta, weights=None, theta_ranges=None, in_size=2, hidden_size=128, num_layers=2, num_coupling_layers=4, nnx_rngs=None)

Initialize the NICE estimator.

Parameters:

Name Type Description Default
theta ndarray

Parameters of the density estimator.

 required
weights ndarray | None

Optional weights for the parameters.

 None
theta_ranges ndarray | None

Optional ranges for the parameters.

 None
in_size int

Input size.

 2
hidden_size int

Sizes of hidden layers.

 128
num_layers int

Number of layers in each coupling network.

 2
num_coupling_layers int

Number of coupling layers.

 4
nnx_rngs dict | None

Optional RNGs for Flax.

 None
Source code in margarine/estimators/nice.py 
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def __init__(
    self,
    theta: jnp.ndarray,
    weights: jnp.ndarray | None = None,
    theta_ranges: jnp.ndarray | None = None,
    in_size: int = 2,
    hidden_size: int = 128,
    num_layers: int = 2,
    num_coupling_layers: int = 4,
    nnx_rngs: dict | None = None,
) -> None:
    """Initialize the NICE estimator.

    Args:
        theta: Parameters of the density estimator.
        weights: Optional weights for the parameters.
        theta_ranges: Optional ranges for the parameters.
        in_size: Input size.
        hidden_size: Sizes of hidden layers.
        num_layers: Number of layers in each coupling network.
        num_coupling_layers: Number of coupling layers.
        nnx_rngs: Optional RNGs for Flax.
    """
    super().__init__()
    self.theta = nnx.data(theta)
    self.weights = nnx.data(weights)
    self.theta_ranges = nnx.data(theta_ranges)

    if self.weights is None:
        self.weights = jnp.ones(len(self.theta))
    if theta_ranges is None:
        self.theta_ranges = approximate_bounds(self.theta, self.weights)

    if nnx_rngs is None:
        nnx_rngs = nnx.Rngs(0)

    self.in_size = in_size
    self.hidden_size = hidden_size
    self.nlayers = num_layers
    self.num_coupling_layers = num_coupling_layers

    self.pass_size = in_size // 2
    self.net_in_size = in_size - self.pass_size

    layers = nnx.List()
    layers.append(
        nnx.Linear(self.net_in_size, self.hidden_size, rngs=nnx_rngs)
    )
    layers.append(lambda x: jax.nn.gelu(x))
    for _ in range(self.nlayers):
        layers.append(
            nnx.Linear(self.hidden_size, self.hidden_size, rngs=nnx_rngs)
        )
        layers.append(lambda x: jax.nn.gelu(x))

    layers.append(
        nnx.Linear(self.hidden_size, self.pass_size, rngs=nnx_rngs)
    )
    self.mlp = nnx.List(
        [nnx.Sequential(*layers) for _ in range(self.num_coupling_layers)]
    )

    self.S = nnx.Param(jnp.zeros(in_size))

forward(x)

NICE forward pass.

This is the forward pass of the NICE coupling layer from samples in the target space to the base distribution space.

Parameters:

Name Type Description Default
x ndarray

Input samples.

 required
Source code in margarine/estimators/nice.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
@jax.jit
def forward(self, x: jnp.ndarray) -> jnp.ndarray:
    """NICE forward pass.

    This is the forward pass of the NICE coupling layer from
    samples in the target space to the base distribution space.

    Args:
        x (jnp.ndarray): Input samples.
    """
    mlp_fns = tuple((lambda xb, mlp=mlp: mlp(xb)) for mlp in self.mlp)

    def body(i: int, x: jnp.ndarray) -> jnp.ndarray:
        """Body function for the coupling layers.

        Args:
            i: Forward counter.
            x: Input samples.

        Returns:
            Transformed samples.
        """
        x = jax.lax.cond(
            i % 2 == 0,
            lambda x: jnp.flip(x, axis=-1),
            lambda x: x,
            x,
        )
        xa, xb = x[:, : self.pass_size], x[:, self.pass_size :]
        m = jax.lax.switch(i, mlp_fns, xb)
        x = jnp.concatenate([xa + m, xb], axis=1)
        return x

    x = jax.lax.fori_loop(0, self.num_coupling_layers, body, x)

    x = jnp.exp(self.S) * x

    if self.num_coupling_layers % 2 == 1:
        x = jnp.flip(x, axis=-1)

    return x

inverse(x)

NICE inverse pass.

Parameters:

Name Type Description Default
x ndarray

Input samples.

 required
Source code in margarine/estimators/nice.py 
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
@jax.jit
def inverse(self, x: jnp.ndarray) -> jnp.ndarray:
    """NICE inverse pass.

    Args:
        x: Input samples.
    """
    if self.num_coupling_layers % 2 == 1:
        x = jnp.flip(x, axis=-1)

    x = x / jnp.exp(self.S)  # Undo scaling first

    mlp_fns = tuple((lambda xb, mlp=mlp: mlp(xb)) for mlp in self.mlp)

    def body(k: int, x: jnp.ndarray) -> jnp.ndarray:
        """Body function for the inverse coupling layers.

        Args:
            k: Forward counter.
            x: Input samples.

        Returns:
            Inverted samples.
        """
        # convert upward counter to downward counter
        i = self.num_coupling_layers - 1 - k

        xa, xb = x[:, : self.pass_size], x[:, self.pass_size :]
        m = jax.lax.switch(i, mlp_fns, xb)
        x = jnp.concatenate([xa - m, xb], axis=1)
        x = jax.lax.cond(
            i % 2 == 0,
            lambda x: jnp.flip(x, axis=-1),
            lambda x: x,
            x,
        )
        return x

    return jax.lax.fori_loop(0, self.num_coupling_layers, body, x)

load(filename) classmethod

Load a trained NICE model from a file.

Parameters:

Name Type Description Default
filename str

Path to the file from which the model will be loaded.

 required

Returns:

Type Description
BaseDensityEstimator | None

Loaded NICE model.
Source code in margarine/estimators/nice.py 
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
@classmethod
def load(cls, filename: str) -> BaseDensityEstimator | None:
    """Load a trained NICE model from a file.

    Args:
        filename (str): Path to the file from which the
            model will be loaded.

    Returns:
        Loaded NICE model.
    """
    zip_path = Path(f"{filename}.marg")
    path = Path(filename + ".tmp").resolve()
    with ZipFile(zip_path) as z:
        # Extract all files to a folder
        z.extractall(path)

    with open(f"{path}/metadata.yaml") as f:
        metadata = yaml.unsafe_load(f)

    version = metadata.get("margarine_version", None)
    if version is None:
        print(
            "Warning: The KDE was saved with a version of margarine ",
            " < 2.0.0. In order to load it you will need to downgrade ",
            "margarine to a version < 2.0.0. e.g. ",
            "pip install margarine<2.0.0",
        )
        return
    if version != _version.__version__:
        print(
            f"Warning: The KDE was saved with margarine version "
            f"{version}, but you are loading it with version "
            f"{_version.__version__}. This may lead to "
            f"incompatibilities."
        )

    with open(f"{path}/config.yaml") as f:
        config = yaml.unsafe_load(f)

    instance = cls(
        theta=jnp.zeros((1, 2)),
        in_size=config["in_size"],
        hidden_size=config["hidden_size"],
        num_layers=config["num_layers"],
        num_coupling_layers=config["num_coupling_layers"],
        theta_ranges=config["theta_ranges"],
    )

    abstract_state = nnx.state(instance)
    checkpointer = orbax.PyTreeCheckpointer()
    state = checkpointer.restore(
        f"{path}/state", item=abstract_state, partial_restore=True
    )
    nnx.update(instance, state)

    instance.theta = metadata["theta"]
    instance.weights = metadata["weights"]
    instance.train_loss = metadata["train_loss"]
    instance.val_loss = metadata["val_loss"]
    instance.train_phi = metadata["train_phi"]
    instance.test_phi = metadata["test_phi"]
    instance.train_weights = metadata["train_weights"]
    instance.test_weights = metadata["test_weights"]
    instance.val_phi = metadata["val_phi"]
    instance.val_weights = metadata["val_weights"]

    shutil.rmtree(path)
    return instance

log_like(x, logevidence, prior_density)

Compute the marginal log-likelihood of given samples.

Parameters:

Name Type Description Default
x ndarray

Samples for which to compute the log-likelihood.

 required
logevidence float

Log-evidence term.

 required
prior_density ndarray | BaseDensityEstimator

Prior density or density estimator.

 required

Returns:

Type Description
ndarray

Log likelihoods of the input data.
Source code in margarine/estimators/nice.py 
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
def log_like(
    self,
    x: jnp.ndarray,
    logevidence: float,
    prior_density: jnp.ndarray | BaseDensityEstimator,
) -> jnp.ndarray:
    """Compute the marginal log-likelihood of given samples.

    Args:
        x: Samples for which to compute the log-likelihood.
        logevidence: Log-evidence term.
        prior_density: Prior density or density estimator.

    Returns:
        Log likelihoods of the input data.
    """
    if isinstance(prior_density, BaseDensityEstimator):
        prior_density = prior_density.log_prob(x)

    return self.log_prob(x) + logevidence - prior_density

log_prob(x)

Compute the log probability of the input data.

Parameters:

Name Type Description Default
x ndarray

Input data.

 required

Returns:

Type Description
ndarray

Log probabilities of the input data.
Source code in margarine/estimators/nice.py 
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
@jax.jit
def log_prob(self, x: jnp.ndarray) -> jnp.ndarray:
    """Compute the log probability of the input data.

    Args:
        x (jnp.ndarray): Input data.

    Returns:
        Log probabilities of the input data.
    """
    transformed_x = forward_transform(
        x, self.theta_ranges[0], self.theta_ranges[1]
    )
    transform_chain = tfb.Chain(
        [
            tfb.Invert(tfb.NormalCDF()),
            tfb.Scale(1 / (self.theta_ranges[1] - self.theta_ranges[0])),
            tfb.Shift(-self.theta_ranges[0]),
        ]
    )

    def norm_jac(y: jnp.ndarray) -> jnp.ndarray:
        """Calculate the normalising jacobian for the transformation."""
        return transform_chain.inverse_log_det_jacobian(y, event_ndims=0)

    correction = -norm_jac(transformed_x).sum(axis=1)

    return self.log_prob_under_NICE(transformed_x) + correction

log_prob_under_NICE(x)

Compute the log probability under the NICE model.

Parameters:

Name Type Description Default
x ndarray

Input data.

 required

Returns:

Type Description
ndarray

Log probabilities of the input data.
Source code in margarine/estimators/nice.py 
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
@jax.jit
def log_prob_under_NICE(self, x: jnp.ndarray) -> jnp.ndarray:
    """Compute the log probability under the NICE model.

    Args:
        x (jnp.ndarray): Input data.

    Returns:
        Log probabilities of the input data.
    """
    # calculate the actual log prob under the NICE model
    # assuming a standard normal base distribution
    z = self.forward(x)
    log_pz = -0.5 * jnp.sum(z**2 + jnp.log(2 * jnp.pi), axis=1)
    log_det_J = jnp.sum(self.S)
    log_prob_under_nice = log_pz + log_det_J
    return log_prob_under_nice

sample(key, num_samples)

Sample from the trained NICE model.

Parameters:

Name Type Description Default
key ndarray

JAX random key.

 required
num_samples int

Number of samples to draw.

 required

Returns:

Type Description
ndarray

Samples drawn from the NICE model.
Source code in margarine/estimators/nice.py 
324
325
326
327
328
329
330
331
332
333
334
335
336
337
def sample(self, key: jnp.ndarray, num_samples: int) -> jnp.ndarray:
    """Sample from the trained NICE model.

    Args:
        key: JAX random key.
        num_samples: Number of samples to draw.

    Returns:
        Samples drawn from the NICE model.
    """
    u = jax.random.uniform(key, shape=(num_samples, self.theta.shape[1]))

    samples = self(u)
    return samples

save(filename)

Save the trained NICE model to a file.

Parameters:

Name Type Description Default
filename str

Path to the file where the model will be saved.

 required
Source code in margarine/estimators/nice.py 
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
def save(self, filename: str) -> None:
    """Save the trained NICE model to a file.

    Args:
        filename: Path to the file where the model will be saved.
    """
    path = Path(filename).resolve()
    if path.exists():
        shutil.rmtree(path)

    state = nnx.state(self)

    checkpointer = orbax.PyTreeCheckpointer()
    checkpointer.save(f"{path}/state", state)

    config = {
        "in_size": self.in_size,
        "hidden_size": self.hidden_size,
        "num_layers": self.nlayers,
        "num_coupling_layers": self.num_coupling_layers,
        "theta_ranges": self.theta_ranges,
    }
    with open(f"{path}/config.yaml", "w") as f:
        yaml.dump(config, f)

    metadata = {
        "theta": self.theta,
        "weights": self.weights,
        "train_loss": self.train_loss,
        "val_loss": self.val_loss,
        "train_phi": self.train_phi,
        "test_phi": self.test_phi,
        "train_weights": self.train_weights,
        "test_weights": self.test_weights,
        "val_phi": self.val_phi,
        "val_weights": self.val_weights,
        "margarine_version": _version.__version__,
    }

    with open(f"{path}/metadata.yaml", "w") as f:
        yaml.dump(metadata, f)

    with ZipFile(filename + ".marg", "w") as z:
        for subpath in path.rglob("*"):
            if subpath.is_file():
                z.write(subpath, arcname=subpath.relative_to(path))

    shutil.rmtree(path)

train(key, learning_rate=0.0001, epochs=1000, patience=50, batch_size=256)

Train the NICE model.

Parameters:

Name Type Description Default
key ndarray

JAX random key for data splitting.

 required
learning_rate float

Learning rate for the optimizer.

 0.0001
epochs int

Number of training epochs.

 1000
patience int

Patience for early stopping.

 50
batch_size int

Batch size for training.

 256
Source code in margarine/estimators/nice.py 
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
def train(
    self,
    key: jnp.ndarray,
    learning_rate: float = 1e-4,
    epochs: int = 1000,
    patience: int = 50,
    batch_size: int = 256,
) -> None:
    """Train the NICE model.

    Args:
        key: JAX random key for data splitting.
        learning_rate: Learning rate for the optimizer.
        epochs: Number of training epochs.
        patience: Patience for early stopping.
        batch_size: Batch size for training.
    """

    @nnx.jit
    def loss_function(
        model: "NICE",
        targets: jnp.ndarray,
        weights: jnp.ndarray,
    ) -> jnp.ndarray:
        """Loss function for training NICE model.

        Args:
            model (NICE): NICE model.
            targets (jnp.ndarray): Samples from the target distribution
                that have been passed through the forward_transform.
            weights (jnp.ndarray): Weights for the samples.

        Returns:
            jnp.ndarray: Computed loss.
        """
        return -jnp.mean(weights * model.log_prob_under_NICE(targets))

    @nnx.jit
    def train_step(
        model: NICE,
        optimizer: nnx.Optimizer,
        train_phi: jnp.ndarray,
        train_weights: jnp.ndarray,
    ) -> None:
        """Single training step for NICE model.

        Args:
            model (NICE): NICE model.
            optimizer (nnx.Optimizer): Optimizer for
                updating model parameters.
            train_phi (jnp.ndarray): Training samples.
            train_weights (jnp.ndarray): Weights for the training samples.
        """
        grad = nnx.grad(loss_function)(model, train_phi, train_weights)
        optimizer.update(model, grad)

    phi = forward_transform(
        self.theta, self.theta_ranges[0], self.theta_ranges[1]
    )
    weights = self.weights / jnp.sum(self.weights)

    key, subkey = jax.random.split(key)
    # need to split the data into training and validation sets
    (
        self.train_phi,
        self.test_phi,
        self.train_weights,
        self.test_weights,
    ) = train_test_split(
        phi,
        weights,
        subkey,
        test_size=0.2,
    )
    key, subkey = jax.random.split(key)
    self.test_phi, self.val_phi, self.test_weights, self.val_weights = (
        train_test_split(
            self.test_phi, self.test_weights, subkey, test_size=0.5
        )
    )

    tx = optax.adam(learning_rate)
    optimizer = nnx.Optimizer(self, tx, wrt=nnx.Param)

    pbar = tqdm.tqdm(range(epochs), desc="Training")

    best_loss = jnp.inf
    best_model = nnx.state(self, nnx.Param)
    c = 0

    data_size = len(self.train_phi)

    tl, vl = [], []
    for _ in pbar:
        data_permutations = jax.random.permutation(
            subkey, jnp.arange(data_size)
        )
        accumulated_train_loss = 0.0
        for i in range(0, len(self.train_phi), batch_size):
            batch_indices = data_permutations[i : i + batch_size]
            batch_phi = self.train_phi[batch_indices]
            batch_weights = self.train_weights[batch_indices]
            loss = loss_function(self, batch_phi, batch_weights)
            train_step(self, optimizer, batch_phi, batch_weights)
            accumulated_train_loss += loss * len(batch_phi)
        loss = accumulated_train_loss / len(self.train_phi)
        tl.append(loss)

        val_loss = loss_function(self, self.val_phi, self.val_weights)
        vl.append(val_loss)

        if val_loss < best_loss:
            best_loss = val_loss
            best_model = nnx.state(self, nnx.Param)
            best_epoch = _
            c = 0
        else:
            c += 1
            if c >= patience:
                print(
                    f"Early stopping at epoch {_}, best "
                    + f"epoch was {best_epoch} with val loss {best_loss}"
                )
                break

        pbar.set_postfix(
            loss=f"{loss:.3e}",
            val_loss=f"{val_loss:.3e}",
            best_loss=f"{best_loss:.3e}",
        )

    self.train_loss = jnp.array(tl)
    self.val_loss = jnp.array(vl)

    if best_model is not None:
        nnx.update(self, best_model)

margarine.estimators.realnvp

Implementation of the RealNVP estimator.

RealNVP

Bases: BaseDensityEstimator, Module

Implementation of the RealNVP architecture for density estimation.

Details in https://arxiv.org/abs/1605.08803.

Source code in margarine/estimators/realnvp.py 
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
class RealNVP(BaseDensityEstimator, nnx.Module):
    """Implementation of the RealNVP architecture for density estimation.

    Details in https://arxiv.org/abs/1605.08803.
    """

    def __init__(
        self,
        theta: jnp.ndarray,
        weights: jnp.ndarray | None = None,
        theta_ranges: jnp.ndarray | None = None,
        in_size: int = 2,
        hidden_size: int = 128,
        num_layers: int = 2,
        num_coupling_layers: int = 4,
        nnx_rngs: dict | None = None,
        permutations_key: jnp.ndarray = jax.random.PRNGKey(0),
    ) -> None:
        """Initialize the RealNVP estimator.

        Args:
            theta: Parameters of the density estimator.
            weights: Optional weights for the parameters.
            theta_ranges: Optional ranges for the parameters.
            in_size: Input size.
            hidden_size: Size of hidden layers.
            num_layers: Number of layers in each coupling network.
            num_coupling_layers: Number of coupling layers.
            nnx_rngs: Optional RNGs for Flax.
            permutations_key: JAX random key for permutations.
        """
        super().__init__()
        self.theta = nnx.data(theta)
        self.weights = nnx.data(weights)
        self.theta_ranges = nnx.data(theta_ranges)

        if self.weights is None:
            self.weights = jnp.ones(len(self.theta))
        if theta_ranges is None:
            self.theta_ranges = approximate_bounds(self.theta, self.weights)

        if nnx_rngs is None:
            nnx_rngs = nnx.Rngs(0)

        self.in_size = in_size
        self.hidden_size = hidden_size
        self.nlayers = num_layers
        self.num_coupling_layers = num_coupling_layers

        self.pass_size = in_size // 2
        self.net_in_size = in_size - self.pass_size

        # like glorot normal but with 0.1 stddev
        kernel_init = nnx.initializers.variance_scaling(
            0.1, "fan_avg", "truncated_normal"
        )

        additive_layers = nnx.List()
        additive_layers.append(
            nnx.Linear(
                self.net_in_size,
                self.hidden_size,
                kernel_init=kernel_init,
                rngs=nnx_rngs,
            )
        )
        additive_layers.append(lambda x: jax.nn.relu(x))
        for _ in range(self.nlayers):
            additive_layers.append(
                nnx.Linear(
                    self.hidden_size,
                    self.hidden_size,
                    kernel_init=kernel_init,
                    rngs=nnx_rngs,
                )
            )
            additive_layers.append(lambda x: jax.nn.relu(x))
        additive_layers.append(
            nnx.Linear(
                self.hidden_size,
                self.pass_size,
                kernel_init=kernel_init,
                rngs=nnx_rngs,
            )
        )
        self.additive_mlp = nnx.List(
            [
                nnx.Sequential(*additive_layers)
                for _ in range(self.num_coupling_layers)
            ]
        )

        scaling_layers = nnx.List()
        scaling_layers.append(
            nnx.Linear(
                self.net_in_size,
                self.hidden_size,
                kernel_init=kernel_init,
                rngs=nnx_rngs,
            )
        )
        scaling_layers.append(lambda x: jax.nn.relu(x))
        for _ in range(self.nlayers):
            scaling_layers.append(
                nnx.Linear(
                    self.hidden_size,
                    self.hidden_size,
                    kernel_init=kernel_init,
                    rngs=nnx_rngs,
                )
            )
            scaling_layers.append(lambda x: jax.nn.relu(x))
        scaling_layers.append(
            nnx.Linear(
                self.hidden_size,
                self.pass_size,
                kernel_init=kernel_init,
                rngs=nnx_rngs,
            )
        )

        self.scaling_mlp = nnx.List(
            [
                nnx.Sequential(*scaling_layers)
                for _ in range(self.num_coupling_layers)
            ]
        )

        self.permutations_key = permutations_key
        key = self.permutations_key
        self.permutations = []
        for _ in range(self.num_coupling_layers):
            key, subkey = jax.random.split(key)
            perm = jax.random.permutation(subkey, in_size)
            self.permutations.append(perm)
        self.permutations = nnx.data(self.permutations)
        self.inverse_permutations = nnx.data(
            [jnp.argsort(p) for p in self.permutations]
        )

    def forward(
        self, x: jnp.ndarray, return_log_det: bool = False
    ) -> jnp.ndarray | tuple[jnp.ndarray, jnp.ndarray]:
        """Forward pass of the RealNVP coupling layer.

        Args:
            x (jnp.ndarray): Input data.
            return_log_det (bool): Whether to return the log-determinant
                of the Jacobian.

        Returns:
            jnp.ndarray | tuple[jnp.ndarray, jnp.ndarray] :
                Transformed data, and optionally the
                log-determinant of the Jacobian.
        """
        log_det = 0.0
        for i in range(self.num_coupling_layers):
            x = x[:, self.permutations[i]]
            xa, xb = x[:, : self.pass_size], x[:, self.pass_size :]
            m = self.additive_mlp[i](xb)
            s = self.scaling_mlp[i](xb)
            xa = xa * jnp.exp(s) + m
            x = jnp.concat([xa, xb], axis=1)  # Concatenate correctly
            log_det += jnp.sum(s, axis=1)  # Sum over the scaling parameters

        if return_log_det:
            return x, log_det
        else:
            return x

    def inverse(self, x: jnp.ndarray) -> jnp.ndarray:
        """Inverse pass of the RealNVP coupling layer.

        From the base distribution to the target distribution.

        Args:
            x (jnp.ndarray): Samples from the base distribution.

        Returns:
            jnp.ndarray: Transformed samples in the target distribution.
        """
        for i in reversed(
            range(self.num_coupling_layers)
        ):  # Reverse the coupling order
            xa, xb = x[:, : self.pass_size], x[:, self.pass_size :]
            m = self.additive_mlp[i](xb)
            s = self.scaling_mlp[i](xb)
            xa = (xa - m) / jnp.exp(s)
            x = jnp.concat([xa, xb], axis=1)  # Concatenate correctly
            x = x[:, self.inverse_permutations[i]]

        return x

    def train(
        self,
        key: jnp.ndarray,
        learning_rate: float = 1e-4,
        epochs: int = 1000,
        patience: int = 50,
        batch_size: int = 256,
    ) -> None:
        """Train the RealNVP model.

        Args:
            key: JAX random key for data splitting.
            learning_rate: Learning rate for the optimizer.
            epochs: Number of training epochs.
            patience: Patience for early stopping.
            batch_size: Size of training batches.
        """

        @jax.jit
        def loss_function(
            model: RealNVP,
            targets: jnp.ndarray,
            weights: jnp.ndarray,
        ) -> jnp.ndarray:
            """Loss function for training RealNVP model.

            Args:
                model (RealNVP): The RealNVP model.
                targets (jnp.ndarray): Target samples.
                weights (jnp.ndarray): Weights for the samples.

            Returns:
                jnp.ndarray: Computed loss.
            """
            return -jnp.mean(weights * model.log_prob_under_RealNVP(targets))

        # @nnx.jit
        def train_step(
            model: RealNVP,
            optimizer: nnx.Optimizer,
            train_phi: jnp.ndarray,
            train_weights: jnp.ndarray,
        ) -> None:
            """Single training step for NICE model.

            Args:
                model (NICE): The NICE model to be trained.
                optimizer (nnx.Optimizer): The optimizer to
                    update model parameters.
                train_phi (jnp.ndarray): Training samples.
                train_weights (jnp.ndarray): Weights for the training samples.
            """
            grad = nnx.grad(loss_function)(model, train_phi, train_weights)
            optimizer.update(model, grad)

        phi = forward_transform(
            self.theta, self.theta_ranges[0], self.theta_ranges[1]
        )
        weights = self.weights / jnp.sum(self.weights)

        key, subkey = jax.random.split(key)
        # need to split the data into training and validation sets
        (
            self.train_phi,
            self.test_phi,
            self.train_weights,
            self.test_weights,
        ) = train_test_split(
            phi,
            weights,
            subkey,
            test_size=0.2,
        )
        key, subkey = jax.random.split(key)
        self.test_phi, self.val_phi, self.test_weights, self.val_weights = (
            train_test_split(
                self.test_phi,
                self.test_weights,
                subkey,
                test_size=0.5,
            )
        )

        tx = optax.adam(learning_rate)
        optimizer = nnx.Optimizer(self, tx, wrt=nnx.Param)

        pbar = tqdm.tqdm(range(epochs), desc="Training")

        best_loss = jnp.inf
        best_model = nnx.state(self, nnx.Param)
        c = 0

        data_size = len(self.train_phi)

        tl, vl = [], []
        for _ in pbar:
            data_permutations = jax.random.permutation(
                subkey, jnp.arange(data_size)
            )
            accumulated_train_loss = 0.0
            for i in range(0, len(self.train_phi), batch_size):
                batch_indices = data_permutations[i : i + batch_size]
                batch_phi = self.train_phi[batch_indices]
                batch_weights = self.train_weights[batch_indices]
                loss = loss_function(self, batch_phi, batch_weights)
                train_step(self, optimizer, batch_phi, batch_weights)
                accumulated_train_loss += loss * len(batch_phi)
            loss = accumulated_train_loss / len(self.train_phi)
            tl.append(loss)

            val_loss = loss_function(self, self.val_phi, self.val_weights)
            vl.append(val_loss)

            if val_loss < best_loss:
                best_loss = val_loss
                best_model = nnx.state(self, nnx.Param)
                best_epoch = _
                c = 0
            else:
                c += 1
                if c >= patience:
                    print(
                        f"Early stopping at epoch {_}, best "
                        + f"epoch was {best_epoch} with val loss {best_loss}"
                    )
                    break

            pbar.set_postfix(
                loss=f"{loss:.3e}",
                val_loss=f"{val_loss:.3e}",
                best_loss=f"{best_loss:.3e}",
            )

        self.train_loss = jnp.array(tl)
        self.val_loss = jnp.array(vl)

        if best_model is not None:
            nnx.update(self, best_model)

    def sample(self, key: jnp.ndarray, num_samples: int) -> jnp.ndarray:
        """Sample from the trained RealNVP model.

        Args:
            key: JAX random key.
            num_samples: Number of samples to draw.

        Returns:
            Samples drawn from the RealNVP model.
        """
        u = jax.random.uniform(key, shape=(num_samples, self.theta.shape[1]))
        samples = self(u)
        return samples

    def __call__(self, u: jnp.ndarray) -> jnp.ndarray:
        """Transform samples from the unit hypercube.

        Args:
            u (jnp.ndarray): Samples from the unit hypercube.

        Returns:
            Log probabilities of the input data.
        """
        # from unit hypercube to standard normal space
        x = forward_transform(u, 0, 1)

        # from standard normal to target space
        x = self.inverse(x)

        # from gaussianized target space to original target space
        x = inverse_transform(x, self.theta_ranges[0], self.theta_ranges[1])
        return x

    @jax.jit
    def log_prob_under_RealNVP(self, x: jnp.ndarray) -> jnp.ndarray:
        """Compute the log probability under the RealNVP model.

        Args:
            x (jnp.ndarray): Input data.

        Returns:
            Log probabilities of the input data.
        """
        # calcualte the actual log prob under the RealNVP model
        # assuming a standard normal base distribution
        z, log_det_J = self.forward(x, return_log_det=True)
        log_pz = -0.5 * jnp.sum(z**2 + jnp.log(2 * jnp.pi), axis=1)
        log_prob_under_realnvp = log_pz + log_det_J
        return log_prob_under_realnvp

    @jax.jit
    def log_prob(self, x: jnp.ndarray) -> jnp.ndarray:
        """Compute the log probability of the input data.

        Args:
            x (jnp.ndarray): Input data.

        Returns:
            Log probabilities of the input data.
        """
        transformed_x = forward_transform(
            x, self.theta_ranges[0], self.theta_ranges[1]
        )
        transform_chain = tfb.Chain(
            [
                tfb.Invert(tfb.NormalCDF()),
                tfb.Scale(1 / (self.theta_ranges[1] - self.theta_ranges[0])),
                tfb.Shift(-self.theta_ranges[0]),
            ]
        )

        def norm_jac(y: jnp.ndarray) -> jnp.ndarray:
            """Calculate the normalising jacobian for the transformation."""
            return transform_chain.inverse_log_det_jacobian(y, event_ndims=0)

        correction = -norm_jac(transformed_x).sum(axis=1)

        return self.log_prob_under_RealNVP(transformed_x) + correction

    def log_like(
        self,
        x: jnp.ndarray,
        logevidence: float,
        prior_density: jnp.ndarray | BaseDensityEstimator,
    ) -> jnp.ndarray:
        """Compute the marginal log-likelihood of given samples.

        Args:
            x: Samples for which to compute the log-likelihood.
            logevidence: Log-evidence term.
            prior_density: Prior density or density estimator.

        Returns:
            Log likelihoods of the input data.
        """
        if isinstance(prior_density, BaseDensityEstimator):
            prior_density = prior_density.log_prob(x)

        return self.log_prob(x) + logevidence - prior_density

    def save(self, filename: str) -> None:
        """Save the trained RealNVP model to a file.

        Args:
            filename: Path to the file where the model will be saved.
        """
        path = Path(filename + ".tmp").resolve()
        if path.exists():
            shutil.rmtree(path)

        state = nnx.state(self)

        checkpointer = orbax.PyTreeCheckpointer()
        checkpointer.save(f"{path}/state", state)

        config = {
            "in_size": self.in_size,
            "hidden_size": self.hidden_size,
            "num_layers": self.nlayers,
            "num_coupling_layers": self.num_coupling_layers,
            "key": self.permutations_key.tolist(),
            "theta_ranges": self.theta_ranges,
        }
        with open(f"{path}/config.yaml", "w") as f:
            yaml.dump(config, f)

        metadata = {
            "theta": self.theta,
            "weights": self.weights,
            "train_loss": self.train_loss,
            "val_loss": self.val_loss,
            "train_phi": self.train_phi,
            "test_phi": self.test_phi,
            "train_weights": self.train_weights,
            "test_weights": self.test_weights,
            "val_phi": self.val_phi,
            "val_weights": self.val_weights,
            "margarine_version": _version.__version__,
        }

        with open(f"{path}/metadata.yaml", "w") as f:
            yaml.dump(metadata, f)

        with ZipFile(filename + ".marg", "w") as z:
            for subpath in path.rglob("*"):
                if subpath.is_file():
                    z.write(subpath, arcname=subpath.relative_to(path))

        shutil.rmtree(path)

    @classmethod
    def load(cls, filename: str) -> BaseDensityEstimator | None:
        """Load a trained RealNVP model from a file.

        Args:
            filename: Path to the file from which the model will be loaded.

        Returns:
            Loaded RealNVP model.
        """
        zip_path = Path(f"{filename}.marg")
        path = Path(filename + ".tmp").resolve()
        with ZipFile(zip_path) as z:
            # Extract all files to a folder
            z.extractall(path)

        with open(f"{path}/metadata.yaml") as f:
            metadata = yaml.unsafe_load(f)

        version = metadata.get("margarine_version", None)
        if version is None:
            print(
                "Warning: The KDE was saved with a version of margarine ",
                " < 2.0.0. In order to load it you will need to downgrade ",
                "margarine to a version < 2.0.0. e.g. ",
                "pip install margarine<2.0.0",
            )
            return
        if version != _version.__version__:
            print(
                f"Warning: The KDE was saved with margarine version "
                f"{version}, but you are loading it with version "
                f"{_version.__version__}. This may lead to "
                f"incompatibilities."
            )

        with open(f"{path}/config.yaml") as f:
            config = yaml.unsafe_load(f)

        instance = cls(
            theta=jnp.zeros((1, 2)),
            in_size=config["in_size"],
            hidden_size=config["hidden_size"],
            num_layers=config["num_layers"],
            num_coupling_layers=config["num_coupling_layers"],
            permutations_key=jnp.array(config["key"], dtype=jnp.uint32),
            theta_ranges=config["theta_ranges"],
        )

        abstract_state = nnx.state(instance)
        checkpointer = orbax.PyTreeCheckpointer()
        state = checkpointer.restore(
            f"{path}/state", item=abstract_state, partial_restore=True
        )
        nnx.update(instance, state)

        instance.theta = metadata["theta"]
        instance.weights = metadata["weights"]
        instance.train_loss = metadata["train_loss"]
        instance.val_loss = metadata["val_loss"]
        instance.train_phi = metadata["train_phi"]
        instance.test_phi = metadata["test_phi"]
        instance.train_weights = metadata["train_weights"]
        instance.test_weights = metadata["test_weights"]
        instance.val_phi = metadata["val_phi"]
        instance.val_weights = metadata["val_weights"]

        shutil.rmtree(path)
        return instance

__call__(u)

Transform samples from the unit hypercube.

Parameters:

Name Type Description Default
u ndarray

Samples from the unit hypercube.

 required

Returns:

Type Description
ndarray

Log probabilities of the input data.
Source code in margarine/estimators/realnvp.py 
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
def __call__(self, u: jnp.ndarray) -> jnp.ndarray:
    """Transform samples from the unit hypercube.

    Args:
        u (jnp.ndarray): Samples from the unit hypercube.

    Returns:
        Log probabilities of the input data.
    """
    # from unit hypercube to standard normal space
    x = forward_transform(u, 0, 1)

    # from standard normal to target space
    x = self.inverse(x)

    # from gaussianized target space to original target space
    x = inverse_transform(x, self.theta_ranges[0], self.theta_ranges[1])
    return x

__init__(theta, weights=None, theta_ranges=None, in_size=2, hidden_size=128, num_layers=2, num_coupling_layers=4, nnx_rngs=None, permutations_key=jax.random.PRNGKey(0))

Initialize the RealNVP estimator.

Parameters:

Name Type Description Default
theta ndarray

Parameters of the density estimator.

 required
weights ndarray | None

Optional weights for the parameters.

 None
theta_ranges ndarray | None

Optional ranges for the parameters.

 None
in_size int

Input size.

 2
hidden_size int

Size of hidden layers.

 128
num_layers int

Number of layers in each coupling network.

 2
num_coupling_layers int

Number of coupling layers.

 4
nnx_rngs dict | None

Optional RNGs for Flax.

 None
permutations_key ndarray

JAX random key for permutations.

 PRNGKey(0)
Source code in margarine/estimators/realnvp.py 
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def __init__(
    self,
    theta: jnp.ndarray,
    weights: jnp.ndarray | None = None,
    theta_ranges: jnp.ndarray | None = None,
    in_size: int = 2,
    hidden_size: int = 128,
    num_layers: int = 2,
    num_coupling_layers: int = 4,
    nnx_rngs: dict | None = None,
    permutations_key: jnp.ndarray = jax.random.PRNGKey(0),
) -> None:
    """Initialize the RealNVP estimator.

    Args:
        theta: Parameters of the density estimator.
        weights: Optional weights for the parameters.
        theta_ranges: Optional ranges for the parameters.
        in_size: Input size.
        hidden_size: Size of hidden layers.
        num_layers: Number of layers in each coupling network.
        num_coupling_layers: Number of coupling layers.
        nnx_rngs: Optional RNGs for Flax.
        permutations_key: JAX random key for permutations.
    """
    super().__init__()
    self.theta = nnx.data(theta)
    self.weights = nnx.data(weights)
    self.theta_ranges = nnx.data(theta_ranges)

    if self.weights is None:
        self.weights = jnp.ones(len(self.theta))
    if theta_ranges is None:
        self.theta_ranges = approximate_bounds(self.theta, self.weights)

    if nnx_rngs is None:
        nnx_rngs = nnx.Rngs(0)

    self.in_size = in_size
    self.hidden_size = hidden_size
    self.nlayers = num_layers
    self.num_coupling_layers = num_coupling_layers

    self.pass_size = in_size // 2
    self.net_in_size = in_size - self.pass_size

    # like glorot normal but with 0.1 stddev
    kernel_init = nnx.initializers.variance_scaling(
        0.1, "fan_avg", "truncated_normal"
    )

    additive_layers = nnx.List()
    additive_layers.append(
        nnx.Linear(
            self.net_in_size,
            self.hidden_size,
            kernel_init=kernel_init,
            rngs=nnx_rngs,
        )
    )
    additive_layers.append(lambda x: jax.nn.relu(x))
    for _ in range(self.nlayers):
        additive_layers.append(
            nnx.Linear(
                self.hidden_size,
                self.hidden_size,
                kernel_init=kernel_init,
                rngs=nnx_rngs,
            )
        )
        additive_layers.append(lambda x: jax.nn.relu(x))
    additive_layers.append(
        nnx.Linear(
            self.hidden_size,
            self.pass_size,
            kernel_init=kernel_init,
            rngs=nnx_rngs,
        )
    )
    self.additive_mlp = nnx.List(
        [
            nnx.Sequential(*additive_layers)
            for _ in range(self.num_coupling_layers)
        ]
    )

    scaling_layers = nnx.List()
    scaling_layers.append(
        nnx.Linear(
            self.net_in_size,
            self.hidden_size,
            kernel_init=kernel_init,
            rngs=nnx_rngs,
        )
    )
    scaling_layers.append(lambda x: jax.nn.relu(x))
    for _ in range(self.nlayers):
        scaling_layers.append(
            nnx.Linear(
                self.hidden_size,
                self.hidden_size,
                kernel_init=kernel_init,
                rngs=nnx_rngs,
            )
        )
        scaling_layers.append(lambda x: jax.nn.relu(x))
    scaling_layers.append(
        nnx.Linear(
            self.hidden_size,
            self.pass_size,
            kernel_init=kernel_init,
            rngs=nnx_rngs,
        )
    )

    self.scaling_mlp = nnx.List(
        [
            nnx.Sequential(*scaling_layers)
            for _ in range(self.num_coupling_layers)
        ]
    )

    self.permutations_key = permutations_key
    key = self.permutations_key
    self.permutations = []
    for _ in range(self.num_coupling_layers):
        key, subkey = jax.random.split(key)
        perm = jax.random.permutation(subkey, in_size)
        self.permutations.append(perm)
    self.permutations = nnx.data(self.permutations)
    self.inverse_permutations = nnx.data(
        [jnp.argsort(p) for p in self.permutations]
    )

forward(x, return_log_det=False)

Forward pass of the RealNVP coupling layer.

Parameters:

Name Type Description Default
x ndarray

Input data.

 required
return_log_det bool

Whether to return the log-determinant of the Jacobian.

 False

Returns:

Type Description
ndarray | tuple[ndarray, ndarray]

jnp.ndarray | tuple[jnp.ndarray, jnp.ndarray] : Transformed data, and optionally the log-determinant of the Jacobian.
Source code in margarine/estimators/realnvp.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
def forward(
    self, x: jnp.ndarray, return_log_det: bool = False
) -> jnp.ndarray | tuple[jnp.ndarray, jnp.ndarray]:
    """Forward pass of the RealNVP coupling layer.

    Args:
        x (jnp.ndarray): Input data.
        return_log_det (bool): Whether to return the log-determinant
            of the Jacobian.

    Returns:
        jnp.ndarray | tuple[jnp.ndarray, jnp.ndarray] :
            Transformed data, and optionally the
            log-determinant of the Jacobian.
    """
    log_det = 0.0
    for i in range(self.num_coupling_layers):
        x = x[:, self.permutations[i]]
        xa, xb = x[:, : self.pass_size], x[:, self.pass_size :]
        m = self.additive_mlp[i](xb)
        s = self.scaling_mlp[i](xb)
        xa = xa * jnp.exp(s) + m
        x = jnp.concat([xa, xb], axis=1)  # Concatenate correctly
        log_det += jnp.sum(s, axis=1)  # Sum over the scaling parameters

    if return_log_det:
        return x, log_det
    else:
        return x

inverse(x)

Inverse pass of the RealNVP coupling layer.

From the base distribution to the target distribution.

Parameters:

Name Type Description Default
x ndarray

Samples from the base distribution.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Transformed samples in the target distribution.
Source code in margarine/estimators/realnvp.py 
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
def inverse(self, x: jnp.ndarray) -> jnp.ndarray:
    """Inverse pass of the RealNVP coupling layer.

    From the base distribution to the target distribution.

    Args:
        x (jnp.ndarray): Samples from the base distribution.

    Returns:
        jnp.ndarray: Transformed samples in the target distribution.
    """
    for i in reversed(
        range(self.num_coupling_layers)
    ):  # Reverse the coupling order
        xa, xb = x[:, : self.pass_size], x[:, self.pass_size :]
        m = self.additive_mlp[i](xb)
        s = self.scaling_mlp[i](xb)
        xa = (xa - m) / jnp.exp(s)
        x = jnp.concat([xa, xb], axis=1)  # Concatenate correctly
        x = x[:, self.inverse_permutations[i]]

    return x

load(filename) classmethod

Load a trained RealNVP model from a file.

Parameters:

Name Type Description Default
filename str

Path to the file from which the model will be loaded.

 required

Returns:

Type Description
BaseDensityEstimator | None

Loaded RealNVP model.
Source code in margarine/estimators/realnvp.py 
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
@classmethod
def load(cls, filename: str) -> BaseDensityEstimator | None:
    """Load a trained RealNVP model from a file.

    Args:
        filename: Path to the file from which the model will be loaded.

    Returns:
        Loaded RealNVP model.
    """
    zip_path = Path(f"{filename}.marg")
    path = Path(filename + ".tmp").resolve()
    with ZipFile(zip_path) as z:
        # Extract all files to a folder
        z.extractall(path)

    with open(f"{path}/metadata.yaml") as f:
        metadata = yaml.unsafe_load(f)

    version = metadata.get("margarine_version", None)
    if version is None:
        print(
            "Warning: The KDE was saved with a version of margarine ",
            " < 2.0.0. In order to load it you will need to downgrade ",
            "margarine to a version < 2.0.0. e.g. ",
            "pip install margarine<2.0.0",
        )
        return
    if version != _version.__version__:
        print(
            f"Warning: The KDE was saved with margarine version "
            f"{version}, but you are loading it with version "
            f"{_version.__version__}. This may lead to "
            f"incompatibilities."
        )

    with open(f"{path}/config.yaml") as f:
        config = yaml.unsafe_load(f)

    instance = cls(
        theta=jnp.zeros((1, 2)),
        in_size=config["in_size"],
        hidden_size=config["hidden_size"],
        num_layers=config["num_layers"],
        num_coupling_layers=config["num_coupling_layers"],
        permutations_key=jnp.array(config["key"], dtype=jnp.uint32),
        theta_ranges=config["theta_ranges"],
    )

    abstract_state = nnx.state(instance)
    checkpointer = orbax.PyTreeCheckpointer()
    state = checkpointer.restore(
        f"{path}/state", item=abstract_state, partial_restore=True
    )
    nnx.update(instance, state)

    instance.theta = metadata["theta"]
    instance.weights = metadata["weights"]
    instance.train_loss = metadata["train_loss"]
    instance.val_loss = metadata["val_loss"]
    instance.train_phi = metadata["train_phi"]
    instance.test_phi = metadata["test_phi"]
    instance.train_weights = metadata["train_weights"]
    instance.test_weights = metadata["test_weights"]
    instance.val_phi = metadata["val_phi"]
    instance.val_weights = metadata["val_weights"]

    shutil.rmtree(path)
    return instance

log_like(x, logevidence, prior_density)

Compute the marginal log-likelihood of given samples.

Parameters:

Name Type Description Default
x ndarray

Samples for which to compute the log-likelihood.

 required
logevidence float

Log-evidence term.

 required
prior_density ndarray | BaseDensityEstimator

Prior density or density estimator.

 required

Returns:

Type Description
ndarray

Log likelihoods of the input data.
Source code in margarine/estimators/realnvp.py 
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
def log_like(
    self,
    x: jnp.ndarray,
    logevidence: float,
    prior_density: jnp.ndarray | BaseDensityEstimator,
) -> jnp.ndarray:
    """Compute the marginal log-likelihood of given samples.

    Args:
        x: Samples for which to compute the log-likelihood.
        logevidence: Log-evidence term.
        prior_density: Prior density or density estimator.

    Returns:
        Log likelihoods of the input data.
    """
    if isinstance(prior_density, BaseDensityEstimator):
        prior_density = prior_density.log_prob(x)

    return self.log_prob(x) + logevidence - prior_density

log_prob(x)

Compute the log probability of the input data.

Parameters:

Name Type Description Default
x ndarray

Input data.

 required

Returns:

Type Description
ndarray

Log probabilities of the input data.
Source code in margarine/estimators/realnvp.py 
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
@jax.jit
def log_prob(self, x: jnp.ndarray) -> jnp.ndarray:
    """Compute the log probability of the input data.

    Args:
        x (jnp.ndarray): Input data.

    Returns:
        Log probabilities of the input data.
    """
    transformed_x = forward_transform(
        x, self.theta_ranges[0], self.theta_ranges[1]
    )
    transform_chain = tfb.Chain(
        [
            tfb.Invert(tfb.NormalCDF()),
            tfb.Scale(1 / (self.theta_ranges[1] - self.theta_ranges[0])),
            tfb.Shift(-self.theta_ranges[0]),
        ]
    )

    def norm_jac(y: jnp.ndarray) -> jnp.ndarray:
        """Calculate the normalising jacobian for the transformation."""
        return transform_chain.inverse_log_det_jacobian(y, event_ndims=0)

    correction = -norm_jac(transformed_x).sum(axis=1)

    return self.log_prob_under_RealNVP(transformed_x) + correction

log_prob_under_RealNVP(x)

Compute the log probability under the RealNVP model.

Parameters:

Name Type Description Default
x ndarray

Input data.

 required

Returns:

Type Description
ndarray

Log probabilities of the input data.
Source code in margarine/estimators/realnvp.py 
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
@jax.jit
def log_prob_under_RealNVP(self, x: jnp.ndarray) -> jnp.ndarray:
    """Compute the log probability under the RealNVP model.

    Args:
        x (jnp.ndarray): Input data.

    Returns:
        Log probabilities of the input data.
    """
    # calcualte the actual log prob under the RealNVP model
    # assuming a standard normal base distribution
    z, log_det_J = self.forward(x, return_log_det=True)
    log_pz = -0.5 * jnp.sum(z**2 + jnp.log(2 * jnp.pi), axis=1)
    log_prob_under_realnvp = log_pz + log_det_J
    return log_prob_under_realnvp

sample(key, num_samples)

Sample from the trained RealNVP model.

Parameters:

Name Type Description Default
key ndarray

JAX random key.

 required
num_samples int

Number of samples to draw.

 required

Returns:

Type Description
ndarray

Samples drawn from the RealNVP model.
Source code in margarine/estimators/realnvp.py 
368
369
370
371
372
373
374
375
376
377
378
379
380
def sample(self, key: jnp.ndarray, num_samples: int) -> jnp.ndarray:
    """Sample from the trained RealNVP model.

    Args:
        key: JAX random key.
        num_samples: Number of samples to draw.

    Returns:
        Samples drawn from the RealNVP model.
    """
    u = jax.random.uniform(key, shape=(num_samples, self.theta.shape[1]))
    samples = self(u)
    return samples

save(filename)

Save the trained RealNVP model to a file.

Parameters:

Name Type Description Default
filename str

Path to the file where the model will be saved.

 required
Source code in margarine/estimators/realnvp.py 
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
def save(self, filename: str) -> None:
    """Save the trained RealNVP model to a file.

    Args:
        filename: Path to the file where the model will be saved.
    """
    path = Path(filename + ".tmp").resolve()
    if path.exists():
        shutil.rmtree(path)

    state = nnx.state(self)

    checkpointer = orbax.PyTreeCheckpointer()
    checkpointer.save(f"{path}/state", state)

    config = {
        "in_size": self.in_size,
        "hidden_size": self.hidden_size,
        "num_layers": self.nlayers,
        "num_coupling_layers": self.num_coupling_layers,
        "key": self.permutations_key.tolist(),
        "theta_ranges": self.theta_ranges,
    }
    with open(f"{path}/config.yaml", "w") as f:
        yaml.dump(config, f)

    metadata = {
        "theta": self.theta,
        "weights": self.weights,
        "train_loss": self.train_loss,
        "val_loss": self.val_loss,
        "train_phi": self.train_phi,
        "test_phi": self.test_phi,
        "train_weights": self.train_weights,
        "test_weights": self.test_weights,
        "val_phi": self.val_phi,
        "val_weights": self.val_weights,
        "margarine_version": _version.__version__,
    }

    with open(f"{path}/metadata.yaml", "w") as f:
        yaml.dump(metadata, f)

    with ZipFile(filename + ".marg", "w") as z:
        for subpath in path.rglob("*"):
            if subpath.is_file():
                z.write(subpath, arcname=subpath.relative_to(path))

    shutil.rmtree(path)

train(key, learning_rate=0.0001, epochs=1000, patience=50, batch_size=256)

Train the RealNVP model.

Parameters:

Name Type Description Default
key ndarray

JAX random key for data splitting.

 required
learning_rate float

Learning rate for the optimizer.

 0.0001
epochs int

Number of training epochs.

 1000
patience int

Patience for early stopping.

 50
batch_size int

Size of training batches.

 256
Source code in margarine/estimators/realnvp.py 
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
def train(
    self,
    key: jnp.ndarray,
    learning_rate: float = 1e-4,
    epochs: int = 1000,
    patience: int = 50,
    batch_size: int = 256,
) -> None:
    """Train the RealNVP model.

    Args:
        key: JAX random key for data splitting.
        learning_rate: Learning rate for the optimizer.
        epochs: Number of training epochs.
        patience: Patience for early stopping.
        batch_size: Size of training batches.
    """

    @jax.jit
    def loss_function(
        model: RealNVP,
        targets: jnp.ndarray,
        weights: jnp.ndarray,
    ) -> jnp.ndarray:
        """Loss function for training RealNVP model.

        Args:
            model (RealNVP): The RealNVP model.
            targets (jnp.ndarray): Target samples.
            weights (jnp.ndarray): Weights for the samples.

        Returns:
            jnp.ndarray: Computed loss.
        """
        return -jnp.mean(weights * model.log_prob_under_RealNVP(targets))

    # @nnx.jit
    def train_step(
        model: RealNVP,
        optimizer: nnx.Optimizer,
        train_phi: jnp.ndarray,
        train_weights: jnp.ndarray,
    ) -> None:
        """Single training step for NICE model.

        Args:
            model (NICE): The NICE model to be trained.
            optimizer (nnx.Optimizer): The optimizer to
                update model parameters.
            train_phi (jnp.ndarray): Training samples.
            train_weights (jnp.ndarray): Weights for the training samples.
        """
        grad = nnx.grad(loss_function)(model, train_phi, train_weights)
        optimizer.update(model, grad)

    phi = forward_transform(
        self.theta, self.theta_ranges[0], self.theta_ranges[1]
    )
    weights = self.weights / jnp.sum(self.weights)

    key, subkey = jax.random.split(key)
    # need to split the data into training and validation sets
    (
        self.train_phi,
        self.test_phi,
        self.train_weights,
        self.test_weights,
    ) = train_test_split(
        phi,
        weights,
        subkey,
        test_size=0.2,
    )
    key, subkey = jax.random.split(key)
    self.test_phi, self.val_phi, self.test_weights, self.val_weights = (
        train_test_split(
            self.test_phi,
            self.test_weights,
            subkey,
            test_size=0.5,
        )
    )

    tx = optax.adam(learning_rate)
    optimizer = nnx.Optimizer(self, tx, wrt=nnx.Param)

    pbar = tqdm.tqdm(range(epochs), desc="Training")

    best_loss = jnp.inf
    best_model = nnx.state(self, nnx.Param)
    c = 0

    data_size = len(self.train_phi)

    tl, vl = [], []
    for _ in pbar:
        data_permutations = jax.random.permutation(
            subkey, jnp.arange(data_size)
        )
        accumulated_train_loss = 0.0
        for i in range(0, len(self.train_phi), batch_size):
            batch_indices = data_permutations[i : i + batch_size]
            batch_phi = self.train_phi[batch_indices]
            batch_weights = self.train_weights[batch_indices]
            loss = loss_function(self, batch_phi, batch_weights)
            train_step(self, optimizer, batch_phi, batch_weights)
            accumulated_train_loss += loss * len(batch_phi)
        loss = accumulated_train_loss / len(self.train_phi)
        tl.append(loss)

        val_loss = loss_function(self, self.val_phi, self.val_weights)
        vl.append(val_loss)

        if val_loss < best_loss:
            best_loss = val_loss
            best_model = nnx.state(self, nnx.Param)
            best_epoch = _
            c = 0
        else:
            c += 1
            if c >= patience:
                print(
                    f"Early stopping at epoch {_}, best "
                    + f"epoch was {best_epoch} with val loss {best_loss}"
                )
                break

        pbar.set_postfix(
            loss=f"{loss:.3e}",
            val_loss=f"{val_loss:.3e}",
            best_loss=f"{best_loss:.3e}",
        )

    self.train_loss = jnp.array(tl)
    self.val_loss = jnp.array(vl)

    if best_model is not None:
        nnx.update(self, best_model)

margarine.estimators.clustered

Module for clustered mixture of density estimators.

cluster

Create clustered mixture of MAFs to model multi-modal distributions.

Source code in margarine/estimators/clustered.py 
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
class cluster:
    """Create clustered mixture of MAFs to model multi-modal distributions."""

    def __init__(
        self,
        theta: jnp.ndarray,
        base_estimator: NICE | KDE | RealNVP,
        weights: jnp.ndarray | None = None,
        theta_ranges: jnp.ndarray | None = None,
        clusters: jnp.ndarray | None | int = None,
        max_cluster_number: int = 10,
        **kwargs: object,
    ) -> None:
        r"""Piecewise normalizing flow built from masked autoregressive flows.

        This class is a wrapper around the MAF class with additional clustering
        functionality. It trains, loads, and
        calls a piecewise density estimator where
        different base estimators are trained on
        different clusters of the sample space.

        Args:
            theta (jnp.ndarray): Samples to train the clustered MAF on.
            base_estimator (NICE | KDE | RealNVP): The base density estimator
                to use for each cluster.
            weights (jnp.ndarray | None, optional): Weights for the samples.
                Defaults to None.
            theta_ranges (jnp.ndarray | None, optional): Ranges for the
                parameters. Should have shape
                (nparams, 2). Defaults to None.
            clusters (jnp.ndarray | None | int, optional): Predefined cluster
                labels for each sample or an integer
                corresponding to the number of expected clusters.
                If None, k-means clustering is used.
                Defaults to None.
            max_cluster_number (int, optional): Maximum number of clusters
                to consider when using k-means clustering. Defaults to 10.
            **kwargs: Additional keyword arguments for the base estimator.
        """
        self.clusters = clusters
        self.base_estimator = base_estimator
        self.kwargs = kwargs

        self.theta = theta
        self.weights = weights
        self.theta_ranges = theta_ranges

        if self.weights is None:
            self.weights = jnp.ones(len(self.theta))

        if self.theta_ranges is None:
            self.theta_ranges = approximate_bounds(self.theta, self.weights)

        if self.clusters is None:
            ks = jnp.arange(2, max_cluster_number + 1)
            losses = []
            for k in ks:
                labels = kmeans(
                    self.theta,
                    k=k,
                )
                losses.append(-silhouette_score(self.theta, labels))
            losses = jnp.array(losses)
            minimum_index = jnp.argmin(losses)
            self.cluster_number = ks[minimum_index]

            self.clusters = kmeans(
                self.theta,
                k=self.cluster_number,
            )
        elif isinstance(self.clusters, int):
            self.cluster_number = self.clusters
            self.clusters = kmeans(
                self.theta,
                k=self.cluster_number,
            )
        else:
            self.cluster_number = len(jnp.unique(self.clusters))

        # count the number of times a cluster label appears in cluster_labels
        self.cluster_count = jnp.bincount(self.clusters)

        # 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]
            self.clusters = kmeans(
                self.theta, k=self.cluster_number, num_iters=25
            )
            self.cluster_count = jnp.bincount(self.clusters)
            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_weights = []
        for i in range(self.cluster_number):
            split_theta.append(self.theta[self.clusters == i])
            split_weights.append(self.weights[self.clusters == i])

        self.split_theta = split_theta
        self.split_weights = split_weights

        self.estimators = []
        for i in range(len(split_theta)):
            self.estimators.append(
                base_estimator(
                    split_theta[i],
                    weights=split_weights[i],
                    theta_ranges=self.theta_ranges,
                    **self.kwargs,
                )
            )

    def train(self, **kwargs: object) -> None:
        r"""Train the cluster estimator.

        Args:
            **kwargs: Additional keyword arguments for training.
        """
        for i in range(len(self.estimators)):
            self.estimators[i].train(**kwargs)

    def log_prob(self, x: jnp.ndarray) -> jnp.ndarray:
        """Log-probability for a given 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.

        Args:
            x (jnp.ndarray): The set of samples for which to
                calculate the log probability.

        Returns:
            jnp.ndarray: The log-probabilities of the provided samples.
        """
        estimator_importance = jnp.array(
            [jnp.sum(weights) for weights in self.split_weights]
        )
        estimator_importance = estimator_importance / jnp.sum(
            estimator_importance
        )

        logprob = []
        for estimator, importance in zip(
            self.estimators, estimator_importance
        ):
            prob = estimator.log_prob(x)
            probs = prob + jnp.log(importance)
            logprob.append(probs)
        logprob = jnp.array(logprob)
        logprob = logsumexp(logprob, axis=0)

        return logprob

    def log_like(
        self,
        x: jnp.ndarray,
        logevidence: float,
        prior_density: jnp.ndarray | NICE | KDE | RealNVP,
    ) -> jnp.ndarray:
        """Compute the marginal log-likelihood of given samples.

        Args:
            x: Samples for which to compute the log-likelihood.
            logevidence: Log-evidence term.
            prior_density: Prior density or density estimator.

        Returns:
            jnp.ndarray: Log-likelihoods of the samples.
        """
        if isinstance(prior_density, NICE | KDE | RealNVP):
            prior_density = prior_density.log_prob(x)

        return self.log_prob(x) + logevidence - prior_density

    def __call__(self, key: jnp.ndarray, u: jnp.ndarray) -> jnp.ndarray:
        r"""Transform samples from the unit hypercube to the cluster.

        This function is used when calling the cluster class to transform
        samples from the unit hypercube to samples on the
        clustered distribution.

        Args:
            key (jnp.ndarray): JAX random key for sampling.
            u (jnp.ndarray): Samples on the unit hypercube.

        Returns:
           jnp.ndarray: Samples transformed to the clusterMAF.
        """
        estimator_importance = jnp.array(
            [jnp.sum(weights) for weights in self.split_weights]
        )
        probabilities = estimator_importance / jnp.sum(estimator_importance)
        options = jnp.arange(0, self.cluster_number)

        choice = jax.random.choice(
            key, options, p=probabilities, shape=(len(u),)
        )

        totals = jnp.array(
            [len(choice[choice == options[i]]) for i in range(len(options))]
        )
        totals = jnp.hstack([0, jnp.cumsum(totals)])

        values = []
        for i in range(len(options)):
            x = self.estimators[i](u[totals[i] : totals[i + 1]])
            values.append(x)

        return jnp.vstack(values)

    def sample(self, key: jnp.ndarray, num_samples: int = 1000) -> jnp.ndarray:
        r"""Generate samples on the cluster.

        Args:
            key (jnp.ndarray): JAX random key for sampling.
            num_samples (int, optional): The number of samples to generate.
                Defaults to 1000.

        Returns:
            jnp.ndarray: Samples generated on the clusterMAF.
        """
        estimator_importance = jnp.array(
            [jnp.sum(weights) for weights in self.split_weights]
        )
        probabilities = estimator_importance / jnp.sum(estimator_importance)
        options = jnp.arange(0, self.cluster_number)

        choice = jax.random.choice(
            key, options, p=probabilities, shape=(num_samples,)
        )

        totals = jnp.array(
            [len(choice[choice == options[i]]) for i in range(len(options))]
        )

        values = []
        for i in range(len(options)):
            key, subkey = jax.random.split(key)
            x = self.estimators[i].sample(subkey, totals[i])
            values.append(x)

        return jnp.vstack(values)

    def save(self, filename: str) -> None:
        """Save the clustered estimator to a file.

        Args:
            filename (str): The name of the file to save the estimator to.
        """
        path = Path(filename).resolve()
        if path.exists():
            shutil.rmtree(path)

        os.makedirs(path)

        config = {
            "base_estimator": self.base_estimator.__name__,
            "kwargs": self.kwargs,
            "theta_ranges": self.theta_ranges,
            "clusters": self.clusters,
        }

        with open(path / "config.yaml", "w") as f:
            yaml.dump(config, f)

        for i, estimator in enumerate(self.estimators):
            estimator.save(str(path) + f"/estimator_{i}")

        metadata = {
            "theta": self.theta,
            "weights": self.weights,
        }

        with open(path / "metadata.yaml", "w") as f:
            yaml.dump(metadata, f)

        with ZipFile(filename + ".clumarg", "w") as z:
            for subpath in path.rglob("*"):
                if subpath.is_file():
                    z.write(subpath, arcname=subpath.relative_to(path))

        shutil.rmtree(path)

    @classmethod
    def load(cls, filename: str) -> "cluster":
        """Load a clustered estimator from a file.

        Args:
            filename (str): The name of the file to load the estimator from.

        Returns:
            cluster: The loaded clustered estimator.
        """
        zip_path = Path(f"{filename}.clumarg")
        path = Path(filename + ".tmp").resolve()
        with ZipFile(zip_path) as z:
            # Extract all files to a folder
            z.extractall(path)

        with open(path / "config.yaml") as f:
            config = yaml.unsafe_load(f)

        with open(path / "metadata.yaml") as f:
            metadata = yaml.unsafe_load(f)

        base_estimator_class = globals()[config["base_estimator"]]

        instance = cls(
            theta=metadata["theta"],  # Placeholder, will be overwritten
            base_estimator=base_estimator_class,
            weights=metadata["weights"],  # Placeholder, will be overwritten
            theta_ranges=config["theta_ranges"],
            clusters=config["clusters"],
            **config["kwargs"],
        )

        instance.estimators = []
        for i in range(instance.cluster_number):
            estimator = base_estimator_class.load(
                str(path) + f"/estimator_{i}"
            )
            instance.estimators.append(estimator)

        shutil.rmtree(path)

        return instance

__call__(key, u)

Transform samples from the unit hypercube to the cluster.

This function is used when calling the cluster class to transform samples from the unit hypercube to samples on the clustered distribution.

Parameters:

Name Type Description Default
key ndarray

JAX random key for sampling.

 required
u ndarray

Samples on the unit hypercube.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Samples transformed to the clusterMAF.
Source code in margarine/estimators/clustered.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
def __call__(self, key: jnp.ndarray, u: jnp.ndarray) -> jnp.ndarray:
    r"""Transform samples from the unit hypercube to the cluster.

    This function is used when calling the cluster class to transform
    samples from the unit hypercube to samples on the
    clustered distribution.

    Args:
        key (jnp.ndarray): JAX random key for sampling.
        u (jnp.ndarray): Samples on the unit hypercube.

    Returns:
       jnp.ndarray: Samples transformed to the clusterMAF.
    """
    estimator_importance = jnp.array(
        [jnp.sum(weights) for weights in self.split_weights]
    )
    probabilities = estimator_importance / jnp.sum(estimator_importance)
    options = jnp.arange(0, self.cluster_number)

    choice = jax.random.choice(
        key, options, p=probabilities, shape=(len(u),)
    )

    totals = jnp.array(
        [len(choice[choice == options[i]]) for i in range(len(options))]
    )
    totals = jnp.hstack([0, jnp.cumsum(totals)])

    values = []
    for i in range(len(options)):
        x = self.estimators[i](u[totals[i] : totals[i + 1]])
        values.append(x)

    return jnp.vstack(values)

__init__(theta, base_estimator, weights=None, theta_ranges=None, clusters=None, max_cluster_number=10, **kwargs)

Piecewise normalizing flow built from masked autoregressive flows.

This class is a wrapper around the MAF class with additional clustering functionality. It trains, loads, and calls a piecewise density estimator where different base estimators are trained on different clusters of the sample space.

Parameters:

Name Type Description Default
theta ndarray

Samples to train the clustered MAF on.

 required
base_estimator NICE | KDE | RealNVP

The base density estimator to use for each cluster.

 required
weights ndarray | None

Weights for the samples. Defaults to None.

 None
theta_ranges ndarray | None

Ranges for the parameters. Should have shape (nparams, 2). Defaults to None.

 None
clusters ndarray | None | int

Predefined cluster labels for each sample or an integer corresponding to the number of expected clusters. If None, k-means clustering is used. Defaults to None.

 None
max_cluster_number int

Maximum number of clusters to consider when using k-means clustering. Defaults to 10.

 10
**kwargs object

Additional keyword arguments for the base estimator.

 {}
Source code in margarine/estimators/clustered.py 
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
def __init__(
    self,
    theta: jnp.ndarray,
    base_estimator: NICE | KDE | RealNVP,
    weights: jnp.ndarray | None = None,
    theta_ranges: jnp.ndarray | None = None,
    clusters: jnp.ndarray | None | int = None,
    max_cluster_number: int = 10,
    **kwargs: object,
) -> None:
    r"""Piecewise normalizing flow built from masked autoregressive flows.

    This class is a wrapper around the MAF class with additional clustering
    functionality. It trains, loads, and
    calls a piecewise density estimator where
    different base estimators are trained on
    different clusters of the sample space.

    Args:
        theta (jnp.ndarray): Samples to train the clustered MAF on.
        base_estimator (NICE | KDE | RealNVP): The base density estimator
            to use for each cluster.
        weights (jnp.ndarray | None, optional): Weights for the samples.
            Defaults to None.
        theta_ranges (jnp.ndarray | None, optional): Ranges for the
            parameters. Should have shape
            (nparams, 2). Defaults to None.
        clusters (jnp.ndarray | None | int, optional): Predefined cluster
            labels for each sample or an integer
            corresponding to the number of expected clusters.
            If None, k-means clustering is used.
            Defaults to None.
        max_cluster_number (int, optional): Maximum number of clusters
            to consider when using k-means clustering. Defaults to 10.
        **kwargs: Additional keyword arguments for the base estimator.
    """
    self.clusters = clusters
    self.base_estimator = base_estimator
    self.kwargs = kwargs

    self.theta = theta
    self.weights = weights
    self.theta_ranges = theta_ranges

    if self.weights is None:
        self.weights = jnp.ones(len(self.theta))

    if self.theta_ranges is None:
        self.theta_ranges = approximate_bounds(self.theta, self.weights)

    if self.clusters is None:
        ks = jnp.arange(2, max_cluster_number + 1)
        losses = []
        for k in ks:
            labels = kmeans(
                self.theta,
                k=k,
            )
            losses.append(-silhouette_score(self.theta, labels))
        losses = jnp.array(losses)
        minimum_index = jnp.argmin(losses)
        self.cluster_number = ks[minimum_index]

        self.clusters = kmeans(
            self.theta,
            k=self.cluster_number,
        )
    elif isinstance(self.clusters, int):
        self.cluster_number = self.clusters
        self.clusters = kmeans(
            self.theta,
            k=self.cluster_number,
        )
    else:
        self.cluster_number = len(jnp.unique(self.clusters))

    # count the number of times a cluster label appears in cluster_labels
    self.cluster_count = jnp.bincount(self.clusters)

    # 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]
        self.clusters = kmeans(
            self.theta, k=self.cluster_number, num_iters=25
        )
        self.cluster_count = jnp.bincount(self.clusters)
        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_weights = []
    for i in range(self.cluster_number):
        split_theta.append(self.theta[self.clusters == i])
        split_weights.append(self.weights[self.clusters == i])

    self.split_theta = split_theta
    self.split_weights = split_weights

    self.estimators = []
    for i in range(len(split_theta)):
        self.estimators.append(
            base_estimator(
                split_theta[i],
                weights=split_weights[i],
                theta_ranges=self.theta_ranges,
                **self.kwargs,
            )
        )

load(filename) classmethod

Load a clustered estimator from a file.

Parameters:

Name Type Description Default
filename str

The name of the file to load the estimator from.

 required

Returns:

Name Type Description
cluster cluster

The loaded clustered estimator.
Source code in margarine/estimators/clustered.py 
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
@classmethod
def load(cls, filename: str) -> "cluster":
    """Load a clustered estimator from a file.

    Args:
        filename (str): The name of the file to load the estimator from.

    Returns:
        cluster: The loaded clustered estimator.
    """
    zip_path = Path(f"{filename}.clumarg")
    path = Path(filename + ".tmp").resolve()
    with ZipFile(zip_path) as z:
        # Extract all files to a folder
        z.extractall(path)

    with open(path / "config.yaml") as f:
        config = yaml.unsafe_load(f)

    with open(path / "metadata.yaml") as f:
        metadata = yaml.unsafe_load(f)

    base_estimator_class = globals()[config["base_estimator"]]

    instance = cls(
        theta=metadata["theta"],  # Placeholder, will be overwritten
        base_estimator=base_estimator_class,
        weights=metadata["weights"],  # Placeholder, will be overwritten
        theta_ranges=config["theta_ranges"],
        clusters=config["clusters"],
        **config["kwargs"],
    )

    instance.estimators = []
    for i in range(instance.cluster_number):
        estimator = base_estimator_class.load(
            str(path) + f"/estimator_{i}"
        )
        instance.estimators.append(estimator)

    shutil.rmtree(path)

    return instance

log_like(x, logevidence, prior_density)

Compute the marginal log-likelihood of given samples.

Parameters:

Name Type Description Default
x ndarray

Samples for which to compute the log-likelihood.

 required
logevidence float

Log-evidence term.

 required
prior_density ndarray | NICE | KDE | RealNVP

Prior density or density estimator.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Log-likelihoods of the samples.
Source code in margarine/estimators/clustered.py 
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
def log_like(
    self,
    x: jnp.ndarray,
    logevidence: float,
    prior_density: jnp.ndarray | NICE | KDE | RealNVP,
) -> jnp.ndarray:
    """Compute the marginal log-likelihood of given samples.

    Args:
        x: Samples for which to compute the log-likelihood.
        logevidence: Log-evidence term.
        prior_density: Prior density or density estimator.

    Returns:
        jnp.ndarray: Log-likelihoods of the samples.
    """
    if isinstance(prior_density, NICE | KDE | RealNVP):
        prior_density = prior_density.log_prob(x)

    return self.log_prob(x) + logevidence - prior_density

log_prob(x)

Log-probability for a given 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:

Name Type Description Default
x ndarray

The set of samples for which to calculate the log probability.

 required

Returns:

Type Description
ndarray

jnp.ndarray: The log-probabilities of the provided samples.
Source code in margarine/estimators/clustered.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def log_prob(self, x: jnp.ndarray) -> jnp.ndarray:
    """Log-probability for a given 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.

    Args:
        x (jnp.ndarray): The set of samples for which to
            calculate the log probability.

    Returns:
        jnp.ndarray: The log-probabilities of the provided samples.
    """
    estimator_importance = jnp.array(
        [jnp.sum(weights) for weights in self.split_weights]
    )
    estimator_importance = estimator_importance / jnp.sum(
        estimator_importance
    )

    logprob = []
    for estimator, importance in zip(
        self.estimators, estimator_importance
    ):
        prob = estimator.log_prob(x)
        probs = prob + jnp.log(importance)
        logprob.append(probs)
    logprob = jnp.array(logprob)
    logprob = logsumexp(logprob, axis=0)

    return logprob

sample(key, num_samples=1000)

Generate samples on the cluster.

Parameters:

Name Type Description Default
key ndarray

JAX random key for sampling.

 required
num_samples int

The number of samples to generate. Defaults to 1000.

 1000

Returns:

Type Description
ndarray

jnp.ndarray: Samples generated on the clusterMAF.
Source code in margarine/estimators/clustered.py 
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
def sample(self, key: jnp.ndarray, num_samples: int = 1000) -> jnp.ndarray:
    r"""Generate samples on the cluster.

    Args:
        key (jnp.ndarray): JAX random key for sampling.
        num_samples (int, optional): The number of samples to generate.
            Defaults to 1000.

    Returns:
        jnp.ndarray: Samples generated on the clusterMAF.
    """
    estimator_importance = jnp.array(
        [jnp.sum(weights) for weights in self.split_weights]
    )
    probabilities = estimator_importance / jnp.sum(estimator_importance)
    options = jnp.arange(0, self.cluster_number)

    choice = jax.random.choice(
        key, options, p=probabilities, shape=(num_samples,)
    )

    totals = jnp.array(
        [len(choice[choice == options[i]]) for i in range(len(options))]
    )

    values = []
    for i in range(len(options)):
        key, subkey = jax.random.split(key)
        x = self.estimators[i].sample(subkey, totals[i])
        values.append(x)

    return jnp.vstack(values)

save(filename)

Save the clustered estimator to a file.

Parameters:

Name Type Description Default
filename str

The name of the file to save the estimator to.

 required
Source code in margarine/estimators/clustered.py 
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
def save(self, filename: str) -> None:
    """Save the clustered estimator to a file.

    Args:
        filename (str): The name of the file to save the estimator to.
    """
    path = Path(filename).resolve()
    if path.exists():
        shutil.rmtree(path)

    os.makedirs(path)

    config = {
        "base_estimator": self.base_estimator.__name__,
        "kwargs": self.kwargs,
        "theta_ranges": self.theta_ranges,
        "clusters": self.clusters,
    }

    with open(path / "config.yaml", "w") as f:
        yaml.dump(config, f)

    for i, estimator in enumerate(self.estimators):
        estimator.save(str(path) + f"/estimator_{i}")

    metadata = {
        "theta": self.theta,
        "weights": self.weights,
    }

    with open(path / "metadata.yaml", "w") as f:
        yaml.dump(metadata, f)

    with ZipFile(filename + ".clumarg", "w") as z:
        for subpath in path.rglob("*"):
            if subpath.is_file():
                z.write(subpath, arcname=subpath.relative_to(path))

    shutil.rmtree(path)

train(**kwargs)

Train the cluster estimator.

Parameters:

Name Type Description Default
**kwargs object

Additional keyword arguments for training.

 {}
Source code in margarine/estimators/clustered.py 
150
151
152
153
154
155
156
157
def train(self, **kwargs: object) -> None:
    r"""Train the cluster estimator.

    Args:
        **kwargs: Additional keyword arguments for training.
    """
    for i in range(len(self.estimators)):
        self.estimators[i].train(**kwargs)

margarine.utils.kmeans

K-Means clustering utility functions.

kmeans(X, k, num_iters=100)

Performs K-Means clustering on the given data.

Parameters:

Name Type Description Default
X ndarray

Input data of shape (n_samples, n_features).

 required
k int | ndarray

Number of clusters.

 required
num_iters int

Number of iterations for the K-Means algorithm.

 100

Returns:

Type Description
ndarray

jnp.ndarray: Cluster labels for each data point.
Source code in margarine/utils/kmeans.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
def kmeans(
    X: jnp.ndarray, k: int | jnp.ndarray, num_iters: int = 100
) -> jnp.ndarray:
    """Performs K-Means clustering on the given data.

    Args:
        X: Input data of shape (n_samples, n_features).
        k: Number of clusters.
        num_iters: Number of iterations for the K-Means algorithm.

    Returns:
        jnp.ndarray: Cluster labels for each data point.
    """

    @jax.jit
    def distance(a: jnp.ndarray, b: jnp.ndarray) -> jnp.ndarray:
        """Calculate Euclidean distance between two sets of points.

        Args:
            a: First set of points.
            b: Second set of points.

        Returns:
            jnp.ndarray: Euclidean distances.
        """
        return jnp.sum((a - b) ** 2, axis=-1)

    initial_centroids = jax.random.multivariate_normal(
        jax.random.PRNGKey(0),
        jnp.mean(X, axis=0),
        jnp.cov(X, rowvar=False),
        shape=(k,),
    )

    vmap_distance = jax.vmap(distance, in_axes=(0, 0))

    def clustering_step(centroids: jnp.ndarray, X: jnp.ndarray) -> jnp.ndarray:
        """Performs a single step of the K-Means clustering algorithm.

        Args:
            centroids: Current centroids of shape (k, n_features).
            X: Input data of shape (n_samples, n_features).

        Returns:
            jnp.ndarray: Updated centroids and labels.
        """
        initial_centroids = jnp.tile(
            centroids[jnp.newaxis, :, :], (X.shape[0], 1, 1)
        )
        distances = vmap_distance(initial_centroids, X)

        labels = jnp.argmin(distances, axis=1)
        new_centroids = jnp.array(
            [jnp.mean(X[labels == i], axis=0) for i in range(k)]
        )
        return new_centroids, labels

    centroids = initial_centroids
    for _ in range(num_iters):
        centroids, labels = clustering_step(centroids, X)

    return labels

silhouette_score(X, labels)

Calculates the silhouette score for the clustering.

Parameters:

Name Type Description Default
X ndarray

Input data of shape (n_samples, n_features).

 required
labels ndarray

Cluster labels for each data point.

 required

Returns:

Name Type Description
float ndarray

Silhouette score.
Source code in margarine/utils/kmeans.py 
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
def silhouette_score(X: jnp.ndarray, labels: jnp.ndarray) -> jnp.ndarray:
    """Calculates the silhouette score for the clustering.

    Args:
        X: Input data of shape (n_samples, n_features).
        labels: Cluster labels for each data point.

    Returns:
        float: Silhouette score.
    """

    def average_distance(x: jnp.ndarray, X: jnp.ndarray) -> jnp.ndarray:
        """Calculate the average distance from point x to all points in X."""
        return jnp.mean((X - x) ** 2)

    vmapped_avg_distance = jax.vmap(average_distance, in_axes=(0, None))

    silhouette_scores = []
    for label in jnp.unique(labels):
        cluster_points = X[labels == label]
        cluster_distance = vmapped_avg_distance(cluster_points, cluster_points)
        seperation_distance = jnp.min(
            jnp.array(
                [
                    vmapped_avg_distance(
                        cluster_points, X[labels == other_label]
                    )
                    for other_label in jnp.unique(labels)
                    if other_label != label
                ]
            )
        )
        # score per point in cluster
        score = (seperation_distance - cluster_distance) / jnp.maximum(
            seperation_distance, cluster_distance
        )
        # score per cluster
        silhouette_scores.append(jnp.mean(score))

    return jnp.mean(jnp.array(silhouette_scores))

margarine.utils.utils

Utility functions for margarine package.

approximate_bounds(theta, weights)

Function to estimate prior bounds from samples.

Sample maximum and minimum are biased estimators of the true bounds of the distribution. This function provides an improved estimate using the weights of the samples. Comes from the expectation value of the maximum and minimum samples of a uniform distribution with an effective number of samples given by Kish's formula.

Parameters:

Name Type Description Default
theta ndarray

samples from the target distribution.

 required
weights ndarray

weights for the samples from the target distribution.

 required
Return

a (jnp.ndarray): estimate of the upper bound on the prior. b (jnp.ndarray): estimate of the lower bound on the prior.

Source code in margarine/utils/utils.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
@jax.jit
def approximate_bounds(
    theta: jnp.ndarray, weights: jnp.ndarray
) -> jnp.ndarray:
    """Function to estimate prior bounds from samples.

    Sample maximum and minimum are biased estimators of the true
    bounds of the distribution. This function provides an improved
    estimate using the weights of the samples. Comes from the expectation
    value of the maximum and minimum samples
    of a uniform distribution with an effective number of samples
    given by Kish's formula.

    Args:
        theta (jnp.ndarray): samples from the target distribution.
        weights (jnp.ndarray): weights for the samples from the
            target distribution.

    Return:
        a (jnp.ndarray): estimate of the upper bound on the prior.
        b (jnp.ndarray): estimate of the lower bound on the prior.
    """
    n = (jnp.sum(weights) ** 2) / (jnp.sum(weights**2))
    sample_max = jnp.max(theta, axis=0)
    sample_min = jnp.min(theta, axis=0)

    a = ((n - 2) * sample_max - sample_min) / (n - 3)
    b = ((n - 2) * sample_min - sample_max) / (n - 3)

    return jnp.stack([b, a])

forward_transform(x, min_val, max_val)

Forward transform input samples.

Normalise between 0 and 1 and then transform onto samples of standard normal distribution.

Parameters:

Name Type Description Default
x ndarray

Samples to be normalised.

 required
min_val float

Minimum value for normalization.

 required
max_val float

Maximum value for normalization.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Transformed samples.
Source code in margarine/utils/utils.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
@jax.jit
def forward_transform(
    x: jnp.ndarray, min_val: float, max_val: float
) -> jnp.ndarray:
    """Forward transform input samples.

    Normalise between 0 and 1 and then transform
    onto samples of standard normal distribution.

    Args:
        x (jnp.ndarray): Samples to be normalised.
        min_val (float): Minimum value for normalization.
        max_val (float): Maximum value for normalization.

    Returns:
        jnp.ndarray: Transformed samples.
    """
    x = tfd.Uniform(min_val, max_val).cdf(x)
    x = tfd.Normal(0, 1).quantile(x)
    return x

inverse_transform(x, min_val, max_val)

Inverse transform output samples.

Inverts the processes in forward_transform.

Parameters:

Name Type Description Default
x ndarray

Samples to be inverse normalised.

 required
min_val float

Minimum value for normalization.

 required
max_val float

Maximum value for normalization.

 required

Returns:

Type Description
ndarray

jnp.ndarray: Inverse transformed samples.
Source code in margarine/utils/utils.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
@jax.jit
def inverse_transform(
    x: jnp.ndarray, min_val: float, max_val: float
) -> jnp.ndarray:
    """Inverse transform output samples.

    Inverts the processes in ``forward_transform``.

    Args:
        x (jnp.ndarray): Samples to be inverse normalised.
        min_val (float): Minimum value for normalization.
        max_val (float): Maximum value for normalization.

    Returns:
        jnp.ndarray: Inverse transformed samples.
    """
    x = tfd.Normal(0, 1).cdf(x)
    x = tfd.Uniform(min_val, max_val).quantile(x)
    return x

train_test_split(a, b, key=None, 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 jax arrays.

Parameters:

Name Type Description Default
a ndarray

First set of data to be split.

 required
b ndarray

Second set of data to be split.

 required
test_size float

Proportion of data to be used for testing.

 0.2
key KeyArray

JAX random key for shuffling.

 None

Returns:

Name Type Description
a_train ndarray

Training set from a.
a_test ndarray

Testing set from a.
b_train ndarray

Training set from b.
b_test ndarray

Testing set from b.
Source code in margarine/utils/utils.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def train_test_split(
    a: jnp.ndarray,
    b: jnp.ndarray,
    key: jnp.ndarray | None = None,
    test_size: float = 0.2,
) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
    """Splitting data into training and testing sets.

    Function is equivalent
    to sklearn.model_selection.train_test_split but a and b
    are jax arrays.

    Args:
        a (jnp.ndarray): First set of data to be split.
        b (jnp.ndarray): Second set of data to be split.
        test_size (float): Proportion of data to be used for testing.
        key (jax.random.KeyArray): JAX random key for shuffling.

    Returns:
        a_train (jnp.ndarray): Training set from a.
        a_test (jnp.ndarray): Testing set from a.
        b_train (jnp.ndarray): Training set from b.
        b_test (jnp.ndarray): Testing set from b.
    """
    if key is None:
        key = jax.random.PRNGKey(0)
    n_samples = a.shape[0]
    n_test = int(n_samples * test_size)

    permuted_indices = jax.random.permutation(key, n_samples)
    test_indices = permuted_indices[:n_test]
    train_indices = permuted_indices[n_test:]

    a_train = a[train_indices]
    a_test = a[test_indices]
    b_train = b[train_indices]
    b_test = b[test_indices]

    return a_train, a_test, b_train, b_test

margarine.statistics

Statistics functions for margarine.

integrate(density_estimator, likelihood, prior, batch_size=1000, sample_size=10000, logzero=-1e+30, key=jax.random.PRNGKey(0))

Importance sampling integration of a likelihood function.

Parameters:

Name Type Description Default
density_estimator BaseDensityEstimator

A density estimator

 required
likelihood Callable

The likelihood function to integrate.

 required
prior BaseDensityEstimator

A density estimator

 required
batch_size int

The number of samples to draw at each iteration.

 1000
sample_size int

The number of samples to draw in total.

 10000
logzero float

The definition of zero for the loglikelihood function.

 -1e+30
key ndarray

JAX random key for sampling.

 PRNGKey(0)

Returns:

Name Type Description
stats dict

Dictionary containing useful statistics
Source code in margarine/statistics.py 
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
def integrate(
    density_estimator: BaseDensityEstimator,
    likelihood: Callable,
    prior: BaseDensityEstimator,
    batch_size: int = 1000,
    sample_size: int = 10000,
    logzero: float = -1e30,
    key: jnp.ndarray = jax.random.PRNGKey(0),
) -> dict:
    """Importance sampling integration of a likelihood function.

    Args:
        density_estimator (BaseDensityEstimator): A density estimator
        for the likelihood function.
        likelihood (Callable): The likelihood function to integrate.
        prior (BaseDensityEstimator): A density estimator
        for the prior probability.
        batch_size (int): The number of samples to draw at each iteration.
        sample_size (int): The number of samples to draw in total.
        logzero (float): The definition of zero for the
            loglikelihood function.
        key (jnp.ndarray): JAX random key for sampling.

    Returns:
        stats (dict): Dictionary containing useful statistics
    """
    xs = jnp.empty((sample_size, density_estimator.theta.shape[-1]))
    fs = jnp.empty(sample_size)
    gs = jnp.empty(sample_size)
    pis = jnp.empty(sample_size)

    n_todo = sample_size
    trials = 0

    with tqdm.tqdm(total=sample_size) as pbar:
        while n_todo > 0:  # draw samples until we have enough accepted
            key, subkey = jax.random.split(key)
            x = density_estimator.sample(subkey, batch_size)
            f = jnp.array(list(map(likelihood, x)))
            g = density_estimator.log_prob(x)
            # determine which samples are accepted based on logzero
            in_bounds = jnp.logical_and(f >= logzero, g >= logzero)
            n_accept = x[in_bounds].shape[0]
            if n_accept <= n_todo:
                # accepted samples
                xs = xs.at[
                    sample_size - n_todo : sample_size - n_todo + n_accept
                ].set(x[in_bounds])
                # corresponding likelihoods
                fs = fs.at[
                    sample_size - n_todo : sample_size - n_todo + n_accept
                ].set(f[in_bounds])
                # corresponding flow log-probs
                gs = gs.at[
                    sample_size - n_todo : sample_size - n_todo + n_accept
                ].set(g[in_bounds])
                # corresponding prior log-probs
                pis = pis.at[
                    sample_size - n_todo : sample_size - n_todo + n_accept
                ].set(prior.log_prob(x[in_bounds]))
                trials += batch_size
            else:
                n_accept = n_todo
                xs = xs.at[sample_size - n_todo :].set(x[in_bounds][:n_accept])
                fs = fs.at[sample_size - n_todo :].set(f[in_bounds][:n_accept])
                gs = gs.at[sample_size - n_todo :].set(g[in_bounds][:n_accept])
                pis = pis.at[sample_size - n_todo :].set(
                    prior.log_prob(x[in_bounds])[:n_accept]
                )
                last_index = in_bounds[-1]
                trials += last_index + 1
            n_todo -= n_accept
            pbar.update(n_accept)
            # prevent excessive calls
            if trials > 10 * sample_size:
                raise ValueError(
                    "Too many unsuccessful trials, this"
                    + "typically indicates mismatch between"
                    + "flow and likelihood"
                )

        # calculate the importance weights = prior * likelihood / flow
        weights = jnp.exp(fs + pis - gs)
        # effective sample size
        eff = jnp.sum(weights) ** 2 / jnp.sum(weights**2) / sample_size
        # estimate of the integral and its standard error
        integral = sample_size / trials * weights.mean()
        log_integral = logsumexp(fs + pis - gs) - jnp.log(trials)

        stderr = jnp.sqrt(
            (jnp.sum(weights**2) / trials - integral**2) / (trials - 1)
        )
        log_stderr = jnp.log(stderr)

    stats = {
        "x": xs,
        "y": fs,
        "weights": weights,
        "efficiency": eff,
        "trials": trials,
        "log_integral": log_integral,
        "log_stderr": log_stderr,
    }
    return stats

kldivergence(density_estimator_p, density_estimator_q, samples_p=None, weights=None, key=jax.random.PRNGKey(0))

Kullback-Leibler divergence between two density estimators.

Parameters:

Name Type Description Default
density_estimator_p BaseDensityEstimator

The first

 required
density_estimator_q BaseDensityEstimator

The second

 required
samples_p ndarray | None

Optional samples from the

 None
weights ndarray | None

Optional weights for the

 None
key ndarray

JAX random key for sampling.

 PRNGKey(0)

Returns:

Name Type Description
kld float

The Kullback-Leibler divergence D_KL(P || Q).
Source code in margarine/statistics.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def kldivergence(
    density_estimator_p: BaseDensityEstimator,
    density_estimator_q: BaseDensityEstimator,
    samples_p: jnp.ndarray | None = None,
    weights: jnp.ndarray | None = None,
    key: jnp.ndarray = jax.random.PRNGKey(0),
) -> float:
    """Kullback-Leibler divergence between two density estimators.

    Args:
        density_estimator_p (BaseDensityEstimator): The first
        density estimator.
        density_estimator_q (BaseDensityEstimator): The second
        density estimator.
        samples_p (jnp.ndarray | None): Optional samples from the
        first density estimator. If None, samples will be drawn
        from density_estimator_p.
        weights (jnp.ndarray | None): Optional weights for the
        samples. If None, uniform weights will be used.
        key (jnp.ndarray): JAX random key for sampling.

    Returns:
        kld (float): The Kullback-Leibler divergence D_KL(P || Q).
    """
    if samples_p is None:
        samples_p = density_estimator_p.sample(key, num_samples=10000)
    log_p = density_estimator_p.log_prob(samples_p)
    log_q = density_estimator_q.log_prob(samples_p)
    kld = jnp.average(log_p - log_q, weights=weights)
    return kld

model_dimensionality(density_estimator_p, density_estimator_q, samples_p=None, weights=None, key=jax.random.PRNGKey(0))

Model dimensionality between two density estimators.

Parameters:

Name Type Description Default
density_estimator_p BaseDensityEstimator

The first

 required
density_estimator_q BaseDensityEstimator

The second

 required
samples_p ndarray | None

Optional samples from the

 None
weights ndarray | None

Optional weights for the

 None
key ndarray

JAX random key for sampling.

 PRNGKey(0)

Returns:

Name Type Description
dim float

The model dimensionality.
Source code in margarine/statistics.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def model_dimensionality(
    density_estimator_p: BaseDensityEstimator,
    density_estimator_q: BaseDensityEstimator,
    samples_p: jnp.ndarray | None = None,
    weights: jnp.ndarray | None = None,
    key: jnp.ndarray = jax.random.PRNGKey(0),
) -> float:
    """Model dimensionality between two density estimators.

    Args:
        density_estimator_p (BaseDensityEstimator): The first
        density estimator.
        density_estimator_q (BaseDensityEstimator): The second
        density estimator.
        samples_p (jnp.ndarray | None): Optional samples from the
        first density estimator. If None, samples will be drawn
        from density_estimator_p.
        weights (jnp.ndarray | None): Optional weights for the
        samples. If None, uniform weights will be used.
        key (jnp.ndarray): JAX random key for sampling.

    Returns:
        dim (float): The model dimensionality.
    """
    if samples_p is None:
        samples_p = density_estimator_p.sample(key, num_samples=10000)
    log_p = density_estimator_p.log_prob(samples_p)
    log_q = density_estimator_q.log_prob(samples_p)
    delta_log = log_p - log_q
    dim = 2 * jnp.cov(delta_log, aweights=weights)
    return dim