# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

import numpy as np
import ipyvolume.pylab as p3
from numpy.testing import assert_array_equal

from ipyvolume.pylab import make_triangles_lines as new_make_triangles_lines
def old_make_triangles_lines(x,y,z,wrapx=False,wrapy=False):
    def dim(x):
        d = 0
        el = x
        while True:
            try:
                el = el[0]
                d += 1
            except Exception:
                break
        return d

    if dim(x) == 2:
        nx, ny = shape = x.shape
    else:
        nx, ny = shape = x[0].shape

    def reshape(ar):
        if dim(ar) == 3:
            return [k.reshape(-1) for k in ar]
        else:
            return ar.reshape(-1)

    x = reshape(x)
    y = reshape(y)
    z = reshape(z)

    mx = nx if wrapx else nx - 1
    my = ny if wrapy else ny - 1
    triangles = np.zeros(((mx) * (my) * 2, 3), dtype=np.uint32)
    lines = np.zeros(((mx) * (my) * 4, 2), dtype=np.uint32)
    
    
      
 
    def index_from2d(i, j):
        xi = (i % nx)
        yi = (j % ny)
        return ny * xi + yi
        """
        ^ ydir
        |
        2 3
        0 1  ---> x dir
        """
 
    for i in range(mx):
        for j in range(my):
            p0 = index_from2d(i, j)
            p1 = index_from2d(i + 1, j)
            p2 = index_from2d(i, j + 1)
            p3 = index_from2d(i + 1, j + 1)
            triangle_index = (i * my) + j
            triangles[triangle_index * 2 + 0, :] = [p0, p1, p3]
            triangles[triangle_index * 2 + 1, :] = [p0, p3, p2]
            lines[triangle_index * 4 + 0, :] = [p0, p1]
            lines[triangle_index * 4 + 1, :] = [p0, p2]
            lines[triangle_index * 4 + 2, :] = [p2, p3]
            lines[triangle_index * 4 + 3, :] = [p1, p3]
    
    return triangles, lines
            

# +

assert_array_equal(tri_new,tri_old)
assert_array_equal(line_new,line_old)

# +
# Small data set.
nx = 100
ny = 50
X = np.linspace(-5, 5, nx)
Y = np.linspace(-5, 5, ny)
X, Y = np.meshgrid(X, Y, indexing='xy')
R = np.sqrt(X**2 + Y**2)
Z = np.sin(R)

tri_new, line_new = new_make_triangles_lines(X,Y,Z)
tri_old, line_old = old_make_triangles_lines(X,Y,Z)

assert_array_equal(tri_new,tri_old)
assert_array_equal(line_new,line_old)


# %timeit new_make_triangles_lines(X,Y,Z)
# %timeit old_make_triangles_lines(X,Y,Z)


# +
# Larger data set.
nx = 400
ny = 425
X = np.linspace(-5, 5, nx)
Y = np.linspace(-5, 5, ny)
X, Y = np.meshgrid(X, Y, indexing='xy')
R = np.sqrt(X**2 + Y**2)
Z = np.sin(R)

tri_new, line_new = new_make_triangles_lines(X,Y,Z)
tri_old, line_old = old_make_triangles_lines(X,Y,Z)

assert_array_equal(tri_new,tri_old)
assert_array_equal(line_new,line_old)

# %timeit tri_new, line_new = new_make_triangles_lines(X,Y,Z)
# %timeit tri_old, line_old = old_make_triangles_lines(X,Y,Z)
# -



StarcoderdataJupyterScriptsDedupFiltered