{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Adaptive PDE discretizations on Cartesian grids\n",
    "## Volume : Divergence form PDEs\n",
    "## Part : Linear elasticity\n",
    "## Chapter : Elastic energy"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We present a new discretization of the Dirichlet energy arising in linear elasticity, defined by a field of positive definite Hooke tensors, and associated to any small displacement field. The approach is based on Selling's decomposition of the Hooke tensor, and for this reason is only applies in dimension two at the moment.\n",
    "\n",
    "The Dirichlet energy of linear elasitcity, defined for (small) displacement maps $v : \\Omega \\to R^d$, reads\n",
    "$$\n",
    "    E(v) := \\frac 1 2 \\int_\\Omega \\sum_{ijkl} c_{ijkl}(x) \\epsilon_{ij}(x) \\epsilon_{kl}(x) \\ dx,\n",
    "$$\n",
    "where the indices $i,j,k,l$ range from $0$ to $d-1$. We denoted by $c_{ijkl}(x)$ the Hooke tensor at a point $x \\in \\Omega$, and introduced the symmetrized gradient of the displacement field, also known as the strain tensor $\\epsilon$\n",
    "$$\n",
    "    \\epsilon_{ij}(x) = \\frac 1 2 \\Big (\\frac{\\partial v_i}{\\partial x_j} + \\frac{\\partial v_j}{\\partial x_i} \\Big).\n",
    "$$\n",
    "\n",
    "**Non staggered grids.** Finite difference schemes based on staggered grids are popular for isotropic equations, and in particular for the elastic wave equation with an isotropic Hooke tensor. In contrast, the approach presented in this notebook may deal with arbitrary anisotropy, but cannot take advantage of staggered grids.\n",
    "\n",
    "**Second order scheme.** The numerical scheme presented in this notebook is second order accurate in space. Substantial modifications are required to implement a fourth order scheme in space."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[**Summary**](Summary.ipynb) of volume Divergence form PDEs, this series of notebooks.\n",
    "\n",
    "[**Main summary**](../Summary.ipynb) of the Adaptive Grid Discretizations \n",
    "\tbook of notebooks, including the other volumes.\n",
    "\n",
    "# Table of contents\n",
    "  * [1. Decomposition of a hooke tensor](#1.-Decomposition-of-a-hooke-tensor)\n",
    "    * [1.1 Generic tensor](#1.1-Generic-tensor)\n",
    "    * [1.2 Isotropic tensor](#1.2-Isotropic-tensor)\n",
    "    * [1.3 VTI tensor](#1.3-VTI-tensor)\n",
    "  * [2. Finite difference energy](#2.-Finite-difference-energy)\n",
    "    * [2.1 CFL condition](#2.1-CFL-condition)\n",
    "    * [2.2 Comparison with automatic differentiation](#2.2-Comparison-with-automatic-differentiation)\n",
    "    * [2.3 Structure of the mass matrix](#2.3-Structure-of-the-mass-matrix)\n",
    "    * [2.4 Isotropic scheme](#2.4-Isotropic-scheme)\n",
    "  * [3. Three dimensions](#3.-Three-dimensions)\n",
    "    * [3.1 Decomposition of a single tensor](#3.1-Decomposition-of-a-single-tensor)\n",
    "    * [3.2 Structure of the sparse matrix](#3.2-Structure-of-the-sparse-matrix)\n",
    "\n",
    "\n",
    "\n",
    "**Acknowledgement.** Some of the experiments presented in these notebooks are part of \n",
    "ongoing research with Ludovic Métivier and Da Chen.\n",
    "\n",
    "Copyright Jean-Marie Mirebeau, Centre Borelli, ENS Paris-Saclay, CNRS, University Paris-Saclay"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Importing the required libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.686125Z",
     "iopub.status.busy": "2024-04-30T09:01:31.685669Z",
     "iopub.status.idle": "2024-04-30T09:01:31.699898Z",
     "shell.execute_reply": "2024-04-30T09:01:31.699260Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "import sys; sys.path.insert(0,\"..\") # Allow import of agd from parent directory (useless if conda package installed)\n",
    "#from Miscellaneous import TocTools; print(TocTools.displayTOC('ElasticEnergy','Div'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.703400Z",
     "iopub.status.busy": "2024-04-30T09:01:31.703109Z",
     "iopub.status.idle": "2024-04-30T09:01:31.963249Z",
     "shell.execute_reply": "2024-04-30T09:01:31.962939Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "ExportCode"
    ]
   },
   "outputs": [],
   "source": [
    "from agd import LinearParallel as lp\n",
    "from agd import FiniteDifferences as fd\n",
    "from agd.Metrics.Seismic import Hooke\n",
    "from agd import AutomaticDifferentiation as ad\n",
    "from agd import Domain\n",
    "from agd.Plotting import savefig; #savefig.dirName = 'Images/ElasticityDirichlet'\n",
    "\n",
    "norm_infinity = ad.Optimization.norm_infinity\n",
    "norm_average = ad.Optimization.norm_average\n",
    "mica,_ = Hooke.mica # Hooke tensor associated to this crystal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.964923Z",
     "iopub.status.busy": "2024-04-30T09:01:31.964811Z",
     "iopub.status.idle": "2024-04-30T09:01:31.966842Z",
     "shell.execute_reply": "2024-04-30T09:01:31.966581Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "ExportCode"
    ]
   },
   "outputs": [],
   "source": [
    "import numpy as np; xp=np; allclose=np.allclose\n",
    "import matplotlib.pyplot as plt\n",
    "from copy import copy\n",
    "import scipy.sparse; import scipy.sparse.linalg\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.968163Z",
     "iopub.status.busy": "2024-04-30T09:01:31.968078Z",
     "iopub.status.idle": "2024-04-30T09:01:31.970009Z",
     "shell.execute_reply": "2024-04-30T09:01:31.969786Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "def ReloadPackages():\n",
    "    from Miscellaneous.rreload import rreload\n",
    "    global lp,fd,Hooke,ad,Domain\n",
    "    lp,fd,Hooke,ad,Domain = rreload([lp,fd,Hooke,ad,Domain],rootdir=\"../..\")\n",
    "    ad.array.caster = xp.asarray"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 0.1 Additional configuration\n",
    "\n",
    "Uncomment the following line to run the notebook on the GPU. (This is purely for compatibility testing, as no intensive computation is involved.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.971490Z",
     "iopub.status.busy": "2024-04-30T09:01:31.971399Z",
     "iopub.status.idle": "2024-04-30T09:01:31.973030Z",
     "shell.execute_reply": "2024-04-30T09:01:31.972812Z"
    },
    "tags": [
     "GPU_config"
    ]
   },
   "outputs": [],
   "source": [
    "#xp,mica,allclose = map(ad.cupy_friendly,(xp,mica,allclose))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Decomposition of a hooke tensor\n",
    "\n",
    "A Hooke tensor defines a quadratic form on the set of symmetric matrices $\\epsilon \\in S_d$\n",
    "$$\n",
    "    c(\\epsilon,\\epsilon) = \\sum_{ijkl} c_{ijkl} \\epsilon_{ij} \\epsilon_{kl}.\n",
    "$$\n",
    "Note that $S_d$ has dimension $D=d (d+1)/2$. We limit our attention to the case $d=2$, since the case $d=1$ is excessively trivial, and the case $d=3$ would require an implementation of the $6$-dimensional Voronoi reduction, which we do not have at this stage.\n",
    "\n",
    "We use Selling's decomposition to rewrite this quadratic form as\n",
    "$$\n",
    "    c(\\epsilon,\\epsilon) = \\sum_r \\rho_r \\mathrm{Tr}(\\epsilon m_r)^2\n",
    "$$\n",
    "where $\\rho_r \\geq 0$, $m_r \\in S_2(Z)$ is a symmetric matrix with integer coordinates, and $0 \\leq r < D (D+1)/2=6$. For that purpose, we rely on Selling's decomposition of the Hooke tensor, which applies in dimension since the linear space of $2\\times 2$ symmetric matrices has dimension $3$.\n",
    "\n",
    "<!--- In the following, we denote as usual $c(\\epsilon) := c(\\epsilon,\\epsilon)$. --->\n",
    "\n",
    "The stress tensor $\\sigma$ depends linearly on the strain tensor $\\epsilon$, for a given hooke tensor $c$, and is characterized by the identity\n",
    "$$\n",
    "    \\mathrm{Tr}(\\sigma \\epsilon) = c(\\epsilon,\\epsilon).\n",
    "$$\n",
    "With the correct index conventions, one has $\\sigma_{ij} = \\sum_{kl} c_{ijkl} \\epsilon_{kl}$, or simply $\\sigma = c \\epsilon$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Generic tensor\n",
    "\n",
    "We illustrate the decomposition on a generic tensor, describing the anisotropic elasticity of a mica rock medium, whose layers are rotated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.974350Z",
     "iopub.status.busy": "2024-04-30T09:01:31.974272Z",
     "iopub.status.idle": "2024-04-30T09:01:31.976034Z",
     "shell.execute_reply": "2024-04-30T09:01:31.975809Z"
    }
   },
   "outputs": [],
   "source": [
    "hooke = mica.extract_xz().rotate_by(0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Selling's decomposition involves $D=6$ weights and offsets, in dimension $d=2$. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.977374Z",
     "iopub.status.busy": "2024-04-30T09:01:31.977296Z",
     "iopub.status.idle": "2024-04-30T09:01:31.979484Z",
     "shell.execute_reply": "2024-04-30T09:01:31.979263Z"
    }
   },
   "outputs": [],
   "source": [
    "coefs,moffsets = hooke.Selling()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.980784Z",
     "iopub.status.busy": "2024-04-30T09:01:31.980705Z",
     "iopub.status.idle": "2024-04-30T09:01:31.983755Z",
     "shell.execute_reply": "2024-04-30T09:01:31.983522Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([ 4.39753907, 31.3882811 ,  3.86972663, 17.65302611, 41.16677899,\n",
       "       29.81855447])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coefs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The offsets are presented as symmetric matrices, with integer entries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.985139Z",
     "iopub.status.busy": "2024-04-30T09:01:31.985060Z",
     "iopub.status.idle": "2024-04-30T09:01:31.987099Z",
     "shell.execute_reply": "2024-04-30T09:01:31.986877Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2, 2, 6)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "moffsets.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.988353Z",
     "iopub.status.busy": "2024-04-30T09:01:31.988275Z",
     "iopub.status.idle": "2024-04-30T09:01:31.990330Z",
     "shell.execute_reply": "2024-04-30T09:01:31.990075Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[[1 1]\n",
      " [1 1]] \n",
      "\n",
      "[[-1 -1]\n",
      " [-1  0]] \n",
      "\n",
      "[[2 1]\n",
      " [1 1]] \n",
      "\n",
      "[[ 0  0]\n",
      " [ 0 -1]] \n",
      "\n",
      "[[-1  0]\n",
      " [ 0  0]] \n",
      "\n",
      "[[1 0]\n",
      " [0 1]] \n",
      "\n"
     ]
    }
   ],
   "source": [
    "for i in range(6): print(moffsets[...,i],\"\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to check the correctness of the decomposition, let us introduce an arbitrary $2\\times 2$ symmetric matrix, and evaluate $c(m)$ using the original expression and the decomposition."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.991653Z",
     "iopub.status.busy": "2024-04-30T09:01:31.991577Z",
     "iopub.status.idle": "2024-04-30T09:01:31.993280Z",
     "shell.execute_reply": "2024-04-30T09:01:31.993057Z"
    }
   },
   "outputs": [],
   "source": [
    "m = xp.array(((1.5,2.3),(2.3,3.9)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.994600Z",
     "iopub.status.busy": "2024-04-30T09:01:31.994520Z",
     "iopub.status.idle": "2024-04-30T09:01:31.996224Z",
     "shell.execute_reply": "2024-04-30T09:01:31.996004Z"
    }
   },
   "outputs": [],
   "source": [
    "sum_hooke = hooke.dot_AA(m) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:31.997518Z",
     "iopub.status.busy": "2024-04-30T09:01:31.997440Z",
     "iopub.status.idle": "2024-04-30T09:01:31.999288Z",
     "shell.execute_reply": "2024-04-30T09:01:31.999056Z"
    }
   },
   "outputs": [],
   "source": [
    "mm = fd.as_field(m,coefs.shape)\n",
    "sum_inner = (coefs*lp.trace(lp.dot_AA(mm,moffsets))**2).sum(axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.000602Z",
     "iopub.status.busy": "2024-04-30T09:01:32.000524Z",
     "iopub.status.idle": "2024-04-30T09:01:32.002032Z",
     "shell.execute_reply": "2024-04-30T09:01:32.001800Z"
    }
   },
   "outputs": [],
   "source": [
    "assert allclose(sum_hooke,sum_inner)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Isotropic tensor\n",
    "\n",
    "Isotropic elasticity tensors only have two degrees of freedom. Without loss of generality, we consider the Lamé parameters. These parameters relate the strain tensor $\\epsilon$ with the stress tensor $\\sigma$\n",
    "$$\n",
    "    \\sigma = 2 \\mu \\epsilon + \\lambda \\mathrm{Tr}(\\epsilon) \\mathrm{Id},\n",
    "$$\n",
    "and the quadratic form reads \n",
    "$$\n",
    "    c(\\epsilon,\\epsilon) = 2 \\mu \\mathrm{Tr}(\\epsilon^2) + \\lambda \\mathrm{Tr}(\\epsilon)^2.\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.003413Z",
     "iopub.status.busy": "2024-04-30T09:01:32.003327Z",
     "iopub.status.idle": "2024-04-30T09:01:32.005400Z",
     "shell.execute_reply": "2024-04-30T09:01:32.005166Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "An isotropic hooke tensor :\n",
      "[[5. 1. 0.]\n",
      " [1. 5. 0.]\n",
      " [0. 0. 2.]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "hooke = Hooke.from_Lame(xp.array(1.),2.) # xp.array is a type hint for GPU\n",
    "print(f\"\"\"An isotropic hooke tensor :\\n{hooke.hooke}\\n\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.006687Z",
     "iopub.status.busy": "2024-04-30T09:01:32.006612Z",
     "iopub.status.idle": "2024-04-30T09:01:32.008658Z",
     "shell.execute_reply": "2024-04-30T09:01:32.008423Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Isotropic Hooke tensors are linear combinations of: \n",
      "[[1. 1. 0.]\n",
      " [1. 1. 0.]\n",
      " [0. 0. 0.]]\n",
      "and \n",
      "[[2. 0. 0.]\n",
      " [0. 2. 0.]\n",
      " [0. 0. 1.]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "print(f\"\"\"Isotropic Hooke tensors are linear combinations of: \\n{Hooke.from_Lame(1.,0.).hooke}\\n\"\"\"\n",
    "      f\"\"\"and \\n{Hooke.from_Lame(0.,1.).hooke}\\n\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As their name suggests, isotropic hooke tensors are invariant under rotations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.010040Z",
     "iopub.status.busy": "2024-04-30T09:01:32.009962Z",
     "iopub.status.idle": "2024-04-30T09:01:32.011868Z",
     "shell.execute_reply": "2024-04-30T09:01:32.011627Z"
    }
   },
   "outputs": [],
   "source": [
    "hooke_rot = copy(hooke).rotate_by(0.5)\n",
    "assert allclose(hooke.hooke,hooke_rot.hooke)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Selling's decomposition of an isotropic Hooke tensor is very structured and predictable. It involves offsets $m_r$ which are independent of the parameters $(\\lambda,\\mu)$, and weights $\\rho_r(\\lambda,\\mu)$ depending linearly on the Lame parameters. In addition, several of these coefficients vanish."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.013187Z",
     "iopub.status.busy": "2024-04-30T09:01:32.013108Z",
     "iopub.status.idle": "2024-04-30T09:01:32.015038Z",
     "shell.execute_reply": "2024-04-30T09:01:32.014803Z"
    }
   },
   "outputs": [],
   "source": [
    "coefs,moffsets = hooke.Selling()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.016339Z",
     "iopub.status.busy": "2024-04-30T09:01:32.016240Z",
     "iopub.status.idle": "2024-04-30T09:01:32.018339Z",
     "shell.execute_reply": "2024-04-30T09:01:32.018104Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([-0., -0.,  2.,  4.,  4.,  1.])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "coefs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.019637Z",
     "iopub.status.busy": "2024-04-30T09:01:32.019558Z",
     "iopub.status.idle": "2024-04-30T09:01:32.021562Z",
     "shell.execute_reply": "2024-04-30T09:01:32.021311Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "coef:2.0, moffset:\n",
      "[[0 1]\n",
      " [1 0]]\n",
      "\n",
      "coef:4.0, moffset:\n",
      "[[ 0  0]\n",
      " [ 0 -1]]\n",
      "\n",
      "coef:4.0, moffset:\n",
      "[[-1  0]\n",
      " [ 0  0]]\n",
      "\n",
      "coef:1.0, moffset:\n",
      "[[1 0]\n",
      " [0 1]]\n",
      "\n"
     ]
    }
   ],
   "source": [
    "for c,o in zip(coefs,np.moveaxis(moffsets,-1,0)): \n",
    "    if c!=0.:\n",
    "        print(f\"coef:{c}, moffset:\\n{o}\\n\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 VTI tensor\n",
    "\n",
    "A VTI Hooke tensor, or slightly more generally a tetragonal Hooke tensor, has a block diagonal form.\n",
    "One of the blocks has size $d \\times d$, and is fully generic in dimension $d=2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.022871Z",
     "iopub.status.busy": "2024-04-30T09:01:32.022793Z",
     "iopub.status.idle": "2024-04-30T09:01:32.024612Z",
     "shell.execute_reply": "2024-04-30T09:01:32.024348Z"
    }
   },
   "outputs": [],
   "source": [
    "hooke = Hooke.from_tetragonal(1.,2.,3.,4.,5.,6.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.025938Z",
     "iopub.status.busy": "2024-04-30T09:01:32.025864Z",
     "iopub.status.idle": "2024-04-30T09:01:32.028056Z",
     "shell.execute_reply": "2024-04-30T09:01:32.027819Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1., 3., 0.],\n",
       "       [3., 4., 0.],\n",
       "       [0., 0., 5.]])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hooke.extract_xz().hooke"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.029366Z",
     "iopub.status.busy": "2024-04-30T09:01:32.029290Z",
     "iopub.status.idle": "2024-04-30T09:01:32.031366Z",
     "shell.execute_reply": "2024-04-30T09:01:32.031111Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1., 2., 3., 0., 0., 0.],\n",
       "       [2., 1., 3., 0., 0., 0.],\n",
       "       [3., 3., 4., 0., 0., 0.],\n",
       "       [0., 0., 0., 5., 0., 0.],\n",
       "       [0., 0., 0., 0., 5., 0.],\n",
       "       [0., 0., 0., 0., 0., 6.]])"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "hooke.hooke"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Finite difference energy\n",
    "\n",
    "We approximate the linear elastic energy using a second order accurate finite differences scheme, which exploits our tensor decomposition. The scheme is based on the identity\n",
    "$$\n",
    "    c(\\epsilon,\\epsilon) = \\sum_r \\rho_r \\mathrm{Tr}(m_r \\nabla v)^2\n",
    "$$\n",
    "where $(\\rho_r,m_r)$ is Selling's decomposition of $c$. We could replace $\\sigma$ with $\\nabla v$ thanks to the symmetry of $m_r$.\n",
    "\n",
    "Then we use the finite difference approximations\n",
    "$$\n",
    "    \\mathrm{Tr}(m \\nabla v) = \\sum_{0 \\leq i < d} \\frac{v_i(x+h s_i m[i] )-v_i(x)}{h s_i} + O(h),\n",
    "$$\n",
    "where $s_i\\in\\{-1,1\\}$, $0\\leq i < d$ are arbitrary signs, and $m[i]$ denotes the $i$-th column of $m$ (which is a symmetric matrix). Squaring this expression, averaging over all possible sign choices, and summing with weights $\\rho_r$, we obtain a second order consistent approximation of the local linear elastic energy.\n",
    "$$\n",
    "    c(\\epsilon,\\epsilon) = \\sum_{0 \\leq r \\leq D (D+1)/2} \\frac{\\rho_r}{2^d} \\sum_{s \\in \\{-1,1\\}^d}  \n",
    "    \\Big(\\sum_{0 \\leq i < d} \\frac{v_i(x+h s_i m_r[i] )-v_i(x)}{h s_i}\\Big)^2.\n",
    "$$\n",
    "\n",
    "**Remark**\n",
    "If the coordinates of $m_r[i]$ are not co-prime, for some $0 \\leq r < D (D+1)/2$ and $0 \\leq i < d$, then one can improve the scheme by taking advantage of this fact in the finite differences. Namely replace $u(x+khe)-u(x)$ with the more local finite difference $k(u(x+he)-u(x))$. \n",
    "\n",
    "**Scheme order.**\n",
    "For symmetry reasons, using first order finite differences yields a second order estimation of the elastic energy. \n",
    "Likewise, using third order finite differences would yield a fourth order estimation of the elastic enery, but this does not seem practical."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "editable": true,
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.032729Z",
     "iopub.status.busy": "2024-04-30T09:01:32.032626Z",
     "iopub.status.idle": "2024-04-30T09:01:32.035287Z",
     "shell.execute_reply": "2024-04-30T09:01:32.035055Z"
    },
    "slideshow": {
     "slide_type": ""
    },
    "tags": [
     "ExportCode"
    ]
   },
   "outputs": [],
   "source": [
    "def ElasticEnergy(v,hooke,dom):\n",
    "    \"\"\"\n",
    "    Finite differences approximation of c(ϵ,ϵ), where c is the Hooke tensor and ϵ the strain tensor,\n",
    "    which is twice the (linearized) elastic energy density.\n",
    "    \"\"\"\n",
    "    d=len(v)\n",
    "    coefs,moffsets = hooke.Selling()\n",
    "    dvp = tuple( dom.DiffUpwind(v[i], moffsets[i]) for i in range(d))\n",
    "    dvm = tuple(-dom.DiffUpwind(v[i],-moffsets[i]) for i in range(d))\n",
    "    \n",
    "    # Consistent approximations of Tr(moffset*grad(v))\n",
    "    dv = ad.array([sum(s) for s in itertools.product(*zip(dvp,dvm))])\n",
    "    dv2 = (dv**2).sum(axis=0)\n",
    "    \n",
    "    coefs = fd.as_field(coefs,v.shape[1:]) * 2**(-d)\n",
    "    return (coefs*dv2).sum(axis=0) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 CFL condition\n",
    "\n",
    "The above discretization of the elastic energy is intended to be used in numerical simulations of the elastic wave equation, which involves a CFL condition. See the [discussion in the one dimension case](Time1D_Div.ipynb). The key to deriving this condition is to upper bound the quadratic form $c(\\epsilon,\\epsilon)$ in terms of a quadratic form with only diagonal coefficients.\n",
    "\n",
    "Recall the Cauchy-Schwartz inequality, or arithmetico-geometric inequality\n",
    "$$\n",
    "    \\Big(\\sum_{1 \\leq k \\leq K} a_k\\Big)^2 \\leq k \\sum_{1 \\leq k \\leq K} a_k^2,\n",
    "$$\n",
    "and observe that $c(\\epsilon,\\epsilon)$ is a sum of squares, each comprising $K = 2 d$ terms. Denoting by $A$ the sparse matrix associated with the integrated elastic energy, and by $D$ its diagonal, one therefore has $A \\preceq K D$.\n",
    "\n",
    "In the case where the hooke tensors are constant, one may be slightly more explicit, observing that \n",
    "$$\n",
    "    \\mathrm{Tr}(c) = \\sum_r \\rho_r \\mathrm{Tr}(m_r^2) \\geq \\sum_r \\rho_r,\n",
    "$$\n",
    "since the matrices $m_r$ have integer coefficients (and are non-zero). By $\\mathrm{Tr}(c)$ we denote the trace of the Mandel form of $c$, which is also its trace as a quadratic form on the space of $2\\times 2$ symmetric matrices equipped with the Frobenius inner product.\n",
    "From this point, one gets $A \\leq L \\mathrm{Tr}(c) \\mathrm{Id}$, where $L = 2 K d = (2d)^2$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.036613Z",
     "iopub.status.busy": "2024-04-30T09:01:32.036536Z",
     "iopub.status.idle": "2024-04-30T09:01:32.038878Z",
     "shell.execute_reply": "2024-04-30T09:01:32.038660Z"
    }
   },
   "outputs": [],
   "source": [
    "hooke = mica.extract_xz().rotate_by(0.5)\n",
    "coefs,moffsets = hooke.Selling()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.040210Z",
     "iopub.status.busy": "2024-04-30T09:01:32.040134Z",
     "iopub.status.idle": "2024-04-30T09:01:32.042094Z",
     "shell.execute_reply": "2024-04-30T09:01:32.041863Z"
    }
   },
   "outputs": [],
   "source": [
    "tr_decomp = np.sum(coefs*lp.trace(lp.dot_AA(moffsets,moffsets)))\n",
    "tr_Mandel = lp.trace(hooke.to_Mandel())\n",
    "assert allclose(tr_decomp,tr_Mandel)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.2 Comparison with automatic differentiation\n",
    "\n",
    "For comparison, we also evaluate the elastic energy using automatic differentiation for the derivatives of $v$, instead of finite differences, which yields an exact expression. This is only possible when $v$ is provided as a differentiable function, rather than a table."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.043519Z",
     "iopub.status.busy": "2024-04-30T09:01:32.043436Z",
     "iopub.status.idle": "2024-04-30T09:01:32.045541Z",
     "shell.execute_reply": "2024-04-30T09:01:32.045297Z"
    }
   },
   "outputs": [],
   "source": [
    "def ElasticEnergy_ad(v,hooke,X):\n",
    "    \"\"\"\n",
    "    Returns c(ϵ,ϵ), where c is the Hooke tensor, and ϵ is the stress tensor.\n",
    "    The displacement field must be provided as a function, compatible with automatic differentiation.\n",
    "    \"\"\"\n",
    "    # Differentiate the displacement field\n",
    "    X_ad = ad.Dense.identity(constant=X,shape_free=(2,))\n",
    "    grad = v(X_ad).gradient()\n",
    "    \n",
    "    # Compute the stress tensor, and the inner product associated with the Hooke tensor.\n",
    "    ϵ = 0.5*(grad+lp.transpose(grad))\n",
    "    return hooke.dot_AA(ϵ) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us observe the convergence of the finite element energy toward the exact energy in a smooth periodic setting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.046915Z",
     "iopub.status.busy": "2024-04-30T09:01:32.046834Z",
     "iopub.status.idle": "2024-04-30T09:01:32.049008Z",
     "shell.execute_reply": "2024-04-30T09:01:32.048786Z"
    }
   },
   "outputs": [],
   "source": [
    "def v(X):\n",
    "    x0,x1 = X*(2.*np.pi)\n",
    "    return ad.array((np.cos(x0) - 2.*np.sin(x1),np.cos(x0+2*x1)))\n",
    "\n",
    "def hooke(X):\n",
    "    x0,x1 = X*(2.*np.pi)\n",
    "    angle = 0.3*np.sin(x0)+0.5*np.cos(x1)    \n",
    "    return mica.extract_xz().rotate_by(angle)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.050318Z",
     "iopub.status.busy": "2024-04-30T09:01:32.050238Z",
     "iopub.status.idle": "2024-04-30T09:01:32.052300Z",
     "shell.execute_reply": "2024-04-30T09:01:32.052076Z"
    }
   },
   "outputs": [],
   "source": [
    "n=20\n",
    "aX,h = xp.linspace(0,1,n,endpoint=False,retstep=True)\n",
    "X=ad.array(np.meshgrid(aX,aX,indexing='ij'))\n",
    "dom = Domain.MockDirichlet(X.shape,h,padding=None) #Periodic domain (wrap instead of pad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.053620Z",
     "iopub.status.busy": "2024-04-30T09:01:32.053542Z",
     "iopub.status.idle": "2024-04-30T09:01:32.058081Z",
     "shell.execute_reply": "2024-04-30T09:01:32.057857Z"
    }
   },
   "outputs": [],
   "source": [
    "energy_density_fd = ElasticEnergy(   v(X),hooke(X),dom) # Uses samples of v\n",
    "energy_density_ad = ElasticEnergy_ad(v,   hooke(X),X)   # Uses function v"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The finite differences scheme is second order accurate. A fourth order scheme could easily be achieved, in the domain interior, by using higher order finite differences. The treatment of non-periodic boundary conditions in another story altogether, that will be addressed later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.059422Z",
     "iopub.status.busy": "2024-04-30T09:01:32.059332Z",
     "iopub.status.idle": "2024-04-30T09:01:32.061546Z",
     "shell.execute_reply": "2024-04-30T09:01:32.061299Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.7490743389566955"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "norm_infinity(energy_density_fd-energy_density_ad) / norm_average(energy_density_ad)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In order to obtain the total elastic energy, the density must be integrated over the domain, not forgetting about the $1/2$ factor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.062887Z",
     "iopub.status.busy": "2024-04-30T09:01:32.062802Z",
     "iopub.status.idle": "2024-04-30T09:01:32.064788Z",
     "shell.execute_reply": "2024-04-30T09:01:32.064574Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Relative error: 0.03070905470908115\n"
     ]
    }
   ],
   "source": [
    "energy_ad = 0.5 * energy_density_ad.sum() * h**2\n",
    "energy_fd = 0.5 * energy_density_fd.sum() * h**2\n",
    "print(f\"Relative error: {(energy_ad-energy_fd)/energy_fd}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.3 Structure of the mass matrix\n",
    "\n",
    "We evaluate the finite difference scheme on a second order sparse AD vector, to find the mass matrix of the elastic energy. At each point in the domain, the energy density depends on a number of neighbor values of the strain tensor, referred to as the stencil."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.066083Z",
     "iopub.status.busy": "2024-04-30T09:01:32.066003Z",
     "iopub.status.idle": "2024-04-30T09:01:32.075327Z",
     "shell.execute_reply": "2024-04-30T09:01:32.075069Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stencil cardinality: 384\n"
     ]
    }
   ],
   "source": [
    "v_sp = ad.Sparse2.identity(constant=np.zeros_like(X))\n",
    "energy_density_sp = ElasticEnergy(v_sp,hooke(X),dom)\n",
    "print(f\"Stencil cardinality: {energy_density_sp.size_ad2}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The *simplification* step compresses the sparse matrix by merging entries corresponding to the same coefficient, and removing zero coefficients.\n",
    "Note that we use `atol=0.` to also remove coefficients arising from cancellation effects: the sum of merged entries vanishes. \n",
    "A more careful implementation of `Energy_fd` may allow to bypass these simplification steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.076782Z",
     "iopub.status.busy": "2024-04-30T09:01:32.076698Z",
     "iopub.status.idle": "2024-04-30T09:01:32.094942Z",
     "shell.execute_reply": "2024-04-30T09:01:32.094674Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stencil cardinality: 70\n"
     ]
    }
   ],
   "source": [
    "energy_density_sp.simplify_ad(atol=True,rtol=True)\n",
    "print(f\"Stencil cardinality: {energy_density_sp.size_ad2}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The energy density at a given point in space is given by a sparse quadratic form, with the above number of non-zero entries at most."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.096381Z",
     "iopub.status.busy": "2024-04-30T09:01:32.096297Z",
     "iopub.status.idle": "2024-04-30T09:01:32.098982Z",
     "shell.execute_reply": "2024-04-30T09:01:32.098744Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(array([ 5.07960333e+03, -5.07960333e+03,  2.43309559e+03,  1.06706077e+02,\n",
       "        -1.06706077e+02, -2.43309559e+03,  6.51713229e+04, -6.51713229e+04,\n",
       "         3.27025834e+03, -3.27025834e+03,  3.27501189e+02, -3.27501189e+02,\n",
       "         1.63750594e+02, -1.63750594e+02, -5.07960333e+03, -6.51713229e+04,\n",
       "        -3.27501189e+02,  1.41156855e+05, -3.27501189e+02, -6.51713229e+04,\n",
       "        -5.07960333e+03, -3.27501189e+02,  3.27501189e+02, -1.63750594e+02,\n",
       "         1.63750594e+02, -6.51713229e+04,  6.51713229e+04, -3.27025834e+03,\n",
       "         3.27025834e+03, -5.07960333e+03,  5.07960333e+03, -2.43309559e+03,\n",
       "        -1.06706077e+02,  1.06706077e+02,  2.43309559e+03,  2.43309559e+03,\n",
       "        -2.43309559e+03,  4.86619118e+03, -4.86619118e+03,  1.06706077e+02,\n",
       "         1.63750594e+02, -1.63750594e+02, -1.06706077e+02,  5.40913343e+02,\n",
       "        -5.40913343e+02,  3.27025834e+03, -3.27025834e+03,  2.13139514e+04,\n",
       "        -2.13139514e+04, -4.86619118e+03, -5.40913343e+02, -2.13139514e+04,\n",
       "         5.34421118e+04, -2.13139514e+04, -5.40913343e+02, -4.86619118e+03,\n",
       "        -3.27025834e+03,  3.27025834e+03, -2.13139514e+04,  2.13139514e+04,\n",
       "        -1.06706077e+02, -1.63750594e+02,  1.63750594e+02,  1.06706077e+02,\n",
       "        -5.40913343e+02,  5.40913343e+02, -2.43309559e+03,  2.43309559e+03,\n",
       "        -4.86619118e+03,  4.86619118e+03]),\n",
       " (array([164, 164, 164, 164, 164, 164, 165, 165, 165, 165, 184, 184, 184,\n",
       "         184, 185, 185, 185, 185, 185, 185, 185, 186, 186, 186, 186, 205,\n",
       "         205, 205, 205, 206, 206, 206, 206, 206, 206, 565, 565, 565, 565,\n",
       "         566, 566, 566, 566, 566, 566, 584, 584, 584, 584, 585, 585, 585,\n",
       "         585, 585, 585, 585, 586, 586, 586, 586, 604, 604, 604, 604, 604,\n",
       "         604, 605, 605, 605, 605]),\n",
       "  array([164, 185, 565, 566, 604, 605, 165, 185, 584, 586, 184, 185, 566,\n",
       "         604, 164, 165, 184, 185, 186, 205, 206, 185, 186, 566, 604, 185,\n",
       "         205, 584, 586, 185, 206, 565, 566, 604, 605, 164, 206, 565, 585,\n",
       "         164, 184, 186, 206, 566, 585, 165, 205, 584, 585, 565, 566, 584,\n",
       "         585, 586, 604, 605, 165, 205, 585, 586, 164, 184, 186, 206, 585,\n",
       "         604, 164, 206, 585, 605])))"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "energy_density_sp[9,5].triplets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.100339Z",
     "iopub.status.busy": "2024-04-30T09:01:32.100262Z",
     "iopub.status.idle": "2024-04-30T09:01:32.102440Z",
     "shell.execute_reply": "2024-04-30T09:01:32.102219Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2465.644648535158"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.max(np.min(np.abs(energy_density_sp.coef2),axis=-1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us construct the sparse symmetric matrix associated to the total energy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.103868Z",
     "iopub.status.busy": "2024-04-30T09:01:32.103793Z",
     "iopub.status.idle": "2024-04-30T09:01:32.106549Z",
     "shell.execute_reply": "2024-04-30T09:01:32.106307Z"
    }
   },
   "outputs": [],
   "source": [
    "energy_sp = 0.5 * energy_density_sp.sum() * h**2\n",
    "energy_hess = energy_sp.hessian_operator()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The number of non-zero entries for each line is approximately 20."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.107998Z",
     "iopub.status.busy": "2024-04-30T09:01:32.107918Z",
     "iopub.status.idle": "2024-04-30T09:01:32.109970Z",
     "shell.execute_reply": "2024-04-30T09:01:32.109745Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "20.755"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "energy_hess.nnz / energy_hess.shape[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since the finite difference energy is quadratic, it is exactly reproduced by the quadratic form defined by is half hessian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.111297Z",
     "iopub.status.busy": "2024-04-30T09:01:32.111221Z",
     "iopub.status.idle": "2024-04-30T09:01:32.112950Z",
     "shell.execute_reply": "2024-04-30T09:01:32.112721Z"
    }
   },
   "outputs": [],
   "source": [
    "v_fl = v(X).flatten()\n",
    "energy_fl = 0.5*np.dot(v_fl,energy_hess*v_fl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.114240Z",
     "iopub.status.busy": "2024-04-30T09:01:32.114160Z",
     "iopub.status.idle": "2024-04-30T09:01:32.115984Z",
     "shell.execute_reply": "2024-04-30T09:01:32.115771Z"
    }
   },
   "outputs": [],
   "source": [
    "assert allclose(energy_sp.as_func(v_fl),energy_fd)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.117313Z",
     "iopub.status.busy": "2024-04-30T09:01:32.117236Z",
     "iopub.status.idle": "2024-04-30T09:01:32.118705Z",
     "shell.execute_reply": "2024-04-30T09:01:32.118454Z"
    }
   },
   "outputs": [],
   "source": [
    "# Fails on cupy 6., works on cupy 7.4 (bug in cupyx sparse matrices)\n",
    "assert allclose(energy_fl,energy_fd)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.4 Isotropic scheme\n",
    "\n",
    "When the material is isotropic, the numerical scheme simplifies substantially. \n",
    "\n",
    "**Staggered grid discretization.** For isotropic materials, there exists scheme taking advantage of staggered grids, which are expected to be more accurate than the one proposed in this notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.120030Z",
     "iopub.status.busy": "2024-04-30T09:01:32.119951Z",
     "iopub.status.idle": "2024-04-30T09:01:32.142930Z",
     "shell.execute_reply": "2024-04-30T09:01:32.142642Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stencil cardinality: 384 (before simplification)\n",
      "Stencil cardinality: 42 (after simplification)\n",
      "Stencil cardinality: 9.0 (linear operator)\n"
     ]
    }
   ],
   "source": [
    "v_sp = ad.Sparse2.identity(constant=np.zeros_like(X))\n",
    "energy_density_sp = ElasticEnergy(v_sp,Hooke.from_Lame(xp.array(1.),2.),dom)\n",
    "print(f\"Stencil cardinality: {energy_density_sp.size_ad2} (before simplification)\")\n",
    "energy_density_sp.simplify_ad(atol=True,rtol=True)\n",
    "print(f\"Stencil cardinality: {energy_density_sp.size_ad2} (after simplification)\")\n",
    "energy_hess = energy_density_sp.hessian_operator()\n",
    "print(f\"Stencil cardinality: {energy_hess.nnz / energy_hess.shape[0]} (linear operator)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Three dimensions\n",
    "\n",
    "The scheme also works in three dimensions."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Decomposition of a single tensor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.144464Z",
     "iopub.status.busy": "2024-04-30T09:01:32.144375Z",
     "iopub.status.idle": "2024-04-30T09:01:32.146422Z",
     "shell.execute_reply": "2024-04-30T09:01:32.146215Z"
    }
   },
   "outputs": [],
   "source": [
    "hooke = mica.rotate_by(0.5, (1,2,3)) # Rotate around some arbitrary axis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.147720Z",
     "iopub.status.busy": "2024-04-30T09:01:32.147640Z",
     "iopub.status.idle": "2024-04-30T09:01:32.161134Z",
     "shell.execute_reply": "2024-04-30T09:01:32.160783Z"
    }
   },
   "outputs": [],
   "source": [
    "coefs,moffsets = hooke.Selling()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.162786Z",
     "iopub.status.busy": "2024-04-30T09:01:32.162671Z",
     "iopub.status.idle": "2024-04-30T09:01:32.165454Z",
     "shell.execute_reply": "2024-04-30T09:01:32.165232Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Coefficients ρ_i :  [  0.19785622   0.20221908   0.52227781   1.15121531   1.28421131\n",
      "   1.81478768   1.84331979   2.02288754   2.47114991   2.53679669\n",
      "   5.08269759   6.34762182   6.59824431   7.2971531    7.331879\n",
      "   9.57493168  18.88768073  23.40618513  36.7647839   69.00465899\n",
      " 129.73804697]\n",
      "Matrix offsets D_i : \n",
      " [[[ 1  0  0]\n",
      "  [ 0 -1 -1]\n",
      "  [ 0 -1  1]]\n",
      "\n",
      " [[ 1  0  0]\n",
      "  [ 0  1  0]\n",
      "  [ 0  0  0]]\n",
      "\n",
      " [[ 1  1  0]\n",
      "  [ 1  0 -1]\n",
      "  [ 0 -1  0]]\n",
      "\n",
      " [[ 0  1  0]\n",
      "  [ 1  1  0]\n",
      "  [ 0  0 -1]]\n",
      "\n",
      " [[ 2 -1 -1]\n",
      "  [-1  0  0]\n",
      "  [-1  0  1]]\n",
      "\n",
      " [[ 2  0 -1]\n",
      "  [ 0  0  0]\n",
      "  [-1  0  0]]\n",
      "\n",
      " [[ 1 -1 -1]\n",
      "  [-1  0  1]\n",
      "  [-1  1  0]]\n",
      "\n",
      " [[ 1  0 -1]\n",
      "  [ 0  0  0]\n",
      "  [-1  0 -1]]\n",
      "\n",
      " [[ 0  0  0]\n",
      "  [ 0  1  1]\n",
      "  [ 0  1  0]]\n",
      "\n",
      " [[ 0  1  0]\n",
      "  [ 1  1  0]\n",
      "  [ 0  0  0]]\n",
      "\n",
      " [[ 0 -1  0]\n",
      "  [-1  0  1]\n",
      "  [ 0  1  0]]\n",
      "\n",
      " [[ 1  0  0]\n",
      "  [ 0  0  0]\n",
      "  [ 0  0  1]]\n",
      "\n",
      " [[ 1  1  0]\n",
      "  [ 1  1  0]\n",
      "  [ 0  0  0]]\n",
      "\n",
      " [[ 2  0 -1]\n",
      "  [ 0  1  0]\n",
      "  [-1  0  0]]\n",
      "\n",
      " [[ 0 -1  0]\n",
      "  [-1  1  1]\n",
      "  [ 0  1  0]]\n",
      "\n",
      " [[ 1  0 -1]\n",
      "  [ 0  0  0]\n",
      "  [-1  0  0]]\n",
      "\n",
      " [[ 1  0  0]\n",
      "  [ 0  1  0]\n",
      "  [ 0  0  1]]\n",
      "\n",
      " [[ 0  0  0]\n",
      "  [ 0  0  0]\n",
      "  [ 0  0  1]]\n",
      "\n",
      " [[ 0  1  0]\n",
      "  [ 1  0  0]\n",
      "  [ 0  0  0]]\n",
      "\n",
      " [[ 1  0  0]\n",
      "  [ 0  0  0]\n",
      "  [ 0  0  0]]\n",
      "\n",
      " [[ 0  0  0]\n",
      "  [ 0  1  0]\n",
      "  [ 0  0  0]]]\n"
     ]
    }
   ],
   "source": [
    "print(\"Coefficients ρ_i : \",coefs)\n",
    "print(\"Matrix offsets D_i : \\n\",np.moveaxis(moffsets,-1,0).astype(int))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.166765Z",
     "iopub.status.busy": "2024-04-30T09:01:32.166663Z",
     "iopub.status.idle": "2024-04-30T09:01:32.168266Z",
     "shell.execute_reply": "2024-04-30T09:01:32.168045Z"
    }
   },
   "outputs": [],
   "source": [
    "if xp is not np: raise ad.DeliberateNotebookError(\n",
    "        \"\"\"End of notebook raises 'Out of memory' on cupy, reasons unknown.\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Structure of the sparse matrix\n",
    "\n",
    "We obtain a matrix fill of 15 for an isotropic material, and approximately 80 for an anisotropic material. (In the latter case, the figure will increase a bit if the material is non-constant.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.169607Z",
     "iopub.status.busy": "2024-04-30T09:01:32.169514Z",
     "iopub.status.idle": "2024-04-30T09:01:32.171677Z",
     "shell.execute_reply": "2024-04-30T09:01:32.171454Z"
    }
   },
   "outputs": [],
   "source": [
    "n=8\n",
    "aX,h = xp.linspace(0,1,n,endpoint=False,retstep=True)\n",
    "X=ad.array(np.meshgrid(aX,aX,aX,indexing='ij'))\n",
    "dom = Domain.MockDirichlet(X.shape,h,padding=None) #Periodic domain (wrap instead of pad)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.172973Z",
     "iopub.status.busy": "2024-04-30T09:01:32.172874Z",
     "iopub.status.idle": "2024-04-30T09:01:32.577563Z",
     "shell.execute_reply": "2024-04-30T09:01:32.577225Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stencil cardinality: 6048 (before simplification)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stencil cardinality: 161 (after simplification)\n",
      "Stencil cardinality: 22.333333333333332 (linear operator)\n"
     ]
    }
   ],
   "source": [
    "v_sp = ad.Sparse2.identity(constant=np.zeros_like(X))\n",
    "energy_density_sp = ElasticEnergy(v_sp,Hooke.from_Lame(xp.array(1.),2.,vdim=3),dom)\n",
    "print(f\"Stencil cardinality: {energy_density_sp.size_ad2} (before simplification)\")\n",
    "energy_density_sp.simplify_ad(atol=True,rtol=True)\n",
    "print(f\"Stencil cardinality: {energy_density_sp.size_ad2} (after simplification)\")\n",
    "energy_hess = energy_density_sp.hessian_operator()\n",
    "print(f\"Stencil cardinality: {energy_hess.nnz / energy_hess.shape[0]} (linear operator)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {
    "execution": {
     "iopub.execute_input": "2024-04-30T09:01:32.579191Z",
     "iopub.status.busy": "2024-04-30T09:01:32.579095Z",
     "iopub.status.idle": "2024-04-30T09:01:33.000038Z",
     "shell.execute_reply": "2024-04-30T09:01:32.999683Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stencil cardinality: 6048 (before simplification)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Stencil cardinality: 363 (after simplification)\n",
      "Stencil cardinality: 79.66666666666667 (linear operator)\n"
     ]
    }
   ],
   "source": [
    "hooke = mica.rotate_by(0.5,(1,2,3))\n",
    "v_sp = ad.Sparse2.identity(constant=np.zeros_like(X))\n",
    "energy_density_sp = ElasticEnergy(v_sp,hooke,dom)\n",
    "print(f\"Stencil cardinality: {energy_density_sp.size_ad2} (before simplification)\")\n",
    "energy_density_sp.simplify_ad(atol=True,rtol=True)\n",
    "print(f\"Stencil cardinality: {energy_density_sp.size_ad2} (after simplification)\")\n",
    "energy_hess = energy_density_sp.hessian_operator()\n",
    "print(f\"Stencil cardinality: {energy_hess.nnz / energy_hess.shape[0]} (linear operator)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "celltoolbar": "Format de la Cellule Texte Brut",
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.8"
  },
  "varInspector": {
   "cols": {
    "lenName": 16,
    "lenType": 16,
    "lenVar": 40
   },
   "kernels_config": {
    "python": {
     "delete_cmd_postfix": "",
     "delete_cmd_prefix": "del ",
     "library": "var_list.py",
     "varRefreshCmd": "print(var_dic_list())"
    },
    "r": {
     "delete_cmd_postfix": ") ",
     "delete_cmd_prefix": "rm(",
     "library": "var_list.r",
     "varRefreshCmd": "cat(var_dic_list()) "
    }
   },
   "types_to_exclude": [
    "module",
    "function",
    "builtin_function_or_method",
    "instance",
    "_Feature"
   ],
   "window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}