{
  "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# POINT_MASS \u2014 single-DOF spring-mass oscillator\n\nTwo nodes connected by a SPRING element. A POINT_MASS element sits on the\nfree node. The first modal frequency is compared to the textbook SDOF\nresult ``\u03c9 = \u221a(k / m)``.\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, VTK_VERTEX\n\nimport femorph_solver\nfrom femorph_solver import ELEMENTS"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Problem data\n``k = 1 kN/m``, ``m = 0.25 kg`` \u2192 ``\u03c9 = \u221a(4000) \u2248 63.2456 rad/s``\n(``f \u2248 10.065 Hz``).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "k = 1000.0\nmass = 0.25\nomega_expected = np.sqrt(k / mass)\nf_expected = omega_expected / (2.0 * np.pi)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Build the model\nTwo element types: SPRING (spring, TYPE 1) and POINT_MASS (point mass,\nTYPE 2). Each gets its own REAL set. The grid carries per-cell\n``ansys_elem_type_num`` / ``ansys_real_constant`` arrays that pin the\nkernel + real-set down element-by-element \u2014 the same routing legacy APDL\nuses but stamped on the grid up front.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "points = np.array([[0.0, 0.0, 0.0], [1.0, 0.0, 0.0]], dtype=np.float64)\ncells = np.array([2, 0, 1, 1, 1], dtype=np.int64)  # VTK_LINE(0,1) + VTK_VERTEX(1)\ncell_types = np.array([VTK_LINE, VTK_VERTEX], dtype=np.uint8)\ngrid = pv.UnstructuredGrid(cells, cell_types, points)\ngrid.cell_data[\"ansys_elem_type_num\"] = np.array([1, 2], dtype=np.int32)\ngrid.cell_data[\"ansys_real_constant\"] = np.array([1, 2], dtype=np.int32)\n\nmdl = femorph_solver.Model.from_grid(grid)\nmdl.assign(ELEMENTS.SPRING, real=(k,), et_id=1, real_id=1)\nmdl.assign(ELEMENTS.POINT_MASS, real=(mass,), et_id=2, real_id=2)\n\nmdl.fix(nodes=[1], dof=\"ALL\")  # clamp the spring base\nmdl.fix(nodes=[2], dof=\"UY\")  # kill transverse rigid-body modes\nmdl.fix(nodes=[2], dof=\"UZ\")\nm = mdl  # alias preserves later references in this script"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Modal solve + analytical comparison\nOnly one physical DOF is free (UX at node 2). ``modal_solve`` returns\na single positive eigenvalue whose square root is \u03c9.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "res = m.solve_modal(n_modes=1)\nomega_computed = float(np.sqrt(res.omega_sq[0]))\nf_computed = float(res.frequency[0])\n\nprint(f\"Expected \u03c9 = {omega_expected:.6f} rad/s, f = {f_expected:.6f} Hz\")\nprint(f\"Computed \u03c9 = {omega_computed:.6f} rad/s, f = {f_computed:.6f} Hz\")\nassert np.isclose(omega_computed, omega_expected, rtol=1e-10)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Plot the mode shape\nDilate the mode shape for visualisation. With POINT_MASS being a\nsingle-node element the mesh contains a VTK_LINE (the spring) and a\nVTK_VERTEX (the mass) \u2014 pyvista happily draws the combined\nunstructured grid.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "dof = m.dof_map()\nmode = res.mode_shapes[:, 0]\ngrid = m.grid.copy()\ndisplacement = 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        displacement[i, int(dof[r, 1])] = mode[r]\n# Scale so the peak is 0.3 m for clarity.\npeak = float(np.max(np.abs(displacement))) or 1.0\ndisplacement *= 0.3 / peak\ngrid.point_data[\"mode1\"] = displacement\n\nwarped = grid.warp_by_vector(\"mode1\", factor=1.0)\nplotter = pv.Plotter(off_screen=True)\nplotter.add_mesh(grid, style=\"wireframe\", color=\"gray\", line_width=3)\nplotter.add_mesh(\n    warped,\n    color=\"tomato\",\n    line_width=6,\n    render_points_as_spheres=True,\n    point_size=18,\n)\nplotter.add_text(f\"f1 = {f_computed:.3f} Hz\", font_size=12)\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
}