thermalCompliance.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e4b43518",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/UW-ERSL/AuTO/blob/main/thermalCompliance.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": "acoustic-assessment",
   "metadata": {
    "id": "acoustic-assessment"
   },
   "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": "macro-destination",
   "metadata": {
    "id": "macro-destination"
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import numpy.matlib\n",
    "\n",
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import jit, grad, random, jacfwd, value_and_grad\n",
    "from functools import partial\n",
    "\n",
    "import time\n",
    "\n",
    "from utilfuncs import ThermalMesher, computeLocalElements, computeFilter\n",
    "from mmaOptimize import optimize\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "rand_key = random.PRNGKey(0);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "wired-lambda",
   "metadata": {
    "id": "wired-lambda"
   },
   "outputs": [],
   "source": [
    "nelx, nely = 40, 30;\n",
    "elemSize = np.array([1., 1.])\n",
    "mesh = {'nelx':nelx, 'nely':nely, 'elemSize':elemSize,\\\n",
    "        'ndof':(nelx+1)*(nely+1), 'numElems':nelx*nely};"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "developing-lyric",
   "metadata": {
    "id": "developing-lyric"
   },
   "outputs": [],
   "source": [
    "material = {'k0':1., 'alpha':1.e-4, 'penal':3}; # alpha expn coeff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "quick-taste",
   "metadata": {
    "id": "quick-taste"
   },
   "outputs": [],
   "source": [
    "filterRadius = 1.5;\n",
    "H, Hs = computeFilter(mesh, filterRadius);\n",
    "ft = {'type':1, 'H':H, 'Hs':Hs};"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ignored-strain",
   "metadata": {
    "id": "ignored-strain"
   },
   "outputs": [],
   "source": [
    "# body load\n",
    "force = 0.01 * np.ones((mesh['ndof'], 1))\n",
    "fixed = int(nely / 2 + 1 - nely / 20);\n",
    "free = np.setdiff1d(np.arange(mesh['ndof']),fixed);\n",
    "bc = {'heat':force, 'fixedTempNodes':fixed, 'freeTempNodes':free};"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "stuck-closure",
   "metadata": {
    "id": "stuck-closure"
   },
   "outputs": [],
   "source": [
    "# setup constraints\n",
    "maxLengthScale = {'isOn':False, 'radius': 6, 'voidVol': 0.05*np.pi*9**2,\\\n",
    "                 'penal':-6.};\n",
    "globalVolumeConstraint = {'isOn':True, 'vf':0.5};"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dangerous-combat",
   "metadata": {
    "id": "dangerous-combat"
   },
   "outputs": [],
   "source": [
    "optimizationParams = {'maxIters':250,'minIters':100,'relTol':0.02};"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "improved-jesus",
   "metadata": {
    "id": "improved-jesus"
   },
   "outputs": [],
   "source": [
    "class ThermalComplianceMinimizer:\n",
    "    def __init__(self, mesh, bc, material, \\\n",
    "                 globalvolCons, maxLengthScaleCons):\n",
    "        self.mesh = mesh;\n",
    "        self.material = material;\n",
    "        self.bc = bc;\n",
    "        M = ThermalMesher();\n",
    "        self.edofMat, self.idx = M.getMeshStructure(mesh);\n",
    "        self.K0 = M.getK0();\n",
    "        self.globalVolumeConstraint = globalvolCons;\n",
    "        self.maxLengthScale = maxLengthScaleCons;\n",
    "        self.objectiveHandle = jit(value_and_grad(self.computeObjective))\n",
    "        self.consHandle = self.computeConstraints;\n",
    "        if(maxLengthScaleCons['isOn']):\n",
    "            self.localElems = computeLocalElements(mesh,\\\n",
    "                               maxLengthScaleCons['radius']);\n",
    "        \n",
    "        self.numConstraints = int(maxLengthScaleCons['isOn'] + \\\n",
    "                     globalvolCons['isOn']);\n",
    "    #-----------------------#\n",
    "    def computeObjective(self, rho):\n",
    "        @jit\n",
    "        def materialModel(rho):\n",
    "            k = self.material['k0']*(rho+0.01)**self.material['penal'];\n",
    "            return k;\n",
    "        E = materialModel(rho);\n",
    "        #-----------------------#\n",
    "        @jit\n",
    "        def assembleK(E):\n",
    "            K = jnp.zeros((self.mesh['ndof'],)*2);\n",
    "            kflat_t = (self.K0.flatten()[np.newaxis]).T \n",
    "            sK = (kflat_t*E).T.flatten();\n",
    "            # K = jax.ops.index_add(K, self.idx, sK); #deprecated\n",
    "            K = K.at[(self.idx)].add(sK); #UPDATED\n",
    "\n",
    "            return K;\n",
    "        K = assembleK(E);\n",
    "        #-----------------------#\n",
    "        @jit\n",
    "        def solve(K): \n",
    "            \n",
    "            Kfree = K[self.bc['freeTempNodes'],:]\\\n",
    "                    [:,self.bc['freeTempNodes']];\n",
    "            ffree = self.bc['heat'][self.bc['freeTempNodes']] \n",
    "    \n",
    "            u_free = jax.scipy.linalg.solve(Kfree, ffree);\n",
    "    \n",
    "            u = jnp.zeros((self.mesh['ndof']));\n",
    "            # u = jax.ops.index_add(u, self.bc['freeTempNodes'],\\\n",
    "            #                       u_free.reshape(-1)); #deprecated\n",
    "            u = u.at[self.bc['freeTempNodes']].set(u_free.reshape(-1)) #UPDATED\n",
    "\n",
    "            \n",
    "            return u;\n",
    "        u = solve(K);\n",
    "        #-----------------------#\n",
    "        @jit\n",
    "        def computeCompliance(K, u):\n",
    "            J = jnp.dot(u.T, jnp.dot(K,u));\n",
    "            return J;\n",
    "        J = computeCompliance(K, u);\n",
    "        return J; \n",
    "    #-----------------------#\n",
    "    def computeConstraints(self, rho, epoch): \n",
    "        \n",
    "        @jit\n",
    "        def computeGlobalVolumeConstraint(rho):\n",
    "            vc = jnp.mean(rho)/self.globalVolumeConstraint['vf'] - 1.;\n",
    "            return vc;\n",
    "        @partial(jit, static_argnums=(1,))\n",
    "        def computeMaxLengthScaleConstraint(rho, epoch):\n",
    "            n =  min(3., 1. + epoch*0.05);\n",
    "            voidVol = jnp.matmul(self.localElems, (1.01-rho)**n);\n",
    "            minVoidVol = jnp.power(\\\n",
    "                       jnp.sum(voidVol**self.maxLengthScale['penal']),\\\n",
    "                             1./self.maxLengthScale['penal']);\n",
    "            consVal = 1.-(minVoidVol/self.maxLengthScale['voidVol']);\n",
    "            return consVal;\n",
    "\n",
    "        if(self.globalVolumeConstraint['isOn']): \n",
    "            c, gradc = value_and_grad(computeGlobalVolumeConstraint)\\\n",
    "                                    (rho);\n",
    "            c, gradc = c.reshape((1,1)), gradc.reshape((1,-1));\n",
    "        \n",
    "        if(self.maxLengthScale['isOn']):\n",
    "            maxls, dmaxls = value_and_grad(computeMaxLengthScaleConstraint)\\\n",
    "                                    (rho, epoch);\n",
    "            c = np.vstack((c, maxls));\n",
    "            gradc = np.vstack((gradc, dmaxls))\n",
    "        return c, gradc\n",
    "    #-----------------------#\n",
    "    def TO(self, optimizationParams, ft):\n",
    "        rho = optimize(self.mesh, optimizationParams, ft, \\\n",
    "             self.objectiveHandle, self.consHandle, self.numConstraints);\n",
    "        return rho\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "outstanding-forge",
   "metadata": {
    "id": "outstanding-forge"
   },
   "outputs": [],
   "source": [
    "Opt = ThermalComplianceMinimizer(mesh, bc, material, \\\n",
    "                globalVolumeConstraint, maxLengthScale);\n",
    "rho = Opt.TO(optimizationParams, ft)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53cce46c-6b49-4383-8331-bd89a0e7d5d0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "thermalCompliance.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
}