Skip to content

Commit

Permalink
Numpy 2.0 compatibility
Browse files Browse the repository at this point in the history
  • Loading branch information
@hexane360
hexane360 committed Jun 20, 2024
1 parent 884751c commit 6203439
Show file tree
Hide file tree
Showing 8 changed files with 41 additions and 41 deletions.
16 changes: 8 additions & 8 deletions atomlib/cell.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -91,12 +91,12 @@ def metric(self) -> LinearTransform3D:
return ortho.T @ ortho

@property
def cell_size(self) -> NDArray[numpy.float_]:
def cell_size(self) -> NDArray[numpy.float64]:
"""Unit cell size."""
return self.get_cell()._cell_size

@property
def cell_angle(self) -> NDArray[numpy.float_]:
def cell_angle(self) -> NDArray[numpy.float64]:
"""Unit cell angles, in radians."""
return self.get_cell()._cell_angle

Expand All @@ -111,7 +111,7 @@ def pbc(self) -> NDArray[numpy.bool_]:
return self.get_cell()._pbc

@property
def ortho_size(self) -> NDArray[numpy.float_]:
def ortho_size(self) -> NDArray[numpy.float64]:
"""
Return size of orthogonal unit cell.
Expand All @@ -120,7 +120,7 @@ def ortho_size(self) -> NDArray[numpy.float_]:
return self.cell_size * numpy.diag(self.ortho.inner)

@property
def box_size(self) -> NDArray[numpy.float_]:
def box_size(self) -> NDArray[numpy.float64]:
"""
Return size of the cell box.
Expand Down Expand Up @@ -369,11 +369,11 @@ def with_cell(self: Cell, cell: Cell) -> Cell:
Orthogonalization transformation. Skews but does not scale the crystal axes to cartesian axes.
"""

_cell_size: NDArray[numpy.float_]
_cell_size: NDArray[numpy.float64]
"""Unit cell size."""
_cell_angle: NDArray[numpy.float_] = field(default_factory=lambda: numpy.full(3, numpy.pi/2.))
_cell_angle: NDArray[numpy.float64] = field(default_factory=lambda: numpy.full(3, numpy.pi/2.))
"""Unit cell angles, in radians."""
_n_cells: NDArray[numpy.int_] = field(default_factory=lambda: numpy.ones(3, numpy.int_))
_n_cells: NDArray[numpy.int64] = field(default_factory=lambda: numpy.ones(3, numpy.int64))
"""Number of unit cells."""
_pbc: NDArray[numpy.bool_] = field(default_factory=lambda: numpy.ones(3, numpy.bool_))
"""Flags indicating the presence of periodic boundary conditions along each axis."""
Expand All @@ -387,7 +387,7 @@ def __init__(self, *,
object.__setattr__(self, '_ortho', LinearTransform3D() if ortho is None else ortho)
object.__setattr__(self, '_cell_size', to_vec3(cell_size))
object.__setattr__(self, '_cell_angle', numpy.full(3, numpy.pi/2.) if cell_angle is None else to_vec3(cell_angle))
object.__setattr__(self, '_n_cells', numpy.ones(3, numpy.int_) if n_cells is None else to_vec3(n_cells, numpy.int_))
object.__setattr__(self, '_n_cells', numpy.ones(3, numpy.int_) if n_cells is None else to_vec3(n_cells, numpy.int64))
object.__setattr__(self, '_pbc', numpy.ones(3, numpy.bool_) if pbc is None else to_vec3(pbc, numpy.bool_))

@staticmethod
Expand Down
14 changes: 7 additions & 7 deletions atomlib/defect.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@
from .vec import norm, dot, perp, split_arr, polygon_solid_angle, polygon_winding


def ellip_pi(n: NDArray[numpy.float_], m: NDArray[numpy.float_]) -> NDArray[numpy.float_]:
def ellip_pi(n: NDArray[numpy.float64], m: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
"""
Complete elliptic integral of the third kind, $\\Pi(n | m)$.
Expand Down Expand Up @@ -267,7 +267,7 @@ def disloc_screw(atoms: HasAtomsT, center: VecLike, b: VecLike, cut: t.Optional[
cut = to_vec3([1., 0., 0.])
else:
# otherwise find plane by rotating around 111
cut = cast(NDArray[numpy.float_], LinearTransform3D.rotate([1., 1., 1.], 2*numpy.pi/3).transform(t))
cut = cast(NDArray[numpy.float64], LinearTransform3D.rotate([1., 1., 1.], 2*numpy.pi/3).transform(t))
else:
cut = to_vec3(cut)
cut /= norm(cut)
Expand Down Expand Up @@ -408,7 +408,7 @@ def disloc_poly_z(atoms: HasAtomsT, b: VecLike, poly: ArrayLike, center: t.Optio
raise ValueError(f"Expected a 2d polygon. Instead got shape {poly.shape}")

frame = atoms.get_atoms('local')
coords: NDArray[numpy.float_] = frame.coords() - center
coords: NDArray[numpy.float64] = frame.coords() - center

branch = None
d = numpy.dot(b_vec, [0, 0, 1])
Expand Down Expand Up @@ -441,8 +441,8 @@ def disloc_poly_z(atoms: HasAtomsT, b: VecLike, poly: ArrayLike, center: t.Optio
return atoms.with_atoms(frame.with_coords(coords + disp + center), 'local')


def _poly_disp_z(pts: NDArray[numpy.float_], b_vec: NDArray[numpy.float_], poly: NDArray[numpy.float_], *,
poisson: float = 0.25, branch: t.Optional[numpy.ndarray] = None) -> NDArray[numpy.float_]:
def _poly_disp_z(pts: NDArray[numpy.float64], b_vec: NDArray[numpy.float64], poly: NDArray[numpy.float64], *,
poisson: float = 0.25, branch: t.Optional[numpy.ndarray] = None) -> NDArray[numpy.float64]:

if branch is None:
branch = numpy.array(1.)
Expand All @@ -458,15 +458,15 @@ def _poly_disp_z(pts: NDArray[numpy.float_], b_vec: NDArray[numpy.float_], poly:
e2 = numpy.cross(eta, r)
e2 /= numpy.linalg.norm(e2, axis=-1, keepdims=True)

def _disp(r: NDArray[numpy.float_]) -> NDArray[numpy.float_]:
def _disp(r: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
r_norm = numpy.linalg.norm(r, axis=-1, keepdims=True)
disps = (1-2*poisson)*numpy.cross(b_vec, eta) * numpy.log(r_norm + dot(r, eta)) - dot(b_vec, e2) * numpy.cross(r / r_norm, e2)
return numpy.sum(disps, axis=-2)

return b_vec * omega[:, None] / (4*numpy.pi) + 1/(8*numpy.pi*(1-poisson)) * (_disp(r_n) - _disp(r))


def _loop_disp_z(pts: NDArray[numpy.float_], b_vec: numpy.ndarray, loop_r: float, *,
def _loop_disp_z(pts: NDArray[numpy.float64], b_vec: numpy.ndarray, loop_r: float, *,
poisson: float = 0.25, branch: t.Optional[numpy.ndarray] = None) -> numpy.ndarray:
from scipy.special import ellipk, ellipe

Expand Down
4 changes: 2 additions & 2 deletions atomlib/io/xyz.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -148,15 +148,15 @@ def write(self, file: FileOrPath, fmt: XYZFormat = 'exyz'):
).strip() + '\n' for row in self.atoms.select(('symbol', 'coords')).rows()
)

def cell_matrix(self) -> t.Optional[NDArray[numpy.float_]]:
def cell_matrix(self) -> t.Optional[NDArray[numpy.float64]]:
if (s := self.params.get('Lattice')) is None:
return None

try:
items = list(map(float, s.split()))
if not len(items) == 9:
raise ValueError("Invalid length")
return numpy.array(items, dtype=numpy.float_).reshape((3, 3)).T
return numpy.array(items, dtype=numpy.float64).reshape((3, 3)).T
except ValueError:
warnings.warn(f"Warning: Invalid format for key 'Lattice=\"{s}\"'")
return None
Expand Down
14 changes: 7 additions & 7 deletions atomlib/transform.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -434,7 +434,7 @@ def __matmul__(self, other: t.Union[Transform3D, ArrayLike, BBox3D]): # type: i
class LinearTransform3D(AffineTransform3D):
def __init__(self, array: t.Optional[ArrayLike] = None):
if array is None:
array = numpy.eye(3, dtype=numpy.float_)
array = numpy.eye(3, dtype=numpy.float64)
self.inner = numpy.broadcast_to(array, (3, 3))

@property
Expand Down Expand Up @@ -514,11 +514,11 @@ def mirror(self, a: t.Union[Num, VecLike],
Can be called as a classmethod or instance method.
"""
if isinstance(a, t.Sized):
v = numpy.array(numpy.broadcast_to(a, 3), dtype=numpy.float_)
v = numpy.array(numpy.broadcast_to(a, 3), dtype=numpy.float64)
if b is not None or c is not None:
raise ValueError("mirror() must be passed a sequence or three numbers.")
else:
v = numpy.array([a, b, c], dtype=numpy.float_)
v = numpy.array([a, b, c], dtype=numpy.float64)
v /= numpy.linalg.norm(v)
mirror = numpy.eye(3) - 2 * numpy.outer(v, v)
return LinearTransform3D(mirror @ self.inner)
Expand Down Expand Up @@ -550,7 +550,7 @@ def rotate(self, v: VecLike, theta: Num) -> LinearTransform3D:
Can be called as a classmethod or instance method.
"""
theta = float(theta)
v = numpy.array(numpy.broadcast_to(v, (3,)), dtype=numpy.float_)
v = numpy.array(numpy.broadcast_to(v, (3,)), dtype=numpy.float64)
l = numpy.linalg.norm(v)
if numpy.isclose(l, 0.):
if numpy.isclose(theta, 0.):
Expand All @@ -562,7 +562,7 @@ def rotate(self, v: VecLike, theta: Num) -> LinearTransform3D:
# Rodrigues rotation formula
w = numpy.array([[ 0., -v[2], v[1]],
[ v[2], 0., -v[0]],
[-v[1], v[0], 0.]], dtype=numpy.float_)
[-v[1], v[0], 0.]], dtype=numpy.float64)
# I + sin(t) W + (1 - cos(t)) W^2 = I + sin(t) W + 2*sin^2(t/2) W^2
a = numpy.eye(3) + numpy.sin(theta) * w + 2 * (numpy.sin(theta / 2)**2) * w @ w
return LinearTransform3D(a @ self.inner)
Expand All @@ -576,13 +576,13 @@ def rotate_euler(self, x: Num = 0., y: Num = 0., z: Num = 0.) -> LinearTransform
Can be called as a classmethod or instance method.
"""
angles = numpy.array([x, y, z], dtype=numpy.float_)
angles = numpy.array([x, y, z], dtype=numpy.float64)
c, s = numpy.cos(angles), numpy.sin(angles)
a = numpy.array([
[c[1]*c[2], s[0]*s[1]*c[2] - c[0]*s[2], c[0]*s[1]*c[2] + s[0]*s[2]],
[c[1]*s[2], s[0]*s[1]*s[2] + c[0]*c[2], c[0]*s[1]*s[2] - s[0]*c[2]],
[-s[1], s[0]*c[1], c[0]*c[1]],
], dtype=numpy.float_)
], dtype=numpy.float64)
return LinearTransform3D(a @ self.inner)

@opt_classmethod
Expand Down
4 changes: 2 additions & 2 deletions atomlib/types.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@


@t.overload
def to_vec3(v: VecLike, dtype: None = None) -> NDArray[numpy.float_]:
def to_vec3(v: VecLike, dtype: None = None) -> NDArray[numpy.float64]:
...

@t.overload
Expand All @@ -67,7 +67,7 @@ def to_vec3(v: VecLike, dtype: t.Optional[t.Type[numpy.generic]] = None) -> NDAr
"""

try:
v = numpy.broadcast_to(v, (3,)).astype(dtype or numpy.float_)
v = numpy.broadcast_to(v, (3,)).astype(dtype or numpy.float64)
except (ValueError, TypeError):
raise TypeError("Expected a vector of 3 elements.") from None
return v
Expand Down
18 changes: 9 additions & 9 deletions atomlib/vec.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -48,7 +48,7 @@ def split_arr(a: NDArray[ScalarT], axis: int = 0) -> t.Iterator[NDArray[ScalarT]


def polygon_solid_angle(poly: ArrayLike, pts: t.Optional[ArrayLike] = None,
winding: t.Optional[ArrayLike] = None) -> NDArray[numpy.float_]:
winding: t.Optional[ArrayLike] = None) -> NDArray[numpy.float64]:
"""
Return the signed solid angle of the polygon ``poly`` in the xy plane, as viewed from ``pts``.
Expand All @@ -57,8 +57,8 @@ def polygon_solid_angle(poly: ArrayLike, pts: t.Optional[ArrayLike] = None,
Returns a ndarray of shape ``broadcast(poly.shape[:-2], pts.shape[:-1])``
"""
poly = numpy.atleast_2d(poly).astype(numpy.float_)
pts = (numpy.array([0., 0., 0.]) if pts is None else numpy.atleast_1d(pts)).astype(numpy.float_)
poly = numpy.atleast_2d(poly).astype(numpy.float64)
pts = (numpy.array([0., 0., 0.]) if pts is None else numpy.atleast_1d(pts)).astype(numpy.float64)

if poly.shape[-1] == 3:
raise ValueError("Only 2d polygons are supported.")
Expand All @@ -79,7 +79,7 @@ def polygon_solid_angle(poly: ArrayLike, pts: t.Optional[ArrayLike] = None,
# normalize polygon points to unit sphere
numpy.divide(poly, numpy.linalg.norm(poly, axis=-1, keepdims=True), out=poly)

def _dot(v1: NDArray[numpy.float_], v2: NDArray[numpy.float_]) -> NDArray[numpy.float_]:
def _dot(v1: NDArray[numpy.float64], v2: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
return numpy.add.reduce(v1 * v2, axis=-1)

# next and previous points in polygon
Expand All @@ -95,7 +95,7 @@ def _dot(v1: NDArray[numpy.float_], v2: NDArray[numpy.float_]) -> NDArray[numpy.
return angle - 2*numpy.pi*winding


def polygon_winding(poly: ArrayLike, pt: t.Optional[ArrayLike] = None) -> NDArray[numpy.int_]:
def polygon_winding(poly: ArrayLike, pt: t.Optional[ArrayLike] = None) -> NDArray[numpy.int64]:
"""
Return the winding number of the given 2d polygon ``poly`` around the point ``pt``.
If ``pt`` is not specified, return the polygon's total winding number (turning number).
Expand All @@ -105,7 +105,7 @@ def polygon_winding(poly: ArrayLike, pt: t.Optional[ArrayLike] = None) -> NDArra
poly = numpy.atleast_2d(poly)
if poly.dtype == object:
raise ValueError("Ragged arrays not supported.")
poly = poly.astype(numpy.float_)
poly = poly.astype(numpy.float64)

if pt is None:
# return polygon's total winding number (turning number)
Expand All @@ -120,7 +120,7 @@ def polygon_winding(poly: ArrayLike, pt: t.Optional[ArrayLike] = None) -> NDArra
numpy.isclose(poly[..., 1], 0., atol=1e-10))
poly = poly[~zero_pts]

pt = numpy.atleast_1d(pt)[..., None, :].astype(numpy.float_)
pt = numpy.atleast_1d(pt)[..., None, :].astype(numpy.float64)

# shift the polygon's origin to `pt`.
poly = poly - pt
Expand All @@ -135,7 +135,7 @@ def polygon_winding(poly: ArrayLike, pt: t.Optional[ArrayLike] = None) -> NDArra
down_crossing = (y > 0) & (yn <= 0) & (x_pos < 0)

# reduce and return
return numpy.sum(up_crossing, axis=-1) - numpy.sum(down_crossing, axis=-1)
return (numpy.sum(up_crossing, axis=-1) - numpy.sum(down_crossing, axis=-1)).astype(numpy.int64)


WindingRule = t.Literal['nonzero', 'evenodd', 'positive', 'negative']
Expand Down Expand Up @@ -174,7 +174,7 @@ def in_polygon(poly: numpy.ndarray, pt: t.Optional[numpy.ndarray] = None, *,
"'nonzero', 'evenodd', 'positive', or 'negative'.")


def reduce_vec(arr: ArrayLike, max_denom: int = 10000) -> NDArray[numpy.int_]:
def reduce_vec(arr: ArrayLike, max_denom: int = 10000) -> NDArray[numpy.int64]:
"""
Reduce a crystallographic vector (int or float) to lowest common terms.
Example: reduce_vec([3, 3, 3]) = [1, 1, 1]
Expand Down
10 changes: 5 additions & 5 deletions atomlib/visualize/__init__.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ def __init__(self, elem: int, fc=None, s: float = 2., **kwargs):

def get_zone(atoms: HasAtoms, zone: t.Optional[VecLike] = None,
plane: t.Optional[VecLike] = None,
default: t.Optional[VecLike] = None) -> NDArray[numpy.float_]:
default: t.Optional[VecLike] = None) -> NDArray[numpy.float64]:
if zone is not None and plane is not None:
raise ValueError("'zone' and 'plane' can't both be specified.")
if plane is not None:
Expand All @@ -117,13 +117,13 @@ def get_zone(atoms: HasAtoms, zone: t.Optional[VecLike] = None,
raise NotImplementedError()
zone = plane
if zone is not None:
return numpy.broadcast_to(zone, 3)
return numpy.broadcast_to(zone, 3).astype(numpy.float64)
if default is not None:
return numpy.broadcast_to(default, 3)
return numpy.array([0., 0., 1.])
return numpy.broadcast_to(default, 3).astype(numpy.float64)
return numpy.array([0., 0., 1.], dtype=numpy.float64)


def get_plot_radii(atoms: HasAtoms, min_r: t.Optional[float] = 1.0, style: AtomStyle = 'small') -> NDArray[numpy.float_]:
def get_plot_radii(atoms: HasAtoms, min_r: t.Optional[float] = 1.0, style: AtomStyle = 'small') -> NDArray[numpy.float64]:
radii = get_radius(atoms['elem']).to_numpy()
if style == 'small':
radii = radii * 0.6
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ classifiers = [
requires-python = ">=3.9"
dependencies = [
"click~=8.1", # for cli
"numpy~=1.22",
"numpy>=1.22,<2.3.0", # tested on 2.0.0
"scipy~=1.8",
"polars~=0.20.9",
"matplotlib~=3.5",
Expand Down

0 comments on commit 6203439

Please sign in to comment.