py: test swift grid writing

test.py

import numpy as np
import h5py as h5
import xarray as xa
from pathlib import Path
import numba

import pytest

SNAP = 12


@pytest.mark.parametrize("name", "Density Vx Vy Vz".split())
def test_run_modes(name):
    """Test that we get the same result independent of the mode we run in
    (threaded vs. N ranks vs. 1 rank)"""

    run_paths = Path("./").glob("run-*")
    grids = {}
    for rp in run_paths:
        with h5.File(str(rp / f"snap_{SNAP:04d}.hdf5"), "r") as fp:
            grids[str(rp).lstrip("run-")] = xa.DataArray(fp[f"PartType1/Grids/{name}"][:])

    baseline = None
    for v in grids.values():
        if baseline is None:
            baseline = v.copy()
        else:
            # Note that I've uped the atol here to account for differences
            # caused by the order of summing from using different numbers of
            # ranks and doing a reduction.
            assert np.all(np.isclose(baseline, v, atol=3e-7))


@numba.jit(nopython=True)
def increment_grid(inds, grid):
    for ii in range(inds.shape[0]):
        grid[(inds[ii][0], inds[ii][1], inds[ii][2])] += 1.0


def ngp(parts, dim, box_size):
    grid = np.zeros([dim] * 3)
    cell_size = box_size / dim
    inds = np.around(parts / cell_size).astype(int) % dim
    increment_grid(inds, grid)
    return grid


def test_total_density():

    run_path = Path("./run-4ranks")
    dims = ("x", "y", "z")

    with h5.File(str(run_path / f"snap_{SNAP:04d}.hdf5"), "r") as fp:
        grid = xa.DataArray(fp["PartType1/Grids/Density"][:], dims=dims)
        mass = fp["PartType1/Masses"][0]
        Np = fp["Header"].attrs["NumPart_Total"][1]
        box_size = fp["Header"].attrs["BoxSize"][0]

    truth = mass * Np / box_size ** 3
    total = grid.sum() / (grid.shape[0] ** 3)

    assert np.isclose(truth, total)


def test_ngp():

    run_path = Path("./run-4ranks")
    dims = ("x", "y", "z")

    with h5.File(str(run_path / f"snap_{SNAP:04d}.hdf5"), "r") as fp:
        density = xa.DataArray(fp["PartType1/Grids/Density"][:], dims=dims)
        particles = fp["PartType1/Coordinates"][:]
        mass = fp["PartType1/Masses"][0]
        box_size = fp["Header"].attrs["BoxSize"][0]

    truth = ngp(particles, density.shape[0], box_size)
    truth *= mass / box_size ** 3 * density.shape[0] ** 3

    assert np.all(np.isclose(density, truth))