py: CIC gridding

cic.py

#!/usr/bin/env python

"""A faster-than-pure-python+numpy cloud-in-cell gridding implementation.

**THIS IS A WORK IN PROGRESS AND NOT FULLY TESTED.**

Several assumptions are made:

1. All particles have the same mass (weight)
2. The volume being gridded is an N-dimensional cartesian
3. Boundary conditions are periodic
4. Cell centres are assumed to occur at {0 ... (N-1)/L}

The code should work for any number of dimensions up to 3 (adding more is simple).

I'm pretty sure there must be a more flexible and elegant way to generalise the code for problems with any
dimensionality.  Any comments, suggestions, or PRs are warmly welcomed!
"""

__author__ = "Simon Mutch"
__date__ = "2019-08-26"

import numpy as np
from numba import njit
import itertools
from typing import Tuple


@njit
def _cic_coeffs(point: np.ndarray, cell_size: float, grid_size: int, perms: np.ndarray) -> Tuple[np.ndarray, ...]:
    """Calculate coefficients required for each point.

    Args:
        point:      the point
        cell_size:  the length of the cell
        grid_size:  the total number of cells along a single dimension
        perms:      the cell adjacency array

    Returns:
        a bunch of coeffecients in the form of numpy arrays
    """
    gp_ind = np.floor(point / cell_size).astype(np.int64) % grid_size
    d = point / cell_size - gp_ind
    t = 1.0 - d
    inds = (gp_ind + perms) % grid_size

    return gp_ind, d, t, inds


@njit
def _cic_1d(points: np.ndarray, mass: float, grid: np.ndarray, cell_size: float, grid_size: int, perms: np.ndarray):
    """1D gridding.

    Args:
        points:     an array of points (Np, dim)
        grid:       the grid being calculated (**modified in-place**)
        cell_size:  the length of the cell
        grid_size:  the total number of cells along a single dimension
        perms:      the cell adjacency array
    """
    for ii in range(points.shape[0]):
        point = points[ii]
        gp_ind, d, t, inds = _cic_coeffs(point, cell_size, grid_size, perms)

        for jj in range(inds.shape[0]):
            grid[inds[jj, 0]] += np.prod(perms[jj] * d + (1 - perms[jj]) * t)

    grid *= mass


@njit
def _cic_2d(points: np.ndarray, mass: float, grid: np.ndarray, cell_size: float, grid_size: int, perms: np.ndarray):
    """2D gridding.

    Args:
        points:     an array of points (Np, dim)
        grid:       the grid being calculated (**modified in-place**)
        cell_size:  the length of the cell
        grid_size:  the total number of cells along a single dimension
        perms:      the cell adjacency array
    """
    for ii in range(points.shape[0]):
        point = points[ii]
        gp_ind, d, t, inds = _cic_coeffs(point, cell_size, grid_size, perms)

        for jj in range(inds.shape[0]):
            grid[inds[jj, 0], inds[jj, 1]] += np.prod(perms[jj] * d + (1 - perms[jj]) * t)

    grid *= mass


@njit
def _cic_3d(points: np.ndarray, mass: float, grid: np.ndarray, cell_size: float, grid_size: int, perms: np.ndarray):
    """3D gridding.

    Args:
        points:     an array of points (Np, dim)
        grid:       the grid being calculated (**modified in-place**)
        cell_size:  the length of the cell
        grid_size:  the total number of cells along a single dimension
        perms:      the cell adjacency array
    """
    for ii in range(points.shape[0]):
        point = points[ii]
        gp_ind, d, t, inds = _cic_coeffs(point, cell_size, grid_size, perms)

        for jj in range(inds.shape[0]):
            grid[inds[jj, 0], inds[jj, 1], inds[jj, 2]] += np.prod(perms[jj] * d + (1 - perms[jj]) * t)

    grid *= mass


def cic(points: np.ndarray, box_size: float, grid_size: int, mass: float) -> np.ndarray:
    """Construct a grid from an array of points using a cloud-in-cell (CIC)
    assignment scheme.

    Args:
        points:     an array of points (Np, dim)
        box_size:   the physical length scale of the box
        grid_size:  the total number of cells along a single dimension
        mass:       the mass of each particle

    Returns:
        The newly constructed grid.
    """
    dim = points.shape[1]
    grid = np.zeros([grid_size] * dim)
    cell_size = box_size / grid_size

    perms = np.array(list(itertools.product([0, 1], repeat=dim)))

    if points.shape[1] == 1:
        _cic_1d(points, mass, grid, cell_size, grid_size, perms)

    if points.shape[1] == 2:
        _cic_2d(points, mass, grid, cell_size, grid_size, perms)

    if points.shape[1] == 3:
        _cic_3d(points, mass, grid, cell_size, grid_size, perms)

    return grid


if __name__ == "__main__":
    # Example use
    box_size = 60.0
    grid_size = 64
    n_points = 100
    points = np.random.uniform(0.0, box_size, (n_points, 3))

    grid = cic(points, box_size, grid_size, 1.0)
    assert np.isclose(grid.sum(), n_points)

test_cic.py

"""
For coverage report, run with `NUMBA_DISABLE_JIT=1`.
"""

import pytest
from cic import cic
import numpy as np


@pytest.fixture(params=[1, 2, 3])
def simple_grid(request):
    box_size = 60.0
    grid_size = 64
    np.random.seed(993)

    points = np.random.uniform(0.0, box_size, (100, request.param))
    return cic(points, box_size, grid_size, 1.0)


def test_mass_conservation(simple_grid):
    assert(np.isclose(simple_grid.sum(), 100))