geometry.py

import numpy as np


def box_surface_grid(size, n_steps):
    size = np.asarray(size)
    half_size = 0.5 * size

    #     +-------+
    #     |       |
    # +---+-------+---+-------+     -
    # |   |       |   |       |     |
    # |   |       |   |       |     b
    # |   |       |   |       |     |
    # +---+-------+---+-------+     -
    #     |       |
    #     +-------+
    #
    # |-c-|---a---|-c-|---a---|

    x_axis_scaling = half_size[0] + size[2]
    y_axis_scaling = half_size[1] + size[2]
    x_coords = np.linspace(-x_axis_scaling, x_axis_scaling + size[0], n_steps)
    y_coords = np.linspace(-y_axis_scaling, y_axis_scaling, n_steps)
    x_grid, y_grid = np.meshgrid(x_coords, y_coords)

    x = np.empty_like(x_grid)
    bottom = x_grid > x_axis_scaling
    not_bottom = np.logical_not(bottom)
    x[not_bottom] = np.clip(x_grid[not_bottom], -half_size[0], half_size[0])
    x[bottom] = -(x_grid[bottom] - x_axis_scaling) + half_size[0]
    y = np.clip(y_grid, -half_size[1], half_size[1])
    z_coords_x = half_size[2] - (np.abs(x_grid) - half_size[0])
    z_coords_y = half_size[2] - (np.abs(y_grid) - half_size[1])
    z = np.clip(np.min((z_coords_x, z_coords_y), axis=0), -half_size[2], half_size[2])

    return x, y, z

plot.py

import matplotlib.pyplot as plt
import geometry as pg
import pytransform3d.plot_utils as ppu


plt.figure()
ax = ppu.make_3d_axis(1)
x, y, z = pg.box_surface_grid([0.1, 1, 0.5], 50)
ax.plot_surface(x, y, z)
plt.show()