jan/masque
1
1
Fork 0
Code Issues Pull requests Releases Wiki Activity

Use ArrayLike and NDArray wherever possible. Some type fixes and some related corner cases

Browse source
This commit is contained in:
jan 2022-02-23 15:47:38 -08:00
commit a4fe3d9e2e
20 changed files with 291 additions and 224 deletions
Download patch file Download diff file Expand all files Collapse all files

59
masque/repetition.py
View file

@ -7,11 +7,11 @@ from typing import Union, Dict, Optional, Sequence, Any
import copy
from abc import ABCMeta, abstractmethod
import numpy # type: ignore
from numpy.typing import ArrayLike
import numpy
from numpy.typing import ArrayLike, NDArray
from .error import PatternError
from .utils import rotation_matrix_2d, vector2, AutoSlots
from .utils import rotation_matrix_2d, AutoSlots
from .traits import LockableImpl, Copyable, Scalable, Rotatable, Mirrorable
@ -23,7 +23,7 @@ class Repetition(Copyable, Rotatable, Mirrorable, Scalable, metaclass=ABCMeta):
@property
@abstractmethod
def displacements(self) -> numpy.ndarray:
def displacements(self) -> NDArray[numpy.float64]:
"""
An Nx2 ndarray specifying all offsets generated by this repetition
"""
@ -44,7 +44,7 @@ class Grid(LockableImpl, Repetition, metaclass=AutoSlots):
'_a_count',
'_b_count')
_a_vector: numpy.ndarray
_a_vector: NDArray[numpy.float64]
""" Vector `[x, y]` specifying the first lattice vector of the grid.
Specifies center-to-center spacing between adjacent elements.
"""
@ -52,7 +52,7 @@ class Grid(LockableImpl, Repetition, metaclass=AutoSlots):
_a_count: int
""" Number of instances along the direction specified by the `a_vector` """
_b_vector: Optional[numpy.ndarray]
_b_vector: Optional[NDArray[numpy.float64]]
""" Vector `[x, y]` specifying a second lattice vector for the grid.
Specifies center-to-center spacing between adjacent elements.
Can be `None` for a 1D array.
@ -100,8 +100,8 @@ class Grid(LockableImpl, Repetition, metaclass=AutoSlots):
raise PatternError(f'Repetition has too-small b_count: {b_count}')
object.__setattr__(self, 'locked', False)
self.a_vector = a_vector
self.b_vector = b_vector
self.a_vector = a_vector # type: ignore # setter handles type conversion
self.b_vector = b_vector # type: ignore # setter handles type conversion
self.a_count = a_count
self.b_count = b_count
self.locked = locked
@ -122,11 +122,11 @@ class Grid(LockableImpl, Repetition, metaclass=AutoSlots):
# a_vector property
@property
def a_vector(self) -> numpy.ndarray:
def a_vector(self) -> NDArray[numpy.float64]:
return self._a_vector
@a_vector.setter
def a_vector(self, val: vector2):
def a_vector(self, val: ArrayLike) -> None:
if not isinstance(val, numpy.ndarray):
val = numpy.array(val, dtype=float)
@ -136,11 +136,11 @@ class Grid(LockableImpl, Repetition, metaclass=AutoSlots):
# b_vector property
@property
def b_vector(self) -> numpy.ndarray:
def b_vector(self) -> Optional[NDArray[numpy.float64]]:
return self._b_vector
@b_vector.setter
def b_vector(self, val: vector2):
def b_vector(self, val: ArrayLike) -> None:
if not isinstance(val, numpy.ndarray):
val = numpy.array(val, dtype=float, copy=True)
@ -154,7 +154,7 @@ class Grid(LockableImpl, Repetition, metaclass=AutoSlots):
return self._a_count
@a_count.setter
def a_count(self, val: int):
def a_count(self, val: int) -> None:
if val != int(val):
raise PatternError('a_count must be convertable to an int!')
self._a_count = int(val)
@ -165,13 +165,16 @@ class Grid(LockableImpl, Repetition, metaclass=AutoSlots):
return self._b_count
@b_count.setter
def b_count(self, val: int):
def b_count(self, val: int) -> None:
if val != int(val):
raise PatternError('b_count must be convertable to an int!')
self._b_count = int(val)
@property
def displacements(self) -> numpy.ndarray:
def displacements(self) -> NDArray[numpy.float64]:
if self.b_vector is None:
return numpy.arange(self.a_count)[:, None] * self.a_vector[None, :]
aa, bb = numpy.meshgrid(numpy.arange(self.a_count), numpy.arange(self.b_count), indexing='ij')
return (aa.flatten()[:, None] * self.a_vector[None, :]
+ bb.flatten()[:, None] * self.b_vector[None, :]) # noqa
@ -207,7 +210,7 @@ class Grid(LockableImpl, Repetition, metaclass=AutoSlots):
self.b_vector[1 - axis] *= -1
return self
def get_bounds(self) -> Optional[numpy.ndarray]:
def get_bounds(self) -> Optional[NDArray[numpy.float64]]:
"""
Return a `numpy.ndarray` containing `[[x_min, y_min], [x_max, y_max]]`, corresponding to the
extent of the `Grid` in each dimension.
@ -216,7 +219,7 @@ class Grid(LockableImpl, Repetition, metaclass=AutoSlots):
`[[x_min, y_min], [x_max, y_max]]` or `None`
"""
a_extent = self.a_vector * self.a_count
b_extent = self.b_vector * self.b_count if self.b_count != 0 else 0
b_extent = self.b_vector * self.b_count if (self.b_vector is not None) else 0 # type: Union[NDArray[numpy.float64], float]
corners = ((0, 0), a_extent, b_extent, a_extent + b_extent)
xy_min = numpy.min(corners, axis=0)
@ -296,24 +299,26 @@ class Arbitrary(LockableImpl, Repetition, metaclass=AutoSlots):
`[[x0, y0], [x1, y1], ...]`
"""
_displacements: numpy.ndarray
_displacements: NDArray[numpy.float64]
""" List of vectors `[[x0, y0], [x1, y1], ...]` specifying the offsets
of the instances.
"""
@property
def displacements(self) -> numpy.ndarray:
def displacements(self) -> Any: # TODO: mypy#3004 NDArray[numpy.float64]:
return self._displacements
@displacements.setter
def displacements(self, val: ArrayLike):
val = numpy.array(val, float)
val = numpy.sort(val.view([('', val.dtype)] * val.shape[1]), 0).view(val.dtype) # sort rows
self._displacements = val
def displacements(self, val: ArrayLike) -> None:
vala: NDArray[numpy.float64] = numpy.array(vala, dtype=float)
vala = numpy.sort(vala.view([('', vala.dtype)] * vala.shape[1]), 0).view(vala.dtype) # sort rows
self._displacements = vala
def __init__(self,
displacements: ArrayLike,
locked: bool = False,):
def __init__(
self,
displacements: ArrayLike,
locked: bool = False,
) -> None:
"""
Args:
displacements: List of vectors (Nx2 ndarray) specifying displacements.
@ -383,7 +388,7 @@ class Arbitrary(LockableImpl, Repetition, metaclass=AutoSlots):
self.displacements[1 - axis] *= -1
return self
def get_bounds(self) -> Optional[numpy.ndarray]:
def get_bounds(self) -> Optional[NDArray[numpy.float64]]:
"""
Return a `numpy.ndarray` containing `[[x_min, y_min], [x_max, y_max]]`, corresponding to the
extent of the `displacements` in each dimension.