elastodynamics.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "elastodynamics.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "_SkQtiQnrCfB"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import zipfile\n",
        "try:\n",
        "  import google.colab\n",
        "except ImportError:\n",
        "  from dolfin import *\n",
        "else:\n",
        "  try:\n",
        "    from dolfin import *\n",
        "  except ImportError:\n",
        "    !wget \"https://fem-on-colab.github.io/releases/fenics-install.sh\" -O \"/tmp/fenics-install.sh\" && bash \"/tmp/fenics-install.sh\"\n",
        "    from dolfin import *\n",
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# 1. Read the mesh for this problem\n",
        "\n",
        "!wget \"https://github.com/niravshah241/elastodynamics_LSTM_PINN/archive/refs/heads/main.zip\" -O \"/tmp/github_files.zip\"\n",
        "\n",
        "zip_ref = zipfile.ZipFile(\"/tmp/github_files.zip\", 'r')\n",
        "zip_ref.extractall(\"/tmp\")\n",
        "zip_ref.close()\n",
        "\n",
        "mesh = Mesh(\"/tmp/elastodynamics_LSTM_PINN-main/mesh_data/mesh.xml\")\n",
        "subdomains = MeshFunction(\"size_t\", mesh, \"/tmp/elastodynamics_LSTM_PINN-main/mesh_data/mesh_physical_region.xml\")\n",
        "boundaries = MeshFunction(\"size_t\", mesh, \"/tmp/elastodynamics_LSTM_PINN-main/mesh_data/mesh_facet_region.xml\")\n",
        "\n",
        "dx = Measure(\"dx\")(subdomain_data=subdomains)\n",
        "ds = Measure(\"ds\")(subdomain_data=boundaries)\n",
        "ds_sym = ds(2) + ds(3) + ds(6)\n",
        "ds_bottom = ds(1)\n",
        "ds_sf = ds(5) + ds(7) + ds(8) + ds(9) + ds(10)\n",
        "ds_top = ds(4) + ds(11) + ds(12) + ds(13) + ds(14)\n",
        "ds_out = ds(15)\n",
        "dx_cc1 = dx(1) \n",
        "dx_cc2 = dx(2)\n",
        "dx_cc2_upper = dx(4)\n",
        "dx_bb_surround = dx(3)\n",
        "dx_cc = dx(6)\n",
        "dx_cb = dx(5)\n",
        "dx_ss = dx(7)"
      ],
      "metadata": {
        "id": "49Nwsts6OYtk"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# 2A. Material properties\n",
        "#Youngs modulus\n",
        "E_cc1 = 9.89E9\n",
        "E_cc2 = E_cc2_upper = 14.2E9\n",
        "E_bb_surround = 14.4E9\n",
        "E_cb = 58.1E9\n",
        "E_ss = 200.0E9\n",
        "E_cc = 1.75E9\n",
        "\n",
        "#Poisson ratio\n",
        "nu_cc1 = 0.3\n",
        "nu_cc2 = nu_cc2_upper = 0.2\n",
        "nu_bb_surround = 0.2\n",
        "nu_cb = 0.08\n",
        "nu_ss = 0.25\n",
        "nu_cc = 0.13\n",
        "\n",
        "#Lam\\'e parameters\n",
        "mu_cc1, lmbda_cc1 = Constant(E_cc1/(2.*(1.+nu_cc1))), Constant(E_cc1*nu_cc1/((1.+nu_cc1)*(1.-2.*nu_cc1)))\n",
        "mu_cc2, lmbda_cc2 = Constant(E_cc2/(2.*(1.+nu_cc2))), Constant(E_cc2*nu_cc2/((1.+nu_cc2)*(1.-2.*nu_cc2)))\n",
        "mu_cc2_upper, lmbda_cc2_upper = Constant(E_cc2_upper/(2.*(1.+nu_cc2_upper))), Constant(E_cc2_upper*nu_cc2_upper/((1.+nu_cc2_upper)*(1.-2.*nu_cc2_upper)))\n",
        "mu_bb_surround, lmbda_bb_surround = Constant(E_bb_surround/(2.*(1.+nu_bb_surround))), Constant(E_bb_surround*nu_bb_surround/((1.+nu_bb_surround)*(1.-2.*nu_bb_surround)))\n",
        "mu_cb, lmbda_cb = Constant(E_cb/(2.*(1.+nu_cb))), Constant(E_cb*nu_cb/((1.+nu_cb)*(1.-2.*nu_cb)))\n",
        "mu_ss, lmbda_ss = Constant(E_ss/(2.*(1.+nu_ss))), Constant(E_ss*nu_ss/((1.+nu_ss)*(1.-2.*nu_ss)))\n",
        "mu_cc, lmbda_cc = Constant(E_cc/(2.*(1.+nu_cc))), Constant(E_cc*nu_cc/((1.+nu_cc)*(1.-2.*nu_cc)))\n",
        "\n",
        "#Material densities\n",
        "rho_cc1 = 1. #3520.0\n",
        "rho_cc2 = rho_cc2_upper = 1. #3200.0\n",
        "rho_bb_surround = 1. #2800.0\n",
        "rho_cb = 1. #1760.0\n",
        "rho_ss = 1. #7900.0\n",
        "rho_cc = 1. #2080.0\n",
        "\n",
        "# 2B. Rayleigh damping coefficient\n",
        "eta_m = Constant(0.0)\n",
        "eta_k = Constant(0.0)\n",
        "\n",
        "# 2C. Generalised alpha method parameters\n",
        "alpha_m = Constant(0.)\n",
        "alpha_f = Constant(0.)\n",
        "gamma = Constant(0.5 + alpha_f - alpha_m)\n",
        "beta = Constant((gamma+0.5)**2/4.)"
      ],
      "metadata": {
        "id": "_vSHYjGZPPmG"
      },
      "execution_count": 23,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# 3A. Function spaces\n",
        "VM = VectorFunctionSpace(mesh,\"CG\",1)\n",
        "VMsig = TensorFunctionSpace(mesh,\"DG\",0)\n",
        "u_, phi = TrialFunction(VM), TestFunction(VM)\n",
        "u = Function(VM,name=\"Displacement\")\n",
        "u_old = Function(VM)\n",
        "v_old = Function(VM)\n",
        "a_old = Function(VM)\n",
        "\n",
        "x0 = Expression(\"x[0]\",element=VM.sub(0).ufl_element())\n",
        "x1 = Expression(\"x[1]\",element=VM.sub(1).ufl_element())\n",
        "\n",
        "# 3B. # Boundary condition\n",
        "bcs = [DirichletBC(VM.sub(1), Constant(0.), boundaries, 1), DirichletBC(VM.sub(0), Constant(0.), boundaries, 2), DirichletBC(VM.sub(0), Constant(0.), boundaries, 3), \\\n",
        "       DirichletBC(VM.sub(0), Constant(0.), boundaries, 6), DirichletBC(VM.sub(1), Constant(0.), boundaries, 4), DirichletBC(VM.sub(1), Constant(0.), boundaries, 11), \\\n",
        "\t\t\t DirichletBC(VM.sub(1), Constant(0.), boundaries, 12), DirichletBC(VM.sub(1), Constant(0.), boundaries, 13), DirichletBC(VM.sub(1), Constant(0.), boundaries, 14)]"
      ],
      "metadata": {
        "id": "9fKoBti_QLP5"
      },
      "execution_count": 24,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# 4A.  Time steps\n",
        "T = 40.\n",
        "Nsteps = 100\n",
        "dt = Constant(T/Nsteps)\n",
        "\n",
        "# 3B. Source term/ Extenal force\n",
        "n = FacetNormal(mesh)\n",
        "rho_liquid = 7874.\n",
        "H = max(mesh.coordinates()[:,1])\n",
        "p = Expression(\"x[1] <= t*H/T ? rho * g * (t*H/T-x[1]) : 0\", t=0, H = H, T = T, rho=rho_liquid, g = 9.81, degree=1)"
      ],
      "metadata": {
        "id": "flqCozXKP4_1"
      },
      "execution_count": 25,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# 5A. Stress tensor\n",
        "def sigma_cc1(r):\n",
        "\treturn 2.*mu_cc1*sym(grad(r)) + lmbda_cc1*tr(sym(grad(r)))*Identity(len(r))\n",
        "def sigma_cc2(r):\n",
        "\treturn 2.*mu_cc2*sym(grad(r)) + lmbda_cc2*tr(sym(grad(r)))*Identity(len(r))\n",
        "def sigma_cc2_upper(r):\n",
        "\treturn 2.*mu_cc2_upper*sym(grad(r)) + lmbda_cc2_upper*tr(sym(grad(r)))*Identity(len(r))\n",
        "def sigma_bb_surround(r):\n",
        "\treturn 2.*mu_bb_surround*sym(grad(r)) + lmbda_bb_surround*tr(sym(grad(r)))*Identity(len(r))\n",
        "def sigma_cb(r):\n",
        "\treturn 2.*mu_cb*sym(grad(r)) + lmbda_cb*tr(sym(grad(r)))*Identity(len(r))\n",
        "def sigma_ss(r):\n",
        "\treturn 2.*mu_ss*sym(grad(r)) + lmbda_ss*tr(sym(grad(r)))*Identity(len(r))\n",
        "def sigma_cc(r):\n",
        "\treturn 2.*mu_cc*sym(grad(r)) + lmbda_cc*tr(sym(grad(r)))*Identity(len(r))\n",
        "\n",
        "# 5B. Mass matrix\n",
        "def m(u_,phi):\n",
        "\treturn rho_cc1 * inner(u_,phi) * dx_cc1 + rho_cc2 * inner(u_,phi) * dx_cc2 + rho_bb_surround * inner(u_,phi) * dx_bb_surround + rho_cc2_upper * inner(u_,phi) * dx_cc2_upper \\\n",
        "\t+ rho_cb * inner(u_,phi) * dx_cb + rho_ss * inner(u_,phi) * dx_ss + rho_cc * inner(u_,phi) * dx_cc\n",
        "\n",
        "# 5C. Stiffnss matrix\n",
        "def k(u_,phi):\n",
        "\treturn inner(sigma_cc1(u_),sym(grad(phi))) * dx_cc1 + inner(sigma_cc2(u_),sym(grad(phi))) * dx_cc2 + inner(sigma_cc2_upper(u_),sym(grad(phi))) * dx_cc2_upper + \\\n",
        "\tinner(sigma_bb_surround(u_),sym(grad(phi))) * dx_bb_surround + inner(sigma_cb(u_),sym(grad(phi))) * dx_cb + inner(sigma_ss(u_),sym(grad(phi))) * dx_ss + \\\n",
        "\tinner(sigma_cc(u_),sym(grad(phi))) * dx_cc\n",
        "\n",
        "# 5D. Damping matrix\n",
        "def c(u_,phi):\n",
        "\treturn eta_m*m(u_,phi) + eta_k*k(u_,phi)\n",
        "\n",
        "# 5E. Work of external forces\n",
        "def Wext(phi,p):\n",
        "\treturn dot(phi,p*n) * ds_sf\n",
        "\n",
        "def update_a(u,u_old,v_old,a_old,ufl=True):\n",
        "\tif ufl:\n",
        "\t\tdt_ = dt\n",
        "\t\tbeta_ = beta\n",
        "\telse:\n",
        "\t\tdt_ = float(dt)\n",
        "\t\tbeta_ = float(beta)\n",
        "\treturn (u - u_old - dt*v_old)/(beta * dt**2) - (1 - 2*beta) * a_old /(2*beta)\n",
        "\n",
        "def update_v(a,u_old,v_old,a_old,ufl=True):\n",
        "\tif ufl:\n",
        "\t\tdt_ = dt\n",
        "\t\tgamma_ = gamma\n",
        "\telse:\n",
        "\t\tdt_ = float(dt)\n",
        "\t\tgamma_ = float(gamma)\n",
        "\treturn v_old + dt_ * ((1-gamma_) * a_old + gamma_ * a)\n",
        "\n",
        "def update_fields(u,u_old,v_old,a_old):\n",
        "\t#Store reference\n",
        "\tu_vec, u0_vec = u.vector(), u_old.vector()\n",
        "\tv0_vec, a0_vec = v_old.vector(), a_old.vector()\n",
        "\t# Updated velocity and accelaration\n",
        "\ta_current = update_a(u_vec,u0_vec,v0_vec,a0_vec,ufl=False)\n",
        "\tv_current = update_v(a_current,u0_vec,v0_vec,a0_vec,ufl=False)\n",
        "\t#Update v_old and a_old\n",
        "\ta_old.vector()[:] = a_current\n",
        "\tv_old.vector()[:] = v_current\n",
        "\tu_old.vector()[:] = u.vector().get_local()\n",
        "\n",
        "def avg(x_old,x_new,alpha):\n",
        "\treturn (1-alpha) * x_new + alpha * x_old"
      ],
      "metadata": {
        "id": "X34uAUVgaZ0n"
      },
      "execution_count": 26,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Residual forms\n",
        "a_new = update_a(u_, u_old, v_old, a_old, ufl=True)\n",
        "v_new = update_v(a_new, u_old, v_old, a_old, ufl=True)\n",
        "res_total = m(avg(a_old,a_new,alpha_m),phi) + c(avg(v_old,v_new,alpha_f),phi) + k(avg(u_old,u_,alpha_f),phi) - Wext(phi,p)\n",
        "a_form = lhs(res_total)\n",
        "L_form = rhs(res_total)\n",
        "K, res = assemble_system(a_form, L_form, bcs)\n",
        "solver = LUSolver(K,\"mumps\")\n",
        "solver.parameters[\"symmetric\"] = True\n",
        "\n",
        "# Time-stepping\n",
        "num_subdomains = 7\n",
        "time = np.linspace(0,T,Nsteps+1)\n",
        "u_max_x = np.zeros((Nsteps+1,num_subdomains))\n",
        "u_min_x = np.zeros((Nsteps+1,num_subdomains))\n",
        "u_max_y = np.zeros((Nsteps+1,num_subdomains))\n",
        "u_min_y = np.zeros((Nsteps+1,num_subdomains))\n",
        "von_mises_max = np.zeros((Nsteps+1,num_subdomains))\n",
        "energies = np.zeros((Nsteps+1,4))\n",
        "E_damp = 0\n",
        "E_ext = 0\n",
        "sig = Function(VMsig,name=\"sigma\")\n",
        "xdmf_file = XDMFFile(\"/content/drive/My Drive/Colab Notebooks/solution_field/elastodynamics_results.xdmf\")\n",
        "xdmf_file.parameters[\"flush_output\"] = True\n",
        "xdmf_file.parameters[\"functions_share_mesh\"] = True\n",
        "xdmf_file.parameters[\"rewrite_function_mesh\"] = False\n",
        "\n",
        "# For stress computation\n",
        "def local_project(u_sol,VMsig,sigma_sol=None):\n",
        "\tsigma_trial = TrialFunction(VMsig)\n",
        "\tsigma_test = TestFunction(VMsig)\n",
        "\ta_proj = inner(sigma_trial,sigma_test) * dx\n",
        "\tb_proj = inner(sigma_cc1(u_sol),sigma_test) * dx_cc1 + inner(sigma_cc2(u_sol),sigma_test) * dx_cc2 + inner(sigma_bb_surround(u_sol),sigma_test) * dx_bb_surround + \\\n",
        "\tinner(sigma_cc2_upper(u_sol),sigma_test) * dx_cc2_upper + inner(sigma_cb(u_sol),sigma_test) * dx_cb + inner(sigma_ss(u_sol),sigma_test) * dx_ss + inner(sigma_cc(u_sol),sigma_test) * dx_cc\n",
        "\tsolver = LocalSolver(a_proj,b_proj)\n",
        "\tsolver.factorize()\n",
        "\tif sigma_sol is None:\n",
        "\t\tsigma_sol = Function(VMsig)\n",
        "\t\tsolver.solve_local_rhs(sigma_sol)\n",
        "\t\treturn sigma_sol\n",
        "\telse:\n",
        "\t\tsolver.solve_local_rhs(sigma_sol)\n",
        "\t\treturn\n",
        "\n",
        "def subdomainwise_displacement_stress(u_sol,VM,sigma_funcs,num_subdomains,volume_measures):\n",
        "\tmax_min_displacement_max_stress = np.zeros([5,num_subdomains]) #u_max_x,u_min_x,u_max_y,u_min_y,sigma_von_mises\n",
        "\tu_sol_subdomain = Function(VM)\n",
        "\tu = TrialFunction(VM)\n",
        "\tv = TestFunction(VM)\n",
        "\tif VM.ufl_element().degree() == 1:\n",
        "\t\tVM_von_mises = FunctionSpace(mesh, 'DG', 0)\n",
        "\telse:\n",
        "\t\tVM_von_mises = FunctionSpace(mesh, 'CG', VM.ufl_element().degree() - 1)\n",
        "\taM = inner(u,v)*dx\n",
        "\tfor i in range(num_subdomains):\n",
        "\t\tbM = inner(u_sol,v) * volume_measures[i]\n",
        "\t\tsolve(aM==bM,u_sol_subdomain)\n",
        "\t\ts = sigma_funcs[i](u_sol_subdomain) - (1./3.)*tr(sigma_funcs[i](u_sol_subdomain))*Identity(2)\n",
        "\t\tvon_Mises = sqrt(3./2.*inner(s, s))\n",
        "\t\tvon_Mises = project(von_Mises, VM_von_mises)\n",
        "\t\tmax_min_displacement_max_stress[0,i] = max(u_sol.sub(0).vector().get_local())\n",
        "\t\tmax_min_displacement_max_stress[1,i] = min(u_sol.sub(0).vector().get_local())\n",
        "\t\tmax_min_displacement_max_stress[2,i] = min(u_sol.sub(1).vector().get_local())\n",
        "\t\tmax_min_displacement_max_stress[3,i] = min(u_sol.sub(1).vector().get_local())\n",
        "\t\tmax_min_displacement_max_stress[4,i] = max(von_Mises.vector().get_local())\n",
        "\treturn max_min_displacement_max_stress\n",
        "\n",
        "for (i,dt) in enumerate(np.diff(time)):\n",
        "\tt = time[i+1]\n",
        "\tprint(\"Time: \",t)\n",
        "\t\n",
        "\tp.t = t - float(alpha_f*dt)\n",
        "\tp_old = p\n",
        "\n",
        "\tres = assemble(L_form)\n",
        "\tbcs[0].apply(res)\n",
        "\tsolver.solve(K, u.vector(), res)\n",
        "\t\n",
        "\tupdate_fields(u, u_old, v_old, a_old)\n",
        "\tALE.move(mesh,u)\n",
        "\txdmf_file.write(u,t)\n",
        "\tALE.move(mesh,project(-u,VM))\n",
        "\t\n",
        "\tlocal_project(u,VMsig,sig)\n",
        "\txdmf_file.write(sig, t)\n",
        "\n",
        "\tp.t = t\n",
        "\t\n",
        "\tdisplacement_stress_results = subdomainwise_displacement_stress(u, VM, [sigma_cc1,sigma_cc2,sigma_bb_surround,sigma_cc2_upper,sigma_cb,sigma_ss,sigma_cc], 7, \\\n",
        "\t                                                                [dx_cc1,dx_cc2,dx_bb_surround,dx_cc2_upper,dx_cb,dx_ss,dx_cc])\n",
        "\tu_max_x[i+1,:], u_min_x[i+1,:], u_max_y[i+1,:], u_min_y[i+1,:], von_mises_max[i+1,:] = displacement_stress_results[0,:], displacement_stress_results[1,:], displacement_stress_results[2,:], displacement_stress_results[3,:], displacement_stress_results[4,:]\n",
        "\t\n",
        "\tE_elas = assemble(0.5*k(u_old,u_old))\n",
        "\tE_kin = assemble(0.5*m(v_old,v_old))\n",
        "\tE_damp += dt*assemble(c(v_old,v_old)) #TODO: check\n",
        "\t#E_ext = assemble(Wext(u_old,p_old)) # TODO\n",
        "\tE_tot = E_elas + E_kin + E_damp #- E_ext\n",
        "\tprint(\"E_elas\",E_elas)\n",
        "\tprint(\"E_kin\",E_kin)\n",
        "\tprint(\"E_damp\",E_damp)\n",
        "\tprint(\"E_tot\",E_tot)\n",
        "\tenergies[i+1,:] = np.array([E_elas,E_kin,E_damp,E_tot])"
      ],
      "metadata": {
        "id": "nKfaci2maaoB"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "# Plot energies evolution\n",
        "fig = plt.figure()\n",
        "plt.plot(time, energies)\n",
        "plt.legend((\"elastic\", \"kinetic\", \"damping\", \"total\"),loc=\"upper left\")\n",
        "plt.xlabel(\"Time (s)\",fontsize=16)\n",
        "plt.ylabel(\"Energies (J/m3)\",fontsize=16)\n",
        "plt.tight_layout()\n",
        "#plt.show()\n",
        "fig.savefig(\"energies.png\")\n",
        "\n",
        "# Plot Von Mises stress\n",
        "fig = plt.figure()\n",
        "for i in range(num_subdomains):\n",
        "  plt.plot(time,von_mises_max[:,i])\n",
        "plt.legend((\"cc1\",\"cc2\",\"bb_surround\",\"cc2_upper\",\"cb\",\"ss\",\"cc\"),loc=\"upper left\")\n",
        "plt.xlabel(\"Time (s)\",fontsize=16)\n",
        "plt.ylabel(\"Von Mises stress\",fontsize=16)\n",
        "plt.tight_layout()\n",
        "#plt.show()\n",
        "fig.savefig(\"von_mises_stress.png\")"
      ],
      "metadata": {
        "id": "dF1A0kH0hws1"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}