A New Option for Scientific Computing: MindSpore SPONGE

August 1, 2022

Foreword - Author: Yang Yi, Associate Researcher of the Shenzhen Bay Laboratory

Molecular dynamics (MD) simulation is a scientific computing method that simulates the microscopic motion of molecules based on the potential energy function described by the molecular force field to calculate the physical and chemical properties of the molecular system.

At present, MD simulation is widely used in chemistry, biology, physics, pharmaceuticals, materials and environments. Early in the 1960s and 1970s, molecular dynamics simulation programs for specific systems have emerged. After decades of development, many general-purpose MD simulation software has been available in the market. Table 1 shows some popular MD simulation software.

1. Inherent Defects of MD Simulation Software

A. The program framework is outdated and has poor flexibility. Before adding a new algorithm to a program, significant modification on the program code must be done. Especially for third-party developers of a program, it is difficult to implement an external algorithm in the program.

B. Most programs are compiled based on CPUs. To run these programs on computing acceleration devices such as GPUs and implement effective acceleration, the programs must be greatly modified. However, the workload is huge, and the algorithms developed by third-party developers may become invalid.

C. Programs are mostly written in C/C++ or even Fortran, which may be incompatible with mainstream AI frameworks that use Python as the front-end language. As a result, some AI-based algorithms developed in recent years cannot be widely used in MD simulation software.

These problems severely restrict the development of MD simulation, especially the incompatibility problem with the AI framework. They are also a technical barrier that hinders the intelligent reconstruction of MD simulation by AI technologies.

2. MD Simulation Software Based on the AI Framework

To fundamentally solve the inherent defects of MD simulation software, MD simulation software must be developed based on the AI framework to implement seamless combination of AI algorithms and MD simulation.

Over the past few decades, IT technologies have developed rapidly, and the scientific computing field has changed accordingly. On the one hand, various new hardware, such as acceleration devices based on GPUs, have been widely used in scientific computing, greatly accelerating scientific computing.

On the other hand, various new algorithms emerge, especially AI-based algorithms represented by AlphaFold. These algorithms demonstrate that AI can achieve many goals that were once impossible or difficult to reach in scientific computing.

Google launched JAX-MD, an MD simulation software based on its JAX framework at the end of 2019. Some universities also developed TorchMD based on the PyTorch framework. In China, Professor Gao Yiqin cooperated with Huawei to develop the next-generation intelligent MD simulation software MindSpore SPONGE based on Huawei all-scenario AI framework MindSpore.

Compared with traditional MD simulation software, MindSpore SPONGE based on MindSpore has the following advantages:

A. With the automatic differentiation function of MindSpore, the code form of the molecular force field can be greatly simplified.

B. Programs compiled based on MindSpore can greatly simplify software adaptation to different hardware devices.

C. MindSpore SPONGE features end-to-end differentiable models.

3. Simplified Molecular Force Field Code Form Based on Automatic Differential

Mathematically, MD simulation calculates the force of each atom using the potential energy function described by the molecular force field to simulate the motion track of the molecule. However, the potential energy function code written in a general programming language cannot be directly used to calculate the atomic force (the negative gradient of the atomic coordinate relative to the potential energy function). Therefore, the traditional MD simulation software not only needs to be used to write the code of the potential energy function, but also to solve the derivative of the potential energy function to the atomic coordinate (the parse expression of the atomic force). Hence, the workload of both theoretical derivation and code writing will increase greatly, and the program will become more complex, which is one of the difficulties in writing traditional MD simulation programs. In addition, the potential energy function needs to be solved manually, which greatly hampers the molecular force field from using a more complex potential energy function in mathematical form. Otherwise, it is difficult to solve the derivative of the function, which also restricts the development of the molecular force field.

The automatic differential function of Huawei MindSpore can directly calculate the derivative of a function. Therefore, to compile a molecular force field based on MindSpore, you only need to compile the code of the potential energy function. The atomic force (the negative gradient of the atomic coordinate to the potential energy function) can be directly calculated using automatic differential, which not only greatly reduces the workload of code writing, but also simplifies the structure of the MD simulation program. As a result, a broader range of mathematical models can be used in the molecular force field.

4. Simplified Software Adaptation to Different Hardware Devices

MD simulation software is a type of scientific computing program with an ultra-large computing amount. It brings about parallel computing or requires using a computing acceleration device such as a GPU for computing.

For a traditional MD simulation program, besides its basic scientific computing functions, the engineering work such as parallelization of the program or porting to GPUs is the most complex and difficult part of the entire software project.

Huawei all-scenario AI framework MindSpore can run on CPUs, GPUs, and Huawei-developed Ascend AI acceleration chips. Programs compiled using MindSpore can be ported to different hardware devices by modifying only one line of code. In addition, Huawei MindSpore features the automatic parallelism capability. Only minor code modification is required to automatically implement parallel computing. Therefore, developing MD simulation programs based on Huawei MindSpore enables scientists to focus on programs' scientific computing functions and leave engineering problems to MindSpore.

5. End-to-End Differentiable MD Simulation

If an AI framework features the high-order automatic differential capability, an end-to-end differentiable MD simulation program can be compiled to implement many functions that cannot be implemented by traditional MD simulation.

Huawei MindSpore features the capability of high-order automatic differential and can automatically solve the high-order derivatives of functions. End-to-end differentiable MD simulation is a new concept proposed by Google when releasing JAX-MD. The whole process of MD simulation from inputting coordinates to calculating power to updating coordinates is differentiable, which is a revolutionary technology for MD simulation.

Traditional MD simulation cannot directly obtain the relationship between the final simulation result and the input coordinates or parameters. Once the simulation result does not meet the expectation, scientific researchers have to continuously adjust the start coordinates or program parameters based on personal experience, and perform simulation repeatedly until the ideal result is obtained.

In end-to-end differentiable MD simulation, the derivative of the output result relative to the input coordinates or parameters can be directly solved. Therefore, scientific researchers can directly optimize the MD simulation process like running common AI optimization algorithms to obtain ideal results, avoiding the process of repeated attempts. This is a meta-optimization process. End-to-end differentiable is considered an important feature of the next-generation MD simulation.

Figure 1 Comparison between deep learning training and molecular dynamics simulation

The procedure on the left shows the neural network parameter training process based on the MindSpore AI framework. The procedure on the right shows the molecular dynamics simulation process based on the MindSpore SPONGE molecular dynamics simulation library. Each block represents a program module, the gray bold arrow indicates the module combination process, and the black arrow indicates the data update process.

6. MD Simulation Software Compiled Based on the AI-like Architecture

Compared with the traditional software architecture, MindSpore SPONGE is simpler, more efficient, and more intelligent. In addition, MindSpore SPONGE can implement many functions that traditional software cannot. You will find that there are many similarities between AI training and MD simulation, which seem to be completely different.

As shown in the figure above, in an AI training process, an optimizer uses a loss function as an optimization target, and calculates the gradient of a neural network parameter using a provided training dataset to update the parameter of a network model.

In MD simulation, the integrator calculates the atomic force based on the potential energy function (the negative gradient of the potential energy function relative to the atomic coordinate) to update the atomic coordinate.

Therefore, MD simulation can be regarded as a special AI training process in some way. The MD simulation program compiled following the AI training process is the MD simulation program of the AI-like framework.

MindSpore SPONGE also features modularization, high flux, and end-to-end differentiable. It is a next-generation intelligent MD simulation software compiled based on this architecture.

7. Summary

MD simulation is an important scientific computing method. After decades of development, traditional MD simulation software has become mature and has a large user base. Therefore, it is difficult to develop brand-new MD simulation software to surpass old software.

However, the next-generation intelligent MD simulation software MindSpore SPONGE has many advantages that traditional software does not have and can implement many functions that traditional software cannot. As a brand new software, MindSpore SPONGE is still in the developing stage.