Geomechanics: Drucker-Prager triaxial compression

This example introduces a small geomechanics application inside the solid-mechanics example family. It implements an example-only Drucker-Prager plasticity assembler for a pressure-dependent soil or weak-rock material and drives it through a simple triaxial-compression strain path.

The goal is deliberately modest: demonstrate TensorMesh conventions for geomechanics without adding a public API. The local assembler lives in examples/solid/geomechanics/drucker_prager_triaxial/drucker_prager_triaxial.py.

Problem

The model is small-strain, associated Drucker-Prager plasticity with linear isotropic hardening. TensorMesh keeps the internal solid-mechanics convention stress tension-positive. For geomechanics reporting, the script prints axial stress and mean pressure as compression-positive quantities.

The yield function is written internally as

\[f(\sigma, \alpha) = q + \eta I_1 - (k + H\alpha) \le 0,\]

where

\[I_1 = \mathrm{tr}(\sigma), \qquad q = \sqrt{\frac{3}{2}\,s:s}, \qquad p = -\frac{I_1}{3}.\]

Because compression gives negative I1 in the tension-positive convention, higher confinement lowers f and delays yielding.

History variables

The example follows the same pattern as the existing J2 plasticity example:

• per-quadrature history variables are stored in self.history[etype];

• previous-step eps_p and alpha are passed through element_data;

• update_state(u) is called after each converged load step under torch.no_grad().

The example remains local to the script so the public API can be discussed after the convention-setting example has been reviewed.

Sanity check

Two confinement levels are run:

• p0 = 0 kPa;

• p0 = 100 kPa.

The script checks that the higher-confinement case reaches the elastic trial yield surface later and that the committed plastic history variable is monotonic.

Running it

cd examples/solid/geomechanics/drucker_prager_triaxial
python drucker_prager_triaxial.py

For a fast numerical-only run without writing the plot:

python drucker_prager_triaxial.py --no-plot --steps 16

The default run writes drucker_prager_triaxial.png with axial stress and plastic-history curves.

Fig. 45 Output of drucker_prager_triaxial.py. The left panel shows the compression-positive axial stress response; the right panel shows the committed plastic history variable. The higher-confinement case yields later, matching the pressure-dependent Drucker-Prager sanity check.

Core implementation

    """Example-only associated Drucker-Prager plasticity assembler.

    This class is deliberately kept inside the example.  It is not a public
    TensorMesh API.  The implementation follows the same high-level lifecycle
    as J2Plasticity:

    1. store per-quadrature history in ``self.history[etype]``;
    2. pass previous-step state through ``element_data`` during energy calls;
    3. call ``update_state(u)`` after each converged load step.

    Notes
    -----
    TensorMesh uses tension-positive stress.  With compression-positive mean
    pressure ``p = -tr(sigma) / 3``, the yield function is written internally as

        f = q + eta * I1 - (k + H * alpha) <= 0,

    where ``I1 = tr(sigma)`` and ``q = sqrt(3/2 s:s)``.  Because compression
    gives negative ``I1``, confinement increases the yield stress.
    """

    def __post_init__(self, params: DruckerPragerParameters | None = None):
        if params is None:
            params = DruckerPragerParameters()

        self.params = params
        self.E = float(params.E)
        self.nu = float(params.nu)
        self.cohesion = float(params.cohesion)
        self.friction_angle_deg = float(params.friction_angle_deg)
        self.H = float(params.H)

        self.mu = self.E / (2.0 * (1.0 + self.nu))
        self.bulk = self.E / (3.0 * (1.0 - 2.0 * self.nu))

        phi = math.radians(self.friction_angle_deg)
        sin_phi = math.sin(phi)
        cos_phi = math.cos(phi)

        # Triaxial-compression meridian fit to Mohr-Coulomb:
        # q = M p + k, p compression-positive.
        # With TensorMesh tension-positive stress, p = -I1 / 3, so
        # f = q + (M/3) I1 - k.
        self.M = 6.0 * sin_phi / (3.0 - sin_phi)
        self.eta = self.M / 3.0
        self.k = 6.0 * self.cohesion * cos_phi / (3.0 - sin_phi)

        # Associated linear Drucker-Prager return denominator for q-based f.
        self.return_denominator = 3.0 * self.mu + 9.0 * self.bulk * self.eta**2 + self.H

        self.history: Dict[str, Dict[str, torch.Tensor]] = {}
        for etype, trans in self.transformation.items():
            n_elem = trans.n_elements
            n_quad = trans.n_quadrature
            eps_p = torch.zeros((n_elem, n_quad, 3, 3), device=self.device, dtype=self.dtype)
            alpha = torch.zeros((n_elem, n_quad), device=self.device, dtype=self.dtype)
            self.history[etype] = {"eps_p": eps_p, "alpha": alpha}

    @staticmethod
    def _small_strain_3d(graddisplacement: torch.Tensor) -> torch.Tensor:
        """Return the 3D small-strain tensor for 2D or 3D input gradients."""
        dim = graddisplacement.shape[-1]
        if dim == 2:
            eps_2d = 0.5 * (graddisplacement + graddisplacement.transpose(-1, -2))
            eps = torch.zeros((3, 3), device=graddisplacement.device, dtype=graddisplacement.dtype)
            eps[:2, :2] = eps_2d
            return eps
        return 0.5 * (graddisplacement + graddisplacement.transpose(-1, -2))

    def _elastic_stress(self, eps_e: torch.Tensor) -> torch.Tensor:
        """Tension-positive isotropic elastic stress from elastic strain."""
        eye = torch.eye(3, device=eps_e.device, dtype=eps_e.dtype)
        tr_eps = eps_e.diagonal(dim1=-2, dim2=-1).sum(-1)
        dev_eps = eps_e - (tr_eps[..., None, None] / 3.0) * eye
        return 2.0 * self.mu * dev_eps + self.bulk * tr_eps[..., None, None] * eye

    @staticmethod
    def _invariants(sigma: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Return I1, deviatoric stress, and q for a tension-positive stress."""
        eye = torch.eye(3, device=sigma.device, dtype=sigma.dtype)
        I1 = sigma.diagonal(dim1=-2, dim2=-1).sum(-1)
        s = sigma - (I1[..., None, None] / 3.0) * eye
        s_contract_s = (s * s).sum(dim=(-2, -1))
        q = torch.sqrt(torch.clamp(1.5 * s_contract_s, min=1.0e-30))
        return I1, s, q

    def element_energy(
        self,
        graddisplacement: torch.Tensor,
        eps_p_n: torch.Tensor,
        alpha_n: torch.Tensor,
    ) -> torch.Tensor:
        """Algorithmic incremental potential density at one quadrature point."""
        eps = self._small_strain_3d(graddisplacement)
        eps_e_trial = eps - eps_p_n
        sigma_trial = self._elastic_stress(eps_e_trial)
        I1_trial, _, q_trial = self._invariants(sigma_trial)

        f_trial = q_trial + self.eta * I1_trial - (self.k + self.H * alpha_n)
        dgamma = torch.clamp(f_trial, min=0.0) / self.return_denominator

        tr_eps_e = eps_e_trial.diagonal(dim1=-2, dim2=-1).sum(-1)
        eye = torch.eye(3, device=eps.device, dtype=eps.dtype)
        dev_eps_e = eps_e_trial - (tr_eps_e / 3.0) * eye
        elastic_energy = 0.5 * self.bulk * tr_eps_e**2 + self.mu * (dev_eps_e * dev_eps_e).sum()

        return elastic_energy - 0.5 * self.return_denominator * dgamma**2

    def update_state(self, u_vec: torch.Tensor) -> None:
        """Commit per-quadrature state after a converged load step."""
        with torch.no_grad():
            for etype, trans in self.transformation.items():
                cells = trans.elements
                u_elem = u_vec[cells]
                grad_u = torch.einsum("bqkx,bku->bqux", trans.shape_grad, u_elem)

                dim = grad_u.shape[-1]
                if dim == 2:
                    eps = torch.zeros(grad_u.shape[:2] + (3, 3), device=u_vec.device, dtype=u_vec.dtype)
                    eps[..., :2, :2] = 0.5 * (grad_u + grad_u.transpose(-1, -2))
                else:
                    eps = 0.5 * (grad_u + grad_u.transpose(-1, -2))

                hist = self.history[etype]
                eps_p_n = hist["eps_p"]
                alpha_n = hist["alpha"]

                eps_e_trial = eps - eps_p_n
                sigma_trial = self._elastic_stress(eps_e_trial)
                I1_trial, s_trial, q_trial = self._invariants(sigma_trial)
                f_trial = q_trial + self.eta * I1_trial - (self.k + self.H * alpha_n)
                dgamma = torch.clamp(f_trial, min=0.0) / self.return_denominator

                q_safe = torch.clamp(q_trial, min=1.0e-30)
                n_dev = 1.5 * s_trial / q_safe[..., None, None]
                eye = torch.eye(3, device=u_vec.device, dtype=u_vec.dtype)
                flow_dir = n_dev + self.eta * eye

                active = (f_trial > 0.0).to(dtype=u_vec.dtype)
                dgamma = dgamma * active

                hist["eps_p"] += dgamma[..., None, None] * flow_dir
                hist["alpha"] += dgamma

    def element_data_from_history(self) -> Dict[str, Dict[str, torch.Tensor]]:
        """Return history in the element_data structure expected by energy()."""
        return {
            "eps_p_n": {etype: h["eps_p"] for etype, h in self.history.items()},
            "alpha_n": {etype: h["alpha"] for etype, h in self.history.items()},
        }

    def mean_alpha(self) -> torch.Tensor:
        """Return the mean committed plastic multiplier across all quadrature points."""
        values = [hist["alpha"].reshape(-1) for hist in self.history.values()]
        return torch.cat(values).mean()

    def max_alpha(self) -> torch.Tensor:
        """Return the maximum committed plastic multiplier."""
        values = [hist["alpha"].reshape(-1) for hist in self.history.values()]
        return torch.cat(values).max()

    def mean_stress(self, u_vec: torch.Tensor) -> torch.Tensor:
        """Average committed Cauchy stress over elements and quadrature points."""
        stresses = []
        with torch.no_grad():
            for etype, trans in self.transformation.items():
                u_elem = u_vec[trans.elements]
                grad_u = torch.einsum("bqkx,bku->bqux", trans.shape_grad, u_elem)
                dim = grad_u.shape[-1]
                if dim == 2:
                    eps = torch.zeros(grad_u.shape[:2] + (3, 3), device=u_vec.device, dtype=u_vec.dtype)
                    eps[..., :2, :2] = 0.5 * (grad_u + grad_u.transpose(-1, -2))
                else:
                    eps = 0.5 * (grad_u + grad_u.transpose(-1, -2))
                sigma = self._elastic_stress(eps - self.history[etype]["eps_p"])
                stresses.append(sigma.reshape(-1, 3, 3))
        return torch.cat(stresses, dim=0).mean(dim=0)