<commit_before>import hoomd
import pytest
import numpy as np

_harmonic_args = {
    'k': [3.0, 10.0, 5.0],
    't0': [np.pi / 2, np.pi / 4, np.pi / 6]
}
_harmonic_arg_list = [
    (hoomd.md.angle.Harmonic,
     dict(zip(_harmonic_args, val)) for val in zip(*_harmonic_args.values()))
]

_cosinesq_args = {
    'k': [3.0, 10.0, 5.0],
    't0': [np.pi / 2, np.pi / 4, np.pi / 6]
}
_cosinesq_arg_list = [
    (hoomd.md.angle.Cosinesq,
     dict(zip(_cosinesq_args, val)) for val in zip(*_cosinesq_args.values()))
]

def get_angle_and_args():
    return _harmonic_arg_list + _cosinesq_arg_list


def get_angle_args_forces_and_energies():
    harmonic_forces = [-1.5708, 2.6180, 2.6180]
    harmonic_energies = [0.4112, 0.3427, 0.6854]
    cosinesq_forces = [-1.29904, 1.7936, 1.58494]
    cosinesq_energies = [0.375, 0.214466, 0.334936]

    harmonic_args_and_vals = []
    cosinesq_args_and_vals = []
    for i in range(3):
        harmonic_args_and_vals.append(_harmonic_arg_list[i][0],
                                      _harmonic_arg_list[i][1],
                                      harmonic_forces[i],
                                      harmonic_energies[i])
        cosinesq_args_and_vals.append(_cosinesq_arg_list[i][0],
                                      _cosinesq_arg_list[i][1],
                                      cosinesq_forces[i],
                                      cosinesq_energies[i])

    return harmonic_args_and_vals + cosinesq_args_and_vals


@pytest.fixture(scope='session')
def triplet_snapshot_factory(device):

    def make_snapshot(d=1.0,
                      theta_deg=60,
                      particle_types=['A'],
                      dimensions=3,
                      L=20):
        theta_rad = theta_deg * (np.pi / 180)
        s = hoomd.Snapshot(device.communicator)
        N = 3
        if s.exists:
            box = [L, L, L, 0, 0, 0]
            if dimensions == 2:
                box[2] = 0
            s.configuration.box = box
            s.particles.N = N
            # shift particle positions slightly in z so MPI tests pass
            s.particles.position[:] = [
                [-d * np.sin(theta_rad / 2), d * np.cos(theta_rad / 2), 0.1],
                [0.0, 0.0, 0.1],
                [d * np.sin(theta_rad / 2), d * np.cos(theta_rad / 2), 0.1]
            ]
            s.particles.types = particle_types
            if dimensions == 2:
                box[2] = 0
                s.particles.position[:] = [[
                    -d * np.sin(theta_rad / 2), d * np.cos(theta_rad / 2) + 0.1,
                    0.0
                ], [0.0, 0.1, 0.0],
                                           [
                                               d * np.sin(theta_rad / 2),
                                               d * np.cos(theta_rad / 2) + 0.1,
                                               0.0
                                           ]]
        return s

    return make_snapshot


@pytest.mark.parametrize("angle_and_args", get_angle_and_args())
def test_before_attaching(angle_and_args):
    angle, args = angle_and_args
    angle_potential = angle()
    angle_potential.params['backbone'] = args
    for key in args.keys():
        np.testing.assert_allclose(angle_potential.params['backbone'][key],
                                   args[key],
                                   rtol=1e-6)


@pytest.mark.parametrize("angle_and_args", get_angle_and_args())
def test_after_attaching(triplet_snapshot_factory, simulation_factory,
                         angle_and_args):
    snap = triplet_snapshot_factory()
    if snap.exists:
        snap.angles.N = 1
        snap.angles.types = ['backbone']
        snap.angles.typeid[0] = 0
        snap.angles.group[0] = (0, 1, 2)
    sim = simulation_factory(snap)

    angle, args = angle_and_args
    angle_potential = angle()
    angle_potential.params['backbone'] = args

    integrator = hoomd.md.Integrator(dt=0.005)

    integrator.forces.append(angle_potential)

    langevin = hoomd.md.methods.Langevin(kT=1,<commit_msg>Parametrize test angle before attaching<commit_after>import hoomd
import pytest
import numpy as np

_harmonic_args = {
    'k': [3.0, 10.0, 5.0],
    't0': [np.pi / 2, np.pi / 4, np.pi / 6]
}
_harmonic_arg_list = [
    (hoomd.md.angle.Harmonic,
     dict(zip(_harmonic_args, val)) for val in zip(*_harmonic_args.values()))
]

_cosinesq_args = {
    'k': [3.0, 10.0, 5.0],
    't0': [np.pi / 2, np.pi / 4, np.pi / 6]
}
_cosinesq_arg_list = [
    (hoomd.md.angle.Cosinesq,
     dict(zip(_cosinesq_args, val)) for val in zip(*_cosinesq_args.values()))
]

def get_angle_and_args():
    return _harmonic_arg_list + _cosinesq_arg_list


def get_angle_args_forces_and_energies():
    harmonic_forces = [-1.5708, 2.6180, 2.6180]
    harmonic_energies = [0.4112, 0.3427, 0.6854]
    cosinesq_forces = [-1.29904, 1.7936, 1.58494]
    cosinesq_energies = [0.375, 0.214466, 0.334936]

    harmonic_args_and_vals = []
    cosinesq_args_and_vals = []
    for i in range(3):
        harmonic_args_and_vals.append(_harmonic_arg_list[i][0],
                                      _harmonic_arg_list[i][1],
                                      harmonic_forces[i],
                                      harmonic_energies[i])
        cosinesq_args_and_vals.append(_cosinesq_arg_list[i][0],
                                      _cosinesq_arg_list[i][1],
                                      cosinesq_forces[i],
                                      cosinesq_energies[i])

    return harmonic_args_and_vals + cosinesq_args_and_vals


@pytest.fixture(scope='session')
def triplet_snapshot_factory(device):

    def make_snapshot(d=1.0,
                      theta_deg=60,
                      particle_types=['A'],
                      dimensions=3,
                      L=20):
        theta_rad = theta_deg * (np.pi / 180)
        s = hoomd.Snapshot(device.communicator)
        N = 3
        if s.exists:
            box = [L, L, L, 0, 0, 0]
            if dimensions == 2:
                box[2] = 0
            s.configuration.box = box
            s.particles.N = N
            # shift particle positions slightly in z so MPI tests pass
            s.particles.position[:] = [
                [-d * np.sin(theta_rad / 2), d * np.cos(theta_rad / 2), 0.1],
                [0.0, 0.0, 0.1],
                [d * np.sin(theta_rad / 2), d * np.cos(theta_rad / 2), 0.1]
            ]
            s.particles.types = particle_types
            if dimensions == 2:
                box[2] = 0
                s.particles.position[:] = [[
                    -d * np.sin(theta_rad / 2), d * np.cos(theta_rad / 2) + 0.1,
                    0.0
                ], [0.0, 0.1, 0.0],
                                           [
                                               d * np.sin(theta_rad / 2),
                                               d * np.cos(theta_rad / 2) + 0.1,
                                               0.0
                                           ]]
        return s

    return make_snapshot


@pytest.mark.parametrize("angle_cls, potential_kwargs", get_angle_and_args())
def test_before_attaching(angle_cls, potential_kwargs):
    angle_potential = angle_cls()
    angle_potential.params['backbone'] = potential_kwargs
    for key in potential_kwargs.keys():
        np.testing.assert_allclose(angle_potential.params['backbone'][key],
                                   potential_kwargs[key],
                                   rtol=1e-6)


@pytest.mark.parametrize("angle_and_args", get_angle_and_args())
def test_after_attaching(triplet_snapshot_factory, simulation_factory,
                         angle_and_args):
    snap = triplet_snapshot_factory()
    if snap.exists:
        snap.angles.N = 1
        snap.angles.types = ['backbone']
        snap.angles.typeid[0] = 0
        snap.angles.group[0] = (0, 1, 2)
    sim = simulation_factory(snap)

    angle, args = angle_and_args
    angle_potential = angle()
    angle_potential.params['backbone'] = args

    integrator = hoomd.md.Integrator(dt=0.005)

    integrator.forces.append(angle_potential)

    langevin = hoomd.md.methods.Langevin(kT=1,
{'fraction_non_alphanumeric': 0.11662379027318388, 'fraction_numerical': 0.03871125811588877, 'mean_word_length': 1.5925690695458876, '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': 34, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1437867', 'n_tokens_mistral': 2924, 'n_tokens_neox': 2636, 'n_words': 470}