Compare commits
No commits in common. "8e7e0edb1f8c070249e301427f0ad36ff3699ec5" and "9ab97e763cc5662df4297c51bc2331652698c345" have entirely different histories.
8e7e0edb1f
...
9ab97e763c
@ -15,8 +15,8 @@ Dependencies:
|
||||
- mpl_toolkits.mplot3d [Grid.visualize_isosurface()]
|
||||
- skimage [Grid.visualize_isosurface()]
|
||||
"""
|
||||
from .error import GridError as GridError
|
||||
from .grid import Grid as Grid
|
||||
from .error import GridError
|
||||
from .grid import Grid
|
||||
|
||||
__author__ = 'Jan Petykiewicz'
|
||||
__version__ = '1.1'
|
||||
|
||||
196
gridlock/base.py
196
gridlock/base.py
@ -1,196 +0,0 @@
|
||||
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)
|
||||
@ -1,14 +1,13 @@
|
||||
"""
|
||||
Drawing-related methods for Grid class
|
||||
"""
|
||||
from collections.abc import Sequence, Callable
|
||||
from typing import Union, Sequence, Callable
|
||||
|
||||
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
|
||||
@ -18,16 +17,15 @@ from .position import GridPosMixin
|
||||
|
||||
|
||||
foreground_callable_t = Callable[[NDArray, NDArray, NDArray], NDArray]
|
||||
foreground_t = float | foreground_callable_t
|
||||
foreground_t = Union[float, foreground_callable_t]
|
||||
|
||||
|
||||
class GridDrawMixin(GridPosMixin):
|
||||
def draw_polygons(
|
||||
def draw_polygons(
|
||||
self,
|
||||
cell_data: NDArray,
|
||||
surface_normal: int,
|
||||
center: ArrayLike,
|
||||
polygons: Sequence[ArrayLike],
|
||||
polygons: Sequence[NDArray],
|
||||
thickness: float,
|
||||
foreground: Sequence[foreground_t] | foreground_t,
|
||||
) -> None:
|
||||
@ -52,20 +50,18 @@ class GridDrawMixin(GridPosMixin):
|
||||
"""
|
||||
if surface_normal not in range(3):
|
||||
raise GridError('Invalid surface_normal direction')
|
||||
|
||||
center = numpy.squeeze(center)
|
||||
poly_list = [numpy.array(poly, copy=False) for poly in polygons]
|
||||
|
||||
# Check polygons, and remove redundant coordinates
|
||||
surface = numpy.delete(range(3), surface_normal)
|
||||
|
||||
for ii in range(len(poly_list)):
|
||||
polygon = poly_list[ii]
|
||||
malformed = f'Malformed polygon: ({ii})'
|
||||
for i, polygon in enumerate(polygons):
|
||||
malformed = f'Malformed polygon: ({i})'
|
||||
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')
|
||||
@ -86,7 +82,7 @@ class GridDrawMixin(GridPosMixin):
|
||||
# 1) Compute outer bounds (bd) of polygons
|
||||
bd_2d_min = numpy.array([0, 0])
|
||||
bd_2d_max = numpy.array([0, 0])
|
||||
for polygon in poly_list:
|
||||
for polygon in polygons:
|
||||
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
|
||||
@ -104,7 +100,7 @@ class GridDrawMixin(GridPosMixin):
|
||||
bdi_max = numpy.minimum(numpy.ceil(bdi_max), self.shape - 1).astype(int)
|
||||
|
||||
# 3) Adjust polygons for center
|
||||
poly_list = [poly + center[surface] for poly in poly_list]
|
||||
polygons = [poly + center[surface] for poly in polygons]
|
||||
|
||||
# ## Generate weighing function
|
||||
def to_3d(vector: NDArray, val: float = 0.0) -> NDArray[numpy.float64]:
|
||||
@ -112,7 +108,6 @@ class GridDrawMixin(GridPosMixin):
|
||||
return numpy.insert(v_2d, surface_normal, (val,))
|
||||
|
||||
# iterate over grids
|
||||
foreground_val: NDArray | float
|
||||
for i, _ in enumerate(cell_data):
|
||||
# ## Evaluate or expand foregrounds[i]
|
||||
foregrounds_i = foregrounds[i]
|
||||
@ -134,7 +129,7 @@ class GridDrawMixin(GridPosMixin):
|
||||
w_xy = numpy.zeros((bdi_max - bdi_min + 1)[surface].astype(int))
|
||||
|
||||
# Draw each polygon separately
|
||||
for polygon in poly_list:
|
||||
for polygon in polygons:
|
||||
|
||||
# Get the boundaries of the polygon
|
||||
pbd_min = polygon.min(axis=0)
|
||||
@ -149,13 +144,13 @@ class GridDrawMixin(GridPosMixin):
|
||||
|
||||
# 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, strict=True)]
|
||||
edge_slices = [numpy.s_[i:f + 2] for i, f in zip(corner_min, corner_max)]
|
||||
# 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], strict=True))
|
||||
corner_max - bdi_min[surface]))
|
||||
|
||||
aa_x, aa_y = (self.shifted_exyz(i)[a][s] for a, s in zip(surface, edge_slices, strict=True))
|
||||
aa_x, aa_y = (self.shifted_exyz(i)[a][s] for a, s in zip(surface, edge_slices))
|
||||
w_xy[centers_slice] += raster(polygon.T, aa_x, aa_y)
|
||||
|
||||
# Clamp overlapping polygons to 1
|
||||
@ -164,7 +159,7 @@ class GridDrawMixin(GridPosMixin):
|
||||
# 2) Generate weights in z-direction
|
||||
w_z = numpy.zeros(((bdi_max - bdi_min + 1)[surface_normal], ))
|
||||
|
||||
def get_zi(offset: float, i=i, w_z=w_z) -> tuple[float, int]: # noqa: ANN001
|
||||
def get_zi(offset, i=i, w_z=w_z):
|
||||
edges = self.shifted_exyz(i)[surface_normal]
|
||||
point = center[surface_normal] + offset
|
||||
grid_coord = numpy.digitize(point, edges) - 1
|
||||
@ -200,7 +195,7 @@ class GridDrawMixin(GridPosMixin):
|
||||
cell_data[g_slice] = (1 - w) * cell_data[g_slice] + w * foreground_val
|
||||
|
||||
|
||||
def draw_polygon(
|
||||
def draw_polygon(
|
||||
self,
|
||||
cell_data: NDArray,
|
||||
surface_normal: int,
|
||||
@ -225,7 +220,7 @@ class GridDrawMixin(GridPosMixin):
|
||||
self.draw_polygons(cell_data, surface_normal, center, [polygon], thickness, foreground)
|
||||
|
||||
|
||||
def draw_slab(
|
||||
def draw_slab(
|
||||
self,
|
||||
cell_data: NDArray,
|
||||
surface_normal: int,
|
||||
@ -276,7 +271,7 @@ class GridDrawMixin(GridPosMixin):
|
||||
self.draw_polygon(cell_data, surface_normal, center_shift, p, thickness, foreground)
|
||||
|
||||
|
||||
def draw_cuboid(
|
||||
def draw_cuboid(
|
||||
self,
|
||||
cell_data: NDArray,
|
||||
center: ArrayLike,
|
||||
@ -302,7 +297,7 @@ class GridDrawMixin(GridPosMixin):
|
||||
self.draw_polygon(cell_data, 2, center, p, thickness, foreground)
|
||||
|
||||
|
||||
def draw_cylinder(
|
||||
def draw_cylinder(
|
||||
self,
|
||||
cell_data: NDArray,
|
||||
surface_normal: int,
|
||||
@ -331,7 +326,7 @@ class GridDrawMixin(GridPosMixin):
|
||||
self.draw_polygon(cell_data, surface_normal, center, polygon, thickness, foreground)
|
||||
|
||||
|
||||
def draw_extrude_rectangle(
|
||||
def draw_extrude_rectangle(
|
||||
self,
|
||||
cell_data: NDArray,
|
||||
rectangle: ArrayLike,
|
||||
@ -382,10 +377,10 @@ class GridDrawMixin(GridPosMixin):
|
||||
ind[direction] += 1 # type: ignore #(known safe)
|
||||
foreground += mult[1] * grid[tuple(ind)]
|
||||
|
||||
def f_foreground(xs, ys, zs, i=i, foreground=foreground) -> NDArray[numpy.int64]: # noqa: ANN001
|
||||
def f_foreground(xs, ys, zs, i=i, foreground=foreground) -> NDArray[numpy.int_]:
|
||||
# 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, strict=True)], dtype=numpy.int64)
|
||||
for qrs in zip(xs.flat, ys.flat, zs.flat)], dtype=int)
|
||||
# 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]
|
||||
|
||||
@ -29,7 +29,7 @@ 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 = [numpy.array([-x for x in half_x[::-1]] + [0] + half_x),
|
||||
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)]
|
||||
eg = Grid(xyz3)
|
||||
@ -37,8 +37,8 @@ if __name__ == '__main__':
|
||||
# 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_points=1000, foreground=1)
|
||||
thickness=10, num_poitns=1000, foreground=1)
|
||||
eg.draw_extrude_rectangle(egc, rectangle=[[-2, 1, -1], [0, 1, 1]],
|
||||
direction=1, polarity=+1, distance=5)
|
||||
direction=1, poalarity=+1, distance=5)
|
||||
eg.visualize_slice(egc, surface_normal=2, center=0, which_shifts=2)
|
||||
eg.visualize_isosurface(egc, which_shifts=2)
|
||||
|
||||
200
gridlock/grid.py
200
gridlock/grid.py
@ -1,5 +1,4 @@
|
||||
from typing import ClassVar, Self
|
||||
from collections.abc import Callable, Sequence
|
||||
from typing import Callable, Sequence, ClassVar, Self
|
||||
|
||||
import numpy
|
||||
from numpy.typing import NDArray, ArrayLike
|
||||
@ -9,15 +8,12 @@ import warnings
|
||||
import copy
|
||||
|
||||
from . import GridError
|
||||
from .draw import GridDrawMixin
|
||||
from .read import GridReadMixin
|
||||
from .position import GridPosMixin
|
||||
|
||||
|
||||
foreground_callable_type = Callable[[NDArray, NDArray, NDArray], NDArray]
|
||||
|
||||
|
||||
class Grid(GridDrawMixin, GridReadMixin, GridPosMixin):
|
||||
class Grid:
|
||||
"""
|
||||
Simulation grid metadata for finite-difference simulations.
|
||||
|
||||
@ -74,6 +70,193 @@ class Grid(GridDrawMixin, GridReadMixin, GridPosMixin):
|
||||
], 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[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.int_]:
|
||||
"""
|
||||
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[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) -> 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=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)
|
||||
else:
|
||||
return numpy.full(self.cell_data_shape, fill_value, dtype=dtype)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pixel_edge_coordinates: Sequence[ArrayLike],
|
||||
@ -94,12 +277,11 @@ class Grid(GridDrawMixin, GridReadMixin, GridPosMixin):
|
||||
Raises:
|
||||
`GridError` on invalid input
|
||||
"""
|
||||
edge_arrs = [numpy.array(cc, copy=False) for cc in pixel_edge_coordinates]
|
||||
self.exyz = [numpy.unique(edges) for edges in edge_arrs]
|
||||
self.exyz = [numpy.unique(pixel_edge_coordinates[i]) for i in range(3)]
|
||||
self.shifts = numpy.array(shifts, dtype=float)
|
||||
|
||||
for i in range(3):
|
||||
if self.exyz[i].size != edge_arrs[i].size:
|
||||
if len(self.exyz[i]) != len(pixel_edge_coordinates[i]):
|
||||
warnings.warn(f'Dimension {i} had duplicate edge coordinates', stacklevel=2)
|
||||
|
||||
if isinstance(periodic, bool):
|
||||
|
||||
@ -5,11 +5,9 @@ import numpy
|
||||
from numpy.typing import NDArray, ArrayLike
|
||||
|
||||
from . import GridError
|
||||
from .base import GridBase
|
||||
|
||||
|
||||
class GridPosMixin(GridBase):
|
||||
def ind2pos(
|
||||
def ind2pos(
|
||||
self,
|
||||
ind: NDArray,
|
||||
which_shifts: int | None = None,
|
||||
@ -61,7 +59,7 @@ class GridPosMixin(GridBase):
|
||||
return numpy.array(position, dtype=float)
|
||||
|
||||
|
||||
def pos2ind(
|
||||
def pos2ind(
|
||||
self,
|
||||
r: ArrayLike,
|
||||
which_shifts: int | None,
|
||||
|
||||
@ -7,7 +7,6 @@ import numpy
|
||||
from numpy.typing import NDArray
|
||||
|
||||
from . import GridError
|
||||
from .position import GridPosMixin
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import matplotlib.axes
|
||||
@ -19,8 +18,7 @@ if TYPE_CHECKING:
|
||||
# .visualize_isosurface uses mpl_toolkits.mplot3d
|
||||
|
||||
|
||||
class GridReadMixin(GridPosMixin):
|
||||
def get_slice(
|
||||
def get_slice(
|
||||
self,
|
||||
cell_data: NDArray,
|
||||
surface_normal: int,
|
||||
@ -74,7 +72,7 @@ class GridReadMixin(GridPosMixin):
|
||||
|
||||
# Extract grid values from planes above and below visualized slice
|
||||
sliced_grid = numpy.zeros(self.shape[surface])
|
||||
for ci, weight in zip(centers, w, strict=True):
|
||||
for ci, weight in zip(centers, w):
|
||||
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)]
|
||||
|
||||
@ -84,7 +82,7 @@ class GridReadMixin(GridPosMixin):
|
||||
return sliced_grid
|
||||
|
||||
|
||||
def visualize_slice(
|
||||
def visualize_slice(
|
||||
self,
|
||||
cell_data: NDArray,
|
||||
surface_normal: int,
|
||||
@ -137,7 +135,7 @@ class GridReadMixin(GridPosMixin):
|
||||
return fig, ax
|
||||
|
||||
|
||||
def visualize_isosurface(
|
||||
def visualize_isosurface(
|
||||
self,
|
||||
cell_data: NDArray,
|
||||
level: float | None = None,
|
||||
@ -164,7 +162,6 @@ class GridReadMixin(GridPosMixin):
|
||||
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]
|
||||
@ -196,7 +193,7 @@ class GridReadMixin(GridPosMixin):
|
||||
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, strict=True):
|
||||
for xb, yb, zb in zip(xbs, ybs, zbs):
|
||||
ax.plot([xb], [yb], [zb], 'w')
|
||||
|
||||
if finalize:
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
import pytest # type: ignore
|
||||
import numpy
|
||||
from numpy.testing import assert_allclose #, assert_array_equal
|
||||
from numpy.testing import assert_allclose, assert_array_equal
|
||||
|
||||
from .. import Grid
|
||||
|
||||
|
||||
@ -53,47 +53,3 @@ visualization-isosurface = [
|
||||
"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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user