microstructureDesign.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2e5c4fa9",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/UW-ERSL/AuTO/blob/main/microstructureDesign.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dutch-beauty",
   "metadata": {
    "id": "dutch-beauty"
   },
   "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": "iraqi-adapter",
   "metadata": {
    "id": "iraqi-adapter"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax.scipy as jsp\n",
    "from jax import jit, grad, value_and_grad\n",
    "\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib import colors\n",
    "\n",
    "import time\n",
    "from microstrutilfuncs import getMeshStructure, assignMeshDofs,\\\n",
    "getK0, getInitialDensity, computeFilter,\\\n",
    "applySensitivityFilter, getBC, oc\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "norwegian-chocolate",
   "metadata": {
    "id": "norwegian-chocolate"
   },
   "outputs": [],
   "source": [
    "nelx, nely = 30, 30\n",
    "elemSize = np.array([1., 1.])\n",
    "mesh = {'nelx':nelx, 'nely':nely, 'elemSize':elemSize,\\\n",
    "        'ndof':2*(nelx+1)*(nely+1), 'numElems':nelx*nely}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "tribal-ottawa",
   "metadata": {
    "id": "tribal-ottawa"
   },
   "outputs": [],
   "source": [
    "matProp = {'Emax':1., 'Emin':1e-3, 'nu':0.3, 'penal':3.}\n",
    "methodType = 'bulkModulus' # 'bulkModulus' # 'shearModulus' # 'poissonRatio' \n",
    "filterRadius = 1.3\n",
    "vf= 0.25"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "political-voltage",
   "metadata": {
    "id": "political-voltage"
   },
   "outputs": [],
   "source": [
    "H, Hs = computeFilter(mesh, filterRadius)\n",
    "ft = {'type':1, 'H':H, 'Hs':Hs}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ready-marketing",
   "metadata": {
    "id": "ready-marketing"
   },
   "outputs": [],
   "source": [
    "class MicrostructuralOptimization:\n",
    "    def __init__(self, mesh, matProp, methodType, vf):\n",
    "        self.mesh = mesh\n",
    "        self.matProp = matProp\n",
    "        self.methodType = methodType\n",
    "        self.objectiveHandle = jit(value_and_grad(self.computeObjective))\n",
    "        self.edofMat, self.idx = getMeshStructure(mesh)\n",
    "        self.dofs = assignMeshDofs(mesh)\n",
    "        self.K0 = getK0(matProp)\n",
    "        self.ufixed, self.wfixed = getBC(mesh)\n",
    "        self.vf = vf\n",
    "    \n",
    "    #--------------------------#\n",
    "    # Code snippet 4.1\n",
    "    def computeObjective(self, rho):\n",
    "        @jit\n",
    "        def materialModel(rho):\n",
    "            E = self.matProp['Emin'] + \\\n",
    "            (self.matProp['Emax']-self.matProp['Emin'])*(rho+0.01)**self.matProp['penal']\n",
    "            return E\n",
    "        E = materialModel(rho)\n",
    "\n",
    "        #--------------------------#\n",
    "        @jit\n",
    "        def assembleK(Y):\n",
    "            K_asm = jnp.zeros((self.mesh['ndof'], self.mesh['ndof']))\n",
    "            K_elem = (self.K0.flatten()[np.newaxis]).T \n",
    "            K_elem = (K_elem*Y).T.flatten()\n",
    "            K_asm = K_asm.at[(self.idx)].add(K_elem) #UPDATED\n",
    "            return K_asm\n",
    "        K = assembleK(E)\n",
    "\n",
    "        #--------------------------#\n",
    "        @jit\n",
    "        def computeSubMatrices(K):\n",
    "            subk = {}\n",
    "            for k1 in ['interior', 'leftBtm', 'rightUp']:\n",
    "                for k2 in ['corner','interior', 'leftBtm', 'rightUp']:\n",
    "                    subk[k1 + '_' + k2] = K[np.ix_(self.dofs[k1], self.dofs[k2])]\n",
    "\n",
    "            Kr = jnp.vstack((jnp.hstack((subk['interior_interior'], \\\n",
    "                     subk['interior_leftBtm']+subk['interior_rightUp'])), \\\n",
    "                     jnp.hstack((subk['leftBtm_interior']+subk['rightUp_interior'],\\\n",
    "                     subk['leftBtm_leftBtm']+subk['rightUp_leftBtm']+\\\n",
    "                     subk['leftBtm_rightUp']+subk['rightUp_rightUp']))))\n",
    "\n",
    "            F = jnp.matmul(-jnp.vstack((subk['interior_corner'], \\\n",
    "                            subk['leftBtm_corner']+subk['rightUp_corner'])), self.ufixed)+ \\\n",
    "                jnp.matmul(-jnp.vstack((subk['interior_rightUp'], \\\n",
    "                        subk['leftBtm_rightUp']+subk['rightUp_rightUp'])), self.wfixed)\n",
    "            \n",
    "            return Kr, F\n",
    "        Kr, F = computeSubMatrices(K)\n",
    "        \n",
    "        #--------------------------#  \n",
    "        @jit\n",
    "        def performFE(Kr, F):\n",
    "            nx, ny = self.mesh['nelx'], self.mesh['nely']\n",
    "            U = jnp.zeros((2*(nx+1)*(ny+1),3))\n",
    "            U23 = jnp.linalg.solve(Kr, F)\n",
    "            # U = jax.ops.index_update(U, np.hstack((self.dofs['interior'],\\\n",
    "            #                                        self.dofs['leftBtm'])), U23)\n",
    "            # U = jax.ops.index_update(U, self.dofs['corner'], self.ufixed)\n",
    "            # U = jax.ops.index_update(U, self.dofs['rightUp'], \\\n",
    "            #                          self.wfixed+U[self.dofs['leftBtm'],:])\n",
    "            U = U.at[np.hstack((self.dofs['interior'],self.dofs['leftBtm']))].set(U23) #UPDATED\n",
    "            U = U.at[(self.dofs['corner'])].set(self.ufixed) #UPDATED\n",
    "            U = U.at[(self.dofs['rightUp'])].set(self.wfixed+U[self.dofs['leftBtm'],:]) #UPDATED\n",
    "\n",
    "            \n",
    "            return U\n",
    "        U = performFE(Kr, F)\n",
    "        \n",
    "        #--------------------------#\n",
    "        @jit\n",
    "        def homogenizedMatrix(U, rho):\n",
    "            nx, ny = self.mesh['nelx'], self.mesh['nely']\n",
    "            E = {}\n",
    "            for i in range(3):\n",
    "                for j in range(3):\n",
    "                    U1 = U[:,i]\n",
    "                    U2 = U[:,j]\n",
    "                    ijstr = '{:d}_{:d}'.format(i,j)\n",
    "                    uk0u = (jnp.dot(U1[self.edofMat].reshape(nx*ny,8),self.K0) *\\\n",
    "                            U2[self.edofMat].reshape(nx*ny,8)).sum(1)/(nx*ny)\n",
    "\n",
    "                    E0 = ((self.matProp['Emax']-self.matProp['Emin'])*\\\n",
    "                          ((rho+0.01)**self.matProp['penal'])*uk0u).sum()\n",
    "                    E[ijstr] = E0\n",
    "            return E\n",
    "        \n",
    "\n",
    "        E = homogenizedMatrix(U, rho)\n",
    "        if self.methodType == 'bulkModulus':\n",
    "          J = -E['0_0']  -E['0_1']  -E['1_1']  -E['1_0'] #bulkModulus\n",
    "        elif self.methodType == 'shearModulus':\n",
    "          J = -E['2_2'] #shearModulus\n",
    "        elif self.methodType == 'poissonRatio':\n",
    "          J = E['0_1'] - (0.8**self.loop)*(E['0_0'] + E['1_1']) #poissonRatio\n",
    "        return J     \n",
    "\n",
    "    #--------------------------#\n",
    "    def optimize(self, maxIter = 200):\n",
    "        rho = jnp.array(getInitialDensity(self.mesh, self.vf))\n",
    "        change, self.loop = 10., 0\n",
    "        t0 = time.perf_counter()\n",
    "        while(change > 0.01 and self.loop < maxIter):\n",
    "            self.loop += 1\n",
    "            c, dc = self.objectiveHandle(rho)\n",
    "\n",
    "            dv = jnp.ones((self.mesh['nelx']*self.mesh['nely']))\n",
    "            dc, dv = applySensitivityFilter(ft, rho, dc, dv)\n",
    "\n",
    "            rho, change = oc(rho, dc, dv, ft, vf)\n",
    "            rho = jnp.array(rho)\n",
    "            status = 'iter {:d} ;  obj {:.2F} ; vol {:.2F}'.format(\\\n",
    "                    self.loop,  c, jnp.mean(rho))\n",
    "            if(self.loop % 20 == 0):\n",
    "                plt.imshow(-rho.reshape((self.mesh['nelx'],self.mesh['nely'])),\\\n",
    "                           cmap = 'gray')\n",
    "                plt.title(status)\n",
    "                plt.show()\n",
    "\n",
    "            print(status, 'change {:.2F}'.format(change))\n",
    "        print('time taken (sec): ', time.perf_counter() - t0)\n",
    "        plt.imshow(-rho.reshape((self.mesh['nelx'],self.mesh['nely'])),\\\n",
    "                       cmap = 'gray')\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "planned-viking",
   "metadata": {
    "id": "planned-viking"
   },
   "outputs": [],
   "source": [
    "M = MicrostructuralOptimization(mesh, matProp, methodType, vf)\n",
    "M.optimize(200)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "global-rabbit",
   "metadata": {
    "id": "global-rabbit"
   },
   "source": [
    "### "
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "microstructureDesign.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
}