<commit_before>"""Test that `hoomd.data.LocalSnapshot` and hoomd.data.LocalSnapshotGPU` work.
"""
from copy import deepcopy
import hoomd
from hoomd.array import HOOMDGPUArray
import numpy as np
import pytest
try:
    # This try block is purely to allow testing locally without mpi4py. We could
    # require it for testing, and simplify the logic here. The CI containers all
    # have mpi4py.
    from mpi4py import MPI
except ImportError:
    skip_mpi4py = True
else:
    skip_mpi4py = False

skip_mpi4py = pytest.mark.skipif(
    skip_mpi4py, reason='mpi4py could not be imported.')

try:
    # We use the CUPY_IMPORTED variable to allow for local GPU testing without
    # CuPy installed. This code could be simplified to only work with CuPy, by
    # requiring its installation for testing. The CI containers already have
    # CuPy installed when build for the GPU.
    import cupy
    CUPY_IMPORTED = True
except ImportError:
    CUPY_IMPORTED = False


"""
_N and _types are distinct in that the local snapshot does not know about them.
We use the underscore to signify this. Those keys are skipped when testing the
local snapshots, though are still used to define the state.
"""

Np = 5
_particle_data = dict(
    _N=Np,
    position=dict(
        np_type=np.floating,
        value=[[-1, -1, -1], [-1, -1, 0], [-1, 0, 0], [1, 1, 1], [1, 0, 0]],
        new_value=[[5, 5, 5]] * Np,
        shape=(Np, 3)),
    velocity=dict(
        np_type=np.floating,
        value=np.linspace(-4, 4, Np * 3).reshape((Np, 3)),
        new_value=np.linspace(4, 8, Np * 3).reshape((Np, 3)),
        shape=(Np, 3)),
    acceleration=dict(
        np_type=np.floating,
        value=np.linspace(-4, 4, Np * 3).reshape((Np, 3)),
        new_value=np.linspace(4, 8, Np * 3).reshape((Np, 3)),
        shape=(Np, 3)),
    typeid=dict(
        np_type=np.integer,
        value=[0, 0, 0, 1, 1],
        new_value=[1, 1, 1, 0, 0],
        shape=(Np,)),
    mass=dict(np_type=np.floating, value=[5, 4, 3, 2, 1],
              new_value=[1, 2, 3, 4, 5], shape=(Np,)),
    charge=dict(np_type=np.floating, value=[1, 2, 3, 2, 1],
                new_value=[-1, -1, -3, -2, -1], shape=(Np,)),
    diameter=dict(np_type=np.floating, value=[5, 2, 3, 2, 5],
                  new_value=[2, 1, 0.5, 1, 2], shape=(Np,)),
    image=dict(
        np_type=np.integer,
        value=np.linspace(-10, 20, Np * 3, dtype=np.int).reshape(Np, 3),
        new_value=np.linspace(-20, 10, Np * 3, dtype=np.int).reshape(Np, 3),
        shape=(Np, 3)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Np,)),
    _types=['p1', 'p2']
)


Nb = 2
_bond_data = dict(
    _N=Nb,
    typeid=dict(np_type=np.unsignedinteger,
                value=[0, 1], new_value=[1, 0], shape=(Nb,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[0, 1], [2, 3]],
        new_value=[[1, 0], [3, 2]],
        shape=(Nb, 2)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Nb,)),
    _types=['b1', 'b2']
)


Na = 2
_angle_data = dict(
    _N=Na,
    typeid=dict(np_type=np.unsignedinteger,
                value=[1, 0], new_value=[0, 1], shape=(Na,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[0, 1, 2], [2, 3, 4]],
        new_value=[[1, 3, 4], [0, 2, 4]],
        shape=(Na, 3)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Na,)),
    _types=['a1', 'a2']
)


Nd = 2
_dihedral_data = dict(
    _N=Nd,
    typeid=dict(np_type=np.unsignedinteger,
                value=[1, 0], new_value=[0, 1], shape=(Nd,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[0, 1, 2, 3], [1, 2, 3, 4]],
        new_value=[[4, 3, 2, 1], [2, 4, 0, 1]],
        shape=(Nd, 4)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Nd,)),
    _types=['d1', 'd2']
)


Ni = 2
_improper_data = dict(
    _N=Ni,
    typeid=dict(
        np_type=np.unsignedinteger, value=[0, 0], shape=(Ni,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[3, 2, 1, 0], [1, 2, 3, 4]],
        new_value=[[1, 2, 3, 0], [4, 2, 3, 1]],
        shape=(Ni, 4)),
    tag=dict(
        np_type=np.unsignedinteger, value=None, shape=(Ni,)),
    _types=['i1']
)

Nc = 3
_constraint_data = dict(
    _N=Nc,
    value=dict(np_type=np.floating, value=[2.5, 0.5, 2.],
               new_value=[3., 1.5, 1.], shape=(Nc,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[0, 1], [2, 3], [1, 3]],
        new_value=[[4, 1], [3, 1], [2, 4]],
        shape=(Nc, 2)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Nc,)),
)


Npa = 2
_pair_data = dict(
    _N=Npa,
    typeid=dict(np_type=np.unsignedinteger,
                value=[0, 1], new_value=[1, 0], shape=(Npa,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[0, 1], [2, 3]],
        new_value=[[4, 1], [0, 3]],
        shape=(Npa, 2)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Npa,)),
    _types=['p1', 'p2']
)


_global_dict = dict(rtag=dict(
    particles=dict(np_type=np.unsignedinteger, value=None, shape=(Np,)),
    bonds=dict(np_type=np.unsignedinteger, value=None, shape=(Nb,)),
    angles=dict(np_type=np.unsignedinteger, value=None, shape=(Na,)),
    dihedrals=dict(np_type=np.unsignedinteger, value=None, shape=(Nd,)),
    impropers=dict(np_type=np.unsignedinteger, value=None, shape=(Ni,)),
    constraints=dict(np_type=np.unsignedinteger, value=None, shape=(Nc,)),
    pairs=dict(np_type=np.unsignedinteger, value=None, shape=(Npa,)),
))


@pytest.fixture(scope='session')
def base_snapshot(device):
    """Defines a snapshot using the data given above."""
    def set_snapshot(snap, data, base):
        """Sets individual sections of snapshot (e.g. particles)."""
        snap_section = getattr(snap, base)
        for k in data:
            if k.startswith('_'):
                setattr(snap_section, k[1:], data[k])
                continue
            elif data[k]['value'] is None:
                continue
            try:
                array = getattr(snap_section, k)
                array[:] = data[k]['value']
            except TypeError:
                setattr(snap_section, k, data[k]['value'])

    snapshot = hoomd.Snapshot(device.comm)

    if snapshot.exists:
        snapshot.configuration.box = [2.1, 2.1, 2.1, 0, 0, 0]
        set_snapshot(snapshot, _particle_data, 'particles')
        set_snapshot(snapshot, _bond_data, 'bonds')
        set_snapshot(snapshot, _angle_data, 'angles')
        set_snapshot(snapshot, _dihedral_data, 'dihedrals')
        set_snapshot(snapshot, _improper_data, 'impropers')
        set_snapshot(snapshot, _constraint_data, 'constraints')
        set_snapshot(snapshot, _pair_data, 'pairs')
    return snapshot


@pytest.fixture(scope='session')
def cpu_simulation_factory(device_cpu):
    """Creates the simulation from the base_snapshot."""

    def make_simulation(snapshot):
        sim = hoomd.Simulation(device_cpu)

        # reduce sorter grid to avoid Hilbert curve overhead in unit tests
        for tuner in sim.operations.tuners:
            if isinstance(tuner, hoomd.tuner.ParticleSorter):
                tuner.grid = 8

        sim.create_state_from_snapshot(snapshot)
        return sim

    return make_simulation


@pytest.fixture(scope='session')
def gpu_simulation_factory(device_gpu):
    """Creates the simulation from the base_snapshot."""

    def make_simulation(snapshot):
        sim = hoomd.Simulation(device_gpu)

        # reduce sorter grid to avoid Hilbert curve overhead in unit tests
        for tuner in sim.operations.tuners:
            if isinstance(tuner, hoomd.tuner.ParticleSorter):
                tuner.grid = 8

        sim.create_state_from_snapshot(snapshot)
        return sim

    return make_simulation


def check_box(local_snapshot, global_box, ranks):
    """General check that ``box`` and ``local_box`` properties work."""
    assert type(local_snapshot.global_box) == hoomd.Box
    assert type(local_snapshot.local_box) == hoomd.Box

    if ranks == 1:
        assert local_snapshot.local_box == global_box
        assert local_snapshot.global_box == global_box
    else:
        assert local_snapshot.local_box != global_box
        assert local_snapshot.global_box == global_box


def test_box_cpu(cpu_simulation_factory, base_snapshot):
    sim = cpu_simulation_factory(base_snapshot)
    with sim.state.cpu_local_snapshot as data:
        check_box(data, sim.state.box, sim.device.num_ranks)

<commit_msg>Improve check_box function for local snapshot tests<commit_after>"""Test that `hoomd.data.LocalSnapshot` and hoomd.data.LocalSnapshotGPU` work.
"""
from copy import deepcopy
import hoomd
from hoomd.array import HOOMDGPUArray
import numpy as np
import pytest
try:
    # This try block is purely to allow testing locally without mpi4py. We could
    # require it for testing, and simplify the logic here. The CI containers all
    # have mpi4py.
    from mpi4py import MPI
except ImportError:
    skip_mpi4py = True
else:
    skip_mpi4py = False

skip_mpi4py = pytest.mark.skipif(
    skip_mpi4py, reason='mpi4py could not be imported.')

try:
    # We use the CUPY_IMPORTED variable to allow for local GPU testing without
    # CuPy installed. This code could be simplified to only work with CuPy, by
    # requiring its installation for testing. The CI containers already have
    # CuPy installed when build for the GPU.
    import cupy
    CUPY_IMPORTED = True
except ImportError:
    CUPY_IMPORTED = False


"""
_N and _types are distinct in that the local snapshot does not know about them.
We use the underscore to signify this. Those keys are skipped when testing the
local snapshots, though are still used to define the state.
"""

Np = 5
_particle_data = dict(
    _N=Np,
    position=dict(
        np_type=np.floating,
        value=[[-1, -1, -1], [-1, -1, 0], [-1, 0, 0], [1, 1, 1], [1, 0, 0]],
        new_value=[[5, 5, 5]] * Np,
        shape=(Np, 3)),
    velocity=dict(
        np_type=np.floating,
        value=np.linspace(-4, 4, Np * 3).reshape((Np, 3)),
        new_value=np.linspace(4, 8, Np * 3).reshape((Np, 3)),
        shape=(Np, 3)),
    acceleration=dict(
        np_type=np.floating,
        value=np.linspace(-4, 4, Np * 3).reshape((Np, 3)),
        new_value=np.linspace(4, 8, Np * 3).reshape((Np, 3)),
        shape=(Np, 3)),
    typeid=dict(
        np_type=np.integer,
        value=[0, 0, 0, 1, 1],
        new_value=[1, 1, 1, 0, 0],
        shape=(Np,)),
    mass=dict(np_type=np.floating, value=[5, 4, 3, 2, 1],
              new_value=[1, 2, 3, 4, 5], shape=(Np,)),
    charge=dict(np_type=np.floating, value=[1, 2, 3, 2, 1],
                new_value=[-1, -1, -3, -2, -1], shape=(Np,)),
    diameter=dict(np_type=np.floating, value=[5, 2, 3, 2, 5],
                  new_value=[2, 1, 0.5, 1, 2], shape=(Np,)),
    image=dict(
        np_type=np.integer,
        value=np.linspace(-10, 20, Np * 3, dtype=np.int).reshape(Np, 3),
        new_value=np.linspace(-20, 10, Np * 3, dtype=np.int).reshape(Np, 3),
        shape=(Np, 3)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Np,)),
    _types=['p1', 'p2']
)


Nb = 2
_bond_data = dict(
    _N=Nb,
    typeid=dict(np_type=np.unsignedinteger,
                value=[0, 1], new_value=[1, 0], shape=(Nb,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[0, 1], [2, 3]],
        new_value=[[1, 0], [3, 2]],
        shape=(Nb, 2)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Nb,)),
    _types=['b1', 'b2']
)


Na = 2
_angle_data = dict(
    _N=Na,
    typeid=dict(np_type=np.unsignedinteger,
                value=[1, 0], new_value=[0, 1], shape=(Na,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[0, 1, 2], [2, 3, 4]],
        new_value=[[1, 3, 4], [0, 2, 4]],
        shape=(Na, 3)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Na,)),
    _types=['a1', 'a2']
)


Nd = 2
_dihedral_data = dict(
    _N=Nd,
    typeid=dict(np_type=np.unsignedinteger,
                value=[1, 0], new_value=[0, 1], shape=(Nd,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[0, 1, 2, 3], [1, 2, 3, 4]],
        new_value=[[4, 3, 2, 1], [2, 4, 0, 1]],
        shape=(Nd, 4)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Nd,)),
    _types=['d1', 'd2']
)


Ni = 2
_improper_data = dict(
    _N=Ni,
    typeid=dict(
        np_type=np.unsignedinteger, value=[0, 0], shape=(Ni,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[3, 2, 1, 0], [1, 2, 3, 4]],
        new_value=[[1, 2, 3, 0], [4, 2, 3, 1]],
        shape=(Ni, 4)),
    tag=dict(
        np_type=np.unsignedinteger, value=None, shape=(Ni,)),
    _types=['i1']
)

Nc = 3
_constraint_data = dict(
    _N=Nc,
    value=dict(np_type=np.floating, value=[2.5, 0.5, 2.],
               new_value=[3., 1.5, 1.], shape=(Nc,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[0, 1], [2, 3], [1, 3]],
        new_value=[[4, 1], [3, 1], [2, 4]],
        shape=(Nc, 2)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Nc,)),
)


Npa = 2
_pair_data = dict(
    _N=Npa,
    typeid=dict(np_type=np.unsignedinteger,
                value=[0, 1], new_value=[1, 0], shape=(Npa,)),
    group=dict(
        np_type=np.unsignedinteger,
        value=[[0, 1], [2, 3]],
        new_value=[[4, 1], [0, 3]],
        shape=(Npa, 2)),
    tag=dict(np_type=np.unsignedinteger, value=None, shape=(Npa,)),
    _types=['p1', 'p2']
)


_global_dict = dict(rtag=dict(
    particles=dict(np_type=np.unsignedinteger, value=None, shape=(Np,)),
    bonds=dict(np_type=np.unsignedinteger, value=None, shape=(Nb,)),
    angles=dict(np_type=np.unsignedinteger, value=None, shape=(Na,)),
    dihedrals=dict(np_type=np.unsignedinteger, value=None, shape=(Nd,)),
    impropers=dict(np_type=np.unsignedinteger, value=None, shape=(Ni,)),
    constraints=dict(np_type=np.unsignedinteger, value=None, shape=(Nc,)),
    pairs=dict(np_type=np.unsignedinteger, value=None, shape=(Npa,)),
))


@pytest.fixture(scope='session')
def base_snapshot(device):
    """Defines a snapshot using the data given above."""
    def set_snapshot(snap, data, base):
        """Sets individual sections of snapshot (e.g. particles)."""
        snap_section = getattr(snap, base)
        for k in data:
            if k.startswith('_'):
                setattr(snap_section, k[1:], data[k])
                continue
            elif data[k]['value'] is None:
                continue
            try:
                array = getattr(snap_section, k)
                array[:] = data[k]['value']
            except TypeError:
                setattr(snap_section, k, data[k]['value'])

    snapshot = hoomd.Snapshot(device.comm)

    if snapshot.exists:
        snapshot.configuration.box = [2.1, 2.1, 2.1, 0, 0, 0]
        set_snapshot(snapshot, _particle_data, 'particles')
        set_snapshot(snapshot, _bond_data, 'bonds')
        set_snapshot(snapshot, _angle_data, 'angles')
        set_snapshot(snapshot, _dihedral_data, 'dihedrals')
        set_snapshot(snapshot, _improper_data, 'impropers')
        set_snapshot(snapshot, _constraint_data, 'constraints')
        set_snapshot(snapshot, _pair_data, 'pairs')
    return snapshot


@pytest.fixture(scope='session')
def cpu_simulation_factory(device_cpu):
    """Creates the simulation from the base_snapshot."""

    def make_simulation(snapshot):
        sim = hoomd.Simulation(device_cpu)

        # reduce sorter grid to avoid Hilbert curve overhead in unit tests
        for tuner in sim.operations.tuners:
            if isinstance(tuner, hoomd.tuner.ParticleSorter):
                tuner.grid = 8

        sim.create_state_from_snapshot(snapshot)
        return sim

    return make_simulation


@pytest.fixture(scope='session')
def gpu_simulation_factory(device_gpu):
    """Creates the simulation from the base_snapshot."""

    def make_simulation(snapshot):
        sim = hoomd.Simulation(device_gpu)

        # reduce sorter grid to avoid Hilbert curve overhead in unit tests
        for tuner in sim.operations.tuners:
            if isinstance(tuner, hoomd.tuner.ParticleSorter):
                tuner.grid = 8

        sim.create_state_from_snapshot(snapshot)
        return sim

    return make_simulation


def check_box(local_snapshot, global_box, ranks):
    """General check that ``box`` and ``local_box`` properties work."""
    assert type(local_snapshot.global_box) == hoomd.Box
    assert type(local_snapshot.local_box) == hoomd.Box

    assert local_snapshot.global_box == global_box
    # The local box and global box are equal if and only if
    # we run on a single rank.
    assert (local_snapshot.local_box == global_box) == (ranks == 1)


def test_box_cpu(cpu_simulation_factory, base_snapshot):
    sim = cpu_simulation_factory(base_snapshot)
    with sim.state.cpu_local_snapshot as data:
        check_box(data, sim.state.box, sim.device.num_ranks)


{'fraction_non_alphanumeric': 0.150528117070868, 'fraction_numerical': 0.027733522157313978, 'mean_word_length': 2.9204256303492944, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3205155', 'n_tokens_mistral': 6374, 'n_tokens_neox': 6069, 'n_words': 1430}