atomlib.atomcell

Atoms with an associated Cell.

This module defines HasAtomCell and the concrete AtomCell, which combines the functionality of HasAtoms and HasCell.

HasAtomCell

Bases: HasAtoms, HasCell, ABC

Source code in atomlib/atomcell.py

class HasAtomCell(HasAtoms, HasCell, abc.ABC):
    @abc.abstractmethod
    def get_frame(self) -> CoordinateFrame:
        """Get the coordinate frame atoms are stored in."""
        ...

    @abc.abstractmethod
    def with_atoms(self, atoms: HasAtoms, frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Replace the atoms in `self`. If no coordinate frame is specified, keep the coordinate frame unchanged.
        """
        ...

    def with_cell(self, cell: Cell) -> Self:
        """
        Replace the cell in `self`, without touching the atomic coordinates.
        """
        return self.to_frame('local').with_cell(cell)

    def get_atomcell(self) -> AtomCell:
        frame = self.get_frame()
        return AtomCell(self.get_atoms(frame), self.get_cell(), frame=frame, keep_frame=True)

    @abc.abstractmethod
    def get_atoms(self, frame: t.Optional[CoordinateFrame] = None) -> Atoms:
        """Get atoms contained in `self`, in the given coordinate frame."""
        ...

    def bbox_atoms(self, frame: t.Optional[CoordinateFrame] = None) -> BBox3D:
        """Return the bounding box of all the atoms in `self`, in the given coordinate frame."""
        return self.get_atoms(frame).bbox()

    def bbox(self, frame: CoordinateFrame = 'local') -> BBox3D:
        """
        Return the combined bounding box of the cell and atoms in the given coordinate system.
        To get the cell or atoms bounding box only, use [`bbox_cell`][atomlib.atomcell.HasAtomCell.bbox_cell] or [`bbox_atoms`][atomlib.atomcell.HasAtomCell.bbox_atoms].
        """
        return self.bbox_atoms(frame) | self.bbox_cell(frame)

    # transformation

    def _transform_atoms_in_frame(self, frame: t.Optional[CoordinateFrame], f: t.Callable[[Atoms], Atoms]) -> Atoms:
        # ugly code
        if frame is None or frame == self.get_frame():
            return f(self.get_atoms())
        return f(self.get_atoms(frame)).transform(self.get_transform(self.get_frame(), frame))

    def to_frame(self, frame: CoordinateFrame) -> Self:
        """Convert the stored Atoms to the given coordinate frame."""
        return self.with_atoms(self.get_atoms(frame), frame)

    def transform_atoms(self, transform: IntoTransform3D, selection: t.Optional[AtomSelection] = None, *,
                        frame: CoordinateFrame = 'local', transform_velocities: bool = False) -> Self:
        """
        Transform the atoms in `self` by `transform`.
        If `selection` is given, only transform the atoms in `selection`.
        """
        transform = self.change_transform(Transform3D.make(transform), self.get_frame(), frame)
        return self.with_atoms(self.get_atoms(self.get_frame()).transform(transform, selection, transform_velocities=transform_velocities))

    def transform_cell(self, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> Self:
        """
        Apply the given transform to the unit cell, without changing atom positions.
        The transform is applied in coordinate frame 'frame'.
        """
        return self.with_cell(self.get_cell().transform_cell(transform, frame=frame))

    def transform(self, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> Self:
        if isinstance(transform, Transform3D) and not isinstance(transform, AffineTransform3D):
            raise ValueError("Non-affine transforms cannot change the box dimensions. Use 'transform_atoms' instead.")
        # TODO: cleanup once tests pass
        # coordinate change the transform into atomic coordinates
        new_cell = self.get_cell().transform_cell(transform, frame)
        transform = self.get_cell().change_transform(transform, self.get_frame(), frame)
        return self.with_atoms(self.get_atoms().transform(transform), self.get_frame()).with_cell(new_cell)

    # crop methods

    def crop(self, x_min: float = -numpy.inf, x_max: float = numpy.inf,
             y_min: float = -numpy.inf, y_max: float = numpy.inf,
             z_min: float = -numpy.inf, z_max: float = numpy.inf, *,
             frame: CoordinateFrame = 'local') -> Self:
        """
        Crop atoms and cell to the given extents. For a non-orthogonal
        cell, this must be specified in cell coordinates. This
        function implicity `explode`s the cell as well.

        To crop atoms only, use `crop_atoms` instead.
        """

        cell = self.get_cell().crop(x_min, x_max, y_min, y_max, z_min, z_max, frame=frame)
        atoms = self._transform_atoms_in_frame(frame, lambda atoms: atoms.crop_atoms(x_min, x_max, y_min, y_max, z_min, z_max))
        return self.with_cell(cell).with_atoms(atoms)

    def crop_atoms(self, x_min: float = -numpy.inf, x_max: float = numpy.inf,
                   y_min: float = -numpy.inf, y_max: float = numpy.inf,
                   z_min: float = -numpy.inf, z_max: float = numpy.inf, *,
                   frame: CoordinateFrame = 'local') -> Self:
        atoms = self._transform_atoms_in_frame(frame, lambda atoms: atoms.crop_atoms(x_min, x_max, y_min, y_max, z_min, z_max))
        return self.with_atoms(atoms)

    def crop_to_box(self, eps: float = 1e-5) -> Self:
        atoms = self._transform_atoms_in_frame('cell_box', lambda atoms: atoms.crop_atoms(*([-eps, 1-eps]*3)))
        return self.with_atoms(atoms)

    def wrap(self, eps: float = 1e-5) -> Self:
        """Wrap atoms around the cell boundaries."""
        return self.with_atoms(self._transform_atoms_in_frame('cell_box', lambda a: a._wrap(eps)))

    def _repeat_to_contain(self, pts: numpy.ndarray, pad: int = 0, frame: CoordinateFrame = 'cell_frac') -> Self:
        #print(f"pts: {pts} in frame {frame}")
        pts = self.get_transform('cell_frac', frame) @ pts

        bbox = BBox3D.unit() | BBox3D.from_pts(pts)
        min_bounds = numpy.floor(bbox.min).astype(int) - pad
        max_bounds = numpy.ceil(bbox.max).astype(int) + pad
        #print(f"tiling to {min_bounds}, {max_bounds}")
        repeat = max_bounds - min_bounds
        cells = numpy.stack(numpy.meshgrid(*map(numpy.arange, repeat))).reshape(3, -1).T.astype(float)

        atoms = self.get_atoms('cell_frac')
        atoms = Atoms.concat([
            atoms.transform(AffineTransform3D.translate(cell))
            for cell in cells
        ])
        #print(f"atoms:\n{atoms}")
        cell = self.get_cell().repeat(repeat) \
            .transform_cell(AffineTransform3D.translate(min_bounds), 'cell_frac')
        return self.with_cell(cell).with_atoms(atoms, 'cell_frac')

    def repeat(self, n: t.Union[int, VecLike]) -> Self:
        """Tile the cell"""
        ns = numpy.broadcast_to(n, 3)
        if not numpy.issubdtype(ns.dtype, numpy.integer):
            raise ValueError("repeat() argument must be an integer or integer array.")

        cells = numpy.stack(numpy.meshgrid(*map(numpy.arange, ns))) \
            .reshape(3, -1).T.astype(float)
        cells = cells * self.box_size

        atoms = self.get_atoms('cell')
        atoms = Atoms.concat([
            atoms.transform(AffineTransform3D.translate(cell))
            for cell in cells
        ]) #.transform(self.cell.get_transform('local', 'cell_frac'))
        return self.with_atoms(atoms, 'cell').with_cell(self.get_cell().repeat(ns))

    def repeat_to(self, size: VecLike, crop: t.Union[bool, t.Sequence[bool]] = False) -> Self:
        """
        Repeat the cell so it is at least `size` along the crystal's axes.

        If `crop`, then crop the cell to exactly `size`. This may break periodicity.
        `crop` may be a vector, in which case you can specify cropping only along some axes.
        """
        size = to_vec3(size)
        cell_size = self.cell_size * self.n_cells
        repeat = numpy.maximum(numpy.ceil(size / cell_size).astype(int), 1)
        atom_cell = self.repeat(repeat)

        crop_v = to_vec3(crop, dtype=numpy.bool_)
        if numpy.any(crop_v):
            crop_x, crop_y, crop_z = crop_v
            return atom_cell.crop(
                x_max = size[0] if crop_x else numpy.inf,
                y_max = size[1] if crop_y else numpy.inf,
                z_max = size[2] if crop_z else numpy.inf,
                frame='cell'
            )

        return atom_cell

    def repeat_x(self, n: int) -> Self:
        """Tile the cell in the x axis."""
        return self.repeat((n, 1, 1))

    def repeat_y(self, n: int) -> Self:
        """Tile the cell in the y axis."""
        return self.repeat((1, n, 1))

    def repeat_z(self, n: int) -> Self:
        """Tile the cell in the z axis."""
        return self.repeat((1, 1, n))

    def repeat_to_x(self, size: float, crop: bool = False) -> Self:
        """Repeat the cell so it is at least size `size` along the x axis."""
        return self.repeat_to([size, 0., 0.], [crop, False, False])

    def repeat_to_y(self, size: float, crop: bool = False) -> Self:
        """Repeat the cell so it is at least size `size` along the y axis."""
        return self.repeat_to([0., size, 0.], [False, crop, False])

    def repeat_to_z(self, size: float, crop: bool = False) -> Self:
        """Repeat the cell so it is at least size `size` along the z axis."""
        return self.repeat_to([0., 0., size], [False, False, crop])

    def repeat_to_aspect(self, plane: t.Literal['xy', 'xz', 'yz'] = 'xy', *,
                         aspect: float = 1., min_size: t.Optional[VecLike] = None,
                         max_size: t.Optional[VecLike] = None) -> Self:
        """
        Repeat to optimize the aspect ratio in `plane`,
        while staying above `min_size` and under `max_size`.
        """
        if min_size is None:
            min_n = numpy.array([1, 1, 1], numpy.int_)
        else:
            min_n = numpy.maximum(numpy.ceil(to_vec3(min_size) / self.box_size), 1).astype(numpy.int_)

        if max_size is None:
            max_n = 3 * min_n
        else:
            max_n = numpy.maximum(numpy.floor(to_vec3(max_size) / self.box_size), 1).astype(numpy.int_)

        if plane == 'xy':
            indices = [0, 1]
        elif plane == 'xz':
            indices = [0, 2]
        elif plane == 'yz':
            indices = [1, 2]
        else:
            raise ValueError(f"Invalid plane '{plane}'. Exepcted 'xy', 'xz', 'or 'yz'.")

        na = numpy.arange(min_n[indices[0]], max_n[indices[0]])
        nb = numpy.arange(min_n[indices[1]], max_n[indices[1]])
        (na, nb) = numpy.meshgrid(na, nb)

        aspects = na * self.box_size[indices[0]] / (nb * self.box_size[indices[1]])
        # cost function: log(aspect)^2  (so cost(0.5) == cost(2))
        min_i = numpy.argmin(numpy.log(aspects / aspect)**2)
        repeat = numpy.array([1, 1, 1], numpy.int_)
        repeat[indices] = na.flatten()[min_i], nb.flatten()[min_i]
        return self.repeat(repeat)

    def explode(self) -> Self:
        """Materialize repeated cells as one supercell."""
        frame = self.get_frame()

        return self.with_atoms(self.get_atoms('local'), 'local') \
            .with_cell(self.get_cell().explode()) \
            .to_frame(frame)

    def periodic_duplicate(self, eps: float = 1e-5) -> Self:
        """
        Add duplicate copies of atoms near periodic boundaries.

        For instance, an atom at a corner will be duplicated into 8 copies.
        This is mostly only useful for visualization.
        """
        frame_save = self.get_frame()
        self = self.to_frame('cell_box').wrap(eps=eps)

        for i in range(3):
            self = self.concat((self,
                self.filter(polars.col('coords').arr.get(i).abs() <= eps, frame='cell_box')
                    .transform_atoms(AffineTransform3D.translate([1. if i == j else 0. for j in range(3)]), frame='cell_box')
            ))

        return self.to_frame(frame_save)

    # add frame to some HasAtoms methods

    @_fwd_atoms_get
    def describe(self, percentiles: t.Union[t.Sequence[float], float, None] = (0.25, 0.5, 0.75), *,
                 interpolation: RollingInterpolationMethod = 'nearest',
                 frame: t.Optional[CoordinateFrame] = None) -> polars.DataFrame:
        """
        Return summary statistics for `self`. See [`DataFrame.describe`][polars.DataFrame.describe] for more information.

        Args:
          percentiles: List of percentiles/quantiles to include. Defaults to 25% (first quartile),
                       50% (median), and 75% (third quartile).

        Returns:
          A dataframe containing summary statistics (mean, std. deviation, percentiles, etc.) for each column.
        """
        ...

    @_fwd_atoms_transform
    def with_columns(self,
                     *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
                     frame: t.Optional[CoordinateFrame] = None,
                     **named_exprs: IntoExpr) -> Self:
        """Return a copy of `self` with the given columns added."""
        ...

    with_column = with_columns

    @_fwd_atoms_get
    def get_column(self, name: str, *, frame: t.Optional[CoordinateFrame] = None) -> polars.Series:
        """
        Get the specified column from `self`, raising [`polars.ColumnNotFoundError`][polars.exceptions.ColumnNotFoundError] if it's not present.

        [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
        """
        ...

    @_fwd_atoms_get
    def get_columns(self, *, frame: t.Optional[CoordinateFrame] = None) -> t.List[polars.Series]:
        """
        Return all columns from `self` as a list of [`Series`][polars.Series].

        [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
        """
        ...

    @_fwd_atoms_get
    def group_by(self, *by: t.Union[IntoExpr, t.Iterable[IntoExpr]],
                 maintain_order: bool = False, frame: t.Optional[CoordinateFrame] = None,
                 **named_by: IntoExpr) -> polars.dataframe.group_by.GroupBy:
        """
        Start a group by operation. See [`DataFrame.group_by`][polars.DataFrame.group_by] for more information.
        """
        ...

    def pipe(self: HasAtomCellT, function: t.Callable[Concatenate[HasAtomCellT, P], T], *args: P.args, **kwargs: P.kwargs) -> T:
        """Apply `function` to `self` (in method-call syntax)."""
        return function(self, *args, **kwargs)

    @_fwd_atoms_transform
    def filter(
        self,
        *predicates: t.Union[None, IntoExprColumn, t.Iterable[IntoExprColumn], bool, t.List[bool], numpy.ndarray],
        frame: t.Optional[CoordinateFrame] = None,
        **constraints: t.Any,
    ) -> Self:
        """Filter `self`, removing rows which evaluate to `False`."""
        ...

    @_fwd_atoms_transform
    def sort(
        self,
        by: t.Union[IntoExpr, t.Iterable[IntoExpr]],
        *more_by: IntoExpr,
        descending: t.Union[bool, t.Sequence[bool]] = False,
        nulls_last: bool = False,
    ) -> Self:
        """
        Sort the atoms in `self` by the given columns/expressions.
        """
        ...

    @_fwd_atoms_transform
    def slice(self, offset: int, length: t.Optional[int] = None, *,
              frame: t.Optional[CoordinateFrame] = None) -> Self:
        """Return a slice of the rows in `self`."""
        ...

    @_fwd_atoms_transform
    def head(self, n: int = 5, *, frame: t.Optional[CoordinateFrame] = None) -> Self:
        """Return the first `n` rows of `self`."""
        ...

    @_fwd_atoms_transform
    def tail(self, n: int = 5, *, frame: t.Optional[CoordinateFrame] = None) -> Self:
        """Return the last `n` rows of `self`."""
        ...

    @_fwd_atoms_transform
    def fill_null(
        self, value: t.Any = None, strategy: t.Optional[FillNullStrategy] = None,
        limit: t.Optional[int] = None, matches_supertype: bool = True,
    ) -> Self:
        ...

    @_fwd_atoms_transform
    def fill_nan(self, value: t.Union[polars.Expr, int, float, None], *,
                 frame: t.Optional[CoordinateFrame] = None) -> Self:
        ...

    # TODO: partition_by

    @_fwd_atoms_get
    def select(
        self, *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
        frame: t.Optional[CoordinateFrame] = None,
        **named_exprs: IntoExpr
    ) -> polars.DataFrame:
        """
        Select `exprs` from `self`, and return as a [`polars.DataFrame`][polars.DataFrame].

        Expressions may either be columns or expressions of columns.

        [polars.DataFrame]: https://docs.pola.rs/py-polars/html/reference/dataframe/index.html
        """
        ...

    @_fwd_atoms_transform
    def select_props(
        self,
        *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
        frame: t.Optional[CoordinateFrame] = None,
        **named_exprs: IntoExpr
    ) -> Self:
        """
        Select `exprs` from `self`, while keeping required columns.
        Doesn't affect the cell.

        Returns:
          A [`HasAtomCell`][atomlib.atomcell.HasAtomCell] filtered to contain
          the specified properties (as well as required columns).
        """
        ...

    @_fwd_atoms_get
    def try_select(
        self, *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
        frame: t.Optional[CoordinateFrame] = None,
        **named_exprs: IntoExpr
    ) -> t.Optional[polars.DataFrame]:
        """
        Try to select `exprs` from `self`, and return as a [`polars.DataFrame`][polars.DataFrame].

        Expressions may either be columns or expressions of columns. Returns `None` if any
        columns are missing.

        [polars.DataFrame]: https://docs.pola.rs/py-polars/html/reference/dataframe/index.html
        """
        ...

    @_fwd_atoms_transform
    def round_near_zero(self, tol: float = 1e-14, *,
                        frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Round atom position values near zero to zero.
        """
        ...

    @_fwd_atoms_get
    def coords(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Optional[CoordinateFrame] = None) -> NDArray[numpy.float64]:
        """
        Return a `(N, 3)` ndarray of atom positions (dtype [`numpy.float64`][numpy.float64])
        in the given coordinate frame.
        """
        ...

    @_fwd_atoms_get
    def velocities(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Optional[CoordinateFrame] = None) -> t.Optional[NDArray[numpy.float64]]:
        """
        Return a `(N, 3)` ndarray of atom velocities (dtype [`numpy.float64`][numpy.float64])
        in the given coordinate frame.
        """
        ...

    @t.overload
    def add_atom(self, elem: t.Union[int, str], x: ArrayLike, /, *,
                 y: None = None, z: None = None, frame: t.Optional[CoordinateFrame] = None,
                 **kwargs: t.Any) -> Self:
        ...

    @t.overload
    def add_atom(self, elem: t.Union[int, str], /,
                 x: float, y: float, z: float, *,
                 frame: t.Optional[CoordinateFrame] = None,
                 **kwargs: t.Any) -> Self:
        ...

    @_fwd_atoms_transform
    def add_atom(self, elem: t.Union[int, str], /,  # type: ignore (spurious)
                 x: t.Union[ArrayLike, float],
                 y: t.Optional[float] = None,
                 z: t.Optional[float] = None, *,
                 frame: t.Optional[CoordinateFrame] = None,
                 **kwargs: t.Any) -> Self:
        """
        Return a copy of `self` with an extra atom.

        By default, all extra columns present in `self` must be specified as `**kwargs`.

        Try to avoid calling this in a loop (Use [`concat`][atomlib.atomcell.HasAtomCell.concat] instead).
        """
        ...

    @_fwd_atoms_transform
    def with_index(self, index: t.Optional[AtomValues] = None, *,
                   frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Returns `self` with a row index added in column 'i' (dtype [`polars.Int64`][polars.datatypes.Int64]).
        If `index` is not specified, defaults to an existing index or a new index.
        """
        ...

    @_fwd_atoms_transform
    def with_wobble(self, wobble: t.Optional[AtomValues] = None, *,
                    frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Return `self` with the given displacements in column 'wobble' (dtype [`polars.Float64`][polars.datatypes.Float64]).
        If `wobble` is not specified, defaults to the already-existing wobbles or 0.
        """
        ...

    @_fwd_atoms_transform
    def with_occupancy(self, frac_occupancy: t.Optional[AtomValues] = None, *,
                       frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Return self with the given fractional occupancies (dtype [`polars.Float64`][polars.datatypes.Float64]).
        If `frac_occupancy` is not specified, defaults to the already-existing occupancies or 1.
        """
        ...

    @_fwd_atoms_transform
    def apply_wobble(self, rng: t.Union[numpy.random.Generator, int, None] = None,
                     frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Displace the atoms in `self` by the amount in the `wobble` column.
        `wobble` is interpretated as a mean-squared displacement, which is distributed
        equally over each axis.
        """
        ...

    @_fwd_atoms_transform
    def with_type(self, types: t.Optional[AtomValues] = None, *,
                  frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Return `self` with the given atom types in column 'type'.
        If `types` is not specified, use the already existing types or auto-assign them.

        When auto-assigning, each symbol is given a unique value, case-sensitive.
        Values are assigned from lowest atomic number to highest.
        For instance: `["Ag+", "Na", "H", "Ag"]` => `[3, 11, 1, 2]`
        """
        ...

    @_fwd_atoms_transform
    def with_mass(self, mass: t.Optional[ArrayLike] = None, *,
                  frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Return `self` with the given atom masses in column `'mass'`.
        If `mass` is not specified, use the already existing masses or auto-assign them.
        """
        ...

    @_fwd_atoms_transform
    def with_symbol(self, symbols: ArrayLike, selection: t.Optional[AtomSelection] = None, *,
                    frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Return `self` with the given atomic symbols.
        """
        ...

    @_fwd_atoms_transform
    def with_coords(self, pts: ArrayLike, selection: t.Optional[AtomSelection] = None, *,
                    frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Return `self` replaced with the given atomic positions.
        """
        ...

    @_fwd_atoms_transform
    def with_velocity(self, pts: t.Optional[ArrayLike] = None,
                      selection: t.Optional[AtomSelection] = None, *,
                      frame: t.Optional[CoordinateFrame] = None) -> Self:
        """
        Return `self` replaced with the given atomic velocities.
        If `pts` is not specified, use the already existing velocities or zero.
        """
        ...

affine property

affine: AffineTransform3D

Affine transformation. Holds transformation from 'ortho' to 'local' coordinates, including rotation away from the standard crystal orientation.

ortho property

ortho: LinearTransform3D

Orthogonalization transformation. Skews but does not scale the crystal axes to cartesian axes.

metric property

metric: LinearTransform3D

Cell metric tensor

Returns the dot product between every combination of basis vectors. :math:\mathbf{a} \cdot \mathbf{b} = a_i M_ij b_j

cell_size property

cell_size: NDArray[float64]

Unit cell size.

cell_angle property

cell_angle: NDArray[float64]

Unit cell angles, in radians.

n_cells property

n_cells: NDArray[int_]

Number of unit cells.

pbc property

pbc: NDArray[bool_]

Flags indicating the presence of periodic boundary conditions along each axis.

ortho_size property

ortho_size: NDArray[float64]

Return size of orthogonal unit cell.

Equivalent to the diagonal of the orthogonalization matrix.

box_size property

box_size: NDArray[float64]

Return size of the cell box.

Equivalent to self.n_cells * self.cell_size.

columns property

columns: List[str]

Return the column names in self.

RETURNS	DESCRIPTION
List[str]	A sequence of column names

dtypes property

dtypes: List[DataType]

Return the datatypes in self.

RETURNS	DESCRIPTION
List[DataType]	A sequence of column DataTypes

schema property

schema: Schema

Return the schema of self.

RETURNS	DESCRIPTION
Schema	A dictionary of column names and DataTypes

unique class-attribute instance-attribute

unique = deduplicate

with_column class-attribute instance-attribute

with_column = with_columns

get_cell abstractmethod

get_cell() -> Cell

Get the cell contained in self. This should be a low cost method.

Source code in atomlib/cell.py

@abc.abstractmethod
def get_cell(self) -> Cell:
    """Get the cell contained in ``self``. This should be a low cost method."""
    ...

get_transform

get_transform(
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> AffineTransform3D

In the two-argument form, get the transform to frame_to from frame_from. In the one-argument form, get the transform from local coordinates to 'frame'.

Source code in atomlib/cell.py

def get_transform(self, frame_to: t.Optional[CoordinateFrame] = None, frame_from: t.Optional[CoordinateFrame] = None) -> AffineTransform3D:
    """
    In the two-argument form, get the transform to `frame_to` from `frame_from`.
    In the one-argument form, get the transform from local coordinates to 'frame'.
    """
    transform_from = self._get_transform_to_local(frame_from) if frame_from is not None else AffineTransform3D()
    transform_to = self._get_transform_to_local(frame_to) if frame_to is not None else AffineTransform3D()
    if frame_from is not None and frame_to is not None and frame_from.lower() == frame_to.lower():
        return AffineTransform3D()
    return transform_to.inverse() @ transform_from

corners

corners(frame: CoordinateFrame = 'local') -> ndarray

Source code in atomlib/cell.py

def corners(self, frame: CoordinateFrame = 'local') -> numpy.ndarray:
    corners = numpy.array(list(itertools.product((0., 1.), repeat=3)))
    return self.get_transform(frame, 'cell_box') @ corners

bbox_cell

bbox_cell(frame: CoordinateFrame = 'local') -> BBox3D

Return the bounding box of the cell box in the given coordinate system.

Source code in atomlib/cell.py

def bbox_cell(self, frame: CoordinateFrame = 'local') -> BBox3D:
    """Return the bounding box of the cell box in the given coordinate system."""
    return BBox3D.from_pts(self.corners(frame))

is_orthogonal

is_orthogonal(tol: float = 1e-08) -> bool

Returns whether this cell is orthogonal (axes are at right angles.)

Source code in atomlib/cell.py

def is_orthogonal(self, tol: float = 1e-8) -> bool:
    """Returns whether this cell is orthogonal (axes are at right angles.)"""
    return self.ortho.is_diagonal(tol=tol)

is_orthogonal_in_local

is_orthogonal_in_local(tol: float = 1e-08) -> bool

Returns whether this cell is orthogonal and aligned with the local coordinate system.

Source code in atomlib/cell.py

def is_orthogonal_in_local(self, tol: float = 1e-8) -> bool:
    """Returns whether this cell is orthogonal and aligned with the local coordinate system."""
    transform = (self.affine @ self.ortho).to_linear()
    if not transform.is_scaled_orthogonal(tol):
        return False
    normed = transform.inner / numpy.linalg.norm(transform.inner, axis=-2, keepdims=True)
    # every row of transform must be a +/- 1 times a basis vector (i, j, or k)
    return all(
        any(numpy.isclose(numpy.abs(numpy.dot(row, v)), 1., atol=tol) for v in numpy.eye(3))
        for row in normed
    )

to_ortho

to_ortho() -> AffineTransform3D

Source code in atomlib/cell.py

def to_ortho(self) -> AffineTransform3D:
    return self.get_transform('local', 'cell_box')

strain_orthogonal

strain_orthogonal() -> HasCellT

Orthogonalize using strain.

Strain is applied such that the x-axis remains fixed, and the y-axis remains in the xy plane. For small displacements, no hydrostatic strain is applied (volume is conserved).

Source code in atomlib/cell.py

def strain_orthogonal(self: HasCellT) -> HasCellT:
    """
    Orthogonalize using strain.

    Strain is applied such that the x-axis remains fixed, and the y-axis remains in the xy plane.
    For small displacements, no hydrostatic strain is applied (volume is conserved).
    """
    return self.with_cell(Cell(
        affine=self.affine,
        ortho=LinearTransform3D(),
        cell_size=self.cell_size,
        n_cells=self.n_cells,
        pbc=self.pbc,
    ))

explode_z

explode_z() -> HasCellT

Materialize repeated cells as one supercell in z.

Source code in atomlib/cell.py

def explode_z(self: HasCellT) -> HasCellT:
    """Materialize repeated cells as one supercell in z."""
    return self.with_cell(Cell(
        affine=self.affine,
        ortho=self.ortho,
        cell_size=self.cell_size*[1, 1, self.n_cells[2]],
        n_cells=[*self.n_cells[:2], 1],
        cell_angle=self.cell_angle,
        pbc=self.pbc,
    ))

change_transform

change_transform(
    transform: AffineTransform3D,
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> AffineTransform3D

change_transform(
    transform: Transform3D,
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> Transform3D

change_transform(
    transform: Transform3D,
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> Transform3D

Coordinate-change a transformation from frame_from into frame_to.

Source code in atomlib/cell.py

def change_transform(self, transform: Transform3D,
                     frame_to: t.Optional[CoordinateFrame] = None,
                     frame_from: t.Optional[CoordinateFrame] = None) -> Transform3D:
    """Coordinate-change a transformation from `frame_from` into `frame_to`."""
    if frame_to == frame_from and frame_to is not None:
        return transform
    coord_change = self.get_transform(frame_to, frame_from)
    return coord_change @ transform @ coord_change.inverse()

assert_equal

assert_equal(other: Any)

Source code in atomlib/atoms.py

def assert_equal(self, other: t.Any):
    assert isinstance(other, HasAtoms)
    assert dict(self.schema) == dict(other.schema)
    for col in self.schema.keys():
        polars.testing.assert_series_equal(self[col], other[col], check_names=False, rtol=1e-3, atol=1e-8)

insert_column

insert_column(index: int, column: Series) -> DataFrame

Source code in atomlib/atoms.py

@_fwd_frame_map
def insert_column(self, index: int, column: polars.Series) -> polars.DataFrame:
    return self._get_frame().insert_column(index, column)

get_column_index

get_column_index(name: str) -> int

Get the index of a column by name, raising polars.ColumnNotFoundError if it's not present.

Source code in atomlib/atoms.py

@_fwd_frame(polars.DataFrame.get_column_index)
def get_column_index(self, name: str) -> int:
    """Get the index of a column by name, raising [`polars.ColumnNotFoundError`][polars.exceptions.ColumnNotFoundError] if it's not present."""
    ...

clone

clone() -> DataFrame

Return a copy of self.

Source code in atomlib/atoms.py

@_fwd_frame_map
def clone(self) -> polars.DataFrame:
    """Return a copy of `self`."""
    return self._get_frame().clone()

drop

drop(
    *columns: Union[str, Iterable[str]], strict: bool = True
) -> DataFrame

Return self with the specified columns removed.

Source code in atomlib/atoms.py

def drop(self, *columns: t.Union[str, t.Iterable[str]], strict: bool = True) -> polars.DataFrame:
    """Return `self` with the specified columns removed."""
    return self._get_frame().drop(*columns, strict=strict)

drop_nulls

drop_nulls(
    subset: Union[str, Collection[str], None] = None
) -> DataFrame

Drop rows that contain nulls in any of columns subset.

Source code in atomlib/atoms.py

@_fwd_frame_map
def drop_nulls(self, subset: t.Union[str, t.Collection[str], None] = None) -> polars.DataFrame:
    """Drop rows that contain nulls in any of columns `subset`."""
    return self._get_frame().drop_nulls(subset)

concat classmethod

concat(
    atoms: Union[
        HasAtomsT,
        IntoAtoms,
        Iterable[Union[HasAtomsT, IntoAtoms]],
    ],
    *,
    rechunk: bool = True,
    how: ConcatMethod = "vertical"
) -> HasAtomsT

Concatenate multiple Atoms together, handling metadata appropriately.

Source code in atomlib/atoms.py

@classmethod
def concat(cls: t.Type[HasAtomsT],
           atoms: t.Union[HasAtomsT, IntoAtoms, t.Iterable[t.Union[HasAtomsT, IntoAtoms]]], *,
           rechunk: bool = True, how: ConcatMethod = 'vertical') -> HasAtomsT:
    """Concatenate multiple `Atoms` together, handling metadata appropriately."""
    # this method is tricky. It needs to accept raw Atoms, as well as HasAtoms of the
    # same type as ``cls``.
    if _is_abstract(cls):
        raise TypeError("concat() must be called on a concrete class.")

    if isinstance(atoms, HasAtoms):
        atoms = (atoms,)
    dfs = [a.get_atoms('local').inner if isinstance(a, HasAtoms) else Atoms(t.cast(IntoAtoms, a)).inner for a in atoms]
    representative = cls._combine_metadata(*(a for a in atoms if isinstance(a, HasAtoms)))

    if len(dfs) == 0:
        return representative.with_atoms(Atoms.empty(), 'local')

    if how in ('vertical', 'vertical_relaxed'):
        # get order from first member
        cols = dfs[0].columns
        dfs = [df.select(cols) for df in dfs]
    elif how == 'inner':
        cols = reduce(operator.and_, (df.schema.keys() for df in dfs))
        schema = OrderedDict((col, dfs[0].schema[col]) for col in cols)
        if len(schema) == 0:
            raise ValueError("Atoms have no columns in common")

        dfs = [_select_schema(df, schema) for df in dfs]
        how = 'vertical'

    return representative.with_atoms(Atoms(polars.concat(dfs, rechunk=rechunk, how=how)), 'local')

partition_by

partition_by(
    by: Union[str, Sequence[str]],
    *more_by: str,
    maintain_order: bool = True,
    include_key: bool = True,
    as_dict: Literal[False] = False
) -> List[Self]

partition_by(
    by: Union[str, Sequence[str]],
    *more_by: str,
    maintain_order: bool = True,
    include_key: bool = True,
    as_dict: Literal[True] = ...
) -> Dict[Any, Self]

partition_by(
    by: Union[str, Sequence[str]],
    *more_by: str,
    maintain_order: bool = True,
    include_key: bool = True,
    as_dict: bool = False
) -> Union[List[Self], Dict[Any, Self]]

Group by the given columns and partition into separate dataframes.

Return the partitions as a dictionary by specifying as_dict=True.

Source code in atomlib/atoms.py

def partition_by(
    self, by: t.Union[str, t.Sequence[str]], *more_by: str,
    maintain_order: bool = True, include_key: bool = True, as_dict: bool = False
) -> t.Union[t.List[Self], t.Dict[t.Any, Self]]:
    """
    Group by the given columns and partition into separate dataframes.

    Return the partitions as a dictionary by specifying `as_dict=True`.
    """
    if as_dict:
        d = self._get_frame().partition_by(by, *more_by, maintain_order=maintain_order, include_key=include_key, as_dict=True)
        return {k: self.with_atoms(Atoms(df, _unchecked=True)) for (k, df) in d.items()}

    return [
        self.with_atoms(Atoms(df, _unchecked=True))
        for df in self._get_frame().partition_by(by, *more_by, maintain_order=maintain_order, include_key=include_key, as_dict=False)
    ]

select_schema

select_schema(schema: SchemaDict) -> DataFrame

Select columns from self and cast to the given schema. Raises TypeError if a column is not found or if it can't be cast.

Source code in atomlib/atoms.py

def select_schema(self, schema: SchemaDict) -> polars.DataFrame:
    """
    Select columns from `self` and cast to the given schema.
    Raises [`TypeError`][TypeError] if a column is not found or if it can't be cast.
    """
    return _select_schema(self, schema)

try_get_column

try_get_column(name: str) -> Optional[Series]

Try to get a column from self, returning None if it doesn't exist.

Source code in atomlib/atoms.py

def try_get_column(self, name: str) -> t.Optional[polars.Series]:
    """Try to get a column from `self`, returning `None` if it doesn't exist."""
    try:
        return self.get_column(name)
    except polars.exceptions.ColumnNotFoundError:
        return None

deduplicate

deduplicate(
    tol: float = 0.001,
    subset: Iterable[str] = ("x", "y", "z", "symbol"),
    keep: UniqueKeepStrategy = "first",
    maintain_order: bool = True,
) -> Self

De-duplicate atoms in self. Atoms of the same symbol that are closer than tolerance to each other (by Euclidian distance) will be removed, leaving only the atom specified by keep (defaults to the first atom).

If subset is specified, only those columns will be included while assessing duplicates. Floating point columns other than 'x', 'y', and 'z' will not by toleranced.

Source code in atomlib/atoms.py

def deduplicate(self, tol: float = 1e-3, subset: t.Iterable[str] = ('x', 'y', 'z', 'symbol'),
                keep: UniqueKeepStrategy = 'first', maintain_order: bool = True) -> Self:
    """
    De-duplicate atoms in `self`. Atoms of the same `symbol` that are closer than `tolerance`
    to each other (by Euclidian distance) will be removed, leaving only the atom specified by
    `keep` (defaults to the first atom).

    If `subset` is specified, only those columns will be included while assessing duplicates.
    Floating point columns other than 'x', 'y', and 'z' will not by toleranced.
    """
    import scipy.spatial

    cols = set((subset,) if isinstance(subset, str) else subset)

    indices = numpy.arange(len(self))

    spatial_cols = cols.intersection(('x', 'y', 'z'))
    cols -= spatial_cols
    if len(spatial_cols) > 0:
        coords = self.select([_coord_expr(col).alias(col) for col in spatial_cols]).to_numpy()
        tree = scipy.spatial.KDTree(coords)

        # TODO This is a bad algorithm
        while True:
            changed = False
            for (i, j) in tree.query_pairs(tol, 2.):
                # whenever we encounter a pair, ensure their index matches
                i_i, i_j = indices[[i, j]]
                if i_i != i_j:
                    indices[i] = indices[j] = min(i_i, i_j)
                    changed = True
            if not changed:
                break

        self = self.with_column(polars.Series('_unique_pts', indices))
        cols.add('_unique_pts')

    frame = self._get_frame().unique(subset=list(cols), keep=keep, maintain_order=maintain_order)
    if len(spatial_cols) > 0:
        frame = frame.drop('_unique_pts')

    return self.with_atoms(Atoms(frame, _unchecked=True))

with_bounds

with_bounds(
    cell_size: Optional[VecLike] = None,
    cell_origin: Optional[VecLike] = None,
) -> "AtomCell"

Return a periodic cell with the given orthogonal cell dimensions.

If cell_size is not specified, it will be assumed (and may be incorrect).

Source code in atomlib/atoms.py

def with_bounds(self, cell_size: t.Optional[VecLike] = None, cell_origin: t.Optional[VecLike] = None) -> 'AtomCell':
    """
    Return a periodic cell with the given orthogonal cell dimensions.

    If cell_size is not specified, it will be assumed (and may be incorrect).
    """
    # TODO: test this
    from .atomcell import AtomCell

    if cell_size is None:
        warnings.warn("Cell boundary unknown. Defaulting to cell BBox")
        cell_size = self.bbox().size
        cell_origin = self.bbox().min

    # TODO test this origin code
    cell = Cell.from_unit_cell(cell_size)
    if cell_origin is not None:
        cell = cell.transform_cell(AffineTransform3D.translate(to_vec3(cell_origin)))

    return AtomCell(self.get_atoms(), cell, frame='local')

x

x() -> Expr

Source code in atomlib/atoms.py

def x(self) -> polars.Expr:
    return polars.col('coords').arr.get(0).alias('x')

y

y() -> Expr

Source code in atomlib/atoms.py

def y(self) -> polars.Expr:
    return polars.col('coords').arr.get(1).alias('y')

z

z() -> Expr

Source code in atomlib/atoms.py

def z(self) -> polars.Expr:
    return polars.col('coords').arr.get(2).alias('z')

types

types() -> Optional[Series]

Returns a Series of atom types (dtype polars.Int32).

Source code in atomlib/atoms.py

def types(self) -> t.Optional[polars.Series]:
    """
    Returns a [`Series`][polars.Series] of atom types (dtype [`polars.Int32`][polars.datatypes.Int32]).

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    return self.try_get_column('type')

masses

masses() -> Optional[Series]

Returns a Series of atom masses (dtype polars.Float32).

Source code in atomlib/atoms.py

def masses(self) -> t.Optional[polars.Series]:
    """
    Returns a [`Series`][polars.Series] of atom masses (dtype [`polars.Float32`][polars.datatypes.Float32]).

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    return self.try_get_column('mass')

pos

pos(
    x: Sequence[Optional[float]],
    /,
    *,
    y: None = None,
    z: None = None,
    tol: float = 1e-06,
    **kwargs: Any,
) -> Expr

pos(
    x: Optional[float] = None,
    y: Optional[float] = None,
    z: Optional[float] = None,
    *,
    tol: float = 1e-06,
    **kwargs: Any
) -> Expr

pos(
    x: Union[Sequence[Optional[float]], float, None] = None,
    y: Optional[float] = None,
    z: Optional[float] = None,
    *,
    tol: float = 1e-06,
    **kwargs: Any
) -> Expr

Select all atoms at a given position.

Formally, returns all atoms within a cube of radius tol centered at (x,y,z), exclusive of the cube's surface.

Additional parameters given as kwargs will be checked as additional parameters (with strict equality).

Source code in atomlib/atoms.py

def pos(self,
        x: t.Union[t.Sequence[t.Optional[float]], float, None] = None,
        y: t.Optional[float] = None, z: t.Optional[float] = None, *,
        tol: float = 1e-6, **kwargs: t.Any) -> polars.Expr:
    """
    Select all atoms at a given position.

    Formally, returns all atoms within a cube of radius ``tol``
    centered at ``(x,y,z)``, exclusive of the cube's surface.

    Additional parameters given as ``kwargs`` will be checked
    as additional parameters (with strict equality).
    """

    if isinstance(x, t.Sequence):
        (x, y, z) = x

    tol = abs(float(tol))
    selection = polars.lit(True)
    if x is not None:
        selection &= self.x().is_between(x - tol, x + tol, closed='none')
    if y is not None:
        selection &= self.y().is_between(y - tol, y + tol, closed='none')
    if z is not None:
        selection &= self.z().is_between(z - tol, z + tol, closed='none')
    for (col, val) in kwargs.items():
        selection &= (polars.col(col) == val)

    return selection

apply_occupancy

apply_occupancy(
    rng: Union[Generator, int, None] = None
) -> Self

For each atom in self, use its frac_occupancy to randomly decide whether to remove it.

Source code in atomlib/atoms.py

def apply_occupancy(self, rng: t.Union[numpy.random.Generator, int, None] = None) -> Self:
    """
    For each atom in `self`, use its `frac_occupancy` to randomly decide whether to remove it.
    """
    if 'frac_occupancy' not in self.columns:
        return self
    rng = numpy.random.default_rng(seed=rng)

    frac = self.select('frac_occupancy').to_series().to_numpy()
    choice = rng.binomial(1, frac).astype(numpy.bool_)
    return self.filter(polars.lit(choice))

get_frame abstractmethod

get_frame() -> CoordinateFrame

Get the coordinate frame atoms are stored in.

Source code in atomlib/atomcell.py

@abc.abstractmethod
def get_frame(self) -> CoordinateFrame:
    """Get the coordinate frame atoms are stored in."""
    ...

with_atoms abstractmethod

with_atoms(
    atoms: HasAtoms, frame: Optional[CoordinateFrame] = None
) -> Self

Replace the atoms in self. If no coordinate frame is specified, keep the coordinate frame unchanged.

Source code in atomlib/atomcell.py

@abc.abstractmethod
def with_atoms(self, atoms: HasAtoms, frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Replace the atoms in `self`. If no coordinate frame is specified, keep the coordinate frame unchanged.
    """
    ...

with_cell

with_cell(cell: Cell) -> Self

Replace the cell in self, without touching the atomic coordinates.

Source code in atomlib/atomcell.py

def with_cell(self, cell: Cell) -> Self:
    """
    Replace the cell in `self`, without touching the atomic coordinates.
    """
    return self.to_frame('local').with_cell(cell)

get_atomcell

get_atomcell() -> AtomCell

Source code in atomlib/atomcell.py

def get_atomcell(self) -> AtomCell:
    frame = self.get_frame()
    return AtomCell(self.get_atoms(frame), self.get_cell(), frame=frame, keep_frame=True)

get_atoms abstractmethod

get_atoms(frame: Optional[CoordinateFrame] = None) -> Atoms

Get atoms contained in self, in the given coordinate frame.

Source code in atomlib/atomcell.py

@abc.abstractmethod
def get_atoms(self, frame: t.Optional[CoordinateFrame] = None) -> Atoms:
    """Get atoms contained in `self`, in the given coordinate frame."""
    ...

bbox_atoms

bbox_atoms(
    frame: Optional[CoordinateFrame] = None,
) -> BBox3D

Return the bounding box of all the atoms in self, in the given coordinate frame.

Source code in atomlib/atomcell.py

def bbox_atoms(self, frame: t.Optional[CoordinateFrame] = None) -> BBox3D:
    """Return the bounding box of all the atoms in `self`, in the given coordinate frame."""
    return self.get_atoms(frame).bbox()

bbox

bbox(frame: CoordinateFrame = 'local') -> BBox3D

Return the combined bounding box of the cell and atoms in the given coordinate system. To get the cell or atoms bounding box only, use bbox_cell or bbox_atoms.

Source code in atomlib/atomcell.py

def bbox(self, frame: CoordinateFrame = 'local') -> BBox3D:
    """
    Return the combined bounding box of the cell and atoms in the given coordinate system.
    To get the cell or atoms bounding box only, use [`bbox_cell`][atomlib.atomcell.HasAtomCell.bbox_cell] or [`bbox_atoms`][atomlib.atomcell.HasAtomCell.bbox_atoms].
    """
    return self.bbox_atoms(frame) | self.bbox_cell(frame)

to_frame

to_frame(frame: CoordinateFrame) -> Self

Convert the stored Atoms to the given coordinate frame.

Source code in atomlib/atomcell.py

def to_frame(self, frame: CoordinateFrame) -> Self:
    """Convert the stored Atoms to the given coordinate frame."""
    return self.with_atoms(self.get_atoms(frame), frame)

transform_atoms

transform_atoms(
    transform: IntoTransform3D,
    selection: Optional[AtomSelection] = None,
    *,
    frame: CoordinateFrame = "local",
    transform_velocities: bool = False
) -> Self

Transform the atoms in self by transform. If selection is given, only transform the atoms in selection.

Source code in atomlib/atomcell.py

def transform_atoms(self, transform: IntoTransform3D, selection: t.Optional[AtomSelection] = None, *,
                    frame: CoordinateFrame = 'local', transform_velocities: bool = False) -> Self:
    """
    Transform the atoms in `self` by `transform`.
    If `selection` is given, only transform the atoms in `selection`.
    """
    transform = self.change_transform(Transform3D.make(transform), self.get_frame(), frame)
    return self.with_atoms(self.get_atoms(self.get_frame()).transform(transform, selection, transform_velocities=transform_velocities))

transform_cell

transform_cell(
    transform: AffineTransform3D,
    frame: CoordinateFrame = "local",
) -> Self

Apply the given transform to the unit cell, without changing atom positions. The transform is applied in coordinate frame 'frame'.

Source code in atomlib/atomcell.py

def transform_cell(self, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> Self:
    """
    Apply the given transform to the unit cell, without changing atom positions.
    The transform is applied in coordinate frame 'frame'.
    """
    return self.with_cell(self.get_cell().transform_cell(transform, frame=frame))

transform

transform(
    transform: AffineTransform3D,
    frame: CoordinateFrame = "local",
) -> Self

Source code in atomlib/atomcell.py

def transform(self, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> Self:
    if isinstance(transform, Transform3D) and not isinstance(transform, AffineTransform3D):
        raise ValueError("Non-affine transforms cannot change the box dimensions. Use 'transform_atoms' instead.")
    # TODO: cleanup once tests pass
    # coordinate change the transform into atomic coordinates
    new_cell = self.get_cell().transform_cell(transform, frame)
    transform = self.get_cell().change_transform(transform, self.get_frame(), frame)
    return self.with_atoms(self.get_atoms().transform(transform), self.get_frame()).with_cell(new_cell)

crop

crop(
    x_min: float = -inf,
    x_max: float = inf,
    y_min: float = -inf,
    y_max: float = inf,
    z_min: float = -inf,
    z_max: float = inf,
    *,
    frame: CoordinateFrame = "local"
) -> Self

Crop atoms and cell to the given extents. For a non-orthogonal cell, this must be specified in cell coordinates. This function implicity explodes the cell as well.

To crop atoms only, use crop_atoms instead.

Source code in atomlib/atomcell.py

def crop(self, x_min: float = -numpy.inf, x_max: float = numpy.inf,
         y_min: float = -numpy.inf, y_max: float = numpy.inf,
         z_min: float = -numpy.inf, z_max: float = numpy.inf, *,
         frame: CoordinateFrame = 'local') -> Self:
    """
    Crop atoms and cell to the given extents. For a non-orthogonal
    cell, this must be specified in cell coordinates. This
    function implicity `explode`s the cell as well.

    To crop atoms only, use `crop_atoms` instead.
    """

    cell = self.get_cell().crop(x_min, x_max, y_min, y_max, z_min, z_max, frame=frame)
    atoms = self._transform_atoms_in_frame(frame, lambda atoms: atoms.crop_atoms(x_min, x_max, y_min, y_max, z_min, z_max))
    return self.with_cell(cell).with_atoms(atoms)

crop_atoms

crop_atoms(
    x_min: float = -inf,
    x_max: float = inf,
    y_min: float = -inf,
    y_max: float = inf,
    z_min: float = -inf,
    z_max: float = inf,
    *,
    frame: CoordinateFrame = "local"
) -> Self

Source code in atomlib/atomcell.py

def crop_atoms(self, x_min: float = -numpy.inf, x_max: float = numpy.inf,
               y_min: float = -numpy.inf, y_max: float = numpy.inf,
               z_min: float = -numpy.inf, z_max: float = numpy.inf, *,
               frame: CoordinateFrame = 'local') -> Self:
    atoms = self._transform_atoms_in_frame(frame, lambda atoms: atoms.crop_atoms(x_min, x_max, y_min, y_max, z_min, z_max))
    return self.with_atoms(atoms)

crop_to_box

crop_to_box(eps: float = 1e-05) -> Self

Source code in atomlib/atomcell.py

def crop_to_box(self, eps: float = 1e-5) -> Self:
    atoms = self._transform_atoms_in_frame('cell_box', lambda atoms: atoms.crop_atoms(*([-eps, 1-eps]*3)))
    return self.with_atoms(atoms)

wrap

wrap(eps: float = 1e-05) -> Self

Wrap atoms around the cell boundaries.

Source code in atomlib/atomcell.py

def wrap(self, eps: float = 1e-5) -> Self:
    """Wrap atoms around the cell boundaries."""
    return self.with_atoms(self._transform_atoms_in_frame('cell_box', lambda a: a._wrap(eps)))

repeat

repeat(n: Union[int, VecLike]) -> Self

Tile the cell

Source code in atomlib/atomcell.py

def repeat(self, n: t.Union[int, VecLike]) -> Self:
    """Tile the cell"""
    ns = numpy.broadcast_to(n, 3)
    if not numpy.issubdtype(ns.dtype, numpy.integer):
        raise ValueError("repeat() argument must be an integer or integer array.")

    cells = numpy.stack(numpy.meshgrid(*map(numpy.arange, ns))) \
        .reshape(3, -1).T.astype(float)
    cells = cells * self.box_size

    atoms = self.get_atoms('cell')
    atoms = Atoms.concat([
        atoms.transform(AffineTransform3D.translate(cell))
        for cell in cells
    ]) #.transform(self.cell.get_transform('local', 'cell_frac'))
    return self.with_atoms(atoms, 'cell').with_cell(self.get_cell().repeat(ns))

repeat_to

repeat_to(
    size: VecLike, crop: Union[bool, Sequence[bool]] = False
) -> Self

Repeat the cell so it is at least size along the crystal's axes.

If crop, then crop the cell to exactly size. This may break periodicity. crop may be a vector, in which case you can specify cropping only along some axes.

Source code in atomlib/atomcell.py

def repeat_to(self, size: VecLike, crop: t.Union[bool, t.Sequence[bool]] = False) -> Self:
    """
    Repeat the cell so it is at least `size` along the crystal's axes.

    If `crop`, then crop the cell to exactly `size`. This may break periodicity.
    `crop` may be a vector, in which case you can specify cropping only along some axes.
    """
    size = to_vec3(size)
    cell_size = self.cell_size * self.n_cells
    repeat = numpy.maximum(numpy.ceil(size / cell_size).astype(int), 1)
    atom_cell = self.repeat(repeat)

    crop_v = to_vec3(crop, dtype=numpy.bool_)
    if numpy.any(crop_v):
        crop_x, crop_y, crop_z = crop_v
        return atom_cell.crop(
            x_max = size[0] if crop_x else numpy.inf,
            y_max = size[1] if crop_y else numpy.inf,
            z_max = size[2] if crop_z else numpy.inf,
            frame='cell'
        )

    return atom_cell

repeat_x

repeat_x(n: int) -> Self

Tile the cell in the x axis.

Source code in atomlib/atomcell.py

def repeat_x(self, n: int) -> Self:
    """Tile the cell in the x axis."""
    return self.repeat((n, 1, 1))

repeat_y

repeat_y(n: int) -> Self

Tile the cell in the y axis.

Source code in atomlib/atomcell.py

def repeat_y(self, n: int) -> Self:
    """Tile the cell in the y axis."""
    return self.repeat((1, n, 1))

repeat_z

repeat_z(n: int) -> Self

Tile the cell in the z axis.

Source code in atomlib/atomcell.py

def repeat_z(self, n: int) -> Self:
    """Tile the cell in the z axis."""
    return self.repeat((1, 1, n))

repeat_to_x

repeat_to_x(size: float, crop: bool = False) -> Self

Repeat the cell so it is at least size size along the x axis.

Source code in atomlib/atomcell.py

def repeat_to_x(self, size: float, crop: bool = False) -> Self:
    """Repeat the cell so it is at least size `size` along the x axis."""
    return self.repeat_to([size, 0., 0.], [crop, False, False])

repeat_to_y

repeat_to_y(size: float, crop: bool = False) -> Self

Repeat the cell so it is at least size size along the y axis.

Source code in atomlib/atomcell.py

def repeat_to_y(self, size: float, crop: bool = False) -> Self:
    """Repeat the cell so it is at least size `size` along the y axis."""
    return self.repeat_to([0., size, 0.], [False, crop, False])

repeat_to_z

repeat_to_z(size: float, crop: bool = False) -> Self

Repeat the cell so it is at least size size along the z axis.

Source code in atomlib/atomcell.py

def repeat_to_z(self, size: float, crop: bool = False) -> Self:
    """Repeat the cell so it is at least size `size` along the z axis."""
    return self.repeat_to([0., 0., size], [False, False, crop])

repeat_to_aspect

repeat_to_aspect(
    plane: Literal["xy", "xz", "yz"] = "xy",
    *,
    aspect: float = 1.0,
    min_size: Optional[VecLike] = None,
    max_size: Optional[VecLike] = None
) -> Self

Repeat to optimize the aspect ratio in plane, while staying above min_size and under max_size.

Source code in atomlib/atomcell.py

def repeat_to_aspect(self, plane: t.Literal['xy', 'xz', 'yz'] = 'xy', *,
                     aspect: float = 1., min_size: t.Optional[VecLike] = None,
                     max_size: t.Optional[VecLike] = None) -> Self:
    """
    Repeat to optimize the aspect ratio in `plane`,
    while staying above `min_size` and under `max_size`.
    """
    if min_size is None:
        min_n = numpy.array([1, 1, 1], numpy.int_)
    else:
        min_n = numpy.maximum(numpy.ceil(to_vec3(min_size) / self.box_size), 1).astype(numpy.int_)

    if max_size is None:
        max_n = 3 * min_n
    else:
        max_n = numpy.maximum(numpy.floor(to_vec3(max_size) / self.box_size), 1).astype(numpy.int_)

    if plane == 'xy':
        indices = [0, 1]
    elif plane == 'xz':
        indices = [0, 2]
    elif plane == 'yz':
        indices = [1, 2]
    else:
        raise ValueError(f"Invalid plane '{plane}'. Exepcted 'xy', 'xz', 'or 'yz'.")

    na = numpy.arange(min_n[indices[0]], max_n[indices[0]])
    nb = numpy.arange(min_n[indices[1]], max_n[indices[1]])
    (na, nb) = numpy.meshgrid(na, nb)

    aspects = na * self.box_size[indices[0]] / (nb * self.box_size[indices[1]])
    # cost function: log(aspect)^2  (so cost(0.5) == cost(2))
    min_i = numpy.argmin(numpy.log(aspects / aspect)**2)
    repeat = numpy.array([1, 1, 1], numpy.int_)
    repeat[indices] = na.flatten()[min_i], nb.flatten()[min_i]
    return self.repeat(repeat)

explode

explode() -> Self

Materialize repeated cells as one supercell.

Source code in atomlib/atomcell.py

def explode(self) -> Self:
    """Materialize repeated cells as one supercell."""
    frame = self.get_frame()

    return self.with_atoms(self.get_atoms('local'), 'local') \
        .with_cell(self.get_cell().explode()) \
        .to_frame(frame)

periodic_duplicate

periodic_duplicate(eps: float = 1e-05) -> Self

Add duplicate copies of atoms near periodic boundaries.

For instance, an atom at a corner will be duplicated into 8 copies. This is mostly only useful for visualization.

Source code in atomlib/atomcell.py

def periodic_duplicate(self, eps: float = 1e-5) -> Self:
    """
    Add duplicate copies of atoms near periodic boundaries.

    For instance, an atom at a corner will be duplicated into 8 copies.
    This is mostly only useful for visualization.
    """
    frame_save = self.get_frame()
    self = self.to_frame('cell_box').wrap(eps=eps)

    for i in range(3):
        self = self.concat((self,
            self.filter(polars.col('coords').arr.get(i).abs() <= eps, frame='cell_box')
                .transform_atoms(AffineTransform3D.translate([1. if i == j else 0. for j in range(3)]), frame='cell_box')
        ))

    return self.to_frame(frame_save)

describe

describe(
    percentiles: Union[Sequence[float], float, None] = (
        0.25,
        0.5,
        0.75,
    ),
    *,
    interpolation: RollingInterpolationMethod = "nearest",
    frame: Optional[CoordinateFrame] = None
) -> DataFrame

Return summary statistics for self. See DataFrame.describe for more information.

PARAMETER	DESCRIPTION
percentiles	List of percentiles/quantiles to include. Defaults to 25% (first quartile), 50% (median), and 75% (third quartile). TYPE: Union[Sequence[float], float, None] DEFAULT: (0.25, 0.5, 0.75)

RETURNS	DESCRIPTION
DataFrame	A dataframe containing summary statistics (mean, std. deviation, percentiles, etc.) for each column.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def describe(self, percentiles: t.Union[t.Sequence[float], float, None] = (0.25, 0.5, 0.75), *,
             interpolation: RollingInterpolationMethod = 'nearest',
             frame: t.Optional[CoordinateFrame] = None) -> polars.DataFrame:
    """
    Return summary statistics for `self`. See [`DataFrame.describe`][polars.DataFrame.describe] for more information.

    Args:
      percentiles: List of percentiles/quantiles to include. Defaults to 25% (first quartile),
                   50% (median), and 75% (third quartile).

    Returns:
      A dataframe containing summary statistics (mean, std. deviation, percentiles, etc.) for each column.
    """
    ...

with_columns

with_columns(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Self

Return a copy of self with the given columns added.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_columns(self,
                 *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
                 frame: t.Optional[CoordinateFrame] = None,
                 **named_exprs: IntoExpr) -> Self:
    """Return a copy of `self` with the given columns added."""
    ...

get_column

get_column(
    name: str, *, frame: Optional[CoordinateFrame] = None
) -> Series

Get the specified column from self, raising polars.ColumnNotFoundError if it's not present.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def get_column(self, name: str, *, frame: t.Optional[CoordinateFrame] = None) -> polars.Series:
    """
    Get the specified column from `self`, raising [`polars.ColumnNotFoundError`][polars.exceptions.ColumnNotFoundError] if it's not present.

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    ...

get_columns

get_columns(
    *, frame: Optional[CoordinateFrame] = None
) -> List[Series]

Return all columns from self as a list of Series.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def get_columns(self, *, frame: t.Optional[CoordinateFrame] = None) -> t.List[polars.Series]:
    """
    Return all columns from `self` as a list of [`Series`][polars.Series].

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    ...

group_by

group_by(
    *by: Union[IntoExpr, Iterable[IntoExpr]],
    maintain_order: bool = False,
    frame: Optional[CoordinateFrame] = None,
    **named_by: IntoExpr
) -> GroupBy

Start a group by operation. See DataFrame.group_by for more information.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def group_by(self, *by: t.Union[IntoExpr, t.Iterable[IntoExpr]],
             maintain_order: bool = False, frame: t.Optional[CoordinateFrame] = None,
             **named_by: IntoExpr) -> polars.dataframe.group_by.GroupBy:
    """
    Start a group by operation. See [`DataFrame.group_by`][polars.DataFrame.group_by] for more information.
    """
    ...

pipe

pipe(
    function: Callable[Concatenate[HasAtomCellT, P], T],
    *args: args,
    **kwargs: kwargs
) -> T

Apply function to self (in method-call syntax).

Source code in atomlib/atomcell.py

def pipe(self: HasAtomCellT, function: t.Callable[Concatenate[HasAtomCellT, P], T], *args: P.args, **kwargs: P.kwargs) -> T:
    """Apply `function` to `self` (in method-call syntax)."""
    return function(self, *args, **kwargs)

filter

filter(
    *predicates: Union[
        None,
        IntoExprColumn,
        Iterable[IntoExprColumn],
        bool,
        List[bool],
        ndarray,
    ],
    frame: Optional[CoordinateFrame] = None,
    **constraints: Any
) -> Self

Filter self, removing rows which evaluate to False.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def filter(
    self,
    *predicates: t.Union[None, IntoExprColumn, t.Iterable[IntoExprColumn], bool, t.List[bool], numpy.ndarray],
    frame: t.Optional[CoordinateFrame] = None,
    **constraints: t.Any,
) -> Self:
    """Filter `self`, removing rows which evaluate to `False`."""
    ...

sort

sort(
    by: Union[IntoExpr, Iterable[IntoExpr]],
    *more_by: IntoExpr,
    descending: Union[bool, Sequence[bool]] = False,
    nulls_last: bool = False
) -> Self

Sort the atoms in self by the given columns/expressions.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def sort(
    self,
    by: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    *more_by: IntoExpr,
    descending: t.Union[bool, t.Sequence[bool]] = False,
    nulls_last: bool = False,
) -> Self:
    """
    Sort the atoms in `self` by the given columns/expressions.
    """
    ...

slice

slice(
    offset: int,
    length: Optional[int] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return a slice of the rows in self.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def slice(self, offset: int, length: t.Optional[int] = None, *,
          frame: t.Optional[CoordinateFrame] = None) -> Self:
    """Return a slice of the rows in `self`."""
    ...

head

head(
    n: int = 5, *, frame: Optional[CoordinateFrame] = None
) -> Self

Return the first n rows of self.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def head(self, n: int = 5, *, frame: t.Optional[CoordinateFrame] = None) -> Self:
    """Return the first `n` rows of `self`."""
    ...

tail

tail(
    n: int = 5, *, frame: Optional[CoordinateFrame] = None
) -> Self

Return the last n rows of self.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def tail(self, n: int = 5, *, frame: t.Optional[CoordinateFrame] = None) -> Self:
    """Return the last `n` rows of `self`."""
    ...

fill_null

fill_null(
    value: Any = None,
    strategy: Optional[FillNullStrategy] = None,
    limit: Optional[int] = None,
    matches_supertype: bool = True,
) -> Self

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def fill_null(
    self, value: t.Any = None, strategy: t.Optional[FillNullStrategy] = None,
    limit: t.Optional[int] = None, matches_supertype: bool = True,
) -> Self:
    ...

fill_nan

fill_nan(
    value: Union[Expr, int, float, None],
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def fill_nan(self, value: t.Union[polars.Expr, int, float, None], *,
             frame: t.Optional[CoordinateFrame] = None) -> Self:
    ...

select

select(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> DataFrame

Select exprs from self, and return as a polars.DataFrame.

Expressions may either be columns or expressions of columns.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def select(
    self, *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    frame: t.Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> polars.DataFrame:
    """
    Select `exprs` from `self`, and return as a [`polars.DataFrame`][polars.DataFrame].

    Expressions may either be columns or expressions of columns.

    [polars.DataFrame]: https://docs.pola.rs/py-polars/html/reference/dataframe/index.html
    """
    ...

select_props

select_props(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Self

Select exprs from self, while keeping required columns. Doesn't affect the cell.

RETURNS	DESCRIPTION
Self	A HasAtomCell filtered to contain
Self	the specified properties (as well as required columns).

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def select_props(
    self,
    *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    frame: t.Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Self:
    """
    Select `exprs` from `self`, while keeping required columns.
    Doesn't affect the cell.

    Returns:
      A [`HasAtomCell`][atomlib.atomcell.HasAtomCell] filtered to contain
      the specified properties (as well as required columns).
    """
    ...

try_select

try_select(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Optional[DataFrame]

Try to select exprs from self, and return as a polars.DataFrame.

Expressions may either be columns or expressions of columns. Returns None if any columns are missing.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def try_select(
    self, *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    frame: t.Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> t.Optional[polars.DataFrame]:
    """
    Try to select `exprs` from `self`, and return as a [`polars.DataFrame`][polars.DataFrame].

    Expressions may either be columns or expressions of columns. Returns `None` if any
    columns are missing.

    [polars.DataFrame]: https://docs.pola.rs/py-polars/html/reference/dataframe/index.html
    """
    ...

round_near_zero

round_near_zero(
    tol: float = 1e-14,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Round atom position values near zero to zero.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def round_near_zero(self, tol: float = 1e-14, *,
                    frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Round atom position values near zero to zero.
    """
    ...

coords

coords(
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> NDArray[float64]

Return a (N, 3) ndarray of atom positions (dtype numpy.float64) in the given coordinate frame.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def coords(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Optional[CoordinateFrame] = None) -> NDArray[numpy.float64]:
    """
    Return a `(N, 3)` ndarray of atom positions (dtype [`numpy.float64`][numpy.float64])
    in the given coordinate frame.
    """
    ...

velocities

velocities(
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Optional[NDArray[float64]]

Return a (N, 3) ndarray of atom velocities (dtype numpy.float64) in the given coordinate frame.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def velocities(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Optional[CoordinateFrame] = None) -> t.Optional[NDArray[numpy.float64]]:
    """
    Return a `(N, 3)` ndarray of atom velocities (dtype [`numpy.float64`][numpy.float64])
    in the given coordinate frame.
    """
    ...

add_atom

add_atom(
    elem: Union[int, str],
    x: ArrayLike,
    /,
    *,
    y: None = None,
    z: None = None,
    frame: Optional[CoordinateFrame] = None,
    **kwargs: Any,
) -> Self

add_atom(
    elem: Union[int, str],
    /,
    x: float,
    y: float,
    z: float,
    *,
    frame: Optional[CoordinateFrame] = None,
    **kwargs: Any,
) -> Self

add_atom(
    elem: Union[int, str],
    /,
    x: Union[ArrayLike, float],
    y: Optional[float] = None,
    z: Optional[float] = None,
    *,
    frame: Optional[CoordinateFrame] = None,
    **kwargs: Any,
) -> Self

Return a copy of self with an extra atom.

By default, all extra columns present in self must be specified as **kwargs.

Try to avoid calling this in a loop (Use concat instead).

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def add_atom(self, elem: t.Union[int, str], /,  # type: ignore (spurious)
             x: t.Union[ArrayLike, float],
             y: t.Optional[float] = None,
             z: t.Optional[float] = None, *,
             frame: t.Optional[CoordinateFrame] = None,
             **kwargs: t.Any) -> Self:
    """
    Return a copy of `self` with an extra atom.

    By default, all extra columns present in `self` must be specified as `**kwargs`.

    Try to avoid calling this in a loop (Use [`concat`][atomlib.atomcell.HasAtomCell.concat] instead).
    """
    ...

with_index

with_index(
    index: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Returns self with a row index added in column 'i' (dtype polars.Int64). If index is not specified, defaults to an existing index or a new index.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_index(self, index: t.Optional[AtomValues] = None, *,
               frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Returns `self` with a row index added in column 'i' (dtype [`polars.Int64`][polars.datatypes.Int64]).
    If `index` is not specified, defaults to an existing index or a new index.
    """
    ...

with_wobble

with_wobble(
    wobble: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given displacements in column 'wobble' (dtype polars.Float64). If wobble is not specified, defaults to the already-existing wobbles or 0.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_wobble(self, wobble: t.Optional[AtomValues] = None, *,
                frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given displacements in column 'wobble' (dtype [`polars.Float64`][polars.datatypes.Float64]).
    If `wobble` is not specified, defaults to the already-existing wobbles or 0.
    """
    ...

with_occupancy

with_occupancy(
    frac_occupancy: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given fractional occupancies (dtype polars.Float64). If frac_occupancy is not specified, defaults to the already-existing occupancies or 1.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_occupancy(self, frac_occupancy: t.Optional[AtomValues] = None, *,
                   frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return self with the given fractional occupancies (dtype [`polars.Float64`][polars.datatypes.Float64]).
    If `frac_occupancy` is not specified, defaults to the already-existing occupancies or 1.
    """
    ...

apply_wobble

apply_wobble(
    rng: Union[Generator, int, None] = None,
    frame: Optional[CoordinateFrame] = None,
) -> Self

Displace the atoms in self by the amount in the wobble column. wobble is interpretated as a mean-squared displacement, which is distributed equally over each axis.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def apply_wobble(self, rng: t.Union[numpy.random.Generator, int, None] = None,
                 frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Displace the atoms in `self` by the amount in the `wobble` column.
    `wobble` is interpretated as a mean-squared displacement, which is distributed
    equally over each axis.
    """
    ...

with_type

with_type(
    types: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given atom types in column 'type'. If types is not specified, use the already existing types or auto-assign them.

When auto-assigning, each symbol is given a unique value, case-sensitive. Values are assigned from lowest atomic number to highest. For instance: ["Ag+", "Na", "H", "Ag"] => [3, 11, 1, 2]

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_type(self, types: t.Optional[AtomValues] = None, *,
              frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given atom types in column 'type'.
    If `types` is not specified, use the already existing types or auto-assign them.

    When auto-assigning, each symbol is given a unique value, case-sensitive.
    Values are assigned from lowest atomic number to highest.
    For instance: `["Ag+", "Na", "H", "Ag"]` => `[3, 11, 1, 2]`
    """
    ...

with_mass

with_mass(
    mass: Optional[ArrayLike] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given atom masses in column 'mass'. If mass is not specified, use the already existing masses or auto-assign them.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_mass(self, mass: t.Optional[ArrayLike] = None, *,
              frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given atom masses in column `'mass'`.
    If `mass` is not specified, use the already existing masses or auto-assign them.
    """
    ...

with_symbol

with_symbol(
    symbols: ArrayLike,
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given atomic symbols.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_symbol(self, symbols: ArrayLike, selection: t.Optional[AtomSelection] = None, *,
                frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given atomic symbols.
    """
    ...

with_coords

with_coords(
    pts: ArrayLike,
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self replaced with the given atomic positions.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_coords(self, pts: ArrayLike, selection: t.Optional[AtomSelection] = None, *,
                frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` replaced with the given atomic positions.
    """
    ...

with_velocity

with_velocity(
    pts: Optional[ArrayLike] = None,
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self replaced with the given atomic velocities. If pts is not specified, use the already existing velocities or zero.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_velocity(self, pts: t.Optional[ArrayLike] = None,
                  selection: t.Optional[AtomSelection] = None, *,
                  frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` replaced with the given atomic velocities.
    If `pts` is not specified, use the already existing velocities or zero.
    """
    ...

AtomCell dataclass

Bases: AtomCellIOMixin, HasAtomCell

Cell of atoms with known size and periodic boundary conditions.

Source code in atomlib/atomcell.py

@dataclass(init=False, repr=False, frozen=True)
class AtomCell(AtomCellIOMixin, HasAtomCell):
    """
    Cell of atoms with known size and periodic boundary conditions.
    """

    atoms: Atoms
    """Atoms in the cell. Stored in 'local' coordinates (i.e. relative to the enclosing group but not relative to box dimensions)."""

    cell: Cell
    """Cell coordinate system."""

    frame: CoordinateFrame = 'local'
    """Coordinate frame 'atoms' are stored in."""

    def get_cell(self) -> Cell:
        return self.cell

    def with_cell(self, cell: Cell) -> Self:
        return self.__class__(self.atoms, cell, frame=self.frame, keep_frame=True)

    def get_atoms(self, frame: t.Optional[CoordinateFrame] = None) -> Atoms:
        """Get atoms contained in ``self``, in the given coordinate frame."""

        if frame is None or frame == self.get_frame():
            return self.atoms
        return self.atoms.transform(self.get_transform(frame, self.get_frame()))

    def with_atoms(self, atoms: HasAtoms, frame: t.Optional[CoordinateFrame] = None) -> Self:
        frame = frame if frame is not None else self.frame
        return self.__class__(atoms.get_atoms(), cell=self.cell, frame=frame, keep_frame=True)
        #return replace(self, atoms=atoms, frame = frame if frame is not None else self.frame, keep_frame=True)

    def get_frame(self) -> CoordinateFrame:
        """Get the coordinate frame atoms are stored in."""
        return self.frame

    @classmethod
    def _combine_metadata(cls: t.Type[AtomCellT], *atoms: HasAtoms, n: t.Optional[int] = None) -> AtomCellT:
        """
        When combining multiple [`HasAtoms`][atomlib.atoms.HasAtoms], check that they are compatible with each other,
        and return a 'representative' which best represents the combined metadata.
        Implementors should treat [`Atoms`][atomlib.atoms.Atoms] as acceptable, but having no metadata.
        """
        if n is not None:
            rep = atoms[n]
            if not isinstance(rep, AtomCell):
                raise ValueError(f"Atoms #{n} has no cell")
        else:
            atom_cells = [a for a in atoms if isinstance(a, AtomCell)]
            if len(atom_cells) == 0:
                raise TypeError("No AtomCells to combine")
            rep = atom_cells[0]
            if not all(a.cell == rep.cell for a in atom_cells[1:]):
                raise TypeError("Can't combine AtomCells with different cells")

        return cls(Atoms.empty(), frame=rep.frame, cell=rep.cell)

    @classmethod
    def from_ortho(cls, atoms: IntoAtoms, ortho: LinearTransform3D, *,
                   n_cells: t.Optional[VecLike] = None,
                   frame: CoordinateFrame = 'local',
                   keep_frame: bool = False):
        """
        Make an atom cell given a list of atoms and an orthogonalization matrix.
        Atoms are assumed to be in the coordinate system `frame`.
        """
        cell = Cell.from_ortho(ortho, n_cells)
        return cls(atoms, cell, frame=frame, keep_frame=keep_frame)

    @classmethod
    def from_unit_cell(cls, atoms: IntoAtoms, cell_size: VecLike,
                       cell_angle: t.Optional[VecLike] = None, *,
                       n_cells: t.Optional[VecLike] = None,
                       frame: CoordinateFrame = 'local',
                       keep_frame: bool = False):
        """
        Make a cell given a list of atoms and unit cell parameters.
        Atoms are assumed to be in the coordinate system `frame`.
        """
        cell = Cell.from_unit_cell(cell_size, cell_angle, n_cells=n_cells)
        return cls(atoms, cell, frame=frame, keep_frame=keep_frame)

    def __init__(self, atoms: IntoAtoms, cell: Cell, *,
                 frame: CoordinateFrame = 'local',
                 keep_frame: bool = False):
        atoms = Atoms(atoms)
        # by default, store in local coordinates
        if not keep_frame and frame != 'local':
            atoms = atoms.transform(cell.get_transform('local', frame))
            frame = 'local'

        object.__setattr__(self, 'atoms', atoms)
        object.__setattr__(self, 'cell', cell)
        object.__setattr__(self, 'frame', frame)

        self.__post_init__()

    def __post_init__(self):
        pass

    def orthogonalize(self) -> OrthoCell:
        if self.is_orthogonal():
            return OrthoCell(self.atoms, self.cell, frame=self.frame)
        raise NotImplementedError()

    def clone(self: AtomCellT) -> AtomCellT:
        """Make a deep copy of `self`."""
        return self.__class__(**{field.name: copy.deepcopy(getattr(self, field.name)) for field in fields(self)})

    def assert_equal(self, other: t.Any):
        """Assert this structure is equal to `other`"""
        assert isinstance(other, AtomCell)
        self.cell.assert_equal(other.cell)
        self.get_atoms('local').assert_equal(other.get_atoms('local'))

    def _str_parts(self) -> t.Iterable[t.Any]:
        return (
            f"Cell size:  {self.cell.cell_size!s}",
            f"Cell angle: {self.cell.cell_angle!s}",
            f"# Cells: {self.cell.n_cells!s}",
            f"Frame: {self.frame}",
            self.atoms,
        )

    def __str__(self) -> str:
        return "\n".join(map(str, self._str_parts()))

    def __repr__(self) -> str:
        return f"{self.__class__.__name__}({self.atoms!r}, cell={self.cell!r}, frame={self.frame})"

    def _repr_pretty_(self, p: t.Any, cycle: bool) -> None:
        p.text(f'{self.__class__.__name__}(...)') if cycle else p.text(str(self))

affine property

affine: AffineTransform3D

Affine transformation. Holds transformation from 'ortho' to 'local' coordinates, including rotation away from the standard crystal orientation.

ortho property

ortho: LinearTransform3D

Orthogonalization transformation. Skews but does not scale the crystal axes to cartesian axes.

metric property

metric: LinearTransform3D

Cell metric tensor

Returns the dot product between every combination of basis vectors. :math:\mathbf{a} \cdot \mathbf{b} = a_i M_ij b_j

cell_size property

cell_size: NDArray[float64]

Unit cell size.

cell_angle property

cell_angle: NDArray[float64]

Unit cell angles, in radians.

n_cells property

n_cells: NDArray[int_]

Number of unit cells.

pbc property

pbc: NDArray[bool_]

Flags indicating the presence of periodic boundary conditions along each axis.

ortho_size property

ortho_size: NDArray[float64]

Return size of orthogonal unit cell.

Equivalent to the diagonal of the orthogonalization matrix.

box_size property

box_size: NDArray[float64]

Return size of the cell box.

Equivalent to self.n_cells * self.cell_size.

columns property

columns: List[str]

Return the column names in self.

RETURNS	DESCRIPTION
List[str]	A sequence of column names

dtypes property

dtypes: List[DataType]

Return the datatypes in self.

RETURNS	DESCRIPTION
List[DataType]	A sequence of column DataTypes

schema property

schema: Schema

Return the schema of self.

RETURNS	DESCRIPTION
Schema	A dictionary of column names and DataTypes

with_column class-attribute instance-attribute

with_column = with_columns

unique class-attribute instance-attribute

unique = deduplicate

atoms instance-attribute

atoms: Atoms

Atoms in the cell. Stored in 'local' coordinates (i.e. relative to the enclosing group but not relative to box dimensions).

cell instance-attribute

cell: Cell

Cell coordinate system.

frame class-attribute instance-attribute

frame: CoordinateFrame = 'local'

Coordinate frame 'atoms' are stored in.

get_transform

get_transform(
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> AffineTransform3D

In the two-argument form, get the transform to frame_to from frame_from. In the one-argument form, get the transform from local coordinates to 'frame'.

Source code in atomlib/cell.py

def get_transform(self, frame_to: t.Optional[CoordinateFrame] = None, frame_from: t.Optional[CoordinateFrame] = None) -> AffineTransform3D:
    """
    In the two-argument form, get the transform to `frame_to` from `frame_from`.
    In the one-argument form, get the transform from local coordinates to 'frame'.
    """
    transform_from = self._get_transform_to_local(frame_from) if frame_from is not None else AffineTransform3D()
    transform_to = self._get_transform_to_local(frame_to) if frame_to is not None else AffineTransform3D()
    if frame_from is not None and frame_to is not None and frame_from.lower() == frame_to.lower():
        return AffineTransform3D()
    return transform_to.inverse() @ transform_from

corners

corners(frame: CoordinateFrame = 'local') -> ndarray

Source code in atomlib/cell.py

def corners(self, frame: CoordinateFrame = 'local') -> numpy.ndarray:
    corners = numpy.array(list(itertools.product((0., 1.), repeat=3)))
    return self.get_transform(frame, 'cell_box') @ corners

bbox_cell

bbox_cell(frame: CoordinateFrame = 'local') -> BBox3D

Return the bounding box of the cell box in the given coordinate system.

Source code in atomlib/cell.py

def bbox_cell(self, frame: CoordinateFrame = 'local') -> BBox3D:
    """Return the bounding box of the cell box in the given coordinate system."""
    return BBox3D.from_pts(self.corners(frame))

bbox

bbox(frame: CoordinateFrame = 'local') -> BBox3D

Return the combined bounding box of the cell and atoms in the given coordinate system. To get the cell or atoms bounding box only, use bbox_cell or bbox_atoms.

Source code in atomlib/atomcell.py

def bbox(self, frame: CoordinateFrame = 'local') -> BBox3D:
    """
    Return the combined bounding box of the cell and atoms in the given coordinate system.
    To get the cell or atoms bounding box only, use [`bbox_cell`][atomlib.atomcell.HasAtomCell.bbox_cell] or [`bbox_atoms`][atomlib.atomcell.HasAtomCell.bbox_atoms].
    """
    return self.bbox_atoms(frame) | self.bbox_cell(frame)

is_orthogonal

is_orthogonal(tol: float = 1e-08) -> bool

Returns whether this cell is orthogonal (axes are at right angles.)

Source code in atomlib/cell.py

def is_orthogonal(self, tol: float = 1e-8) -> bool:
    """Returns whether this cell is orthogonal (axes are at right angles.)"""
    return self.ortho.is_diagonal(tol=tol)

is_orthogonal_in_local

is_orthogonal_in_local(tol: float = 1e-08) -> bool

Returns whether this cell is orthogonal and aligned with the local coordinate system.

Source code in atomlib/cell.py

def is_orthogonal_in_local(self, tol: float = 1e-8) -> bool:
    """Returns whether this cell is orthogonal and aligned with the local coordinate system."""
    transform = (self.affine @ self.ortho).to_linear()
    if not transform.is_scaled_orthogonal(tol):
        return False
    normed = transform.inner / numpy.linalg.norm(transform.inner, axis=-2, keepdims=True)
    # every row of transform must be a +/- 1 times a basis vector (i, j, or k)
    return all(
        any(numpy.isclose(numpy.abs(numpy.dot(row, v)), 1., atol=tol) for v in numpy.eye(3))
        for row in normed
    )

to_ortho

to_ortho() -> AffineTransform3D

Source code in atomlib/cell.py

def to_ortho(self) -> AffineTransform3D:
    return self.get_transform('local', 'cell_box')

transform_cell

transform_cell(
    transform: AffineTransform3D,
    frame: CoordinateFrame = "local",
) -> Self

Apply the given transform to the unit cell, without changing atom positions. The transform is applied in coordinate frame 'frame'.

Source code in atomlib/atomcell.py

def transform_cell(self, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> Self:
    """
    Apply the given transform to the unit cell, without changing atom positions.
    The transform is applied in coordinate frame 'frame'.
    """
    return self.with_cell(self.get_cell().transform_cell(transform, frame=frame))

strain_orthogonal

strain_orthogonal() -> HasCellT

Orthogonalize using strain.

Strain is applied such that the x-axis remains fixed, and the y-axis remains in the xy plane. For small displacements, no hydrostatic strain is applied (volume is conserved).

Source code in atomlib/cell.py

def strain_orthogonal(self: HasCellT) -> HasCellT:
    """
    Orthogonalize using strain.

    Strain is applied such that the x-axis remains fixed, and the y-axis remains in the xy plane.
    For small displacements, no hydrostatic strain is applied (volume is conserved).
    """
    return self.with_cell(Cell(
        affine=self.affine,
        ortho=LinearTransform3D(),
        cell_size=self.cell_size,
        n_cells=self.n_cells,
        pbc=self.pbc,
    ))

repeat

repeat(n: Union[int, VecLike]) -> Self

Tile the cell

Source code in atomlib/atomcell.py

def repeat(self, n: t.Union[int, VecLike]) -> Self:
    """Tile the cell"""
    ns = numpy.broadcast_to(n, 3)
    if not numpy.issubdtype(ns.dtype, numpy.integer):
        raise ValueError("repeat() argument must be an integer or integer array.")

    cells = numpy.stack(numpy.meshgrid(*map(numpy.arange, ns))) \
        .reshape(3, -1).T.astype(float)
    cells = cells * self.box_size

    atoms = self.get_atoms('cell')
    atoms = Atoms.concat([
        atoms.transform(AffineTransform3D.translate(cell))
        for cell in cells
    ]) #.transform(self.cell.get_transform('local', 'cell_frac'))
    return self.with_atoms(atoms, 'cell').with_cell(self.get_cell().repeat(ns))

explode

explode() -> Self

Materialize repeated cells as one supercell.

Source code in atomlib/atomcell.py

def explode(self) -> Self:
    """Materialize repeated cells as one supercell."""
    frame = self.get_frame()

    return self.with_atoms(self.get_atoms('local'), 'local') \
        .with_cell(self.get_cell().explode()) \
        .to_frame(frame)

explode_z

explode_z() -> HasCellT

Materialize repeated cells as one supercell in z.

Source code in atomlib/cell.py

def explode_z(self: HasCellT) -> HasCellT:
    """Materialize repeated cells as one supercell in z."""
    return self.with_cell(Cell(
        affine=self.affine,
        ortho=self.ortho,
        cell_size=self.cell_size*[1, 1, self.n_cells[2]],
        n_cells=[*self.n_cells[:2], 1],
        cell_angle=self.cell_angle,
        pbc=self.pbc,
    ))

crop

crop(
    x_min: float = -inf,
    x_max: float = inf,
    y_min: float = -inf,
    y_max: float = inf,
    z_min: float = -inf,
    z_max: float = inf,
    *,
    frame: CoordinateFrame = "local"
) -> Self

Crop atoms and cell to the given extents. For a non-orthogonal cell, this must be specified in cell coordinates. This function implicity explodes the cell as well.

To crop atoms only, use crop_atoms instead.

Source code in atomlib/atomcell.py

def crop(self, x_min: float = -numpy.inf, x_max: float = numpy.inf,
         y_min: float = -numpy.inf, y_max: float = numpy.inf,
         z_min: float = -numpy.inf, z_max: float = numpy.inf, *,
         frame: CoordinateFrame = 'local') -> Self:
    """
    Crop atoms and cell to the given extents. For a non-orthogonal
    cell, this must be specified in cell coordinates. This
    function implicity `explode`s the cell as well.

    To crop atoms only, use `crop_atoms` instead.
    """

    cell = self.get_cell().crop(x_min, x_max, y_min, y_max, z_min, z_max, frame=frame)
    atoms = self._transform_atoms_in_frame(frame, lambda atoms: atoms.crop_atoms(x_min, x_max, y_min, y_max, z_min, z_max))
    return self.with_cell(cell).with_atoms(atoms)

change_transform

change_transform(
    transform: AffineTransform3D,
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> AffineTransform3D

change_transform(
    transform: Transform3D,
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> Transform3D

change_transform(
    transform: Transform3D,
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> Transform3D

Coordinate-change a transformation from frame_from into frame_to.

Source code in atomlib/cell.py

def change_transform(self, transform: Transform3D,
                     frame_to: t.Optional[CoordinateFrame] = None,
                     frame_from: t.Optional[CoordinateFrame] = None) -> Transform3D:
    """Coordinate-change a transformation from `frame_from` into `frame_to`."""
    if frame_to == frame_from and frame_to is not None:
        return transform
    coord_change = self.get_transform(frame_to, frame_from)
    return coord_change @ transform @ coord_change.inverse()

describe

describe(
    percentiles: Union[Sequence[float], float, None] = (
        0.25,
        0.5,
        0.75,
    ),
    *,
    interpolation: RollingInterpolationMethod = "nearest",
    frame: Optional[CoordinateFrame] = None
) -> DataFrame

Return summary statistics for self. See DataFrame.describe for more information.

PARAMETER	DESCRIPTION
percentiles	List of percentiles/quantiles to include. Defaults to 25% (first quartile), 50% (median), and 75% (third quartile). TYPE: Union[Sequence[float], float, None] DEFAULT: (0.25, 0.5, 0.75)

RETURNS	DESCRIPTION
DataFrame	A dataframe containing summary statistics (mean, std. deviation, percentiles, etc.) for each column.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def describe(self, percentiles: t.Union[t.Sequence[float], float, None] = (0.25, 0.5, 0.75), *,
             interpolation: RollingInterpolationMethod = 'nearest',
             frame: t.Optional[CoordinateFrame] = None) -> polars.DataFrame:
    """
    Return summary statistics for `self`. See [`DataFrame.describe`][polars.DataFrame.describe] for more information.

    Args:
      percentiles: List of percentiles/quantiles to include. Defaults to 25% (first quartile),
                   50% (median), and 75% (third quartile).

    Returns:
      A dataframe containing summary statistics (mean, std. deviation, percentiles, etc.) for each column.
    """
    ...

with_columns

with_columns(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Self

Return a copy of self with the given columns added.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_columns(self,
                 *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
                 frame: t.Optional[CoordinateFrame] = None,
                 **named_exprs: IntoExpr) -> Self:
    """Return a copy of `self` with the given columns added."""
    ...

insert_column

insert_column(index: int, column: Series) -> DataFrame

Source code in atomlib/atoms.py

@_fwd_frame_map
def insert_column(self, index: int, column: polars.Series) -> polars.DataFrame:
    return self._get_frame().insert_column(index, column)

get_column

get_column(
    name: str, *, frame: Optional[CoordinateFrame] = None
) -> Series

Get the specified column from self, raising polars.ColumnNotFoundError if it's not present.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def get_column(self, name: str, *, frame: t.Optional[CoordinateFrame] = None) -> polars.Series:
    """
    Get the specified column from `self`, raising [`polars.ColumnNotFoundError`][polars.exceptions.ColumnNotFoundError] if it's not present.

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    ...

get_columns

get_columns(
    *, frame: Optional[CoordinateFrame] = None
) -> List[Series]

Return all columns from self as a list of Series.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def get_columns(self, *, frame: t.Optional[CoordinateFrame] = None) -> t.List[polars.Series]:
    """
    Return all columns from `self` as a list of [`Series`][polars.Series].

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    ...

get_column_index

get_column_index(name: str) -> int

Get the index of a column by name, raising polars.ColumnNotFoundError if it's not present.

Source code in atomlib/atoms.py

@_fwd_frame(polars.DataFrame.get_column_index)
def get_column_index(self, name: str) -> int:
    """Get the index of a column by name, raising [`polars.ColumnNotFoundError`][polars.exceptions.ColumnNotFoundError] if it's not present."""
    ...

group_by

group_by(
    *by: Union[IntoExpr, Iterable[IntoExpr]],
    maintain_order: bool = False,
    frame: Optional[CoordinateFrame] = None,
    **named_by: IntoExpr
) -> GroupBy

Start a group by operation. See DataFrame.group_by for more information.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def group_by(self, *by: t.Union[IntoExpr, t.Iterable[IntoExpr]],
             maintain_order: bool = False, frame: t.Optional[CoordinateFrame] = None,
             **named_by: IntoExpr) -> polars.dataframe.group_by.GroupBy:
    """
    Start a group by operation. See [`DataFrame.group_by`][polars.DataFrame.group_by] for more information.
    """
    ...

pipe

pipe(
    function: Callable[Concatenate[HasAtomCellT, P], T],
    *args: args,
    **kwargs: kwargs
) -> T

Apply function to self (in method-call syntax).

Source code in atomlib/atomcell.py

def pipe(self: HasAtomCellT, function: t.Callable[Concatenate[HasAtomCellT, P], T], *args: P.args, **kwargs: P.kwargs) -> T:
    """Apply `function` to `self` (in method-call syntax)."""
    return function(self, *args, **kwargs)

drop

drop(
    *columns: Union[str, Iterable[str]], strict: bool = True
) -> DataFrame

Return self with the specified columns removed.

Source code in atomlib/atoms.py

def drop(self, *columns: t.Union[str, t.Iterable[str]], strict: bool = True) -> polars.DataFrame:
    """Return `self` with the specified columns removed."""
    return self._get_frame().drop(*columns, strict=strict)

filter

filter(
    *predicates: Union[
        None,
        IntoExprColumn,
        Iterable[IntoExprColumn],
        bool,
        List[bool],
        ndarray,
    ],
    frame: Optional[CoordinateFrame] = None,
    **constraints: Any
) -> Self

Filter self, removing rows which evaluate to False.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def filter(
    self,
    *predicates: t.Union[None, IntoExprColumn, t.Iterable[IntoExprColumn], bool, t.List[bool], numpy.ndarray],
    frame: t.Optional[CoordinateFrame] = None,
    **constraints: t.Any,
) -> Self:
    """Filter `self`, removing rows which evaluate to `False`."""
    ...

sort

sort(
    by: Union[IntoExpr, Iterable[IntoExpr]],
    *more_by: IntoExpr,
    descending: Union[bool, Sequence[bool]] = False,
    nulls_last: bool = False
) -> Self

Sort the atoms in self by the given columns/expressions.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def sort(
    self,
    by: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    *more_by: IntoExpr,
    descending: t.Union[bool, t.Sequence[bool]] = False,
    nulls_last: bool = False,
) -> Self:
    """
    Sort the atoms in `self` by the given columns/expressions.
    """
    ...

slice

slice(
    offset: int,
    length: Optional[int] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return a slice of the rows in self.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def slice(self, offset: int, length: t.Optional[int] = None, *,
          frame: t.Optional[CoordinateFrame] = None) -> Self:
    """Return a slice of the rows in `self`."""
    ...

head

head(
    n: int = 5, *, frame: Optional[CoordinateFrame] = None
) -> Self

Return the first n rows of self.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def head(self, n: int = 5, *, frame: t.Optional[CoordinateFrame] = None) -> Self:
    """Return the first `n` rows of `self`."""
    ...

tail

tail(
    n: int = 5, *, frame: Optional[CoordinateFrame] = None
) -> Self

Return the last n rows of self.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def tail(self, n: int = 5, *, frame: t.Optional[CoordinateFrame] = None) -> Self:
    """Return the last `n` rows of `self`."""
    ...

drop_nulls

drop_nulls(
    subset: Union[str, Collection[str], None] = None
) -> DataFrame

Drop rows that contain nulls in any of columns subset.

Source code in atomlib/atoms.py

@_fwd_frame_map
def drop_nulls(self, subset: t.Union[str, t.Collection[str], None] = None) -> polars.DataFrame:
    """Drop rows that contain nulls in any of columns `subset`."""
    return self._get_frame().drop_nulls(subset)

fill_null

fill_null(
    value: Any = None,
    strategy: Optional[FillNullStrategy] = None,
    limit: Optional[int] = None,
    matches_supertype: bool = True,
) -> Self

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def fill_null(
    self, value: t.Any = None, strategy: t.Optional[FillNullStrategy] = None,
    limit: t.Optional[int] = None, matches_supertype: bool = True,
) -> Self:
    ...

fill_nan

fill_nan(
    value: Union[Expr, int, float, None],
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def fill_nan(self, value: t.Union[polars.Expr, int, float, None], *,
             frame: t.Optional[CoordinateFrame] = None) -> Self:
    ...

concat classmethod

concat(
    atoms: Union[
        HasAtomsT,
        IntoAtoms,
        Iterable[Union[HasAtomsT, IntoAtoms]],
    ],
    *,
    rechunk: bool = True,
    how: ConcatMethod = "vertical"
) -> HasAtomsT

Concatenate multiple Atoms together, handling metadata appropriately.

Source code in atomlib/atoms.py

@classmethod
def concat(cls: t.Type[HasAtomsT],
           atoms: t.Union[HasAtomsT, IntoAtoms, t.Iterable[t.Union[HasAtomsT, IntoAtoms]]], *,
           rechunk: bool = True, how: ConcatMethod = 'vertical') -> HasAtomsT:
    """Concatenate multiple `Atoms` together, handling metadata appropriately."""
    # this method is tricky. It needs to accept raw Atoms, as well as HasAtoms of the
    # same type as ``cls``.
    if _is_abstract(cls):
        raise TypeError("concat() must be called on a concrete class.")

    if isinstance(atoms, HasAtoms):
        atoms = (atoms,)
    dfs = [a.get_atoms('local').inner if isinstance(a, HasAtoms) else Atoms(t.cast(IntoAtoms, a)).inner for a in atoms]
    representative = cls._combine_metadata(*(a for a in atoms if isinstance(a, HasAtoms)))

    if len(dfs) == 0:
        return representative.with_atoms(Atoms.empty(), 'local')

    if how in ('vertical', 'vertical_relaxed'):
        # get order from first member
        cols = dfs[0].columns
        dfs = [df.select(cols) for df in dfs]
    elif how == 'inner':
        cols = reduce(operator.and_, (df.schema.keys() for df in dfs))
        schema = OrderedDict((col, dfs[0].schema[col]) for col in cols)
        if len(schema) == 0:
            raise ValueError("Atoms have no columns in common")

        dfs = [_select_schema(df, schema) for df in dfs]
        how = 'vertical'

    return representative.with_atoms(Atoms(polars.concat(dfs, rechunk=rechunk, how=how)), 'local')

partition_by

partition_by(
    by: Union[str, Sequence[str]],
    *more_by: str,
    maintain_order: bool = True,
    include_key: bool = True,
    as_dict: Literal[False] = False
) -> List[Self]

partition_by(
    by: Union[str, Sequence[str]],
    *more_by: str,
    maintain_order: bool = True,
    include_key: bool = True,
    as_dict: Literal[True] = ...
) -> Dict[Any, Self]

partition_by(
    by: Union[str, Sequence[str]],
    *more_by: str,
    maintain_order: bool = True,
    include_key: bool = True,
    as_dict: bool = False
) -> Union[List[Self], Dict[Any, Self]]

Group by the given columns and partition into separate dataframes.

Return the partitions as a dictionary by specifying as_dict=True.

Source code in atomlib/atoms.py

def partition_by(
    self, by: t.Union[str, t.Sequence[str]], *more_by: str,
    maintain_order: bool = True, include_key: bool = True, as_dict: bool = False
) -> t.Union[t.List[Self], t.Dict[t.Any, Self]]:
    """
    Group by the given columns and partition into separate dataframes.

    Return the partitions as a dictionary by specifying `as_dict=True`.
    """
    if as_dict:
        d = self._get_frame().partition_by(by, *more_by, maintain_order=maintain_order, include_key=include_key, as_dict=True)
        return {k: self.with_atoms(Atoms(df, _unchecked=True)) for (k, df) in d.items()}

    return [
        self.with_atoms(Atoms(df, _unchecked=True))
        for df in self._get_frame().partition_by(by, *more_by, maintain_order=maintain_order, include_key=include_key, as_dict=False)
    ]

select

select(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> DataFrame

Select exprs from self, and return as a polars.DataFrame.

Expressions may either be columns or expressions of columns.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def select(
    self, *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    frame: t.Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> polars.DataFrame:
    """
    Select `exprs` from `self`, and return as a [`polars.DataFrame`][polars.DataFrame].

    Expressions may either be columns or expressions of columns.

    [polars.DataFrame]: https://docs.pola.rs/py-polars/html/reference/dataframe/index.html
    """
    ...

select_schema

select_schema(schema: SchemaDict) -> DataFrame

Select columns from self and cast to the given schema. Raises TypeError if a column is not found or if it can't be cast.

Source code in atomlib/atoms.py

def select_schema(self, schema: SchemaDict) -> polars.DataFrame:
    """
    Select columns from `self` and cast to the given schema.
    Raises [`TypeError`][TypeError] if a column is not found or if it can't be cast.
    """
    return _select_schema(self, schema)

select_props

select_props(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Self

Select exprs from self, while keeping required columns. Doesn't affect the cell.

RETURNS	DESCRIPTION
Self	A HasAtomCell filtered to contain
Self	the specified properties (as well as required columns).

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def select_props(
    self,
    *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    frame: t.Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Self:
    """
    Select `exprs` from `self`, while keeping required columns.
    Doesn't affect the cell.

    Returns:
      A [`HasAtomCell`][atomlib.atomcell.HasAtomCell] filtered to contain
      the specified properties (as well as required columns).
    """
    ...

try_select

try_select(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Optional[DataFrame]

Try to select exprs from self, and return as a polars.DataFrame.

Expressions may either be columns or expressions of columns. Returns None if any columns are missing.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def try_select(
    self, *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    frame: t.Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> t.Optional[polars.DataFrame]:
    """
    Try to select `exprs` from `self`, and return as a [`polars.DataFrame`][polars.DataFrame].

    Expressions may either be columns or expressions of columns. Returns `None` if any
    columns are missing.

    [polars.DataFrame]: https://docs.pola.rs/py-polars/html/reference/dataframe/index.html
    """
    ...

try_get_column

try_get_column(name: str) -> Optional[Series]

Try to get a column from self, returning None if it doesn't exist.

Source code in atomlib/atoms.py

def try_get_column(self, name: str) -> t.Optional[polars.Series]:
    """Try to get a column from `self`, returning `None` if it doesn't exist."""
    try:
        return self.get_column(name)
    except polars.exceptions.ColumnNotFoundError:
        return None

bbox_atoms

bbox_atoms(
    frame: Optional[CoordinateFrame] = None,
) -> BBox3D

Return the bounding box of all the atoms in self, in the given coordinate frame.

Source code in atomlib/atomcell.py

def bbox_atoms(self, frame: t.Optional[CoordinateFrame] = None) -> BBox3D:
    """Return the bounding box of all the atoms in `self`, in the given coordinate frame."""
    return self.get_atoms(frame).bbox()

transform_atoms

transform_atoms(
    transform: IntoTransform3D,
    selection: Optional[AtomSelection] = None,
    *,
    frame: CoordinateFrame = "local",
    transform_velocities: bool = False
) -> Self

Transform the atoms in self by transform. If selection is given, only transform the atoms in selection.

Source code in atomlib/atomcell.py

def transform_atoms(self, transform: IntoTransform3D, selection: t.Optional[AtomSelection] = None, *,
                    frame: CoordinateFrame = 'local', transform_velocities: bool = False) -> Self:
    """
    Transform the atoms in `self` by `transform`.
    If `selection` is given, only transform the atoms in `selection`.
    """
    transform = self.change_transform(Transform3D.make(transform), self.get_frame(), frame)
    return self.with_atoms(self.get_atoms(self.get_frame()).transform(transform, selection, transform_velocities=transform_velocities))

transform

transform(
    transform: AffineTransform3D,
    frame: CoordinateFrame = "local",
) -> Self

Source code in atomlib/atomcell.py

def transform(self, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> Self:
    if isinstance(transform, Transform3D) and not isinstance(transform, AffineTransform3D):
        raise ValueError("Non-affine transforms cannot change the box dimensions. Use 'transform_atoms' instead.")
    # TODO: cleanup once tests pass
    # coordinate change the transform into atomic coordinates
    new_cell = self.get_cell().transform_cell(transform, frame)
    transform = self.get_cell().change_transform(transform, self.get_frame(), frame)
    return self.with_atoms(self.get_atoms().transform(transform), self.get_frame()).with_cell(new_cell)

round_near_zero

round_near_zero(
    tol: float = 1e-14,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Round atom position values near zero to zero.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def round_near_zero(self, tol: float = 1e-14, *,
                    frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Round atom position values near zero to zero.
    """
    ...

crop_atoms

crop_atoms(
    x_min: float = -inf,
    x_max: float = inf,
    y_min: float = -inf,
    y_max: float = inf,
    z_min: float = -inf,
    z_max: float = inf,
    *,
    frame: CoordinateFrame = "local"
) -> Self

Source code in atomlib/atomcell.py

def crop_atoms(self, x_min: float = -numpy.inf, x_max: float = numpy.inf,
               y_min: float = -numpy.inf, y_max: float = numpy.inf,
               z_min: float = -numpy.inf, z_max: float = numpy.inf, *,
               frame: CoordinateFrame = 'local') -> Self:
    atoms = self._transform_atoms_in_frame(frame, lambda atoms: atoms.crop_atoms(x_min, x_max, y_min, y_max, z_min, z_max))
    return self.with_atoms(atoms)

deduplicate

deduplicate(
    tol: float = 0.001,
    subset: Iterable[str] = ("x", "y", "z", "symbol"),
    keep: UniqueKeepStrategy = "first",
    maintain_order: bool = True,
) -> Self

De-duplicate atoms in self. Atoms of the same symbol that are closer than tolerance to each other (by Euclidian distance) will be removed, leaving only the atom specified by keep (defaults to the first atom).

If subset is specified, only those columns will be included while assessing duplicates. Floating point columns other than 'x', 'y', and 'z' will not by toleranced.

Source code in atomlib/atoms.py

def deduplicate(self, tol: float = 1e-3, subset: t.Iterable[str] = ('x', 'y', 'z', 'symbol'),
                keep: UniqueKeepStrategy = 'first', maintain_order: bool = True) -> Self:
    """
    De-duplicate atoms in `self`. Atoms of the same `symbol` that are closer than `tolerance`
    to each other (by Euclidian distance) will be removed, leaving only the atom specified by
    `keep` (defaults to the first atom).

    If `subset` is specified, only those columns will be included while assessing duplicates.
    Floating point columns other than 'x', 'y', and 'z' will not by toleranced.
    """
    import scipy.spatial

    cols = set((subset,) if isinstance(subset, str) else subset)

    indices = numpy.arange(len(self))

    spatial_cols = cols.intersection(('x', 'y', 'z'))
    cols -= spatial_cols
    if len(spatial_cols) > 0:
        coords = self.select([_coord_expr(col).alias(col) for col in spatial_cols]).to_numpy()
        tree = scipy.spatial.KDTree(coords)

        # TODO This is a bad algorithm
        while True:
            changed = False
            for (i, j) in tree.query_pairs(tol, 2.):
                # whenever we encounter a pair, ensure their index matches
                i_i, i_j = indices[[i, j]]
                if i_i != i_j:
                    indices[i] = indices[j] = min(i_i, i_j)
                    changed = True
            if not changed:
                break

        self = self.with_column(polars.Series('_unique_pts', indices))
        cols.add('_unique_pts')

    frame = self._get_frame().unique(subset=list(cols), keep=keep, maintain_order=maintain_order)
    if len(spatial_cols) > 0:
        frame = frame.drop('_unique_pts')

    return self.with_atoms(Atoms(frame, _unchecked=True))

with_bounds

with_bounds(
    cell_size: Optional[VecLike] = None,
    cell_origin: Optional[VecLike] = None,
) -> "AtomCell"

Return a periodic cell with the given orthogonal cell dimensions.

If cell_size is not specified, it will be assumed (and may be incorrect).

Source code in atomlib/atoms.py

def with_bounds(self, cell_size: t.Optional[VecLike] = None, cell_origin: t.Optional[VecLike] = None) -> 'AtomCell':
    """
    Return a periodic cell with the given orthogonal cell dimensions.

    If cell_size is not specified, it will be assumed (and may be incorrect).
    """
    # TODO: test this
    from .atomcell import AtomCell

    if cell_size is None:
        warnings.warn("Cell boundary unknown. Defaulting to cell BBox")
        cell_size = self.bbox().size
        cell_origin = self.bbox().min

    # TODO test this origin code
    cell = Cell.from_unit_cell(cell_size)
    if cell_origin is not None:
        cell = cell.transform_cell(AffineTransform3D.translate(to_vec3(cell_origin)))

    return AtomCell(self.get_atoms(), cell, frame='local')

coords

coords(
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> NDArray[float64]

Return a (N, 3) ndarray of atom positions (dtype numpy.float64) in the given coordinate frame.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def coords(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Optional[CoordinateFrame] = None) -> NDArray[numpy.float64]:
    """
    Return a `(N, 3)` ndarray of atom positions (dtype [`numpy.float64`][numpy.float64])
    in the given coordinate frame.
    """
    ...

x

x() -> Expr

Source code in atomlib/atoms.py

def x(self) -> polars.Expr:
    return polars.col('coords').arr.get(0).alias('x')

y

y() -> Expr

Source code in atomlib/atoms.py

def y(self) -> polars.Expr:
    return polars.col('coords').arr.get(1).alias('y')

z

z() -> Expr

Source code in atomlib/atoms.py

def z(self) -> polars.Expr:
    return polars.col('coords').arr.get(2).alias('z')

velocities

velocities(
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Optional[NDArray[float64]]

Return a (N, 3) ndarray of atom velocities (dtype numpy.float64) in the given coordinate frame.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def velocities(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Optional[CoordinateFrame] = None) -> t.Optional[NDArray[numpy.float64]]:
    """
    Return a `(N, 3)` ndarray of atom velocities (dtype [`numpy.float64`][numpy.float64])
    in the given coordinate frame.
    """
    ...

types

types() -> Optional[Series]

Returns a Series of atom types (dtype polars.Int32).

Source code in atomlib/atoms.py

def types(self) -> t.Optional[polars.Series]:
    """
    Returns a [`Series`][polars.Series] of atom types (dtype [`polars.Int32`][polars.datatypes.Int32]).

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    return self.try_get_column('type')

masses

masses() -> Optional[Series]

Returns a Series of atom masses (dtype polars.Float32).

Source code in atomlib/atoms.py

def masses(self) -> t.Optional[polars.Series]:
    """
    Returns a [`Series`][polars.Series] of atom masses (dtype [`polars.Float32`][polars.datatypes.Float32]).

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    return self.try_get_column('mass')

add_atom

add_atom(
    elem: Union[int, str],
    x: ArrayLike,
    /,
    *,
    y: None = None,
    z: None = None,
    frame: Optional[CoordinateFrame] = None,
    **kwargs: Any,
) -> Self

add_atom(
    elem: Union[int, str],
    /,
    x: float,
    y: float,
    z: float,
    *,
    frame: Optional[CoordinateFrame] = None,
    **kwargs: Any,
) -> Self

add_atom(
    elem: Union[int, str],
    /,
    x: Union[ArrayLike, float],
    y: Optional[float] = None,
    z: Optional[float] = None,
    *,
    frame: Optional[CoordinateFrame] = None,
    **kwargs: Any,
) -> Self

Return a copy of self with an extra atom.

By default, all extra columns present in self must be specified as **kwargs.

Try to avoid calling this in a loop (Use concat instead).

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def add_atom(self, elem: t.Union[int, str], /,  # type: ignore (spurious)
             x: t.Union[ArrayLike, float],
             y: t.Optional[float] = None,
             z: t.Optional[float] = None, *,
             frame: t.Optional[CoordinateFrame] = None,
             **kwargs: t.Any) -> Self:
    """
    Return a copy of `self` with an extra atom.

    By default, all extra columns present in `self` must be specified as `**kwargs`.

    Try to avoid calling this in a loop (Use [`concat`][atomlib.atomcell.HasAtomCell.concat] instead).
    """
    ...

pos

pos(
    x: Sequence[Optional[float]],
    /,
    *,
    y: None = None,
    z: None = None,
    tol: float = 1e-06,
    **kwargs: Any,
) -> Expr

pos(
    x: Optional[float] = None,
    y: Optional[float] = None,
    z: Optional[float] = None,
    *,
    tol: float = 1e-06,
    **kwargs: Any
) -> Expr

pos(
    x: Union[Sequence[Optional[float]], float, None] = None,
    y: Optional[float] = None,
    z: Optional[float] = None,
    *,
    tol: float = 1e-06,
    **kwargs: Any
) -> Expr

Select all atoms at a given position.

Formally, returns all atoms within a cube of radius tol centered at (x,y,z), exclusive of the cube's surface.

Additional parameters given as kwargs will be checked as additional parameters (with strict equality).

Source code in atomlib/atoms.py

def pos(self,
        x: t.Union[t.Sequence[t.Optional[float]], float, None] = None,
        y: t.Optional[float] = None, z: t.Optional[float] = None, *,
        tol: float = 1e-6, **kwargs: t.Any) -> polars.Expr:
    """
    Select all atoms at a given position.

    Formally, returns all atoms within a cube of radius ``tol``
    centered at ``(x,y,z)``, exclusive of the cube's surface.

    Additional parameters given as ``kwargs`` will be checked
    as additional parameters (with strict equality).
    """

    if isinstance(x, t.Sequence):
        (x, y, z) = x

    tol = abs(float(tol))
    selection = polars.lit(True)
    if x is not None:
        selection &= self.x().is_between(x - tol, x + tol, closed='none')
    if y is not None:
        selection &= self.y().is_between(y - tol, y + tol, closed='none')
    if z is not None:
        selection &= self.z().is_between(z - tol, z + tol, closed='none')
    for (col, val) in kwargs.items():
        selection &= (polars.col(col) == val)

    return selection

with_index

with_index(
    index: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Returns self with a row index added in column 'i' (dtype polars.Int64). If index is not specified, defaults to an existing index or a new index.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_index(self, index: t.Optional[AtomValues] = None, *,
               frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Returns `self` with a row index added in column 'i' (dtype [`polars.Int64`][polars.datatypes.Int64]).
    If `index` is not specified, defaults to an existing index or a new index.
    """
    ...

with_wobble

with_wobble(
    wobble: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given displacements in column 'wobble' (dtype polars.Float64). If wobble is not specified, defaults to the already-existing wobbles or 0.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_wobble(self, wobble: t.Optional[AtomValues] = None, *,
                frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given displacements in column 'wobble' (dtype [`polars.Float64`][polars.datatypes.Float64]).
    If `wobble` is not specified, defaults to the already-existing wobbles or 0.
    """
    ...

with_occupancy

with_occupancy(
    frac_occupancy: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given fractional occupancies (dtype polars.Float64). If frac_occupancy is not specified, defaults to the already-existing occupancies or 1.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_occupancy(self, frac_occupancy: t.Optional[AtomValues] = None, *,
                   frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return self with the given fractional occupancies (dtype [`polars.Float64`][polars.datatypes.Float64]).
    If `frac_occupancy` is not specified, defaults to the already-existing occupancies or 1.
    """
    ...

apply_wobble

apply_wobble(
    rng: Union[Generator, int, None] = None,
    frame: Optional[CoordinateFrame] = None,
) -> Self

Displace the atoms in self by the amount in the wobble column. wobble is interpretated as a mean-squared displacement, which is distributed equally over each axis.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def apply_wobble(self, rng: t.Union[numpy.random.Generator, int, None] = None,
                 frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Displace the atoms in `self` by the amount in the `wobble` column.
    `wobble` is interpretated as a mean-squared displacement, which is distributed
    equally over each axis.
    """
    ...

apply_occupancy

apply_occupancy(
    rng: Union[Generator, int, None] = None
) -> Self

For each atom in self, use its frac_occupancy to randomly decide whether to remove it.

Source code in atomlib/atoms.py

def apply_occupancy(self, rng: t.Union[numpy.random.Generator, int, None] = None) -> Self:
    """
    For each atom in `self`, use its `frac_occupancy` to randomly decide whether to remove it.
    """
    if 'frac_occupancy' not in self.columns:
        return self
    rng = numpy.random.default_rng(seed=rng)

    frac = self.select('frac_occupancy').to_series().to_numpy()
    choice = rng.binomial(1, frac).astype(numpy.bool_)
    return self.filter(polars.lit(choice))

with_type

with_type(
    types: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given atom types in column 'type'. If types is not specified, use the already existing types or auto-assign them.

When auto-assigning, each symbol is given a unique value, case-sensitive. Values are assigned from lowest atomic number to highest. For instance: ["Ag+", "Na", "H", "Ag"] => [3, 11, 1, 2]

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_type(self, types: t.Optional[AtomValues] = None, *,
              frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given atom types in column 'type'.
    If `types` is not specified, use the already existing types or auto-assign them.

    When auto-assigning, each symbol is given a unique value, case-sensitive.
    Values are assigned from lowest atomic number to highest.
    For instance: `["Ag+", "Na", "H", "Ag"]` => `[3, 11, 1, 2]`
    """
    ...

with_mass

with_mass(
    mass: Optional[ArrayLike] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given atom masses in column 'mass'. If mass is not specified, use the already existing masses or auto-assign them.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_mass(self, mass: t.Optional[ArrayLike] = None, *,
              frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given atom masses in column `'mass'`.
    If `mass` is not specified, use the already existing masses or auto-assign them.
    """
    ...

with_symbol

with_symbol(
    symbols: ArrayLike,
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given atomic symbols.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_symbol(self, symbols: ArrayLike, selection: t.Optional[AtomSelection] = None, *,
                frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given atomic symbols.
    """
    ...

with_coords

with_coords(
    pts: ArrayLike,
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self replaced with the given atomic positions.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_coords(self, pts: ArrayLike, selection: t.Optional[AtomSelection] = None, *,
                frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` replaced with the given atomic positions.
    """
    ...

with_velocity

with_velocity(
    pts: Optional[ArrayLike] = None,
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self replaced with the given atomic velocities. If pts is not specified, use the already existing velocities or zero.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_velocity(self, pts: t.Optional[ArrayLike] = None,
                  selection: t.Optional[AtomSelection] = None, *,
                  frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` replaced with the given atomic velocities.
    If `pts` is not specified, use the already existing velocities or zero.
    """
    ...

get_atomcell

get_atomcell() -> AtomCell

Source code in atomlib/atomcell.py

def get_atomcell(self) -> AtomCell:
    frame = self.get_frame()
    return AtomCell(self.get_atoms(frame), self.get_cell(), frame=frame, keep_frame=True)

to_frame

to_frame(frame: CoordinateFrame) -> Self

Convert the stored Atoms to the given coordinate frame.

Source code in atomlib/atomcell.py

def to_frame(self, frame: CoordinateFrame) -> Self:
    """Convert the stored Atoms to the given coordinate frame."""
    return self.with_atoms(self.get_atoms(frame), frame)

crop_to_box

crop_to_box(eps: float = 1e-05) -> Self

Source code in atomlib/atomcell.py

def crop_to_box(self, eps: float = 1e-5) -> Self:
    atoms = self._transform_atoms_in_frame('cell_box', lambda atoms: atoms.crop_atoms(*([-eps, 1-eps]*3)))
    return self.with_atoms(atoms)

wrap

wrap(eps: float = 1e-05) -> Self

Wrap atoms around the cell boundaries.

Source code in atomlib/atomcell.py

def wrap(self, eps: float = 1e-5) -> Self:
    """Wrap atoms around the cell boundaries."""
    return self.with_atoms(self._transform_atoms_in_frame('cell_box', lambda a: a._wrap(eps)))

repeat_to

repeat_to(
    size: VecLike, crop: Union[bool, Sequence[bool]] = False
) -> Self

Repeat the cell so it is at least size along the crystal's axes.

If crop, then crop the cell to exactly size. This may break periodicity. crop may be a vector, in which case you can specify cropping only along some axes.

Source code in atomlib/atomcell.py

def repeat_to(self, size: VecLike, crop: t.Union[bool, t.Sequence[bool]] = False) -> Self:
    """
    Repeat the cell so it is at least `size` along the crystal's axes.

    If `crop`, then crop the cell to exactly `size`. This may break periodicity.
    `crop` may be a vector, in which case you can specify cropping only along some axes.
    """
    size = to_vec3(size)
    cell_size = self.cell_size * self.n_cells
    repeat = numpy.maximum(numpy.ceil(size / cell_size).astype(int), 1)
    atom_cell = self.repeat(repeat)

    crop_v = to_vec3(crop, dtype=numpy.bool_)
    if numpy.any(crop_v):
        crop_x, crop_y, crop_z = crop_v
        return atom_cell.crop(
            x_max = size[0] if crop_x else numpy.inf,
            y_max = size[1] if crop_y else numpy.inf,
            z_max = size[2] if crop_z else numpy.inf,
            frame='cell'
        )

    return atom_cell

repeat_x

repeat_x(n: int) -> Self

Tile the cell in the x axis.

Source code in atomlib/atomcell.py

def repeat_x(self, n: int) -> Self:
    """Tile the cell in the x axis."""
    return self.repeat((n, 1, 1))

repeat_y

repeat_y(n: int) -> Self

Tile the cell in the y axis.

Source code in atomlib/atomcell.py

def repeat_y(self, n: int) -> Self:
    """Tile the cell in the y axis."""
    return self.repeat((1, n, 1))

repeat_z

repeat_z(n: int) -> Self

Tile the cell in the z axis.

Source code in atomlib/atomcell.py

def repeat_z(self, n: int) -> Self:
    """Tile the cell in the z axis."""
    return self.repeat((1, 1, n))

repeat_to_x

repeat_to_x(size: float, crop: bool = False) -> Self

Repeat the cell so it is at least size size along the x axis.

Source code in atomlib/atomcell.py

def repeat_to_x(self, size: float, crop: bool = False) -> Self:
    """Repeat the cell so it is at least size `size` along the x axis."""
    return self.repeat_to([size, 0., 0.], [crop, False, False])

repeat_to_y

repeat_to_y(size: float, crop: bool = False) -> Self

Repeat the cell so it is at least size size along the y axis.

Source code in atomlib/atomcell.py

def repeat_to_y(self, size: float, crop: bool = False) -> Self:
    """Repeat the cell so it is at least size `size` along the y axis."""
    return self.repeat_to([0., size, 0.], [False, crop, False])

repeat_to_z

repeat_to_z(size: float, crop: bool = False) -> Self

Repeat the cell so it is at least size size along the z axis.

Source code in atomlib/atomcell.py

def repeat_to_z(self, size: float, crop: bool = False) -> Self:
    """Repeat the cell so it is at least size `size` along the z axis."""
    return self.repeat_to([0., 0., size], [False, False, crop])

repeat_to_aspect

repeat_to_aspect(
    plane: Literal["xy", "xz", "yz"] = "xy",
    *,
    aspect: float = 1.0,
    min_size: Optional[VecLike] = None,
    max_size: Optional[VecLike] = None
) -> Self

Repeat to optimize the aspect ratio in plane, while staying above min_size and under max_size.

Source code in atomlib/atomcell.py

def repeat_to_aspect(self, plane: t.Literal['xy', 'xz', 'yz'] = 'xy', *,
                     aspect: float = 1., min_size: t.Optional[VecLike] = None,
                     max_size: t.Optional[VecLike] = None) -> Self:
    """
    Repeat to optimize the aspect ratio in `plane`,
    while staying above `min_size` and under `max_size`.
    """
    if min_size is None:
        min_n = numpy.array([1, 1, 1], numpy.int_)
    else:
        min_n = numpy.maximum(numpy.ceil(to_vec3(min_size) / self.box_size), 1).astype(numpy.int_)

    if max_size is None:
        max_n = 3 * min_n
    else:
        max_n = numpy.maximum(numpy.floor(to_vec3(max_size) / self.box_size), 1).astype(numpy.int_)

    if plane == 'xy':
        indices = [0, 1]
    elif plane == 'xz':
        indices = [0, 2]
    elif plane == 'yz':
        indices = [1, 2]
    else:
        raise ValueError(f"Invalid plane '{plane}'. Exepcted 'xy', 'xz', 'or 'yz'.")

    na = numpy.arange(min_n[indices[0]], max_n[indices[0]])
    nb = numpy.arange(min_n[indices[1]], max_n[indices[1]])
    (na, nb) = numpy.meshgrid(na, nb)

    aspects = na * self.box_size[indices[0]] / (nb * self.box_size[indices[1]])
    # cost function: log(aspect)^2  (so cost(0.5) == cost(2))
    min_i = numpy.argmin(numpy.log(aspects / aspect)**2)
    repeat = numpy.array([1, 1, 1], numpy.int_)
    repeat[indices] = na.flatten()[min_i], nb.flatten()[min_i]
    return self.repeat(repeat)

periodic_duplicate

periodic_duplicate(eps: float = 1e-05) -> Self

Add duplicate copies of atoms near periodic boundaries.

For instance, an atom at a corner will be duplicated into 8 copies. This is mostly only useful for visualization.

Source code in atomlib/atomcell.py

def periodic_duplicate(self, eps: float = 1e-5) -> Self:
    """
    Add duplicate copies of atoms near periodic boundaries.

    For instance, an atom at a corner will be duplicated into 8 copies.
    This is mostly only useful for visualization.
    """
    frame_save = self.get_frame()
    self = self.to_frame('cell_box').wrap(eps=eps)

    for i in range(3):
        self = self.concat((self,
            self.filter(polars.col('coords').arr.get(i).abs() <= eps, frame='cell_box')
                .transform_atoms(AffineTransform3D.translate([1. if i == j else 0. for j in range(3)]), frame='cell_box')
        ))

    return self.to_frame(frame_save)

read classmethod

read(path: FileOrPath, ty: FileType) -> HasAtomsT

read(
    path: Union[str, Path, TextIO], ty: Literal[None] = None
) -> HasAtomsT

read(
    path: FileOrPath, ty: Optional[FileType] = None
) -> HasAtomsT

Read a structure from a file.

Supported types can be found in the io module. If no ty is specified, it is inferred from the file's extension.

Source code in atomlib/mixins.py

@classmethod
def read(cls: t.Type[HasAtomsT], path: FileOrPath, ty: t.Optional[FileType] = None) -> HasAtomsT:
    """
    Read a structure from a file.

    Supported types can be found in the [io][atomlib.io] module.
    If no `ty` is specified, it is inferred from the file's extension.
    """
    from .io import read
    return _cast_atoms(read(path, ty), cls)  # type: ignore

read_cif classmethod

read_cif(
    f: Union[FileOrPath, CIF, CIFDataBlock],
    block: Union[int, str, None] = None,
) -> HasAtomsT

Read a structure from a CIF file.

If block is specified, read data from the given block of the CIF file (index or name).

Source code in atomlib/mixins.py

@classmethod
def read_cif(cls: t.Type[HasAtomsT], f: t.Union[FileOrPath, CIF, CIFDataBlock], block: t.Union[int, str, None] = None) -> HasAtomsT:
    """
    Read a structure from a CIF file.

    If `block` is specified, read data from the given block of the CIF file (index or name).
    """
    from .io import read_cif
    return _cast_atoms(read_cif(f, block), cls)

read_xyz classmethod

read_xyz(f: Union[FileOrPath, XYZ]) -> HasAtomsT

Read a structure from an XYZ file.

Source code in atomlib/mixins.py

@classmethod
def read_xyz(cls: t.Type[HasAtomsT], f: t.Union[FileOrPath, XYZ]) -> HasAtomsT:
    """Read a structure from an XYZ file."""
    from .io import read_xyz
    return _cast_atoms(read_xyz(f), cls)

read_xsf classmethod

read_xsf(f: Union[FileOrPath, XSF]) -> HasAtomsT

Read a structure from an XSF file.

Source code in atomlib/mixins.py

@classmethod
def read_xsf(cls: t.Type[HasAtomsT], f: t.Union[FileOrPath, XSF]) -> HasAtomsT:
    """Read a structure from an XSF file."""
    from .io import read_xsf
    return _cast_atoms(read_xsf(f), cls)

read_cfg classmethod

read_cfg(f: Union[FileOrPath, CFG]) -> HasAtomsT

Read a structure from a CFG file.

Source code in atomlib/mixins.py

@classmethod
def read_cfg(cls: t.Type[HasAtomsT], f: t.Union[FileOrPath, CFG]) -> HasAtomsT:
    """Read a structure from a CFG file."""
    from .io import read_cfg
    return _cast_atoms(read_cfg(f), cls)

read_lmp classmethod

read_lmp(
    f: Union[FileOrPath, LMP],
    type_map: Optional[Dict[int, Union[str, int]]] = None,
) -> HasAtomsT

Read a structure from a LAAMPS data file.

Source code in atomlib/mixins.py

@classmethod
def read_lmp(cls: t.Type[HasAtomsT], f: t.Union[FileOrPath, LMP], type_map: t.Optional[t.Dict[int, t.Union[str, int]]] = None) -> HasAtomsT:
    """Read a structure from a LAAMPS data file."""
    from .io import read_lmp
    return _cast_atoms(read_lmp(f, type_map=type_map), cls)

write_cif

write_cif(f: FileOrPath)

Source code in atomlib/mixins.py

def write_cif(self, f: FileOrPath):
    from .io import write_cif
    write_cif(self, f)

write_xyz

write_xyz(f: FileOrPath, fmt: XYZFormat = 'exyz')

Source code in atomlib/mixins.py

def write_xyz(self, f: FileOrPath, fmt: XYZFormat = 'exyz'):
    from .io import write_xyz
    write_xyz(self, f, fmt)

write_xsf

write_xsf(f: FileOrPath)

Source code in atomlib/mixins.py

def write_xsf(self, f: FileOrPath):
    from .io import write_xsf
    write_xsf(self, f)

write_cfg

write_cfg(f: FileOrPath)

Source code in atomlib/mixins.py

def write_cfg(self, f: FileOrPath):
    from .io import write_cfg
    write_cfg(self, f)

write_lmp

write_lmp(f: FileOrPath)

Source code in atomlib/mixins.py

def write_lmp(self, f: FileOrPath):
    from .io import write_lmp
    write_lmp(self, f)

write

write(path: FileOrPath, ty: FileType)

write(
    path: Union[str, Path, TextIO], ty: Literal[None] = None
)

write(path: FileOrPath, ty: Optional[FileType] = None)

Write this structure to a file.

A file type may be specified using ty. If no ty is specified, it is inferred from the path's extension.

Source code in atomlib/mixins.py

def write(self, path: FileOrPath, ty: t.Optional[FileType] = None):
    """
    Write this structure to a file.

    A file type may be specified using `ty`.
    If no `ty` is specified, it is inferred from the path's extension.
    """
    from .io import write
    write(self, path, ty)  # type: ignore

write_mslice

write_mslice(
    f: BinaryFileOrPath,
    template: Optional[MSliceFile] = None,
    *,
    slice_thickness: Optional[float] = None,
    scan_points: Optional[ArrayLike] = None,
    scan_extent: Optional[ArrayLike] = None,
    noise_sigma: Optional[float] = None,
    conv_angle: Optional[float] = None,
    energy: Optional[float] = None,
    defocus: Optional[float] = None,
    tilt: Optional[Tuple[float, float]] = None,
    tds: Optional[bool] = None,
    n_cells: Optional[ArrayLike] = None
)

Write a structure to an mslice file.

template may be a file, path, or ElementTree containing an existing mslice file. Its structure will be modified to make the final output. If not specified, a default template will be used.

Additional options modify simulation properties. If an option is not specified, the template's properties are used.

Source code in atomlib/mixins.py

def write_mslice(self, f: BinaryFileOrPath, template: t.Optional[MSliceFile] = None, *,
             slice_thickness: t.Optional[float] = None,  # angstrom
             scan_points: t.Optional[ArrayLike] = None,
             scan_extent: t.Optional[ArrayLike] = None,
             noise_sigma: t.Optional[float] = None,  # angstrom
             conv_angle: t.Optional[float] = None,  # mrad
             energy: t.Optional[float] = None,  # keV
             defocus: t.Optional[float] = None,  # angstrom
             tilt: t.Optional[t.Tuple[float, float]] = None,  # (mrad, mrad)
             tds: t.Optional[bool] = None,
             n_cells: t.Optional[ArrayLike] = None):
    """
    Write a structure to an mslice file.

    `template` may be a file, path, or `ElementTree` containing an existing mslice file.
    Its structure will be modified to make the final output. If not specified, a default
    template will be used.

    Additional options modify simulation properties. If an option is not specified, the
    template's properties are used.
    """
    from .io import write_mslice
    return write_mslice(self, f, template, slice_thickness=slice_thickness,
                        scan_points=scan_points, scan_extent=scan_extent,
                        conv_angle=conv_angle, energy=energy, defocus=defocus,
                        noise_sigma=noise_sigma, tilt=tilt, tds=tds, n_cells=n_cells)

write_qe

write_qe(
    f: FileOrPath,
    pseudo: Optional[Mapping[str, str]] = None,
)

Write a structure to a Quantum Espresso pw.x file.

PARAMETER	DESCRIPTION
f	File or path to write to TYPE: FileOrPath
pseudo	Mapping from atom symbol TYPE: Optional[Mapping[str, str]] DEFAULT: None

Source code in atomlib/mixins.py

def write_qe(self, f: FileOrPath, pseudo: t.Optional[t.Mapping[str, str]] = None):
    """
    Write a structure to a Quantum Espresso pw.x file.

    Args:
      f: File or path to write to
      pseudo: Mapping from atom symbol
    """
    from .io import write_qe
    write_qe(self, f, pseudo)

get_cell

get_cell() -> Cell

Source code in atomlib/atomcell.py

def get_cell(self) -> Cell:
    return self.cell

with_cell

with_cell(cell: Cell) -> Self

Source code in atomlib/atomcell.py

def with_cell(self, cell: Cell) -> Self:
    return self.__class__(self.atoms, cell, frame=self.frame, keep_frame=True)

get_atoms

get_atoms(frame: Optional[CoordinateFrame] = None) -> Atoms

Get atoms contained in self, in the given coordinate frame.

Source code in atomlib/atomcell.py

def get_atoms(self, frame: t.Optional[CoordinateFrame] = None) -> Atoms:
    """Get atoms contained in ``self``, in the given coordinate frame."""

    if frame is None or frame == self.get_frame():
        return self.atoms
    return self.atoms.transform(self.get_transform(frame, self.get_frame()))

with_atoms

with_atoms(
    atoms: HasAtoms, frame: Optional[CoordinateFrame] = None
) -> Self

Source code in atomlib/atomcell.py

def with_atoms(self, atoms: HasAtoms, frame: t.Optional[CoordinateFrame] = None) -> Self:
    frame = frame if frame is not None else self.frame
    return self.__class__(atoms.get_atoms(), cell=self.cell, frame=frame, keep_frame=True)

get_frame

get_frame() -> CoordinateFrame

Get the coordinate frame atoms are stored in.

Source code in atomlib/atomcell.py

def get_frame(self) -> CoordinateFrame:
    """Get the coordinate frame atoms are stored in."""
    return self.frame

from_ortho classmethod

from_ortho(
    atoms: IntoAtoms,
    ortho: LinearTransform3D,
    *,
    n_cells: Optional[VecLike] = None,
    frame: CoordinateFrame = "local",
    keep_frame: bool = False
)

Make an atom cell given a list of atoms and an orthogonalization matrix. Atoms are assumed to be in the coordinate system frame.

Source code in atomlib/atomcell.py

@classmethod
def from_ortho(cls, atoms: IntoAtoms, ortho: LinearTransform3D, *,
               n_cells: t.Optional[VecLike] = None,
               frame: CoordinateFrame = 'local',
               keep_frame: bool = False):
    """
    Make an atom cell given a list of atoms and an orthogonalization matrix.
    Atoms are assumed to be in the coordinate system `frame`.
    """
    cell = Cell.from_ortho(ortho, n_cells)
    return cls(atoms, cell, frame=frame, keep_frame=keep_frame)

from_unit_cell classmethod

from_unit_cell(
    atoms: IntoAtoms,
    cell_size: VecLike,
    cell_angle: Optional[VecLike] = None,
    *,
    n_cells: Optional[VecLike] = None,
    frame: CoordinateFrame = "local",
    keep_frame: bool = False
)

Make a cell given a list of atoms and unit cell parameters. Atoms are assumed to be in the coordinate system frame.

Source code in atomlib/atomcell.py

@classmethod
def from_unit_cell(cls, atoms: IntoAtoms, cell_size: VecLike,
                   cell_angle: t.Optional[VecLike] = None, *,
                   n_cells: t.Optional[VecLike] = None,
                   frame: CoordinateFrame = 'local',
                   keep_frame: bool = False):
    """
    Make a cell given a list of atoms and unit cell parameters.
    Atoms are assumed to be in the coordinate system `frame`.
    """
    cell = Cell.from_unit_cell(cell_size, cell_angle, n_cells=n_cells)
    return cls(atoms, cell, frame=frame, keep_frame=keep_frame)

orthogonalize

orthogonalize() -> OrthoCell

Source code in atomlib/atomcell.py

def orthogonalize(self) -> OrthoCell:
    if self.is_orthogonal():
        return OrthoCell(self.atoms, self.cell, frame=self.frame)
    raise NotImplementedError()

clone

clone() -> AtomCellT

Make a deep copy of self.

Source code in atomlib/atomcell.py

def clone(self: AtomCellT) -> AtomCellT:
    """Make a deep copy of `self`."""
    return self.__class__(**{field.name: copy.deepcopy(getattr(self, field.name)) for field in fields(self)})

assert_equal

assert_equal(other: Any)

Assert this structure is equal to other

Source code in atomlib/atomcell.py

def assert_equal(self, other: t.Any):
    """Assert this structure is equal to `other`"""
    assert isinstance(other, AtomCell)
    self.cell.assert_equal(other.cell)
    self.get_atoms('local').assert_equal(other.get_atoms('local'))

OrthoCell dataclass

Bases: AtomCell

Source code in atomlib/atomcell.py

class OrthoCell(AtomCell):
    def __post_init__(self):
        if not numpy.allclose(self.cell.cell_angle, numpy.pi/2.):
            raise ValueError(f"OrthoCell constructed with non-orthogonal angles: {self.cell.cell_angle}")

    def is_orthogonal(self, tol: float = 1e-8) -> t.Literal[True]:
        """Returns whether this cell is orthogonal (axes are at right angles.)"""
        return True

affine property

affine: AffineTransform3D

Affine transformation. Holds transformation from 'ortho' to 'local' coordinates, including rotation away from the standard crystal orientation.

ortho property

ortho: LinearTransform3D

Orthogonalization transformation. Skews but does not scale the crystal axes to cartesian axes.

metric property

metric: LinearTransform3D

Cell metric tensor

Returns the dot product between every combination of basis vectors. :math:\mathbf{a} \cdot \mathbf{b} = a_i M_ij b_j

cell_size property

cell_size: NDArray[float64]

Unit cell size.

cell_angle property

cell_angle: NDArray[float64]

Unit cell angles, in radians.

n_cells property

n_cells: NDArray[int_]

Number of unit cells.

pbc property

pbc: NDArray[bool_]

Flags indicating the presence of periodic boundary conditions along each axis.

ortho_size property

ortho_size: NDArray[float64]

Return size of orthogonal unit cell.

Equivalent to the diagonal of the orthogonalization matrix.

box_size property

box_size: NDArray[float64]

Return size of the cell box.

Equivalent to self.n_cells * self.cell_size.

columns property

columns: List[str]

Return the column names in self.

RETURNS	DESCRIPTION
List[str]	A sequence of column names

dtypes property

dtypes: List[DataType]

Return the datatypes in self.

RETURNS	DESCRIPTION
List[DataType]	A sequence of column DataTypes

schema property

schema: Schema

Return the schema of self.

RETURNS	DESCRIPTION
Schema	A dictionary of column names and DataTypes

with_column class-attribute instance-attribute

with_column = with_columns

unique class-attribute instance-attribute

unique = deduplicate

atoms instance-attribute

atoms: Atoms

Atoms in the cell. Stored in 'local' coordinates (i.e. relative to the enclosing group but not relative to box dimensions).

cell instance-attribute

cell: Cell

Cell coordinate system.

frame class-attribute instance-attribute

frame: CoordinateFrame = 'local'

Coordinate frame 'atoms' are stored in.

get_cell

get_cell() -> Cell

Source code in atomlib/atomcell.py

def get_cell(self) -> Cell:
    return self.cell

with_cell

with_cell(cell: Cell) -> Self

Source code in atomlib/atomcell.py

def with_cell(self, cell: Cell) -> Self:
    return self.__class__(self.atoms, cell, frame=self.frame, keep_frame=True)

get_transform

get_transform(
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> AffineTransform3D

In the two-argument form, get the transform to frame_to from frame_from. In the one-argument form, get the transform from local coordinates to 'frame'.

Source code in atomlib/cell.py

def get_transform(self, frame_to: t.Optional[CoordinateFrame] = None, frame_from: t.Optional[CoordinateFrame] = None) -> AffineTransform3D:
    """
    In the two-argument form, get the transform to `frame_to` from `frame_from`.
    In the one-argument form, get the transform from local coordinates to 'frame'.
    """
    transform_from = self._get_transform_to_local(frame_from) if frame_from is not None else AffineTransform3D()
    transform_to = self._get_transform_to_local(frame_to) if frame_to is not None else AffineTransform3D()
    if frame_from is not None and frame_to is not None and frame_from.lower() == frame_to.lower():
        return AffineTransform3D()
    return transform_to.inverse() @ transform_from

corners

corners(frame: CoordinateFrame = 'local') -> ndarray

Source code in atomlib/cell.py

def corners(self, frame: CoordinateFrame = 'local') -> numpy.ndarray:
    corners = numpy.array(list(itertools.product((0., 1.), repeat=3)))
    return self.get_transform(frame, 'cell_box') @ corners

bbox_cell

bbox_cell(frame: CoordinateFrame = 'local') -> BBox3D

Return the bounding box of the cell box in the given coordinate system.

Source code in atomlib/cell.py

def bbox_cell(self, frame: CoordinateFrame = 'local') -> BBox3D:
    """Return the bounding box of the cell box in the given coordinate system."""
    return BBox3D.from_pts(self.corners(frame))

bbox

bbox(frame: CoordinateFrame = 'local') -> BBox3D

Return the combined bounding box of the cell and atoms in the given coordinate system. To get the cell or atoms bounding box only, use bbox_cell or bbox_atoms.

Source code in atomlib/atomcell.py

def bbox(self, frame: CoordinateFrame = 'local') -> BBox3D:
    """
    Return the combined bounding box of the cell and atoms in the given coordinate system.
    To get the cell or atoms bounding box only, use [`bbox_cell`][atomlib.atomcell.HasAtomCell.bbox_cell] or [`bbox_atoms`][atomlib.atomcell.HasAtomCell.bbox_atoms].
    """
    return self.bbox_atoms(frame) | self.bbox_cell(frame)

is_orthogonal_in_local

is_orthogonal_in_local(tol: float = 1e-08) -> bool

Returns whether this cell is orthogonal and aligned with the local coordinate system.

Source code in atomlib/cell.py

def is_orthogonal_in_local(self, tol: float = 1e-8) -> bool:
    """Returns whether this cell is orthogonal and aligned with the local coordinate system."""
    transform = (self.affine @ self.ortho).to_linear()
    if not transform.is_scaled_orthogonal(tol):
        return False
    normed = transform.inner / numpy.linalg.norm(transform.inner, axis=-2, keepdims=True)
    # every row of transform must be a +/- 1 times a basis vector (i, j, or k)
    return all(
        any(numpy.isclose(numpy.abs(numpy.dot(row, v)), 1., atol=tol) for v in numpy.eye(3))
        for row in normed
    )

to_ortho

to_ortho() -> AffineTransform3D

Source code in atomlib/cell.py

def to_ortho(self) -> AffineTransform3D:
    return self.get_transform('local', 'cell_box')

transform_cell

transform_cell(
    transform: AffineTransform3D,
    frame: CoordinateFrame = "local",
) -> Self

Apply the given transform to the unit cell, without changing atom positions. The transform is applied in coordinate frame 'frame'.

Source code in atomlib/atomcell.py

def transform_cell(self, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> Self:
    """
    Apply the given transform to the unit cell, without changing atom positions.
    The transform is applied in coordinate frame 'frame'.
    """
    return self.with_cell(self.get_cell().transform_cell(transform, frame=frame))

strain_orthogonal

strain_orthogonal() -> HasCellT

Orthogonalize using strain.

Strain is applied such that the x-axis remains fixed, and the y-axis remains in the xy plane. For small displacements, no hydrostatic strain is applied (volume is conserved).

Source code in atomlib/cell.py

def strain_orthogonal(self: HasCellT) -> HasCellT:
    """
    Orthogonalize using strain.

    Strain is applied such that the x-axis remains fixed, and the y-axis remains in the xy plane.
    For small displacements, no hydrostatic strain is applied (volume is conserved).
    """
    return self.with_cell(Cell(
        affine=self.affine,
        ortho=LinearTransform3D(),
        cell_size=self.cell_size,
        n_cells=self.n_cells,
        pbc=self.pbc,
    ))

repeat

repeat(n: Union[int, VecLike]) -> Self

Tile the cell

Source code in atomlib/atomcell.py

def repeat(self, n: t.Union[int, VecLike]) -> Self:
    """Tile the cell"""
    ns = numpy.broadcast_to(n, 3)
    if not numpy.issubdtype(ns.dtype, numpy.integer):
        raise ValueError("repeat() argument must be an integer or integer array.")

    cells = numpy.stack(numpy.meshgrid(*map(numpy.arange, ns))) \
        .reshape(3, -1).T.astype(float)
    cells = cells * self.box_size

    atoms = self.get_atoms('cell')
    atoms = Atoms.concat([
        atoms.transform(AffineTransform3D.translate(cell))
        for cell in cells
    ]) #.transform(self.cell.get_transform('local', 'cell_frac'))
    return self.with_atoms(atoms, 'cell').with_cell(self.get_cell().repeat(ns))

explode

explode() -> Self

Materialize repeated cells as one supercell.

Source code in atomlib/atomcell.py

def explode(self) -> Self:
    """Materialize repeated cells as one supercell."""
    frame = self.get_frame()

    return self.with_atoms(self.get_atoms('local'), 'local') \
        .with_cell(self.get_cell().explode()) \
        .to_frame(frame)

explode_z

explode_z() -> HasCellT

Materialize repeated cells as one supercell in z.

Source code in atomlib/cell.py

def explode_z(self: HasCellT) -> HasCellT:
    """Materialize repeated cells as one supercell in z."""
    return self.with_cell(Cell(
        affine=self.affine,
        ortho=self.ortho,
        cell_size=self.cell_size*[1, 1, self.n_cells[2]],
        n_cells=[*self.n_cells[:2], 1],
        cell_angle=self.cell_angle,
        pbc=self.pbc,
    ))

crop

crop(
    x_min: float = -inf,
    x_max: float = inf,
    y_min: float = -inf,
    y_max: float = inf,
    z_min: float = -inf,
    z_max: float = inf,
    *,
    frame: CoordinateFrame = "local"
) -> Self

Crop atoms and cell to the given extents. For a non-orthogonal cell, this must be specified in cell coordinates. This function implicity explodes the cell as well.

To crop atoms only, use crop_atoms instead.

Source code in atomlib/atomcell.py

def crop(self, x_min: float = -numpy.inf, x_max: float = numpy.inf,
         y_min: float = -numpy.inf, y_max: float = numpy.inf,
         z_min: float = -numpy.inf, z_max: float = numpy.inf, *,
         frame: CoordinateFrame = 'local') -> Self:
    """
    Crop atoms and cell to the given extents. For a non-orthogonal
    cell, this must be specified in cell coordinates. This
    function implicity `explode`s the cell as well.

    To crop atoms only, use `crop_atoms` instead.
    """

    cell = self.get_cell().crop(x_min, x_max, y_min, y_max, z_min, z_max, frame=frame)
    atoms = self._transform_atoms_in_frame(frame, lambda atoms: atoms.crop_atoms(x_min, x_max, y_min, y_max, z_min, z_max))
    return self.with_cell(cell).with_atoms(atoms)

change_transform

change_transform(
    transform: AffineTransform3D,
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> AffineTransform3D

change_transform(
    transform: Transform3D,
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> Transform3D

change_transform(
    transform: Transform3D,
    frame_to: Optional[CoordinateFrame] = None,
    frame_from: Optional[CoordinateFrame] = None,
) -> Transform3D

Coordinate-change a transformation from frame_from into frame_to.

Source code in atomlib/cell.py

def change_transform(self, transform: Transform3D,
                     frame_to: t.Optional[CoordinateFrame] = None,
                     frame_from: t.Optional[CoordinateFrame] = None) -> Transform3D:
    """Coordinate-change a transformation from `frame_from` into `frame_to`."""
    if frame_to == frame_from and frame_to is not None:
        return transform
    coord_change = self.get_transform(frame_to, frame_from)
    return coord_change @ transform @ coord_change.inverse()

assert_equal

assert_equal(other: Any)

Assert this structure is equal to other

Source code in atomlib/atomcell.py

def assert_equal(self, other: t.Any):
    """Assert this structure is equal to `other`"""
    assert isinstance(other, AtomCell)
    self.cell.assert_equal(other.cell)
    self.get_atoms('local').assert_equal(other.get_atoms('local'))

get_atoms

get_atoms(frame: Optional[CoordinateFrame] = None) -> Atoms

Get atoms contained in self, in the given coordinate frame.

Source code in atomlib/atomcell.py

def get_atoms(self, frame: t.Optional[CoordinateFrame] = None) -> Atoms:
    """Get atoms contained in ``self``, in the given coordinate frame."""

    if frame is None or frame == self.get_frame():
        return self.atoms
    return self.atoms.transform(self.get_transform(frame, self.get_frame()))

with_atoms

with_atoms(
    atoms: HasAtoms, frame: Optional[CoordinateFrame] = None
) -> Self

Source code in atomlib/atomcell.py

def with_atoms(self, atoms: HasAtoms, frame: t.Optional[CoordinateFrame] = None) -> Self:
    frame = frame if frame is not None else self.frame
    return self.__class__(atoms.get_atoms(), cell=self.cell, frame=frame, keep_frame=True)

describe

describe(
    percentiles: Union[Sequence[float], float, None] = (
        0.25,
        0.5,
        0.75,
    ),
    *,
    interpolation: RollingInterpolationMethod = "nearest",
    frame: Optional[CoordinateFrame] = None
) -> DataFrame

Return summary statistics for self. See DataFrame.describe for more information.

PARAMETER	DESCRIPTION
percentiles	List of percentiles/quantiles to include. Defaults to 25% (first quartile), 50% (median), and 75% (third quartile). TYPE: Union[Sequence[float], float, None] DEFAULT: (0.25, 0.5, 0.75)

RETURNS	DESCRIPTION
DataFrame	A dataframe containing summary statistics (mean, std. deviation, percentiles, etc.) for each column.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def describe(self, percentiles: t.Union[t.Sequence[float], float, None] = (0.25, 0.5, 0.75), *,
             interpolation: RollingInterpolationMethod = 'nearest',
             frame: t.Optional[CoordinateFrame] = None) -> polars.DataFrame:
    """
    Return summary statistics for `self`. See [`DataFrame.describe`][polars.DataFrame.describe] for more information.

    Args:
      percentiles: List of percentiles/quantiles to include. Defaults to 25% (first quartile),
                   50% (median), and 75% (third quartile).

    Returns:
      A dataframe containing summary statistics (mean, std. deviation, percentiles, etc.) for each column.
    """
    ...

with_columns

with_columns(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Self

Return a copy of self with the given columns added.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_columns(self,
                 *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
                 frame: t.Optional[CoordinateFrame] = None,
                 **named_exprs: IntoExpr) -> Self:
    """Return a copy of `self` with the given columns added."""
    ...

insert_column

insert_column(index: int, column: Series) -> DataFrame

Source code in atomlib/atoms.py

@_fwd_frame_map
def insert_column(self, index: int, column: polars.Series) -> polars.DataFrame:
    return self._get_frame().insert_column(index, column)

get_column

get_column(
    name: str, *, frame: Optional[CoordinateFrame] = None
) -> Series

Get the specified column from self, raising polars.ColumnNotFoundError if it's not present.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def get_column(self, name: str, *, frame: t.Optional[CoordinateFrame] = None) -> polars.Series:
    """
    Get the specified column from `self`, raising [`polars.ColumnNotFoundError`][polars.exceptions.ColumnNotFoundError] if it's not present.

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    ...

get_columns

get_columns(
    *, frame: Optional[CoordinateFrame] = None
) -> List[Series]

Return all columns from self as a list of Series.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def get_columns(self, *, frame: t.Optional[CoordinateFrame] = None) -> t.List[polars.Series]:
    """
    Return all columns from `self` as a list of [`Series`][polars.Series].

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    ...

get_column_index

get_column_index(name: str) -> int

Get the index of a column by name, raising polars.ColumnNotFoundError if it's not present.

Source code in atomlib/atoms.py

@_fwd_frame(polars.DataFrame.get_column_index)
def get_column_index(self, name: str) -> int:
    """Get the index of a column by name, raising [`polars.ColumnNotFoundError`][polars.exceptions.ColumnNotFoundError] if it's not present."""
    ...

group_by

group_by(
    *by: Union[IntoExpr, Iterable[IntoExpr]],
    maintain_order: bool = False,
    frame: Optional[CoordinateFrame] = None,
    **named_by: IntoExpr
) -> GroupBy

Start a group by operation. See DataFrame.group_by for more information.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def group_by(self, *by: t.Union[IntoExpr, t.Iterable[IntoExpr]],
             maintain_order: bool = False, frame: t.Optional[CoordinateFrame] = None,
             **named_by: IntoExpr) -> polars.dataframe.group_by.GroupBy:
    """
    Start a group by operation. See [`DataFrame.group_by`][polars.DataFrame.group_by] for more information.
    """
    ...

pipe

pipe(
    function: Callable[Concatenate[HasAtomCellT, P], T],
    *args: args,
    **kwargs: kwargs
) -> T

Apply function to self (in method-call syntax).

Source code in atomlib/atomcell.py

def pipe(self: HasAtomCellT, function: t.Callable[Concatenate[HasAtomCellT, P], T], *args: P.args, **kwargs: P.kwargs) -> T:
    """Apply `function` to `self` (in method-call syntax)."""
    return function(self, *args, **kwargs)

clone

clone() -> AtomCellT

Make a deep copy of self.

Source code in atomlib/atomcell.py

def clone(self: AtomCellT) -> AtomCellT:
    """Make a deep copy of `self`."""
    return self.__class__(**{field.name: copy.deepcopy(getattr(self, field.name)) for field in fields(self)})

drop

drop(
    *columns: Union[str, Iterable[str]], strict: bool = True
) -> DataFrame

Return self with the specified columns removed.

Source code in atomlib/atoms.py

def drop(self, *columns: t.Union[str, t.Iterable[str]], strict: bool = True) -> polars.DataFrame:
    """Return `self` with the specified columns removed."""
    return self._get_frame().drop(*columns, strict=strict)

filter

filter(
    *predicates: Union[
        None,
        IntoExprColumn,
        Iterable[IntoExprColumn],
        bool,
        List[bool],
        ndarray,
    ],
    frame: Optional[CoordinateFrame] = None,
    **constraints: Any
) -> Self

Filter self, removing rows which evaluate to False.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def filter(
    self,
    *predicates: t.Union[None, IntoExprColumn, t.Iterable[IntoExprColumn], bool, t.List[bool], numpy.ndarray],
    frame: t.Optional[CoordinateFrame] = None,
    **constraints: t.Any,
) -> Self:
    """Filter `self`, removing rows which evaluate to `False`."""
    ...

sort

sort(
    by: Union[IntoExpr, Iterable[IntoExpr]],
    *more_by: IntoExpr,
    descending: Union[bool, Sequence[bool]] = False,
    nulls_last: bool = False
) -> Self

Sort the atoms in self by the given columns/expressions.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def sort(
    self,
    by: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    *more_by: IntoExpr,
    descending: t.Union[bool, t.Sequence[bool]] = False,
    nulls_last: bool = False,
) -> Self:
    """
    Sort the atoms in `self` by the given columns/expressions.
    """
    ...

slice

slice(
    offset: int,
    length: Optional[int] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return a slice of the rows in self.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def slice(self, offset: int, length: t.Optional[int] = None, *,
          frame: t.Optional[CoordinateFrame] = None) -> Self:
    """Return a slice of the rows in `self`."""
    ...

head

head(
    n: int = 5, *, frame: Optional[CoordinateFrame] = None
) -> Self

Return the first n rows of self.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def head(self, n: int = 5, *, frame: t.Optional[CoordinateFrame] = None) -> Self:
    """Return the first `n` rows of `self`."""
    ...

tail

tail(
    n: int = 5, *, frame: Optional[CoordinateFrame] = None
) -> Self

Return the last n rows of self.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def tail(self, n: int = 5, *, frame: t.Optional[CoordinateFrame] = None) -> Self:
    """Return the last `n` rows of `self`."""
    ...

drop_nulls

drop_nulls(
    subset: Union[str, Collection[str], None] = None
) -> DataFrame

Drop rows that contain nulls in any of columns subset.

Source code in atomlib/atoms.py

@_fwd_frame_map
def drop_nulls(self, subset: t.Union[str, t.Collection[str], None] = None) -> polars.DataFrame:
    """Drop rows that contain nulls in any of columns `subset`."""
    return self._get_frame().drop_nulls(subset)

fill_null

fill_null(
    value: Any = None,
    strategy: Optional[FillNullStrategy] = None,
    limit: Optional[int] = None,
    matches_supertype: bool = True,
) -> Self

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def fill_null(
    self, value: t.Any = None, strategy: t.Optional[FillNullStrategy] = None,
    limit: t.Optional[int] = None, matches_supertype: bool = True,
) -> Self:
    ...

fill_nan

fill_nan(
    value: Union[Expr, int, float, None],
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def fill_nan(self, value: t.Union[polars.Expr, int, float, None], *,
             frame: t.Optional[CoordinateFrame] = None) -> Self:
    ...

concat classmethod

concat(
    atoms: Union[
        HasAtomsT,
        IntoAtoms,
        Iterable[Union[HasAtomsT, IntoAtoms]],
    ],
    *,
    rechunk: bool = True,
    how: ConcatMethod = "vertical"
) -> HasAtomsT

Concatenate multiple Atoms together, handling metadata appropriately.

Source code in atomlib/atoms.py

@classmethod
def concat(cls: t.Type[HasAtomsT],
           atoms: t.Union[HasAtomsT, IntoAtoms, t.Iterable[t.Union[HasAtomsT, IntoAtoms]]], *,
           rechunk: bool = True, how: ConcatMethod = 'vertical') -> HasAtomsT:
    """Concatenate multiple `Atoms` together, handling metadata appropriately."""
    # this method is tricky. It needs to accept raw Atoms, as well as HasAtoms of the
    # same type as ``cls``.
    if _is_abstract(cls):
        raise TypeError("concat() must be called on a concrete class.")

    if isinstance(atoms, HasAtoms):
        atoms = (atoms,)
    dfs = [a.get_atoms('local').inner if isinstance(a, HasAtoms) else Atoms(t.cast(IntoAtoms, a)).inner for a in atoms]
    representative = cls._combine_metadata(*(a for a in atoms if isinstance(a, HasAtoms)))

    if len(dfs) == 0:
        return representative.with_atoms(Atoms.empty(), 'local')

    if how in ('vertical', 'vertical_relaxed'):
        # get order from first member
        cols = dfs[0].columns
        dfs = [df.select(cols) for df in dfs]
    elif how == 'inner':
        cols = reduce(operator.and_, (df.schema.keys() for df in dfs))
        schema = OrderedDict((col, dfs[0].schema[col]) for col in cols)
        if len(schema) == 0:
            raise ValueError("Atoms have no columns in common")

        dfs = [_select_schema(df, schema) for df in dfs]
        how = 'vertical'

    return representative.with_atoms(Atoms(polars.concat(dfs, rechunk=rechunk, how=how)), 'local')

partition_by

partition_by(
    by: Union[str, Sequence[str]],
    *more_by: str,
    maintain_order: bool = True,
    include_key: bool = True,
    as_dict: Literal[False] = False
) -> List[Self]

partition_by(
    by: Union[str, Sequence[str]],
    *more_by: str,
    maintain_order: bool = True,
    include_key: bool = True,
    as_dict: Literal[True] = ...
) -> Dict[Any, Self]

partition_by(
    by: Union[str, Sequence[str]],
    *more_by: str,
    maintain_order: bool = True,
    include_key: bool = True,
    as_dict: bool = False
) -> Union[List[Self], Dict[Any, Self]]

Group by the given columns and partition into separate dataframes.

Return the partitions as a dictionary by specifying as_dict=True.

Source code in atomlib/atoms.py

def partition_by(
    self, by: t.Union[str, t.Sequence[str]], *more_by: str,
    maintain_order: bool = True, include_key: bool = True, as_dict: bool = False
) -> t.Union[t.List[Self], t.Dict[t.Any, Self]]:
    """
    Group by the given columns and partition into separate dataframes.

    Return the partitions as a dictionary by specifying `as_dict=True`.
    """
    if as_dict:
        d = self._get_frame().partition_by(by, *more_by, maintain_order=maintain_order, include_key=include_key, as_dict=True)
        return {k: self.with_atoms(Atoms(df, _unchecked=True)) for (k, df) in d.items()}

    return [
        self.with_atoms(Atoms(df, _unchecked=True))
        for df in self._get_frame().partition_by(by, *more_by, maintain_order=maintain_order, include_key=include_key, as_dict=False)
    ]

select

select(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> DataFrame

Select exprs from self, and return as a polars.DataFrame.

Expressions may either be columns or expressions of columns.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def select(
    self, *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    frame: t.Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> polars.DataFrame:
    """
    Select `exprs` from `self`, and return as a [`polars.DataFrame`][polars.DataFrame].

    Expressions may either be columns or expressions of columns.

    [polars.DataFrame]: https://docs.pola.rs/py-polars/html/reference/dataframe/index.html
    """
    ...

select_schema

select_schema(schema: SchemaDict) -> DataFrame

Select columns from self and cast to the given schema. Raises TypeError if a column is not found or if it can't be cast.

Source code in atomlib/atoms.py

def select_schema(self, schema: SchemaDict) -> polars.DataFrame:
    """
    Select columns from `self` and cast to the given schema.
    Raises [`TypeError`][TypeError] if a column is not found or if it can't be cast.
    """
    return _select_schema(self, schema)

select_props

select_props(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Self

Select exprs from self, while keeping required columns. Doesn't affect the cell.

RETURNS	DESCRIPTION
Self	A HasAtomCell filtered to contain
Self	the specified properties (as well as required columns).

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def select_props(
    self,
    *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    frame: t.Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Self:
    """
    Select `exprs` from `self`, while keeping required columns.
    Doesn't affect the cell.

    Returns:
      A [`HasAtomCell`][atomlib.atomcell.HasAtomCell] filtered to contain
      the specified properties (as well as required columns).
    """
    ...

try_select

try_select(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    frame: Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> Optional[DataFrame]

Try to select exprs from self, and return as a polars.DataFrame.

Expressions may either be columns or expressions of columns. Returns None if any columns are missing.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def try_select(
    self, *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    frame: t.Optional[CoordinateFrame] = None,
    **named_exprs: IntoExpr
) -> t.Optional[polars.DataFrame]:
    """
    Try to select `exprs` from `self`, and return as a [`polars.DataFrame`][polars.DataFrame].

    Expressions may either be columns or expressions of columns. Returns `None` if any
    columns are missing.

    [polars.DataFrame]: https://docs.pola.rs/py-polars/html/reference/dataframe/index.html
    """
    ...

try_get_column

try_get_column(name: str) -> Optional[Series]

Try to get a column from self, returning None if it doesn't exist.

Source code in atomlib/atoms.py

def try_get_column(self, name: str) -> t.Optional[polars.Series]:
    """Try to get a column from `self`, returning `None` if it doesn't exist."""
    try:
        return self.get_column(name)
    except polars.exceptions.ColumnNotFoundError:
        return None

bbox_atoms

bbox_atoms(
    frame: Optional[CoordinateFrame] = None,
) -> BBox3D

Return the bounding box of all the atoms in self, in the given coordinate frame.

Source code in atomlib/atomcell.py

def bbox_atoms(self, frame: t.Optional[CoordinateFrame] = None) -> BBox3D:
    """Return the bounding box of all the atoms in `self`, in the given coordinate frame."""
    return self.get_atoms(frame).bbox()

transform_atoms

transform_atoms(
    transform: IntoTransform3D,
    selection: Optional[AtomSelection] = None,
    *,
    frame: CoordinateFrame = "local",
    transform_velocities: bool = False
) -> Self

Transform the atoms in self by transform. If selection is given, only transform the atoms in selection.

Source code in atomlib/atomcell.py

def transform_atoms(self, transform: IntoTransform3D, selection: t.Optional[AtomSelection] = None, *,
                    frame: CoordinateFrame = 'local', transform_velocities: bool = False) -> Self:
    """
    Transform the atoms in `self` by `transform`.
    If `selection` is given, only transform the atoms in `selection`.
    """
    transform = self.change_transform(Transform3D.make(transform), self.get_frame(), frame)
    return self.with_atoms(self.get_atoms(self.get_frame()).transform(transform, selection, transform_velocities=transform_velocities))

transform

transform(
    transform: AffineTransform3D,
    frame: CoordinateFrame = "local",
) -> Self

Source code in atomlib/atomcell.py

def transform(self, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> Self:
    if isinstance(transform, Transform3D) and not isinstance(transform, AffineTransform3D):
        raise ValueError("Non-affine transforms cannot change the box dimensions. Use 'transform_atoms' instead.")
    # TODO: cleanup once tests pass
    # coordinate change the transform into atomic coordinates
    new_cell = self.get_cell().transform_cell(transform, frame)
    transform = self.get_cell().change_transform(transform, self.get_frame(), frame)
    return self.with_atoms(self.get_atoms().transform(transform), self.get_frame()).with_cell(new_cell)

round_near_zero

round_near_zero(
    tol: float = 1e-14,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Round atom position values near zero to zero.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def round_near_zero(self, tol: float = 1e-14, *,
                    frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Round atom position values near zero to zero.
    """
    ...

crop_atoms

crop_atoms(
    x_min: float = -inf,
    x_max: float = inf,
    y_min: float = -inf,
    y_max: float = inf,
    z_min: float = -inf,
    z_max: float = inf,
    *,
    frame: CoordinateFrame = "local"
) -> Self

Source code in atomlib/atomcell.py

def crop_atoms(self, x_min: float = -numpy.inf, x_max: float = numpy.inf,
               y_min: float = -numpy.inf, y_max: float = numpy.inf,
               z_min: float = -numpy.inf, z_max: float = numpy.inf, *,
               frame: CoordinateFrame = 'local') -> Self:
    atoms = self._transform_atoms_in_frame(frame, lambda atoms: atoms.crop_atoms(x_min, x_max, y_min, y_max, z_min, z_max))
    return self.with_atoms(atoms)

deduplicate

deduplicate(
    tol: float = 0.001,
    subset: Iterable[str] = ("x", "y", "z", "symbol"),
    keep: UniqueKeepStrategy = "first",
    maintain_order: bool = True,
) -> Self

De-duplicate atoms in self. Atoms of the same symbol that are closer than tolerance to each other (by Euclidian distance) will be removed, leaving only the atom specified by keep (defaults to the first atom).

If subset is specified, only those columns will be included while assessing duplicates. Floating point columns other than 'x', 'y', and 'z' will not by toleranced.

Source code in atomlib/atoms.py

def deduplicate(self, tol: float = 1e-3, subset: t.Iterable[str] = ('x', 'y', 'z', 'symbol'),
                keep: UniqueKeepStrategy = 'first', maintain_order: bool = True) -> Self:
    """
    De-duplicate atoms in `self`. Atoms of the same `symbol` that are closer than `tolerance`
    to each other (by Euclidian distance) will be removed, leaving only the atom specified by
    `keep` (defaults to the first atom).

    If `subset` is specified, only those columns will be included while assessing duplicates.
    Floating point columns other than 'x', 'y', and 'z' will not by toleranced.
    """
    import scipy.spatial

    cols = set((subset,) if isinstance(subset, str) else subset)

    indices = numpy.arange(len(self))

    spatial_cols = cols.intersection(('x', 'y', 'z'))
    cols -= spatial_cols
    if len(spatial_cols) > 0:
        coords = self.select([_coord_expr(col).alias(col) for col in spatial_cols]).to_numpy()
        tree = scipy.spatial.KDTree(coords)

        # TODO This is a bad algorithm
        while True:
            changed = False
            for (i, j) in tree.query_pairs(tol, 2.):
                # whenever we encounter a pair, ensure their index matches
                i_i, i_j = indices[[i, j]]
                if i_i != i_j:
                    indices[i] = indices[j] = min(i_i, i_j)
                    changed = True
            if not changed:
                break

        self = self.with_column(polars.Series('_unique_pts', indices))
        cols.add('_unique_pts')

    frame = self._get_frame().unique(subset=list(cols), keep=keep, maintain_order=maintain_order)
    if len(spatial_cols) > 0:
        frame = frame.drop('_unique_pts')

    return self.with_atoms(Atoms(frame, _unchecked=True))

with_bounds

with_bounds(
    cell_size: Optional[VecLike] = None,
    cell_origin: Optional[VecLike] = None,
) -> "AtomCell"

Return a periodic cell with the given orthogonal cell dimensions.

If cell_size is not specified, it will be assumed (and may be incorrect).

Source code in atomlib/atoms.py

def with_bounds(self, cell_size: t.Optional[VecLike] = None, cell_origin: t.Optional[VecLike] = None) -> 'AtomCell':
    """
    Return a periodic cell with the given orthogonal cell dimensions.

    If cell_size is not specified, it will be assumed (and may be incorrect).
    """
    # TODO: test this
    from .atomcell import AtomCell

    if cell_size is None:
        warnings.warn("Cell boundary unknown. Defaulting to cell BBox")
        cell_size = self.bbox().size
        cell_origin = self.bbox().min

    # TODO test this origin code
    cell = Cell.from_unit_cell(cell_size)
    if cell_origin is not None:
        cell = cell.transform_cell(AffineTransform3D.translate(to_vec3(cell_origin)))

    return AtomCell(self.get_atoms(), cell, frame='local')

coords

coords(
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> NDArray[float64]

Return a (N, 3) ndarray of atom positions (dtype numpy.float64) in the given coordinate frame.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def coords(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Optional[CoordinateFrame] = None) -> NDArray[numpy.float64]:
    """
    Return a `(N, 3)` ndarray of atom positions (dtype [`numpy.float64`][numpy.float64])
    in the given coordinate frame.
    """
    ...

x

x() -> Expr

Source code in atomlib/atoms.py

def x(self) -> polars.Expr:
    return polars.col('coords').arr.get(0).alias('x')

y

y() -> Expr

Source code in atomlib/atoms.py

def y(self) -> polars.Expr:
    return polars.col('coords').arr.get(1).alias('y')

z

z() -> Expr

Source code in atomlib/atoms.py

def z(self) -> polars.Expr:
    return polars.col('coords').arr.get(2).alias('z')

velocities

velocities(
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Optional[NDArray[float64]]

Return a (N, 3) ndarray of atom velocities (dtype numpy.float64) in the given coordinate frame.

Source code in atomlib/atomcell.py

@_fwd_atoms_get
def velocities(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Optional[CoordinateFrame] = None) -> t.Optional[NDArray[numpy.float64]]:
    """
    Return a `(N, 3)` ndarray of atom velocities (dtype [`numpy.float64`][numpy.float64])
    in the given coordinate frame.
    """
    ...

types

types() -> Optional[Series]

Returns a Series of atom types (dtype polars.Int32).

Source code in atomlib/atoms.py

def types(self) -> t.Optional[polars.Series]:
    """
    Returns a [`Series`][polars.Series] of atom types (dtype [`polars.Int32`][polars.datatypes.Int32]).

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    return self.try_get_column('type')

masses

masses() -> Optional[Series]

Returns a Series of atom masses (dtype polars.Float32).

Source code in atomlib/atoms.py

def masses(self) -> t.Optional[polars.Series]:
    """
    Returns a [`Series`][polars.Series] of atom masses (dtype [`polars.Float32`][polars.datatypes.Float32]).

    [polars.Series]: https://docs.pola.rs/py-polars/html/reference/series/index.html
    """
    return self.try_get_column('mass')

add_atom

add_atom(
    elem: Union[int, str],
    x: ArrayLike,
    /,
    *,
    y: None = None,
    z: None = None,
    frame: Optional[CoordinateFrame] = None,
    **kwargs: Any,
) -> Self

add_atom(
    elem: Union[int, str],
    /,
    x: float,
    y: float,
    z: float,
    *,
    frame: Optional[CoordinateFrame] = None,
    **kwargs: Any,
) -> Self

add_atom(
    elem: Union[int, str],
    /,
    x: Union[ArrayLike, float],
    y: Optional[float] = None,
    z: Optional[float] = None,
    *,
    frame: Optional[CoordinateFrame] = None,
    **kwargs: Any,
) -> Self

Return a copy of self with an extra atom.

By default, all extra columns present in self must be specified as **kwargs.

Try to avoid calling this in a loop (Use concat instead).

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def add_atom(self, elem: t.Union[int, str], /,  # type: ignore (spurious)
             x: t.Union[ArrayLike, float],
             y: t.Optional[float] = None,
             z: t.Optional[float] = None, *,
             frame: t.Optional[CoordinateFrame] = None,
             **kwargs: t.Any) -> Self:
    """
    Return a copy of `self` with an extra atom.

    By default, all extra columns present in `self` must be specified as `**kwargs`.

    Try to avoid calling this in a loop (Use [`concat`][atomlib.atomcell.HasAtomCell.concat] instead).
    """
    ...

pos

pos(
    x: Sequence[Optional[float]],
    /,
    *,
    y: None = None,
    z: None = None,
    tol: float = 1e-06,
    **kwargs: Any,
) -> Expr

pos(
    x: Optional[float] = None,
    y: Optional[float] = None,
    z: Optional[float] = None,
    *,
    tol: float = 1e-06,
    **kwargs: Any
) -> Expr

pos(
    x: Union[Sequence[Optional[float]], float, None] = None,
    y: Optional[float] = None,
    z: Optional[float] = None,
    *,
    tol: float = 1e-06,
    **kwargs: Any
) -> Expr

Select all atoms at a given position.

Formally, returns all atoms within a cube of radius tol centered at (x,y,z), exclusive of the cube's surface.

Additional parameters given as kwargs will be checked as additional parameters (with strict equality).

Source code in atomlib/atoms.py

def pos(self,
        x: t.Union[t.Sequence[t.Optional[float]], float, None] = None,
        y: t.Optional[float] = None, z: t.Optional[float] = None, *,
        tol: float = 1e-6, **kwargs: t.Any) -> polars.Expr:
    """
    Select all atoms at a given position.

    Formally, returns all atoms within a cube of radius ``tol``
    centered at ``(x,y,z)``, exclusive of the cube's surface.

    Additional parameters given as ``kwargs`` will be checked
    as additional parameters (with strict equality).
    """

    if isinstance(x, t.Sequence):
        (x, y, z) = x

    tol = abs(float(tol))
    selection = polars.lit(True)
    if x is not None:
        selection &= self.x().is_between(x - tol, x + tol, closed='none')
    if y is not None:
        selection &= self.y().is_between(y - tol, y + tol, closed='none')
    if z is not None:
        selection &= self.z().is_between(z - tol, z + tol, closed='none')
    for (col, val) in kwargs.items():
        selection &= (polars.col(col) == val)

    return selection

with_index

with_index(
    index: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Returns self with a row index added in column 'i' (dtype polars.Int64). If index is not specified, defaults to an existing index or a new index.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_index(self, index: t.Optional[AtomValues] = None, *,
               frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Returns `self` with a row index added in column 'i' (dtype [`polars.Int64`][polars.datatypes.Int64]).
    If `index` is not specified, defaults to an existing index or a new index.
    """
    ...

with_wobble

with_wobble(
    wobble: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given displacements in column 'wobble' (dtype polars.Float64). If wobble is not specified, defaults to the already-existing wobbles or 0.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_wobble(self, wobble: t.Optional[AtomValues] = None, *,
                frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given displacements in column 'wobble' (dtype [`polars.Float64`][polars.datatypes.Float64]).
    If `wobble` is not specified, defaults to the already-existing wobbles or 0.
    """
    ...

with_occupancy

with_occupancy(
    frac_occupancy: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given fractional occupancies (dtype polars.Float64). If frac_occupancy is not specified, defaults to the already-existing occupancies or 1.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_occupancy(self, frac_occupancy: t.Optional[AtomValues] = None, *,
                   frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return self with the given fractional occupancies (dtype [`polars.Float64`][polars.datatypes.Float64]).
    If `frac_occupancy` is not specified, defaults to the already-existing occupancies or 1.
    """
    ...

apply_wobble

apply_wobble(
    rng: Union[Generator, int, None] = None,
    frame: Optional[CoordinateFrame] = None,
) -> Self

Displace the atoms in self by the amount in the wobble column. wobble is interpretated as a mean-squared displacement, which is distributed equally over each axis.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def apply_wobble(self, rng: t.Union[numpy.random.Generator, int, None] = None,
                 frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Displace the atoms in `self` by the amount in the `wobble` column.
    `wobble` is interpretated as a mean-squared displacement, which is distributed
    equally over each axis.
    """
    ...

apply_occupancy

apply_occupancy(
    rng: Union[Generator, int, None] = None
) -> Self

For each atom in self, use its frac_occupancy to randomly decide whether to remove it.

Source code in atomlib/atoms.py

def apply_occupancy(self, rng: t.Union[numpy.random.Generator, int, None] = None) -> Self:
    """
    For each atom in `self`, use its `frac_occupancy` to randomly decide whether to remove it.
    """
    if 'frac_occupancy' not in self.columns:
        return self
    rng = numpy.random.default_rng(seed=rng)

    frac = self.select('frac_occupancy').to_series().to_numpy()
    choice = rng.binomial(1, frac).astype(numpy.bool_)
    return self.filter(polars.lit(choice))

with_type

with_type(
    types: Optional[AtomValues] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given atom types in column 'type'. If types is not specified, use the already existing types or auto-assign them.

When auto-assigning, each symbol is given a unique value, case-sensitive. Values are assigned from lowest atomic number to highest. For instance: ["Ag+", "Na", "H", "Ag"] => [3, 11, 1, 2]

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_type(self, types: t.Optional[AtomValues] = None, *,
              frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given atom types in column 'type'.
    If `types` is not specified, use the already existing types or auto-assign them.

    When auto-assigning, each symbol is given a unique value, case-sensitive.
    Values are assigned from lowest atomic number to highest.
    For instance: `["Ag+", "Na", "H", "Ag"]` => `[3, 11, 1, 2]`
    """
    ...

with_mass

with_mass(
    mass: Optional[ArrayLike] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given atom masses in column 'mass'. If mass is not specified, use the already existing masses or auto-assign them.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_mass(self, mass: t.Optional[ArrayLike] = None, *,
              frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given atom masses in column `'mass'`.
    If `mass` is not specified, use the already existing masses or auto-assign them.
    """
    ...

with_symbol

with_symbol(
    symbols: ArrayLike,
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self with the given atomic symbols.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_symbol(self, symbols: ArrayLike, selection: t.Optional[AtomSelection] = None, *,
                frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` with the given atomic symbols.
    """
    ...

with_coords

with_coords(
    pts: ArrayLike,
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self replaced with the given atomic positions.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_coords(self, pts: ArrayLike, selection: t.Optional[AtomSelection] = None, *,
                frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` replaced with the given atomic positions.
    """
    ...

with_velocity

with_velocity(
    pts: Optional[ArrayLike] = None,
    selection: Optional[AtomSelection] = None,
    *,
    frame: Optional[CoordinateFrame] = None
) -> Self

Return self replaced with the given atomic velocities. If pts is not specified, use the already existing velocities or zero.

Source code in atomlib/atomcell.py

@_fwd_atoms_transform
def with_velocity(self, pts: t.Optional[ArrayLike] = None,
                  selection: t.Optional[AtomSelection] = None, *,
                  frame: t.Optional[CoordinateFrame] = None) -> Self:
    """
    Return `self` replaced with the given atomic velocities.
    If `pts` is not specified, use the already existing velocities or zero.
    """
    ...

get_frame

get_frame() -> CoordinateFrame

Get the coordinate frame atoms are stored in.

Source code in atomlib/atomcell.py

def get_frame(self) -> CoordinateFrame:
    """Get the coordinate frame atoms are stored in."""
    return self.frame

get_atomcell

get_atomcell() -> AtomCell

Source code in atomlib/atomcell.py

def get_atomcell(self) -> AtomCell:
    frame = self.get_frame()
    return AtomCell(self.get_atoms(frame), self.get_cell(), frame=frame, keep_frame=True)

to_frame

to_frame(frame: CoordinateFrame) -> Self

Convert the stored Atoms to the given coordinate frame.

Source code in atomlib/atomcell.py

def to_frame(self, frame: CoordinateFrame) -> Self:
    """Convert the stored Atoms to the given coordinate frame."""
    return self.with_atoms(self.get_atoms(frame), frame)

crop_to_box

crop_to_box(eps: float = 1e-05) -> Self

Source code in atomlib/atomcell.py

def crop_to_box(self, eps: float = 1e-5) -> Self:
    atoms = self._transform_atoms_in_frame('cell_box', lambda atoms: atoms.crop_atoms(*([-eps, 1-eps]*3)))
    return self.with_atoms(atoms)

wrap

wrap(eps: float = 1e-05) -> Self

Wrap atoms around the cell boundaries.

Source code in atomlib/atomcell.py

def wrap(self, eps: float = 1e-5) -> Self:
    """Wrap atoms around the cell boundaries."""
    return self.with_atoms(self._transform_atoms_in_frame('cell_box', lambda a: a._wrap(eps)))

repeat_to

repeat_to(
    size: VecLike, crop: Union[bool, Sequence[bool]] = False
) -> Self

Repeat the cell so it is at least size along the crystal's axes.

If crop, then crop the cell to exactly size. This may break periodicity. crop may be a vector, in which case you can specify cropping only along some axes.

Source code in atomlib/atomcell.py

def repeat_to(self, size: VecLike, crop: t.Union[bool, t.Sequence[bool]] = False) -> Self:
    """
    Repeat the cell so it is at least `size` along the crystal's axes.

    If `crop`, then crop the cell to exactly `size`. This may break periodicity.
    `crop` may be a vector, in which case you can specify cropping only along some axes.
    """
    size = to_vec3(size)
    cell_size = self.cell_size * self.n_cells
    repeat = numpy.maximum(numpy.ceil(size / cell_size).astype(int), 1)
    atom_cell = self.repeat(repeat)

    crop_v = to_vec3(crop, dtype=numpy.bool_)
    if numpy.any(crop_v):
        crop_x, crop_y, crop_z = crop_v
        return atom_cell.crop(
            x_max = size[0] if crop_x else numpy.inf,
            y_max = size[1] if crop_y else numpy.inf,
            z_max = size[2] if crop_z else numpy.inf,
            frame='cell'
        )

    return atom_cell

repeat_x

repeat_x(n: int) -> Self

Tile the cell in the x axis.

Source code in atomlib/atomcell.py

def repeat_x(self, n: int) -> Self:
    """Tile the cell in the x axis."""
    return self.repeat((n, 1, 1))

repeat_y

repeat_y(n: int) -> Self

Tile the cell in the y axis.

Source code in atomlib/atomcell.py

def repeat_y(self, n: int) -> Self:
    """Tile the cell in the y axis."""
    return self.repeat((1, n, 1))

repeat_z

repeat_z(n: int) -> Self

Tile the cell in the z axis.

Source code in atomlib/atomcell.py

def repeat_z(self, n: int) -> Self:
    """Tile the cell in the z axis."""
    return self.repeat((1, 1, n))

repeat_to_x

repeat_to_x(size: float, crop: bool = False) -> Self

Repeat the cell so it is at least size size along the x axis.

Source code in atomlib/atomcell.py

def repeat_to_x(self, size: float, crop: bool = False) -> Self:
    """Repeat the cell so it is at least size `size` along the x axis."""
    return self.repeat_to([size, 0., 0.], [crop, False, False])

repeat_to_y

repeat_to_y(size: float, crop: bool = False) -> Self

Repeat the cell so it is at least size size along the y axis.

Source code in atomlib/atomcell.py

def repeat_to_y(self, size: float, crop: bool = False) -> Self:
    """Repeat the cell so it is at least size `size` along the y axis."""
    return self.repeat_to([0., size, 0.], [False, crop, False])

repeat_to_z

repeat_to_z(size: float, crop: bool = False) -> Self

Repeat the cell so it is at least size size along the z axis.

Source code in atomlib/atomcell.py

def repeat_to_z(self, size: float, crop: bool = False) -> Self:
    """Repeat the cell so it is at least size `size` along the z axis."""
    return self.repeat_to([0., 0., size], [False, False, crop])

repeat_to_aspect

repeat_to_aspect(
    plane: Literal["xy", "xz", "yz"] = "xy",
    *,
    aspect: float = 1.0,
    min_size: Optional[VecLike] = None,
    max_size: Optional[VecLike] = None
) -> Self

Repeat to optimize the aspect ratio in plane, while staying above min_size and under max_size.

Source code in atomlib/atomcell.py

def repeat_to_aspect(self, plane: t.Literal['xy', 'xz', 'yz'] = 'xy', *,
                     aspect: float = 1., min_size: t.Optional[VecLike] = None,
                     max_size: t.Optional[VecLike] = None) -> Self:
    """
    Repeat to optimize the aspect ratio in `plane`,
    while staying above `min_size` and under `max_size`.
    """
    if min_size is None:
        min_n = numpy.array([1, 1, 1], numpy.int_)
    else:
        min_n = numpy.maximum(numpy.ceil(to_vec3(min_size) / self.box_size), 1).astype(numpy.int_)

    if max_size is None:
        max_n = 3 * min_n
    else:
        max_n = numpy.maximum(numpy.floor(to_vec3(max_size) / self.box_size), 1).astype(numpy.int_)

    if plane == 'xy':
        indices = [0, 1]
    elif plane == 'xz':
        indices = [0, 2]
    elif plane == 'yz':
        indices = [1, 2]
    else:
        raise ValueError(f"Invalid plane '{plane}'. Exepcted 'xy', 'xz', 'or 'yz'.")

    na = numpy.arange(min_n[indices[0]], max_n[indices[0]])
    nb = numpy.arange(min_n[indices[1]], max_n[indices[1]])
    (na, nb) = numpy.meshgrid(na, nb)

    aspects = na * self.box_size[indices[0]] / (nb * self.box_size[indices[1]])
    # cost function: log(aspect)^2  (so cost(0.5) == cost(2))
    min_i = numpy.argmin(numpy.log(aspects / aspect)**2)
    repeat = numpy.array([1, 1, 1], numpy.int_)
    repeat[indices] = na.flatten()[min_i], nb.flatten()[min_i]
    return self.repeat(repeat)

periodic_duplicate

periodic_duplicate(eps: float = 1e-05) -> Self

Add duplicate copies of atoms near periodic boundaries.

For instance, an atom at a corner will be duplicated into 8 copies. This is mostly only useful for visualization.

Source code in atomlib/atomcell.py

def periodic_duplicate(self, eps: float = 1e-5) -> Self:
    """
    Add duplicate copies of atoms near periodic boundaries.

    For instance, an atom at a corner will be duplicated into 8 copies.
    This is mostly only useful for visualization.
    """
    frame_save = self.get_frame()
    self = self.to_frame('cell_box').wrap(eps=eps)

    for i in range(3):
        self = self.concat((self,
            self.filter(polars.col('coords').arr.get(i).abs() <= eps, frame='cell_box')
                .transform_atoms(AffineTransform3D.translate([1. if i == j else 0. for j in range(3)]), frame='cell_box')
        ))

    return self.to_frame(frame_save)

read classmethod

read(path: FileOrPath, ty: FileType) -> HasAtomsT

read(
    path: Union[str, Path, TextIO], ty: Literal[None] = None
) -> HasAtomsT

read(
    path: FileOrPath, ty: Optional[FileType] = None
) -> HasAtomsT

Read a structure from a file.

Supported types can be found in the io module. If no ty is specified, it is inferred from the file's extension.

Source code in atomlib/mixins.py

@classmethod
def read(cls: t.Type[HasAtomsT], path: FileOrPath, ty: t.Optional[FileType] = None) -> HasAtomsT:
    """
    Read a structure from a file.

    Supported types can be found in the [io][atomlib.io] module.
    If no `ty` is specified, it is inferred from the file's extension.
    """
    from .io import read
    return _cast_atoms(read(path, ty), cls)  # type: ignore

read_cif classmethod

read_cif(
    f: Union[FileOrPath, CIF, CIFDataBlock],
    block: Union[int, str, None] = None,
) -> HasAtomsT

Read a structure from a CIF file.

If block is specified, read data from the given block of the CIF file (index or name).

Source code in atomlib/mixins.py

@classmethod
def read_cif(cls: t.Type[HasAtomsT], f: t.Union[FileOrPath, CIF, CIFDataBlock], block: t.Union[int, str, None] = None) -> HasAtomsT:
    """
    Read a structure from a CIF file.

    If `block` is specified, read data from the given block of the CIF file (index or name).
    """
    from .io import read_cif
    return _cast_atoms(read_cif(f, block), cls)

read_xyz classmethod

read_xyz(f: Union[FileOrPath, XYZ]) -> HasAtomsT

Read a structure from an XYZ file.

Source code in atomlib/mixins.py

@classmethod
def read_xyz(cls: t.Type[HasAtomsT], f: t.Union[FileOrPath, XYZ]) -> HasAtomsT:
    """Read a structure from an XYZ file."""
    from .io import read_xyz
    return _cast_atoms(read_xyz(f), cls)

read_xsf classmethod

read_xsf(f: Union[FileOrPath, XSF]) -> HasAtomsT

Read a structure from an XSF file.

Source code in atomlib/mixins.py

@classmethod
def read_xsf(cls: t.Type[HasAtomsT], f: t.Union[FileOrPath, XSF]) -> HasAtomsT:
    """Read a structure from an XSF file."""
    from .io import read_xsf
    return _cast_atoms(read_xsf(f), cls)

read_cfg classmethod

read_cfg(f: Union[FileOrPath, CFG]) -> HasAtomsT

Read a structure from a CFG file.

Source code in atomlib/mixins.py

@classmethod
def read_cfg(cls: t.Type[HasAtomsT], f: t.Union[FileOrPath, CFG]) -> HasAtomsT:
    """Read a structure from a CFG file."""
    from .io import read_cfg
    return _cast_atoms(read_cfg(f), cls)

read_lmp classmethod

read_lmp(
    f: Union[FileOrPath, LMP],
    type_map: Optional[Dict[int, Union[str, int]]] = None,
) -> HasAtomsT

Read a structure from a LAAMPS data file.

Source code in atomlib/mixins.py

@classmethod
def read_lmp(cls: t.Type[HasAtomsT], f: t.Union[FileOrPath, LMP], type_map: t.Optional[t.Dict[int, t.Union[str, int]]] = None) -> HasAtomsT:
    """Read a structure from a LAAMPS data file."""
    from .io import read_lmp
    return _cast_atoms(read_lmp(f, type_map=type_map), cls)

write_cif

write_cif(f: FileOrPath)

Source code in atomlib/mixins.py

def write_cif(self, f: FileOrPath):
    from .io import write_cif
    write_cif(self, f)

write_xyz

write_xyz(f: FileOrPath, fmt: XYZFormat = 'exyz')

Source code in atomlib/mixins.py

def write_xyz(self, f: FileOrPath, fmt: XYZFormat = 'exyz'):
    from .io import write_xyz
    write_xyz(self, f, fmt)

write_xsf

write_xsf(f: FileOrPath)

Source code in atomlib/mixins.py

def write_xsf(self, f: FileOrPath):
    from .io import write_xsf
    write_xsf(self, f)

write_cfg

write_cfg(f: FileOrPath)

Source code in atomlib/mixins.py

def write_cfg(self, f: FileOrPath):
    from .io import write_cfg
    write_cfg(self, f)

write_lmp

write_lmp(f: FileOrPath)

Source code in atomlib/mixins.py

def write_lmp(self, f: FileOrPath):
    from .io import write_lmp
    write_lmp(self, f)

write

write(path: FileOrPath, ty: FileType)

write(
    path: Union[str, Path, TextIO], ty: Literal[None] = None
)

write(path: FileOrPath, ty: Optional[FileType] = None)

Write this structure to a file.

A file type may be specified using ty. If no ty is specified, it is inferred from the path's extension.

Source code in atomlib/mixins.py

def write(self, path: FileOrPath, ty: t.Optional[FileType] = None):
    """
    Write this structure to a file.

    A file type may be specified using `ty`.
    If no `ty` is specified, it is inferred from the path's extension.
    """
    from .io import write
    write(self, path, ty)  # type: ignore

write_mslice

write_mslice(
    f: BinaryFileOrPath,
    template: Optional[MSliceFile] = None,
    *,
    slice_thickness: Optional[float] = None,
    scan_points: Optional[ArrayLike] = None,
    scan_extent: Optional[ArrayLike] = None,
    noise_sigma: Optional[float] = None,
    conv_angle: Optional[float] = None,
    energy: Optional[float] = None,
    defocus: Optional[float] = None,
    tilt: Optional[Tuple[float, float]] = None,
    tds: Optional[bool] = None,
    n_cells: Optional[ArrayLike] = None
)

Write a structure to an mslice file.

template may be a file, path, or ElementTree containing an existing mslice file. Its structure will be modified to make the final output. If not specified, a default template will be used.

Additional options modify simulation properties. If an option is not specified, the template's properties are used.

Source code in atomlib/mixins.py

def write_mslice(self, f: BinaryFileOrPath, template: t.Optional[MSliceFile] = None, *,
             slice_thickness: t.Optional[float] = None,  # angstrom
             scan_points: t.Optional[ArrayLike] = None,
             scan_extent: t.Optional[ArrayLike] = None,
             noise_sigma: t.Optional[float] = None,  # angstrom
             conv_angle: t.Optional[float] = None,  # mrad
             energy: t.Optional[float] = None,  # keV
             defocus: t.Optional[float] = None,  # angstrom
             tilt: t.Optional[t.Tuple[float, float]] = None,  # (mrad, mrad)
             tds: t.Optional[bool] = None,
             n_cells: t.Optional[ArrayLike] = None):
    """
    Write a structure to an mslice file.

    `template` may be a file, path, or `ElementTree` containing an existing mslice file.
    Its structure will be modified to make the final output. If not specified, a default
    template will be used.

    Additional options modify simulation properties. If an option is not specified, the
    template's properties are used.
    """
    from .io import write_mslice
    return write_mslice(self, f, template, slice_thickness=slice_thickness,
                        scan_points=scan_points, scan_extent=scan_extent,
                        conv_angle=conv_angle, energy=energy, defocus=defocus,
                        noise_sigma=noise_sigma, tilt=tilt, tds=tds, n_cells=n_cells)

write_qe

write_qe(
    f: FileOrPath,
    pseudo: Optional[Mapping[str, str]] = None,
)

Write a structure to a Quantum Espresso pw.x file.

PARAMETER	DESCRIPTION
f	File or path to write to TYPE: FileOrPath
pseudo	Mapping from atom symbol TYPE: Optional[Mapping[str, str]] DEFAULT: None

Source code in atomlib/mixins.py

def write_qe(self, f: FileOrPath, pseudo: t.Optional[t.Mapping[str, str]] = None):
    """
    Write a structure to a Quantum Espresso pw.x file.

    Args:
      f: File or path to write to
      pseudo: Mapping from atom symbol
    """
    from .io import write_qe
    write_qe(self, f, pseudo)

from_ortho classmethod

from_ortho(
    atoms: IntoAtoms,
    ortho: LinearTransform3D,
    *,
    n_cells: Optional[VecLike] = None,
    frame: CoordinateFrame = "local",
    keep_frame: bool = False
)

Make an atom cell given a list of atoms and an orthogonalization matrix. Atoms are assumed to be in the coordinate system frame.

Source code in atomlib/atomcell.py

@classmethod
def from_ortho(cls, atoms: IntoAtoms, ortho: LinearTransform3D, *,
               n_cells: t.Optional[VecLike] = None,
               frame: CoordinateFrame = 'local',
               keep_frame: bool = False):
    """
    Make an atom cell given a list of atoms and an orthogonalization matrix.
    Atoms are assumed to be in the coordinate system `frame`.
    """
    cell = Cell.from_ortho(ortho, n_cells)
    return cls(atoms, cell, frame=frame, keep_frame=keep_frame)

from_unit_cell classmethod

from_unit_cell(
    atoms: IntoAtoms,
    cell_size: VecLike,
    cell_angle: Optional[VecLike] = None,
    *,
    n_cells: Optional[VecLike] = None,
    frame: CoordinateFrame = "local",
    keep_frame: bool = False
)

Make a cell given a list of atoms and unit cell parameters. Atoms are assumed to be in the coordinate system frame.

Source code in atomlib/atomcell.py

@classmethod
def from_unit_cell(cls, atoms: IntoAtoms, cell_size: VecLike,
                   cell_angle: t.Optional[VecLike] = None, *,
                   n_cells: t.Optional[VecLike] = None,
                   frame: CoordinateFrame = 'local',
                   keep_frame: bool = False):
    """
    Make a cell given a list of atoms and unit cell parameters.
    Atoms are assumed to be in the coordinate system `frame`.
    """
    cell = Cell.from_unit_cell(cell_size, cell_angle, n_cells=n_cells)
    return cls(atoms, cell, frame=frame, keep_frame=keep_frame)

orthogonalize

orthogonalize() -> OrthoCell

Source code in atomlib/atomcell.py

def orthogonalize(self) -> OrthoCell:
    if self.is_orthogonal():
        return OrthoCell(self.atoms, self.cell, frame=self.frame)
    raise NotImplementedError()

is_orthogonal

is_orthogonal(tol: float = 1e-08) -> Literal[True]

Returns whether this cell is orthogonal (axes are at right angles.)

Source code in atomlib/atomcell.py

def is_orthogonal(self, tol: float = 1e-8) -> t.Literal[True]:
    """Returns whether this cell is orthogonal (axes are at right angles.)"""
    return True