{
  "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# SPRING \u2014 two springs in series\n\nThree co-linear nodes joined by two SPRING springs of stiffness\n``k1`` and ``k2``. With the far end fixed and a tip load ``F`` the\nfree-end displacement is ``F / k_eq`` where\n``k_eq = k1 \u00b7 k2 / (k1 + k2)``.\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\nTwo springs. ``k1`` at half the stiffness of ``k2``; their series\ncombination sits in between.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "k1 = 1000.0  # N/m\nk2 = 2000.0  # N/m\nF = 50.0  # N tip load\n\nk_eq = k1 * k2 / (k1 + k2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Build the model\nThree nodes along x. SPRING only carries an axial spring force; the\ntransverse DOFs have zero stiffness and the solver's zero-pivot guard\npins them automatically.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "points = np.array(\n    [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [2.0, 0.0, 0.0]],\n    dtype=np.float64,\n)\ncells = np.array([2, 0, 1, 2, 1, 2], dtype=np.int64)\ncell_types = np.array([VTK_LINE, VTK_LINE], dtype=np.uint8)\ngrid = pv.UnstructuredGrid(cells, cell_types, points)\n# Two real-constant ids: cell 0 uses real-id 1 (k1), cell 1 uses real-id 2 (k2).\ngrid.cell_data[\"ansys_real_constant\"] = np.array([1, 2], dtype=np.int32)\n\nm = femorph_solver.Model.from_grid(grid)\nm.assign(ELEMENTS.SPRING, real=(k1,), real_id=1)\nm.assign(ELEMENTS.SPRING, real=(k2,), real_id=2)\n\nm.fix(nodes=[1], dof=\"ALL\")  # clamp node 1\n# Restrain transverse translations on the free nodes so the DOFs the\n# solver actually keeps are just UX at nodes 2 and 3.\nfor nn in (2, 3):\n    m.fix(nodes=[nn], dof=\"UY\")\n    m.fix(nodes=[nn], dof=\"UZ\")\n\nm.apply_force(3, fx=F)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Static solve + analytical comparison\nWith both springs in equilibrium the tip displacement must equal\n``F / k_eq``. Because the springs are in series the force is the\nsame in each and we can check the intermediate displacement too\n(``u2 = F / k1``).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "res = m.solve_static()\ndof = m.dof_map()\n\nu2 = res.displacement[np.where((dof[:, 0] == 2) & (dof[:, 1] == 0))[0][0]]\nu3 = res.displacement[np.where((dof[:, 0] == 3) & (dof[:, 1] == 0))[0][0]]\n\nprint(f\"u at node 2 = {u2:.6e} m (expected {F / k1:.6e})\")\nprint(f\"u at node 3 = {u3:.6e} m (expected {F / k_eq:.6e})\")\n\nassert np.isclose(u2, F / k1, rtol=1e-12)\nassert np.isclose(u3, F / k_eq, rtol=1e-12)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Visualise the deformation\nBuild a point_data displacement vector and warp the mesh. Because the\nreal deflection is ~7.5 cm on a 2 m baseline a modest exaggeration\nmakes the difference between the two springs visible.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "grid = 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])] = res.displacement[r]\ngrid.point_data[\"displacement\"] = displacement\n\nwarped = grid.warp_by_vector(\"displacement\", 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    scalars=np.linalg.norm(displacement, axis=1),\n    line_width=6,\n    render_points_as_spheres=True,\n    point_size=14,\n    scalar_bar_args={\"title\": \"|u| [m]\"},\n)\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
}