Source code for rheojax.models.flow.herschel_bulkley

"""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), ) finite_mask = np.isfinite(log_gamma) & np.isfinite(log_stress) if finite_mask.sum() < 2: raise ValueError("Insufficient finite data points for power-law fit") coeffs = np.polyfit(log_gamma[finite_mask], log_stress[finite_mask], 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"]