"""Herschel-Bulkley model for non-Newtonian flow with yield stress.
This module implements the Herschel-Bulkley model, which combines yield stress
behavior with power-law flow. This is the most general viscoplastic model,
reducing to simpler models as special cases (ROTATION test mode).
Theory:
σ(γ̇) = σ_y + K ``|γ̇|`` ^n for σ > σ_y
γ̇ = 0 for σ ≤ σ_y
- σ_y: Yield stress (material flows only when σ > σ_y)
- K: Consistency index (viscosity-like parameter)
- n: Flow behavior index (power-law exponent)
References:
- Herschel, W.H., Bulkley, R. (1926). Proc. ASTM 26, 621-633.
"""
from __future__ import annotations
from rheojax.core.jax_config import safe_import_jax
jax, jnp = safe_import_jax()
import numpy as np
from rheojax.core.base import BaseModel, ParameterSet
from rheojax.core.data import RheoData
from rheojax.core.inventory import Protocol
from rheojax.core.registry import ModelRegistry
from rheojax.core.test_modes import DeformationMode, TestMode, detect_test_mode
from rheojax.logging import get_logger, log_fit
# Module logger
logger = get_logger(__name__)
[docs]
@ModelRegistry.register(
"herschel_bulkley",
protocols=[Protocol.FLOW_CURVE],
deformation_modes=[DeformationMode.SHEAR],
)
class HerschelBulkley(BaseModel):
"""Herschel-Bulkley model for viscoplastic flow (ROTATION only).
The Herschel-Bulkley model describes materials that require a minimum
stress (yield stress σ_y) to flow, and then exhibit power-law behavior.
This is widely used for pastes, slurries, and suspensions.
Parameters:
sigma_y: Yield stress (Pa), minimum stress required for flow
K: Consistency index (Pa·s^n), controls viscosity magnitude
n: Flow behavior index (dimensionless), power-law exponent
Constitutive Equation:
σ(γ̇) = σ_y + K ``|γ̇|`` ^n for ``|σ|`` > σ_y
γ̇ = 0 for ``|σ|`` ≤ σ_y
Special Cases:
σ_y = 0: Reduces to Power Law model
n = 1: Reduces to Bingham model (linear viscoplastic)
σ_y = 0, n = 1: Newtonian fluid
Test Mode:
ROTATION (steady shear) only
"""
[docs]
def __init__(self):
"""Initialize Herschel-Bulkley model."""
super().__init__()
self.parameters = ParameterSet()
self.parameters.add(
name="sigma_y",
value=10.0,
bounds=(0.0, 1e6),
units="Pa",
description="Yield stress",
)
self.parameters.add(
name="K",
value=1.0,
bounds=(1e-6, 1e6),
units="Pa·s^n",
description="Consistency index",
)
self.parameters.add(
name="n",
value=0.5,
bounds=(0.01, 2.0),
units="dimensionless",
description="Flow behavior index",
)
def _fit(self, X: np.ndarray, y: np.ndarray, **kwargs) -> HerschelBulkley:
"""Fit Herschel-Bulkley parameters to data.
Args:
X: Shear rate data (γ̇)
y: Stress data (σ)
**kwargs: Additional fitting options
Returns:
self for method chaining
"""
# P3-FLOW-001: Cache test_mode for Bayesian _resolve_test_mode() consistency
self._test_mode = kwargs.get("test_mode", "rotation")
data_shape = (len(X),) if hasattr(X, "__len__") else None
with log_fit(
logger, model="HerschelBulkley", data_shape=data_shape, test_mode="rotation"
) as ctx:
logger.debug(
"Starting Herschel-Bulkley model fit",
n_points=data_shape[0] if data_shape else None,
initial_sigma_y=self.parameters.get_value("sigma_y"),
initial_K=self.parameters.get_value("K"),
initial_n=self.parameters.get_value("n"),
)
try:
# Sort by shear rate
sort_idx = np.argsort(X)
X_sorted = X[sort_idx]
y_sorted = y[sort_idx]
logger.debug("Data sorted by shear rate", n_points=len(X_sorted))
# Estimate yield stress from low shear rate extrapolation
# σ_y ≈ stress at γ̇ → 0
sigma_y_est = np.min(y_sorted[: len(y_sorted) // 10 + 1])
if sigma_y_est < 0:
sigma_y_est = 0.0
logger.debug(
"Yield stress estimated from low shear rate region",
sigma_y_estimate=sigma_y_est,
)
# Subtract yield stress for power-law fitting
y_corrected = y_sorted - sigma_y_est
y_corrected = np.maximum(y_corrected, 1e-10) # Avoid log(0)
# Fit power law to corrected data: log(σ - σ_y) = log(K) + n*log(γ̇)
# Use middle to high shear rate region
start_idx = len(X_sorted) // 4
log_gamma = np.log(np.maximum(np.abs(X_sorted[start_idx:]), 1e-30))
log_stress = np.log(y_corrected[start_idx:])
logger.debug(
"Performing power-law fit on yield-corrected data",
fit_region_start_idx=start_idx,
n_fit_points=len(log_gamma),
)
coeffs = np.polyfit(log_gamma, log_stress, 1)
n_est = coeffs[0]
K_est = np.exp(coeffs[1])
logger.debug(
"Power-law regression results",
slope=coeffs[0],
intercept=coeffs[1],
n_raw=n_est,
K_raw=K_est,
)
# Clip to bounds
sigma_y_est = np.clip(sigma_y_est, 0.0, 1e6)
K_est = np.clip(K_est, 1e-6, 1e6)
n_est = np.clip(n_est, 0.01, 2.0)
self.parameters.set_value("sigma_y", float(sigma_y_est))
self.parameters.set_value("K", float(K_est))
self.parameters.set_value("n", float(n_est))
# Determine material type based on parameters
if sigma_y_est > 0 and abs(n_est - 1.0) < 0.05:
material_type = "Bingham plastic"
elif sigma_y_est > 0:
material_type = "yield stress fluid"
else:
material_type = "power-law fluid"
logger.debug(
"Herschel-Bulkley fit completed successfully",
fitted_sigma_y=float(sigma_y_est),
fitted_K=float(K_est),
fitted_n=float(n_est),
material_type=material_type,
)
ctx["sigma_y"] = float(sigma_y_est)
ctx["K"] = float(K_est)
ctx["n"] = float(n_est)
except Exception as e:
logger.error(
"Herschel-Bulkley fit failed",
error_type=type(e).__name__,
error_message=str(e),
exc_info=True,
)
raise
return self
def _predict(self, X: np.ndarray, **kwargs) -> np.ndarray:
"""Predict stress for given shear rates.
Args:
X: Shear rate data (γ̇)
**kwargs: Additional keyword arguments (deformation_mode, etc.)
Returns:
Predicted stress σ(γ̇)
"""
sigma_y = self.parameters.get_value("sigma_y")
K = self.parameters.get_value("K")
n = self.parameters.get_value("n")
# Convert to JAX for computation
gamma_dot = jnp.array(X)
# Compute stress
stress = self._predict_stress(gamma_dot, sigma_y, K, n)
# Convert back to numpy
return np.array(stress)
[docs]
def model_function(self, X, params, test_mode=None, **kwargs):
"""Model function for Bayesian inference.
This method is required by BayesianMixin for NumPyro NUTS sampling.
It computes predictions given input X and a parameter array.
Args:
X: Independent variable (shear rate γ̇)
params: Array of parameter values [sigma_y, K, n]
Returns:
Model predictions as JAX array
"""
# Extract parameters from array (in order they were added to ParameterSet)
sigma_y = params[0]
K = params[1]
n = params[2]
# Flow model only supports ROTATION test mode
# Compute prediction using the internal JAX method
return self._predict_stress(X, sigma_y, K, n)
@staticmethod
@jax.jit
def _predict_stress(
gamma_dot: jnp.ndarray,
sigma_y: float,
K: float,
n: float,
threshold: float = 1e-9,
) -> jnp.ndarray:
"""Compute shear stress using Herschel-Bulkley model.
Args:
gamma_dot: Shear rate (s^-1)
sigma_y: Yield stress (Pa)
K: Consistency index (Pa·s^n)
n: Flow behavior index
threshold: Threshold shear rate for yield (default: 1e-9)
Returns:
Shear stress (Pa)
"""
# σ(γ̇) = σ_y + K |γ̇|^n for |γ̇| > threshold
# σ(γ̇) = 0 for |γ̇| ≤ threshold (below yield)
abs_gamma_dot = jnp.abs(gamma_dot)
# Compute stress above yield.
# R5-GRAD-001: Add epsilon to avoid infinite gradient at abs_gamma_dot=0
# when n < 1. jnp.where evaluates both branches, so the power-law term
# is differentiated even in the below-yield branch.
stress_above_yield = sigma_y + K * jnp.power(abs_gamma_dot + 1e-30, n)
# Apply yield condition using jnp.where
return jnp.where(abs_gamma_dot > threshold, stress_above_yield, 0.0)
@staticmethod
@jax.jit
def _predict_viscosity(
gamma_dot: jnp.ndarray,
sigma_y: float,
K: float,
n: float,
threshold: float = 1e-9,
) -> jnp.ndarray:
"""Compute apparent viscosity using Herschel-Bulkley model.
Args:
gamma_dot: Shear rate (s^-1)
sigma_y: Yield stress (Pa)
K: Consistency index (Pa·s^n)
n: Flow behavior index
threshold: Threshold shear rate for yield (default: 1e-9)
Returns:
Apparent viscosity (Pa·s)
"""
# η_app(γ̇) = σ(γ̇) / γ̇ = σ_y/|γ̇| + K |γ̇|^(n-1)
abs_gamma_dot = jnp.abs(gamma_dot)
# Compute viscosity above yield
# R5-GRAD-001: epsilon guards infinite gradient at abs_gamma_dot=0 for n < 1
viscosity_above_yield = sigma_y / (abs_gamma_dot + threshold) + K * jnp.power(
abs_gamma_dot + 1e-30, n - 1.0
)
# Apply yield condition.
# R5-JAX-004: jnp.inf in the false branch corrupts gradients under NUTS
# because jnp.where evaluates both branches before selecting.
# Use a large finite sentinel (1e30 Pa·s) instead — physically equivalent
# (zero-shear-rate apparent viscosity is not physically meaningful) and
# gradient-safe because 1e30 has finite autodiff.
return jnp.where(abs_gamma_dot > threshold, viscosity_above_yield, 1e30)
[docs]
def predict_viscosity(self, gamma_dot: np.ndarray) -> np.ndarray:
"""Predict apparent viscosity for given shear rates.
Args:
gamma_dot: Shear rate data (γ̇)
Returns:
Predicted apparent viscosity η_app(γ̇)
"""
sigma_y = self.parameters.get_value("sigma_y")
K = self.parameters.get_value("K")
n = self.parameters.get_value("n")
# Convert to JAX for computation
gamma_dot_jax = jnp.array(gamma_dot)
# Compute viscosity
viscosity = self._predict_viscosity(gamma_dot_jax, sigma_y, K, n)
# Convert back to numpy
return np.array(viscosity)
[docs]
def predict_rheo(
self,
rheo_data: RheoData,
test_mode: TestMode | None = None,
output: str = "stress",
) -> RheoData:
"""Predict rheological response for RheoData.
Args:
rheo_data: Input rheological data
test_mode: Test mode (must be ROTATION)
output: Output type ('stress' or 'viscosity')
Returns:
RheoData with predictions
Raises:
ValueError: If test mode is not ROTATION
"""
# Detect test mode if not provided
if test_mode is None:
test_mode = detect_test_mode(rheo_data)
# Validate test mode
if test_mode != TestMode.ROTATION:
raise ValueError(
f"Herschel-Bulkley model only supports ROTATION test mode, got {test_mode}"
)
# Get shear rate data
gamma_dot = rheo_data.x
# Get parameters
sigma_y = self.parameters.get_value("sigma_y")
K = self.parameters.get_value("K")
n = self.parameters.get_value("n")
# Convert to JAX
gamma_dot_jax = jnp.array(gamma_dot)
# Compute prediction based on output type
if output == "stress":
y_pred = self._predict_stress(gamma_dot_jax, sigma_y, K, n)
y_units = "Pa"
elif output == "viscosity":
y_pred = self._predict_viscosity(gamma_dot_jax, sigma_y, K, n)
y_units = "Pa·s"
else:
raise ValueError(
f"Invalid output type: {output}. Must be 'stress' or 'viscosity'"
)
# Convert back to numpy
y_pred = np.array(y_pred)
# Create output RheoData
return RheoData(
x=np.array(gamma_dot),
y=y_pred,
x_units=rheo_data.x_units or "1/s",
y_units=y_units,
domain="time",
metadata={
"model": "HerschelBulkley",
"test_mode": TestMode.ROTATION,
"output": output,
"sigma_y": sigma_y,
"K": K,
"n": n,
},
validate=False,
)
[docs]
def __repr__(self) -> str:
"""String representation."""
sigma_y = self.parameters.get_value("sigma_y")
K = self.parameters.get_value("K")
n = self.parameters.get_value("n")
return f"HerschelBulkley(sigma_y={sigma_y:.3e}, K={K:.3e}, n={n:.3f})"
__all__ = ["HerschelBulkley"]