Bayesian Models

Overview

The piblin_jax.bayesian module provides a powerful framework for Bayesian inference and uncertainty quantification in piblin_jax. Built on NumPyro, it enables probabilistic modeling with Markov Chain Monte Carlo (MCMC) sampling for robust parameter estimation and predictive uncertainty.

The Bayesian approach offers several advantages over traditional curve fitting:

  • Full Uncertainty Quantification: Instead of point estimates, Bayesian methods provide complete posterior distributions over parameters. This captures parameter correlations and enables principled uncertainty propagation to predictions.

  • Prior Knowledge Integration: Domain expertise and physical constraints can be incorporated through informative priors, improving estimates when data is limited or noisy.

  • Model Comparison: Bayesian model evidence enables rigorous comparison of competing models, helping you select the most appropriate model for your data.

  • Predictive Distributions: Obtain not just predicted values but full predictive distributions, capturing both aleatoric (data) and epistemic (parameter) uncertainty.

The module includes pre-built models for common rheological relationships (Power Law, Cross, Carreau-Yasuda) and thermal activation (Arrhenius), with a flexible base class for creating custom Bayesian models. All models are JIT-compiled for efficient MCMC sampling.

Quick Examples

Fitting a Power Law Model

Fit a rheological power law model with MCMC:

from piblin_jax.bayesian import PowerLawModel
import numpy as np

# Prepare data
shear_rate = np.logspace(-2, 3, 50)
viscosity = np.array([...])  # Your measurement data

# Create and fit model
model = PowerLawModel()
model.fit(
    x=shear_rate,
    y=viscosity,
    num_warmup=1000,
    num_samples=2000
)

# Get parameter posteriors
samples = model.get_samples()
print(f"Consistency index K: {samples['K'].mean():.3f} +/- {samples['K'].std():.3f}")
print(f"Flow index n: {samples['n'].mean():.3f} +/- {samples['n'].std():.3f}")

Making Predictions with Uncertainty

Generate predictions with full uncertainty quantification:

# Predict on new data
x_new = np.logspace(-2, 3, 100)
predictions = model.predict(x_new)

# Extract statistics
y_mean = predictions['mean']
y_std = predictions['std']
y_lower = predictions['quantiles'][0.025]  # 2.5th percentile
y_upper = predictions['quantiles'][0.975]  # 97.5th percentile

# Plot with uncertainty bands
import matplotlib.pyplot as plt
plt.fill_between(x_new, y_lower, y_upper, alpha=0.3, label='95% CI')
plt.plot(x_new, y_mean, label='Mean prediction')
plt.scatter(shear_rate, viscosity, label='Data')
plt.legend()

Custom Priors and Advanced Usage

Customize priors based on domain knowledge:

from piblin_jax.bayesian import ArrheniusModel

# Create model with custom priors
model = ArrheniusModel()

# Fit with custom MCMC settings
model.fit(
    x=temperature,
    y=reaction_rate,
    num_warmup=2000,
    num_samples=5000,
    num_chains=4,  # Parallel chains for convergence diagnostics
    target_accept_prob=0.9  # Higher acceptance for difficult posteriors
)

# Access full MCMC diagnostics
mcmc_info = model.get_mcmc_info()
print(f"Effective sample size: {mcmc_info['ess']}")
print(f"R-hat (convergence): {mcmc_info['r_hat']}")

See Also

API Reference

Module Contents

The piblin_jax.bayesian module provides the following classes:

  • BayesianModel - Base class for all Bayesian models

  • PowerLawModel - Power-law rheological model

  • CrossModel - Cross rheological model

  • CarreauYasudaModel - Carreau-Yasuda rheological model

  • ArrheniusModel - Arrhenius thermal activation model

Base Classes

Base class for Bayesian models using NumPyro.

This module provides the abstract base class for all Bayesian models in piblin-jax. Models use NumPyro for MCMC sampling and uncertainty quantification.

class piblin_jax.bayesian.base.BayesianModel(n_samples=1000, n_warmup=500, n_chains=2, random_seed=0)[source]

Bases: ABC

Abstract base class for Bayesian models using NumPyro.

This class provides the infrastructure for Bayesian inference using MCMC sampling via NumPyro. Subclasses implement the model() method to define the probabilistic model structure.

Parameters:

  • n_samples (int, optional) – Number of MCMC samples to draw (default: 1000)

  • n_warmup (int, optional) – Number of warmup samples for MCMC (default: 500)

  • n_chains (int, optional) – Number of MCMC chains to run (default: 2)

  • random_seed (int, optional) – Random seed for reproducibility (default: 0)

n_samples

Number of MCMC samples

Type:

int

n_warmup

Number of warmup samples

Type:

int

n_chains

Number of MCMC chains

Type:

int

random_seed

Random seed for PRNG

Type:

int

samples

Posterior samples from MCMC (None before fitting)

Type:

dict[str, array] | None

Examples

>>> import numpy as np
>>> import numpyro
>>> import numpyro.distributions as dist
>>> from piblin_jax.bayesian.base import BayesianModel
>>>
>>> class LinearRegressionModel(BayesianModel):
...     def model(self, x, y=None):
...         # Define priors
...         slope = numpyro.sample('slope', dist.Normal(0, 10))
...         intercept = numpyro.sample('intercept', dist.Normal(0, 10))
...         sigma = numpyro.sample('sigma', dist.HalfNormal(1))
...
...         # Define likelihood
...         mu = slope * x + intercept
...         with numpyro.plate('data', x.shape[0]):
...             numpyro.sample('obs', dist.Normal(mu, sigma), obs=y)
...
...     def predict(self, x, credible_interval=0.95):
...         # Generate predictions from posterior
...         slope_samples = self._samples['slope']
...         intercept_samples = self._samples['intercept']
...         predictions = slope_samples[:, None] * x + intercept_samples[:, None]
...         return {
...             'mean': np.mean(predictions, axis=0),
...             'lower': np.percentile(predictions, 2.5, axis=0),
...             'upper': np.percentile(predictions, 97.5, axis=0)
...         }
>>>
>>> # Create and fit model
>>> x = np.linspace(0, 10, 50)
>>> y = 2.0 * x + 1.0 + 0.1 * np.random.randn(len(x))
>>> model = LinearRegressionModel(n_samples=1000, n_warmup=500)
>>> model.fit(x, y)
>>> predictions = model.predict(x)

Notes

This class cannot be instantiated directly. Subclasses must implement: - model(x, y=None, **kwargs): Define the NumPyro probabilistic model - predict(x, credible_interval=0.95): Generate predictions with uncertainty

The MCMC sampler uses the No-U-Turn Sampler (NUTS) algorithm, which is an efficient Hamiltonian Monte Carlo variant that automatically tunes the step size and number of steps.

Attributes:
samples

Get posterior samples from MCMC.

Methods

fit(x, y[, use_nlsq_init])

Fit the model using MCMC sampling.

get_credible_intervals(param_name[, level, ...])

Get credible intervals for a parameter.

model(x[, y])

Define the NumPyro probabilistic model.

predict(x[, credible_interval])

Generate predictions with uncertainty.

summary()

Get summary statistics for all parameters.
__init__(n_samples=1000, n_warmup=500, n_chains=2, random_seed=0)[source]

Initialize BayesianModel.

Parameters:

  • n_samples (int, optional) – Number of MCMC samples (default: 1000)

  • n_warmup (int, optional) – Number of warmup samples (default: 500)

  • n_chains (int, optional) – Number of MCMC chains (default: 2)

  • random_seed (int, optional) – Random seed (default: 0)

abstractmethod model(x, y=None, **kwargs)[source]

Define the NumPyro probabilistic model.

Subclasses must implement this method to specify the model structure, including priors and likelihood.

Parameters:

  • x (array_like) – Independent variable (input data)

  • y (array_like | None, optional) – Dependent variable (observations, None for prediction)

  • **kwargs (dict) – Additional model-specific parameters

Notes

This method should use NumPyro’s sample primitive to define random variables. When y is not None, it should be used as the observation in a sample call with obs=y.

Examples

>>> def model(self, x, y=None):
...     # Define priors
...     slope = numpyro.sample('slope', dist.Normal(0, 10))
...     intercept = numpyro.sample('intercept', dist.Normal(0, 10))
...     sigma = numpyro.sample('sigma', dist.HalfNormal(1))
...
...     # Define likelihood
...     mu = slope * x + intercept
...     numpyro.sample('obs', dist.Normal(mu, sigma), obs=y)
fit(x, y, use_nlsq_init=False, **kwargs)[source]

Fit the model using MCMC sampling.

Runs MCMC using the No-U-Turn Sampler (NUTS) to sample from the posterior distribution. Stores the posterior samples internally for later prediction and inference.

Parameters:

  • x (array_like) – Independent variable (input data)

  • y (array_like) – Dependent variable (observations)

  • use_nlsq_init (bool, default False) – If True, use NLSQ to get initial parameter estimates for better prior centering (experimental feature).

  • **kwargs (dict) – Additional model-specific parameters passed to model()

Returns:

Returns self for method chaining

Return type:

BayesianModel

Examples

>>> model = MyBayesianModel(n_samples=1000, n_warmup=500)
>>> model.fit(x_data, y_data)
>>> predictions = model.predict(x_new)

Notes

The use_nlsq_init parameter is experimental and may not work for all model types. It attempts to use nonlinear least squares to find good initial parameter estimates, which can help with MCMC convergence.

abstractmethod predict(x, credible_interval=0.95)[source]

Generate predictions with uncertainty.

Subclasses must implement this method to generate predictions using the posterior samples from MCMC.

Parameters:

  • x (array_like) – Points to predict at

  • credible_interval (float, optional) – Credible interval level (default: 0.95)

Returns:

Dictionary with keys: - ‘mean’: Mean prediction - ‘lower’: Lower credible bound - ‘upper’: Upper credible bound - ‘samples’: Full posterior predictive samples (optional)

Return type:

dict

Raises:

RuntimeError – If model has not been fit yet

Examples

>>> predictions = model.predict(x_new, credible_interval=0.95)
>>> mean = predictions['mean']
>>> lower = predictions['lower']
>>> upper = predictions['upper']
property samples: dict[str, ndarray] | None

Get posterior samples from MCMC.

No-index:

Returns:

Dictionary mapping parameter names to arrays of samples, or None if model has not been fit yet.

Return type:

dict[str, np.ndarray] | None

Examples

>>> model.fit(x, y)
>>> samples = model.samples
>>> slope_samples = samples['slope']
>>> print(f"Slope mean: {np.mean(slope_samples)}")
get_credible_intervals(param_name, level=0.95, method='eti')[source]

Get credible intervals for a parameter.

Computes credible intervals from the posterior samples using either equal-tailed intervals (ETI) or highest posterior density (HPD).

Parameters:

  • param_name (str) – Name of the parameter (must match a sample name from model)

  • level (float, optional) – Credible interval level between 0 and 1 (default: 0.95)

  • method (str, optional) – Method for computing intervals: - ‘eti’: Equal-tailed interval (default) - ‘hpd’: Highest posterior density interval

Returns:

(lower_bound, upper_bound) of the credible interval

Return type:

tuple[float, float]

Raises:

Examples

>>> model.fit(x, y)
>>> lower, upper = model.get_credible_intervals('slope', level=0.95)
>>> print(f"95% credible interval for slope: [{lower:.3f}, {upper:.3f}]")

>>> # Use 68% interval (approximately 1 sigma)
>>> lower, upper = model.get_credible_intervals('slope', level=0.68)

Notes

The ETI method uses percentiles and is simple to compute but may not give the shortest interval for skewed distributions.

The HPD method finds the shortest interval containing the specified probability mass, but this is a simplified implementation. For full HPD computation, consider using the arviz library.

summary()[source]

Get summary statistics for all parameters.

Returns:

Dictionary mapping parameter names to summary statistics: - ‘mean’: Posterior mean - ‘std’: Posterior standard deviation - ‘q_2.5’: 2.5th percentile - ‘q_50’: Median (50th percentile) - ‘q_97.5’: 97.5th percentile

Return type:

dict

Raises:

RuntimeError – If model has not been fit yet

Examples

>>> model.fit(x, y)
>>> summary = model.summary()
>>> print(summary['slope'])
{'mean': 2.01, 'std': 0.15, 'q_2.5': 1.72, 'q_50': 2.00, 'q_97.5': 2.31}

Rheological Models

Power Law Model

Power-law viscosity model for rheological analysis.

This module implements the power-law (Ostwald-de Waele) model for shear-thinning and shear-thickening fluids using Bayesian inference.

class piblin_jax.bayesian.models.power_law.PowerLawModel(n_samples=1000, n_warmup=500, n_chains=2, random_seed=0)[source]

Bases: BayesianModel

Power-law viscosity model for non-Newtonian fluids.

The power-law model describes the relationship between viscosity and shear rate:

η(γ̇) = K * γ̇^(n-1)

where:

  • η is the viscosity (Pa·s)

  • γ̇ is the shear rate (s⁻¹)

  • K is the consistency index (Pa·s^n)

  • n is the power-law index (dimensionless)

The power-law index n characterizes the flow behavior:

  • n < 1: Shear-thinning (pseudoplastic) behavior

  • n = 1: Newtonian behavior

  • n > 1: Shear-thickening (dilatant) behavior

Parameters:

  • n_samples (int, optional) – Number of MCMC samples to draw (default: 1000)

  • n_warmup (int, optional) – Number of warmup samples for MCMC (default: 500)

  • n_chains (int, optional) – Number of MCMC chains to run (default: 2)

  • random_seed (int, optional) – Random seed for reproducibility (default: 0)

samples

Posterior samples from MCMC containing: - ‘K’: Consistency index samples - ‘n’: Power-law index samples - ‘sigma’: Observation noise samples

Type:

dict[str, array] | None

Examples

>>> import numpy as np
>>> from piblin_jax.bayesian.models import PowerLawModel
>>>
>>> # Generate synthetic power-law data
>>> shear_rate = np.logspace(-1, 2, 30)  # 0.1 to 100 s^-1
>>> viscosity = 5.0 * shear_rate ** (0.6 - 1)  # K=5, n=0.6 (shear-thinning)
>>>
>>> # Fit model
>>> model = PowerLawModel(n_samples=1000, n_warmup=500)
>>> model.fit(shear_rate, viscosity)
>>>
>>> # Get parameter estimates
>>> summary = model.summary()
>>> print(f"K: {summary['K']['mean']:.2f} +/- {summary['K']['std']:.2f}")
>>> print(f"n: {summary['n']['mean']:.2f} +/- {summary['n']['std']:.2f}")
>>>
>>> # Predict with uncertainty
>>> shear_rate_new = np.array([1.0, 10.0, 50.0])
>>> predictions = model.predict(shear_rate_new, credible_interval=0.95)
>>> print(f"Predicted viscosity at γ̇=10: {predictions['mean'][1]:.2f}")
>>> print(f"95% CI: [{predictions['lower'][1]:.2f}, {predictions['upper'][1]:.2f}]")

Notes

The model uses the following priors:

  • K ~ LogNormal(0, 2): Ensures positive consistency index

  • n ~ Normal(0.5, 0.5): Centers around typical shear-thinning behavior

  • sigma ~ HalfNormal(1): Observation noise

The power-law model is simple but effective for many non-Newtonian fluids. However, it has limitations:

  • Predicts infinite viscosity at zero shear rate (for n < 1)

  • Predicts zero viscosity at infinite shear rate (for n < 1)

  • Does not capture zero-shear or infinite-shear plateaus

For fluids with plateau regions, consider using CrossModel or CarreauYasudaModel.

References

[1]

Ostwald, W. (1925). “Über die rechnerische Darstellung des Strukturgebietes der Viskosität.” Kolloid-Zeitschrift, 36, 99-117.

[2]

de Waele, A. (1923). “Viscometry and plastometry.” Journal of the Oil and Colour Chemists Association, 6, 33-69.

Attributes:
samples

Get posterior samples from MCMC.

Methods

fit(x, y[, use_nlsq_init])

Fit the model using MCMC sampling.

get_credible_intervals(param_name[, level, ...])

Get credible intervals for a parameter.

model(x[, y])

Define the NumPyro probabilistic model for power-law viscosity.

predict(shear_rate[, credible_interval])

Predict viscosity with uncertainty at given shear rates.

summary()

Get summary statistics for all parameters.
model(x, y=None, **kwargs)[source]

Define the NumPyro probabilistic model for power-law viscosity.

This method is called internally by the MCMC inference engine and defines the probabilistic generative process for power-law viscosity data.

Parameters:

  • x (array_like) – Shear rate data (γ̇) in s⁻¹

  • y (array_like | None, optional) – Viscosity observations (η) in Pa·s. If None, generates prior samples.

  • **kwargs (dict) – Additional model parameters (unused)

Examples

This method is typically not called directly. Instead, use the fit() method:

>>> model = PowerLawModel()
>>> model.fit(shear_rate, viscosity)  # Internally calls model()

Notes

This method is called internally by fit() and should not be called directly. It defines the generative model η = K * γ̇^(n-1) + ε where ε ~ Normal(0, σ).

predict(shear_rate, credible_interval=0.95)[source]

Predict viscosity with uncertainty at given shear rates.

Uses posterior samples from MCMC to generate predictions with credible intervals.

Parameters:

  • shear_rate (array_like) – Shear rate values (γ̇) in s⁻¹ at which to predict viscosity

  • credible_interval (float, optional) – Credible interval level between 0 and 1 (default: 0.95)

Returns:

Dictionary containing: - ‘mean’: Mean predicted viscosity (array) - ‘lower’: Lower credible bound (array) - ‘upper’: Upper credible bound (array) - ‘samples’: Full posterior predictive samples (2D array)

Return type:

dict

Raises:

RuntimeError – If model has not been fit yet

Examples

>>> model = PowerLawModel()
>>> model.fit(shear_rate_data, viscosity_data)
>>> predictions = model.predict(np.array([1.0, 10.0, 100.0]))
>>> print(predictions['mean'])
[5.23 2.41 1.11]
>>> print(predictions['lower'])
[4.89 2.21 1.01]
>>> print(predictions['upper'])
[5.61 2.65 1.23]

Cross Model

Cross viscosity model for rheological analysis.

This module implements the Cross model for shear-thinning fluids with zero-shear and infinite-shear plateaus using Bayesian inference.

class piblin_jax.bayesian.models.cross.CrossModel(n_samples=1000, n_warmup=500, n_chains=2, random_seed=0)[source]

Bases: BayesianModel

Cross viscosity model for shear-thinning fluids.

The Cross model describes the transition from zero-shear viscosity to infinite-shear viscosity:

η(γ̇) = η∞ + (η₀ - η∞) / (1 + (λγ̇)^m)

where:

  • η is the viscosity (Pa·s)

  • γ̇ is the shear rate (s⁻¹)

  • η₀ is the zero-shear viscosity (Pa·s)

  • η∞ is the infinite-shear viscosity (Pa·s)

  • λ is the time constant (s)

  • m is the power-law exponent (dimensionless)

At low shear rates (when λγ̇ << 1): η → η₀ (zero-shear plateau) At high shear rates (when λγ̇ >> 1): η → η∞ (infinite-shear plateau)

Parameters:

  • n_samples (int, optional) – Number of MCMC samples to draw (default: 1000)

  • n_warmup (int, optional) – Number of warmup samples for MCMC (default: 500)

  • n_chains (int, optional) – Number of MCMC chains to run (default: 2)

  • random_seed (int, optional) – Random seed for reproducibility (default: 0)

samples

Posterior samples from MCMC containing: - ‘eta0’: Zero-shear viscosity samples - ‘eta_inf’: Infinite-shear viscosity samples - 'lambda_': Time constant samples - ‘m’: Power-law exponent samples - ‘sigma’: Observation noise samples

Type:

dict[str, array] | None

Examples

>>> import numpy as np
>>> from piblin_jax.bayesian.models import CrossModel
>>>
>>> # Shear-thinning data with plateaus
>>> shear_rate = np.logspace(-2, 3, 50)  # 0.01 to 1000 s^-1
>>> # Simulate Cross model: eta0=100, eta_inf=1, lambda_=1, m=0.7
>>> viscosity = 1 + (100 - 1) / (1 + (1.0 * shear_rate)**0.7)
>>>
>>> # Fit Cross model
>>> model = CrossModel(n_samples=1000, n_warmup=500)
>>> model.fit(shear_rate, viscosity)
>>>
>>> # Get parameter estimates
>>> summary = model.summary()
>>> print(f"Zero-shear viscosity: {summary['eta0']['mean']:.1f} Pa·s")
>>> print(f"Infinite-shear viscosity: {summary['eta_inf']['mean']:.2f} Pa·s")
>>>
>>> # Predict with uncertainty
>>> shear_rate_new = np.array([0.1, 1.0, 10.0, 100.0])
>>> predictions = model.predict(shear_rate_new)

Notes

The model uses the following priors:

  • eta0 ~ LogNormal(4, 2): Zero-shear viscosity (centered around e^4 ≈ 55)

  • eta_inf ~ LogNormal(0, 2): Infinite-shear viscosity (centered around 1)

  • lambda_ ~ LogNormal(0, 2): Time constant (centered around 1)

  • m ~ Normal(0.7, 0.3): Power-law exponent (typical range 0.3-1.0)

  • sigma ~ HalfNormal(scale): Observation noise

The Cross model advantages:

  • Captures both low and high shear rate plateaus

  • More physically realistic than simple power-law

  • Four parameters provide flexibility

The model assumes:

  • Single relaxation time

  • No yield stress

  • Monotonic shear-thinning behavior

References

[1]

Cross, M. M. (1965). “Rheology of non-Newtonian fluids: A new flow equation for pseudoplastic systems.” Journal of Colloid Science, 20(5), 417-437.

[2]

Morrison, F. A. (2001). “Understanding Rheology.” Oxford University Press.

Attributes:
samples

Get posterior samples from MCMC.

Methods

fit(x, y[, use_nlsq_init])

Fit the model using MCMC sampling.

get_credible_intervals(param_name[, level, ...])

Get credible intervals for a parameter.

model(x[, y])

Define the NumPyro probabilistic model for Cross viscosity.

predict(shear_rate[, credible_interval])

Predict viscosity with uncertainty at given shear rates.

summary()

Get summary statistics for all parameters.
model(x, y=None, **kwargs)[source]

Define the NumPyro probabilistic model for Cross viscosity.

Parameters:

  • x (array_like) – Shear rate data (γ̇) in s⁻¹

  • y (array_like | None, optional) – Viscosity observations (η) in Pa·s. If None, generates prior samples.

  • **kwargs (dict) – Additional model parameters (unused)

Notes

This method is called internally by fit() and should not be called directly.

predict(shear_rate, credible_interval=0.95)[source]

Predict viscosity with uncertainty at given shear rates.

Uses posterior samples from MCMC to generate predictions with credible intervals.

Parameters:

  • shear_rate (array_like) – Shear rate values (γ̇) in s⁻¹ at which to predict viscosity

  • credible_interval (float, optional) – Credible interval level between 0 and 1 (default: 0.95)

Returns:

Dictionary containing: - ‘mean’: Mean predicted viscosity (array) - ‘lower’: Lower credible bound (array) - ‘upper’: Upper credible bound (array) - ‘samples’: Full posterior predictive samples (2D array)

Return type:

dict

Raises:

RuntimeError – If model has not been fit yet

Examples

>>> model = CrossModel()
>>> model.fit(shear_rate_data, viscosity_data)
>>> predictions = model.predict(np.array([0.01, 0.1, 1.0, 10.0]))
>>> print(predictions['mean'])
[98.5 85.3 45.2 12.1]

Carreau-Yasuda Model

Carreau-Yasuda viscosity model for rheological analysis.

This module implements the Carreau-Yasuda model, which generalizes the Carreau model for more flexible fitting of shear-thinning fluids.

class piblin_jax.bayesian.models.carreau_yasuda.CarreauYasudaModel(n_samples=1000, n_warmup=500, n_chains=2, random_seed=0)[source]

Bases: BayesianModel

Carreau-Yasuda viscosity model for shear-thinning fluids.

The Carreau-Yasuda model is a generalization of the Carreau model that provides better flexibility in the transition region:

η(γ̇) = η∞ + (η₀ - η∞) * [1 + (λγ̇)^a]^((n-1)/a)

where:

  • η is the viscosity (Pa·s)

  • γ̇ is the shear rate (s⁻¹)

  • η₀ is the zero-shear viscosity (Pa·s)

  • η∞ is the infinite-shear viscosity (Pa·s)

  • λ is the relaxation time (s)

  • a is the transition parameter (dimensionless)

  • n is the power-law index (dimensionless)

Special cases:

  • a → ∞: Reduces to power-law model in transition region

  • a = 2: Carreau model

  • a = 1: Cross model (approximately)

Parameters:

  • n_samples (int, optional) – Number of MCMC samples to draw (default: 1000)

  • n_warmup (int, optional) – Number of warmup samples for MCMC (default: 500)

  • n_chains (int, optional) – Number of MCMC chains to run (default: 2)

  • random_seed (int, optional) – Random seed for reproducibility (default: 0)

samples

Posterior samples from MCMC containing: - ‘eta0’: Zero-shear viscosity samples - ‘eta_inf’: Infinite-shear viscosity samples - 'lambda_': Relaxation time samples - ‘a’: Transition parameter samples - ‘n’: Power-law index samples - ‘sigma’: Observation noise samples

Type:

dict[str, array] | None

Examples

>>> import numpy as np
>>> from piblin_jax.bayesian.models import CarreauYasudaModel
>>>
>>> # Generate synthetic Carreau-Yasuda data
>>> shear_rate = np.logspace(-2, 3, 50)
>>> eta0, eta_inf, lam, a, n = 100.0, 1.0, 1.0, 2.0, 0.5
>>> viscosity = eta_inf + (eta0 - eta_inf) * (1 + (lam * shear_rate)**a)**((n-1)/a)
>>>
>>> # Fit model
>>> model = CarreauYasudaModel(n_samples=1000, n_warmup=500)
>>> model.fit(shear_rate, viscosity)
>>>
>>> # Get parameter estimates
>>> summary = model.summary()
>>> print(f"Zero-shear viscosity: {summary['eta0']['mean']:.1f} Pa·s")
>>> print(f"Power-law index: {summary['n']['mean']:.2f}")
>>>
>>> # Predict
>>> predictions = model.predict(np.array([0.1, 1.0, 10.0]))

Notes

The model uses the following priors:

  • eta0 ~ LogNormal(4, 2): Zero-shear viscosity

  • eta_inf ~ LogNormal(0, 2): Infinite-shear viscosity

  • lambda_ ~ LogNormal(0, 2): Relaxation time

  • a ~ LogNormal(0.69, 0.5): Transition parameter (centered around 2)

  • n ~ Normal(0.5, 0.3): Power-law index (shear-thinning range)

  • sigma ~ HalfNormal(scale): Observation noise

The Carreau-Yasuda model:

  • Five parameters provide maximum flexibility

  • Captures smooth transition between plateaus

  • Parameter ‘a’ controls transition sharpness

  • Reduces to simpler models as special cases

The model assumes:

  • Monotonic shear-thinning behavior

  • No yield stress

  • Single relaxation mechanism

For materials with yield stress, consider the Herschel-Bulkley model. For simpler analysis with fewer parameters, use Cross or Carreau models.

References

[1]

Yasuda, K., Armstrong, R. C., & Cohen, R. E. (1981). “Shear flow properties of concentrated solutions of linear and star branched polystyrenes.” Rheologica Acta, 20(2), 163-178.

[2]

Carreau, P. J. (1972). “Rheological equations from molecular network theories.” Transactions of the Society of Rheology, 16(1), 99-127.

[3]

Bird, R. B., Armstrong, R. C., & Hassager, O. (1987). “Dynamics of Polymeric Liquids. Vol. 1: Fluid Mechanics.” Wiley, New York.

Attributes:
samples

Get posterior samples from MCMC.

Methods

fit(x, y[, use_nlsq_init])

Fit the model using MCMC sampling.

get_credible_intervals(param_name[, level, ...])

Get credible intervals for a parameter.

model(x[, y])

Define the NumPyro probabilistic model for Carreau-Yasuda viscosity.

predict(shear_rate[, credible_interval])

Predict viscosity with uncertainty at given shear rates.

summary()

Get summary statistics for all parameters.
model(x, y=None, **kwargs)[source]

Define the NumPyro probabilistic model for Carreau-Yasuda viscosity.

Parameters:

  • x (array_like) – Shear rate data (γ̇) in s⁻¹

  • y (array_like | None, optional) – Viscosity observations (η) in Pa·s. If None, generates prior samples.

  • **kwargs (dict) – Additional model parameters (unused)

Notes

This method is called internally by fit() and should not be called directly.

predict(shear_rate, credible_interval=0.95)[source]

Predict viscosity with uncertainty at given shear rates.

Uses posterior samples from MCMC to generate predictions with credible intervals.

Parameters:

  • shear_rate (array_like) – Shear rate values (γ̇) in s⁻¹ at which to predict viscosity

  • credible_interval (float, optional) – Credible interval level between 0 and 1 (default: 0.95)

Returns:

Dictionary containing: - ‘mean’: Mean predicted viscosity (array) - ‘lower’: Lower credible bound (array) - ‘upper’: Upper credible bound (array) - ‘samples’: Full posterior predictive samples (2D array)

Return type:

dict

Raises:

RuntimeError – If model has not been fit yet

Examples

>>> model = CarreauYasudaModel()
>>> model.fit(shear_rate_data, viscosity_data)
>>> predictions = model.predict(np.array([0.1, 1.0, 10.0, 100.0]))
>>> print(predictions['mean'])
[95.3 48.2 12.5  2.8]

Thermal Models

Arrhenius Model

Arrhenius temperature-viscosity model for rheological analysis.

This module implements the Arrhenius equation for modeling temperature-dependent viscosity using Bayesian inference.

class piblin_jax.bayesian.models.arrhenius.ArrheniusModel(n_samples=1000, n_warmup=500, n_chains=2, random_seed=0)[source]

Bases: BayesianModel

Arrhenius temperature-viscosity model.

The Arrhenius equation describes how viscosity changes with temperature:

η(T) = A * exp(Ea / (R*T))

where:

  • η is the viscosity (Pa·s)

  • T is the absolute temperature (K)

  • A is the pre-exponential factor (Pa·s)

  • Ea is the activation energy (J/mol)

  • R is the universal gas constant (8.314 J/(mol·K))

This model is widely used for polymer melts, glass-forming liquids, and other materials where viscosity decreases exponentially with temperature.

Parameters:

  • n_samples (int, optional) – Number of MCMC samples to draw (default: 1000)

  • n_warmup (int, optional) – Number of warmup samples for MCMC (default: 500)

  • n_chains (int, optional) – Number of MCMC chains to run (default: 2)

  • random_seed (int, optional) – Random seed for reproducibility (default: 0)

samples

Posterior samples from MCMC containing: - ‘A’: Pre-exponential factor samples - ‘Ea’: Activation energy samples - ‘sigma’: Observation noise samples

Type:

dict[str, array] | None

R

Universal gas constant (8.314 J/(mol·K))

No-index:

Type:

float

Examples

>>> import numpy as np
>>> from piblin_jax.bayesian.models import ArrheniusModel
>>>
>>> # Temperature-dependent viscosity data
>>> temperature = np.array([300, 320, 340, 360, 380, 400])  # K
>>> viscosity = np.array([1000, 450, 220, 120, 70, 45])  # Pa·s
>>>
>>> # Fit Arrhenius model
>>> model = ArrheniusModel(n_samples=1000, n_warmup=500)
>>> model.fit(temperature, viscosity)
>>>
>>> # Get activation energy
>>> summary = model.summary()
>>> Ea_mean = summary['Ea']['mean']
>>> print(f"Activation energy: {Ea_mean/1000:.1f} kJ/mol")
>>>
>>> # Predict at new temperature
>>> temp_new = np.array([350])
>>> predictions = model.predict(temp_new)
>>> print(f"Predicted viscosity at 350K: {predictions['mean'][0]:.1f} Pa·s")

Notes

The model uses the following priors:

  • A ~ LogNormal(-10, 5): Wide prior on pre-exponential factor

  • Ea ~ Normal(50000, 30000): Prior centered around typical activation energies

  • sigma ~ HalfNormal(scale): Observation noise (scale = 10% of mean viscosity)

The Arrhenius equation assumes:

  • Activation energy is constant over the temperature range

  • Single relaxation mechanism (no structural transitions)

  • Newtonian behavior at each temperature

For materials with glass transitions or multiple relaxation processes, consider the Williams-Landel-Ferry (WLF) equation or Vogel-Fulcher-Tammann (VFT) equation.

References

[1]

Arrhenius, S. (1889). “Über die Reaktionsgeschwindigkeit bei der Inversion von Rohrzucker durch Säuren.” Zeitschrift für Physikalische Chemie, 4, 226-248.

[2]

Ferry, J. D. (1980). “Viscoelastic Properties of Polymers,” 3rd ed. Wiley, New York.

Attributes:
samples

Get posterior samples from MCMC.

Methods

fit(x, y[, use_nlsq_init])

Fit the model using MCMC sampling.

get_credible_intervals(param_name[, level, ...])

Get credible intervals for a parameter.

model(x[, y])

Define the NumPyro probabilistic model for Arrhenius viscosity.

predict(temperature[, credible_interval])

Predict viscosity with uncertainty at given temperatures.

summary()

Get summary statistics for all parameters.
R = 8.314
model(x, y=None, **kwargs)[source]

Define the NumPyro probabilistic model for Arrhenius viscosity.

This method is called internally by the MCMC inference engine and defines the probabilistic generative process for temperature-dependent viscosity.

Parameters:

  • x (array_like) – Temperature data (T) in Kelvin

  • y (array_like | None, optional) – Viscosity observations (η) in Pa·s. If None, generates prior samples.

  • **kwargs (dict) – Additional model parameters (unused)

Examples

This method is typically not called directly. Instead, use the fit() method:

>>> model = ArrheniusModel()
>>> model.fit(temperature, rate_constant)  # Internally calls model()

Notes

This method is called internally by fit() and should not be called directly. It defines the generative model η(T) = A * exp(Ea/(R*T)) + ε where ε ~ Normal(0, σ).

predict(temperature, credible_interval=0.95)[source]

Predict viscosity with uncertainty at given temperatures.

Uses posterior samples from MCMC to generate predictions with credible intervals.

Parameters:

  • temperature (array_like) – Temperature values (T) in Kelvin at which to predict viscosity

  • credible_interval (float, optional) – Credible interval level between 0 and 1 (default: 0.95)

Returns:

Dictionary containing: - ‘mean’: Mean predicted viscosity (array) - ‘lower’: Lower credible bound (array) - ‘upper’: Upper credible bound (array) - ‘samples’: Full posterior predictive samples (2D array)

Return type:

dict

Raises:

RuntimeError – If model has not been fit yet

Examples

>>> model = ArrheniusModel()
>>> model.fit(temp_data, viscosity_data)
>>> predictions = model.predict(np.array([300, 350, 400]))
>>> print(predictions['mean'])
[980.5 165.3  48.2]