{
  "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# QUAD4_PLANE \u2014 plane-stress uniaxial tension\n\nA square piece of plate is fixed on its left edge and pulled on its\nright edge. We check the tip displacement against Hooke's law\n``u_x = \u03c3 L / E`` and the Poisson contraction ``u_y = \u2212\u03bd u_x`` in\nplane stress, then render the deformed 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_QUAD\n\nimport femorph_solver\nfrom femorph_solver import ELEMENTS"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Problem set-up\nUnit square of steel, 1 mm thick, pulled with 1 MN total along +x.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "E = 2.0e11  # Pa\nNU = 0.3\nRHO = 7850.0  # kg/m\u00b3\nTHK = 1.0e-3  # m\nF_TOTAL = 1.0e6  # N\n\ncoords = np.array(\n    [(0.0, 0.0, 0.0), (1.0, 0.0, 0.0), (1.0, 1.0, 0.0), (0.0, 1.0, 0.0)],\n    dtype=np.float64,\n)\ncells = np.array([4, 0, 1, 2, 3], dtype=np.int64)\ncell_types = np.array([VTK_QUAD], dtype=np.uint8)\ngrid = pv.UnstructuredGrid(cells, cell_types, coords)\n\nm = femorph_solver.Model.from_grid(grid)\nm.assign(\n    ELEMENTS.QUAD4_PLANE,\n    material={\"EX\": E, \"PRXY\": NU, \"DENS\": RHO},\n    real=(THK,),\n)\n# Plane-stress KEYOPT(3)=0 is the default; annotate it explicitly:\nm.materials[1][\"_PLANE_MODE\"] = \"stress\"\n\n# Clamp the x=0 edge in UX, and kill rigid-body UY without locking Poisson.\nfor nn in (1, 4):\n    m.fix(nodes=[nn], dof=\"UX\")\nfor nn in (1, 2):\n    m.fix(nodes=[nn], dof=\"UY\")\n\nF_each = F_TOTAL / 2.0\nfor nn in (2, 3):\n    m.apply_force(nn, fx=F_each)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Static solve + analytical comparison\nOut-of-plane stress is zero (plane stress), so the effective axial stress\non the x=1 edge is ``\u03c3 = F_TOTAL / (A \u00b7 thickness)`` with edge length\n1 m, giving strain ``\u03c3 / E`` and tip displacement ``\u03c3 \u00b7 L / E``.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "res = m.solve_static()\ndof = m.dof_map()\n\nsigma = F_TOTAL / (1.0 * THK)\nux_expected = sigma / E  # L = 1 m\nuy_expected = -NU * ux_expected\n\nprint(f\"Expected u_x (+x edge) = {ux_expected:.6e} m\")\nfor nn in (2, 3):\n    row = np.where((dof[:, 0] == nn) & (dof[:, 1] == 0))[0][0]\n    print(f\"  node {nn} UX         = {res.displacement[row]:.6e}\")\n    assert np.isclose(res.displacement[row], ux_expected, rtol=1e-8)\n\nprint(f\"Expected u_y (+y edge) = {uy_expected:.6e} m\")\nfor nn in (3, 4):\n    row = np.where((dof[:, 0] == nn) & (dof[:, 1] == 1))[0][0]\n    print(f\"  node {nn} UY         = {res.displacement[row]:.6e}\")\n    assert np.isclose(res.displacement[row], uy_expected, rtol=1e-8)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plot the deformed quad, coloured by displacement magnitude\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "grid = m.grid.copy()\ndisp = np.zeros((grid.n_points, 3), dtype=np.float64)\nfor i, nn in enumerate(grid.point_data[\"ansys_node_num\"]):\n    rows = np.where(dof[:, 0] == int(nn))[0]\n    for r in rows:\n        disp[i, int(dof[r, 1])] = res.displacement[r]\ngrid.point_data[\"displacement\"] = disp\n\nplotter = pv.Plotter(off_screen=True)\nwarped = grid.warp_by_vector(\"displacement\", factor=50.0)\nplotter.add_mesh(\n    warped,\n    scalars=np.linalg.norm(disp, axis=1),\n    scalar_bar_args={\"title\": \"disp [m]\"},\n    show_edges=True,\n    cmap=\"viridis\",\n)\nplotter.add_mesh(grid, style=\"wireframe\", color=\"gray\", opacity=0.4)\nplotter.view_xy()\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
}