def test_pow(): # note hat a**2 is converted to a*a and a**0.5 to sqrt(a) @python2shader_and_validate def compute_shader( index_xyz: ("input", "GlobalInvocationId", ivec3), data1: ("buffer", 0, Array(f32)), data2: ("buffer", 1, Array(vec4)), ): index = index_xyz.x a = data1[index] data2[index] = vec4(a**2, a**0.5, a**3.0, a**3.1) skip_if_no_wgpu() values1 = [i - 5 for i in range(10)] inp_arrays = {0: (ctypes.c_float * 10)(*values1)} out_arrays = {1: ctypes.c_float * 40} out = compute_with_buffers(inp_arrays, out_arrays, compute_shader) res = list(out[1]) assert res[0::4] == [i**2 for i in values1] assert iters_close(res[1::4], [i**0.5 for i in values1]) assert res[2::4] == [i**3 for i in values1] assert iters_close(res[3::4], [i**3.1 for i in values1])
def test_mix(): @python2shader_and_validate def compute_shader( index_xyz: ("input", "GlobalInvocationId", ivec3), data1: ("buffer", 0, Array(vec4)), data2: ("buffer", 1, Array(vec4)), ): index = index_xyz.x v = data1[index] v1 = mix(v.x, v.y, v.z) v2 = mix(vec2(v.x, v.x), vec2(v.y, v.y), v.z) data2[index] = vec4(v1, v2.x, v2.y, 0.0) skip_if_no_wgpu() values1 = [-4, -3, -2, -1, +0, +0, +1, +2, +3, +4] values2 = [-2, -5, -5, +2, +2, -1, +3, +1, +1, -6] weights = [0.1 * i for i in range(10)] stubs = [0] * 10 values = sum(zip(values1, values2, weights, stubs), ()) inp_arrays = {0: (ctypes.c_float * 40)(*values)} out_arrays = {1: ctypes.c_float * 40} out = compute_with_buffers(inp_arrays, out_arrays, compute_shader, n=10) res = list(out[1]) ref = [ values1[i] * (1 - w) + values2[i] * w for i, w in enumerate(weights) ] assert iters_close(res[0::4], ref) assert iters_close(res[1::4], ref) assert iters_close(res[2::4], ref)
def test_abs(): @python2shader_and_validate def compute_shader( index_xyz: ("input", "GlobalInvocationId", ivec3), data1: ("buffer", 0, Array(f32)), data2: ("buffer", 1, Array(i32)), data3: ("buffer", 2, Array(vec2)), ): index = index_xyz.x v1 = abs(data1[index]) # float v2 = abs(data2[index]) # int data3[index] = vec2(f32(v1), v2) skip_if_no_wgpu() values1 = [random.uniform(-2, 2) for i in range(10)] values2 = [random.randint(-100, 100) for i in range(10)] inp_arrays = { 0: (ctypes.c_float * 10)(*values1), 1: (ctypes.c_int * 10)(*values2) } out_arrays = {2: ctypes.c_float * 20} out = compute_with_buffers(inp_arrays, out_arrays, compute_shader, n=10) res = list(out[2]) assert iters_close(res[0::2], [abs(v) for v in values1]) assert res[1::2] == [abs(v) for v in values2]
def test_normalize(): @python2shader_and_validate def compute_shader( index_xyz: ("input", "GlobalInvocationId", ivec3), data1: ("buffer", 0, Array(f32)), data2: ("buffer", 1, Array(vec2)), ): index = index_xyz.x v = data1[index] data2[index] = normalize(vec2(v, v)) skip_if_no_wgpu() values1 = [i - 5 for i in range(10)] inp_arrays = {0: (ctypes.c_float * 10)(*values1)} out_arrays = {1: ctypes.c_float * 20} out = compute_with_buffers(inp_arrays, out_arrays, compute_shader, n=10) res = list(out[1]) assert iters_close(res[:10], [-(2**0.5) / 2 for i in range(10)]) assert iters_close(res[-8:], [+(2**0.5) / 2 for i in range(8)]) assert math.isnan(res[10]) and math.isnan( res[11]) # or can this also be inf?
def test_sqrt(): @python2shader_and_validate def compute_shader( index_xyz: ("input", "GlobalInvocationId", ivec3), data1: ("buffer", 0, Array(f32)), data2: ("buffer", 1, Array(vec4)), ): index = index_xyz.x a = data1[index] data2[index] = vec4(a**0.5, math.sqrt(a), stdlib.sqrt(a), 0.0) skip_if_no_wgpu() values1 = [i for i in range(10)] inp_arrays = {0: (ctypes.c_float * 10)(*values1)} out_arrays = {1: ctypes.c_float * 40} out = compute_with_buffers(inp_arrays, out_arrays, compute_shader) res = list(out[1]) ref = [i**0.5 for i in values1] assert iters_close(res[0::4], ref) assert iters_close(res[1::4], ref) assert iters_close(res[2::4], ref)
def test_math_constants(): @python2shader_and_validate def compute_shader( index_xyz: ("input", "GlobalInvocationId", ivec3), data2: ("buffer", 1, Array(f32)), ): index = index_xyz.x if index % 2 == 0: data2[index] = math.pi else: data2[index] = math.e skip_if_no_wgpu() out = compute_with_buffers({}, {1: ctypes.c_float * 10}, compute_shader, n=10) res = list(out[1]) assert iters_close(res, [math.pi, math.e] * 5)
def test_length(): @python2shader_and_validate def compute_shader( index_xyz: ("input", "GlobalInvocationId", ivec3), data1: ("buffer", 0, Array(vec2)), data2: ("buffer", 1, Array(f32)), ): index = index_xyz.x data2[index] = length(data1[index]) skip_if_no_wgpu() values1 = [random.uniform(-2, 2) for i in range(20)] inp_arrays = {0: (ctypes.c_float * 20)(*values1)} out_arrays = {1: ctypes.c_float * 10} out = compute_with_buffers(inp_arrays, out_arrays, compute_shader, n=10) res = list(out[1]) ref = [(values1[i * 2]**2 + values1[i * 2 + 1]**2)**0.5 for i in range(10)] assert iters_close(res, ref)