{
  "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\n# TRUSS2 \u2014 axial truss under end load\n\nSingle TRUSS2 spar fixed at node 1 and pulled along the x-axis at\nnode 2. The tip displacement is compared to the closed-form rod\nsolution ``u = P L / (E A)`` and the axial stress is plotted on the\ndeformed mesh.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "from __future__ import annotations\n\nimport numpy as np\nimport pyvista as pv\nfrom vtkmodules.util.vtkConstants import VTK_LINE\n\nimport femorph_solver\nfrom femorph_solver import ELEMENTS"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Problem data\nSteel rod, 1 m long, 100 mm\u00b2 cross-section, pulled with a 1 kN tip load.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "E = 2.1e11  # Pa\nA = 1.0e-4  # m\u00b2 (100 mm\u00b2)\nL = 1.0  # m\nP = 1.0e3  # N (tensile)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Build the model\n``femorph_solver.Model.assign`` collapses the legacy ``et / mp / r``\ntrio into one call.  A TRUSS2 has three translational DOFs per node,\nso we only fix UX at node 2's support and zero-out the transverse\nDOFs to kill the transverse rigid modes.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "points = np.array([[0.0, 0.0, 0.0], [L, 0.0, 0.0]], dtype=np.float64)\ncells = np.array([2, 0, 1], dtype=np.int64)\ncell_types = np.array([VTK_LINE], dtype=np.uint8)\ngrid = pv.UnstructuredGrid(cells, cell_types, points)\n\nm = femorph_solver.Model.from_grid(grid)\nm.assign(\n    ELEMENTS.TRUSS2,\n    material={\"EX\": E, \"DENS\": 7850.0},\n    real=(A,),\n)\n\nm.fix(nodes=[1], dof=\"ALL\")  # clamp all translational DOFs at node 1\nm.fix(nodes=[2], dof=\"UY\")  # suppress transverse rigid-body motion\nm.fix(nodes=[2], dof=\"UZ\")\nm.apply_force(2, fx=P)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Static solve + analytical comparison\nThe rod equation gives ``u_tip = P L / (E A)``. With a single TRUSS2\nthis is exact (linear shape functions are sufficient for a prismatic\nbar under end load).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "res = m.solve_static()\n\ndof = m.dof_map()\ntip_ux_row = np.where((dof[:, 0] == 2) & (dof[:, 1] == 0))[0][0]\nu_tip = res.displacement[tip_ux_row]\nu_expected = P * L / (E * A)\n\nprint(f\"TRUSS2 tip UX        = {u_tip:.6e} m\")\nprint(f\"Analytical PL/(EA)    = {u_expected:.6e} m\")\nassert np.isclose(u_tip, u_expected, rtol=1e-10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Post-processing: axial stress\nFor a single TRUSS2 the axial stress is uniform: ``\u03c3 = E \u00b7 (\u0394u / L)``\nwhere ``\u0394u`` is the elongation. Carry it as cell data on the mesh for\nplotting.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "sigma_axial = E * (u_tip / L)\nprint(f\"Axial stress          = {sigma_axial:.3e} Pa (= P/A = {P / A:.3e})\")\n\ngrid = m.grid.copy()\ndisplacement = np.zeros((grid.n_points, 3), dtype=np.float64)\nfor i, node_num in enumerate(grid.point_data[\"ansys_node_num\"]):\n    rows = np.where(dof[:, 0] == int(node_num))[0]\n    for r in rows:\n        displacement[i, int(dof[r, 1])] = res.displacement[r]\ngrid.point_data[\"displacement\"] = displacement\ngrid.cell_data[\"sigma_axial\"] = np.array([sigma_axial])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plot the deformed truss coloured by axial stress\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "warped = grid.warp_by_vector(\"displacement\", factor=1.0e5)\nplotter = pv.Plotter(off_screen=True)\nplotter.add_mesh(\n    grid,\n    style=\"wireframe\",\n    color=\"gray\",\n    line_width=3,\n    label=\"undeformed\",\n)\nplotter.add_mesh(\n    warped,\n    scalars=\"sigma_axial\",\n    line_width=6,\n    scalar_bar_args={\"title\": \"axial stress [Pa]\"},\n    label=\"deformed (\u00d71e5)\",\n)\nplotter.add_legend()\nplotter.add_axes()\nplotter.show()"
      ]
    }
  ],
  "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
}