(39e) Performant Python Extensions in HOOMD-Blue Version 3.0

In this talk, we describe how we enable performant Python extensions in the core of HOOMD-blue version 3.0, and showcase examples highlighting these features. HOOMD-blue is a molecular simulation engine for molecular dynamics and hard particle Monte Carlo simulations. Our new major release, version 3.0, increases the extensibility, flexibility, and usability of the package by enabling extensions to HOOMD-blue’s run loop in Python. With this flexibility, users and developers can build plug-ins and packages written in Python for use with the main package. Python plug-ins provide researchers with quicker prototyping, easier compatibility (no linker or compiler issues), and simpler distribution (conda-forge, PyPI) than compiled C++ plug-ins. This extensibility is accomplished, in part, through custom actions. Custom actions are Python classes that can perform simulation protocols such as writing out system data, resizing the simulation box, or tuning a neighbor list’s hyper-parameters for performance. However, pure Python code can be much slower than compiled languages like C++. For this reason, when implementing custom actions, we worked to minimize any performance penalties associated with using the feature. One way we accomplish this is through the implementation of zero-copy access to HOOMD-blue’s main data buffers (e.g. particle positions, bonded particles, etc.) both on the CPU and GPU. This allows packages like NumPy, SciPy, and others to efficiently compute on the CPU and packages like PyTorch, CuPy, and others to efficiently compute on the GPU. Via Python’s extensions system and packages like Numba and Cython, the custom action’s computationally expensive tasks can even be compiled for further performance gains. In many cases, pure Python plug-ins can be made to be sufficiently efficient for research-quality simulations. To further increase extensibility, performance, and facile plug-in creation, we are enabling just-in-time compilation (JIT) for new use cases in HOOMD-blue version 3.0 that will continue to drive and facilitate new research.