ALLEN INSTITUTE

# Source code for bmtk.tests.builder.auxi.test_node_params

    import os
    import pytest
    import numpy as np

    from bmtk.builder.auxi.node_params import positions_columinar, positions_cuboid, positions_list, positions_rect_prism, positions_ellipsoid, positions_density_matrix, positions_nrrd

    try:
        from neuron import h
        from bmtk.simulator.bionet.nrn import load_neuron_modules
        load_neuron_modules(
            mechanisms_dir='bmtk/tests/simulator/bionet/components/mechanisms',
            templates_dir=None,
            default_templates=True
        )
        vecstim = h.VecStim()
        has_mechanism = True
    except (AttributeError, ModuleNotFoundError):
        has_mechanism = False

    def test_positions_columinar():
        np.random.seed(1)
        points = positions_columinar(N=10, center=[0.0, 0.0, 0.0], height=10.0, min_radius=1.0, max_radius=2.0)
        assert(points.shape == (10, 3))

        # check height restriction
        ys = points[:, 1]
        assert(np.min(ys) >= -5.0)
        assert(np.max(ys) <= 5.0)

        # check radius restrictions
        xz_plane = points[:,[0,2]]
        dists = np.linalg.norm(xz_plane, axis=1)
        assert(np.min(dists) >= 1.0)
        assert(np.max(dists) <= 2.0)

    def test_positions_cuboid():
        np.random.seed(1)
        points = positions_cuboid(N=10, center=[0.0, 0.0, 0.0], height=5.0, xside_length=2.0, yside_length=2.0,
                                  min_dist=1.0)
        assert(points.shape == (10, 3))

    def test_positions_list():
        assert(positions_list().shape == (2, 3))

    def test_positions_rect_prism():
        np.random.seed(1)
        points = positions_rect_prism(N=10, center=[10.0, 20.0, 30.0], height=15.0, x_length=3.0, z_length=8.0)
        assert(points.shape == (10, 3))

         # check height, width, length restriction
        ys = points[:, 1]
        assert(np.min(ys) >= 20.0-15.0)
        assert(np.max(ys) <= 20.0+15.0)
        xs = points[:, 0]
        assert(np.min(xs) >= 10.0-3.0)
        assert(np.max(xs) <= 10.0+3.0)
        zs = points[:, 2]
        assert(np.min(zs) >= 30.0-8.0)
        assert(np.max(zs) <= 30.0+8.0)

    def test_positions_ellipsoid():
        np.random.seed(1)
        points = positions_ellipsoid(N=10, center=[10.0, 20.0, 30.0], height=15.0, x_length=3.0, z_length=8.0)
        assert(points.shape == (10, 3))

         # check height, width, length restriction
        dxs = points[:, 0] - 10.0
        dys = points[:, 1] - 20.0
        dzs = points[:, 2] - 30.0

        assert(np.max(dxs**2/3.0**2+dys**2/15.0**2+dzs**2/8.0**2)<1)

    def test_positions_density_matrix():
        np.random.seed(1)
        mat = 100000*np.array([[[0,1,5],[0,0.5,0],[0,0,0]],
                    [[0,1,5],[0,0.5,0],[0,0,0]],
                    [[0,1,1],[0,1,0],[0,0,5]]])

        position_scale=np.array([[25.0,0,0],[0,25.0,0],[0,0,25.0]])

        points = positions_density_matrix(mat, position_scale)

        assert(points.shape == (34,3))
        assert(np.max(abs(points[:,0]))<25*3)
        assert(np.max(abs(points[:,1]))<25*3)
        assert(np.max(abs(points[:,2]))<25*3)

    @pytest.mark.skipif(not has_mechanism, reason='Mechanisms has not been compiled, run nrnivmodl mechanisms.')
    def test_positions_nrrd():
        np.random.seed(1)

        CURR_DIR = os.path.dirname(os.path.realpath(__file__))
        nrrd_filename = os.path.join(CURR_DIR, 'structure_721.nrrd')
        points = positions_nrrd(nrrd_filename, 2000)
        assert(points.shape==(2048, 3))

        points = positions_nrrd(nrrd_filename, 2000, split_bilateral='z')
        assert(points.shape==(1024, 3))

    if __name__ == '__main__':
        # test_positions_columinar()
        # test_positions_cuboid()
        test_positions_list()
        test_positions_rect_prism()