1"""The regular grid that covers the bounding box of a plate, restricted to the
2plate's outline by a mask. Arrays over the grid have shape ``(ny, nx)``, indexed
3``[j, i]`` with ``i`` along x. Node coordinates are relative to the center of the
4bounding box.
6The edge is a staircase of whole cells: a cell belongs to the plate when its
7center lies inside the outline, and a node when it is a corner of such a cell.
8Deciding membership by cell centers keeps the covered area unbiased. A hole is
9handled by the same rule.
10"""
12from __future__ import annotations
14from dataclasses import dataclass
15from functools import cached_property
16from typing import Callable
18import numpy as np
20InsideFn = Callable[[np.ndarray, np.ndarray, np.ndarray, np.ndarray], np.ndarray]
21"""``inside(x, y, r, theta) -> bool array``, evaluated on whole grid matrices."""
25class Grid:
26 nx: int
27 ny: int
28 hx: float
29 hy: float
30 mask: np.ndarray
31 """``(ny, nx)`` bool: True for the nodes that belong to the plate."""
34 def n(self) -> int:
35 """Number of grid nodes, inside or not."""
36 return self.nx * self.ny
39 def shape(self) -> tuple[int, int]:
40 return (self.ny, self.nx)
43 def count(self) -> int:
44 """Number of nodes in the plate."""
45 return int(self.mask.sum())
48 def x(self) -> np.ndarray:
49 """Node x coordinates ``(nx,)``, relative to the center of the bounding box."""
50 return (np.arange(self.nx) - (self.nx - 1) / 2) * self.hx
53 def y(self) -> np.ndarray:
54 return (np.arange(self.ny) - (self.ny - 1) / 2) * self.hy
57 def X(self) -> np.ndarray:
58 """Node x coordinates on the grid, ``(ny, nx)``."""
59 return np.broadcast_to(self.x[None, :], self.shape).copy()
62 def Y(self) -> np.ndarray:
63 return np.broadcast_to(self.y[:, None], self.shape).copy()
66 def cell_mask(self) -> np.ndarray:
67 """``(ny - 1, nx - 1)`` bool: cells whose four corners belong to the plate."""
68 m = self.mask
69 return m[:-1, :-1] & m[:-1, 1:] & m[1:, 1:] & m[1:, :-1]
72 def node_weight(self) -> np.ndarray:
73 """Quadrature and lumped-mass weight of each node: the fraction of its
74 four cells that belong to the plate (1 inside, 1/2 on an edge, 1/4 at a
75 corner of a rectangle, 0 outside)."""
76 return corners_to_nodes(self.cell_mask.astype(float)) / 4
79 def node_area(self) -> np.ndarray:
80 """Area attributed to each node, hx hy times its weight [m^2]."""
81 return self.hx * self.hy * self.node_weight
84 def area(self) -> float:
85 """Area of the staircase plate [m^2]."""
86 return float(self.node_area.sum())
88 def cell_centers(self) -> tuple[np.ndarray, np.ndarray]:
89 """Coordinates of the cell centers, two ``(ny - 1, nx - 1)`` arrays."""
90 return np.meshgrid(self.x[:-1] + self.hx / 2, self.y[:-1] + self.hy / 2)
92 def neighbor_masks(self) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
93 """Whether the left, right, lower (j - 1) and upper (j + 1) neighbor of each node is in the plate."""
94 m = self.mask
95 left = np.zeros_like(m)
96 left[:, 1:] = m[:, :-1]
97 right = np.zeros_like(m)
98 right[:, :-1] = m[:, 1:]
99 down = np.zeros_like(m)
100 down[1:, :] = m[:-1, :]
101 up = np.zeros_like(m)
102 up[:-1, :] = m[1:, :]
103 return left & m, right & m, down & m, up & m
106def corners_to_nodes(c: np.ndarray) -> np.ndarray:
107 """Sums a per-cell array ``(ny - 1, nx - 1)`` onto the four corners of each cell, giving ``(ny, nx)``."""
108 out = np.zeros((c.shape[0] + 1, c.shape[1] + 1), dtype=c.dtype)
109 out[:-1, :-1] += c
110 out[:-1, 1:] += c
111 out[1:, 1:] += c
112 out[1:, :-1] += c
113 return out
116def make_grid(Lx: float, Ly: float, spacing: float, inside: InsideFn | None = None) -> Grid:
117 """The grid covering the box ``Lx`` by ``Ly`` with about the requested
118 spacing (at least 6 nodes per side), restricted to ``inside`` when given."""
119 nx = max(6, round(Lx / spacing) + 1)
120 ny = max(6, round(Ly / spacing) + 1)
121 g = Grid(nx, ny, Lx / (nx - 1), Ly / (ny - 1), np.ones((ny, nx), dtype=bool))
122 if inside is None:
123 return g
124 XC, YC = g.cell_centers()
125 cells = np.asarray(inside(XC, YC, np.hypot(XC, YC), np.arctan2(YC, XC)), dtype=bool)
126 if cells.shape != (ny - 1, nx - 1):
127 raise ValueError(f"inside returned shape {cells.shape} for a {ny - 1} x {nx - 1} array of cells")
128 return with_cell_mask(g, cells)
131def with_cell_mask(g: Grid, cells: np.ndarray) -> Grid:
132 """Restricts a grid to the cells flagged in ``cells`` ``(ny - 1, nx - 1)``:
133 a node belongs to the plate when it is a corner of an included cell. Any
134 node left with no cell of its own would carry no mass and is removed."""
135 mask = corners_to_nodes(cells.astype(np.int8)) > 0
136 while True:
137 gm = Grid(g.nx, g.ny, g.hx, g.hy, mask)
138 bad = mask & (gm.node_weight == 0)
139 if not bad.any():
140 return gm
141 mask = mask & ~bad