Created
March 10, 2020 15:25
-
-
Save Teque5/e8ec4cd5aab5d4353db08af8418a0e66 to your computer and use it in GitHub Desktop.
Computer Language Benchmarks Game :: N-Body Python3+Numba
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/usr/bin/env python3 | |
| ''' | |
| The Computer Language Benchmarks Game | |
| https://salsa.debian.org/benchmarksgame-team/benchmarksgame/ | |
| Original rust version by Ilia Schelokov | |
| Pretty decent python3+numba version by Teque5 | |
| **Numba** (`pip install numba`) uses the LLVM compiler as a JIT solution for python3. I based this off the nice **Rust** version from Ilia Schelokov (rust-8). | |
| | | n-body | | |
| |--------------|:-------:| | |
| | rust-8 | 2.3 s | | |
| | py3 | 5722 s | | |
| | py3+numba | 175.6 s | | |
| | py3+numba+fm | 169.9 s | | |
| It run evens faster if you add `cache=True` to the decorator, but it only helps upon running the program a second time. **fm** indicates `fastmath=True`. | |
| This implementation is 74x slower than Rust, but 35x faster than normal python3. | |
| ### Run | |
| ```bash | |
| $ python3 nbody.py 50000000 | |
| -0.169075164 | |
| -0.169059907 | |
| ``` | |
| ''' | |
| import sys | |
| from itertools import combinations | |
| import numba | |
| import numpy as np | |
| SOLAR_MASS = 4 * np.pi * np.pi | |
| BODIES_COUNT = 5 | |
| DAYS_PER_YEAR = 365.24 | |
| COMBOS = tuple(combinations(range(BODIES_COUNT), 2)) | |
| INTERACTIONS = len(COMBOS) | |
| @numba.njit(fastmath=True) | |
| def magnitude(vec: np.ndarray, dt: np.float64) -> np.float64: | |
| '''N-Dimensional Vector Magnitude''' | |
| acc = sum_squares(vec) | |
| return dt / (acc * np.sqrt(acc)) | |
| @numba.njit(fastmath=True) | |
| def sum_squares(vec: np.ndarray) -> np.float64: | |
| '''N-Dimensional Sum of Squares''' | |
| acc = 0 | |
| for val in vec: | |
| acc += val**2 | |
| return acc | |
| @numba.njit(fastmath=True) | |
| def compute_energy(bodies: np.ndarray) -> np.float64: | |
| '''calculate system energy''' | |
| energy = 0 | |
| for cdx in range(BODIES_COUNT): | |
| energy += .5 * bodies[cdx][6] * sum_squares(bodies[cdx][3:6]) | |
| for idx, jdx in COMBOS: | |
| d_pos = bodies[idx][0:3] - bodies[jdx][0:3] | |
| energy -= bodies[idx][6] * bodies[jdx][6] / np.sqrt(sum_squares(d_pos)) | |
| return energy | |
| @numba.njit(fastmath=True) | |
| def offset_momentum(bodies: np.ndarray): | |
| '''Adjust the Sun's velocity to offset system momentum''' | |
| for bdx in range(BODIES_COUNT): | |
| if bdx is 0: continue | |
| bodies[0][3:6] -= bodies[bdx][3:6] * (bodies[bdx][6] / SOLAR_MASS) | |
| @numba.njit(fastmath=True) | |
| def advance(bodies: np.ndarray, dt: np.float64, steps: np.int): | |
| '''Steps the simulation forward by one time-step''' | |
| d_positions = np.empty((INTERACTIONS, 3), dtype=np.float64) | |
| magnitudes = np.empty(INTERACTIONS, dtype=np.float64) | |
| for _ in range(steps): | |
| # Vectors between each pair of bodies | |
| for cdx, (idx, jdx) in enumerate(COMBOS): | |
| d_positions[cdx] = bodies[idx][0:3] - bodies[jdx][0:3] | |
| # Magnitude between each pair of bodies | |
| for cdx in range(INTERACTIONS): | |
| magnitudes[cdx] = magnitude(d_positions[cdx], dt) | |
| # Apply every other body's gravitation to each body's velocity | |
| for cdx, (idx, jdx) in enumerate(COMBOS): | |
| bodies[idx][3:6] -= d_positions[cdx] * (bodies[jdx][6] * magnitudes[cdx]) | |
| bodies[jdx][3:6] += d_positions[cdx] * (bodies[idx][6] * magnitudes[cdx]) | |
| for cdx in range(BODIES_COUNT): | |
| # position += velocity * dt | |
| bodies[cdx][0:3] += bodies[cdx][3:6] * dt | |
| if __name__ == '__main__': | |
| bodies = np.array([ | |
| # Sun [Position, Velocity, Mass] | |
| [0, 0, 0, 0, 0, 0, SOLAR_MASS], | |
| # Jupiter | |
| [4.841_431_442_464_72e0, | |
| -1.160_320_044_027_428_4e0, | |
| -1.036_220_444_711_231_1e-1, | |
| 1.660_076_642_744_037e-3 * DAYS_PER_YEAR, | |
| 7.699_011_184_197_404e-3 * DAYS_PER_YEAR, | |
| -6.904_600_169_720_63e-5 * DAYS_PER_YEAR, | |
| 9.547_919_384_243_266e-4 * SOLAR_MASS], | |
| # Saturn | |
| [8.343_366_718_244_58e0, | |
| 4.124_798_564_124_305e0, | |
| -4.035_234_171_143_214e-1, | |
| -2.767_425_107_268_624e-3 * DAYS_PER_YEAR, | |
| 4.998_528_012_349_172e-3 * DAYS_PER_YEAR, | |
| 2.304_172_975_737_639_3e-5 * DAYS_PER_YEAR, | |
| 2.858_859_806_661_308e-4 * SOLAR_MASS], | |
| # Uranus | |
| [1.289_436_956_213_913_1e1, | |
| -1.511_115_140_169_863_1e1, | |
| -2.233_075_788_926_557_3e-1, | |
| 2.964_601_375_647_616e-3 * DAYS_PER_YEAR, | |
| 2.378_471_739_594_809_5e-3 * DAYS_PER_YEAR, | |
| -2.965_895_685_402_375_6e-5 * DAYS_PER_YEAR, | |
| 4.366_244_043_351_563e-5 * SOLAR_MASS], | |
| # Neptune | |
| [1.537_969_711_485_091_1e1, | |
| -2.591_931_460_998_796_4e1, | |
| 1.792_587_729_503_711_8e-1, | |
| 2.680_677_724_903_893_2e-3 * DAYS_PER_YEAR, | |
| 1.628_241_700_382_423e-3 * DAYS_PER_YEAR, | |
| -9.515_922_545_197_159e-5 * DAYS_PER_YEAR, | |
| 5.151_389_020_466_114_5e-5 * SOLAR_MASS] | |
| ], dtype=np.float64) | |
| ncycles = int(sys.argv[1]) | |
| offset_momentum(bodies) | |
| print('{:.9f}'.format(compute_energy(bodies))) | |
| advance(bodies, 0.01, ncycles) | |
| print('{:.9f}'.format(compute_energy(bodies))) |
Author
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Not that anybody will ever find this, but my implementation here is still much faster than the posted score for python on the n-body benchmark website.
The maintainer refused to accept my submission b/c he is too lazy to install numba - what a waste. 😞