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 thenumba 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 onmv_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.