CODEWITHSUNDEEP
pythonnumpyoptimizationmemoryplane
swiss_knight
swiss_knight

Reputation: 7901

Build a 2D array representing a 3D plane (storing its Z-values) as defined by 3 points and the desired size of the array

Short

I need to find an optimized way to build a 2D numpy array representing a 3D plane (namely storing its Z-values) given three points (to mathematically define the plane) and the desired array size.

Details

Given three points p0, p1, p2, and a size variable:

import numpy as np
import pylab as plt

p0 = [0.2, -0.4, -0.2]
p1 = [-0.6, 0.1, -0.8]
p2 = [-0.1, 0.3, -0.6]
size = (12000, 17000)

The following function shows the idea of how to build a plane out of these 3 points:

def plane(p0, p1, p2, size):
    p0 = np.array(p0)
    p1 = np.array(p1)
    p2 = np.array(p2)
    ux, uy, uz = [p1[0] - p0[0], p1[1] - p0[1], p1[2] - p0[2]]
    vx, vy, vz = [p2[0] - p0[0], p2[1] - p0[1], p2[2] - p0[2]]
    uv = [uy * vz - uz * vy,
          uz * vx - ux * vz,
          ux * vy - uy * vx]
    norm = np.array(uv)    
    d = -np.array(p0).dot(norm)    
    xx, yy = np.meshgrid(range(size[1]), range(size[0]))    
    z = np.array((-norm[0] * xx - norm[1] * yy - d) * 1. / norm[2])
    return(z)

The resulting z looks like this:

resulting image

This implementation eats my memory in a few seconds:

memory usage

Question

How would you optimize this to achieve the exact same result with the lowest memory usage as possible (being fast is also not optional...)?

Note: size may be 10x larger in both dimensions -> which obviously crashes my Python on a mainstream 2016 laptop.

(Using the most common scientific libraries is perfectly fine)

I'm using Python 3.6.9 (default, Jan 26 2021, 15:33:00) [GCC 8.4.0] on Ubuntu linux 18.04.5.

Upvotes: 4

Views: 611

Answers (1)

Mad Physicist
Mad Physicist

Reputation: 114440

Up to the computation of d, your arrays are small. Everything after that is creating a separate 1.5GiB (12k * 17k * 8bytes/float) temp-array for each part of the operation. You can get rid of most of these arrays entirely, and do the rest in-place. For example:

xx, yy = np.meshgrid(...)

You really don't need the meshgrid: broadcasting on np.arange(size[0])[:, None] and np.arange(size[1]) will make all the arrays you need. That reduces your size from ~3GiB to ~230KiB for that portion of the data:

xx = np.arange(size[1], dtype=float)
yy = np.arange(size[0], dtype=float)[:, None]

The expression z = ... has approximately seven temporary arrays in it, each one 1.5GiB in size:

  1. t1 = -norm[0] * xx
  2. t2 = norm[1] * yy
  3. t3 = t1 - t2
  4. t4 = t3 - d
  5. t5 = t4 * 1
  6. t6 = t5 / norm[2]
  7. t7 = np.array(t6)

Some of these arrays may be garbage collected during the computation, but they all need to get allocated at some point. You can get rid of all of them using broadcasting and in-place operators or the out parameter to the appropriate ufuncs:

So:

d /= norm[2]
norm /= norm[2]
xx *= -norm[0]
yy *= -norm[1]
yy -= d
z = xx + yy

Even in the compact part of the code, you are doing a bunch of unnecessary operations. For example, you call np.array(p0) twice. Here is a nicer version of your function:

def plane2(p0, p1, p2, size):
    u = np.subtract(p1, p0)
    v = np.subtract(p2, p0)
    uxv = np.cross(u, v)
    d = -np.dot(p0, uxv)
    d /= uxv[2]
    uxv /= uxv[2]
    xx = np.arange(size[1], dtype=float)
    xx *= -uxv[0]
    yy = np.arange(size[0], dtype=float)[:, None]
    yy *= -uxv[1]
    yy -= d
    return xx + yy

For sizes 10x smaller than you proposed:

%timeit plane(p0, p1, p2, size)
70.1 ms ± 329 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit plane2(p0, p1, p2, size)
8.29 ms ± 302 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

When I increase the size to the one in your question, your implementation crashes while mine completes in a reasonable amount of time:

%timeit plane(p0, p1, p2, size)
...
MemoryError: Unable to allocate 1.52 GiB for an array with shape (12000, 17000) and data type float64

%timeit plane2(p0, p1, p2, size)
819 ms ± 37.6 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Depending on how much RAM you have, and what the settings for free memory in the python process are, you may have trouble increasing your array size by 100x to 152GiB, not matter how efficiently you perform your operations.

Upvotes: 3

Related Questions