time_compliance_AutodiffvsAnalytical.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "middle-honey",
   "metadata": {
    "id": "middle-honey"
   },
   "source": [
    "### Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "standing-feedback",
   "metadata": {
    "id": "standing-feedback"
   },
   "outputs": [],
   "source": [
    "# # run this first time to clone the directory \n",
    "!git clone https://github.com/UW-ERSL/AuTO.git\n",
    "%cd AuTO/models/\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "functioning-simple",
   "metadata": {
    "id": "functioning-simple"
   },
   "outputs": [],
   "source": [
    "# We begin by importing the necessary libraries\n",
    "import numpy as np\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import jit, grad, random, jacfwd, value_and_grad\n",
    "import functools\n",
    "from functools import partial\n",
    "import time\n",
    "from utilfuncs import Mesher, computeLocalElements, computeFilter\n",
    "from mmaOptimize import optimize\n",
    "from jax import make_jaxpr\n",
    "import matplotlib.pyplot as plt\n",
    "rand_key = random.PRNGKey(0); # reproducibility\n",
    "import matplotlib as mpl\n",
    "mpl.rcParams['figure.dpi'] = 100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "developed-theta",
   "metadata": {
    "id": "developed-theta"
   },
   "outputs": [],
   "source": [
    "class ComplianceMinimizer:\n",
    "  def __init__(self, mesh, bc, material, \\\n",
    "                globalvolCons, AD = True):\n",
    "    self.mesh = mesh\n",
    "    self.material = material\n",
    "    self.bc = bc\n",
    "    M = Mesher()\n",
    "    self.edofMat, self.idx = M.getMeshStructure(mesh)\n",
    "    self.D0 = M.getK0(self.material)\n",
    "    self.globalVolumeConstraint = globalvolCons\n",
    "    self.isADon = AD\n",
    "    if(AD):\n",
    "        self.objectiveHandle = jit(value_and_grad(self.computeObjective))\n",
    "    else:\n",
    "        self.objectiveHandle = self.analyticalObjAndSensHandle\n",
    "    self.consHandle = self.computeConstraints\n",
    "    self.numConstraints = 1\n",
    "    self.projection = {'isOn':False, 'beta':4, 'c0':0.5}\n",
    "  #-----------------------#\n",
    "  @functools.partial(jit, static_argnums = 0) \n",
    "  def analyticalObjAndSensHandle(self, rho):\n",
    "    Y = self.SIMPMaterial(rho)\n",
    "    K = self.assembleK(Y)\n",
    "    u = self.solve(K)\n",
    "    J = self.computeCompliance(K, u)\n",
    "    nelx, nely = self.mesh['nelx'], self.mesh['nely'];\n",
    "    J0elem = self.material['Emin'] + self.material['Emax']*\\\n",
    "        (jnp.dot(u[self.edofMat].reshape(nelx*nely,8),self.D0)\\\n",
    "                  * u[self.edofMat].reshape(nelx*nely,8) ).sum(1);\n",
    "    dJ = -self.material['penal']*\\\n",
    "            jnp.multiply(J0elem, (rho)**(self.material['penal']-1) )\n",
    "    return J, dJ\n",
    "  #-----------------------#\n",
    "  @functools.partial(jit, static_argnums = 0) \n",
    "  def projectionFilter(self, rho):\n",
    "    if(self.projection['isOn']):\n",
    "        v1 = np.tanh(self.projection['c0']*self.projection['beta'])\n",
    "        nm = v1 + jnp.tanh(self.projection['beta']*(rho-self.projection['c0']))\n",
    "        dnm = v1 + jnp.tanh(self.projection['beta']*(1.-self.projection['c0']))\n",
    "        return nm/dnm\n",
    "    else:\n",
    "        return rho\n",
    "  #-----------------------#\n",
    "  @functools.partial(jit, static_argnums = 0)\n",
    "  def SIMPMaterial(self, rho):\n",
    "    Y = self.material['Emin'] + \\\n",
    "          (self.material['Emax']-self.material['Emin'])*\\\n",
    "                          (rho)**self.material['penal']\n",
    "    return Y\n",
    "  #-----------------------#\n",
    "  @functools.partial(jit, static_argnums = 0)\n",
    "  def assembleK(self, Y):\n",
    "    K = jnp.zeros((self.mesh['ndof'], self.mesh['ndof']))\n",
    "    kflat_t = (self.D0.flatten()[np.newaxis]).T \n",
    "    sK = (kflat_t*Y).T.flatten()\n",
    "    K = K.at[(self.idx)].add(sK); #UPDATED\n",
    "    return K\n",
    "  #-----------------------#\n",
    "  @functools.partial(jit, static_argnums = 0)\n",
    "  def solve(self, K): \n",
    "    u_free = jax.scipy.linalg.solve\\\n",
    "            (K[self.bc['free'],:][:,self.bc['free']], \\\n",
    "            self.bc['force'][self.bc['free']], \\\n",
    "              sym_pos = True, check_finite=False);\n",
    "    u = jnp.zeros((self.mesh['ndof']));\n",
    "    u = u.at[self.bc['free']].set(u_free.reshape(-1))\n",
    "\n",
    "    self.u = u;\n",
    "    return u;\n",
    "  #-----------------------#\n",
    "  @functools.partial(jit, static_argnums = 0)\n",
    "  def computeCompliance(self, K, u):\n",
    "    J = jnp.dot(self.bc['force'].T, u)[0]# jnp.dot(u.T, jnp.dot(K,u));\n",
    "    return J; \n",
    "  #-----------------------#\n",
    "  def computeObjective(self, rho):\n",
    "    Y = self.SIMPMaterial(rho);\n",
    "    K = self.assembleK(Y)\n",
    "    u = self.solve(K);\n",
    "    J = self.computeCompliance(K, u);\n",
    "    return J;\n",
    "  #-----------------------#\n",
    "  def computeConstraints(self, rho, epoch): \n",
    "    @jit\n",
    "    def computeGlobalVolumeConstraint(rho):\n",
    "        vc = jnp.mean(rho)/self.globalVolumeConstraint['vf'] - 1.;\n",
    "        return vc;\n",
    "\n",
    "    c, gradc = value_and_grad(computeGlobalVolumeConstraint)\\\n",
    "                                (rho);\n",
    "    c, gradc = c.reshape((1,1)), gradc.reshape((1,-1));\n",
    "    return c, gradc\n",
    "  #-----------------------#\n",
    "  def TO(self, optimizationParams, ft):\n",
    "    rhoHistory = optimize(self.mesh, optimizationParams, ft, \\\n",
    "          self.objectiveHandle, self.consHandle, self.numConstraints);\n",
    "    return rhoHistory;"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "conservative-coordinate",
   "metadata": {
    "id": "conservative-coordinate",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "material = {'Emax':1., 'Emin':1e-2, 'nu':0.3, 'penal':3.}\n",
    "filterRadius = 3.\n",
    "elemSize = np.array([1., 1.])\n",
    "globalVolumeConstraint = {'isOn':True, 'vf':0.5}\n",
    "optimizationParams = {'maxIters':100,'minIters':100,'relTol':0.02}\n",
    "\n",
    "nx = [10, 20, 40, 60]#, 80, 100, 180]\n",
    "ny = [6, 12, 20, 30]#, 40, 50, 90]\n",
    "n = len(nx)\n",
    "numelems = []\n",
    "for num1, num2 in zip(nx, ny):\n",
    "\tnumelems.append(num1 * num2)\n",
    "\n",
    "timeTaken = [[],[]]\n",
    "rhoHistory = [[],[]]\n",
    "for i in range(n):\n",
    "  for idx, ch in enumerate([False, True]):\n",
    "    print('AD :', ch)\n",
    "    nelx, nely = nx[i], ny[i]\n",
    "\n",
    "    mesh = {'nelx':nelx, 'nely':nely, 'elemSize':elemSize,\\\n",
    "            'ndof':2*(nelx+1)*(nely+1), 'numElems':nelx*nely}\n",
    "    H, Hs = computeFilter(mesh, filterRadius)\n",
    "    ft = {'type':1, 'H':H, 'Hs':Hs}\n",
    "    \n",
    "    ndof = 2*(nelx+1)*(nely+1)\n",
    "    force = np.zeros((mesh['ndof'],1))\n",
    "    dofs=np.arange(mesh['ndof'])\n",
    "    fixed = dofs[0:2*(nely+1):1]\n",
    "    free = jnp.setdiff1d(np.arange(mesh['ndof']),fixed)\n",
    "    force[2*(nelx+1)*(nely+1)- (nely+1), 0 ] = -1;\n",
    "    symXAxis = True\n",
    "    symYAxis = False\n",
    "    bc = {'force':force, 'fixed':fixed,'free':free,\\\n",
    "              'symXAxis':symXAxis, 'symYAxis':symYAxis}\n",
    "\n",
    "    Opt = ComplianceMinimizer(mesh, bc, material, \\\n",
    "                    globalVolumeConstraint, ch)\n",
    "    \n",
    "\n",
    "    J, dJ = Opt.objectiveHandle(np.ones((nelx*nely)))\n",
    "    start = time.perf_counter()\n",
    "    rhoHistory[idx] = Opt.TO(optimizationParams, ft)\n",
    "    ttkn = time.perf_counter() - start\n",
    "    print('numelems : {:d} \\t time: {:.2F} \\n'.format(numelems[i], ttkn))\n",
    "    timeTaken[idx].append(ttkn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "criminal-albany",
   "metadata": {
    "id": "criminal-albany"
   },
   "outputs": [],
   "source": [
    "plt.figure();\n",
    "plt.loglog(numelems, timeTaken[0], 'k--', marker = '*', label = 'Analytical')\n",
    "plt.loglog(numelems, timeTaken[1], 'r-', marker = 's', label = 'AutoDiff')\n",
    "plt.xlabel('numElems')\n",
    "plt.ylabel('time (sec)')\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3ceedab-c8e1-4ca7-932e-5c1ecfc0221c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "machine_shape": "hm",
   "name": "time_compliance_AutodiffvsAnalytical.ipynb",
   "provenance": []
  },
  "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.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}