def example_curvilinear_grid(nC, exType):
if not isinstance(nC, list):
raise TypeError("nC must be a list containing the number of nodes")
if len(nC) != 2 and len(nC) != 3:
raise ValueError("nC must either two or three dimensions")
exType = exType.lower()
possibleTypes = ["rect", "rotate"]
if exType not in possibleTypes:
raise TypeError("Not a possible example type.")
if exType == "rect":
return list(
ndgrid([np.cumsum(np.r_[0, np.ones(nx) / nx]) for nx in nC],
vector=False))
elif exType == "rotate":
if len(nC) == 2:
X, Y = ndgrid([np.cumsum(np.r_[0, np.ones(nx) / nx]) for nx in nC],
vector=False)
amt = 0.5 - np.sqrt((X - 0.5)**2 + (Y - 0.5)**2)
amt[amt < 0] = 0
return [X + (-(Y - 0.5)) * amt, Y + (+(X - 0.5)) * amt]
elif len(nC) == 3:
X, Y, Z = ndgrid(
[np.cumsum(np.r_[0, np.ones(nx) / nx]) for nx in nC],
vector=False)
amt = 0.5 - np.sqrt((X - 0.5)**2 + (Y - 0.5)**2 + (Z - 0.5)**2)
amt[amt < 0] = 0
return [
X + (-(Y - 0.5)) * amt,
Y + (-(Z - 0.5)) * amt,
Z + (-(X - 0.5)) * amt,
]
def example_curvilinear_grid(nC, exType):
"""Creates and returns the gridded node locations for a curvilinear mesh.
Parameters
----------
nC : list of int
list of number of cells in each dimension. Must be length 2 or 3
exType : {"rect", "rotate", "sphere"}
String specifying the style of example curvilinear mesh.
Returns
-------
list of numpy.ndarray
List containing the gridded x, y (and z) node locations for the
curvilinear mesh.
"""
if not isinstance(nC, list):
raise TypeError("nC must be a list containing the number of nodes")
if len(nC) != 2 and len(nC) != 3:
raise ValueError("nC must either two or three dimensions")
exType = exType.lower()
possibleTypes = ["rect", "rotate", "sphere"]
if exType not in possibleTypes:
raise TypeError("Not a possible example type.")
if exType == "rect":
return list(
ndgrid([np.cumsum(np.r_[0, np.ones(nx) / nx]) for nx in nC],
vector=False))
elif exType == "sphere":
nodes = list(
ndgrid([np.cumsum(np.r_[0, np.ones(nx) / nx]) - 0.5 for nx in nC],
vector=False))
nodes = np.stack(nodes, axis=-1)
nodes = 2 * nodes
# L_inf distance to center
r0 = np.linalg.norm(nodes, ord=np.inf, axis=-1)
# L2 distance to center
r2 = np.linalg.norm(nodes, axis=-1)
r0[r0 == 0.0] = 1.0
r2[r2 == 0.0] = 1.0
scale = r0 / r2
nodes = nodes * scale[..., None]
nodes = np.transpose(nodes, (-1, *np.arange(len(nC))))
nodes = [node for node in nodes] # turn it into a list
return nodes
elif exType == "rotate":
if len(nC) == 2:
X, Y = ndgrid([np.cumsum(np.r_[0, np.ones(nx) / nx]) for nx in nC],
vector=False)
amt = 0.5 - np.sqrt((X - 0.5)**2 + (Y - 0.5)**2)
amt[amt < 0] = 0
return [X + (-(Y - 0.5)) * amt, Y + (+(X - 0.5)) * amt]
elif len(nC) == 3:
X, Y, Z = ndgrid(
[np.cumsum(np.r_[0, np.ones(nx) / nx]) for nx in nC],
vector=False)
amt = 0.5 - np.sqrt((X - 0.5)**2 + (Y - 0.5)**2 + (Z - 0.5)**2)
amt[amt < 0] = 0
return [
X + (-(Y - 0.5)) * amt,
Y + (-(Z - 0.5)) * amt,
Z + (-(X - 0.5)) * amt,
]
def index_cube(nodes, grid_shape, n=None):
"""Returns the index of nodes on a tensor (or curvilinear) mesh.
For 2D tensor meshes, each cell is defined by nodes
*A, B, C* and *D*. And for 3D tensor meshes, each cell
is defined by nodes *A* through *H* (see below). *index_cube*
outputs the indices for the specified node(s) for all
cells in the mesh.
TWO DIMENSIONS::
node(i,j+1) node(i+i,j+1)
B -------------- C
| |
| cell(i,j) |
| I |
| |
A -------------- D
node(i,j) node(i+1,j)
THREE DIMENSIONS::
node(i,j+1,k+1) node(i+1,j+1,k+1)
F ---------------- G
/| / |
/ | / |
/ | / |
node(i,j,k+1) node(i+1,j,k+1)
E --------------- H |
| B -----------|---- C
| / cell(i,j,k) | /
| / I | /
| / | /
A --------------- D
node(i,j,k) node(i+1,j,k)
Parameters
----------
nodes : str
String specifying which nodes to return. For 2D meshes,
*nodes* must be a string containing combinations of the characters 'A', 'B',
'C', or 'D'. For 3D meshes, *nodes* can also be 'E', 'F', 'G', or 'H'. Note that
order is preserved. E.g. if we want to return the C, D and A node indices in
that particular order, we input *nodes* = 'CDA'.
grid_shape : list of int
Number of nodes along the i,j,k directions; e.g. [ni,nj,nk]
nc : list of int
Number of cells along the i,j,k directions; e.g. [nci,ncj,nck]
Returns
-------
index : tuple of numpy.ndarray
Each entry of the tuple is a 1D :class:`numpy.ndarray` containing the indices of
the nodes specified in the input *nodes* in the order asked;
e.g. if *nodes* = 'DCBA', the tuple returned is ordered (D,C,B,A).
Examples
--------
Here, we construct a small 2D tensor mesh
(works for a curvilinear mesh as well) and use *index_cube*
to find the indices of the 'A' and 'C' nodes. We then
plot the mesh, as well as the 'A' and 'C' node locations.
>>> from discretize import TensorMesh
>>> from discretize.utils import index_cube
>>> from matplotlib import pyplot as plt
>>> import numpy as np
Create a simple tensor mesh.
>>> n_cells = 5
>>> h = 2*np.ones(n_cells)
>>> mesh = TensorMesh([h, h], x0='00')
Get indices of 'A' and 'C' nodes for all cells.
>>> A, C = index_cube('AC', [n_cells+1, n_cells+1])
Plot mesh and the locations of the A and C nodes
.. collapse:: Expand to show scripting for plot
>>> fig1 = plt.figure(figsize=(5, 5))
>>> ax1 = fig1.add_axes([0.1, 0.1, 0.8, 0.8])
>>> mesh.plot_grid(ax=ax1)
>>> ax1.scatter(mesh.nodes[A, 0], mesh.nodes[A, 1], 100, 'r', marker='^')
>>> ax1.scatter(mesh.nodes[C, 0], mesh.nodes[C, 1], 100, 'g', marker='v')
>>> ax1.set_title('A nodes (red) and C nodes (green)')
>>> plt.show()
"""
if not isinstance(nodes, str):
raise TypeError("Nodes must be a str variable: e.g. 'ABCD'")
nodes = nodes.upper()
try:
dim = len(grid_shape)
if n is None:
n = tuple(x - 1 for x in grid_shape)
except TypeError:
return TypeError("grid_shape must be iterable")
# Make sure that we choose from the possible nodes.
possibleNodes = "ABCD" if dim == 2 else "ABCDEFGH"
for node in nodes:
if node not in possibleNodes:
raise ValueError(
"Nodes must be chosen from: '{0!s}'".format(possibleNodes))
if dim == 2:
ij = ndgrid(np.arange(n[0]), np.arange(n[1]))
i, j = ij[:, 0], ij[:, 1]
elif dim == 3:
ijk = ndgrid(np.arange(n[0]), np.arange(n[1]), np.arange(n[2]))
i, j, k = ijk[:, 0], ijk[:, 1], ijk[:, 2]
else:
raise Exception("Only 2 and 3 dimensions supported.")
nodeMap = {
"A": [0, 0, 0],
"B": [0, 1, 0],
"C": [1, 1, 0],
"D": [1, 0, 0],
"E": [0, 0, 1],
"F": [0, 1, 1],
"G": [1, 1, 1],
"H": [1, 0, 1],
}
out = ()
for node in nodes:
shift = nodeMap[node]
if dim == 2:
out += (sub2ind(grid_shape, np.c_[i + shift[0],
j + shift[1]]).flatten(), )
elif dim == 3:
out += (sub2ind(grid_shape, np.c_[i + shift[0], j + shift[1],
k + shift[2]]).flatten(), )
return out
def index_cube(nodes, grid_size, n=None):
"""
Returns the index of nodes on the mesh.
Input:
nodes - string of which nodes to return. e.g. 'ABCD'
grid_size - size of the nodal grid
n - number of nodes each i,j,k direction: [ni,nj,nk]
Output:
index - index in the order asked e.g. 'ABCD' --> (A,B,C,D)
TWO DIMENSIONS::
node(i,j) node(i,j+1)
A -------------- B
| |
| cell(i,j) |
| I |
| |
D -------------- C
node(i+1,j) node(i+1,j+1)
THREE DIMENSIONS::
node(i,j,k+1) node(i,j+1,k+1)
E --------------- F
/| / |
/ | / |
/ | / |
node(i,j,k) node(i,j+1,k)
A -------------- B |
| H ----------|---- G
| /cell(i,j) | /
| / I | /
| / | /
D -------------- C
node(i+1,j,k) node(i+1,j+1,k)
"""
if not isinstance(nodes, str):
raise TypeError("Nodes must be a str variable: e.g. 'ABCD'")
nodes = nodes.upper()
try:
dim = len(grid_size)
if n is None:
n = tuple(x - 1 for x in grid_size)
except TypeError:
return TypeError("grid_size must be iterable")
# Make sure that we choose from the possible nodes.
possibleNodes = "ABCD" if dim == 2 else "ABCDEFGH"
for node in nodes:
if node not in possibleNodes:
raise ValueError(
"Nodes must be chosen from: '{0!s}'".format(possibleNodes))
if dim == 2:
ij = ndgrid(np.arange(n[0]), np.arange(n[1]))
i, j = ij[:, 0], ij[:, 1]
elif dim == 3:
ijk = ndgrid(np.arange(n[0]), np.arange(n[1]), np.arange(n[2]))
i, j, k = ijk[:, 0], ijk[:, 1], ijk[:, 2]
else:
raise Exception("Only 2 and 3 dimensions supported.")
nodeMap = {
"A": [0, 0, 0],
"B": [0, 1, 0],
"C": [1, 1, 0],
"D": [1, 0, 0],
"E": [0, 0, 1],
"F": [0, 1, 1],
"G": [1, 1, 1],
"H": [1, 0, 1],
}
out = ()
for node in nodes:
shift = nodeMap[node]
if dim == 2:
out += (sub2ind(grid_size, np.c_[i + shift[0],
j + shift[1]]).flatten(), )
elif dim == 3:
out += (sub2ind(grid_size, np.c_[i + shift[0], j + shift[1],
k + shift[2]]).flatten(), )
return out