Pyglasma3d_Numba

Pyglasma3d Numba and Cuda Development Version #

TODO: a list of all dependencies consolidated and sectioned according to which packages use them. this also includes driver-utilities (nvidia-smi) and the like

old readme contents #

*docstring of default_jit()

default_jit() #

When using cuda (by setting the environment variable
'MY_NUMBA_TARGET' to 'cuda'), numba.jit() will automatically
compile cuda device functions, when 'jitted' functions
are called from cuda kernels.

*project root readme

pyglasma3d_cuda_devel #

This is a port of the pyglasma3d project to Python3 with numba and CUDA acceleration.
The present version is still in development and does lack lots of features when compared to pyglasma3d.

In addition to this ‘master’ branch, there is the ‘cython_module_porting’ branch. It provides a development path, which was used for porting the present modules from the original ‘pyglasma3d’ version to Python3. This path is built upon the possibility to swap out single functions of the original ‘pyglasma3d’ code for ‘new’ code of the ‘pyglasma3d_numba’ version.
The main advantage of this method is, that there is always a fully functional version of the code that produces comparable results. This allows to test the ‘development’ version against the original after each change made. Continuously testing the ‘development’ version provides a very simple ‘test driven development’ method which makes it easy to find and fix bugs and other coding errors the moment they are introduced.

Supplementary to the ‘cython_module_porting’ branch is the tag: ‘initial_code_swapping_tests’. It is a collection of early tests used to develop the code swapping method. The code is slightly outdated, but the simple examples are much more illustrative for explaining the working principles of that method.

Contents #

pyglasma3d_numba/

  • This will eventually become the fully functional port of pyglasma3d with CUDA support

pyglasma3d_original/

  • This is the unmodified original version of pyglasma3d used for testing the results of the ‘pyglasma3d_numba’ version

test_consistency/

  • This is a Python package with a module to compare files containing numerical data. It is used to test the ‘pyglasma3d_numba’ version’s results against ‘pyglasma3d_original’.

Implemented Features #

The present version of ‘pyglasma3d_numba’ only includes the McLerran-Venugopalan model initial conditions and an explicit solver (./pyglasma3d_numba/pyglasma3d_numba_source/solver_leapfrog.py).

Numba acceleration is used for parallel loops and just in time (jit) compilation for further optimizations. It is possible to specify the execution device: cpu only or cpu with gpu acceleration (only CUDA capable GPUs).

*pyglasma3d_numba/readme.md

pyglasma3d_numba #

A 3D glasma simulation implemented in Python using the

numba framework. #

The code is based on the pyglasma3d project which is implemented in Python and uses Cython for acceleration.

Contents #

./examples/: #

  • mv.py:
    A script for Au-Au collisions in the MV-model

  • mv_simple.py:
    A simplified version of the above script (mv.py) that works on a smaller lattice with fewer iteration steps.

  • mv_simple_gauss.py:
    Based on mv_simple.py, but does not produce an output file. Instead only prints information about the gauss constraints and gauss violation.

./pyglasma3d_numba_source/ #

  • __init__.py:
    Package initialization and docstring

  • core.py:
    The main module containing the ‘Simulation’-Class, which defines the working scheme

  • data_io.py:
    A helper module used to read binary files and convert them to plaintext format

  • energy.py:
    A module for computing components of the energy-momentum tensor

  • gauge.py:
    A module used by the unittests defined in ./unittests/tests.py

  • gauss.py:
    A module containing gauss-constraint functions

  • interpolate.py:
    A module for interpolating charge planes and charge/current movement

  • lattice.py:
    The module containing the basic lattice algebra and operations

  • mv.py:
    The module containing the McLerran-Venugopalan initial conditions

  • numba_target.py:
    The module used to control the numba integration

  • solver_leapfrog.py:
    The module implementing the leapfrog Yang-Mills solver

  • unused_code/

    • solver_implicit_common.py:
      A module for commonly used operations on (mainly implicit) solvers
      This module does not get used by the present code

    • su2.py::
      An incomplete module used for testing
      This module does not get used by the present code

./unittests/ #

  • __init__.py:
    Package initialization and docstring

  • TestRunner.py:
    A module providing nicely formatted output for unittests

  • tests.py:
    A module defining Python unittests

How To: #

Dependencies #

This code depends on the following software to work. The version tags refer to the versions this code was tested with. Newer versions of the software might work too:

  • python 3.7.4
  • numpy 1.16.4
  • numba 0.44.1
  • scipy 1.3.0

Running the Simulation #

To try out the ‘mv_simple.py’ simulation setup execute from the main directory of this project (pyglasma3d_numba/):

python3 -m examples.mv_simple

After the simulation run finished, the results will be available in the pyglasma3d_numba/output/ folder. It will contain a binary and a plaintext version of the energy-momentum tensor calculated throughout the simulation.

Implemented Features #

The code is based on the pyglasma3d project but does not provide all the features of that version.

There is only one solver available (./pyglasma3d_numba_source/leapfrog_solver.py) which uses an explicit scheme.

Also there ore only the McLerran-Venugopalan model initial conditions (./pyglasma3d_numba_source/mv.py) available, which work with the explicit scheme.

The implementation of numba follows the example provided by the curraun project. However, CUDA acceleration is not implemented, despite the option being available. All the other options implemented in the curraun code are not working either.

It is possible though, to specify this environment variable to omit any numba integration and use pure Python. Doing this will significantly slow down the performance. Inside the target execution shell use:

export MY_NUMBA_TARGET=python

By default the setup is configured to take full advantage of the numba acceleration features implemented.

Using numba #

The main configuration of how numba is used in this code is located inside ./pyglasma3d_numba_source/numba_target.py. When starting the simulation this module gets imported and checks if the numba framework is available. There will also be a line printed to the execution shell that tells which mode (Python or numba) will be used as the default.

The ‘numba_target’ module also provides a custom decorator default_numba_jit() that basically is a wrapper for the numba.jit() decorator. Its main purpose is to define a default set of parameters to pass to the numba.jit() decorator:

  • “nopython”: True
  • “parallel”: False
  • “nogil”: True
  • “fastmath”: True

default_numba_jit() will pass any arguments supplied on to the numba.jit() decorator, so that any wrong arguments will be handled accordingly. If any of the default parameters set above is supplied, the default value will be overwritten.

Two features of the numba framework are implemented in this code:

  • Every function (if it is possible) will be just in time (jit) compiled and optimized before its first call.
    This effectively causes the code to be compiled during runtime. Furthermore, the functions will be optimized for the specific use case.

  • The module ‘numba_target.py’ defines the custom parallel loop ‘loop_parallel_over()’ which is based on the numba.prange() loop.

‘loop_parallel_over()’ #

The basic functionality is to loop over a supplied ‘kernel_function’ in parallel. Three positional arguments are required:

  • The pointer of the function to loop over: ‘kernel_function’
  • The starting value for the iteration index
  • The range of the loop following the same scheme as the default Python range()

Any other positional arguments after those three will be passed on to the ‘kernel_function’.
Additionally the following keyword arguments are available:

  • force_dev=None
  • returns=None

The first one allows to overwrite the global setting for the execution method (Python or numba) on a per call basis.
The second one takes the value of an integer. It is used to specify if and how many variables are present for reduction in the parallel loop. The return values of the ‘kernel_function’ are summed up, with respect to their position (eg. all first return values are summed up, all second return values are summed up, …). Thus the integer value has to match the amount of values the ‘kernel_function’ returns. If ‘returns’ is set, ‘loop_parallel_over()’ will return a tuple of the reduced variables (aka the sum) with as many elements as set in ‘returns’.
(Reduction is the process of accumulating a sum in the context of a loop.)

Internally, the ‘kernel_function’ will be called with the first argument being the current iteration index of the loop. Following that will be the additional positional arguments supplied to the ‘loop_parallel_over()’ function in their respective order.

Any references or code specific to CUDA acceleration should be ignored. The present version of ‘pyglasma3d_numba’ does not support CUDA.

Example Scripts #

Three example scripts for configuring and running the simulation are provided:

./examples/:

  • mv.py
  • mv_simple.py
  • mv_simple_gauss.py

The second setup is a simplified version of the first. The grid size is reduced to 128 x 32 x 32 instead of 1024*3 x 32 x 32 and the number of evolution steps is reduced to 42 from 2048.

As their last steps these two scripts convert the initially binary output file of the simulation run to plaintext format. This is done using the ./pyglasma3d_numba_source/data_io.py module. Note though, that this module should not be used for gigantic files, because during the conversion, the whole file is loaded into memory.

The conversion of the output files ensures that the ‘test_consistency’ package works as expected.

The third setup is based on the mv_simple.py setup script. However, it does not produce an output file but instead only prints information to the execution shell. The output only contains information about the gauss constraints and gauss violation.

Unittests #

A set of unittest is available in the ./unittests/ folder. These tests are taken from the original pyglasma3d version and are modified to work with the unittest framework in Python3.

Additionally the output of the standard unittests is reformatted and presented in a nice table.

To run the provided test suite execute from the main directory of this project (pyglasma3d_numba/):

python3 -m unittests.tests

The following tests are conducted:

  • test_gauge_link_update:
    test leapfrog.evolve_u() for correct gauge link update
  • test_gauss_constraint:
    test gauss module against analytic results
  • test_index:
    test lattice.get_index() and lattice.get_point() at random positions
  • test_inv_gauge_link_update:
    test leapfrog.evolve_u_inv() for correct inverse gauge link update
  • test_plaquettes:
    test plaquette calculation against numpy methods
  • test_shift:
    test lattice.shift() at boundaries and random positions
  • test_shift2:
    test lattice.shift2() at boundaries and random positions
  • test_su2_mexp_proj:
    test SU(2) group algebra against numpy mexp and proj
  • test_su2_multiplication:
    test SU(2) group algebra against numpy matrix multiplication

These tests include validation of the custom algebra operations included in the simulation code against existing libraries like scipy and numpy. There are also tests for the leapfrog and gauss modules against analytical results.