Splitting Python MeshGrid Into Cells

Question

Problem Statement

Need to split an N-Dimensional MeshGrid into "cubes":

Ex) 2-D Case:

(-1,1) |(0,1) |(1,1)

(-1,0) |(0,0) |(1,0)

(-1,-1)|(0,-1)|(1,-1)

There will be 4 cells, each with 2^D points:

I want to be able to process the mesh, placing the coordinate points of each cell into a container to do further processing.

Cells =   [{(-1,1) (0,1)(-1,0),(0,0)},

          {(0,1),(1,1),(0,0),(1,0)},

          {(-1,0),(0,0)(-1,-1),(0,-1)}

          {(0,0),(1,0)(0,-1),(1,-1)}]

I use the following to generate the mesh for an arbitrary dimension d:

grid = [np.linspace(-1.0 , 1.0, num = K+1) for i in range(d)]
res_to_unpack = np.meshgrid(*grid,indexing = 'ij')

Which has output:

[array([[-1., -1., -1.],
   [ 0.,  0.,  0.],
   [ 1.,  1.,  1.]]), array([[-1.,  0.,  1.],
   [-1.,  0.,  1.],
   [-1.,  0.,  1.]])]

So I want to be able to generate the above cells container for a given D dimensional mesh grid. Split on a given K that is a power of 2.

I need this container so for each cell, I need to reference all 2^D points associated and calculate the distance from an origin.

Edit For Clarification

K should partition the grid into K ** D number of cells, with (K+1) ** D points. Each Cell should have 2 ** D number of points. Each "cell" will have volume (2/K)^D.

So for K = 4, D = 2

Cells = [ {(-1,1),(-0.5,1),(-1,0.5),(-0.5,0.5)},
          {(-0.5,1),(-0.5,0.5)(0.0,1.0),(0,0.5)},
            ...
          {(0.0,-0.5),(0.5,-0.5),(0.0,-1.0),(0.5,-1.0)},
          {(0.5,-1.0),(0.5,-1.0),(1.0,-0.5),(1.0,-1.0)}]

This is output for TopLeft, TopLeft + Right Over, Bottom Left, Bottom Left + Over Left. There will we 16 cells in this set, each with four coordinates each. For increasing K, say K = 8. There will be 64 cells, each with four points.

Answer 1

This should give you what you need:

from itertools import product
import numpy as np

def splitcubes(K, d):
    coords = [np.linspace(-1.0 , 1.0, num=K + 1) for i in range(d)]
    grid = np.stack(np.meshgrid(*coords)).T

    ks = list(range(1, K))
    for slices in product(*([[slice(b,e) for b,e in zip([None] + ks, [k+1 for k in ks] + [None])]]*d)):
        yield grid[slices]

def cubesets(K, d):
    if (K & (K - 1)) or K < 2:
        raise ValueError('K must be a positive power of 2. K: %s' % K)

    return [set(tuple(p.tolist()) for p in c.reshape(-1, d)) for c in splitcubes(K, d)]

Demonstration of 2D case

Here's a little demonstration of the 2D case:

import matplotlib.pyplot as plt

def assemblecube(c, spread=.03):
    c = np.array(list(c))
    c = c[np.lexsort(c.T[::-1])]

    d = int(np.log2(c.size))
    for i in range(d):
        c[2**i:2**i + 2] = c[2**i + 1:2**i - 1:-1]

    # get the point farthest from the origin
    sp = c[np.argmax((c**2).sum(axis=1)**.5)]
    # shift all points a small distance towards that farthest point
    c += sp * .1 #np.copysign(np.ones(sp.size)*spread, sp)

    # create several different orderings of the same points so that matplotlib will draw a closed shape
    return [(np.roll(c, i, axis=1) - (np.roll(c, i, axis=1)[0] - c[0])[None,:]).T for i in range(d)]

fig = plt.figure(figsize=(6,6))
ax = fig.gca()

for i,c in enumerate(cubesets(4, 2)):
    for cdata in assemblecube(c):
        p = ax.plot(*cdata, c='C%d' % (i % 9))

ax.set_aspect('equal', 'box')
fig.show()

Output:

The cubes have been shifted slightly apart for visualization purposes (so they don't overlap and cover each other up).

Demonstration of 3D case

Here's the same thing for the 3D case:

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

fig = plt.figure(figsize=(6,6))
ax = fig.add_subplot(111, projection='3d')

for i,c in enumerate(cubesets(2,3)):
    for cdata in assemblecube(c, spread=.05):
        ax.plot(*cdata, c=('C%d' % (i % 9)))

plt.gcf().gca().set_aspect('equal', 'box')
plt.show()

Output:

Demonstration for K=4

Here's the output for the same 2D and 3D demonstrations as above, but with K=4: