VLB Transient Network Models

Quick Reference

Model Summary

Model Classes

VLBLocal, VLBMultiNetwork

Physics

Statistically-based transient network with distribution tensor \(\boldsymbol{\mu}\)

Key Parameters

\(G_0\) (network modulus), \(k_d\) (dissociation rate)

Protocols

FLOW_CURVE, STARTUP, RELAXATION, CREEP, OSCILLATION, LAOS

Key Features

Molecular foundation, all-analytical (single network), uniaxial extension

Reference

Vernerey, Long & Brighenti (2017). JMPS 107, 1-20

Import:

from rheojax.models import VLBLocal, VLBMultiNetwork

Basic Usage:

# Single transient network
model = VLBLocal()
model.fit(omega, G_star, test_mode='oscillation')

# Multi-network (generalized Maxwell via VLB)
model = VLBMultiNetwork(n_modes=3, include_permanent=True)
model.fit(omega, G_star, test_mode='oscillation')

Notation Guide

Symbol

Description

Units

\(\boldsymbol{\mu}\)

Distribution tensor (second moment of chain end-to-end vector)

dimensionless

\(\varphi(\mathbf{r},t)\)

Chain end-to-end vector distribution function

1/m3

\(G_0\)

Network modulus (\(= c k_B T\) for Gaussian chains)

Pa

\(k_d\)

Bond dissociation (detachment) rate

1/s

\(k_a\)

Bond association (attachment) rate (= \(k_d\) at equilibrium)

1/s

\(t_R\)

Relaxation time (\(= 1/k_d\))

s

\(\eta_0\)

Zero-shear viscosity (\(= G_0/k_d\))

Pa·s

\(G_e\)

Permanent (equilibrium) network modulus

Pa

\(\eta_s\)

Solvent viscosity

Pa·s

\(\mathbf{L}\)

Velocity gradient tensor

1/s

\(\mathbf{D}\)

Rate-of-deformation tensor (\(= (\mathbf{L} + \mathbf{L}^T)/2\))

1/s

\(\mathbf{W}\)

Vorticity tensor (\(= (\mathbf{L} - \mathbf{L}^T)/2\))

1/s

\(c\)

Number density of elastically active chains

1/m3

\(\text{Wi}\)

Weissenberg number (\(= \dot{\gamma}/k_d\))

dimensionless

\(\text{De}\)

Deborah number (\(= \omega/k_d\) or \(= 1/(k_d \cdot t_{obs})\))

dimensionless

\(N_1\)

First normal stress difference (\(= \sigma_{xx} - \sigma_{yy}\))

Pa

\(J(t)\)

Creep compliance

1/Pa

\(\dot{\varepsilon}\)

Extensional strain rate

1/s

\(\eta_E\)

Extensional (Trouton) viscosity

Pa·s

Overview & Historical Context

Physical picture. Many soft materials — hydrogels, vitrimers, self-healing polymers, telechelic networks, supramolecular assemblies — derive their mechanical response from reversible (dynamic) cross-links that break and reform under thermal fluctuations and mechanical load. At equilibrium the creation and destruction of bonds balance; under deformation the chain configuration evolves and generates stress.

Historical development:

  1. Green & Tobolsky (1946) introduced the concept of a transient network where chains continuously break and reform. Under the assumption of instantaneous reformation in the unstressed state and constant destruction rate, the macroscopic response is Maxwell-like with a single exponential relaxation.

  2. Tanaka & Edwards (1992) formalized the network theory using the conformation tensor \(\mathbf{S} = \langle \mathbf{Q Q} \rangle\) and derived ODE evolution equations. This is the basis for the TNT family in RheoJAX.

  3. Vernerey, Long & Brighenti (2017) returned to the full chain distribution function \(\varphi(\mathbf{r},t)\) and derived the distribution tensor \(\boldsymbol{\mu}\) as its second moment, providing a molecular-statistical foundation that naturally connects to entropy, free energy, and dissipation. This is the VLB framework.

Key insight. At the Gaussian-chain level with constant \(k_d\), the VLB and TNT formulations are mathematically equivalent — both reduce to Maxwell viscoelasticity. The VLB route is preferred when one wishes to incorporate molecular extensions (Langevin finite extensibility, force-dependent \(k_d\), entropic arguments) because the distribution tensor \(\boldsymbol{\mu}\) has a clear statistical-mechanical interpretation.

Physical Foundations

Chain Distribution Function

Consider a network of elastically active chains, each described by its end-to-end vector \(\mathbf{r}\). The chain distribution function \(\varphi(\mathbf{r},t)\) gives the number density of chains with end-to-end vector \(\mathbf{r}\) at time \(t\). Its evolution is:

\[\frac{\partial \varphi}{\partial t} + \nabla_r \cdot (\dot{\mathbf{r}} \, \varphi) = k_a \varphi_0(\mathbf{r}) - k_d \varphi(\mathbf{r},t)\]

where:

  • \(\dot{\mathbf{r}} = \mathbf{L} \cdot \mathbf{r}\) is the affine convection of the end-to-end vector

  • \(k_a \varphi_0(\mathbf{r})\) represents creation of new chains in the equilibrium (isotropic Gaussian) distribution

  • \(k_d \varphi\) represents destruction of existing chains

At equilibrium (\(\mathbf{L} = 0\)): \(k_a \varphi_0 = k_d \varphi_{eq}\), hence \(k_a = k_d\).

Distribution Tensor

The distribution tensor is the normalized second moment:

\[\boldsymbol{\mu} \equiv \frac{\langle \mathbf{r} \otimes \mathbf{r} \rangle} {\langle r_0^2 \rangle / 3} = \frac{1}{c} \frac{3}{\langle r_0^2 \rangle} \int \mathbf{r} \otimes \mathbf{r} \, \varphi(\mathbf{r},t) \, d^3\!r\]

where \(c = \int \varphi \, d^3\!r\) is the total chain number density and \(\langle r_0^2 \rangle\) is the mean-square end-to-end distance at equilibrium.

Properties:

  • \(\boldsymbol{\mu}\) is symmetric and positive-definite

  • At equilibrium: \(\boldsymbol{\mu}_{eq} = \mathbf{I}\)

  • \(\text{tr}(\boldsymbol{\mu})/3\) measures average chain stretch relative to equilibrium

Eigenvalue interpretation:

Since \(\boldsymbol{\mu}\) is symmetric and positive-definite, it has three real positive eigenvalues \(\lambda_1^2, \lambda_2^2, \lambda_3^2\). Each eigenvalue represents the normalized mean-square stretch in the corresponding principal direction:

\[\lambda_i^2 = \frac{\langle r_i^2 \rangle}{\langle r_{0,i}^2 \rangle}\]

where \(r_i\) is the projection of the end-to-end vector onto the \(i\)-th principal axis.

  • \(\lambda_i > 1\): chains are stretched beyond equilibrium in direction \(i\)

  • \(\lambda_i < 1\): chains are compressed relative to equilibrium

  • \(\lambda_i = 1\): equilibrium configuration

The eigenvectors of \(\boldsymbol{\mu}\) define the principal axes of anisotropy in the chain distribution — the directions along which the network is most and least stretched. For simple shear, the principal axes rotate by an angle \(\theta = \frac{1}{2}\arctan(2\mu_{xy}/(\mu_{xx}-\mu_{yy}))\) from the flow direction.

Governing Equations

Evolution of the Distribution Tensor

By taking the second moment of the Smoluchowski equation for \(\varphi(\mathbf{r},t)\):

(1)\[\dot{\boldsymbol{\mu}} = k_d(\mathbf{I} - \boldsymbol{\mu}) + \mathbf{L} \cdot \boldsymbol{\mu} + \boldsymbol{\mu} \cdot \mathbf{L}^T\]

This is the workhorse equation of the VLB model. The three terms represent:

  1. Bond kinetics \(k_d(\mathbf{I} - \boldsymbol{\mu})\): drives \(\boldsymbol{\mu}\) toward equilibrium \(\mathbf{I}\) at rate \(k_d\).

  2. Affine deformation \(\mathbf{L} \cdot \boldsymbol{\mu} + \boldsymbol{\mu} \cdot \mathbf{L}^T\): stretches and rotates chains according to the macroscopic flow.

Note

Equation (1) uses the full velocity gradient \(\mathbf{L}\), not the symmetric part \(\mathbf{D}\). In simple shear with \(L_{12} = \dot{\gamma}\), the components are:

\[\begin{split}\dot{\mu}_{xx} &= k_d(1 - \mu_{xx}) + 2\dot{\gamma}\mu_{xy} \\ \dot{\mu}_{yy} &= k_d(1 - \mu_{yy}) \\ \dot{\mu}_{zz} &= k_d(1 - \mu_{zz}) \\ \dot{\mu}_{xy} &= -k_d \mu_{xy} + \dot{\gamma}\mu_{yy}\end{split}\]

The asymmetry (\(\dot{\gamma}\) appears only via \(\mu_{xy}\) and \(\mu_{yy}\)) arises because the velocity gradient \(\mathbf{L}\) is not symmetric in simple shear.

Cauchy Stress

For Gaussian chains the free energy per chain is \(\frac{3}{2}k_BT \frac{r^2}{\langle r_0^2 \rangle}\), giving the network stress:

(2)\[\boldsymbol{\sigma} = G_0 (\boldsymbol{\mu} - \mathbf{I}) + p\mathbf{I}\]

where \(G_0 = c k_B T\) is the network modulus.

Shear stress:

\[\sigma_{12} = G_0 \mu_{xy}\]

First normal stress difference:

\[N_1 = \sigma_{xx} - \sigma_{yy} = G_0(\mu_{xx} - \mu_{yy})\]

Stored Energy and Dissipation

The Helmholtz free energy density of the network is:

\[\Psi = \frac{1}{2} G_0 \bigl[\text{tr}(\boldsymbol{\mu}) - 3 - \ln \det(\boldsymbol{\mu})\bigr]\]

The mechanical dissipation rate is:

\[\mathcal{D} = G_0 k_d \bigl[\text{tr}(\boldsymbol{\mu}) - 3 - \ln \det(\boldsymbol{\mu})\bigr] \geq 0\]

which is non-negative by the convexity of \(f(x) = x - \ln x - 1\) for \(x > 0\), guaranteeing thermodynamic consistency.

Entropy

The entropy density of the network (per unit volume) is:

\[s = s_0 + \gamma_v \ln\!\left(\frac{T}{T_0}\right) - \frac{1}{2} c k_B \bigl[\text{tr}(\boldsymbol{\mu}) - 3\bigr]\]

where \(s_0\) is the reference entropy, \(\gamma_v\) is the volumetric heat capacity coefficient, and \(c k_B = G_0/T\) is the entropic modulus.

The second term captures the entropic penalty of chain stretching: chains that are stretched beyond equilibrium (\(\text{tr}(\boldsymbol{\mu}) > 3\)) have reduced conformational entropy. This is the molecular origin of rubber elasticity in the VLB framework — the restoring force is entropic, not energetic.

At equilibrium (\(\boldsymbol{\mu} = \mathbf{I}\)), the chain contribution vanishes (\(\text{tr}(\mathbf{I}) - 3 = 0\)), recovering the maximum entropy state.

Parameters

VLBLocal Parameters (2)

Name

Default

Bounds

Units

Description

\(G_0\)

1000.0

(1, 108)

Pa

Network modulus. Product of chain density and thermal energy: \(G_0 = c k_B T\).

\(k_d\)

1.0

(10-6, 106)

1/s

Dissociation rate. Inverse of the characteristic network relaxation time: \(t_R = 1/k_d\).

Derived quantities:

Property

Expression

Physical Meaning

Relaxation time

\(t_R = 1/k_d\)

Time for stress to relax to \(1/e\) of initial value

Zero-shear viscosity

\(\eta_0 = G_0/k_d\)

Newtonian plateau viscosity

Crossover frequency

\(\omega_c = k_d\)

Frequency where \(G' = G''\)

VLBMultiNetwork Parameters (2M + 1 or 2M + 2)

For M transient modes:

Name

Default

Bounds

Units

Description

\(G_I\)

log-spaced

(1, 108)

Pa

Network modulus for mode I (I = 0..M-1)

\(k_{d,I}\)

log-spaced

(10-6, 106)

1/s

Dissociation rate for mode I

\(\eta_s\)

0.0

(0, 104)

Pa·s

Solvent viscosity (always present)

\(G_e\)

0.0

(0, 108)

Pa

Permanent network modulus (only if include_permanent=True)

Total parameters: 2M + 1 (without permanent) or 2M + 2 (with permanent).

Special Cases

The VLB model reduces to several well-known models under special conditions:

Condition

Resulting Model

Details

Single mode, constant \(k_d\)

Maxwell

\(t_R = 1/k_d\), \(\eta = G_0/k_d\)

\(k_d \to 0\)

Neo-Hookean solid

Permanent network, no relaxation

\(k_d \to \infty\)

Newtonian fluid

Instantaneous relaxation, \(\eta = G_0/k_d \to 0\)

M modes + \(G_e\)

Standard linear solid (M=1)

Retardation + relaxation times

M modes + \(\eta_s\)

Generalized Maxwell (Prony series)

\(G(t) = \eta_s \delta(t) + \sum G_I e^{-k_{d,I} t}\)

M modes + \(G_e\) + \(\eta_s\)

Oldroyd-B (M=1)

Solvent + single viscoelastic mode + equilibrium

UCM Equivalence

The single-mode VLB with constant \(k_d\) is mathematically identical to the Upper-Convected Maxwell (UCM) model. To see this, define the polymer extra stress \(\boldsymbol{\tau} = G_0(\boldsymbol{\mu} - \mathbf{I})\) and substitute into the VLB evolution equation (1):

\[\dot{\boldsymbol{\tau}} + k_d \boldsymbol{\tau} = G_0\bigl(\mathbf{L} \cdot \boldsymbol{\mu} + \boldsymbol{\mu} \cdot \mathbf{L}^T\bigr)\]

In the UCM form with relaxation time \(\lambda = 1/k_d\) and modulus \(G = G_0\):

\[\boldsymbol{\tau} + \lambda \stackrel{\nabla}{\boldsymbol{\tau}} = 2 G \lambda \mathbf{D}\]

where \(\stackrel{\nabla}{\boldsymbol{\tau}}\) is the upper-convected derivative. These are the same equation. This equivalence guarantees that:

  • All standard UCM results (Pipkin diagram, asymptotic limits) apply directly

  • VLB inherits the UCM extensional singularity at \(\text{Wi}_{ext} = 1/2\)

  • Existing UCM validation benchmarks serve as VLB test cases

Protocol Summary

For complete step-by-step derivations including the full ODE solutions, see VLB Protocol Equations & Derivations.

Protocol

Single Network Result

Multi-Network Generalization

Flow Curve

\(\sigma = G_0 \dot{\gamma} / k_d\) (Newtonian)

\(\sigma = \bigl(\sum G_I/k_{d,I} + \eta_s\bigr)\dot{\gamma}\)

Startup

\(\sigma_{12}(t) = \frac{G_0 \dot{\gamma}}{k_d}(1-e^{-k_d t})\)

Superposition of exponentials + \(\eta_s \dot{\gamma}\)

Relaxation

\(G(t) = G_0 e^{-k_d t}\) (single exponential)

\(G(t) = G_e + \sum G_I e^{-k_{d,I} t}\)

Creep

\(J(t) = (1 + k_d t)/G_0\) (Maxwell)

SLS analytical (M=1+perm); general: ODE

SAOS

\(G'=G_0\omega^2 t_R^2/(1+\omega^2 t_R^2)\)

Sum of Maxwell modes + \(G_e\) + \(\eta_s \omega\)

LAOS

\(\sigma_{12}\) exactly sinusoidal; \(N_1\) has \(2\omega\)

ODE integration required

Extension

Singularity at \(\dot{\varepsilon} = k_d/2\); \(\text{Tr} \to 3\) at low rate

Sum over modes

Key signatures of constant \(k_d\):

  • Flow curve is Newtonian. Non-Newtonian behavior requires force-dependent \(k_d(\boldsymbol{\mu})\) (see VLB Advanced Theory & Numerical Methods).

  • Startup is monotonic (no overshoot).

  • LAOS \(\sigma_{12}\) has no higher harmonics (\(I_3/I_1 = 0\)).

  • Extension diverges at \(\dot{\varepsilon} = k_d/2\); Langevin finite extensibility regularizes this (see VLB Advanced Theory & Numerical Methods).

Multi-Network Model

Physical Picture

Real polymers often have multiple populations of chains with different lifetimes, or a combination of reversible and permanent cross-links. The VLBMultiNetwork model captures this via:

\[\boldsymbol{\sigma} = \sum_{I=0}^{M-1} G_I (\boldsymbol{\mu}_I - \mathbf{I}) + G_e (\boldsymbol{\mu}_\infty - \mathbf{I}) + 2\eta_s \mathbf{D}\]

where each transient mode \(I\) has its own distribution tensor \(\boldsymbol{\mu}_I\) evolving with rate \(k_{d,I}\), the permanent network (\(k_d = 0\)) maintains equilibrium strain, and the solvent contributes Newtonian stress.

Relaxation Spectrum

The relaxation modulus is a Prony series:

\[G(t) = G_e + \sum_{I=0}^{M-1} G_I \, e^{-k_{d,I} t}\]

Fitting strategy:

  1. Start with \(M = 1\) and increase until residuals plateau

  2. Initialize modes at log-spaced \(k_d\) values spanning the experimental frequency range

  3. Use SAOS data (broadest frequency window) as the primary fitting target

  4. Validate with relaxation and/or startup data

Validity & Assumptions

Assumption

Details & Limitations

Gaussian chains

Chains follow Gaussian statistics (\(P(r) \propto \exp(-3r^2/2\langle r_0^2 \rangle)\)). Breaks down for highly stretched chains. See Langevin extension in VLB Advanced Theory & Numerical Methods.

Constant \(k_d\)

Bond lifetime is independent of chain stretch or force. Results in Newtonian flow curve and linear LAOS. Force-dependent \(k_d\) introduces shear thinning (see VLB Advanced Theory & Numerical Methods).

Affine deformation

Chains deform affinely with the macroscopic flow (\(\dot{\mathbf{r}} = \mathbf{L} \cdot \mathbf{r}\)). Non-affine effects (fluctuations, excluded volume) are neglected.

Incompressibility

Pressure \(p\) is a Lagrange multiplier; material is assumed incompressible.

Monodisperse chains

All chains in a given mode have the same \(G_0\) and \(k_d\). Polydispersity requires multiple modes.

Isothermal

No temperature dependence. Temperature enters through \(G_0 = c k_B T\) and \(k_d = k_d^0 \exp(-E_a/k_BT)\).

No chain entanglement

Chains interact only through cross-links. Entanglement effects (reptation) are not included.

When to Use VLB

For a decision table comparing all VLB variants (Local, MultiNetwork, Variant, Nonlocal), see the VLB Transient Network Models.

For detailed material classification by \(k_d\) regime and diagnostic signatures, see VLB — What You Can Learn.

What You Can Learn

From VLBLocal Parameters

Parameter

Typical Range

Physical Insight

\(G_0\)

10 - 106 Pa

Network stiffness. \(G_0 = c k_B T\), so higher \(G_0\) means more active chains. Compare with rubber elasticity theory.

\(k_d\)

10-3 - 103 1/s

Bond kinetics. Small \(k_d\) = long-lived bonds (permanent-like). Large \(k_d\) = fast turnover (liquid-like).

For material classification by \(k_d\) regime, see the kinetic regimes table in VLB — What You Can Learn.

From Multi-Network Spectrum

The relaxation spectrum \(\{(G_I, t_{R,I})\}\) encodes the distribution of bond lifetimes in the network:

  • Well-separated modes: distinct bond populations with different chemistry

  • Closely-spaced modes: quasi-continuous distribution (polydispersity)

  • Dominant mode: controls the terminal relaxation

  • \(G_e > 0\): permanent cross-links present (solid-like long-time behavior)

  • \(\eta_s > 0\): un-networked polymer or solvent background

Cross-Protocol Validation

For cross-protocol consistency checks and the recommended multi-protocol validation workflow, see Cross-Protocol Validation Workflow in VLB — What You Can Learn.

API Reference

class rheojax.models.vlb.VLBLocal[source]

Single transient network VLB model (2 params: G0, k_d).

Implements the VLB framework for a single population of dynamic crosslinks with constant dissociation rate. Analytically equivalent to the Maxwell model but with molecular-statistical foundations.

The distribution tensor mu has equilibrium mu = I and evolves via:

dmu/dt = k_d*(I - mu) + L·mu + mu·L^T

Cauchy stress: sigma = G0*(mu - I)

Parameters:

  • G0 (float) – Network modulus (Pa), physically G0 = c*k_B*T where c is chain density

  • k_d (float) – Dissociation rate (1/s), inverse of relaxation time t_R = 1/k_d

parameters

Model parameters (G0, k_d)

Type:

ParameterSet

fitted_

Whether the model has been fitted

Type:

bool

See also

VLBMultiNetwork

Multi-network VLB with N transient + permanent + solvent

__init__()[source]

Initialize single-network VLB model.

property G0: float

Get network modulus G0 (Pa).

property k_d: float

Get dissociation rate k_d (1/s).

property relaxation_time: float

Get relaxation time t_R = 1/k_d (s).

property viscosity: float

Get zero-shear viscosity eta_0 = G0/k_d (Pa*s).

model_function(X, params, test_mode=None, **kwargs)[source]

NumPyro/BayesianMixin model function.

Routes to appropriate prediction based on test_mode. This is the stateless function used for both NLSQ optimization and Bayesian inference.

Parameters:

  • X (array-like) – Independent variable

  • params (array-like) – Parameter values in ParameterSet order: [G0, k_d]

  • test_mode (str, optional) – Override stored test mode

  • **kwargs – Protocol-specific parameters: gamma_dot, sigma_applied, gamma_0, omega

Returns:

Predicted response

Return type:

jnp.ndarray

predict_flow_curve(gamma_dot, return_components=False)[source]

Predict steady shear stress (and optionally viscosity, N1).

Newtonian: sigma = G0*gamma_dot/k_d.

Parameters:

  • gamma_dot (ndarray) – Shear rate array (1/s)

  • return_components (bool) – If True, return (sigma, eta, N1)

Returns:

Stress, or (stress, viscosity, N1) if return_components=True

Return type:

ndarray | tuple[ndarray, ndarray, ndarray]

predict_saos(omega, return_components=True)[source]

Predict SAOS storage and loss moduli.

Parameters:

  • omega (ndarray) – Angular frequency array (rad/s)

  • return_components (bool) – If True, return (G’, G’’)

Returns:

(G’, G’’) if return_components=True, else |G*|

Return type:

tuple[ndarray, ndarray] | ndarray

predict_normal_stresses(gamma_dot)[source]

Predict steady-state first and second normal stress differences.

N1 = 2*G0*(gamma_dot/k_d)^2, N2 = 0 (upper-convected Maxwell).

Parameters:

gamma_dot (ndarray) – Shear rate array (1/s)

Returns:

(N1, N2) in Pa. N2 is always zero for upper-convected VLB.

Return type:

tuple[ndarray, ndarray]

simulate_startup(t, gamma_dot, return_full=False)[source]

Simulate startup flow at constant shear rate.

Parameters:

  • t (ndarray) – Time array (s)

  • gamma_dot (float) – Applied shear rate (1/s)

  • return_full (bool) – If True, return (sigma, N1, strain)

Returns:

Stress, or (stress, N1, strain) if return_full=True

Return type:

ndarray | tuple[ndarray, ndarray, ndarray]

simulate_relaxation(t)[source]

Simulate stress relaxation G(t) = G0*exp(-k_d*t).

Parameters:

t (ndarray) – Time after cessation of flow (s)

Returns:

Relaxation modulus G(t) (Pa)

Return type:

ndarray

simulate_creep(t, sigma_0, return_full=False)[source]

Simulate creep under constant applied stress.

Parameters:

  • t (ndarray) – Time array (s)

  • sigma_0 (float) – Applied stress (Pa)

  • return_full (bool) – If True, return (strain, compliance)

Returns:

Strain gamma(t), or (gamma, J) if return_full=True

Return type:

ndarray | tuple[ndarray, ndarray]

simulate_laos(t, gamma_0, omega)[source]

Simulate large-amplitude oscillatory shear (LAOS).

Parameters:

  • t (ndarray) – Time array (s), should span at least 3-5 full cycles

  • gamma_0 (float) – Strain amplitude

  • omega (float) – Angular frequency (rad/s)

Returns:

Keys: ‘time’, ‘strain’, ‘stress’, ‘N1’, ‘gamma_dot’

Return type:

dict

predict_uniaxial_extension(epsilon_dot, return_trouton=False)[source]

Predict steady-state uniaxial extensional stress.

Parameters:

  • epsilon_dot (ndarray) – Extensional strain rate (1/s)

  • return_trouton (bool) – If True, also return Trouton ratio

Returns:

Extensional stress, or (stress, Trouton_ratio)

Return type:

ndarray | tuple[ndarray, ndarray]

simulate_uniaxial_extension(t, epsilon_dot)[source]

Simulate transient uniaxial extension.

Parameters:

  • t (ndarray) – Time array (s)

  • epsilon_dot (float) – Extensional strain rate (1/s)

Returns:

(extensional_stress, extensional_viscosity) as functions of time

Return type:

tuple[ndarray, ndarray]

get_relaxation_spectrum()[source]

Get relaxation spectrum as list of (G, tau) pairs.

Returns:

[(G0, 1/k_d)] — single mode

Return type:

list[tuple[float, float]]

extract_laos_harmonics(laos_result, n_harmonics=5)[source]

Extract Fourier harmonics from LAOS results.

Parameters:

  • laos_result (dict) – Output from simulate_laos()

  • n_harmonics (int) – Number of harmonics to extract

Returns:

Keys: ‘harmonic_index’, ‘sigma_harmonics’, ‘N1_harmonics’

Return type:

dict

__repr__()

Return string representation.

Return type:

str

static compute_stretch(mu_xx, mu_yy, mu_zz)

Compute stretch ratio from distribution tensor.

stretch = sqrt(tr(mu)/3)

At equilibrium (mu=I): stretch = 1. stretch > 1 indicates chains are extended beyond equilibrium.

Parameters:

  • mu_xx (float) – Diagonal distribution tensor components

  • mu_yy (float) – Diagonal distribution tensor components

  • mu_zz (float) – Diagonal distribution tensor components

Returns:

Stretch ratio (dimensionless, >= 0)

Return type:

float

deborah_number(omega)

Compute Deborah number De = t_R * omega.

Parameters:

omega (float) – Angular frequency (rad/s)

Returns:

Deborah number (dimensionless)

Return type:

float

dissociation_rate(mu_xx, mu_yy, mu_zz)

Compute dissociation rate (possibly state-dependent).

Base implementation returns constant k_d. Subclasses (e.g., VLBBell) can override this for force-dependent kinetics.

Parameters:

  • mu_xx (float) – Distribution tensor diagonal components

  • mu_yy (float) – Distribution tensor diagonal components

  • mu_zz (float) – Distribution tensor diagonal components

Returns:

Dissociation rate (1/s)

Return type:

float

fit(X, y=None, method='nlsq', check_compatibility=False, use_log_residuals=None, use_multi_start=None, n_starts=5, perturb_factor=0.3, deformation_mode=None, poisson_ratio=0.5, auto_init=False, return_result=False, check_physics=False, uncertainty=None, **kwargs)

Fit the model to data using NLSQ optimization.

This method uses NLSQ (GPU-accelerated nonlinear least squares) by default for fast point estimation. The optimization result is stored for potential warm-starting of Bayesian inference.

For very wide frequency ranges (>10 decades), multi-start optimization is automatically enabled to escape local minima.

Parameters:

  • X (numpy.typing.ArrayLike) – Input features

  • y (Optional[numpy.typing.ArrayLike]) – Target values

  • method (str) – Optimization method (‘nlsq’ by default for compatibility)

  • check_compatibility (bool) – Whether to check model-data compatibility before fitting. If True, warns when model may not be appropriate for data. Default is False for backward compatibility.

  • use_log_residuals (bool | None) – Whether to use log-space residuals for fitting. Recommended for wide frequency ranges (>8 decades) to prevent optimizer bias. If None (default), automatically detected based on data range. Explicit True/False overrides auto-detection.

  • use_multi_start (bool | None) – Whether to use multi-start optimization to escape local minima. Recommended for very wide ranges (>10 decades). If None (default), automatically enabled for >10 decades.

  • n_starts (int) – Number of random starts for multi-start optimization (default: 5)

  • perturb_factor (float) – Perturbation magnitude for multi-start random starts (default: 0.3). Parameters are perturbed by ± perturb_factor * (value or range). Larger values (0.7-0.9) explore wider parameter space.

  • auto_init (bool) – If True, calls auto_p0() to estimate initial parameters from data before running the optimizer (default: False).

  • return_result (bool) – If True, returns a FitResult instead of self. This intentionally breaks method chaining for workflows that need structured result objects (default: False).

  • check_physics (bool) – If True, runs post-fit physics validation and emits RheoJaxPhysicsWarning for any violations (default: False).

  • uncertainty (str | None) – Post-fit uncertainty method. "hessian" for fast Cramér-Rao bounds, "bootstrap" for residual bootstrap CIs, or None to skip (default: None).

  • **kwargs – Additional fitting options passed to _fit()

Return type:

BaseModel | Any

Returns:

self for method chaining (default), or FitResult if return_result=True.

Example

>>> model = Maxwell()
>>> model.fit(t, G_data)  # Uses NLSQ by default
>>> model.fit(t, G_data, method='nlsq', max_iter=1000)
>>> model.fit(t, G_data, check_compatibility=True)  # Check compatibility
>>> model.fit(omega, G_star, use_log_residuals=True)  # Force log-residuals
>>> model.fit(mastercurve, None, use_multi_start=True, n_starts=10)  # Multi-start
>>> result = model.fit(t, G_data, return_result=True)  # Structured result
>>> result = model.fit(t, G_data, auto_init=True, check_physics=True,
...                    return_result=True)  # Full pipeline
fit_bayesian(X, y=None, num_warmup=1000, num_samples=2000, num_chains=4, initial_values=None, test_mode=None, seed=None, deformation_mode=None, poisson_ratio=0.5, **nuts_kwargs)

Perform Bayesian inference using NumPyro NUTS sampler.

This method delegates to BayesianMixin.fit_bayesian() to run NUTS sampling for Bayesian parameter estimation. If initial_values is not provided and the model has been previously fitted with fit(), the NLSQ point estimates are automatically used for warm-starting.

Multi-chain sampling is enabled by default (num_chains=4) to provide reliable convergence diagnostics (R-hat, ESS) and parallel execution on multi-GPU systems.

Parameters:

  • X (numpy.typing.ArrayLike) – Independent variable data (input features) or RheoData object

  • y (Optional[numpy.typing.ArrayLike]) – Dependent variable data (observations to fit). If X is RheoData, y is ignored and extracted from X.

  • num_warmup (int) – Number of warmup/burn-in iterations (default: 1000)

  • num_samples (int) – Number of posterior samples per chain (default: 2000)

  • num_chains (int) – Number of MCMC chains (default: 4). Multiple chains enable proper R-hat computation and parallel execution. Chain method is auto-selected: ‘parallel’ on multi-GPU, ‘vectorized’ on single GPU/CPU.

  • initial_values (dict[str, float] | None) – Optional dict of initial parameter values for warm-start. If None and model is fitted, uses NLSQ estimates.

  • test_mode (str | None) – Explicit test mode (e.g., ‘relaxation’, ‘creep’, ‘oscillation’). If None, inferred from RheoData.metadata[‘test_mode’] or defaults to ‘relaxation’. Overrides RheoData metadata if provided.

  • seed (int | None) – Random seed for reproducibility. If None, uses seed=0 for deterministic results. Set to different values for independent runs.

  • **nuts_kwargs – Additional arguments passed to NUTS sampler (e.g., target_accept_prob, chain_method)

Return type:

BayesianResult

Returns:

BayesianResult containing posterior samples, summary statistics, and convergence diagnostics (R-hat, ESS, divergences)

Example

>>> model = Maxwell()
>>> # Warm-start from NLSQ with explicit mode
>>> model.fit(t, G_data, test_mode='relaxation')  # NLSQ optimization
>>> result = model.fit_bayesian(t, G_data, test_mode='relaxation')
>>>
>>> # RheoData with embedded mode (recommended)
>>> rheo_data = RheoData(x=omega, y=G_star, metadata={'test_mode': 'oscillation'})
>>> result = model.fit_bayesian(rheo_data)
>>>
>>> # Or provide explicit initial values
>>> result = model.fit_bayesian(
...     t, G_data,
...     initial_values={'G0': 1e5, 'eta': 1e3},
...     test_mode='creep'
... )
fit_predict(X, y, **kwargs)

Fit model and return predictions.

Parameters:

  • X (numpy.typing.ArrayLike) – Input features

  • y (numpy.typing.ArrayLike) – Target values

  • **kwargs – Additional fitting options

Return type:

numpy.typing.ArrayLike

Returns:

Model predictions on training data

classmethod from_dict(data)

Create model from dictionary.

Parameters:

data (dict[str, Any]) – Dictionary representation

Return type:

BaseModel

Returns:

Model instance

get_bayesian_result()

Get stored Bayesian inference result.

Return type:

BayesianResult | None

Returns:

BayesianResult from fit_bayesian(), or None if not run

Example

>>> model.fit_bayesian(t, G_data)
>>> result = model.get_bayesian_result()
>>> print(result.diagnostics['r_hat'])
get_credible_intervals(posterior_samples, credibility=0.95)

Compute highest density intervals (HDI) for posterior samples.

Computes the highest posterior density interval for each parameter, which is the shortest interval containing the specified probability mass. This is preferred over equal-tailed intervals for skewed distributions.

Parameters:

  • posterior_samples (dict[str, ndarray]) – Dictionary mapping parameter names to posterior sample arrays (from BayesianResult.posterior_samples)

  • credibility (float) – Probability mass to include in interval (default: 0.95) Common values: 0.68 (1 sigma), 0.95 (2 sigma), 0.997 (3 sigma)

Return type:

dict[str, tuple[float, float]]

Returns:

Dictionary mapping parameter names to (lower, upper) credible interval tuples. All values are float64.

Example

>>> result = model.fit_bayesian(X, y)
>>> intervals_95 = model.get_credible_intervals(
...     result.posterior_samples, credibility=0.95
... )
>>> print(f"95% CI for a: {intervals_95['a']}")
static get_equilibrium_distribution()

Return equilibrium distribution tensor mu_eq = I.

In the absence of flow, chains adopt their equilibrium configuration: mu = I (isotropic, unstretched).

Returns:

Equilibrium state [mu_xx, mu_yy, mu_zz, mu_xy] = [1, 1, 1, 0]

Return type:

Array

get_nlsq_result()

Get stored NLSQ optimization result.

Returns:

OptimizationResult from NLSQ fit, or None if not fitted

Example

>>> model.fit(t, G_data)
>>> result = model.get_nlsq_result()
>>> if result:
...     print(f"Converged: {result.success}")
get_parameter_dict()

Get all parameters as a dictionary.

Returns:

Dictionary of parameter name -> value

Return type:

dict[str, float]

get_parameter_uncertainties()

Get standard errors for fitted parameters from NLSQ covariance.

Returns:

std_error}, or None if covariance unavailable

Return type:

dict of {param_name

get_params(deep=True)

Get model parameters.

Parameters:

deep (bool) – If True, return parameters of sub-objects

Return type:

dict[str, Any]

Returns:

Dictionary of parameter names and values

initialize_from_creep(t, gamma, sigma_applied)

Initialize parameters from creep data γ(t) = σ₀·(1 + k_d·t)/G₀.

Parameters:

  • t (ndarray) – Time array (s)

  • gamma (ndarray) – Strain array γ(t)

  • sigma_applied (float) – Applied stress σ₀ (Pa)

Return type:

None

initialize_from_flow_curve(gamma_dot, sigma)

Initialize parameters from flow curve data.

Parameters:

  • gamma_dot (ndarray) – Shear rate array (1/s)

  • sigma (ndarray) – Stress array (Pa)

Return type:

None

initialize_from_relaxation(t, G)

Initialize parameters from stress relaxation data G(t) = G₀ exp(-k_d·t).

Parameters:

  • t (ndarray) – Time array (s)

  • G (ndarray) – Relaxation modulus G(t) (Pa)

Return type:

None

initialize_from_saos(omega, G_prime, G_double_prime)

Initialize parameters from SAOS data.

Uses the crossover frequency and modulus to estimate k_d and G0. In the linear regime, VLB reduces to Maxwell: crossover at omega_c = k_d.

Parameters:

  • omega (ndarray) – Angular frequency array (rad/s)

  • G_prime (ndarray) – Storage modulus G’ (Pa)

  • G_double_prime (ndarray) – Loss modulus G’’ (Pa)

Return type:

None

initialize_from_startup(t, sigma, gamma_dot)

Initialize parameters from startup shear data.

Uses the steady-state stress (σ_∞ = G₀/k_d · γ̇) and initial stress rate (dσ/dt|₀ = G₀ · γ̇) to estimate G₀ and k_d.

Parameters:

  • t (ndarray) – Time array (s)

  • sigma (ndarray) – Shear stress array (Pa)

  • gamma_dot (float) – Applied shear rate (1/s)

Return type:

None

property pcov_

Parameter covariance matrix from NLSQ fit.

Returns:

ndarray of shape (n_params, n_params), or None if not fitted

property popt_

Optimal parameter values from NLSQ fit.

Returns:

ndarray of shape (n_params,), or None if not fitted

precompile(test_mode='relaxation', X=None, y=None)

Precompile NLSQ residual functions to eliminate JIT cold-start.

Triggers JIT compilation by running a minimal fit (max_iter=1) with dummy data. The model parameters are reset afterwards so the model is left in its original state.

This is useful for interactive sessions or benchmarks where the ~870ms first-fit JIT overhead should be excluded.

Parameters:

  • test_mode (str) – Test mode to precompile for (default: ‘relaxation’).

  • X (Optional[numpy.typing.ArrayLike]) – Optional input data for shape inference. If None, uses a 10-point logspace array.

  • y (Optional[numpy.typing.ArrayLike]) – Optional output data. If None, generates ones matching X.

Return type:

float

Returns:

Compilation time in seconds.

Example

>>> model = Maxwell()
>>> t = model.precompile(test_mode='relaxation')
>>> print(f"Compiled in {t:.2f}s")
>>> model.fit(X, y)  # No JIT overhead
precompile_bayesian(X=None, y=None, test_mode=None, num_chains=4)

Precompile NUTS kernel to eliminate JIT overhead in subsequent calls.

Triggers JIT compilation of the NumPyro model by running a minimal sampling (1 warmup, 1 sample). This caches the compiled kernel so that subsequent fit_bayesian() calls are 2-5x faster.

Parameters:

  • X (ndarray | RheoData | None) – Sample input data for determining array shapes. If None, uses a default 10-point linspace [0.01, 100].

  • y (ndarray | None) – Sample output data. If None, generates dummy data.

  • test_mode (str | TestModeEnum | None) – Test mode to precompile for. If None, defaults to ‘relaxation’.

Returns:

Compilation time in seconds.

Return type:

float

Example

>>> model = Maxwell()
>>> compile_time = model.precompile_bayesian(test_mode='relaxation')
>>> print(f"Compiled in {compile_time:.1f}s")
>>> # Now fit_bayesian() will be faster
>>> result = model.fit_bayesian(X, y)  # No compilation overhead
predict(X, test_mode=None, deformation_mode=None, poisson_ratio=None, **kwargs)

Make predictions.

Parameters:

  • X (numpy.typing.ArrayLike) – Input features

  • test_mode (str | None) – Optional test mode (‘oscillation’, ‘relaxation’, ‘creep’, ‘flow’). If provided, sets model’s test_mode before prediction. Useful for data generation without fitting.

  • deformation_mode (str | DeformationMode | None) – Optional deformation mode for output conversion. If None, uses the mode stored from fit(). If tensile, converts G* predictions to E* space.

  • poisson_ratio (float | None) – Poisson’s ratio for conversion. If None, uses value stored from fit() (default 0.5).

  • **kwargs – Additional arguments passed to the internal _predict method.

Return type:

numpy.typing.ArrayLike

Returns:

Model predictions (in E* space if deformation_mode is tensile)

sample_prior(num_samples=1000, seed=None)

Sample from prior distributions over parameter bounds.

Samples from uniform prior distributions defined by parameter bounds. This is useful for prior predictive checks and understanding the prior’s influence on the posterior.

Parameters:

  • num_samples (int) – Number of samples to draw from prior (default: 1000)

  • seed (int | None) – Random seed for reproducibility (default: None)

Return type:

dict[str, ndarray]

Returns:

Dictionary mapping parameter names to arrays of prior samples. Each array has shape (num_samples,) and dtype float64.

Raises:

Example

>>> model = MyModel()
>>> prior_samples = model.sample_prior(num_samples=500, seed=42)
>>> print(prior_samples["a"].shape)  # (500,)
score(X, y)

Compute model score (R² by default).

Parameters:

  • X (numpy.typing.ArrayLike) – Input features

  • y (numpy.typing.ArrayLike) – True target values

Return type:

float

Returns:

Model score (R² coefficient)

set_parameter_dict(params)

Set parameters from a dictionary.

Parameters:

params (dict[str, float]) – Dictionary of parameter name -> value

Return type:

None

set_params(**params)

Set model parameters.

Parameters:

**params – Parameter names and values

Return type:

BaseModel

Returns:

self for method chaining

to_dict()

Serialize model to dictionary.

Return type:

dict[str, Any]

Returns:

Dictionary representation of model

weissenberg_number(gamma_dot)

Compute Weissenberg number Wi = t_R * gamma_dot.

Parameters:

gamma_dot (float) – Shear rate (1/s)

Returns:

Weissenberg number (dimensionless)

Return type:

float

class rheojax.models.vlb.VLBMultiNetwork(n_modes=2, include_permanent=False)[source]

Multi-network VLB model: M transient + optional permanent + solvent.

Implements a network with N independent transient crosslink populations, each with modulus G_i and dissociation rate k_d_i. The total stress is a superposition of N Maxwell modes.

Parameters:

  • n_modes (int) – Number of distinct transient network populations (N >= 1)

  • include_permanent (bool) – Whether to include a permanent (elastic) network (G_e)

parameters

Model parameters: [G_0, k_d_0, G_1, k_d_1, …, eta_s, (G_e)]

Type:

ParameterSet

fitted_

Whether the model has been fitted

Type:

bool

_n_modes

Number of transient modes

Type:

int

Notes

Parameter ordering: [G_0, k_d_0, G_1, k_d_1, …, G_{N-1}, k_d_{N-1}, eta_s, (G_e)] Total parameter count: 2N + 1 (without permanent) or 2N + 2 (with permanent)

See also

VLBLocal

Single transient network (2 parameters)

__init__(n_modes=2, include_permanent=False)[source]

Initialize multi-network VLB model.

Parameters:

  • n_modes (int) – Number of transient network populations (must be >= 1)

  • include_permanent (bool) – Include permanent elastic network

property n_modes: int

Number of transient network modes.

property include_permanent: bool

Whether a permanent network is included.

property G_e: float

Permanent network modulus (Pa). 0 if not included.

property eta_s: float

Solvent viscosity (Pa*s).

property G_total: float

sum G_i + G_e.

Type:

Total modulus

property eta_0: float

sum G_i/k_d_i + eta_s.

Type:

Zero-shear viscosity

model_function(X, params, test_mode=None, **kwargs)[source]

NumPyro/BayesianMixin model function.

Routes to appropriate prediction based on test_mode.

Parameters:

  • X (array-like) – Independent variable

  • params (array-like) – Parameter values: [G_0, k_d_0, …, G_{N-1}, k_d_{N-1}, eta_s, (G_e)]

  • test_mode (str, optional) – Override stored test mode

  • **kwargs – Protocol-specific parameters

Returns:

Predicted response

Return type:

jnp.ndarray

predict_saos(omega, return_components=True)[source]

Predict multi-mode SAOS moduli.

Parameters:

  • omega (ndarray) – Angular frequency (rad/s)

  • return_components (bool) – If True, return (G’, G’’)

Returns:

(G’, G’’) or |G*|

Return type:

tuple[ndarray, ndarray] | ndarray

get_relaxation_spectrum()[source]

Get relaxation spectrum as list of (G, tau) pairs.

Returns:

[(G_i, 1/k_d_i)] for each transient mode

Return type:

list[tuple[float, float]]

__repr__()[source]

Return string representation.

Return type:

str

static compute_stretch(mu_xx, mu_yy, mu_zz)

Compute stretch ratio from distribution tensor.

stretch = sqrt(tr(mu)/3)

At equilibrium (mu=I): stretch = 1. stretch > 1 indicates chains are extended beyond equilibrium.

Parameters:

  • mu_xx (float) – Diagonal distribution tensor components

  • mu_yy (float) – Diagonal distribution tensor components

  • mu_zz (float) – Diagonal distribution tensor components

Returns:

Stretch ratio (dimensionless, >= 0)

Return type:

float

deborah_number(omega)

Compute Deborah number De = t_R * omega.

Parameters:

omega (float) – Angular frequency (rad/s)

Returns:

Deborah number (dimensionless)

Return type:

float

dissociation_rate(mu_xx, mu_yy, mu_zz)

Compute dissociation rate (possibly state-dependent).

Base implementation returns constant k_d. Subclasses (e.g., VLBBell) can override this for force-dependent kinetics.

Parameters:

  • mu_xx (float) – Distribution tensor diagonal components

  • mu_yy (float) – Distribution tensor diagonal components

  • mu_zz (float) – Distribution tensor diagonal components

Returns:

Dissociation rate (1/s)

Return type:

float

fit(X, y=None, method='nlsq', check_compatibility=False, use_log_residuals=None, use_multi_start=None, n_starts=5, perturb_factor=0.3, deformation_mode=None, poisson_ratio=0.5, auto_init=False, return_result=False, check_physics=False, uncertainty=None, **kwargs)

Fit the model to data using NLSQ optimization.

This method uses NLSQ (GPU-accelerated nonlinear least squares) by default for fast point estimation. The optimization result is stored for potential warm-starting of Bayesian inference.

For very wide frequency ranges (>10 decades), multi-start optimization is automatically enabled to escape local minima.

Parameters:

  • X (numpy.typing.ArrayLike) – Input features

  • y (Optional[numpy.typing.ArrayLike]) – Target values

  • method (str) – Optimization method (‘nlsq’ by default for compatibility)

  • check_compatibility (bool) – Whether to check model-data compatibility before fitting. If True, warns when model may not be appropriate for data. Default is False for backward compatibility.

  • use_log_residuals (bool | None) – Whether to use log-space residuals for fitting. Recommended for wide frequency ranges (>8 decades) to prevent optimizer bias. If None (default), automatically detected based on data range. Explicit True/False overrides auto-detection.

  • use_multi_start (bool | None) – Whether to use multi-start optimization to escape local minima. Recommended for very wide ranges (>10 decades). If None (default), automatically enabled for >10 decades.

  • n_starts (int) – Number of random starts for multi-start optimization (default: 5)

  • perturb_factor (float) – Perturbation magnitude for multi-start random starts (default: 0.3). Parameters are perturbed by ± perturb_factor * (value or range). Larger values (0.7-0.9) explore wider parameter space.

  • auto_init (bool) – If True, calls auto_p0() to estimate initial parameters from data before running the optimizer (default: False).

  • return_result (bool) – If True, returns a FitResult instead of self. This intentionally breaks method chaining for workflows that need structured result objects (default: False).

  • check_physics (bool) – If True, runs post-fit physics validation and emits RheoJaxPhysicsWarning for any violations (default: False).

  • uncertainty (str | None) – Post-fit uncertainty method. "hessian" for fast Cramér-Rao bounds, "bootstrap" for residual bootstrap CIs, or None to skip (default: None).

  • **kwargs – Additional fitting options passed to _fit()

Return type:

BaseModel | Any

Returns:

self for method chaining (default), or FitResult if return_result=True.

Example

>>> model = Maxwell()
>>> model.fit(t, G_data)  # Uses NLSQ by default
>>> model.fit(t, G_data, method='nlsq', max_iter=1000)
>>> model.fit(t, G_data, check_compatibility=True)  # Check compatibility
>>> model.fit(omega, G_star, use_log_residuals=True)  # Force log-residuals
>>> model.fit(mastercurve, None, use_multi_start=True, n_starts=10)  # Multi-start
>>> result = model.fit(t, G_data, return_result=True)  # Structured result
>>> result = model.fit(t, G_data, auto_init=True, check_physics=True,
...                    return_result=True)  # Full pipeline
fit_bayesian(X, y=None, num_warmup=1000, num_samples=2000, num_chains=4, initial_values=None, test_mode=None, seed=None, deformation_mode=None, poisson_ratio=0.5, **nuts_kwargs)

Perform Bayesian inference using NumPyro NUTS sampler.

This method delegates to BayesianMixin.fit_bayesian() to run NUTS sampling for Bayesian parameter estimation. If initial_values is not provided and the model has been previously fitted with fit(), the NLSQ point estimates are automatically used for warm-starting.

Multi-chain sampling is enabled by default (num_chains=4) to provide reliable convergence diagnostics (R-hat, ESS) and parallel execution on multi-GPU systems.

Parameters:

  • X (numpy.typing.ArrayLike) – Independent variable data (input features) or RheoData object

  • y (Optional[numpy.typing.ArrayLike]) – Dependent variable data (observations to fit). If X is RheoData, y is ignored and extracted from X.

  • num_warmup (int) – Number of warmup/burn-in iterations (default: 1000)

  • num_samples (int) – Number of posterior samples per chain (default: 2000)

  • num_chains (int) – Number of MCMC chains (default: 4). Multiple chains enable proper R-hat computation and parallel execution. Chain method is auto-selected: ‘parallel’ on multi-GPU, ‘vectorized’ on single GPU/CPU.

  • initial_values (dict[str, float] | None) – Optional dict of initial parameter values for warm-start. If None and model is fitted, uses NLSQ estimates.

  • test_mode (str | None) – Explicit test mode (e.g., ‘relaxation’, ‘creep’, ‘oscillation’). If None, inferred from RheoData.metadata[‘test_mode’] or defaults to ‘relaxation’. Overrides RheoData metadata if provided.

  • seed (int | None) – Random seed for reproducibility. If None, uses seed=0 for deterministic results. Set to different values for independent runs.

  • **nuts_kwargs – Additional arguments passed to NUTS sampler (e.g., target_accept_prob, chain_method)

Return type:

BayesianResult

Returns:

BayesianResult containing posterior samples, summary statistics, and convergence diagnostics (R-hat, ESS, divergences)

Example

>>> model = Maxwell()
>>> # Warm-start from NLSQ with explicit mode
>>> model.fit(t, G_data, test_mode='relaxation')  # NLSQ optimization
>>> result = model.fit_bayesian(t, G_data, test_mode='relaxation')
>>>
>>> # RheoData with embedded mode (recommended)
>>> rheo_data = RheoData(x=omega, y=G_star, metadata={'test_mode': 'oscillation'})
>>> result = model.fit_bayesian(rheo_data)
>>>
>>> # Or provide explicit initial values
>>> result = model.fit_bayesian(
...     t, G_data,
...     initial_values={'G0': 1e5, 'eta': 1e3},
...     test_mode='creep'
... )
fit_predict(X, y, **kwargs)

Fit model and return predictions.

Parameters:

  • X (numpy.typing.ArrayLike) – Input features

  • y (numpy.typing.ArrayLike) – Target values

  • **kwargs – Additional fitting options

Return type:

numpy.typing.ArrayLike

Returns:

Model predictions on training data

classmethod from_dict(data)

Create model from dictionary.

Parameters:

data (dict[str, Any]) – Dictionary representation

Return type:

BaseModel

Returns:

Model instance

get_bayesian_result()

Get stored Bayesian inference result.

Return type:

BayesianResult | None

Returns:

BayesianResult from fit_bayesian(), or None if not run

Example

>>> model.fit_bayesian(t, G_data)
>>> result = model.get_bayesian_result()
>>> print(result.diagnostics['r_hat'])
get_credible_intervals(posterior_samples, credibility=0.95)

Compute highest density intervals (HDI) for posterior samples.

Computes the highest posterior density interval for each parameter, which is the shortest interval containing the specified probability mass. This is preferred over equal-tailed intervals for skewed distributions.

Parameters:

  • posterior_samples (dict[str, ndarray]) – Dictionary mapping parameter names to posterior sample arrays (from BayesianResult.posterior_samples)

  • credibility (float) – Probability mass to include in interval (default: 0.95) Common values: 0.68 (1 sigma), 0.95 (2 sigma), 0.997 (3 sigma)

Return type:

dict[str, tuple[float, float]]

Returns:

Dictionary mapping parameter names to (lower, upper) credible interval tuples. All values are float64.

Example

>>> result = model.fit_bayesian(X, y)
>>> intervals_95 = model.get_credible_intervals(
...     result.posterior_samples, credibility=0.95
... )
>>> print(f"95% CI for a: {intervals_95['a']}")
static get_equilibrium_distribution()

Return equilibrium distribution tensor mu_eq = I.

In the absence of flow, chains adopt their equilibrium configuration: mu = I (isotropic, unstretched).

Returns:

Equilibrium state [mu_xx, mu_yy, mu_zz, mu_xy] = [1, 1, 1, 0]

Return type:

Array

get_nlsq_result()

Get stored NLSQ optimization result.

Returns:

OptimizationResult from NLSQ fit, or None if not fitted

Example

>>> model.fit(t, G_data)
>>> result = model.get_nlsq_result()
>>> if result:
...     print(f"Converged: {result.success}")
get_parameter_dict()

Get all parameters as a dictionary.

Returns:

Dictionary of parameter name -> value

Return type:

dict[str, float]

get_parameter_uncertainties()

Get standard errors for fitted parameters from NLSQ covariance.

Returns:

std_error}, or None if covariance unavailable

Return type:

dict of {param_name

get_params(deep=True)

Get model parameters.

Parameters:

deep (bool) – If True, return parameters of sub-objects

Return type:

dict[str, Any]

Returns:

Dictionary of parameter names and values

initialize_from_creep(t, gamma, sigma_applied)

Initialize parameters from creep data γ(t) = σ₀·(1 + k_d·t)/G₀.

Parameters:

  • t (ndarray) – Time array (s)

  • gamma (ndarray) – Strain array γ(t)

  • sigma_applied (float) – Applied stress σ₀ (Pa)

Return type:

None

initialize_from_flow_curve(gamma_dot, sigma)

Initialize parameters from flow curve data.

Parameters:

  • gamma_dot (ndarray) – Shear rate array (1/s)

  • sigma (ndarray) – Stress array (Pa)

Return type:

None

initialize_from_relaxation(t, G)

Initialize parameters from stress relaxation data G(t) = G₀ exp(-k_d·t).

Parameters:

  • t (ndarray) – Time array (s)

  • G (ndarray) – Relaxation modulus G(t) (Pa)

Return type:

None

initialize_from_saos(omega, G_prime, G_double_prime)

Initialize parameters from SAOS data.

Uses the crossover frequency and modulus to estimate k_d and G0. In the linear regime, VLB reduces to Maxwell: crossover at omega_c = k_d.

Parameters:

  • omega (ndarray) – Angular frequency array (rad/s)

  • G_prime (ndarray) – Storage modulus G’ (Pa)

  • G_double_prime (ndarray) – Loss modulus G’’ (Pa)

Return type:

None

initialize_from_startup(t, sigma, gamma_dot)

Initialize parameters from startup shear data.

Uses the steady-state stress (σ_∞ = G₀/k_d · γ̇) and initial stress rate (dσ/dt|₀ = G₀ · γ̇) to estimate G₀ and k_d.

Parameters:

  • t (ndarray) – Time array (s)

  • sigma (ndarray) – Shear stress array (Pa)

  • gamma_dot (float) – Applied shear rate (1/s)

Return type:

None

property pcov_

Parameter covariance matrix from NLSQ fit.

Returns:

ndarray of shape (n_params, n_params), or None if not fitted

property popt_

Optimal parameter values from NLSQ fit.

Returns:

ndarray of shape (n_params,), or None if not fitted

precompile(test_mode='relaxation', X=None, y=None)

Precompile NLSQ residual functions to eliminate JIT cold-start.

Triggers JIT compilation by running a minimal fit (max_iter=1) with dummy data. The model parameters are reset afterwards so the model is left in its original state.

This is useful for interactive sessions or benchmarks where the ~870ms first-fit JIT overhead should be excluded.

Parameters:

  • test_mode (str) – Test mode to precompile for (default: ‘relaxation’).

  • X (Optional[numpy.typing.ArrayLike]) – Optional input data for shape inference. If None, uses a 10-point logspace array.

  • y (Optional[numpy.typing.ArrayLike]) – Optional output data. If None, generates ones matching X.

Return type:

float

Returns:

Compilation time in seconds.

Example

>>> model = Maxwell()
>>> t = model.precompile(test_mode='relaxation')
>>> print(f"Compiled in {t:.2f}s")
>>> model.fit(X, y)  # No JIT overhead
precompile_bayesian(X=None, y=None, test_mode=None, num_chains=4)

Precompile NUTS kernel to eliminate JIT overhead in subsequent calls.

Triggers JIT compilation of the NumPyro model by running a minimal sampling (1 warmup, 1 sample). This caches the compiled kernel so that subsequent fit_bayesian() calls are 2-5x faster.

Parameters:

  • X (ndarray | RheoData | None) – Sample input data for determining array shapes. If None, uses a default 10-point linspace [0.01, 100].

  • y (ndarray | None) – Sample output data. If None, generates dummy data.

  • test_mode (str | TestModeEnum | None) – Test mode to precompile for. If None, defaults to ‘relaxation’.

Returns:

Compilation time in seconds.

Return type:

float

Example

>>> model = Maxwell()
>>> compile_time = model.precompile_bayesian(test_mode='relaxation')
>>> print(f"Compiled in {compile_time:.1f}s")
>>> # Now fit_bayesian() will be faster
>>> result = model.fit_bayesian(X, y)  # No compilation overhead
predict(X, test_mode=None, deformation_mode=None, poisson_ratio=None, **kwargs)

Make predictions.

Parameters:

  • X (numpy.typing.ArrayLike) – Input features

  • test_mode (str | None) – Optional test mode (‘oscillation’, ‘relaxation’, ‘creep’, ‘flow’). If provided, sets model’s test_mode before prediction. Useful for data generation without fitting.

  • deformation_mode (str | DeformationMode | None) – Optional deformation mode for output conversion. If None, uses the mode stored from fit(). If tensile, converts G* predictions to E* space.

  • poisson_ratio (float | None) – Poisson’s ratio for conversion. If None, uses value stored from fit() (default 0.5).

  • **kwargs – Additional arguments passed to the internal _predict method.

Return type:

numpy.typing.ArrayLike

Returns:

Model predictions (in E* space if deformation_mode is tensile)

sample_prior(num_samples=1000, seed=None)

Sample from prior distributions over parameter bounds.

Samples from uniform prior distributions defined by parameter bounds. This is useful for prior predictive checks and understanding the prior’s influence on the posterior.

Parameters:

  • num_samples (int) – Number of samples to draw from prior (default: 1000)

  • seed (int | None) – Random seed for reproducibility (default: None)

Return type:

dict[str, ndarray]

Returns:

Dictionary mapping parameter names to arrays of prior samples. Each array has shape (num_samples,) and dtype float64.

Raises:

Example

>>> model = MyModel()
>>> prior_samples = model.sample_prior(num_samples=500, seed=42)
>>> print(prior_samples["a"].shape)  # (500,)
score(X, y)

Compute model score (R² by default).

Parameters:

  • X (numpy.typing.ArrayLike) – Input features

  • y (numpy.typing.ArrayLike) – True target values

Return type:

float

Returns:

Model score (R² coefficient)

set_parameter_dict(params)

Set parameters from a dictionary.

Parameters:

params (dict[str, float]) – Dictionary of parameter name -> value

Return type:

None

set_params(**params)

Set model parameters.

Parameters:

**params – Parameter names and values

Return type:

BaseModel

Returns:

self for method chaining

to_dict()

Serialize model to dictionary.

Return type:

dict[str, Any]

Returns:

Dictionary representation of model

weissenberg_number(gamma_dot)

Compute Weissenberg number Wi = t_R * gamma_dot.

Parameters:

gamma_dot (float) – Shear rate (1/s)

Returns:

Weissenberg number (dimensionless)

Return type:

float

References

  1. Vernerey, F.J., Long, R. & Brighenti, R. (2017). “A statistically-based continuum theory for polymers with transient networks.” J. Mech. Phys. Solids, 107, 1-20. https://doi.org/10.1016/j.jmps.2017.05.016

  2. Green, M.S. & Tobolsky, A.V. (1946). “A New Approach to the Theory of Relaxing Polymeric Media.” J. Chem. Phys., 14(2), 80-92. https://doi.org/10.1063/1.1724109

  3. Tanaka, F. & Edwards, S.F. (1992). “Viscoelastic properties of physically crosslinked networks.” J. Non-Newtonian Fluid Mech., 43(2-3), 247-271. https://doi.org/10.1016/0377-0257(92)80027-U

  4. Vernerey, F.J. (2018). “Transient response of nonlinear polymer networks: A kinetic theory.” J. Mech. Phys. Solids, 115, 230-247. https://doi.org/10.1016/j.jmps.2018.02.018 PDF

  5. Long, R., Qi, H.J. & Dunn, M.L. (2013). “Modeling the mechanics of covalently adaptable polymer networks with temperature-dependent bond exchange reactions.” Soft Matter, 9, 4083-4096. https://doi.org/10.1039/C3SM27945F