Compare commits

..

24 Commits

Author SHA1 Message Date
c598978543 bump version to v1.2 2024-07-31 22:51:34 -07:00
c7ad0f0e37 comment out unused pytest import 2024-07-31 22:51:19 -07:00
4218f529ea Copy pixel edge coordinates 2024-07-31 22:51:07 -07:00
c95341c9b9 be clearer about floats 2024-07-31 22:50:34 -07:00
045b0c0228 enable numpy 2.0 2024-07-29 02:12:37 -07:00
e1e6134ec0 use asarray (since copy=False meaning changes in numpy 2.0) 2024-07-29 02:11:04 -07:00
8e7e0edb1f add ruff and mypy configs 2024-07-29 01:57:57 -07:00
e5fdc3ce23 drop unused imports 2024-07-29 01:57:48 -07:00
646911c4b5 type annotation improvements 2024-07-29 01:57:39 -07:00
e256f56f2b fix handling of 3d polys 2024-07-29 01:47:12 -07:00
c32d94ed85 fix typos in arg names in example 2024-07-29 01:38:42 -07:00
8c33a39c02 refactor to avoid class-scoped imports 2024-07-29 01:37:58 -07:00
f84a75f35a comment unused import 2024-07-29 00:50:08 -07:00
5a20339eab del axes3d to clarify it's unused on purpose 2024-07-29 00:49:59 -07:00
e29c0901bd use strict zip 2024-07-29 00:46:21 -07:00
a15e4bc05e repeat re-exported names 2024-07-29 00:44:26 -07:00
9ab97e763c bump min python version to 3.11 due to Self type 2024-07-18 00:20:18 -07:00
d44e02e2f7 return figure and axes after plotting 2024-07-18 00:17:58 -07:00
3e4e6eead3 flake8 fixup 2024-07-18 00:17:45 -07:00
a94c2cae67 type hint modernization 2024-07-18 00:17:20 -07:00
73d07bbfe0 disaambiguate some variables for typing purposes 2022-10-18 19:44:47 -07:00
ec5c77e018 typing and formatting updates 2022-10-18 19:44:30 -07:00
7d3b2272bc bump version to v1.1 2022-08-29 13:11:12 -07:00
e1303b8a5c move to hatch-based builds 2022-08-29 13:07:46 -07:00
16 changed files with 974 additions and 837 deletions

29
.flake8 Normal file
View File

@ -0,0 +1,29 @@
[flake8]
ignore =
# E501 line too long
E501,
# W391 newlines at EOF
W391,
# E241 multiple spaces after comma
E241,
# E302 expected 2 newlines
E302,
# W503 line break before binary operator (to be deprecated)
W503,
# E265 block comment should start with '# '
E265,
# E123 closing bracket does not match indentation of opening bracket's line
E123,
# E124 closing bracket does not match visual indentation
E124,
# E221 multiple spaces before operator
E221,
# E201 whitespace after '['
E201,
# E741 ambiguous variable name 'I'
E741,
per-file-ignores =
# F401 import without use
*/__init__.py: F401,

View File

@ -1,2 +0,0 @@
include README.md
include LICENSE.md

View File

@ -14,7 +14,7 @@ the coordinates of the boundary points along each axis).
## Installation
Requirements:
* python 3 (written and tested with 3.9)
* python >3.11 (written and tested with 3.12)
* numpy
* [float_raster](https://mpxd.net/code/jan/float_raster)
* matplotlib (optional, used for visualization functions)

1
gridlock/LICENSE.md Symbolic link
View File

@ -0,0 +1 @@
../LICENSE.md

1
gridlock/README.md Symbolic link
View File

@ -0,0 +1 @@
../README.md

View File

@ -1,4 +0,0 @@
""" VERSION defintion. THIS FILE IS MANUALLY PARSED BY setup.py and REQUIRES A SPECIFIC FORMAT """
__version__ = '''
1.0
'''.strip()

View File

@ -15,10 +15,9 @@ Dependencies:
- mpl_toolkits.mplot3d [Grid.visualize_isosurface()]
- skimage [Grid.visualize_isosurface()]
"""
from .error import GridError
from .grid import Grid
from .error import GridError as GridError
from .grid import Grid as Grid
__author__ = 'Jan Petykiewicz'
from .VERSION import __version__
__version__ = '1.2'
version = __version__

196
gridlock/base.py Normal file
View File

@ -0,0 +1,196 @@
from typing import Protocol
import numpy
from numpy.typing import NDArray
from . import GridError
class GridBase(Protocol):
exyz: list[NDArray]
"""Cell edges. Monotonically increasing without duplicates."""
periodic: list[bool]
"""For each axis, determines how far the rightmost boundary gets shifted. """
shifts: NDArray
"""Offsets `[[x0, y0, z0], [x1, y1, z1], ...]` for grid `0,1,...`"""
@property
def dxyz(self) -> list[NDArray]:
"""
Cell sizes for each axis, no shifts applied
Returns:
List of 3 ndarrays of cell sizes
"""
return [numpy.diff(ee) for ee in self.exyz]
@property
def xyz(self) -> list[NDArray]:
"""
Cell centers for each axis, no shifts applied
Returns:
List of 3 ndarrays of cell edges
"""
return [self.exyz[a][:-1] + self.dxyz[a] / 2.0 for a in range(3)]
@property
def shape(self) -> NDArray[numpy.intp]:
"""
The number of cells in x, y, and z
Returns:
ndarray of [x_centers.size, y_centers.size, z_centers.size]
"""
return numpy.array([coord.size - 1 for coord in self.exyz], dtype=int)
@property
def num_grids(self) -> int:
"""
The number of grids (number of shifts)
"""
return self.shifts.shape[0]
@property
def cell_data_shape(self) -> NDArray[numpy.intp]:
"""
The shape of the cell_data ndarray (num_grids, *self.shape).
"""
return numpy.hstack((self.num_grids, self.shape))
@property
def dxyz_with_ghost(self) -> list[NDArray]:
"""
Gives dxyz with an additional 'ghost' cell at the end, whose value depends
on whether or not the axis has periodic boundary conditions. See main description
above to learn why this is necessary.
If periodic, final edge shifts same amount as first
Otherwise, final edge shifts same amount as second-to-last
Returns:
list of [dxs, dys, dzs] with each element same length as elements of `self.xyz`
"""
el = [0 if p else -1 for p in self.periodic]
return [numpy.hstack((self.dxyz[a], self.dxyz[a][e])) for a, e in zip(range(3), el, strict=True)]
@property
def center(self) -> NDArray[numpy.float64]:
"""
Center position of the entire grid, no shifts applied
Returns:
ndarray of [x_center, y_center, z_center]
"""
# center is just average of first and last xyz, which is just the average of the
# first two and last two exyz
centers = [(self.exyz[a][:2] + self.exyz[a][-2:]).sum() / 4.0 for a in range(3)]
return numpy.array(centers, dtype=float)
@property
def dxyz_limits(self) -> tuple[NDArray, NDArray]:
"""
Returns the minimum and maximum cell size for each axis, as a tuple of two 3-element
ndarrays. No shifts are applied, so these are extreme bounds on these values (as a
weighted average is performed when shifting).
Returns:
Tuple of 2 ndarrays, `d_min=[min(dx), min(dy), min(dz)]` and `d_max=[...]`
"""
d_min = numpy.array([min(self.dxyz[a]) for a in range(3)], dtype=float)
d_max = numpy.array([max(self.dxyz[a]) for a in range(3)], dtype=float)
return d_min, d_max
def shifted_exyz(self, which_shifts: int | None) -> list[NDArray]:
"""
Returns edges for which_shifts.
Args:
which_shifts: Which grid (which shifts) to use, or `None` for unshifted
Returns:
List of 3 ndarrays of cell edges
"""
if which_shifts is None:
return self.exyz
dxyz = self.dxyz_with_ghost
shifts = self.shifts[which_shifts, :]
# If shift is negative, use left cell's dx to determine shift
for a in range(3):
if shifts[a] < 0:
dxyz[a] = numpy.roll(dxyz[a], 1)
return [self.exyz[a] + dxyz[a] * shifts[a] for a in range(3)]
def shifted_dxyz(self, which_shifts: int | None) -> list[NDArray]:
"""
Returns cell sizes for `which_shifts`.
Args:
which_shifts: Which grid (which shifts) to use, or `None` for unshifted
Returns:
List of 3 ndarrays of cell sizes
"""
if which_shifts is None:
return self.dxyz
shifts = self.shifts[which_shifts, :]
dxyz = self.dxyz_with_ghost
# If shift is negative, use left cell's dx to determine size
sdxyz = []
for a in range(3):
if shifts[a] < 0:
roll_dxyz = numpy.roll(dxyz[a], 1)
abs_shift = numpy.abs(shifts[a])
sdxyz.append(roll_dxyz[:-1] * abs_shift + roll_dxyz[1:] * (1 - abs_shift))
else:
sdxyz.append(dxyz[a][:-1] * (1 - shifts[a]) + dxyz[a][1:] * shifts[a])
return sdxyz
def shifted_xyz(self, which_shifts: int | None) -> list[NDArray[numpy.float64]]:
"""
Returns cell centers for `which_shifts`.
Args:
which_shifts: Which grid (which shifts) to use, or `None` for unshifted
Returns:
List of 3 ndarrays of cell centers
"""
if which_shifts is None:
return self.xyz
exyz = self.shifted_exyz(which_shifts)
dxyz = self.shifted_dxyz(which_shifts)
return [exyz[a][:-1] + dxyz[a] / 2.0 for a in range(3)]
def autoshifted_dxyz(self) -> list[NDArray[numpy.float64]]:
"""
Return cell widths, with each dimension shifted by the corresponding shifts.
Returns:
`[grid.shifted_dxyz(which_shifts=a)[a] for a in range(3)]`
"""
if self.num_grids != 3:
raise GridError('Autoshifting requires exactly 3 grids')
return [self.shifted_dxyz(which_shifts=a)[a] for a in range(3)]
def allocate(self, fill_value: float | None = 1.0, dtype: type[numpy.number] = numpy.float32) -> NDArray:
"""
Allocate an ndarray for storing grid data.
Args:
fill_value: Value to initialize the grid to. If None, an
uninitialized array is returned.
dtype: Numpy dtype for the array. Default is `numpy.float32`.
Returns:
The allocated array
"""
if fill_value is None:
return numpy.empty(self.cell_data_shape, dtype=dtype)
return numpy.full(self.cell_data_shape, fill_value, dtype=dtype)

View File

@ -1,12 +1,14 @@
"""
Drawing-related methods for Grid class
"""
from typing import List, Optional, Union, Sequence, Callable
from collections.abc import Sequence, Callable
import numpy # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
from float_raster import raster
from . import GridError
from .position import GridPosMixin
# NOTE: Maybe it would make sense to create a GridDrawer class
@ -15,16 +17,19 @@ from . import GridError
# without having to pass `cell_data` again each time?
foreground_callable_t = Callable[[numpy.ndarray, numpy.ndarray, numpy.ndarray], numpy.ndarray]
foreground_callable_t = Callable[[NDArray, NDArray, NDArray], NDArray]
foreground_t = float | foreground_callable_t
def draw_polygons(self,
cell_data: numpy.ndarray,
class GridDrawMixin(GridPosMixin):
def draw_polygons(
self,
cell_data: NDArray,
surface_normal: int,
center: numpy.ndarray,
polygons: Sequence[numpy.ndarray],
center: ArrayLike,
polygons: Sequence[ArrayLike],
thickness: float,
foreground: Union[Sequence[Union[float, foreground_callable_t]], float, foreground_callable_t],
foreground: Sequence[foreground_t] | foreground_t,
) -> None:
"""
Draw polygons on an axis-aligned plane.
@ -47,18 +52,20 @@ def draw_polygons(self,
"""
if surface_normal not in range(3):
raise GridError('Invalid surface_normal direction')
center = numpy.squeeze(center)
poly_list = [numpy.asarray(poly) for poly in polygons]
# Check polygons, and remove redundant coordinates
surface = numpy.delete(range(3), surface_normal)
for i, polygon in enumerate(polygons):
malformed = f'Malformed polygon: ({i})'
for ii in range(len(poly_list)):
polygon = poly_list[ii]
malformed = f'Malformed polygon: ({ii})'
if polygon.shape[1] not in (2, 3):
raise GridError(malformed + 'must be a Nx2 or Nx3 ndarray')
if polygon.shape[1] == 3:
polygon = polygon[surface, :]
poly_list[ii] = polygon
if not polygon.shape[0] > 2:
raise GridError(malformed + 'must consist of more than 2 points')
@ -67,16 +74,19 @@ def draw_polygons(self,
+ 'xyz'[surface_normal])
# Broadcast foreground where necessary
if numpy.size(foreground) == 1:
foreground = [foreground] * len(cell_data)
foregrounds: Sequence[foreground_callable_t] | Sequence[float]
if numpy.size(foreground) == 1: # type: ignore
foregrounds = [foreground] * len(cell_data) # type: ignore
elif isinstance(foreground, numpy.ndarray):
raise GridError('ndarray not supported for foreground')
else:
foregrounds = foreground # type: ignore
# ## Compute sub-domain of the grid occupied by polygons
# 1) Compute outer bounds (bd) of polygons
bd_2d_min = [0, 0]
bd_2d_max = [0, 0]
for polygon in polygons:
bd_2d_min = numpy.array([0, 0])
bd_2d_max = numpy.array([0, 0])
for polygon in poly_list:
bd_2d_min = numpy.minimum(bd_2d_min, polygon.min(axis=0))
bd_2d_max = numpy.maximum(bd_2d_max, polygon.max(axis=0))
bd_min = numpy.insert(bd_2d_min, surface_normal, -thickness / 2.0) + center
@ -94,35 +104,37 @@ def draw_polygons(self,
bdi_max = numpy.minimum(numpy.ceil(bdi_max), self.shape - 1).astype(int)
# 3) Adjust polygons for center
polygons = [poly + center[surface] for poly in polygons]
poly_list = [poly + center[surface] for poly in poly_list]
# ## Generate weighing function
def to_3d(vector: numpy.ndarray, val: float = 0.0) -> numpy.ndarray:
def to_3d(vector: NDArray, val: float = 0.0) -> NDArray[numpy.float64]:
v_2d = numpy.array(vector, dtype=float)
return numpy.insert(v_2d, surface_normal, (val,))
# iterate over grids
for i, grid in enumerate(cell_data):
# ## Evaluate or expand foreground[i]
if callable(foreground[i]):
foreground_val: NDArray | float
for i, _ in enumerate(cell_data):
# ## Evaluate or expand foregrounds[i]
foregrounds_i = foregrounds[i]
if callable(foregrounds_i):
# meshgrid over the (shifted) domain
domain = [self.shifted_xyz(i)[k][bdi_min[k]:bdi_max[k]+1] for k in range(3)]
domain = [self.shifted_xyz(i)[k][bdi_min[k]:bdi_max[k] + 1] for k in range(3)]
(x0, y0, z0) = numpy.meshgrid(*domain, indexing='ij')
# evaluate on the meshgrid
foreground_i = foreground[i](x0, y0, z0)
if not numpy.isfinite(foreground_i).all():
foreground_val = foregrounds_i(x0, y0, z0)
if not numpy.isfinite(foreground_val).all():
raise GridError(f'Non-finite values in foreground[{i}]')
elif numpy.size(foreground[i]) != 1:
raise GridError(f'Unsupported foreground[{i}]: {type(foreground[i])}')
elif numpy.size(foregrounds_i) != 1:
raise GridError(f'Unsupported foreground[{i}]: {type(foregrounds_i)}')
else:
# foreground[i] is scalar non-callable
foreground_i = foreground[i]
foreground_val = foregrounds_i
w_xy = numpy.zeros((bdi_max - bdi_min + 1)[surface].astype(int))
# Draw each polygon separately
for polygon in polygons:
for polygon in poly_list:
# Get the boundaries of the polygon
pbd_min = polygon.min(axis=0)
@ -137,13 +149,13 @@ def draw_polygons(self,
# Find indices in w_xy which are modified by polygon
# First for the edge coordinates (+1 since we're indexing edges)
edge_slices = [numpy.s_[i:f + 2] for i, f in zip(corner_min, corner_max)]
edge_slices = [numpy.s_[i:f + 2] for i, f in zip(corner_min, corner_max, strict=True)]
# Then for the pixel centers (-bdi_min since we're
# calculating weights within a subspace)
centers_slice = tuple(numpy.s_[i:f + 1] for i, f in zip(corner_min - bdi_min[surface],
corner_max - bdi_min[surface]))
corner_max - bdi_min[surface], strict=True))
aa_x, aa_y = (self.shifted_exyz(i)[a][s] for a, s in zip(surface, edge_slices))
aa_x, aa_y = (self.shifted_exyz(i)[a][s] for a, s in zip(surface, edge_slices, strict=True))
w_xy[centers_slice] += raster(polygon.T, aa_x, aa_y)
# Clamp overlapping polygons to 1
@ -152,7 +164,7 @@ def draw_polygons(self,
# 2) Generate weights in z-direction
w_z = numpy.zeros(((bdi_max - bdi_min + 1)[surface_normal], ))
def get_zi(offset, i=i, w_z=w_z):
def get_zi(offset: float, i=i, w_z=w_z) -> tuple[float, int]: # noqa: ANN001
edges = self.shifted_exyz(i)[surface_normal]
point = center[surface_normal] + offset
grid_coord = numpy.digitize(point, edges) - 1
@ -185,16 +197,17 @@ def draw_polygons(self,
# ## Modify the grid
g_slice = (i,) + tuple(numpy.s_[bdi_min[a]:bdi_max[a] + 1] for a in range(3))
cell_data[g_slice] = (1 - w) * cell_data[g_slice] + w * foreground_i
cell_data[g_slice] = (1 - w) * cell_data[g_slice] + w * foreground_val
def draw_polygon(self,
cell_data: numpy.ndarray,
def draw_polygon(
self,
cell_data: NDArray,
surface_normal: int,
center: numpy.ndarray,
polygon: numpy.ndarray,
center: ArrayLike,
polygon: ArrayLike,
thickness: float,
foreground: Union[Sequence[Union[float, foreground_callable_t]], float, foreground_callable_t],
foreground: Sequence[foreground_t] | foreground_t,
) -> None:
"""
Draw a polygon on an axis-aligned plane.
@ -212,12 +225,13 @@ def draw_polygon(self,
self.draw_polygons(cell_data, surface_normal, center, [polygon], thickness, foreground)
def draw_slab(self,
cell_data: numpy.ndarray,
def draw_slab(
self,
cell_data: NDArray,
surface_normal: int,
center: numpy.ndarray,
center: ArrayLike,
thickness: float,
foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t],
foreground: Sequence[foreground_t] | foreground_t,
) -> None:
"""
Draw an axis-aligned infinite slab.
@ -262,11 +276,12 @@ def draw_slab(self,
self.draw_polygon(cell_data, surface_normal, center_shift, p, thickness, foreground)
def draw_cuboid(self,
cell_data: numpy.ndarray,
center: numpy.ndarray,
dimensions: numpy.ndarray,
foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t],
def draw_cuboid(
self,
cell_data: NDArray,
center: ArrayLike,
dimensions: ArrayLike,
foreground: Sequence[foreground_t] | foreground_t,
) -> None:
"""
Draw an axis-aligned cuboid
@ -278,22 +293,24 @@ def draw_cuboid(self,
sizes of the cuboid
foreground: Value to draw with ('brush color'). See `draw_polygons()` for details.
"""
dimensions = numpy.asarray(dimensions)
p = numpy.array([[-dimensions[0], +dimensions[1]],
[+dimensions[0], +dimensions[1]],
[+dimensions[0], -dimensions[1]],
[-dimensions[0], -dimensions[1]]], dtype=float) / 2.0
[-dimensions[0], -dimensions[1]]], dtype=float) * 0.5
thickness = dimensions[2]
self.draw_polygon(cell_data, 2, center, p, thickness, foreground)
def draw_cylinder(self,
cell_data: numpy.ndarray,
def draw_cylinder(
self,
cell_data: NDArray,
surface_normal: int,
center: numpy.ndarray,
center: ArrayLike,
radius: float,
thickness: float,
num_points: int,
foreground: Union[List[Union[float, foreground_callable_t]], float, foreground_callable_t],
foreground: Sequence[foreground_t] | foreground_t,
) -> None:
"""
Draw an axis-aligned cylinder. Approximated by a num_points-gon
@ -307,16 +324,17 @@ def draw_cylinder(self,
num_points: The circle is approximated by a polygon with `num_points` vertices
foreground: Value to draw with ('brush color'). See `draw_polygons()` for details.
"""
theta = numpy.linspace(0, 2*numpy.pi, num_points, endpoint=False)
theta = numpy.linspace(0, 2 * numpy.pi, num_points, endpoint=False)
x = radius * numpy.sin(theta)
y = radius * numpy.cos(theta)
polygon = numpy.hstack((x[:, None], y[:, None]))
self.draw_polygon(cell_data, surface_normal, center, polygon, thickness, foreground)
def draw_extrude_rectangle(self,
cell_data: numpy.ndarray,
rectangle: numpy.ndarray,
def draw_extrude_rectangle(
self,
cell_data: NDArray,
rectangle: ArrayLike,
direction: int,
polarity: int,
distance: float,
@ -347,8 +365,8 @@ def draw_extrude_rectangle(self,
surface = numpy.delete(range(3), direction)
dim = numpy.fabs(numpy.diff(rectangle, axis=0).T)[surface]
p = numpy.vstack((numpy.array([-1, -1, 1, 1], dtype=float) * dim[0]/2.0,
numpy.array([-1, 1, 1, -1], dtype=float) * dim[1]/2.0)).T
p = numpy.vstack((numpy.array([-1, -1, 1, 1], dtype=float) * dim[0] * 0.5,
numpy.array([-1, 1, 1, -1], dtype=float) * dim[1] * 0.5)).T
thickness = distance
foreground_func = []
@ -358,16 +376,16 @@ def draw_extrude_rectangle(self,
ind = [int(numpy.floor(z)) if i == direction else slice(None) for i in range(3)]
fpart = z - numpy.floor(z)
mult = [1-fpart, fpart][::s] # reverses if s negative
mult = [1 - fpart, fpart][::s] # reverses if s negative
foreground = mult[0] * grid[tuple(ind)]
ind[direction] += 1
ind[direction] += 1 # type: ignore #(known safe)
foreground += mult[1] * grid[tuple(ind)]
def f_foreground(xs, ys, zs, i=i, foreground=foreground) -> numpy.ndarray:
def f_foreground(xs, ys, zs, i=i, foreground=foreground) -> NDArray[numpy.int64]: # noqa: ANN001
# transform from natural position to index
xyzi = numpy.array([self.pos2ind(qrs, which_shifts=i)
for qrs in zip(xs.flat, ys.flat, zs.flat)], dtype=int)
for qrs in zip(xs.flat, ys.flat, zs.flat, strict=True)], dtype=numpy.int64)
# reshape to original shape and keep only in-plane components
qi, ri = (numpy.reshape(xyzi[:, k], xs.shape) for k in surface)
return foreground[qi, ri]

View File

@ -1,4 +1,4 @@
import numpy # type: ignore
import numpy
from gridlock import Grid
@ -29,16 +29,16 @@ if __name__ == '__main__':
# numpy.linspace(-5.5, 5.5, 10)]
half_x = [.25, .5, 0.75, 1, 1.25, 1.5, 2, 2.5, 3, 3.5]
xyz3 = [[-x for x in half_x[::-1]] + [0] + half_x,
numpy.linspace(-5.5, 5.5, 10),
numpy.linspace(-5.5, 5.5, 10)]
xyz3 = [numpy.array([-x for x in half_x[::-1]] + [0] + half_x, dtype=float),
numpy.linspace(-5.5, 5.5, 10, dtype=float),
numpy.linspace(-5.5, 5.5, 10, dtype=float)]
eg = Grid(xyz3)
egc = eg.allocate(0)
# eg.draw_slab(Direction.z, 0, 10, 2)
eg.save('/home/jan/Desktop/test.pickle')
eg.draw_cylinder(egc, surface_normal=2, center=[0, 0, 0], radius=2.0,
thickness=10, num_poitns=1000, foreground=1)
thickness=10, num_points=1000, foreground=1)
eg.draw_extrude_rectangle(egc, rectangle=[[-2, 1, -1], [0, 1, 1]],
direction=1, poalarity=+1, distance=5)
direction=1, polarity=+1, distance=5)
eg.visualize_slice(egc, surface_normal=2, center=0, which_shifts=2)
eg.visualize_isosurface(egc, which_shifts=2)

View File

@ -1,20 +1,23 @@
from typing import List, Tuple, Callable, Dict, Optional, Union, Sequence, ClassVar, TypeVar
from typing import ClassVar, Self
from collections.abc import Callable, Sequence
import numpy # type: ignore
from numpy import diff, floor, ceil, zeros, hstack, newaxis
import numpy
from numpy.typing import NDArray, ArrayLike
import pickle
import warnings
import copy
from . import GridError
from .draw import GridDrawMixin
from .read import GridReadMixin
from .position import GridPosMixin
foreground_callable_type = Callable[[numpy.ndarray, numpy.ndarray, numpy.ndarray], numpy.ndarray]
T = TypeVar('T', bound='Grid')
foreground_callable_type = Callable[[NDArray, NDArray, NDArray], NDArray]
class Grid:
class Grid(GridDrawMixin, GridReadMixin, GridPosMixin):
"""
Simulation grid metadata for finite-difference simulations.
@ -48,216 +51,34 @@ class Grid:
Because of this, we either assume this 'ghost' cell is the same size as the last
real cell, or, if `self.periodic[a]` is set to `True`, the same size as the first cell.
"""
exyz: List[numpy.ndarray]
exyz: list[NDArray]
"""Cell edges. Monotonically increasing without duplicates."""
periodic: List[bool]
periodic: list[bool]
"""For each axis, determines how far the rightmost boundary gets shifted. """
shifts: numpy.ndarray
shifts: NDArray
"""Offsets `[[x0, y0, z0], [x1, y1, z1], ...]` for grid `0,1,...`"""
Yee_Shifts_E: ClassVar[numpy.ndarray] = 0.5 * numpy.array([[1, 0, 0],
Yee_Shifts_E: ClassVar[NDArray] = 0.5 * numpy.array([
[1, 0, 0],
[0, 1, 0],
[0, 0, 1]], dtype=float)
[0, 0, 1],
], dtype=float)
"""Default shifts for Yee grid E-field"""
Yee_Shifts_H: ClassVar[numpy.ndarray] = 0.5 * numpy.array([[0, 1, 1],
Yee_Shifts_H: ClassVar[NDArray] = 0.5 * numpy.array([
[0, 1, 1],
[1, 0, 1],
[1, 1, 0]], dtype=float)
[1, 1, 0],
], dtype=float)
"""Default shifts for Yee grid H-field"""
from .draw import (
draw_polygons, draw_polygon, draw_slab, draw_cuboid,
draw_cylinder, draw_extrude_rectangle,
)
from .read import get_slice, visualize_slice, visualize_isosurface
from .position import ind2pos, pos2ind
@property
def dxyz(self) -> List[numpy.ndarray]:
"""
Cell sizes for each axis, no shifts applied
Returns:
List of 3 ndarrays of cell sizes
"""
return [numpy.diff(ee) for ee in self.exyz]
@property
def xyz(self) -> List[numpy.ndarray]:
"""
Cell centers for each axis, no shifts applied
Returns:
List of 3 ndarrays of cell edges
"""
return [self.exyz[a][:-1] + self.dxyz[a] / 2.0 for a in range(3)]
@property
def shape(self) -> numpy.ndarray:
"""
The number of cells in x, y, and z
Returns:
ndarray of [x_centers.size, y_centers.size, z_centers.size]
"""
return numpy.array([coord.size - 1 for coord in self.exyz], dtype=int)
@property
def num_grids(self) -> int:
"""
The number of grids (number of shifts)
"""
return self.shifts.shape[0]
@property
def cell_data_shape(self):
"""
The shape of the cell_data ndarray (num_grids, *self.shape).
"""
return numpy.hstack((self.num_grids, self.shape))
@property
def dxyz_with_ghost(self) -> List[numpy.ndarray]:
"""
Gives dxyz with an additional 'ghost' cell at the end, whose value depends
on whether or not the axis has periodic boundary conditions. See main description
above to learn why this is necessary.
If periodic, final edge shifts same amount as first
Otherwise, final edge shifts same amount as second-to-last
Returns:
list of [dxs, dys, dzs] with each element same length as elements of `self.xyz`
"""
el = [0 if p else -1 for p in self.periodic]
return [numpy.hstack((self.dxyz[a], self.dxyz[a][e])) for a, e in zip(range(3), el)]
@property
def center(self) -> numpy.ndarray:
"""
Center position of the entire grid, no shifts applied
Returns:
ndarray of [x_center, y_center, z_center]
"""
# center is just average of first and last xyz, which is just the average of the
# first two and last two exyz
centers = [(self.exyz[a][:2] + self.exyz[a][-2:]).sum() / 4.0 for a in range(3)]
return numpy.array(centers, dtype=float)
@property
def dxyz_limits(self) -> Tuple[numpy.ndarray, numpy.ndarray]:
"""
Returns the minimum and maximum cell size for each axis, as a tuple of two 3-element
ndarrays. No shifts are applied, so these are extreme bounds on these values (as a
weighted average is performed when shifting).
Returns:
Tuple of 2 ndarrays, `d_min=[min(dx), min(dy), min(dz)]` and `d_max=[...]`
"""
d_min = numpy.array([min(self.dxyz[a]) for a in range(3)], dtype=float)
d_max = numpy.array([max(self.dxyz[a]) for a in range(3)], dtype=float)
return d_min, d_max
def shifted_exyz(self, which_shifts: Optional[int]) -> List[numpy.ndarray]:
"""
Returns edges for which_shifts.
Args:
which_shifts: Which grid (which shifts) to use, or `None` for unshifted
Returns:
List of 3 ndarrays of cell edges
"""
if which_shifts is None:
return self.exyz
dxyz = self.dxyz_with_ghost
shifts = self.shifts[which_shifts, :]
# If shift is negative, use left cell's dx to determine shift
for a in range(3):
if shifts[a] < 0:
dxyz[a] = numpy.roll(dxyz[a], 1)
return [self.exyz[a] + dxyz[a] * shifts[a] for a in range(3)]
def shifted_dxyz(self, which_shifts: Optional[int]) -> List[numpy.ndarray]:
"""
Returns cell sizes for `which_shifts`.
Args:
which_shifts: Which grid (which shifts) to use, or `None` for unshifted
Returns:
List of 3 ndarrays of cell sizes
"""
if which_shifts is None:
return self.dxyz
shifts = self.shifts[which_shifts, :]
dxyz = self.dxyz_with_ghost
# If shift is negative, use left cell's dx to determine size
sdxyz = []
for a in range(3):
if shifts[a] < 0:
roll_dxyz = numpy.roll(dxyz[a], 1)
abs_shift = numpy.abs(shifts[a])
sdxyz.append(roll_dxyz[:-1] * abs_shift + roll_dxyz[1:] * (1 - abs_shift))
else:
sdxyz.append(dxyz[a][:-1] * (1 - shifts[a]) + dxyz[a][1:] * shifts[a])
return sdxyz
def shifted_xyz(self, which_shifts: Optional[int]) -> List[numpy.ndarray]:
"""
Returns cell centers for `which_shifts`.
Args:
which_shifts: Which grid (which shifts) to use, or `None` for unshifted
Returns:
List of 3 ndarrays of cell centers
"""
if which_shifts is None:
return self.xyz
exyz = self.shifted_exyz(which_shifts)
dxyz = self.shifted_dxyz(which_shifts)
return [exyz[a][:-1] + dxyz[a] / 2.0 for a in range(3)]
def autoshifted_dxyz(self) -> List[numpy.ndarray]:
"""
Return cell widths, with each dimension shifted by the corresponding shifts.
Returns:
`[grid.shifted_dxyz(which_shifts=a)[a] for a in range(3)]`
"""
if self.num_grids != 3:
raise GridError('Autoshifting requires exactly 3 grids')
return [self.shifted_dxyz(which_shifts=a)[a] for a in range(3)]
def allocate(self, fill_value: Optional[float] = 1.0, dtype=numpy.float32) -> numpy.ndarray:
"""
Allocate an ndarray for storing grid data.
Args:
fill_value: Value to initialize the grid to. If None, an
uninitialized array is returned.
dtype: Numpy dtype for the array. Default is `numpy.float32`.
Returns:
The allocated array
"""
if fill_value is None:
return numpy.empty(self.cell_data_shape, dtype=dtype)
else:
return numpy.full(self.cell_data_shape, fill_value, dtype=dtype)
def __init__(self,
pixel_edge_coordinates: Sequence[numpy.ndarray],
shifts: numpy.ndarray = Yee_Shifts_E,
periodic: Union[bool, Sequence[bool]] = False,
def __init__(
self,
pixel_edge_coordinates: Sequence[ArrayLike],
shifts: ArrayLike = Yee_Shifts_E,
periodic: bool | Sequence[bool] = False,
) -> None:
"""
Args:
@ -273,11 +94,12 @@ class Grid:
Raises:
`GridError` on invalid input
"""
self.exyz = [numpy.unique(pixel_edge_coordinates[i]) for i in range(3)]
edge_arrs = [numpy.array(cc) for cc in pixel_edge_coordinates]
self.exyz = [numpy.unique(edges) for edges in edge_arrs]
self.shifts = numpy.array(shifts, dtype=float)
for i in range(3):
if len(self.exyz[i]) != len(pixel_edge_coordinates[i]):
if self.exyz[i].size != edge_arrs[i].size:
warnings.warn(f'Dimension {i} had duplicate edge coordinates', stacklevel=2)
if isinstance(periodic, bool):
@ -314,7 +136,7 @@ class Grid:
g.__dict__.update(tmp_dict)
return g
def save(self: T, filename: str) -> T:
def save(self, filename: str) -> Self:
"""
Save to file.
@ -328,7 +150,7 @@ class Grid:
pickle.dump(self.__dict__, f, protocol=2)
return self
def copy(self: T) -> T:
def copy(self) -> Self:
"""
Returns:
Deep copy of the grid.

View File

@ -1,19 +1,21 @@
"""
Position-related methods for Grid class
"""
from typing import List, Optional
import numpy # type: ignore
import numpy
from numpy.typing import NDArray, ArrayLike
from . import GridError
from .base import GridBase
def ind2pos(self,
ind: numpy.ndarray,
which_shifts: Optional[int] = None,
class GridPosMixin(GridBase):
def ind2pos(
self,
ind: NDArray,
which_shifts: int | None = None,
round_ind: bool = True,
check_bounds: bool = True
) -> numpy.ndarray:
) -> NDArray[numpy.float64]:
"""
Returns the natural position corresponding to the specified cell center indices.
The resulting position is clipped to the bounds of the grid
@ -59,12 +61,13 @@ def ind2pos(self,
return numpy.array(position, dtype=float)
def pos2ind(self,
r: numpy.ndarray,
which_shifts: Optional[int],
def pos2ind(
self,
r: ArrayLike,
which_shifts: int | None,
round_ind: bool = True,
check_bounds: bool = True
) -> numpy.ndarray:
) -> NDArray[numpy.float64]:
"""
Returns the cell-center indices corresponding to the specified natural position.
The resulting position is clipped to within the outer centers of the grid.

View File

@ -1,24 +1,33 @@
"""
Readback and visualization methods for Grid class
"""
from typing import Dict, Optional, Union, Any
from typing import Any, TYPE_CHECKING
import numpy # type: ignore
import numpy
from numpy.typing import NDArray
from . import GridError
from .position import GridPosMixin
if TYPE_CHECKING:
import matplotlib.axes
import matplotlib.figure
# .visualize_* uses matplotlib
# .visualize_isosurface uses skimage
# .visualize_isosurface uses mpl_toolkits.mplot3d
def get_slice(self,
cell_data: numpy.ndarray,
class GridReadMixin(GridPosMixin):
def get_slice(
self,
cell_data: NDArray,
surface_normal: int,
center: float,
which_shifts: int = 0,
sample_period: int = 1
) -> numpy.ndarray:
) -> NDArray:
"""
Retrieve a slice of a grid.
Interpolates if given a position between two planes.
@ -65,7 +74,7 @@ def get_slice(self,
# Extract grid values from planes above and below visualized slice
sliced_grid = numpy.zeros(self.shape[surface])
for ci, weight in zip(centers, w):
for ci, weight in zip(centers, w, strict=True):
s = tuple(ci if a == surface_normal else numpy.s_[::sp] for a in range(3))
sliced_grid += weight * cell_data[which_shifts][tuple(s)]
@ -75,15 +84,16 @@ def get_slice(self,
return sliced_grid
def visualize_slice(self,
cell_data: numpy.ndarray,
def visualize_slice(
self,
cell_data: NDArray,
surface_normal: int,
center: float,
which_shifts: int = 0,
sample_period: int = 1,
finalize: bool = True,
pcolormesh_args: Optional[Dict[str, Any]] = None,
) -> None:
pcolormesh_args: dict[str, Any] | None = None,
) -> tuple['matplotlib.axes.Axes', 'matplotlib.figure.Figure']:
"""
Visualize a slice of a grid.
Interpolates if given a position between two planes.
@ -94,6 +104,9 @@ def visualize_slice(self,
which_shifts: Which grid to display. Default is the first grid (0).
sample_period: Period for down-sampling the image. Default 1 (disabled)
finalize: Whether to call `pyplot.show()` after constructing the plot. Default `True`
Returns:
(Figure, Axes)
"""
from matplotlib import pyplot
@ -112,24 +125,27 @@ def visualize_slice(self,
xmesh, ymesh = numpy.meshgrid(x, y, indexing='ij')
x_label, y_label = ('xyz'[a] for a in surface)
pyplot.figure()
pyplot.pcolormesh(xmesh, ymesh, grid_slice, **pcolormesh_args)
pyplot.colorbar()
pyplot.gca().set_aspect('equal', adjustable='box')
pyplot.xlabel(x_label)
pyplot.ylabel(y_label)
fig, ax = pyplot.subplots()
mappable = ax.pcolormesh(xmesh, ymesh, grid_slice, **pcolormesh_args)
fig.colorbar(mappable)
ax.set_aspect('equal', adjustable='box')
ax.set_xlabel(x_label)
ax.set_ylabel(y_label)
if finalize:
pyplot.show()
return fig, ax
def visualize_isosurface(self,
cell_data: numpy.ndarray,
level: Optional[float] = None,
def visualize_isosurface(
self,
cell_data: NDArray,
level: float | None = None,
which_shifts: int = 0,
sample_period: int = 1,
show_edges: bool = True,
finalize: bool = True,
) -> None:
) -> tuple['matplotlib.axes.Axes', 'matplotlib.figure.Figure']:
"""
Draw an isosurface plot of the device.
@ -140,11 +156,15 @@ def visualize_isosurface(self,
sample_period: Period for down-sampling the image. Default 1 (disabled)
show_edges: Whether to draw triangle edges. Default `True`
finalize: Whether to call `pyplot.show()` after constructing the plot. Default `True`
Returns:
(Figure, Axes)
"""
from matplotlib import pyplot
import skimage.measure
# Claims to be unused, but needed for subplot(projection='3d')
from mpl_toolkits.mplot3d import Axes3D
del Axes3D # imported for side effects only
# Get data from cell_data
grid = cell_data[which_shifts][::sample_period, ::sample_period, ::sample_period]
@ -176,8 +196,10 @@ def visualize_isosurface(self,
ybs = 0.5 * max_range * mg[1].flatten() + 0.5 * (ys.max() + ys.min())
zbs = 0.5 * max_range * mg[2].flatten() + 0.5 * (zs.max() + zs.min())
# Comment or uncomment following both lines to test the fake bounding box:
for xb, yb, zb in zip(xbs, ybs, zbs):
for xb, yb, zb in zip(xbs, ybs, zbs, strict=True):
ax.plot([xb], [yb], [zb], 'w')
if finalize:
pyplot.show()
return fig, ax

View File

@ -1,6 +1,6 @@
import pytest # type: ignore
import numpy # type: ignore
from numpy.testing import assert_allclose, assert_array_equal # type: ignore
# import pytest
import numpy
from numpy.testing import assert_allclose #, assert_array_equal
from .. import Grid

99
pyproject.toml Normal file
View File

@ -0,0 +1,99 @@
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "gridlock"
description = "Coupled gridding library"
readme = "README.md"
license = { file = "LICENSE.md" }
authors = [
{ name="Jan Petykiewicz", email="jan@mpxd.net" },
]
homepage = "https://mpxd.net/code/jan/gridlock"
repository = "https://mpxd.net/code/jan/gridlock"
keywords = [
"FDTD",
"gridding",
"simulation",
"nonuniform",
"FDFD",
"finite",
"difference",
]
classifiers = [
"Programming Language :: Python :: 3",
"Development Status :: 4 - Beta",
"Intended Audience :: Developers",
"Intended Audience :: Science/Research",
"License :: OSI Approved :: GNU Affero General Public License v3",
"Topic :: Multimedia :: Graphics :: 3D Rendering",
"Topic :: Scientific/Engineering :: Electronic Design Automation (EDA)",
"Topic :: Scientific/Engineering :: Physics",
"Topic :: Scientific/Engineering :: Visualization",
]
requires-python = ">=3.11"
include = [
"LICENSE.md"
]
dynamic = ["version"]
dependencies = [
"numpy>=1.26",
"float_raster>=0.8",
]
[tool.hatch.version]
path = "gridlock/__init__.py"
[project.optional-dependencies]
visualization = ["matplotlib"]
visualization-isosurface = [
"matplotlib",
"skimage>=0.13",
"mpl_toolkits",
]
[tool.ruff]
exclude = [
".git",
"dist",
]
line-length = 145
indent-width = 4
lint.dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
lint.select = [
"NPY", "E", "F", "W", "B", "ANN", "UP", "SLOT", "SIM", "LOG",
"C4", "ISC", "PIE", "PT", "RET", "TCH", "PTH", "INT",
"ARG", "PL", "R", "TRY",
"G010", "G101", "G201", "G202",
"Q002", "Q003", "Q004",
]
lint.ignore = [
#"ANN001", # No annotation
"ANN002", # *args
"ANN003", # **kwargs
"ANN401", # Any
"ANN101", # self: Self
"SIM108", # single-line if / else assignment
"RET504", # x=y+z; return x
"PIE790", # unnecessary pass
"ISC003", # non-implicit string concatenation
"C408", # dict(x=y) instead of {'x': y}
"PLR09", # Too many xxx
"PLR2004", # magic number
"PLC0414", # import x as x
"TRY003", # Long exception message
"PTH123", # open()
]
[[tool.mypy.overrides]]
module = [
"matplotlib",
"matplotlib.axes",
"matplotlib.figure",
"mpl_toolkits.mplot3d",
]
ignore_missing_imports = true

View File

@ -1,47 +0,0 @@
#!/usr/bin/env python3
from setuptools import setup, find_packages
with open('README.md', 'r') as f:
long_description = f.read()
with open('gridlock/VERSION.py', 'rt') as f:
version = f.readlines()[2].strip()
setup(name='gridlock',
version=version,
description='Coupled gridding library',
long_description=long_description,
long_description_content_type='text/markdown',
author='Jan Petykiewicz',
author_email='jan@mpxd.net',
url='https://mpxd.net/code/jan/gridlock',
packages=find_packages(),
package_data={
'gridlock': ['py.typed'],
},
install_requires=[
'numpy',
'float_raster',
],
extras_require={
'visualization': ['matplotlib'],
'visualization-isosurface': [
'matplotlib',
'skimage>=0.13',
'mpl_toolkits',
],
},
classifiers=[
'Programming Language :: Python :: 3',
'Development Status :: 4 - Beta',
'Intended Audience :: Developers',
'Intended Audience :: Science/Research',
'License :: OSI Approved :: GNU Affero General Public License v3',
'Topic :: Multimedia :: Graphics :: 3D Rendering',
'Topic :: Scientific/Engineering :: Electronic Design Automation (EDA)',
'Topic :: Scientific/Engineering :: Physics',
'Topic :: Scientific/Engineering :: Visualization',
],
)