src/StructDiffusion/utils/np_speed_test.py

import numpy as np
import time

def numpy_test(w, h, r):
    v, u = np.indices((h, w))
    theta = (u - w / 2.) * 2 * np.pi / w
    phi = (v - h / 2.) * np.pi / h
    cos_phi = np.cos(phi)
    x = cos_phi * np.sin(theta)
    y = np.sin(phi)
    z = cos_phi * np.cos(theta)
    
    ray = np.dstack((x, y, z))
    ray = ray.reshape(-1, 3).dot(r.T)
    ray.shape = (h, w, 3)
    
    x, y, z = np.dsplit(ray, 3)
    theta = np.arctan2(x, z)
    phi = np.arcsin(y)
    u = theta * w / 2 / np.pi + w / 2.
    v = phi * h / np.pi + h / 2.
    xymap = np.dstack((u, v)).astype(np.float32)
    
    return xymap

def matrix_multiplication():
    for i in range(100):
        np.random.random((1000, 1000)) @ np.random.random((1000, 1000))


if __name__ == "__main__":
    w = 3584
    h = int(w / 2)
    r = np.array([
    [0.61566148, -0.78369395, 0.08236955],
    [0.78801075, 0.61228882, -0.06435415],
    [0, 0.10452846, 0.9945219],])
    begin_time = time.time()
    # numpy_test(w, h, r)
    matrix_multiplication()
    print(time.time() - begin_time)