[FIXED] Optimisation of spherical to cartesian notation Python

Issue

I am preprocessing data with a relatively small dataset (4000 instances) but the code I have written takes 9 seconds per instance on my laptop resulting in a 10-hour run time. The code is here below and if anyone could help optimize it so I can use a larger dataset I would really appreciate it.

This code is used to process Lidar data. The data Variable is of shape (4000,N,5), but truncated to (4000,N,3) in this code to represent the xyz data of each point in each scan. Some pre-processing was done to create data_2d which is of shape (4000,N,3) representing the polar coordinates of each point, with the first column being the elevation angle, the second the azimuth angle and the third the distance (the latter used in another part of my process). The code bins the average xyz point based on the theta and phi scale (elevation and azimuth angles) by finding the index in data_2d/ chosen_scan of points with elevation and azimuth angles within each range and using it to find the average xyz in data/chosen_scan_data.

new_2d_scans_3d = np.zeros((4000,80,640,3)) 

min_max_theta =[np.min(data_2d[0][:,0]),np.max(data_2d[0][:,0])]

for i in data_2d:
    current_min = np.min(i[:,0])
    current_max = np.max(i[:,0])
    if  current_min<min_max_theta [0]:
        min_max_theta [0] = current_min
    if  current_max< min_max_theta [1]:
        min_max_theta [1] = current_max

# min_max_theta = [1.3845561488601013, 2.1061203926592453]
min_max_phi = [-3.1414179913970317, 3.1415925661738253]

theta_scale = np.linspace(min_max_theta[0], min_max_theta[1],80)
phi_scale = np.linspace(min_max_phi[0], min_max_phi[1],640)

for k in range(data_2d.shape[0]):
    st = time.time()
    chosen_scan_data = data[k][:,:3]
    chosen_scan = data_2d[k]
    current_theta = theta_scale[0]
    current_phi = phi_scale[0]
    for i in range(1,640):
        for j in range(1,80):
            idx = np.where((chosen_scan[:,0] > current_theta) & (chosen_scan[:,0] < theta_scale[j]) & (chosen_scan[:,1] > current_phi) & (chosen_scan[:,1] <phi_scale[i]))
            if idx[0].size != 0:
                avg_point =np.average(chosen_scan_data[idx[0]],axis=0)
                new_2d_scans_3d[k][j-1][i-1] = avg_point
            else:
                new_2d_scans_3d[k][j-1][i-1] = np.array([0,0,0])
            current_theta = theta_scale[j]
        current_phi = phi_scale[i]
    et = time.time()
    elapsed_time = et - st
    print('Executing Image ', k ,'/4000: ', elapsed_time, 'seconds')

Some random exemplar data for data and data_2d can be obtained as follows:

data = np.random.rand(4000,35000,3)
data_2d = []
for scan in data:
    scan_2d = []
    for point in scan:
        phi = math.atan2(point[1],point[0])
        theta = math.acos(point[2]/math.sqrt(point[0]**2 + point[1]**2 + point[2]**2))
        d =math.sqrt(point[0]**2+point[1]**2)
        scan_2d.append([theta,phi,d])
    data_2d.append(scan_2d)
data_2d = np.array(data_2d)

Solution

Assuming I understood what you wanted (see my comment), here is some piece of code that I believe does the job. I wrote comments in the code, I hope it is enough.

I invite you to read the documentation of the numpy functions "unique", "digitize" and "nonzero".

import numpy as np
import time

### Generate sample data
rng = np.random.default_rng(seed=12345678)
data = 2 * rng.random((4000, 1000, 3)) - 1
d = np.sqrt(np.sum(data**2, axis=2)) # vector norm
data_2d = np.array( # theta, phi, d
    [np.arccos(data[:, :, 2]/d), np.arctan2(data[:, :, 1], data[:, :, 0]), d]
    )
data_2d = np.moveaxis(data_2d, 0, -1)


### Proposed solution
ts0 = time.time_ns()/10e8

## Compute theta bins
nb_bin_theta = 80 # number of bins
theta_min, theta_max = 0, np.pi # extremal values
theta_bins = np.linspace(theta_min, theta_max, nb_bin_theta)
# Compute the bin to which each theta value belongs
theta_idx = np.digitize(data_2d[:, :, 0], theta_bins, right=False) # care with the edge cases !
# Get the unique bin indices, and the mapping to each index
# Note that np.unique works on the flattened array
theta_bin_idx, theta_bin_inv = np.unique(theta_idx, return_inverse=True)

## Do the same with phi
nb_bin_phi = 640
phi_min, phi_max = -np.pi, np.pi
phi_bins = np.linspace(phi_min, phi_max, nb_bin_phi)
phi_idx = np.digitize(data_2d[:, :, 1], phi_bins, right=False)
phi_bin_idx, phi_bin_inv = np.unique(phi_idx, return_inverse=True)

# Since np.unique works on the flattened array, we flatten the data array
# while retaining the cartesian coordinate structure
flat_data = data.reshape((-1, 3))

# Initialize the output, by default all values are zero
output = np.zeros((nb_bin_theta, nb_bin_phi, 3), dtype=float) # (bin_theta, bin_phi, xyz-coordinates)

# Iterate over the bins that contain at least one value
# we know the output will be zero elsewhere
for i, i_theta in enumerate(theta_bin_idx):
    for j, j_phi in enumerate(phi_bin_idx):
        # get the (flattened) data indices such that both theta and phi
        # belong to the corresponding theta and phi bins
        bin_idx, = np.nonzero(np.logical_and(
            theta_bin_inv == i, phi_bin_inv == j))
        if bin_idx.size > 0: # if the retrieved array of data indices is not empty
            # compute the average
            output[i_theta, j_phi, :] = np.average(
                flat_data[bin_idx, :], axis=0)
            # if bin_idx were to be empty, taking the average would return a nan!

ts1 = time.time_ns()/10e8
print("computation in", ts1-ts0, "s")

The code runs in a few seconds on my computer for 4000x1000 points. Note that I was not careful in managing the edge cases in the bins, so it might require some additional work.

Answered By - johnc