pytreegrav

pytreegrav: A fast Python gravity solver - Published in JOSS (2021)

https://github.com/mikegrudic/pytreegrav

Repository

Fast N-body gravitational force and potential in Python

README.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Introduction\n",
    "pytreegrav is a package for computing the gravitational potential and/or field of a set of particles. It includes methods for brute-force direction summation and for the fast, approximate Barnes-Hut treecode method. For the Barnes-Hut method we implement an oct-tree as a numba jitclass to achieve much higher peformance than the equivalent pure Python implementation, without writing a single line of C or Cython."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Walkthrough\n",
    "First let's import the stuff we want and generate some particle positions and masses - these would be your particle data for whatever your problem is."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from pytreegrav import Accel, Potential"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10**5  # number of particles\n",
    "x = np.random.rand(N, 3)  # positions randomly sampled in the unit cube\n",
    "m = np.repeat(1.0 / N, N)  # masses - let the system have unit mass\n",
    "h = np.repeat(0.01, N)  # softening radii - these are optional, assumed 0 if not provided to the frontend functions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can use the ``Accel`` and ``Potential`` functions to compute the gravitational field and potential:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(Accel(x, m, h))\n",
    "print(Potential(x, m, h))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, pytreegrav will try to make the optimal choice between brute-force and tree methods for speed, but we can also force it to use one method or another. Let's try both and compare their runtimes:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import time\n",
    "\n",
    "t = time()\n",
    "# tree gravitational acceleration\n",
    "accel_tree = Accel(x, m, h, method=\"tree\")\n",
    "print(\"Tree accel runtime: %gs\" % (time() - t))\n",
    "t = time()\n",
    "\n",
    "accel_bruteforce = Accel(x, m, h, method=\"bruteforce\")\n",
    "print(\"Brute force accel runtime: %gs\" % (time() - t))\n",
    "t = time()\n",
    "\n",
    "phi_tree = Potential(x, m, h, method=\"tree\")\n",
    "print(\"Tree potential runtime: %gs\" % (time() - t))\n",
    "t = time()\n",
    "\n",
    "phi_bruteforce = Potential(x, m, h, method=\"bruteforce\")\n",
    "print(\"Brute force potential runtime: %gs\" % (time() - t))\n",
    "t = time()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As you can see, the tree-based methods can be much faster than the brute-force methods, especially for particle counts exceeding $10^4$. Here's an example of how much faster the treecode is when run on a Plummer sphere with a variable number of particles, on a single core of an Intel i9 9900k workstation: \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "But there's no free lunch here: the tree methods are approximate. Let's quantify the RMS errors of the stuff we just computed, compared to the exact brute-force solutions:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "acc_error = np.sqrt(np.mean(np.sum((accel_tree - accel_bruteforce) ** 2, axis=1)))  # RMS force error\n",
    "print(\"RMS force error: \", acc_error)\n",
    "phi_error = np.std(phi_tree - phi_bruteforce)\n",
    "print(\"RMS potential error: \", phi_error)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The above errors are typical for default settings: $\\sim 1\\%$ force error and $\\sim 0.1\\%$ potential error. The error in the tree approximation is controlled by the Barnes-Hut opening angle $\\Theta$, set to 0.7 by default. Smaller $\\Theta$ gives higher accuracy, but also runs slower:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thetas = 0.1, 0.2, 0.4, 0.8  # different thetas to try\n",
    "for theta in thetas:\n",
    "    t = time()\n",
    "    accel_tree = Accel(x, m, h, method=\"tree\", theta=theta)\n",
    "    acc_error = np.sqrt(np.mean(np.sum((accel_tree - accel_bruteforce) ** 2, axis=1)))\n",
    "    print(\"theta=%g Runtime: %gs RMS force error: %g\" % (theta, time() - t, acc_error))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Both brute-force and tree-based calculations can be parallelized across all available logical cores via OpenMP, by specifying ``parallel=True``. This can speed things up considerably, with parallel scaling that will vary with your core and particle number:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from time import time\n",
    "\n",
    "t = time()\n",
    "# tree gravitational acceleration\n",
    "accel_tree = Accel(x, m, h, method=\"tree\", parallel=True)\n",
    "print(\"Tree accel runtime in parallel: %gs\" % (time() - t))\n",
    "t = time()\n",
    "\n",
    "accel_bruteforce = Accel(x, m, h, method=\"bruteforce\", parallel=True)\n",
    "print(\"Brute force accel runtime in parallel: %gs\" % (time() - t))\n",
    "t = time()\n",
    "\n",
    "phi_tree = Potential(x, m, h, method=\"tree\", parallel=True)\n",
    "print(\"Tree potential runtime in parallel: %gs\" % (time() - t))\n",
    "t = time()\n",
    "\n",
    "phi_bruteforce = Potential(x, m, h, method=\"bruteforce\", parallel=True)\n",
    "print(\"Brute force potential runtime in parallel: %gs\" % (time() - t))\n",
    "t = time()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# What if I want to evaluate the fields at different points than where the particles are?\n",
    "\n",
    "We got you covered. The ``Target`` methods do exactly this: you specify separate sets of points for the particle positions and the field evaluation, and everything otherwise works exactly the same (including optional parallelization and choice of solver):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytreegrav import AccelTarget, PotentialTarget\n",
    "\n",
    "# generate a separate set of \"target\" positions where we want to know the potential and field\n",
    "N_target = 10**4\n",
    "x_target = np.random.rand(N_target, 3)\n",
    "h_target = np.repeat(\n",
    "    0.01, N_target\n",
    ")  # optional \"target\" softening: this sets a floor on the softening length of all forces/potentials computed\n",
    "\n",
    "accel_tree = AccelTarget(\n",
    "    x_target, x, m, h_target=h_target, h_source=h, method=\"tree\"\n",
    ")  # we provide the points/masses/softenings we generated before as the \"source\" particles\n",
    "accel_bruteforce = AccelTarget(x_target, x, m, h_source=h, method=\"bruteforce\")\n",
    "\n",
    "acc_error = np.sqrt(np.mean(np.sum((accel_tree - accel_bruteforce) ** 2, axis=1)))  # RMS force error\n",
    "print(\"RMS force error: \", acc_error)\n",
    "\n",
    "phi_tree = PotentialTarget(\n",
    "    x_target, x, m, h_target=h_target, h_source=h, method=\"tree\"\n",
    ")  # we provide the points/masses/softenings we generated before as the \"source\" particles\n",
    "phi_bruteforce = PotentialTarget(x_target, x, m, h_target=h_target, h_source=h, method=\"bruteforce\")\n",
    "\n",
    "phi_error = np.std(phi_tree - phi_bruteforce)\n",
    "print(\"RMS potential error: \", phi_error)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}

Owner

JOSS Publication

pytreegrav: A fast Python gravity solver

Published

December 28, 2021

DOI

10.21105/joss.03675

Volume 6, Issue 68, Page 3675

Authors

Michael Y. Grudić
NASA Hubble Fellow, Carnegie Observatories, Department of Physics & Astronomy and CIERA, Northwestern University

Alexander B. Gurvich
Department of Physics & Astronomy and CIERA, Northwestern University

Editor

Dan Foreman-Mackey

Tags

physics gravity simulations

View PDF Review Thread Software Archive

GitHub Events

Committers

Last synced: 7 months ago

Issues and Pull Requests

Last synced: 6 months ago

Packages