atomlib.vec

Helper functions for spatial vectors.

WindingRule module-attribute

WindingRule = Literal[
    "nonzero", "evenodd", "positive", "negative"
]

to_vec3

to_vec3(v: VecLike, dtype: None = None) -> NDArray[float64]

to_vec3(
    v: VecLike, dtype: Type[ScalarT]
) -> NDArray[ScalarT]

to_vec3(
    v: VecLike, dtype: Optional[Type[generic]] = None
) -> NDArray[generic]

Broadcast and coerce v to a Vec3 of type dtype.

Source code in atomlib/types.py

def to_vec3(v: VecLike, dtype: t.Optional[t.Type[numpy.generic]] = None) -> NDArray[numpy.generic]:
    """
    Broadcast and coerce `v` to a [`Vec3`][atomlib.types.Vec3] of type `dtype`.
    """

    try:
        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

dot

dot(
    v1: ArrayLike,
    v2: ArrayLike,
    axis: int = -1,
    keepdims: bool = True,
) -> NDArray[floating]

Take the dot product between two vectors, along axis axis.

Source code in atomlib/vec.py

def dot(v1: ArrayLike, v2: ArrayLike, axis: int = -1, keepdims: bool = True) -> NDArray[numpy.floating]:
    """
    Take the dot product between two vectors, along axis `axis`.
    """
    return numpy.add.reduce(numpy.atleast_1d(v1) * numpy.atleast_1d(v2), axis=axis, keepdims=keepdims)

norm

norm(v: ArrayLike) -> floating

Return the norm of the vector v.

Source code in atomlib/vec.py

def norm(v: ArrayLike) -> numpy.floating:
    """
    Return the norm of the vector `v`.
    """
    return numpy.linalg.norm(v)

perp

perp(v1: ArrayLike, v2: ArrayLike) -> NDArray[floating]

Return the component of v1 perpendicular to v2.

Source code in atomlib/vec.py

def perp(v1: ArrayLike, v2: ArrayLike) -> NDArray[numpy.floating]:
    """Return the component of `v1` perpendicular to `v2`."""
    v1 = numpy.atleast_1d(v1)
    v2 = numpy.atleast_1d(v2)
    v2 /= norm(v2)
    return v1 - v2 * dot(v1, v2)

para

para(v1: ArrayLike, v2: ArrayLike) -> NDArray[floating]

Return the component of v1 parallel to v2.

Source code in atomlib/vec.py

def para(v1: ArrayLike, v2: ArrayLike) -> NDArray[numpy.floating]:
    """Return the component of `v1` parallel to `v2`."""
    v1 = numpy.atleast_1d(v1)
    v2 = numpy.atleast_1d(v2)
    v2 /= norm(v2)
    return v2 * dot(v1, v2)

is_diagonal

is_diagonal(matrix: ndarray, tol: float = 1e-10) -> bool

Return if matrix is diagonal, to tolerance tol.

Source code in atomlib/vec.py

def is_diagonal(matrix: numpy.ndarray, tol: float = 1e-10) -> bool:
    """
    Return if `matrix` is diagonal, to tolerance `tol`.
    """
    d = matrix.shape[0]
    assert matrix.shape == (d, d)
    p, q = matrix.strides
    offdiag = numpy.lib.stride_tricks.as_strided(matrix[:, 1:], (d-1, d), (p+q, q))
    return bool((numpy.abs(offdiag) < tol).all())

split_arr

split_arr(
    a: NDArray[ScalarT], axis: int = 0
) -> Iterator[NDArray[ScalarT]]

Split the array along the axis axis.

Source code in atomlib/vec.py

def split_arr(a: NDArray[ScalarT], axis: int = 0) -> t.Iterator[NDArray[ScalarT]]:
    """
    Split the array along the axis `axis`.
    """
    return (numpy.squeeze(sub_a, axis) for sub_a in numpy.split(a, a.shape[axis], axis))

polygon_solid_angle

polygon_solid_angle(
    poly: ArrayLike,
    pts: Optional[ArrayLike] = None,
    winding: Optional[ArrayLike] = None,
) -> NDArray[float64]

Return the signed solid angle of the polygon poly in the xy plane, as viewed from pts.

PARAMETER	DESCRIPTION
poly	Polygon(s) to compute the angle of, array of shape (..., N, 2) TYPE: ArrayLike
pts	Point(s) to view the polygons from, array of shape (..., 3) TYPE: Optional[ArrayLike] DEFAULT: None

RETURNS	DESCRIPTION
NDArray[float64]	A ndarray of shape broadcast(poly.shape[:-2], pts.shape[:-1]), containing signed solid angles.

Source code in atomlib/vec.py

def polygon_solid_angle(poly: ArrayLike, pts: t.Optional[ArrayLike] = None,
                        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`.

    Args:
      poly: Polygon(s) to compute the angle of, array of shape `(..., N, 2)`
      pts: Point(s) to view the polygons from, array of shape `(..., 3)`

    Returns:
      A ndarray of shape `broadcast(poly.shape[:-2], pts.shape[:-1])`, containing signed solid angles.
    """
    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.")
    if poly.shape[-1] != 2:
        raise ValueError("`poly` must be a list of 2d points.")
    if winding is None:
        # calculate winding
        winding = polygon_winding(poly)
    else:
        winding = numpy.asarray(winding, dtype=int)
    # extend to 3d
    poly = numpy.concatenate((poly, numpy.zeros_like(poly, shape=(*poly.shape[:-1], 1))), axis=-1)

    if pts.shape[-1] != 3:
        raise ValueError("`pts` must be a list of 3d points.")

    poly = poly - pts[..., 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.float64], v2: NDArray[numpy.float64]) -> NDArray[numpy.float64]:
        return numpy.add.reduce(v1 * v2, axis=-1)

    # next and previous points in polygon
    poly_n = numpy.roll(poly, -1, axis=-2)
    poly_p = numpy.roll(poly, 1, axis=-2)

    # spherical angle is 2*pi - sum(atan2(-|v1v2v3|, v1 dot v2 * v2 dot v3 - v1 dot v3))
    angles = numpy.arctan2(_dot(poly_p, numpy.cross(poly, poly_n)), _dot(poly_p, poly) * _dot(poly, poly_n) - _dot(poly_p, poly_n))
    angle = numpy.sum(angles, axis=-1)

    # when winding is nonzero, we have to offset the calculated angle by the angle created by winding.
    numpy.mod(angle, 4*numpy.pi*winding, out=angle, where=(winding != 0))
    return angle - 2*numpy.pi*winding

polygon_winding

polygon_winding(
    poly: ArrayLike, pt: Optional[ArrayLike] = None
) -> NDArray[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).

Vectorized. CCW winding is defined as positive.

Source code in atomlib/vec.py

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).

    Vectorized. CCW winding is defined as positive.
    """
    poly = numpy.atleast_2d(poly)
    if poly.dtype == object:
        raise ValueError("Ragged arrays not supported.")
    poly = poly.astype(numpy.float64)

    if pt is None:
        # return polygon's total winding number (turning number)
        poly_next = numpy.roll(poly, -1, axis=-2)
        # equivalent to the turning number of velocity vectors (difference vectors)
        poly = poly_next - poly
        # about the origin
        pt = numpy.array([0., 0.])

        # remove points at the origin (duplicate points)
        zero_pts = (numpy.isclose(poly[..., 0], 0., atol=1e-10) & 
                    numpy.isclose(poly[..., 1], 0., atol=1e-10))
        poly = poly[~zero_pts]

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

    # shift the polygon's origin to `pt`.
    poly = poly - pt
    poly_next = numpy.roll(poly, -1, axis=-2)
    (x, y) = split_arr(poly, axis=-1)
    (xn, yn) = split_arr(poly_next, axis=-1)

    # |p1 cross (p2 - p1)| -> (p2 - p1) to right or left of origin
    x_pos = x*(yn - y) - y*(xn - x)  # type: ignore
    # count up crossings and down crossings
    up_crossing = (y <= 0) & (yn > 0) & (x_pos > 0)
    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)).astype(numpy.int64)

in_polygon

in_polygon(
    poly: ndarray,
    pt: Literal[None] = None,
    *,
    rule: WindingRule = "evenodd"
) -> Callable[[ndarray], NDArray[bool_]]

in_polygon(
    poly: ndarray,
    pt: ndarray,
    *,
    rule: WindingRule = "evenodd"
) -> NDArray[bool_]

in_polygon(
    poly: ndarray,
    pt: Optional[ndarray] = None,
    *,
    rule: WindingRule = "evenodd"
) -> Union[
    NDArray[bool_], Callable[[ndarray], NDArray[bool_]]
]

Return whether pt is in poly, under the given winding rule. In the one-argument form, return a closure which tests poly for the given point.

Source code in atomlib/vec.py

def in_polygon(poly: numpy.ndarray, pt: t.Optional[numpy.ndarray] = None, *,
               rule: WindingRule = 'evenodd') -> t.Union[NDArray[numpy.bool_], t.Callable[[numpy.ndarray], NDArray[numpy.bool_]]]:
    """
    Return whether `pt` is in `poly`, under the given winding rule.
    In the one-argument form, return a closure which tests `poly` for the given point.
    """
    if pt is None:
        return lambda pt: in_polygon(poly, pt, rule=rule)
    winding = polygon_winding(poly, pt)

    rule = t.cast(WindingRule, rule.lower())
    if rule == 'nonzero':
        return winding.astype(numpy.bool_)
    elif rule == 'evenodd':
        return (winding & 1) > 0
    elif rule == 'positive':
        return winding > 0
    elif rule == 'negative':
        return winding < 0
    raise ValueError(f"Unknown winding rule '{rule}'. Expected one of "
                     "'nonzero', 'evenodd', 'positive', or 'negative'.")

reduce_vec

reduce_vec(
    arr: ArrayLike, max_denom: int = 10000
) -> NDArray[int64]

Reduce a crystallographic vector (int or float) to lowest common terms.

Examples

>>> reduce_vec([3, 6, 9])
[1, 2, 3]
>>> reduce_vec([0.5, 0.25, 0.25])
[2, 1, 1]

Source code in atomlib/vec.py

def reduce_vec(arr: ArrayLike, max_denom: int = 10000) -> NDArray[numpy.int64]:
    """
    Reduce a crystallographic vector (int or float) to lowest common terms.

    # Examples
    ```python
    >>> reduce_vec([3, 6, 9])
    [1, 2, 3]
    >>> reduce_vec([0.5, 0.25, 0.25])
    [2, 1, 1]
    ```
    """

    a = numpy.atleast_1d(arr)
    if not numpy.issubdtype(a.dtype, numpy.floating):
        return a // numpy.gcd.reduce(a, axis=-1, keepdims=True)

    a = a / numpy.max(numpy.abs(a))

    n = numpy.empty(shape=a.shape, dtype=numpy.int64)
    d = numpy.empty(shape=a.shape, dtype=numpy.int64)
    with numpy.nditer([a, n, d], ['refs_ok'], [['readonly'], ['writeonly'], ['writeonly']]) as it:  # type: ignore
        for (v, n_, d_) in it:
            (n_[()], d_[()]) = Fraction(float(v)).limit_denominator(max_denom).as_integer_ratio()

    # reduce to common denominator
    factors = numpy.lcm.reduce(d, axis=-1, keepdims=True) // d
    n *= factors
    # and then reduce numerators
    return n // numpy.gcd.reduce(n, axis=-1, keepdims=True)

miller_4_to_3_vec

miller_4_to_3_vec(
    a: NDArray[number],
    reduce: bool = True,
    max_denom: int = 10000,
) -> NDArray[number]

Convert a vector in 4-axis Miller-Bravais notation to 3-axis Miller notation.

Source code in atomlib/vec.py

def miller_4_to_3_vec(a: NDArray[numpy.number], reduce: bool = True, max_denom: int = 10000) -> NDArray[numpy.number]:
    """Convert a vector in 4-axis Miller-Bravais notation to 3-axis Miller notation."""
    a = numpy.atleast_1d(a)
    assert a.shape[-1] == 4
    U, V, T, W = numpy.split(a, 4, axis=-1)
    assert numpy.allclose(-T, U + V, equal_nan=True)
    out = numpy.concatenate((2*U + V, 2*V + U, W), axis=-1)
    return reduce_vec(out, max_denom) if reduce else out

miller_3_to_4_vec

miller_3_to_4_vec(
    a: NDArray[number],
    reduce: bool = True,
    max_denom: int = 10000,
) -> NDArray[number]

Convert a vector in 3-axis Miller notation to 4-axis Miller-Bravais notation.

Source code in atomlib/vec.py

def miller_3_to_4_vec(a: NDArray[numpy.number], reduce: bool = True, max_denom: int = 10000) -> NDArray[numpy.number]:
    """Convert a vector in 3-axis Miller notation to 4-axis Miller-Bravais notation."""
    a = numpy.atleast_1d(a)
    assert a.shape[-1] == 3
    u, v, w = numpy.split(a, 3, axis=-1)
    U = 2*u - v
    V = 2*v - u
    W = 3*w
    out = numpy.concatenate((U, V, -(U + V), W), axis=-1)
    return reduce_vec(out, max_denom) if reduce else out

miller_4_to_3_plane

miller_4_to_3_plane(
    a: NDArray[number],
    reduce: bool = True,
    max_denom: int = 10000,
) -> NDArray[number]

Convert a plane in 4-axis Miller-Bravais notation to 3-axis Miller notation.

Source code in atomlib/vec.py

def miller_4_to_3_plane(a: NDArray[numpy.number], reduce: bool = True, max_denom: int = 10000) -> NDArray[numpy.number]:
    """Convert a plane in 4-axis Miller-Bravais notation to 3-axis Miller notation."""
    a = numpy.atleast_1d(a)
    assert a.shape[-1] == 4
    h, k, i, l = numpy.split(a, 4, axis=-1)  # noqa: E741
    assert numpy.allclose(-i, h + k, equal_nan=True)
    out = numpy.concatenate((h, k, l), axis=-1)
    return reduce_vec(out, max_denom) if reduce else out

miller_3_to_4_plane

miller_3_to_4_plane(
    a: NDArray[number],
    reduce: bool = True,
    max_denom: int = 10000,
) -> NDArray[number]

Convert a plane in 3-axis Miller notation to 4-axis Miller-Bravais notation.

Source code in atomlib/vec.py

def miller_3_to_4_plane(a: NDArray[numpy.number], reduce: bool = True, max_denom: int = 10000) -> NDArray[numpy.number]:
    """Convert a plane in 3-axis Miller notation to 4-axis Miller-Bravais notation."""
    a = numpy.atleast_1d(a)
    assert a.shape[-1] == 3
    h, k, l = numpy.split(a, 3, axis=-1)  # noqa: E741
    out = numpy.concatenate((h, k, -(h + k), l), axis=-1)
    return reduce_vec(out, max_denom) if reduce else out