diff --git a/gridlock/draw.py b/gridlock/draw.py index 15a6bff..3a8ec8a 100644 --- a/gridlock/draw.py +++ b/gridlock/draw.py @@ -4,8 +4,6 @@ Drawing-related methods for Grid class from typing import List, Optional, Union, Sequence, Callable import numpy # type: ignore -from numpy import diff, floor, ceil, zeros, hstack, newaxis - from float_raster import raster from . import GridError, Direction @@ -91,8 +89,8 @@ def draw_polygons(self, s_max = self.shifts.max(axis=0) bdi_min = self.pos2ind(bd_min + s_min, None, round_ind=False, check_bounds=False) - buf bdi_max = self.pos2ind(bd_max + s_max, None, round_ind=False, check_bounds=False) + buf - bdi_min = numpy.maximum(floor(bdi_min), 0).astype(int) - bdi_max = numpy.minimum(ceil(bdi_max), self.shape - 1).astype(int) + bdi_min = numpy.maximum(numpy.floor(bdi_min), 0).astype(int) + 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] @@ -120,7 +118,7 @@ def draw_polygons(self, # eps[i] is scalar non-callable eps_i = eps[i] - w_xy = zeros((bdi_max - bdi_min + 1)[surface].astype(int)) + w_xy = numpy.zeros((bdi_max - bdi_min + 1)[surface].astype(int)) # Draw each polygon separately for polygon in polygons: @@ -182,7 +180,7 @@ def draw_polygons(self, w_z[zi_bot] = zi_top_f - zi_bot_f # 3) Generate total weight function - w = (w_xy[:, :, newaxis] * w_z).transpose(numpy.insert([0, 1], surface_normal, (2,))) + w = (w_xy[:, :, None] * w_z).transpose(numpy.insert([0, 1], surface_normal, (2,))) # ## Modify the grid g_slice = (i,) + tuple(numpy.s_[bdi_min[a]:bdi_max[a] + 1] for a in range(3)) @@ -308,8 +306,8 @@ def draw_cylinder(self, theta = numpy.linspace(0, 2*numpy.pi, num_points, endpoint=False) x = radius * numpy.sin(theta) y = radius * numpy.cos(theta) - polygon = hstack((x[:, newaxis], y[:, newaxis])) self.draw_polygon(surface_normal, center, polygon, thickness, eps) + polygon = numpy.hstack((x[:, None], y[:, None])) def draw_extrude_rectangle(self, @@ -347,7 +345,7 @@ def draw_extrude_rectangle(self, surface = numpy.delete(range(3), direction) - dim = numpy.fabs(diff(rectangle, axis=0).T)[surface] + 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 thickness = distance @@ -356,9 +354,9 @@ def draw_extrude_rectangle(self, for i, grid in enumerate(self.grids): z = self.pos2ind(rectangle[0, :], i, round_ind=False, check_bounds=False)[direction] - ind = [int(floor(z)) if i == direction else slice(None) for i in range(3)] + ind = [int(numpy.floor(z)) if i == direction else slice(None) for i in range(3)] - fpart = z - floor(z) + fpart = z - numpy.floor(z) mult = [1-fpart, fpart][::s] # reverses if s negative eps = mult[0] * grid[tuple(ind)] diff --git a/gridlock/grid.py b/gridlock/grid.py index 7dc62c6..3c1c4a8 100644 --- a/gridlock/grid.py +++ b/gridlock/grid.py @@ -81,7 +81,7 @@ class Grid: Returns: List of 3 ndarrays of cell sizes """ - return [diff(self.exyz[a]) for a in range(3)] + return [numpy.diff(ee) for ee in self.exyz] @property def xyz(self) -> List[numpy.ndarray]: @@ -117,7 +117,7 @@ class Grid: 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 [hstack((self.dxyz[a], self.dxyz[a][e])) for a, e in zip(range(3), el)] + return [numpy.hstack((self.dxyz[a], self.dxyz[a][e])) for a, e in zip(range(3), el)] @property def center(self) -> numpy.ndarray: diff --git a/gridlock/position.py b/gridlock/position.py index 282824f..02856e8 100644 --- a/gridlock/position.py +++ b/gridlock/position.py @@ -4,7 +4,6 @@ Position-related methods for Grid class from typing import List, Optional import numpy # type: ignore -from numpy import zeros from . import GridError @@ -97,7 +96,7 @@ def pos2ind(self, if self.shape[a] > 1 and (r[a] < sexyz[a][0] or r[a] > sexyz[a][-1]): raise GridError('Position[{}] outside of grid!'.format(a)) - grid_pos = zeros((3,)) + grid_pos = numpy.zeros((3,)) for a in range(3): xi = numpy.digitize(r[a], sexyz[a]) - 1 # Figure out which cell we're in xi_clipped = numpy.clip(xi, 0, sexyz[a].size - 2) # Clip back into grid bounds diff --git a/gridlock/read.py b/gridlock/read.py index 32e145e..e049d77 100644 --- a/gridlock/read.py +++ b/gridlock/read.py @@ -4,7 +4,6 @@ Readback and visualization methods for Grid class from typing import Dict, Optional, Union, Any import numpy # type: ignore -from numpy import floor, ceil, zeros from . import GridError, Direction from ._helpers import is_scalar @@ -56,9 +55,9 @@ def get_slice(self, center3 = numpy.insert([0, 0], surface_normal, (center,)) center_index = self.pos2ind(center3, which_shifts, round_ind=False, check_bounds=False)[surface_normal] - centers = numpy.unique([floor(center_index), ceil(center_index)]).astype(int) + centers = numpy.unique([numpy.floor(center_index), numpy.ceil(center_index)]).astype(int) if len(centers) == 2: - fpart = center_index - floor(center_index) + fpart = center_index - numpy.floor(center_index) w = [1 - fpart, fpart] # longer distance -> less weight else: w = [1] @@ -68,7 +67,7 @@ def get_slice(self, raise GridError('Coordinate of selected plane must be within simulation domain') # Extract grid values from planes above and below visualized slice - sliced_grid = zeros(self.shape[surface]) + sliced_grid = numpy.zeros(self.shape[surface]) 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 * self.grids[which_shifts][tuple(s)]