{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#Assembly\n",
    "\n",
    "Let's make two element model and assemble it. We split the previous one element model to 2 quadrilaterals, make assembly and solve it. Small modifications to functions, I think it's better that they don't allocate memory but do in place operations. \n",
    "\n",
    "**TODO**\n",
    "- Tangent stiffness is calculated using forward finite difference. I think we should try ReverseDiffSparse for it's sparse matrix support, but I don't know how to use it. Or alternatively use FAD like before and assemble after linearization. It would be nice experiment to try linearization *after* assembly, would it work?\n",
    "- Verify calculations using some well known FEM software.\n",
    "\n",
    "\n",
    "Author: Jukka Aho\n",
    "\n",
    "Email: <jukka.aho@kapsi.fi>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0x00000c95556709c2"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tic()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dict{Any,Any} with 0 entries"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type Node\n",
    "    coords\n",
    "end\n",
    "type Element\n",
    "    node_ids\n",
    "end\n",
    "elements = Dict()\n",
    "nodes = Dict()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Element([5,6,3,2])"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ndim = 2\n",
    "nnodes = 6\n",
    "\n",
    "nodes[1] = Node([0, 1, 0])\n",
    "nodes[2] = Node([5, 1, 0])\n",
    "nodes[3] = Node([10, 1, 0])\n",
    "nodes[4] = Node([0, 0, 0])\n",
    "nodes[5] = Node([5, 0, 0])\n",
    "nodes[6] = Node([10, 0, 0])\n",
    "elements[1] = Element([4, 5, 2, 1])\n",
    "elements[2] = Element([5, 6, 3, 2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(36.0,24.0)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Partial derivatives of bilinear Lagrange polynomials\n",
    "dNdξ(ξ) = [[-(1-ξ[2])/4.0    -(1-ξ[1])/4.0],\n",
    "           [ (1-ξ[2])/4.0    -(1+ξ[1])/4.0],\n",
    "           [ (1+ξ[2])/4.0     (1+ξ[1])/4.0],\n",
    "           [-(1+ξ[2])/4.0     (1-ξ[1])/4.0]]    \n",
    "\n",
    "E = 90\n",
    "ν = 0.25\n",
    "μ = E/(2*(1+ν))\n",
    "λ = E*ν/((1+ν)*(1-2*ν))\n",
    "λ = 2*λ*μ/(λ + 2*μ)\n",
    "μ, λ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "calculate_internal_energy! (generic function with 2 methods)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function calculate_internal_energy!(X, u, Wint, dNdξ, λ, μ, dim=2)\n",
    "    \"\"\"Calculate internal energy for a single element.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    X : array [dim x nodes]\n",
    "    u : array [dim x nodes]\n",
    "    dNdξ : shape function derivatives\n",
    "    λ : float\n",
    "    μ : float\n",
    "    dim : integer, optinal\n",
    "\n",
    "    Returns\n",
    "    -------\n",
    "    Nothing, this is inplace function\n",
    "    \n",
    "    \"\"\"\n",
    "    I = eye(dim)\n",
    "    \n",
    "    function J(ξ)\n",
    "        Jᵀ = X*dNdξ(ξ)\n",
    "        ∇N = inv(Jᵀ)*dNdξ(ξ)'\n",
    "        ∇u = u*∇N'\n",
    "        F = I + ∇u  # Deformation gradient\n",
    "        E = 1/2*(∇u' + ∇u + ∇u'*∇u)  # Green-Lagrange strain tensor\n",
    "        P = λ*trace(E)*I + 2*μ*E  # PK1 stress tensor\n",
    "        S = F*P  # PK2 stress tensor\n",
    "        return S*∇N*det(Jᵀ)\n",
    "    end\n",
    "\n",
    "    a = 1/sqrt(3)\n",
    "    ipoints = [[-a -a], [a -a], [a a], [-a a]]\n",
    "    iweights = [1 1 1 1]\n",
    "\n",
    "    for m = 1:length(iweights)\n",
    "        w = iweights[m]\n",
    "        ξ = ipoints[m, :]\n",
    "        Wint[:,:] += w*J(ξ)\n",
    "    end\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2x4 Array{Float64,2}:\n",
       " -60.6667  -1.33333  133.333  -71.3333\n",
       " -24.0     -8.0       35.0     -3.0   "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Wint = zeros(2, 4)\n",
    "X = [0 0; 1 0; 1 1; 0 1]'\n",
    "u = [0 0; 0 0; 1 0; 0 0]'\n",
    "calculate_internal_energy!(X, u, Wint, dNdξ, λ, μ)\n",
    "Wint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "assemble! (generic function with 1 method)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function assemble!(u, R)\n",
    "    \"\"\" Assemble global residual vector R = T - F\n",
    "    \"\"\"\n",
    "    R[:] = 0.0\n",
    "    u = reshape(u, ndim, nnodes)\n",
    "    R = reshape(R, ndim, nnodes)\n",
    "    \n",
    "    Xe = zeros(ndim, 4)\n",
    "    Winte = zeros(ndim, 4)\n",
    "    \n",
    "    # Internal forces, T\n",
    "    for i=1:length(elements)  # loop through elements\n",
    "        Xe[:,:] = 0.0  # FIXME: how to efficiently empty array?\n",
    "        el = elements[i]\n",
    "        nids = el.node_ids\n",
    "        for i=1:length(nids)  # loop through nodes\n",
    "            Xe[:,i] = nodes[nids[i]].coords[1:2]\n",
    "        end\n",
    "        Winte[:,:] = 0.0\n",
    "        calculate_internal_energy!(Xe, u[:,nids], Winte, dNdξ, λ, μ)\n",
    "        R[:,nids] += Winte\n",
    "    end\n",
    "\n",
    "    # External forces, F\n",
    "    # T - F = T + (-F)\n",
    "    R[2, 3] += 2  # Force to the tip of härveli\n",
    "\n",
    "    u = reshape(u, ndim*nnodes)\n",
    "    R = reshape(R, ndim*nnodes)\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1x12 Array{Float64,2}:\n",
       " 0.0  0.0  34.8853  16.2  140.715  78.6  …  -64.1013  -36.0  -111.499  -56.8"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "R = zeros(ndim*nnodes)\n",
    "u = [0 0; 0 0; 1 0; 0 0; 0 0; 0 0]'\n",
    "#u = [0 0; 0 0; 1 0; 0 0]'\n",
    "assemble!(reshape(u, ndim*nnodes), R)\n",
    "R'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We should use ReverseDiffSparse because of it's sparse matrix support"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# FIXME: I don't know how to get these working!\n",
    "#using ReverseDiffSparse\n",
    "#using ForwardDiff\n",
    "#Kt! = forwarddiff_jacobian!(assemble!, Float64, fadtype=:dual, n=12, m=12)\n",
    "#Kt!(reshape(u, 12), Kt)\n",
    "#Kt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Lin (generic function with 2 methods)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# So we go to plan B. FIXME: change this to analytical version\n",
    "\n",
    "function Lin(f!, h=1.0e-6)\n",
    "\n",
    "    function D!(x, J)\n",
    "        J[:,:] = 0\n",
    "        N = length(x)\n",
    "        Δx = zeros(N)\n",
    "        y = zeros(N)\n",
    "        Δy = zeros(N)\n",
    "        f!(x, y)  # Evaluate function f in x and store results to y\n",
    "        for i=1:N\n",
    "            Δx[:] = 0.0\n",
    "            Δx[i] += h\n",
    "            f!(x+Δx, Δy)  # Evaluate function f in x+Δx and store results to Δy\n",
    "            J[i, :] = (Δy-y) / h\n",
    "        end\n",
    "    end\n",
    "\n",
    "    return D!\n",
    "\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Handling homogeneous Dirichlet conditions, using elimination.\n",
    "\n",
    "**INFO**: We could try something like this: http://www.code-aster.org/V2/doc/default/en/man_r/r3/r3.03.01.pdf\n",
    "\n",
    "Here's an idea how to make a very simply elimination"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2x6 Array{Int64,2}:\n",
       " 1  0  0  1  0  0\n",
       " 1  0  0  1  0  0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fixed_dofs = integer(zeros(ndim, nnodes))\n",
    "fixed_dofs[:,1] = fixed_dofs[:,4] = 1\n",
    "fixed_dofs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2x6 Array{Int64,2}:\n",
       " 0  1  1  0  1  1\n",
       " 0  1  1  0  1  1"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "free_dofs = integer(ones(ndim, nnodes)) - fixed_dofs\n",
    "free_dofs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1x8 Array{Int64,2}:\n",
       " 3  4  5  6  9  10  11  12"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "free_dofs = find(free_dofs)\n",
    "free_dofs'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Starting Newton iterations\n",
      "Iteration 1, norm = 12.031651257897487\n",
      "Iteration 2, norm = 3.141275068747564\n",
      "Iteration 3, norm = 1.1530769747882408\n",
      "Iteration 4, norm = 0.2665784114781298\n",
      "Iteration 5, norm = 0.035146969229970175\n",
      "Iteration 6, norm = 0.003125584242156337\n",
      "Iteration 7, norm = 2.227371173806294e-6\n",
      "Iteration 8, norm = 2.1397676611257455e-10\n",
      "Converged.\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "2x6 Array{Float64,2}:\n",
       " 0.0  -0.106192  -1.69115  0.0  -0.761252  -2.48604\n",
       " 0.0  -2.10509   -6.00728  0.0  -1.89221   -5.5914 "
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function solve!(u, free_dofs; max_iterations=10, eps=1.0e-7)\n",
    "    ∇R = Lin(assemble!)\n",
    "    Kt = zeros(ndim*nnodes, ndim*nnodes)\n",
    "    R = zeros(size(u))\n",
    "    println(\"Starting Newton iterations\")\n",
    "    for i=1:max_iterations\n",
    "        print(\"Iteration \",i, \", \")\n",
    "        R[:] = 0.0\n",
    "        assemble!(u, R)  # Calculate internal energy in nodes and store results to R\n",
    "        ∇R(u, Kt)  # Linearize residual in u and save result to Kt\n",
    "        # Solve !\n",
    "        du = Kt[free_dofs, free_dofs] \\ -R[free_dofs]\n",
    "        u[free_dofs] += du\n",
    "        println(\"norm = \",norm(du))\n",
    "        if norm(du) < eps\n",
    "            println(\"Converged.\")\n",
    "            break\n",
    "        end\n",
    "    end\n",
    "    return u\n",
    "end\n",
    "\n",
    "u = zeros(ndim*nnodes)\n",
    "solve!(u, free_dofs)\n",
    "u = reshape(u, ndim, nnodes)\n",
    "u"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "collapsed": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "elapsed time: 4.541729719 seconds\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "4.541729719"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "toc()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 0.3.8",
   "language": "julia",
   "name": "julia-0.3"
  },
  "language_info": {
   "name": "julia",
   "version": "0.3.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}