Best data structure for 3D data?

Question

I have some data, which is a function of two variables, let's say a and b.

f(a,b)

In maths terms, we can think of f as a surface viewed in 3D. The problem is, what is a good data structure to store this data in? I need to plot f as a function of a for constant b and f as a function of b for constant a. Currently I'm using a dict of arrays, something like this:

  f['a1'] = [b1,b2,b3]

but then if I now want to plot f as a function of a with b constant I have to re-make another dict by hand:

  f['b1'] = [a1,a2,a3]

which is proving extremely cumbersome and makes my code unreadable. Is there a nice way to store such 3D data in some numpy data structure or using built in python data structures?

Answer 1

f['a1'] = [b1, b2, b3]

I do not see how this relates to your initial description. You said f was a function of two variables... I'm assuming you misleadingly labelled these 'b'.

Example with regular Python lists

You can easily store a two-dimensional array as a list of lists (see Iliffe vector). That would give you the following:

_f_values = [
        [1, 4, 3],
        [2, 8, 6],
        [0, 7, 5],
    ]
def f(a, b):
    return _f_values[a][b]

Using numpy's ndarray

numpy has a specialized type for multidimensional homogeneous arrays: ndarray. This will definitely be faster. You can access whole rows or columns using :

_f_array = numpy.array([
        [1, 4, 3],
        [2, 8, 6],
        [0, 7, 5],
    ])
def f(a, b):
    return _f_array[a, b]
def f_row(a):
    return _f_array[a, :]
def f_col(b):
    return _f_array[:, b]

Answer 2

If your data is really a surface in 3D space, the natural way of storing would be to have two vectors of unique values of your two independent variables, a and b, and a 2D array with the values for the calls to f() with the Cartesian product of your two vectors. As a dummy example:

def f(a, b) :
    return a + b

a = np.array([1, 3, 9, 15])
b = np.array([6, 18, 32])

# if your function f is vectorized
f_ab = f(a[:, np.newaxis], b)
# if it is not
f_ab = np.empty((len(a), len(b)), dtype=np.array(f(a[0], b[0])).dtype)
for i, a_ in enumerate(a):
    for j, b_ in enumerate(b):
        f_ab[i, j] = f(a_, b_)

You can now get a slice of your data corresponding to a specific value of a, or to a specific value of b as:

>>> f_ab[a == 9, :]
array([[15, 27, 41]])

>>> f_ab[:, b==32]
array([[33],
       [35],
       [41],
       [47]])