utils.py

import numpy as np

def generate_gridpoints(shape: tuple, min_corner: np.ndarray, max_corner: np.ndarray, endpoint=True) -> np.ndarray:
	X = np.linspace(min_corner[0], max_corner[0], shape[-1], endpoint=endpoint)
	Y = np.linspace(min_corner[1], max_corner[1], shape[-2], endpoint=endpoint)

	if len(shape) == 2:
		X, Y = np.meshgrid(X, Y, indexing='xy')
		points = np.stack((X.flatten(), Y.flatten()), axis=1)
	else:
		Z = np.linspace(min_corner[2], max_corner[2], shape[-3])
		X, Y, Z = np.meshgrid(X, Y, Z, indexing='ij')
		X = X.transpose(2, 1, 0)
		Y = Y.transpose(2, 1, 0)
		Z = Z.transpose(2, 1, 0)
		points = np.stack((X.flatten(), Y.flatten(), Z.flatten()), axis=1)

	return points

def compute_gradients(grid: np.ndarray, dx: float) -> np.ndarray:
	dim = len(grid.shape)
	grads = np.zeros(grid.shape + (dim,))

	for i in range(dim):
		grads[..., i] = np.gradient(grid, dx, axis=dim-i-1)

	return grads