JaxTOuNN.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/UW-ERSL/JAXTOuNN/blob/main/JaxTOuNN.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "upset-passage",
      "metadata": {
        "id": "upset-passage"
      },
      "source": [
        "# JaxTOuNN\n",
        "\n",
        "In this notebook we seek to implement TOuNN using Jax. The program is end-end differentiable. This includes the FE solver that has previously been outside the computation chain.\n",
        "\n",
        "Aaditya Chandrasekhar, Krishnan Suresh\n",
        "\n",
        "UW Madison\n",
        "\n",
        "achandrasek3@wisc.edu , ksuresh@wisc.edu"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "vTM54jX_ExTN",
      "metadata": {
        "id": "vTM54jX_ExTN"
      },
      "source": [
        "### Imports"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "coordinated-natural",
      "metadata": {
        "id": "coordinated-natural"
      },
      "outputs": [],
      "source": [
        "# We begin by importing the necessary libraries\n",
        "import numpy as np\n",
        "\n",
        "import jax\n",
        "import jax.numpy as jnp\n",
        "from jax import jit, grad, random, jacfwd, value_and_grad\n",
        "import jax.example_libraries.stax as stax\n",
        "import jax.example_libraries.optimizers as optimizers\n",
        "from functools import partial\n",
        "import jax.nn as nn\n",
        "import matplotlib\n",
        "import matplotlib.pyplot as plt\n",
        "from matplotlib import colors\n",
        "\n",
        "rand_key = random.PRNGKey(0) # reproducibility"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "curious-monte",
      "metadata": {
        "id": "curious-monte"
      },
      "source": [
        "### Mesh\n",
        "\n",
        "For illustration and simplicity, we only consider simple rectangular meshes in the implementation."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "revolutionary-metropolitan",
      "metadata": {
        "id": "revolutionary-metropolitan"
      },
      "outputs": [],
      "source": [
        "nelx, nely = 40, 20\n",
        "elemSize = np.array([1., 1.])\n",
        "mesh = {'ndim':2, 'nelx':nelx, 'nely':nely, 'elemSize':elemSize,\\\n",
        "        'ndof':2*(nelx+1)*(nely+1), 'numElems':nelx*nely}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "engaging-sterling",
      "metadata": {
        "id": "engaging-sterling"
      },
      "outputs": [],
      "source": [
        "# Next we populate the domain with points.\n",
        "# These points are processed by the NN\n",
        "def generatePointsInsideDomain(mesh, res = 1):\n",
        "  # args: Mesh is dictionary containing nelx, nely, elemSize...\n",
        "  # res is the number of points per elem\n",
        "  # returns an array of size (numpts X 2)\n",
        "  xy = np.zeros((res*mesh['nelx']*res*mesh['nely'],2))\n",
        "  ctr = 0\n",
        "  for i in range(res*mesh['nelx']):\n",
        "    for j in range(res*mesh['nely']):\n",
        "      xy[ctr,0] = (i + 0.5)/(res*mesh['elemSize'][0])\n",
        "      xy[ctr,1] = (j + 0.5)/(res*mesh['elemSize'][1])\n",
        "      ctr += 1\n",
        "  return xy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "understanding-swimming",
      "metadata": {
        "id": "understanding-swimming"
      },
      "outputs": [],
      "source": [
        "# observe that xyElems is an array from jax.\n",
        "# This makes tracking the variable possible\n",
        "xyElems = jnp.array(generatePointsInsideDomain(mesh))\n",
        "print(xyElems.shape)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "heard-denver",
      "metadata": {
        "id": "heard-denver"
      },
      "outputs": [],
      "source": [
        "# with the mesh sprecified, we can now calc\n",
        "# info wrt connectivity and node numbering\n",
        "def getMeshStructure(mesh):\n",
        "  # returns edofMat: array of size (numElemsX8) with\n",
        "  # the global dof of each elem\n",
        "  # idx: A tuple informing the position for assembly of computed entries\n",
        "  edofMat=np.zeros((mesh['nelx']*mesh['nely'],8),dtype=int)\n",
        "  for elx in range(mesh['nelx']):\n",
        "    for ely in range(mesh['nely']):\n",
        "      el = ely+elx*mesh['nely']\n",
        "      n1=(mesh['nely']+1)*elx+ely\n",
        "      n2=(mesh['nely']+1)*(elx+1)+ely\n",
        "      edofMat[el,:]=np.array([2*n1+2, 2*n1+3, 2*n2+2,\\\n",
        "                      2*n2+3,2*n2, 2*n2+1, 2*n1, 2*n1+1])\n",
        "  iK = tuple(np.kron(edofMat,np.ones((8,1))).flatten().astype(int))\n",
        "  jK = tuple(np.kron(edofMat,np.ones((1,8))).flatten().astype(int))\n",
        "  idx = (iK,jK)\n",
        "  return edofMat, idx\n",
        "#-----------------------#\n",
        "def getBMatrix(mesh, xi, eta):\n",
        "  dx, dy = mesh['elemSize'][0], mesh['elemSize'][1];\n",
        "  B = np.zeros((3,8));\n",
        "  r1 = np.array([(2.*(eta/4. - 1./4.))/dx, -(2.*(eta/4. - 1./4))/dx,\\\n",
        "                  (2.*(eta/4. + 1./4))/dx,\\\n",
        "                  -(2.*(eta/4. + 1./4))/dx]).reshape(-1);\n",
        "  r2 = np.array([(2.*(xi/4. - 1./4))/dy, -(2.*(xi/4. + 1./4))/dy,\\\n",
        "                  (2.*(xi/4. + 1./4))/dy, -(2.*(xi/4. - 1./4))/dy])\n",
        "\n",
        "  B = [[r1[0], 0., r1[1], 0., r1[2], 0., r1[3], 0.],\\\n",
        "        [0., r2[0], 0., r2[1], 0., r2[2], 0., r2[3]],\\\n",
        "        [r2[0], r1[0], r2[1], r1[1], r2[2], r1[2], r2[3], r1[3]]];\n",
        "\n",
        "  return jnp.array(B)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "agreed-singer",
      "metadata": {
        "id": "agreed-singer"
      },
      "outputs": [],
      "source": [
        "edofMat, idx = getMeshStructure(mesh)\n",
        "B = getBMatrix(mesh, 0., 0.) # eval B matrix at elem center"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "suspended-baptist",
      "metadata": {
        "id": "suspended-baptist"
      },
      "source": [
        "### Material"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "young-train",
      "metadata": {
        "id": "young-train"
      },
      "outputs": [],
      "source": [
        "#Next we define the relevant material property.\n",
        "# We are concerned only with structural mech\n",
        "# at the moment. penal here refers to the SIMP penalization constant\n",
        "material = {'Emax':1., 'Emin':1e-3, 'nu':0.3, 'penal':1.}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "institutional-messenger",
      "metadata": {
        "id": "institutional-messenger"
      },
      "outputs": [],
      "source": [
        "# with the material defined, we can now calculate the base constitutive matrix\n",
        "def getD0(material):\n",
        "  # the base constitutive matrix assumes unit\n",
        "  #area element with E = 1. and nu prescribed.\n",
        "  # the material is also assumed to be isotropic.\n",
        "  # returns a matrix of size (8X8)\n",
        "  E = 1.\n",
        "  nu = material['nu'];\n",
        "  k = np.array([1/2-nu/6,1/8+nu/8,-1/4-nu/12,-1/8+3*nu/8,\\\n",
        "                  -1/4+nu/12,-1/8-nu/8,nu/6,1/8-3*nu/8])\n",
        "  KE = \\\n",
        "  E/(1-nu**2)*np.array([ [k[0], k[1], k[2], k[3], k[4], k[5], k[6], k[7]],\n",
        "  [k[1], k[0], k[7], k[6], k[5], k[4], k[3], k[2]],\n",
        "  [k[2], k[7], k[0], k[5], k[6], k[3], k[4], k[1]],\n",
        "  [k[3], k[6], k[5], k[0], k[7], k[2], k[1], k[4]],\n",
        "  [k[4], k[5], k[6], k[7], k[0], k[1], k[2], k[3]],\n",
        "  [k[5], k[4], k[3], k[2], k[1], k[0], k[7], k[6]],\n",
        "  [k[6], k[3], k[4], k[1], k[2], k[7], k[0], k[5]],\n",
        "  [k[7], k[2], k[1], k[4], k[3], k[6], k[5], k[0]] ])\n",
        "  return KE"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "_4Ghn_ba3qk4",
      "metadata": {
        "id": "_4Ghn_ba3qk4"
      },
      "outputs": [],
      "source": [
        "  material['D0'] = getD0(material)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ZQbzIQKH3cNg",
      "metadata": {
        "id": "ZQbzIQKH3cNg"
      },
      "source": [
        "### BC"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "m3MiVTVR3lQ-",
      "metadata": {
        "id": "m3MiVTVR3lQ-"
      },
      "outputs": [],
      "source": [
        "# Next we get the BCs for the problem.\n",
        "# We encode the BC for some standard probs in examples.py\n",
        "# eg 1- tipcant , 2-midcant, 3- MBB, 4- Michell, 5-DistributedBeam\n",
        "\n",
        "# bc is a dictionary with keys\n",
        "#['exampleName', 'force', 'fixed', 'symXAxis', 'symYAxis']\n",
        "\n",
        "example = 2\n",
        "nelx, nely = mesh['nelx'], mesh['nely']\n",
        "if(example == 1): # tip cantilever\n",
        "  ndof = 2*(nelx+1)*(nely+1)\n",
        "  force = np.zeros((ndof,1))\n",
        "  dofs=np.arange(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)-2*nely+1, 0 ] = -1\n",
        "  symXAxis = False\n",
        "  symYAxis = False\n",
        "elif(example == 2): # mid cantilever\n",
        "  ndof = 2*(nelx+1)*(nely+1)\n",
        "  force = np.zeros((ndof,1))\n",
        "  dofs=np.arange(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",
        "  # force[int(2*nelx*(nely+1)/2)+1 ,0]= -10; # TEST- mid load\n",
        "  symXAxis = True\n",
        "  symYAxis = False\n",
        "bc = {'force':force, 'fixed':fixed,'free':free,\\\n",
        "      'symXAxis':symXAxis, 'symYAxis':symYAxis}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "independent-voluntary",
      "metadata": {
        "id": "independent-voluntary"
      },
      "source": [
        "### Symmetry\n",
        "\n",
        "The resulting structure might be symmetric about an axis. However, owing to the nonlinearity of the NN this may not be enforced implicitly. We therefore explicitly enforce symmetry by transforming the coordinates\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "further-allowance",
      "metadata": {
        "id": "further-allowance"
      },
      "outputs": [],
      "source": [
        "def imposeSymmetry(xy, mesh, bc):\n",
        "  if(bc['symYAxis']):\n",
        "    xv = xy[:,0].at[:].set(0.5*mesh['nelx']\\\n",
        "                          + jnp.abs( xy[:,0] - 0.5*mesh['nelx']) )\n",
        "  else:\n",
        "    xv = xy[:,0]\n",
        "  if(bc['symXAxis']):\n",
        "    yv = xy[:,1].at[:].set(0.5*mesh['nely']\\\n",
        "                          + jnp.abs( xy[:,1] - 0.5*mesh['nely']) )\n",
        "  else:\n",
        "    yv = xy[:,1]\n",
        "  xy = jnp.stack((xv,yv)).T\n",
        "  return xy"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "superior-scheme",
      "metadata": {
        "id": "superior-scheme"
      },
      "outputs": [],
      "source": [
        "# transform the coordinates of the elem centres based on if symm is enforced by the prob.\n",
        "xyElems = imposeSymmetry(xyElems, mesh, bc)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "pregnant-killing",
      "metadata": {
        "id": "pregnant-killing"
      },
      "source": [
        "### Neural Network"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "becoming-venture",
      "metadata": {
        "id": "becoming-venture"
      },
      "outputs": [],
      "source": [
        "# Let us now define the actual NN. We consider a fully connected NN\n",
        "# with LeakyRelu as the activation and a sigmoid in the output layer\n",
        "def elementwise(fun, **fun_kwargs):\n",
        "    \"\"\"Layer that applies a scalar function elementwise on its inputs.\"\"\"\n",
        "    init_fun = lambda rng, input_shape: (input_shape, ())\n",
        "    apply_fun = lambda params, inputs, **kwargs: fun(inputs, **fun_kwargs)\n",
        "    return init_fun, apply_fun\n",
        "Swish = elementwise(nn.swish)\n",
        "\n",
        "def makeNetwork(nnspec):\n",
        "  # JAX network definition\n",
        "  layers = []\n",
        "  for i in range(nnspec['numLayers']-1):\n",
        "    layers.append(stax.Dense(nnspec['numNeuronsPerLayer']))\n",
        "    layers.append(Swish)#(stax.LeakyRelu)\n",
        "  layers.append(stax.Dense(nnspec['outputDim']))\n",
        "  layers.append(stax.Sigmoid)\n",
        "  return stax.serial(*layers)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "familiar-differential",
      "metadata": {
        "id": "familiar-differential"
      },
      "outputs": [],
      "source": [
        "nnspec = {'outputDim':1, 'numNeuronsPerLayer':20,  'numLayers':2}"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "zwqISfoKDjE0",
      "metadata": {
        "id": "zwqISfoKDjE0"
      },
      "source": [
        "### FE Solver\n",
        "\n",
        "We now turn our attention to defining functions that are needed for solving the system. We use jit to speed up the computation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "g5g-nYPZDjFH",
      "metadata": {
        "id": "g5g-nYPZDjFH"
      },
      "outputs": [],
      "source": [
        "class FESolver:\n",
        "  def __init__(self, mesh, material, bc):\n",
        "    self.mesh = mesh\n",
        "    self.bc = bc\n",
        "    self.material = material\n",
        "    self.objectiveHandle = jit(self.objective)\n",
        "  #-----------------------#\n",
        "  def objective(self, Y):\n",
        "    @jit\n",
        "    def assembleK(Y):\n",
        "      K = jnp.zeros((self.mesh['ndof'], self.mesh['ndof']))\n",
        "      kflat_t = (self.material['D0'].flatten()[np.newaxis]).T\n",
        "      sK = (kflat_t*Y).T.flatten()\n",
        "      K = K.at[idx].add(sK)\n",
        "      return K\n",
        "    #-----------------------#\n",
        "    @jit\n",
        "    def solve(K):\n",
        "      # eliminate fixed dofs for solving sys of eqns\n",
        "      u_free = jax.scipy.linalg.solve(K[self.bc['free'],:][:,self.bc['free']], \\\n",
        "              self.bc['force'][self.bc['free']], assume_a = 'pos', check_finite=False);\n",
        "      u = jnp.zeros((self.mesh['ndof']))\n",
        "      u = u.at[self.bc['free']].add(u_free.reshape(-1)) # homog bc wherev fixed\n",
        "      return u\n",
        "    #-----------------------#\n",
        "    @jit\n",
        "    def computeCompliance(K, u):\n",
        "      J = jnp.dot(self.bc['force'].reshape(-1).T, u)\n",
        "      return J\n",
        "    #-----------------------#\n",
        "    K = assembleK(Y)\n",
        "    u = solve(K)\n",
        "    J = computeCompliance(K, u)\n",
        "    return J"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "AleBM0gVQxVY",
      "metadata": {
        "id": "AleBM0gVQxVY"
      },
      "source": [
        "# Opt"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "grand-bearing",
      "metadata": {
        "id": "grand-bearing"
      },
      "source": [
        "### Projections\n",
        "\n",
        "Input and output projections help us define among many geometric, manufacturing constraints."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "legitimate-departure",
      "metadata": {
        "id": "legitimate-departure"
      },
      "outputs": [],
      "source": [
        "#-------FOURIER LENGTH SCALE-----------#\n",
        "def computeFourierMap(mesh, fourierMap):\n",
        "  # compute the map\n",
        "  coordnMapSize = (mesh['ndim'], fourierMap['numTerms']);\n",
        "  freqSign = np.random.choice([-1.,1.], coordnMapSize)\n",
        "  stdUniform = np.random.uniform(0.,1., coordnMapSize)\n",
        "  wmin = 1./(2*fourierMap['maxRadius']*mesh['elemSize'][0])\n",
        "  wmax = 1./(2*fourierMap['minRadius']*mesh['elemSize'][0]) # w~1/R\n",
        "  wu = wmin +  (wmax - wmin)*stdUniform\n",
        "  coordnMap = np.einsum('ij,ij->ij', freqSign, wu)\n",
        "  return coordnMap\n",
        "#-----------------#\n",
        "def applyFourierMap(xy, fourierMap):\n",
        "  if(fourierMap['isOn']):\n",
        "    c = jnp.cos(2*np.pi*jnp.einsum('ij,jk->ik', xyElems, fourierMap['map']))\n",
        "    s = jnp.sin(2*np.pi*jnp.einsum('ij,jk->ik', xyElems, fourierMap['map']))\n",
        "    xy = jnp.concatenate((c, s), axis = 1)\n",
        "  return xy"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "figured-fleet",
      "metadata": {
        "id": "figured-fleet"
      },
      "source": [
        "### Optimization\n",
        "Finally, we are now ready to express the optimization problem"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "planned-sacramento",
      "metadata": {
        "id": "planned-sacramento"
      },
      "outputs": [],
      "source": [
        "# Optimization params\n",
        "lossMethod = {'type':'penalty', 'alpha0':0.05, 'delAlpha':0.05}\n",
        "#lossMethod = {'type':'logBarrier', 't0':3, 'mu':1.1};\n",
        "\n",
        "fourierMap = {'isOn': True, 'minRadius':4., \\\n",
        "              'maxRadius':80., 'numTerms':  200}\n",
        "\n",
        "fourierMap['map'] = computeFourierMap(mesh, fourierMap)\n",
        "\n",
        "\n",
        "optimizationParams = {'maxEpochs':450, 'learningRate':0.01, 'desiredVolumeFraction':0.5,\\\n",
        "                     'lossMethod':lossMethod}"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "contained-coffee",
      "metadata": {
        "id": "contained-coffee"
      },
      "outputs": [],
      "source": [
        "def optimizeDesign(xy, optParams, mesh, material, bc, fourierMap):\n",
        "  FE = FESolver(mesh, material, bc)\n",
        "  # input projection\n",
        "  if(fourierMap['isOn']):\n",
        "   xy = applyFourierMap(xy, fourierMap)\n",
        "  # make the NN\n",
        "  init_fn, applyNN = makeNetwork(nnspec);\n",
        "  fwdNN = jit(lambda nnwts, x: applyNN(nnwts, x))\n",
        "  _, params = init_fn(rand_key, (-1, xy.shape[1]))\n",
        "  # optimizer\n",
        "  opt_init, opt_update, get_params = optimizers.adam(optParams['learningRate'])\n",
        "  opt_state = opt_init(params)\n",
        "  opt_update = jit(opt_update)\n",
        "\n",
        "  # fwd once to get J0-scaling param\n",
        "  density0  = fwdNN(get_params(opt_state), xy)\n",
        "  J0 = FE.objectiveHandle(density0.reshape(-1))\n",
        "\n",
        "  def getYoungsModulus(density):\n",
        "    material['penal'] = min(8., 1. + epoch*0.02)\n",
        "    Y = material['Emin'] + \\\n",
        "          (material['Emax']-material['Emin'])*(density+0.001)**material['penal']\n",
        "    return Y\n",
        "  #-----------------------#\n",
        "  # loss function\n",
        "  def computeLoss(nnwts):\n",
        "    density  = 0.01 + fwdNN(nnwts, xy)\n",
        "    Y = getYoungsModulus(density)\n",
        "    volcons = (jnp.mean(density)/optParams['desiredVolumeFraction'])- 1.\n",
        "    J = FE.objectiveHandle(Y.reshape(-1))\n",
        "\n",
        "    if(optParams['lossMethod']['type'] == 'penalty'):\n",
        "      alpha = optParams['lossMethod']['alpha0'] + \\\n",
        "              epoch*optParams['lossMethod']['delAlpha'] # penalty method\n",
        "      loss = J/J0 + alpha*volcons**2;\n",
        "    if(optParams['lossMethod']['type'] == 'logBarrier'):\n",
        "      t = optParams['lossMethod']['t0']* \\\n",
        "                        optParams['lossMethod']['mu']**epoch\n",
        "      if(volcons < (-1/t**2)):\n",
        "        psi = -jnp.log(-volcons)/t\n",
        "      else:\n",
        "        psi = t*volcons - jnp.log(1/t**2)/t + 1/t\n",
        "      loss = J/J0 + psi\n",
        "\n",
        "    return loss;\n",
        "\n",
        "  # optimization loop\n",
        "  for epoch in range(optParams['maxEpochs']):\n",
        "    opt_state = opt_update(epoch, \\\n",
        "                optimizers.clip_grads(jax.grad(computeLoss)(get_params(opt_state)), 1.),\\\n",
        "                opt_state)\n",
        "\n",
        "    if(epoch%10 == 0):\n",
        "      density = fwdNN(get_params(opt_state), xy)\n",
        "      Y = getYoungsModulus(density)\n",
        "      J = FE.objectiveHandle(Y.reshape(-1))\n",
        "      volf= jnp.mean(density)\n",
        "      if(epoch == 10):\n",
        "        J0 = J;\n",
        "      status = 'epoch {:d}, J {:.2E}, vf {:.2F}'.format(epoch, J/J0, volf);\n",
        "      print(status)\n",
        "      if(epoch%30 == 0):\n",
        "        plt.figure();\n",
        "        plt.imshow(-jnp.flipud(density.reshape((nelx, nely)).T),\\\n",
        "                  cmap='gray')\n",
        "        plt.title(status)\n",
        "        plt.show();\n",
        "\n",
        "  return fwdNN, get_params(opt_state)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "Wy8yn5Ylx13o",
      "metadata": {
        "id": "Wy8yn5Ylx13o"
      },
      "source": [
        "# Run"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "fewer-electron",
      "metadata": {
        "id": "fewer-electron"
      },
      "outputs": [],
      "source": [
        "network, nnwts = optimizeDesign(xyElems, optimizationParams, mesh, material, bc, fourierMap)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "beca5a79-04ac-47eb-bf3a-df24d84e0c64",
      "metadata": {
        "id": "beca5a79-04ac-47eb-bf3a-df24d84e0c64"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "JaxTOuNN.ipynb",
      "provenance": [],
      "include_colab_link": true
    },
    "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.9.7"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}