best-inference 0.1.0__tar.gz

best_inference-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Andreas Nygaard
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

best_inference-0.1.0/PKG-INFO

@@ -0,0 +1,223 @@
+Metadata-Version: 2.4
+Name: best-inference
+Version: 0.1.0
+Summary: Batched Emulator Sampling with Tensorflow
+Author-email: Andreas Nygaard <Andreas@phys.au.dk>
+License-Expression: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy<2.0,>=1.23
+Requires-Dist: tensorflow<2.21,>=2.16
+Requires-Dist: tensorflow-probability>=0.23
+Requires-Dist: tf-keras<2.21,>=2.16
+Requires-Dist: hypersphere-sampler
+Dynamic: license-file
+
+# BEST
+
+(**B**atched **E**mulator **S**ampling with **T**ensorFlow)
+
+A TensorFlow-based inference framework for high-performance Markov Chain Monte Carlo (MCMC) sampling, including support for neural likelihood emulators ("client emulators") and adaptive covariance estimation.
+
+---
+
+## Overview
+
+`best` provides a unified interface for sampling using multiple MCMC algorithms with GPU acceleration via TensorFlow:
+
+### Supported samplers
+
+- Metropolis-Hastings (MH)
+- Affine Invariant Ensemble Sampler (AIES)
+- Hamiltonian Monte Carlo (HMC)
+- No-U-Turn Sampler (NUTS)
+- Metropolis Adjusted Langevin Algorithm (MALA)
+
+### Key features
+
+- TensorFlow / GPU acceleration
+- Automatic covariance matrix adaptation
+- Bounded parameter inference
+- Parallelized multi-chain sampling
+- Pretrained neural likelihood emulators
+- JIT compilation (XLA support)
+
+---
+
+## Installation
+
+### From PyPI
+
+```bash
+pip install best-inference
+```
+### From source
+
+```bash
+git clone https://github.com/AndreasNygaard/best-inference.git
+cd best-inference
+pip install .
+```
+---
+## Quick start
+
+```python
+import best
+import tensorflow as tf
+
+def log_prob(x):
+    return -0.5 * tf.reduce_sum(x**2, axis=-1)
+
+sampler = best.Sampler(log_prob, bounds=([-5, -5], [5, 5]))
+
+results = sampler.sample(
+    method="hmc",
+    n_steps=2000,
+    n_chains=50,
+    initial_distribution="uniform",
+    num_burnin_steps=1000
+)
+
+print(results.samples.shape)
+```
+---
+## Sampler API
+
+### Initialisation
+
+```python
+best.Sampler(
+    log_prob_fn,
+    bounds=None,
+    enforce_boundaries=True,
+    covmat=None,
+    initial_state=None,
+    n_chains=10,
+    initial_distribution="repeat"
+)
+```
+
+### Sampling
+
+```python
+results = sampler.sample(
+    method="mh" | "aies" | "hmc" | "nuts" | "mala",
+    n_steps=1000,
+    n_chains=10,
+    initial_state=None,
+    initial_distribution="repeat" | "uniform" | "gaussian",
+    bounds=None,
+    covmat=None,
+    num_burnin_steps=100,
+    num_covmat_updates=None,
+    sampler_kwargs={},
+    burnin_kwargs={},
+    get_individual_chains=True,
+    jit_compile=True
+)
+```
+
+### Output
+
+```python
+results.samples
+results.acceptance_rate
+results.evaluations
+
+results.burnin_samples
+results.burnin_acceptance_rates
+results.burnin_evaluations
+results.covmat_estimate
+```
+
+## Client emulators
+
+BEST includes pretrained neural likelihood emulators for cosmology-inspired inference problems.
+
+### Available models
+
+ - lcdm
+ - sterile_neutrino
+
+### Load a model
+
+```python
+from best.client_emulators import load_model_and_scalers
+
+log_prob_fn, lower_bounds, upper_bounds = load_model_and_scalers("lcdm")
+```
+
+### Example: emulator-based inference
+
+```
+import best
+from best.client_emulators import load_model_and_scalers
+
+log_prob_fn, lower, upper = load_model_and_scalers("lcdm")
+
+sampler = best.Sampler(log_prob_fn, bounds=(lower, upper))
+
+results = sampler.sample(
+    method="aies",
+    n_steps=5000,
+    n_chains=100,
+    initial_distribution="uniform",
+    num_burnin_steps=2000,
+    num_covmat_updates=1
+)
+```
+
+## Supported Algorithms
+### Metropolis-Hastings (MH)
+Random-walk MCMC with optional adaptive covariance scaling.
+### Affine Invariant Ensemble Sampler (AIES)
+Efficient for highly anisotropic or correlated parameter spaces.
+### Hamiltonian Monte Carlo (HMC)
+Gradient-based sampling with leapfrog integration.
+### No-U-Turn Sampler (NUTS)
+Adaptive HMC variant with automatic trajectory length selection.
+### Metropolis Adjusted Langevin Algorithm (MALA)
+Gradient-informed diffusion-based sampler.
+
+## Performance Notes
+ - TensorFlow enables GPU acceleration where available
+ - JIT compilation (XLA) improves performance for large chains
+ - Vectorized multi-chain execution is used throughout
+ - Covariance estimation is performed during burn-in when enabled
+
+## Example: Multi-sampler comparison
+
+```python
+sampler.set_initial_state(initial_state=means, covmat=covmat, initial_distribution="gaussian")
+res_aies = sampler.sample(method="aies", n_steps=5000, n_chains=100)
+res_hmc  = sampler.sample(method="hmc",  n_steps=5000, n_chains=100)
+res_nuts = sampler.sample(method="nuts", n_steps=5000, n_chains=100)
+res_mh   = sampler.sample(method="mh",   n_steps=5000, n_chains=100)
+res_mala = sampler.sample(method="mala", n_steps=5000, n_chains=100)
+```
+
+## Requirements
+ - Python ≥ 3.10
+ - TensorFlow ≥ 2.17
+ - TensorFlow Probability ≥ 0.24
+ - NumPy
+ - tf-keras
+ - hypersphere-sampler
+
+## Citation
+If you use this package, please cite:
+
+Citation will be added once the accompanying paper is available (arXiv preprint in preparation).
+
+## Contributing
+Contributions are welcome.
+###Steps:
+ - Fork repository
+ - Create feature branch
+ - Add tests in ```tests/```
+ - Submit pull request
+
+## License
+#### MIT License
+Copyright (c) 2026 Andreas Nygaard

best_inference-0.1.0/README.md

@@ -0,0 +1,207 @@
+# BEST
+
+(**B**atched **E**mulator **S**ampling with **T**ensorFlow)
+
+A TensorFlow-based inference framework for high-performance Markov Chain Monte Carlo (MCMC) sampling, including support for neural likelihood emulators ("client emulators") and adaptive covariance estimation.
+
+---
+
+## Overview
+
+`best` provides a unified interface for sampling using multiple MCMC algorithms with GPU acceleration via TensorFlow:
+
+### Supported samplers
+
+- Metropolis-Hastings (MH)
+- Affine Invariant Ensemble Sampler (AIES)
+- Hamiltonian Monte Carlo (HMC)
+- No-U-Turn Sampler (NUTS)
+- Metropolis Adjusted Langevin Algorithm (MALA)
+
+### Key features
+
+- TensorFlow / GPU acceleration
+- Automatic covariance matrix adaptation
+- Bounded parameter inference
+- Parallelized multi-chain sampling
+- Pretrained neural likelihood emulators
+- JIT compilation (XLA support)
+
+---
+
+## Installation
+
+### From PyPI
+
+```bash
+pip install best-inference
+```
+### From source
+
+```bash
+git clone https://github.com/AndreasNygaard/best-inference.git
+cd best-inference
+pip install .
+```
+---
+## Quick start
+
+```python
+import best
+import tensorflow as tf
+
+def log_prob(x):
+    return -0.5 * tf.reduce_sum(x**2, axis=-1)
+
+sampler = best.Sampler(log_prob, bounds=([-5, -5], [5, 5]))
+
+results = sampler.sample(
+    method="hmc",
+    n_steps=2000,
+    n_chains=50,
+    initial_distribution="uniform",
+    num_burnin_steps=1000
+)
+
+print(results.samples.shape)
+```
+---
+## Sampler API
+
+### Initialisation
+
+```python
+best.Sampler(
+    log_prob_fn,
+    bounds=None,
+    enforce_boundaries=True,
+    covmat=None,
+    initial_state=None,
+    n_chains=10,
+    initial_distribution="repeat"
+)
+```
+
+### Sampling
+
+```python
+results = sampler.sample(
+    method="mh" | "aies" | "hmc" | "nuts" | "mala",
+    n_steps=1000,
+    n_chains=10,
+    initial_state=None,
+    initial_distribution="repeat" | "uniform" | "gaussian",
+    bounds=None,
+    covmat=None,
+    num_burnin_steps=100,
+    num_covmat_updates=None,
+    sampler_kwargs={},
+    burnin_kwargs={},
+    get_individual_chains=True,
+    jit_compile=True
+)
+```
+
+### Output
+
+```python
+results.samples
+results.acceptance_rate
+results.evaluations
+
+results.burnin_samples
+results.burnin_acceptance_rates
+results.burnin_evaluations
+results.covmat_estimate
+```
+
+## Client emulators
+
+BEST includes pretrained neural likelihood emulators for cosmology-inspired inference problems.
+
+### Available models
+
+ - lcdm
+ - sterile_neutrino
+
+### Load a model
+
+```python
+from best.client_emulators import load_model_and_scalers
+
+log_prob_fn, lower_bounds, upper_bounds = load_model_and_scalers("lcdm")
+```
+
+### Example: emulator-based inference
+
+```
+import best
+from best.client_emulators import load_model_and_scalers
+
+log_prob_fn, lower, upper = load_model_and_scalers("lcdm")
+
+sampler = best.Sampler(log_prob_fn, bounds=(lower, upper))
+
+results = sampler.sample(
+    method="aies",
+    n_steps=5000,
+    n_chains=100,
+    initial_distribution="uniform",
+    num_burnin_steps=2000,
+    num_covmat_updates=1
+)
+```
+
+## Supported Algorithms
+### Metropolis-Hastings (MH)
+Random-walk MCMC with optional adaptive covariance scaling.
+### Affine Invariant Ensemble Sampler (AIES)
+Efficient for highly anisotropic or correlated parameter spaces.
+### Hamiltonian Monte Carlo (HMC)
+Gradient-based sampling with leapfrog integration.
+### No-U-Turn Sampler (NUTS)
+Adaptive HMC variant with automatic trajectory length selection.
+### Metropolis Adjusted Langevin Algorithm (MALA)
+Gradient-informed diffusion-based sampler.
+
+## Performance Notes
+ - TensorFlow enables GPU acceleration where available
+ - JIT compilation (XLA) improves performance for large chains
+ - Vectorized multi-chain execution is used throughout
+ - Covariance estimation is performed during burn-in when enabled
+
+## Example: Multi-sampler comparison
+
+```python
+sampler.set_initial_state(initial_state=means, covmat=covmat, initial_distribution="gaussian")
+res_aies = sampler.sample(method="aies", n_steps=5000, n_chains=100)
+res_hmc  = sampler.sample(method="hmc",  n_steps=5000, n_chains=100)
+res_nuts = sampler.sample(method="nuts", n_steps=5000, n_chains=100)
+res_mh   = sampler.sample(method="mh",   n_steps=5000, n_chains=100)
+res_mala = sampler.sample(method="mala", n_steps=5000, n_chains=100)
+```
+
+## Requirements
+ - Python ≥ 3.10
+ - TensorFlow ≥ 2.17
+ - TensorFlow Probability ≥ 0.24
+ - NumPy
+ - tf-keras
+ - hypersphere-sampler
+
+## Citation
+If you use this package, please cite:
+
+Citation will be added once the accompanying paper is available (arXiv preprint in preparation).
+
+## Contributing
+Contributions are welcome.
+###Steps:
+ - Fork repository
+ - Create feature branch
+ - Add tests in ```tests/```
+ - Submit pull request
+
+## License
+#### MIT License
+Copyright (c) 2026 Andreas Nygaard

best_inference-0.1.0/best/__init__.py

@@ -0,0 +1,13 @@
+import os
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'
+os.environ["ABSL_MIN_LOG_LEVEL"] = "3"
+
+from .run_sampling import Sampler
+from . import mcmc_methods, tools, client_emulators
+
+__all__ = [
+    "Sampler",
+    "mcmc_methods",
+    "tools",
+    "client_emulators"
+]

best_inference-0.1.0/best/client_emulators/__init__.py

@@ -0,0 +1,7 @@
+from .load_client_network import load_model_and_scalers
+from . import custom_objects
+
+__all__ = [
+    "load_model_and_scalers",
+    "custom_objects"
+]

best_inference-0.1.0/best/client_emulators/custom_objects/__init__.py

@@ -0,0 +1,87 @@
+import tensorflow as tf
+import numpy as np
+
+class CustomTanh(tf.keras.layers.Layer):
+    def __init__(self, initial_alpha=1.0, **kwargs):
+        super().__init__(**kwargs)
+        self.initial_alpha = initial_alpha
+
+    def build(self, input_shape):
+        self.alpha = self.add_weight(
+            name='alpha',
+            shape=(1,),
+            initializer=tf.keras.initializers.Constant(self.initial_alpha),
+            trainable=True
+        )
+        super().build(input_shape)
+
+    @tf.function(jit_compile=True)
+    def call(self, inputs):
+        return tf.math.tanh(self.alpha * inputs)
+    
+    def get_config(self):
+        config = super().get_config()
+        config.update({
+            "initial_alpha": self.initial_alpha
+        })
+        return config
+    
+class Alsing(tf.keras.layers.Layer):
+    def __init__(self, initial_beta=1.0, initial_gamma=0.0, **kwargs):
+        super().__init__(**kwargs)
+        self.initial_beta = initial_beta
+        self.initial_gamma = initial_gamma
+
+    def build(self, input_shape):
+        self.beta = self.add_weight(
+            name="beta",
+            shape=(1,),
+            initializer=tf.keras.initializers.Constant(self.initial_beta),
+            trainable=True
+        )
+        self.gamma = self.add_weight(
+            name="gamma",
+            shape=(1,),
+            initializer=tf.keras.initializers.Constant(self.initial_gamma),
+            trainable=True
+        )
+        super().build(input_shape)
+
+    @tf.function(jit_compile=True)
+    def call(self, inputs):
+        return (self.gamma + (1 - self.gamma) / (1 + tf.exp(-self.beta * inputs))) * inputs
+
+    def get_config(self):
+        config = super().get_config()
+        config.update({
+            "initial_beta": self.initial_beta,
+            "initial_gamma": self.initial_gamma
+        })
+        return config
+
+def delta_chi2_from_k(kappa, n):
+    kappa = tf.cast(kappa, tf.float32)
+    n = tf.cast(n, tf.float32)
+
+    mu = 1.0 - 2.0 / (9.0 * n)
+    sigma = tf.sqrt(2.0 / (9.0 * n))
+    return n * tf.pow(mu + kappa * sigma, 3.0)
+
+def create_msre_loss(y_global_max, kappa, n, y_std):
+    delta_chi2_k = delta_chi2_from_k(kappa, n)
+    delta = -y_global_max - 0.5 * (delta_chi2_k / y_std)
+    
+    def mean_square_relative_error(y_true, y_pred):
+        denominator = y_true + delta
+        relative_error = (y_pred - y_true) / denominator
+        return tf.reduce_mean(tf.square(relative_error))
+    
+    return mean_square_relative_error
+
+
+def get_loss_function(loss_name, y_global_max=None, kappa=None, n=None, y_std=None):
+    if loss_name == 'msre':
+        return create_msre_loss(y_global_max, kappa, n, y_std)
+    if loss_name in ['mse', 'mae']:
+        return loss_name
+    return loss_name

best_inference-0.1.0/best/client_emulators/load_client_network.py

@@ -0,0 +1,51 @@
+import os
+import pickle as pkl
+
+import tensorflow as tf
+import numpy as np
+
+from best.client_emulators.custom_objects import Alsing, CustomTanh, create_msre_loss
+
+def load_model_and_scalers(emulator_dir, root='best/client_emulators'):
+    model_path = os.path.join(root, emulator_dir, 'trained_model.keras')
+    x_scaler_path = os.path.join(root, emulator_dir, 'x_scaler.pkl')
+    y_scaler_path = os.path.join(root, emulator_dir, 'y_scaler.pkl')
+
+    with open(x_scaler_path, 'rb') as f:
+        x_scaler = pkl.load(f)
+    with open(y_scaler_path, 'rb') as f:
+        y_scaler = pkl.load(f)
+
+    x_mean = tf.constant(x_scaler.mean_, dtype=tf.float32)
+    x_scale = tf.constant(x_scaler.scale_, dtype=tf.float32)
+    y_mean = tf.constant(y_scaler.mean_, dtype=tf.float32)
+    y_scale = tf.constant(y_scaler.scale_, dtype=tf.float32)
+
+    mean_square_relative_error = create_msre_loss(y_global_max=10.0, kappa=3.0, n=27.0, y_std=y_scale)
+    custom_objects={'Alsing': Alsing, 'CustomTanh': CustomTanh, 'mean_square_relative_error': mean_square_relative_error}
+
+    model = tf.keras.models.load_model(model_path, custom_objects=custom_objects)
+    
+    ## define box function that is exponentially increasing outside limits. It should be auto-differentiable.
+    lower = tf.constant([
+        0, 0, 0, 0, 0, 0.004, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.9
+    ], dtype=tf.float32)
+
+    upper = tf.constant([
+        3, 0.5, 2, 5, 2, 1, 0.8, 3.0, 200, 1, 10, 400, 400, 400, 400, 10, 50, 50, 100, 400, 10, 10, 10, 10, 10, 10, 3000, 3000, 1.1
+    ], dtype=tf.float32)
+
+    # Box is defined for the sterile neutrino model. If LCDM is loaded instead, two dimensions should be removed
+    if 'lcdm' in emulator_dir:
+        # remove indices 6 and 7 from lower and upper bounds for LCDM model
+        lower = tf.concat([lower[:6], lower[8:]], axis=0)
+        upper = tf.concat([upper[:6], upper[8:]], axis=0)
+
+    @tf.function(reduce_retracing=True)
+    def log_prob_fn(x):
+        inp = (x - x_mean) / x_scale
+        y_scaled = model(inp)
+        log_like = tf.reshape(y_scaled * y_scale + y_mean, [-1])
+        return log_like
+    
+    return log_prob_fn, lower, upper