{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import pyvista\npyvista.OFF_SCREEN = True\npyvista.set_jupyter_backend('static')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Nodal stress recovery + invariants \u2014 Lam\u00e9 thick-walled cylinder\n\nStress is the most-asked-for derived quantity in linear FEA, and\nthe recovery from a displacement solution is where most of the\n\"why is my stress field wrong\" questions actually live.  This\nexample walks through the recovery pipeline on a problem with a\nclean closed-form reference \u2014 the **Lam\u00e9 thick cylinder** under\ninternal pressure (Timoshenko & Goodier 1970 \u00a733) \u2014 so every\nintermediate quantity has an analytical value to compare against.\n\nTwo public post-processing helpers carry the work:\n\n* :func:`femorph_solver.recover.compute_nodal_stress` \u2014 for each\n  node, take the unweighted arithmetic mean of the per-element\n  stress contributions evaluated at that node.  Returns a\n  ``(n_points, 6)`` Voigt array\n  $(\\sigma_{xx}, \\sigma_{yy}, \\sigma_{zz}, \\sigma_{xy},\n  \\sigma_{yz}, \\sigma_{zx})$.  The same call ``Result.stress``\n  drives.\n* :func:`femorph_solver.recover.stress_invariants` \u2014 derive von\n  Mises, hydrostatic, deviatoric, and principal stresses from a\n  Voigt-stress array.  Vectorised over the leading axis so\n  ``stress_invariants(sigma_avg)`` returns a per-node table.\n\nThe closed-form Lam\u00e9 hoop stress at the bore is\n\n\\begin{align}\\sigma_{\\theta}(a) \\;=\\; p_{i}\\,\\frac{a^{2} + b^{2}}{b^{2} - a^{2}},\\end{align}\n\n(see `sphx_glr_gallery_verification_example_verify_lame_cylinder.py`\nfor the full derivation, and\n`sphx_glr_gallery_verification_example_verify_convergence_lame.py`\nfor the convergence plot).\n\nThis example reports the bore-edge $\\sigma_{\\theta}$ for a\nthree-point mesh-refinement ladder, so the role of mesh density on\nrecovered stress is visible at a glance.\n\n## Implementation\n\nQuarter-annulus HEX8 EAS plane-strain mesh \u2014 same setup as the\nverification benchmark.  After each static solve we call:\n\n* :func:`compute_nodal_stress` for the averaged Voigt-stress field.\n* :func:`stress_invariants` to obtain the von-Mises / principal\n  views of the same field.\n\n## References\n\n* Timoshenko, S. P. and Goodier, J. N. (1970) *Theory of\n  Elasticity*, 3rd ed., McGraw-Hill, \u00a733.\n* Cook, R. D., Malkus, D. S., Plesha, M. E., Witt, R. J. (2002)\n  *Concepts and Applications of Finite Element Analysis*, 4th\n  ed., Wiley, \u00a76.12 \u2014 superconvergent stress recovery.\n* Zienkiewicz, O. C. and Zhu, J. Z. (1992) \"The superconvergent\n  patch recovery and a-posteriori error estimates,\"\n  *International Journal for Numerical Methods in Engineering*\n  33 (7), 1331\u20131364 \u2014 the canonical \"do better than averaging\"\n  paper.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from __future__ import annotations\n\nimport math\n\nimport numpy as np\nimport pyvista as pv\n\nimport femorph_solver\nfrom femorph_solver import ELEMENTS\nfrom femorph_solver.recover import compute_nodal_stress, stress_invariants"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Problem data + closed-form reference\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "a = 0.10  # bore radius [m]\nb = 0.20  # outer radius [m]\nt_axial = 0.02  # plane-strain slab thickness [m]\np_i = 1.0e7  # internal pressure [Pa]\nE = 2.0e11\nNU = 0.30\nRHO = 7850.0\n\nsigma_theta_a_pub = p_i * (a**2 + b**2) / (b**2 - a**2)\nsigma_r_a_pub = -p_i\n\nprint(\"Lam\u00e9 thick cylinder \u2014 bore stress recovery via compute_nodal_stress\")\nprint(f\"  reference \u03c3_\u03b8(a) = {sigma_theta_a_pub / 1e6:7.3f} MPa  (Timoshenko & Goodier \u00a733)\")\nprint(f\"  reference \u03c3_r(a) = {sigma_r_a_pub / 1e6:7.3f} MPa  (= -p_i, exact)\")\n\n\ndef build_solve_recover(\n    n_theta: int, n_rad: int\n) -> tuple[\n    pv.UnstructuredGrid,\n    np.ndarray,\n    np.ndarray,\n    np.ndarray,\n]:\n    \"\"\"Build mesh, solve, and recover (\u03c3_avg, \u03c3_invariants, |u|).\"\"\"\n    theta = np.linspace(0.0, 0.5 * math.pi, n_theta + 1)\n    r = np.linspace(a, b, n_rad + 1)\n\n    pts_list: list[list[float]] = []\n    for kz in (0.0, t_axial):\n        for ti in theta:\n            for rj in r:\n                pts_list.append([rj * math.cos(ti), rj * math.sin(ti), kz])\n    pts_arr = np.array(pts_list, dtype=np.float64)\n\n    nx_plane = (n_theta + 1) * (n_rad + 1)\n    n_cells = n_theta * n_rad\n    cells = np.empty((n_cells, 9), dtype=np.int64)\n    cells[:, 0] = 8\n    c = 0\n    for i in range(n_theta):\n        for j in range(n_rad):\n            n00b = i * (n_rad + 1) + j\n            n10b = i * (n_rad + 1) + (j + 1)\n            n11b = (i + 1) * (n_rad + 1) + (j + 1)\n            n01b = (i + 1) * (n_rad + 1) + j\n            cells[c, 1:] = (\n                n00b,\n                n10b,\n                n11b,\n                n01b,\n                n00b + nx_plane,\n                n10b + nx_plane,\n                n11b + nx_plane,\n                n01b + nx_plane,\n            )\n            c += 1\n    grid = pv.UnstructuredGrid(\n        cells.ravel(),\n        np.full(n_cells, 12, dtype=np.uint8),\n        pts_arr,\n    )\n\n    m = femorph_solver.Model.from_grid(grid)\n    m.assign(\n        ELEMENTS.HEX8(integration=\"enhanced_strain\"),\n        material={\"EX\": E, \"PRXY\": NU, \"DENS\": RHO},\n    )\n\n    # Plane-strain BCs + symmetry pins.\n    for k, p in enumerate(pts_arr):\n        if p[0] < 1e-9:\n            m.fix(nodes=[int(k + 1)], dof=\"UX\")\n        if p[1] < 1e-9:\n            m.fix(nodes=[int(k + 1)], dof=\"UY\")\n        m.fix(nodes=[int(k + 1)], dof=\"UZ\")\n\n    # Internal-pressure load on the bore face.\n    fx_acc: dict[int, float] = {}\n    fy_acc: dict[int, float] = {}\n    for kz in (0, 1):\n        inner = [kz * nx_plane + i * (n_rad + 1) + 0 for i in range(n_theta + 1)]\n        for seg in range(n_theta):\n            ai, bi = inner[seg], inner[seg + 1]\n            ds = float(np.linalg.norm(pts_arr[bi] - pts_arr[ai]))\n            mid = 0.5 * (pts_arr[ai] + pts_arr[bi])\n            rxy = np.array([mid[0], mid[1]])\n            nrm = float(np.linalg.norm(rxy))\n            outward = rxy / nrm if nrm > 1e-12 else np.zeros(2)\n            f_seg = p_i * ds * (t_axial / 2.0)\n            for n_idx in (ai, bi):\n                fx_acc[n_idx] = fx_acc.get(n_idx, 0.0) + 0.5 * f_seg * outward[0]\n                fy_acc[n_idx] = fy_acc.get(n_idx, 0.0) + 0.5 * f_seg * outward[1]\n    for n_idx, fx in fx_acc.items():\n        fy = fy_acc[n_idx]\n        m.apply_force(int(n_idx + 1), fx=fx, fy=fy)\n\n    res = m.solve_static()\n    u = np.asarray(res.displacement, dtype=np.float64).ravel()\n    sigma = compute_nodal_stress(m, u)\n    invariants = stress_invariants(sigma)\n    u_mag = np.linalg.norm(u.reshape(-1, 3), axis=1)\n    grid_with_data = grid.copy()\n    grid_with_data.point_data[\"sigma_xx\"] = sigma[:, 0]\n    grid_with_data.point_data[\"sigma_yy\"] = sigma[:, 1]\n    grid_with_data.point_data[\"sigma_vm\"] = invariants[\"von_mises\"]\n    grid_with_data.point_data[\"s1\"] = invariants[\"s1\"]\n    grid_with_data.point_data[\"|u|\"] = u_mag\n    grid_with_data.point_data[\"displacement\"] = u.reshape(-1, 3)\n    return grid_with_data, sigma, invariants, pts_arr"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Mesh-refinement ladder\n\nThe recovered $\\sigma_{\\theta}$ at the bore converges from\nbelow as the mesh refines.  At the equator (``\u03b8=0``) the radial\ndirection is $+x$ and the hoop is $+y$, so we read\n$\\sigma_{\\theta}$ off as $\\sigma_{yy}$.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print()\nprint(\n    f\"  {'(N_\u03b8, N_r)':>11}  {'\u03c3_\u03b8 FE [MPa]':>14}  {'err vs pub':>11}  \"\n    f\"{'\u03c3_r FE [MPa]':>14}  {'\u03c3_VM at bore [MPa]':>20}\"\n)\nprint(f\"  {'-' * 11:>11}  {'-' * 14:>14}  {'-' * 11:>11}  {'-' * 14:>14}  {'-' * 20:>20}\")\n\nladder = ((12, 8), (24, 16), (36, 24))\nlast_grid = None\nlast_pts = None\nfor n_th, n_r in ladder:\n    g_, sigma_, inv_, pts_ = build_solve_recover(n_th, n_r)\n    i_bore = int(np.argmin(np.linalg.norm(pts_ - np.array([a, 0.0, 0.0]), axis=1)))\n    s_yy = float(sigma_[i_bore, 1])\n    s_xx = float(sigma_[i_bore, 0])\n    s_vm = float(inv_[\"von_mises\"][i_bore])\n    err = (s_yy - sigma_theta_a_pub) / sigma_theta_a_pub * 100.0\n    print(\n        f\"  ({n_th:>3}, {n_r:>3})  {s_yy / 1e6:>12.3f}  \"\n        f\"{err:>+10.3f}%  {s_xx / 1e6:>12.3f}  {s_vm / 1e6:>18.3f}\"\n    )\n    last_grid = g_\n    last_pts = pts_"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Render the \u03c3_VM (von Mises) field on the finest mesh\n\n``stress_invariants`` returns the full set of derived scalars in\none pass; here we render the von-Mises field which is the\ncombined-stress measure design codes most often gate against.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "assert last_grid is not None\nassert last_pts is not None\ni_bore = int(np.argmin(np.linalg.norm(last_pts - np.array([a, 0.0, 0.0]), axis=1)))\nwarped = last_grid.warp_by_vector(\"displacement\", factor=2.0e3)\n\nplotter = pv.Plotter(off_screen=True, window_size=(720, 540))\nplotter.add_mesh(last_grid, style=\"wireframe\", color=\"grey\", opacity=0.35, line_width=1)\nplotter.add_mesh(\n    warped,\n    scalars=\"sigma_vm\",\n    cmap=\"plasma\",\n    show_edges=False,\n    scalar_bar_args={\"title\": \"\u03c3_VM  [Pa]  (deformed \u00d72 000)\"},\n)\ndisp = np.asarray(last_grid.point_data[\"displacement\"])\nplotter.add_points(\n    last_pts[i_bore : i_bore + 1] + 2.0e3 * disp[i_bore],\n    render_points_as_spheres=True,\n    point_size=18,\n    color=\"#d62728\",\n    label=f\"bore \u2014 \u03c3_VM = {last_grid.point_data['sigma_vm'][i_bore] / 1e6:.3f} MPa\",\n)\nplotter.add_legend()\nplotter.view_xy()\nplotter.camera.zoom(1.05)\nplotter.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Take-aways\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "print()\nprint(\"Take-aways:\")\nprint(\n    \"  \u2022 compute_nodal_stress(m, u) returns a (n_points, 6) Voigt array of \"\n    \"node-averaged stress.  Same call Result.stress drives.\"\n)\nprint(\n    \"  \u2022 stress_invariants(sigma) is vectorised \u2014 pass a per-node array, \"\n    \"get a dict of per-node von-Mises / principal / hydrostatic / deviatoric scalars.\"\n)\nprint(\n    \"  \u2022 Recovered \u03c3_\u03b8 converges to the closed form from below \u2014 coarse meshes \"\n    \"under-predict at the bore because the steep radial gradient is poorly \"\n    \"represented by trilinear shape functions on the inner ring of elements.\"\n)\nprint(\n    \"  \u2022 Stress is always less accurate than displacement on the same mesh \"\n    \"(d = O(h\u00b2), \u03c3 = O(h) for HEX8); see \"\n    \":ref:`sphx_glr_gallery_verification_example_verify_convergence_lame.py` \"\n    \"for the full convergence curve.\"\n)"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}