Exploiting AI Libraries for Solving Discretised Systems of Partial Differential Equations (AI4HFM): Demonstrated for Stokes Flow

  • Heaney, Claire (Imperial College London)
  • Chen, Boyang (Imperial College London)
  • Xiang, Jiansheng (Imperial College London)
  • Latham, John-Paul (Imperial College London)
  • Pain, Christopher (Imperial College London)

Please login to view abstract download link

This paper will demonstrate a new way of implementing the discretised Stokes Equations using AI libraries such as PyTorch and TensorFlow. Although demonstrated here for Stokes flow, these methods could be applied to any discretised partial differential equations (PDEs). Standard discretisations can be written as a number of coefficients (a stencil) multiplying some nodal values or grid values. This type of operation can also be expressed by a convolution layer operating on a set of neurons in a neural network: if the weights of the filter are set according to the discretisation and the neuron input values are the nodal values at the current time level, then the output values will be the nodal values at the future time step. No training is required as the weights are pre-determined. In this work we use a certain convolutional autoencoder, the U-Net, to mimic the behaviour of a multi-grid solver. Writing discretisations as filters will bring to computational physics the advantages that machine learning programmers enjoy, such as interoperability (the code runs on CPUs, GPUs and soon, AI computers with little effort from the programmer); and separation (parts of the code relating to the architecture type are hidden from the user who deals with issues such as the architecture of the neural network). These advantages arise due to the large amount of work that has already been put into the development of PyTorch and TensorFlow. For developers of computational science codes, effort has been much less targeted with typically only a handful of developers working on one code, meaning that much more effort is required to run this software on different architectures. For flows through packed beds (a low Reynolds number flow), it is challenging to calculate the macroscopic proporties of the packed bed based on microscale simulations. With this efficient AI4HFM code, we demonstrate how the microscale simulations can be performed quickly for flow past pellets.