{
  "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# BEAM2 \u2014 cantilever tip deflection and first mode\n\nSlender steel cantilever modelled with a line of BEAM2 elements. Two\nvalidations:\n\n1. Static tip deflection matches ``P L\u00b3 / (3 E I)`` (Euler\u2013Bernoulli).\n2. First natural frequency matches ``(\u03b2\u2081 L)\u00b2 \u221a(E I / (\u03c1 A L\u2074))`` with\n   ``\u03b2\u2081 L = 1.87510407``.\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, square 50 \u00d7 50 mm section, 1 m span, 1 kN tip load.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "E = 2.1e11  # Pa\nNU = 0.3\nRHO = 7850.0  # kg/m\u00b3\nb = 0.050  # m (square side)\nA = b * b\nI = b**4 / 12.0  # about either principal axis (square)\nJ = 2.0 * I  # thin-square torsion approximation\nL = 1.0  # m\nP = 1.0e3  # N (transverse tip load, +y)\n\nN_ELEM = 10  # discretisation along the span"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Build the model\nBEAM2 has 6 DOFs per node; ``d(node, \"ALL\")`` on the clamped end\nfixes all six. Hermite-cubic shape functions recover Euler\u2013Bernoulli\nexactly for prismatic beams, so 10 elements is already machine-precise\non the static answer.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "points = np.array(\n    [[i * L / N_ELEM, 0.0, 0.0] for i in range(N_ELEM + 1)],\n    dtype=np.float64,\n)\ncells_list: list[int] = []\nfor i in range(N_ELEM):\n    cells_list.extend([2, i, i + 1])\ncells = np.asarray(cells_list, dtype=np.int64)\ncell_types = np.full(N_ELEM, VTK_LINE, dtype=np.uint8)\ngrid = pv.UnstructuredGrid(cells, cell_types, points)\n\nm = femorph_solver.Model.from_grid(grid)\nm.assign(\n    ELEMENTS.BEAM2,\n    material={\"EX\": E, \"PRXY\": NU, \"DENS\": RHO},\n    real=(A, I, I, J),\n)\n\nm.fix(nodes=[1], dof=\"ALL\")  # fully clamp node 1\nm.apply_force(N_ELEM + 1, fy=P)  # tip load in +y"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Static solve + analytical comparison\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "static = m.solve_static()\ndof = m.dof_map()\ntip_uy = np.where((dof[:, 0] == N_ELEM + 1) & (dof[:, 1] == 1))[0][0]\nu_tip = static.displacement[tip_uy]\nu_expected = P * L**3 / (3.0 * E * I)\n\nprint(f\"BEAM2 tip UY       = {u_tip:.6e} m\")\nprint(f\"Analytical PL\u00b3/(3EI) = {u_expected:.6e} m\")\nassert np.isclose(u_tip, u_expected, rtol=1e-8)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Modal solve + analytical comparison\nThe transverse DOFs in two planes give two (degenerate) first bending\nmodes. Because ``I_y = I_z`` here, femorph-solver's eigensolver will return them\nboth at the same frequency.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "modal = m.solve_modal(n_modes=4)\nBETA_L = 1.87510407  # dimensionless cantilever eigenvalue (first mode)\nomega_expected = BETA_L**2 * np.sqrt(E * I / (RHO * A * L**4))\nomega1 = float(np.sqrt(modal.omega_sq[0]))\nomega2 = float(np.sqrt(modal.omega_sq[1]))\n\nprint(f\"Expected \u03c9\u2081  = {omega_expected:.4f} rad/s\")\nprint(f\"Computed \u03c9\u2081  = {omega1:.4f} rad/s\")\nprint(f\"Computed \u03c9\u2082  = {omega2:.4f} rad/s (bending in the other plane)\")\nassert np.isclose(omega1, omega_expected, rtol=5e-3)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plot the deflected shape and the first mode\nScatter the static displacement onto the grid for visualisation, then\noverlay the first mode shape coloured by transverse amplitude.\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)\nmode_disp = np.zeros_like(disp)\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        d_idx = int(dof[r, 1])\n        if d_idx < 3:  # translations only for warping\n            disp[i, d_idx] = static.displacement[r]\n            mode_disp[i, d_idx] = modal.mode_shapes[r, 0]\n\n# Scale mode to unit peak\npeak = float(np.max(np.abs(mode_disp))) or 1.0\nmode_disp /= peak\n\ngrid.point_data[\"static_disp\"] = disp\ngrid.point_data[\"mode1_disp\"] = mode_disp\n\nplotter = pv.Plotter(shape=(1, 2), off_screen=True)\nplotter.subplot(0, 0)\nplotter.add_text(\"Static: PL\u00b3/(3EI)\", font_size=10)\nplotter.add_mesh(grid, style=\"wireframe\", color=\"gray\", line_width=2)\nplotter.add_mesh(\n    grid.warp_by_vector(\"static_disp\", factor=50.0),\n    scalars=np.linalg.norm(disp, axis=1),\n    line_width=5,\n    scalar_bar_args={\"title\": \"|u| [m]\"},\n)\nplotter.add_axes()\n\nplotter.subplot(0, 1)\nplotter.add_text(f\"Mode 1: f = {modal.frequency[0]:.1f} Hz\", font_size=10)\nplotter.add_mesh(grid, style=\"wireframe\", color=\"gray\", line_width=2)\nplotter.add_mesh(\n    grid.warp_by_vector(\"mode1_disp\", factor=0.3),\n    scalars=np.linalg.norm(mode_disp, axis=1),\n    line_width=5,\n    cmap=\"coolwarm\",\n    scalar_bar_args={\"title\": \"|\u03c6|\"},\n)\nplotter.add_axes()\nplotter.link_views()\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
}