How can I speed up the compositing raycasting function?

Question

I'm currently working on a volume rendering project in python where I use a compositing ray casting function to produce an image, given a 3D volume consisting of voxels. The function (which I show below) works correctly, but has a very long runtime. Do you guys have tips on how to make this function faster? The code is Python 3.6.8 and uses various numpy arrays.

    def render_compositing(self, view_matrix: np.ndarray, volume: Volume, image_size: int, image: np.ndarray):
        # Clear the image
        self.clear_image()

        # U, V, View vectors. See documentation in parent's class
        u_vector = view_matrix[0:3]
        v_vector = view_matrix[4:7]
        view_vector = view_matrix[8:11]

        # Center of the image. Image is squared
        image_center = image_size / 2

        # Center of the volume (3-dimensional)
        volume_center = [volume.dim_x / 2, volume.dim_y / 2, volume.dim_z / 2]

        # Define a step size to make the loop faster
        step = 2 if self.interactive_mode else 1

        for i in range(0, image_size, step):
            for j in range(0, image_size, step):
                sum_color = TFColor(0, 0, 0, 0)
                for k in range(0, image_size, step):
                    # Get the voxel coordinate X
                    voxel_coordinate_x = u_vector[0] * (i - image_center) + v_vector[0] * (j - image_center) + \
                                         view_vector[0] * (k - image_center) + volume_center[0]

                    # Get the voxel coordinate Y
                    voxel_coordinate_y = u_vector[1] * (i - image_center) + v_vector[1] * (j - image_center) + \
                                         view_vector[1] * (k - image_center) + volume_center[1]

                    # Get the voxel coordinate Y
                    voxel_coordinate_z = u_vector[2] * (i - image_center) + v_vector[2] * (j - image_center) + \
                                         view_vector[2] * (k - image_center) + volume_center[2]

                    color = self.tfunc.get_color(
                        get_voxel(volume, voxel_coordinate_x, voxel_coordinate_y, voxel_coordinate_z))

                    sum_color.r = color.a * color.r + (1 - color.a) * sum_color.r
                    sum_color.g = color.a * color.g + (1 - color.a) * sum_color.g
                    sum_color.b = color.a * color.b + (1 - color.a) * sum_color.b
                    sum_color.a = color.a + (1 - color.a) * sum_color.a

                red = sum_color.r
                green = sum_color.g
                blue = sum_color.b
                alpha = sum_color.a

                # Compute the color value (0...255)
                red = math.floor(red * 255) if red < 255 else 255
                green = math.floor(green * 255) if green < 255 else 255
                blue = math.floor(blue * 255) if blue < 255 else 255
                alpha = math.floor(alpha * 255) if alpha < 255 else 255

                # Assign color to the pixel i, j
                image[(j * image_size + i) * 4] = red
                image[(j * image_size + i) * 4 + 1] = green
                image[(j * image_size + i) * 4 + 2] = blue
                image[(j * image_size + i) * 4 + 3] = alpha

Answer 1

I don't understand why you want to use python for this code. Isn't using a shader the better approach if you are concerned about speed?

Anyways here are few things that can be done in the current code.

1. voxel coordinates can be calculated using a numpy. you can make a 3 channel 2d image and compute the x,y,z coordinates for an entire slice(k) in a single shot.
2. Above step can be further optimized by storing an image of x,y,z coordinated of first slice(k=0) and a constant view_directionstep (step_size). Now every other slice can be simply calculated by (XYZ@k=0) + kstep_size.
3. Use early ray termination by thresholding alpha value to 0.999 or 0.99. This does not look like much but gives a lot of speed gain.