atomlib.atoms

Raw atoms collection

This module defines HasAtoms and the concrete Atoms, which holds a collection of atoms with no cell or periodicity. Atoms is essentially a wrapper around a polars.DataFrame.

SchemaDict module-attribute

SchemaDict: TypeAlias = OrderedDict[str, DataType]

IntoExprColumn module-attribute

IntoExprColumn: TypeAlias = IntoExprColumn

IntoExpr module-attribute

IntoExpr: TypeAlias = IntoExpr

UniqueKeepStrategy module-attribute

UniqueKeepStrategy: TypeAlias = UniqueKeepStrategy

FillNullStrategy module-attribute

FillNullStrategy: TypeAlias = FillNullStrategy

RollingInterpolationMethod module-attribute

RollingInterpolationMethod: TypeAlias = (
    RollingInterpolationMethod
)

ConcatMethod module-attribute

ConcatMethod: TypeAlias = Literal[
    "horizontal", "vertical", "diagonal", "inner", "align"
]

IntoAtoms module-attribute

IntoAtoms: TypeAlias = Union[
    Dict[str, Sequence[Any]],
    Sequence[Any],
    ndarray,
    DataFrame,
    "Atoms",
]

A type convertible into an Atoms.

AtomSelection module-attribute

AtomSelection: TypeAlias = Union[
    IntoExprColumn,
    NDArray[bool_],
    ArrayLike,
    Mapping[str, Any],
]

Polars expression selecting a subset of atoms. Can be used with many Atoms methods.

AtomValues module-attribute

AtomValues: TypeAlias = Union[
    IntoExprColumn,
    NDArray[generic],
    ArrayLike,
    Mapping[str, Any],
]

Array, value, or polars expression mapping atom symbols to values. Can be used with with_* methods on Atoms

HasAtoms

Bases: ABC

Abstract class representing any (possibly compound) collection of atoms.

Source code in atomlib/atoms.py

class HasAtoms(abc.ABC):
    """Abstract class representing any (possibly compound) collection of atoms."""

    # abstract methods

    @abc.abstractmethod
    def get_atoms(self, frame: t.Literal['local'] = 'local') -> Atoms:
        """
        Get atoms contained in `self`. This should be a low cost method.

        Args:
          frame: Coordinate frame to return atoms in. For a plain [`HasAtoms`][atomlib.atoms.HasAtoms],
                 only `'local'` is supported.

        Return:
          The contained atoms
        """
        ...

    @abc.abstractmethod
    def with_atoms(self, atoms: HasAtoms, frame: t.Literal['local'] = 'local') -> Self:
        """
        Return a copy of self with the inner [`Atoms`][atomlib.atoms.Atoms] replaced.

        Args:
          atoms: [`HasAtoms`][atomlib.atoms.HasAtoms] to replace these with.
          frame: Coordinate frame inside atoms are in. For a plain [`HasAtoms`][atomlib.atoms.HasAtoms],
                 only `'local'` is supported.

        Return:
          A copy of `self` updated with the given atoms
        """
        ...

    @classmethod
    @abc.abstractmethod
    def _combine_metadata(cls: t.Type[HasAtomsT], *atoms: HasAtoms) -> HasAtomsT:
        """
        When combining multiple `HasAtoms`, check that they are compatible with each other,
        and return a 'representative' which best represents the combined metadata.
        Implementors should treat `Atoms` as acceptable, but having no metadata.
        """
        ...

    def _get_frame(self) -> polars.DataFrame:
        return self.get_atoms().inner

    # dataframe methods

    @property
    @_fwd_frame(lambda df: df.columns)
    def columns(self) -> t.List[str]:
        """
        Return the column names in `self`.

        Returns:
          A sequence of column names
        """
        ...

    @property
    @_fwd_frame(lambda df: df.dtypes)
    def dtypes(self) -> t.List[polars.DataType]:
        """
        Return the datatypes in `self`.

        Returns:
          A sequence of column [`DataType`][polars.datatypes.DataType]s
        """
        ...

    @property
    @_fwd_frame(lambda df: df.schema)  # type: ignore
    def schema(self) -> Schema:
        """
        Return the schema of `self`.

        Returns:
          A dictionary of column names and [`DataType`][polars.datatypes.DataType]s
        """
        ...

    @_fwd_frame(polars.DataFrame.describe)
    def describe(self, percentiles: t.Union[t.Sequence[float], float, None] = (0.25, 0.5, 0.75), *,
                 interpolation: RollingInterpolationMethod = 'nearest') -> 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_frame_map
    def with_columns(self,
                     *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
                     **named_exprs: IntoExpr) -> polars.DataFrame:
        """Return a copy of `self` with the given columns added."""
        return self._get_frame().with_columns(*exprs, **named_exprs)

    with_column = with_columns

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

    @_fwd_frame(lambda df, name: df.get_column(name))
    def get_column(self, name: str) -> 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_frame(polars.DataFrame.get_columns)
    def get_columns(self) -> 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_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."""
        ...

    @_fwd_frame(polars.DataFrame.group_by)
    def group_by(self, *by: t.Union[IntoExpr, t.Iterable[IntoExpr]], maintain_order: bool = False,
                 **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: HasAtomsT, function: t.Callable[Concatenate[HasAtomsT, P], T], *args: P.args, **kwargs: P.kwargs) -> T:
        """Apply `function` to `self` (in method-call syntax)."""
        return function(self, *args, **kwargs)

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

    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)

    # row-wise operations

    def filter(
        self,
        *predicates: t.Union[None, IntoExprColumn, t.Iterable[IntoExprColumn], bool, t.List[bool], numpy.ndarray],
        **constraints: t.Any,
    ) -> Self:
        """Filter `self`, removing rows which evaluate to `False`."""
        # TODO clean up
        preds_not_none = tuple(filter(lambda p: p is not None, predicates))
        if not len(preds_not_none) and not len(constraints):
            return self
        return self.with_atoms(Atoms(self._get_frame().filter(*preds_not_none, **constraints), _unchecked=True))  # type: ignore

    @_fwd_frame_map
    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,
    ) -> polars.DataFrame:
        """
        Sort the atoms in `self` by the given columns/expressions.
        """
        return self._get_frame().sort(
            by, *more_by, descending=descending, nulls_last=nulls_last
        )

    @_fwd_frame_map
    def slice(self, offset: int, length: t.Optional[int] = None) -> polars.DataFrame:
        """Return a slice of the rows in `self`."""
        return self._get_frame().slice(offset, length)

    @_fwd_frame_map
    def head(self, n: int = 5) -> polars.DataFrame:
        """Return the first `n` rows of `self`."""
        return self._get_frame().head(n)

    @_fwd_frame_map
    def tail(self, n: int = 5) -> polars.DataFrame:
        """Return the last `n` rows of `self`."""
        return self._get_frame().tail(n)

    @_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)

    @_fwd_frame_map
    def fill_null(
        self, value: t.Any = None, strategy: t.Optional[FillNullStrategy] = None,
        limit: t.Optional[int] = None, matches_supertype: bool = True,
    ) -> polars.DataFrame:
        """Fill null values in `self`, using the specified value or strategy."""
        return self._get_frame().fill_null(value, strategy, limit, matches_supertype=matches_supertype)

    @_fwd_frame_map
    def fill_nan(self, value: t.Union[polars.Expr, int, float, None]) -> polars.DataFrame:
        """Fill floating-point NaN values in `self`."""
        return self._get_frame().fill_nan(value)

    @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')

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

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

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

    # column-wise operations

    @_fwd_frame(polars.DataFrame.select)
    def select(
        self,
        *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
        **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
        """
        ...

    # some helpers we add

    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)

    def select_props(
        self,
        *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
        **named_exprs: IntoExpr
    ) -> Self:
        """
        Select `exprs` from `self`, while keeping required columns.

        Returns:
          A [`HasAtoms`][atomlib.atoms.HasAtoms] filtered to contain the
          specified properties (as well as required columns).
        """
        props = self._get_frame().lazy().select(*exprs, **named_exprs).drop(_REQUIRED_COLUMNS, strict=False).collect(_eager=True)
        return self.with_atoms(
            Atoms(self._get_frame().select(_REQUIRED_COLUMNS).hstack(props), _unchecked=False)
        )

    def try_select(
        self,
        *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
        **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:
            return self._get_frame().select(*exprs, **named_exprs)
        except polars.ColumnNotFoundError:
            return None

    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

    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)

    # dunders

    @_fwd_frame(polars.DataFrame.__len__)
    def __len__(self) -> int:
        """Return the number of atoms in `self`."""
        ...

    @_fwd_frame(polars.DataFrame.__contains__)
    def __contains__(self, key: str) -> bool:
        """Return whether `self` contains the given column."""
        ...

    def __add__(self, other: IntoAtoms) -> HasAtoms:
        return self.__class__.concat((self, other), how='inner')

    def __radd__(self, other: IntoAtoms) -> HasAtoms:
        return self.__class__.concat((other, self), how='inner')

    def __getitem__(self, column: str) -> polars.Series:
        try:
            return self.get_column(column)
        except polars.exceptions.ColumnNotFoundError:
            if column in ('x', 'y', 'z'):
                return self.select(_coord_expr(column)).to_series()
            raise

    @_fwd_frame(polars.DataFrame.__dataframe__)
    def __dataframe__(self, nan_as_null: bool = False, allow_copy: bool = True) -> polars.interchange.dataframe.PolarsDataFrame:
        ...

    # atoms-specific methods

    def bbox_atoms(self) -> BBox3D:
        """Return the bounding box of all the atoms in ``self``."""
        return BBox3D.from_pts(self.coords())

    bbox = bbox_atoms

    def transform_atoms(self, transform: IntoTransform3D, selection: t.Optional[AtomSelection] = None, *,
                        transform_velocities: bool = False) -> Self:
        """
        Transform the atoms in `self` by `transform`.
        If `selection` is given, only transform the atoms in `selection`.
        """
        transform = Transform3D.make(transform)
        selection = _selection_to_numpy(self, selection)
        transformed = self.with_coords(Transform3D.make(transform) @ self.coords(selection), selection)
        # try to transform velocities as well
        if transform_velocities and (velocities := self.velocities(selection)) is not None:
            return transformed.with_velocity(transform.transform_vec(velocities), selection)
        return transformed

    transform = transform_atoms

    def round_near_zero(self, tol: float = 1e-14) -> Self:
        """
        Round atom position values near zero to zero.
        """
        return self.with_columns(coords=polars.concat_list(
            polars.when(_coord_expr(col).abs() >= tol).then(_coord_expr(col)).otherwise(polars.lit(0.))
            for col in range(3)
        ).list.to_array(3))

    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) -> Self:
        """
        Crop, removing all atoms outside of the specified region, inclusive.
        """

        return self.filter(
            self.x().is_between(x_min, x_max, closed='both'),
            self.y().is_between(y_min, y_max, closed='both'),
            self.z().is_between(z_min, z_max, closed='both'),
        )

    crop_atoms = crop

    def _wrap(self, eps: float = 1e-5) -> Self:
        coords = (self.coords() + eps) % 1. - eps
        return self.with_coords(coords)

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

    unique = deduplicate

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

    # property getters and setters

    def coords(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Literal['local'] = 'local') -> NDArray[numpy.float64]:
        """Return a `(N, 3)` ndarray of atom coordinates (dtype [`numpy.float64`][numpy.float64])."""
        df = self if selection is None else self.filter(_selection_to_expr(self, selection))
        return df.get_column('coords').to_numpy().astype(numpy.float64)

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

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

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

    def velocities(self, selection: t.Optional[AtomSelection] = None) -> t.Optional[NDArray[numpy.float64]]:
        """Return a `(N, 3)` ndarray of atom velocities (dtype [`numpy.float64`][numpy.float64])."""
        if 'velocity' not in self:
            return None

        df = self if selection is None else self.filter(_selection_to_expr(self, selection))
        return df.get_column('velocity').to_numpy().astype(numpy.float64)

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

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

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

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

    def add_atom(self, elem: t.Union[int, str], /,
                 x: t.Union[ArrayLike, float],
                 y: t.Optional[float] = None,
                 z: t.Optional[float] = 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 [`HasAtoms.concat`][atomlib.atoms.HasAtoms.concat] instead).
        """
        if isinstance(elem, int):
            kwargs.update(elem=elem)
        else:
            kwargs.update(symbol=elem)
        if hasattr(x, '__len__') and len(x) > 1:  # type: ignore
            (x, y, z) = to_vec3(x)
        elif y is None or z is None:
            raise ValueError("Must specify vector of positions or x, y, & z.")

        sym = get_sym(elem) if isinstance(elem, int) else elem
        d: t.Dict[str, t.Any] = {'x': x, 'y': y, 'z': z, 'symbol': sym, **kwargs}
        return self.concat(
            (self, Atoms(d).select_schema(self.schema)),
            how='vertical'
        )

    @t.overload
    def pos(self, x: t.Sequence[t.Optional[float]], /, *,
            y: None = None, z: None = None,
            tol: float = 1e-6, **kwargs: t.Any) -> polars.Expr:
        ...

    @t.overload
    def pos(self, x: t.Optional[float] = None, y: t.Optional[float] = None, z: t.Optional[float] = None, *,
            tol: float = 1e-6, **kwargs: t.Any) -> polars.Expr:
        ...

    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

    def with_index(self, index: t.Optional[AtomValues] = 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.
        """
        if index is None and 'i' in self.columns:
            return self
        if index is None:
            index = numpy.arange(len(self), dtype=numpy.int64)
        return self.with_column(_values_to_expr(self, index, polars.Int64).alias('i'))

    def with_wobble(self, wobble: t.Optional[AtomValues] = 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.
        """
        if wobble is None and 'wobble' in self.columns:
            return self
        wobble = 0. if wobble is None else wobble
        return self.with_column(_values_to_expr(self, wobble, polars.Float64).alias('wobble'))

    def with_occupancy(self, frac_occupancy: t.Optional[AtomValues] = 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.
        """
        if frac_occupancy is None and 'frac_occupancy' in self.columns:
            return self
        frac_occupancy = 1. if frac_occupancy is None else frac_occupancy
        return self.with_column(_values_to_expr(self, frac_occupancy, polars.Float64).alias('frac_occupancy'))

    def apply_wobble(self, rng: t.Union[numpy.random.Generator, int, None] = 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.
        """
        if 'wobble' not in self.columns:
            return self
        rng = numpy.random.default_rng(seed=rng)

        stddev = self.select((polars.col('wobble') / 3.).sqrt()).to_series().to_numpy()
        coords = self.coords()
        coords += stddev[:, None] * rng.standard_normal(coords.shape)
        return self.with_coords(coords)

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

    def with_type(self, types: t.Optional[AtomValues] = 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]`
        """
        if types is not None:
            return self.with_columns(type=_values_to_expr(self, types, polars.Int32))
        if 'type' in self.columns:
            return self

        unique = Atoms(self._get_frame().unique(maintain_order=False, subset=['elem', 'symbol']).sort(['elem', 'symbol']), _unchecked=True)
        new = self.with_column(polars.Series('type', values=numpy.zeros(len(self)), dtype=polars.Int32))

        logging.warning("Auto-assigning element types")
        for (i, (elem, sym)) in enumerate(unique.select(('elem', 'symbol')).rows()):
            print(f"Assigning type {i+1} to element '{sym}'")
            new = new.with_column(polars.when((polars.col('elem') == elem) & (polars.col('symbol') == sym))
                                        .then(polars.lit(i+1))
                                        .otherwise(polars.col('type'))
                                        .alias('type'))

        assert (new.get_column('type') == 0).sum() == 0
        return new

    def with_mass(self, mass: t.Optional[ArrayLike] = 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.
        """
        if mass is not None:
            return self.with_column(_values_to_expr(self, mass, polars.Float32).alias('mass'))
        if 'mass' in self.columns:
            return self

        unique_elems = self.get_column('elem').unique()
        new = self.with_column(polars.Series('mass', values=numpy.zeros(len(self)), dtype=polars.Float32))

        logging.warning("Auto-assigning element masses")
        for elem in unique_elems:
            new = new.with_column(polars.when(polars.col('elem') == elem)
                                        .then(polars.lit(get_mass(elem)))
                                        .otherwise(polars.col('mass'))
                                        .alias('mass'))

        assert (new.get_column('mass').abs() < 1e-10).sum() == 0
        return new

    def with_symbol(self, symbols: ArrayLike, selection: t.Optional[AtomSelection] = None) -> Self:
        """
        Return `self` with the given atomic symbols.
        """
        if selection is not None:
            selection = _selection_to_numpy(self, selection)
            new_symbols = self.get_column('symbol')
            new_symbols[selection] = polars.Series(list(numpy.broadcast_to(symbols, len(selection))), dtype=polars.Utf8)
            symbols = new_symbols

        # TODO better cast here
        symbols = polars.Series('symbol', list(numpy.broadcast_to(symbols, len(self))), dtype=polars.Utf8)
        return self.with_columns((symbols, get_elem(symbols)))

    def with_coords(self, pts: ArrayLike, selection: t.Optional[AtomSelection] = None, *, frame: t.Literal['local'] = 'local') -> Self:
        """
        Return `self` replaced with the given atomic positions.
        """
        if selection is not None:
            selection = _selection_to_numpy(self, selection)
            new_pts = self.coords()
            pts = numpy.atleast_2d(pts)
            assert pts.shape[-1] == 3
            new_pts[selection] = pts
            pts = new_pts

        # https://github.com/pola-rs/polars/issues/18369
        pts = numpy.broadcast_to(pts, (len(self), 3)) if len(self) else []
        return self.with_columns(polars.Series('coords', pts, polars.Array(polars.Float64, 3)))

    def with_velocity(self, pts: t.Optional[ArrayLike] = None,
                      selection: t.Optional[AtomSelection] = None) -> Self:
        """
        Return `self` replaced with the given atomic velocities.
        If `pts` is not specified, use the already existing velocities or zero.
        """
        if pts is None:
            if 'velocity' in self:
                return self
            all_pts = numpy.zeros((len(self), 3))
        else:
            all_pts = self['velocity'].to_numpy()

        if selection is None:
            all_pts = pts or all_pts
        elif pts is not None:
            selection = _selection_to_numpy(self, selection)
            all_pts = numpy.require(all_pts, requirements=['WRITEABLE'])
            pts = numpy.atleast_2d(pts)
            assert pts.shape[-1] == 3
            all_pts[selection] = pts

        all_pts = numpy.broadcast_to(all_pts, (len(self), 3))
        return self.with_columns(polars.Series('velocity', all_pts, polars.Array(polars.Float64, 3)))

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

bbox class-attribute instance-attribute

bbox = bbox_atoms

transform class-attribute instance-attribute

transform = transform_atoms

crop_atoms class-attribute instance-attribute

crop_atoms = crop

unique class-attribute instance-attribute

unique = deduplicate

get_atoms abstractmethod

get_atoms(frame: Literal['local'] = 'local') -> Atoms

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

PARAMETER	DESCRIPTION
frame	Coordinate frame to return atoms in. For a plain HasAtoms, only 'local' is supported. TYPE: Literal['local'] DEFAULT: 'local'

Return

The contained atoms

Source code in atomlib/atoms.py

@abc.abstractmethod
def get_atoms(self, frame: t.Literal['local'] = 'local') -> Atoms:
    """
    Get atoms contained in `self`. This should be a low cost method.

    Args:
      frame: Coordinate frame to return atoms in. For a plain [`HasAtoms`][atomlib.atoms.HasAtoms],
             only `'local'` is supported.

    Return:
      The contained atoms
    """
    ...

with_atoms abstractmethod

with_atoms(
    atoms: HasAtoms, frame: Literal["local"] = "local"
) -> Self

Return a copy of self with the inner Atoms replaced.

PARAMETER	DESCRIPTION
atoms	HasAtoms to replace these with. TYPE: HasAtoms
frame	Coordinate frame inside atoms are in. For a plain HasAtoms, only 'local' is supported. TYPE: Literal['local'] DEFAULT: 'local'

Return

A copy of self updated with the given atoms

Source code in atomlib/atoms.py

@abc.abstractmethod
def with_atoms(self, atoms: HasAtoms, frame: t.Literal['local'] = 'local') -> Self:
    """
    Return a copy of self with the inner [`Atoms`][atomlib.atoms.Atoms] replaced.

    Args:
      atoms: [`HasAtoms`][atomlib.atoms.HasAtoms] to replace these with.
      frame: Coordinate frame inside atoms are in. For a plain [`HasAtoms`][atomlib.atoms.HasAtoms],
             only `'local'` is supported.

    Return:
      A copy of `self` updated with the given atoms
    """
    ...

describe

describe(
    percentiles: Union[Sequence[float], float, None] = (
        0.25,
        0.5,
        0.75,
    ),
    *,
    interpolation: RollingInterpolationMethod = "nearest"
) -> 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/atoms.py

@_fwd_frame(polars.DataFrame.describe)
def describe(self, percentiles: t.Union[t.Sequence[float], float, None] = (0.25, 0.5, 0.75), *,
             interpolation: RollingInterpolationMethod = 'nearest') -> 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]],
    **named_exprs: IntoExpr
) -> DataFrame

Return a copy of self with the given columns added.

Source code in atomlib/atoms.py

@_fwd_frame_map
def with_columns(self,
                 *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
                 **named_exprs: IntoExpr) -> polars.DataFrame:
    """Return a copy of `self` with the given columns added."""
    return self._get_frame().with_columns(*exprs, **named_exprs)

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

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

Source code in atomlib/atoms.py

@_fwd_frame(lambda df, name: df.get_column(name))
def get_column(self, name: str) -> 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() -> List[Series]

Return all columns from self as a list of Series.

Source code in atomlib/atoms.py

@_fwd_frame(polars.DataFrame.get_columns)
def get_columns(self) -> 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,
    **named_by: IntoExpr
) -> GroupBy

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

Source code in atomlib/atoms.py

@_fwd_frame(polars.DataFrame.group_by)
def group_by(self, *by: t.Union[IntoExpr, t.Iterable[IntoExpr]], maintain_order: bool = False,
             **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[HasAtomsT, P], T],
    *args: args,
    **kwargs: kwargs
) -> T

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

Source code in atomlib/atoms.py

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

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)

filter

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

Filter self, removing rows which evaluate to False.

Source code in atomlib/atoms.py

def filter(
    self,
    *predicates: t.Union[None, IntoExprColumn, t.Iterable[IntoExprColumn], bool, t.List[bool], numpy.ndarray],
    **constraints: t.Any,
) -> Self:
    """Filter `self`, removing rows which evaluate to `False`."""
    # TODO clean up
    preds_not_none = tuple(filter(lambda p: p is not None, predicates))
    if not len(preds_not_none) and not len(constraints):
        return self
    return self.with_atoms(Atoms(self._get_frame().filter(*preds_not_none, **constraints), _unchecked=True))  # type: ignore

sort

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

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

Source code in atomlib/atoms.py

@_fwd_frame_map
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,
) -> polars.DataFrame:
    """
    Sort the atoms in `self` by the given columns/expressions.
    """
    return self._get_frame().sort(
        by, *more_by, descending=descending, nulls_last=nulls_last
    )

slice

slice(
    offset: int, length: Optional[int] = None
) -> DataFrame

Return a slice of the rows in self.

Source code in atomlib/atoms.py

@_fwd_frame_map
def slice(self, offset: int, length: t.Optional[int] = None) -> polars.DataFrame:
    """Return a slice of the rows in `self`."""
    return self._get_frame().slice(offset, length)

head

head(n: int = 5) -> DataFrame

Return the first n rows of self.

Source code in atomlib/atoms.py

@_fwd_frame_map
def head(self, n: int = 5) -> polars.DataFrame:
    """Return the first `n` rows of `self`."""
    return self._get_frame().head(n)

tail

tail(n: int = 5) -> DataFrame

Return the last n rows of self.

Source code in atomlib/atoms.py

@_fwd_frame_map
def tail(self, n: int = 5) -> polars.DataFrame:
    """Return the last `n` rows of `self`."""
    return self._get_frame().tail(n)

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,
) -> DataFrame

Fill null values in self, using the specified value or strategy.

Source code in atomlib/atoms.py

@_fwd_frame_map
def fill_null(
    self, value: t.Any = None, strategy: t.Optional[FillNullStrategy] = None,
    limit: t.Optional[int] = None, matches_supertype: bool = True,
) -> polars.DataFrame:
    """Fill null values in `self`, using the specified value or strategy."""
    return self._get_frame().fill_null(value, strategy, limit, matches_supertype=matches_supertype)

fill_nan

fill_nan(value: Union[Expr, int, float, None]) -> DataFrame

Fill floating-point NaN values in self.

Source code in atomlib/atoms.py

@_fwd_frame_map
def fill_nan(self, value: t.Union[polars.Expr, int, float, None]) -> polars.DataFrame:
    """Fill floating-point NaN values in `self`."""
    return self._get_frame().fill_nan(value)

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]],
    **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/atoms.py

@_fwd_frame(polars.DataFrame.select)
def select(
    self,
    *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    **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]],
    **named_exprs: IntoExpr
) -> Self

Select exprs from self, while keeping required columns.

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

Source code in atomlib/atoms.py

def select_props(
    self,
    *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    **named_exprs: IntoExpr
) -> Self:
    """
    Select `exprs` from `self`, while keeping required columns.

    Returns:
      A [`HasAtoms`][atomlib.atoms.HasAtoms] filtered to contain the
      specified properties (as well as required columns).
    """
    props = self._get_frame().lazy().select(*exprs, **named_exprs).drop(_REQUIRED_COLUMNS, strict=False).collect(_eager=True)
    return self.with_atoms(
        Atoms(self._get_frame().select(_REQUIRED_COLUMNS).hstack(props), _unchecked=False)
    )

try_select

try_select(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    **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/atoms.py

def try_select(
    self,
    *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    **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:
        return self._get_frame().select(*exprs, **named_exprs)
    except polars.ColumnNotFoundError:
        return None

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

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)

bbox_atoms

bbox_atoms() -> BBox3D

Return the bounding box of all the atoms in self.

Source code in atomlib/atoms.py

def bbox_atoms(self) -> BBox3D:
    """Return the bounding box of all the atoms in ``self``."""
    return BBox3D.from_pts(self.coords())

transform_atoms

transform_atoms(
    transform: IntoTransform3D,
    selection: Optional[AtomSelection] = None,
    *,
    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/atoms.py

def transform_atoms(self, transform: IntoTransform3D, selection: t.Optional[AtomSelection] = None, *,
                    transform_velocities: bool = False) -> Self:
    """
    Transform the atoms in `self` by `transform`.
    If `selection` is given, only transform the atoms in `selection`.
    """
    transform = Transform3D.make(transform)
    selection = _selection_to_numpy(self, selection)
    transformed = self.with_coords(Transform3D.make(transform) @ self.coords(selection), selection)
    # try to transform velocities as well
    if transform_velocities and (velocities := self.velocities(selection)) is not None:
        return transformed.with_velocity(transform.transform_vec(velocities), selection)
    return transformed

round_near_zero

round_near_zero(tol: float = 1e-14) -> Self

Round atom position values near zero to zero.

Source code in atomlib/atoms.py

def round_near_zero(self, tol: float = 1e-14) -> Self:
    """
    Round atom position values near zero to zero.
    """
    return self.with_columns(coords=polars.concat_list(
        polars.when(_coord_expr(col).abs() >= tol).then(_coord_expr(col)).otherwise(polars.lit(0.))
        for col in range(3)
    ).list.to_array(3))

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,
) -> Self

Crop, removing all atoms outside of the specified region, inclusive.

Source code in atomlib/atoms.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) -> Self:
    """
    Crop, removing all atoms outside of the specified region, inclusive.
    """

    return self.filter(
        self.x().is_between(x_min, x_max, closed='both'),
        self.y().is_between(y_min, y_max, closed='both'),
        self.z().is_between(z_min, z_max, closed='both'),
    )

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: Literal["local"] = "local"
) -> NDArray[float64]

Return a (N, 3) ndarray of atom coordinates (dtype numpy.float64).

Source code in atomlib/atoms.py

def coords(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Literal['local'] = 'local') -> NDArray[numpy.float64]:
    """Return a `(N, 3)` ndarray of atom coordinates (dtype [`numpy.float64`][numpy.float64])."""
    df = self if selection is None else self.filter(_selection_to_expr(self, selection))
    return df.get_column('coords').to_numpy().astype(numpy.float64)

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,
) -> Optional[NDArray[float64]]

Return a (N, 3) ndarray of atom velocities (dtype numpy.float64).

Source code in atomlib/atoms.py

def velocities(self, selection: t.Optional[AtomSelection] = None) -> t.Optional[NDArray[numpy.float64]]:
    """Return a `(N, 3)` ndarray of atom velocities (dtype [`numpy.float64`][numpy.float64])."""
    if 'velocity' not in self:
        return None

    df = self if selection is None else self.filter(_selection_to_expr(self, selection))
    return df.get_column('velocity').to_numpy().astype(numpy.float64)

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,
    **kwargs: Any,
) -> Self

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

add_atom(
    elem: Union[int, str],
    /,
    x: Union[ArrayLike, float],
    y: Optional[float] = None,
    z: Optional[float] = 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 HasAtoms.concat instead).

Source code in atomlib/atoms.py

def add_atom(self, elem: t.Union[int, str], /,
             x: t.Union[ArrayLike, float],
             y: t.Optional[float] = None,
             z: t.Optional[float] = 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 [`HasAtoms.concat`][atomlib.atoms.HasAtoms.concat] instead).
    """
    if isinstance(elem, int):
        kwargs.update(elem=elem)
    else:
        kwargs.update(symbol=elem)
    if hasattr(x, '__len__') and len(x) > 1:  # type: ignore
        (x, y, z) = to_vec3(x)
    elif y is None or z is None:
        raise ValueError("Must specify vector of positions or x, y, & z.")

    sym = get_sym(elem) if isinstance(elem, int) else elem
    d: t.Dict[str, t.Any] = {'x': x, 'y': y, 'z': z, 'symbol': sym, **kwargs}
    return self.concat(
        (self, Atoms(d).select_schema(self.schema)),
        how='vertical'
    )

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) -> 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/atoms.py

def with_index(self, index: t.Optional[AtomValues] = 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.
    """
    if index is None and 'i' in self.columns:
        return self
    if index is None:
        index = numpy.arange(len(self), dtype=numpy.int64)
    return self.with_column(_values_to_expr(self, index, polars.Int64).alias('i'))

with_wobble

with_wobble(wobble: Optional[AtomValues] = 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/atoms.py

def with_wobble(self, wobble: t.Optional[AtomValues] = 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.
    """
    if wobble is None and 'wobble' in self.columns:
        return self
    wobble = 0. if wobble is None else wobble
    return self.with_column(_values_to_expr(self, wobble, polars.Float64).alias('wobble'))

with_occupancy

with_occupancy(
    frac_occupancy: Optional[AtomValues] = 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/atoms.py

def with_occupancy(self, frac_occupancy: t.Optional[AtomValues] = 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.
    """
    if frac_occupancy is None and 'frac_occupancy' in self.columns:
        return self
    frac_occupancy = 1. if frac_occupancy is None else frac_occupancy
    return self.with_column(_values_to_expr(self, frac_occupancy, polars.Float64).alias('frac_occupancy'))

apply_wobble

apply_wobble(
    rng: Union[Generator, int, None] = 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/atoms.py

def apply_wobble(self, rng: t.Union[numpy.random.Generator, int, None] = 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.
    """
    if 'wobble' not in self.columns:
        return self
    rng = numpy.random.default_rng(seed=rng)

    stddev = self.select((polars.col('wobble') / 3.).sqrt()).to_series().to_numpy()
    coords = self.coords()
    coords += stddev[:, None] * rng.standard_normal(coords.shape)
    return self.with_coords(coords)

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) -> 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/atoms.py

def with_type(self, types: t.Optional[AtomValues] = 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]`
    """
    if types is not None:
        return self.with_columns(type=_values_to_expr(self, types, polars.Int32))
    if 'type' in self.columns:
        return self

    unique = Atoms(self._get_frame().unique(maintain_order=False, subset=['elem', 'symbol']).sort(['elem', 'symbol']), _unchecked=True)
    new = self.with_column(polars.Series('type', values=numpy.zeros(len(self)), dtype=polars.Int32))

    logging.warning("Auto-assigning element types")
    for (i, (elem, sym)) in enumerate(unique.select(('elem', 'symbol')).rows()):
        print(f"Assigning type {i+1} to element '{sym}'")
        new = new.with_column(polars.when((polars.col('elem') == elem) & (polars.col('symbol') == sym))
                                    .then(polars.lit(i+1))
                                    .otherwise(polars.col('type'))
                                    .alias('type'))

    assert (new.get_column('type') == 0).sum() == 0
    return new

with_mass

with_mass(mass: Optional[ArrayLike] = 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/atoms.py

def with_mass(self, mass: t.Optional[ArrayLike] = 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.
    """
    if mass is not None:
        return self.with_column(_values_to_expr(self, mass, polars.Float32).alias('mass'))
    if 'mass' in self.columns:
        return self

    unique_elems = self.get_column('elem').unique()
    new = self.with_column(polars.Series('mass', values=numpy.zeros(len(self)), dtype=polars.Float32))

    logging.warning("Auto-assigning element masses")
    for elem in unique_elems:
        new = new.with_column(polars.when(polars.col('elem') == elem)
                                    .then(polars.lit(get_mass(elem)))
                                    .otherwise(polars.col('mass'))
                                    .alias('mass'))

    assert (new.get_column('mass').abs() < 1e-10).sum() == 0
    return new

with_symbol

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

Return self with the given atomic symbols.

Source code in atomlib/atoms.py

def with_symbol(self, symbols: ArrayLike, selection: t.Optional[AtomSelection] = None) -> Self:
    """
    Return `self` with the given atomic symbols.
    """
    if selection is not None:
        selection = _selection_to_numpy(self, selection)
        new_symbols = self.get_column('symbol')
        new_symbols[selection] = polars.Series(list(numpy.broadcast_to(symbols, len(selection))), dtype=polars.Utf8)
        symbols = new_symbols

    # TODO better cast here
    symbols = polars.Series('symbol', list(numpy.broadcast_to(symbols, len(self))), dtype=polars.Utf8)
    return self.with_columns((symbols, get_elem(symbols)))

with_coords

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

Return self replaced with the given atomic positions.

Source code in atomlib/atoms.py

def with_coords(self, pts: ArrayLike, selection: t.Optional[AtomSelection] = None, *, frame: t.Literal['local'] = 'local') -> Self:
    """
    Return `self` replaced with the given atomic positions.
    """
    if selection is not None:
        selection = _selection_to_numpy(self, selection)
        new_pts = self.coords()
        pts = numpy.atleast_2d(pts)
        assert pts.shape[-1] == 3
        new_pts[selection] = pts
        pts = new_pts

    # https://github.com/pola-rs/polars/issues/18369
    pts = numpy.broadcast_to(pts, (len(self), 3)) if len(self) else []
    return self.with_columns(polars.Series('coords', pts, polars.Array(polars.Float64, 3)))

with_velocity

with_velocity(
    pts: Optional[ArrayLike] = None,
    selection: Optional[AtomSelection] = 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/atoms.py

def with_velocity(self, pts: t.Optional[ArrayLike] = None,
                  selection: t.Optional[AtomSelection] = None) -> Self:
    """
    Return `self` replaced with the given atomic velocities.
    If `pts` is not specified, use the already existing velocities or zero.
    """
    if pts is None:
        if 'velocity' in self:
            return self
        all_pts = numpy.zeros((len(self), 3))
    else:
        all_pts = self['velocity'].to_numpy()

    if selection is None:
        all_pts = pts or all_pts
    elif pts is not None:
        selection = _selection_to_numpy(self, selection)
        all_pts = numpy.require(all_pts, requirements=['WRITEABLE'])
        pts = numpy.atleast_2d(pts)
        assert pts.shape[-1] == 3
        all_pts[selection] = pts

    all_pts = numpy.broadcast_to(all_pts, (len(self), 3))
    return self.with_columns(polars.Series('velocity', all_pts, polars.Array(polars.Float64, 3)))

Atoms

Bases: AtomsIOMixin, HasAtoms

A collection of atoms, absent any implied coordinate system. Implemented as a wrapper around a polars.DataFrame.

Must contain the following columns:

• coords: array of [x, y, z] positions, float
• elem: atomic number, int
• symbol: atomic symbol (may contain charges)

In addition, it commonly contains the following columns:

• i: Initial atom number
• wobble: Isotropic Debye-Waller mean-squared deviation (\(\left<u^2\right> = B \cdot \frac{3}{8 \pi^2}\), dimensions of [Length^2])
• frac_occupancy: Fractional occupancy, in the range [0., 1.]
• mass: Atomic mass, in g/mol (approx. Da)
• velocity: array of [x, y, z] velocities, float, dimensions of length/time
• type: Numeric atom type, as used by programs like LAMMPS

Source code in atomlib/atoms.py

class Atoms(AtomsIOMixin, HasAtoms):
    r"""
    A collection of atoms, absent any implied coordinate system.
    Implemented as a wrapper around a [`polars.DataFrame`][polars.DataFrame].

    Must contain the following columns:

    - coords: array of `[x, y, z]` positions, float
    - elem: atomic number, int
    - symbol: atomic symbol (may contain charges)

    In addition, it commonly contains the following columns:

    - i: Initial atom number
    - wobble: Isotropic Debye-Waller mean-squared deviation ($\left<u^2\right> = B \cdot \frac{3}{8 \pi^2}$, dimensions of [Length^2])
    - frac_occupancy: Fractional occupancy, in the range [0., 1.]
    - mass: Atomic mass, in g/mol (approx. Da)
    - velocity: array of `[x, y, z]` velocities, float, dimensions of length/time
    - type: Numeric atom type, as used by programs like LAMMPS

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

    def __init__(self, data: t.Optional[IntoAtoms] = None, columns: t.Optional[t.Sequence[str]] = None,
                 orient: t.Union[t.Literal['row'], t.Literal['col'], None] = None,
                 _unchecked: bool = False):
        self._bbox: t.Optional[BBox3D] = None
        self.inner: polars.DataFrame

        if data is None:
            assert columns is None
            self.inner = polars.DataFrame([
                polars.Series('coords', (), dtype=polars.Array(polars.Float64, 3)),
                polars.Series('elem', (), dtype=polars.Int8),
                polars.Series('symbol', (), dtype=polars.Utf8),
            ])
        elif isinstance(data, polars.DataFrame):
            self.inner = data
        elif isinstance(data, Atoms):
            self.inner = data.inner
            _unchecked = True
        else:
            self.inner = polars.DataFrame(data, schema=columns, orient=orient)

        if not _unchecked:
            # stack ('x', 'y', 'z') -> 'coords'
            self.inner = _with_columns_stacked(self.inner, ('x', 'y', 'z'), 'coords')
            self.inner = _with_columns_stacked(self.inner, ('v_x', 'v_y', 'v_z'), 'velocity')

            missing: t.Tuple[str, ...] = tuple(set(['symbol', 'elem']) - set(self.columns))
            if len(missing) > 1:
                raise ValueError("'Atoms' missing columns 'elem' and/or 'symbol'.")
            # fill 'symbol' from 'elem' or vice-versa
            if missing == ('symbol',):
                self.inner = self.inner.with_columns(get_sym(self.inner['elem']))
            elif missing == ('elem',):
                # by convention, add before 'symbol' column
                self.inner = self.inner.insert_column(
                    self.inner.get_column_index('symbol'),
                    get_elem(self.inner['symbol']),
                )

            # cast to standard dtypes
            self.inner = self.inner.with_columns([
                self.inner[col].cast(dtype)
                for (col, dtype) in _COLUMN_DTYPES.items() if col in self.inner
            ])

            self._validate_atoms()

    @staticmethod
    def empty() -> Atoms:
        """
        Return an empty Atoms with only the mandatory columns.
        """
        return Atoms()

    def _validate_atoms(self):
        missing = [col for col in _REQUIRED_COLUMNS if col not in self.columns]
        if len(missing):
            raise ValueError(f"'Atoms' missing column(s) {', '.join(map(repr, missing))}")

    def get_atoms(self, frame: t.Literal['local'] = 'local') -> Atoms:
        if frame != 'local':
            raise ValueError(f"Atoms without a cell only support the 'local' coordinate frame, not '{frame}'.")
        return self

    def with_atoms(self, atoms: HasAtoms, frame: t.Literal['local'] = 'local') -> Atoms:
        if frame != 'local':
            raise ValueError(f"Atoms without a cell only support the 'local' coordinate frame, not '{frame}'.")
        return atoms.get_atoms()

    @classmethod
    def _combine_metadata(cls: t.Type[Atoms], *atoms: HasAtoms) -> Atoms:
        return cls.empty()

    def bbox(self) -> BBox3D:
        """Return the bounding box of all the points in `self`."""
        if self._bbox is None:
            self._bbox = BBox3D.from_pts(self.coords())

        return self._bbox

    def __str__(self) -> str:
        return f"Atoms, {self.inner!s}"

    def __repr__(self) -> str:
        buf = StringIO()
        buf.write("Atoms([\n")

        for series in self.inner.to_dict().values():
            buf.write(f"    Series({series.name!r}, {series.to_list()!r}, dtype={series.dtype!r}),\n")

        buf.write("])\n")
        return buf.getvalue()

    def _repr_pretty_(self, p: t.Any, cycle: bool) -> None:
        p.text('Atoms(...)') if cycle else p.text(str(self))

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

transform class-attribute instance-attribute

transform = transform_atoms

crop_atoms class-attribute instance-attribute

crop_atoms = crop

unique class-attribute instance-attribute

unique = deduplicate

inner instance-attribute

inner: DataFrame

describe

describe(
    percentiles: Union[Sequence[float], float, None] = (
        0.25,
        0.5,
        0.75,
    ),
    *,
    interpolation: RollingInterpolationMethod = "nearest"
) -> 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/atoms.py

@_fwd_frame(polars.DataFrame.describe)
def describe(self, percentiles: t.Union[t.Sequence[float], float, None] = (0.25, 0.5, 0.75), *,
             interpolation: RollingInterpolationMethod = 'nearest') -> 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]],
    **named_exprs: IntoExpr
) -> DataFrame

Return a copy of self with the given columns added.

Source code in atomlib/atoms.py

@_fwd_frame_map
def with_columns(self,
                 *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
                 **named_exprs: IntoExpr) -> polars.DataFrame:
    """Return a copy of `self` with the given columns added."""
    return self._get_frame().with_columns(*exprs, **named_exprs)

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

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

Source code in atomlib/atoms.py

@_fwd_frame(lambda df, name: df.get_column(name))
def get_column(self, name: str) -> 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() -> List[Series]

Return all columns from self as a list of Series.

Source code in atomlib/atoms.py

@_fwd_frame(polars.DataFrame.get_columns)
def get_columns(self) -> 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,
    **named_by: IntoExpr
) -> GroupBy

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

Source code in atomlib/atoms.py

@_fwd_frame(polars.DataFrame.group_by)
def group_by(self, *by: t.Union[IntoExpr, t.Iterable[IntoExpr]], maintain_order: bool = False,
             **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[HasAtomsT, P], T],
    *args: args,
    **kwargs: kwargs
) -> T

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

Source code in atomlib/atoms.py

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

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)

filter

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

Filter self, removing rows which evaluate to False.

Source code in atomlib/atoms.py

def filter(
    self,
    *predicates: t.Union[None, IntoExprColumn, t.Iterable[IntoExprColumn], bool, t.List[bool], numpy.ndarray],
    **constraints: t.Any,
) -> Self:
    """Filter `self`, removing rows which evaluate to `False`."""
    # TODO clean up
    preds_not_none = tuple(filter(lambda p: p is not None, predicates))
    if not len(preds_not_none) and not len(constraints):
        return self
    return self.with_atoms(Atoms(self._get_frame().filter(*preds_not_none, **constraints), _unchecked=True))  # type: ignore

sort

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

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

Source code in atomlib/atoms.py

@_fwd_frame_map
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,
) -> polars.DataFrame:
    """
    Sort the atoms in `self` by the given columns/expressions.
    """
    return self._get_frame().sort(
        by, *more_by, descending=descending, nulls_last=nulls_last
    )

slice

slice(
    offset: int, length: Optional[int] = None
) -> DataFrame

Return a slice of the rows in self.

Source code in atomlib/atoms.py

@_fwd_frame_map
def slice(self, offset: int, length: t.Optional[int] = None) -> polars.DataFrame:
    """Return a slice of the rows in `self`."""
    return self._get_frame().slice(offset, length)

head

head(n: int = 5) -> DataFrame

Return the first n rows of self.

Source code in atomlib/atoms.py

@_fwd_frame_map
def head(self, n: int = 5) -> polars.DataFrame:
    """Return the first `n` rows of `self`."""
    return self._get_frame().head(n)

tail

tail(n: int = 5) -> DataFrame

Return the last n rows of self.

Source code in atomlib/atoms.py

@_fwd_frame_map
def tail(self, n: int = 5) -> polars.DataFrame:
    """Return the last `n` rows of `self`."""
    return self._get_frame().tail(n)

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,
) -> DataFrame

Fill null values in self, using the specified value or strategy.

Source code in atomlib/atoms.py

@_fwd_frame_map
def fill_null(
    self, value: t.Any = None, strategy: t.Optional[FillNullStrategy] = None,
    limit: t.Optional[int] = None, matches_supertype: bool = True,
) -> polars.DataFrame:
    """Fill null values in `self`, using the specified value or strategy."""
    return self._get_frame().fill_null(value, strategy, limit, matches_supertype=matches_supertype)

fill_nan

fill_nan(value: Union[Expr, int, float, None]) -> DataFrame

Fill floating-point NaN values in self.

Source code in atomlib/atoms.py

@_fwd_frame_map
def fill_nan(self, value: t.Union[polars.Expr, int, float, None]) -> polars.DataFrame:
    """Fill floating-point NaN values in `self`."""
    return self._get_frame().fill_nan(value)

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]],
    **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/atoms.py

@_fwd_frame(polars.DataFrame.select)
def select(
    self,
    *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    **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]],
    **named_exprs: IntoExpr
) -> Self

Select exprs from self, while keeping required columns.

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

Source code in atomlib/atoms.py

def select_props(
    self,
    *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    **named_exprs: IntoExpr
) -> Self:
    """
    Select `exprs` from `self`, while keeping required columns.

    Returns:
      A [`HasAtoms`][atomlib.atoms.HasAtoms] filtered to contain the
      specified properties (as well as required columns).
    """
    props = self._get_frame().lazy().select(*exprs, **named_exprs).drop(_REQUIRED_COLUMNS, strict=False).collect(_eager=True)
    return self.with_atoms(
        Atoms(self._get_frame().select(_REQUIRED_COLUMNS).hstack(props), _unchecked=False)
    )

try_select

try_select(
    *exprs: Union[IntoExpr, Iterable[IntoExpr]],
    **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/atoms.py

def try_select(
    self,
    *exprs: t.Union[IntoExpr, t.Iterable[IntoExpr]],
    **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:
        return self._get_frame().select(*exprs, **named_exprs)
    except polars.ColumnNotFoundError:
        return None

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

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)

bbox_atoms

bbox_atoms() -> BBox3D

Return the bounding box of all the atoms in self.

Source code in atomlib/atoms.py

def bbox_atoms(self) -> BBox3D:
    """Return the bounding box of all the atoms in ``self``."""
    return BBox3D.from_pts(self.coords())

transform_atoms

transform_atoms(
    transform: IntoTransform3D,
    selection: Optional[AtomSelection] = None,
    *,
    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/atoms.py

def transform_atoms(self, transform: IntoTransform3D, selection: t.Optional[AtomSelection] = None, *,
                    transform_velocities: bool = False) -> Self:
    """
    Transform the atoms in `self` by `transform`.
    If `selection` is given, only transform the atoms in `selection`.
    """
    transform = Transform3D.make(transform)
    selection = _selection_to_numpy(self, selection)
    transformed = self.with_coords(Transform3D.make(transform) @ self.coords(selection), selection)
    # try to transform velocities as well
    if transform_velocities and (velocities := self.velocities(selection)) is not None:
        return transformed.with_velocity(transform.transform_vec(velocities), selection)
    return transformed

round_near_zero

round_near_zero(tol: float = 1e-14) -> Self

Round atom position values near zero to zero.

Source code in atomlib/atoms.py

def round_near_zero(self, tol: float = 1e-14) -> Self:
    """
    Round atom position values near zero to zero.
    """
    return self.with_columns(coords=polars.concat_list(
        polars.when(_coord_expr(col).abs() >= tol).then(_coord_expr(col)).otherwise(polars.lit(0.))
        for col in range(3)
    ).list.to_array(3))

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,
) -> Self

Crop, removing all atoms outside of the specified region, inclusive.

Source code in atomlib/atoms.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) -> Self:
    """
    Crop, removing all atoms outside of the specified region, inclusive.
    """

    return self.filter(
        self.x().is_between(x_min, x_max, closed='both'),
        self.y().is_between(y_min, y_max, closed='both'),
        self.z().is_between(z_min, z_max, closed='both'),
    )

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: Literal["local"] = "local"
) -> NDArray[float64]

Return a (N, 3) ndarray of atom coordinates (dtype numpy.float64).

Source code in atomlib/atoms.py

def coords(self, selection: t.Optional[AtomSelection] = None, *, frame: t.Literal['local'] = 'local') -> NDArray[numpy.float64]:
    """Return a `(N, 3)` ndarray of atom coordinates (dtype [`numpy.float64`][numpy.float64])."""
    df = self if selection is None else self.filter(_selection_to_expr(self, selection))
    return df.get_column('coords').to_numpy().astype(numpy.float64)

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,
) -> Optional[NDArray[float64]]

Return a (N, 3) ndarray of atom velocities (dtype numpy.float64).

Source code in atomlib/atoms.py

def velocities(self, selection: t.Optional[AtomSelection] = None) -> t.Optional[NDArray[numpy.float64]]:
    """Return a `(N, 3)` ndarray of atom velocities (dtype [`numpy.float64`][numpy.float64])."""
    if 'velocity' not in self:
        return None

    df = self if selection is None else self.filter(_selection_to_expr(self, selection))
    return df.get_column('velocity').to_numpy().astype(numpy.float64)

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,
    **kwargs: Any,
) -> Self

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

add_atom(
    elem: Union[int, str],
    /,
    x: Union[ArrayLike, float],
    y: Optional[float] = None,
    z: Optional[float] = 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 HasAtoms.concat instead).

Source code in atomlib/atoms.py

def add_atom(self, elem: t.Union[int, str], /,
             x: t.Union[ArrayLike, float],
             y: t.Optional[float] = None,
             z: t.Optional[float] = 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 [`HasAtoms.concat`][atomlib.atoms.HasAtoms.concat] instead).
    """
    if isinstance(elem, int):
        kwargs.update(elem=elem)
    else:
        kwargs.update(symbol=elem)
    if hasattr(x, '__len__') and len(x) > 1:  # type: ignore
        (x, y, z) = to_vec3(x)
    elif y is None or z is None:
        raise ValueError("Must specify vector of positions or x, y, & z.")

    sym = get_sym(elem) if isinstance(elem, int) else elem
    d: t.Dict[str, t.Any] = {'x': x, 'y': y, 'z': z, 'symbol': sym, **kwargs}
    return self.concat(
        (self, Atoms(d).select_schema(self.schema)),
        how='vertical'
    )

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) -> 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/atoms.py

def with_index(self, index: t.Optional[AtomValues] = 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.
    """
    if index is None and 'i' in self.columns:
        return self
    if index is None:
        index = numpy.arange(len(self), dtype=numpy.int64)
    return self.with_column(_values_to_expr(self, index, polars.Int64).alias('i'))

with_wobble

with_wobble(wobble: Optional[AtomValues] = 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/atoms.py

def with_wobble(self, wobble: t.Optional[AtomValues] = 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.
    """
    if wobble is None and 'wobble' in self.columns:
        return self
    wobble = 0. if wobble is None else wobble
    return self.with_column(_values_to_expr(self, wobble, polars.Float64).alias('wobble'))

with_occupancy

with_occupancy(
    frac_occupancy: Optional[AtomValues] = 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/atoms.py

def with_occupancy(self, frac_occupancy: t.Optional[AtomValues] = 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.
    """
    if frac_occupancy is None and 'frac_occupancy' in self.columns:
        return self
    frac_occupancy = 1. if frac_occupancy is None else frac_occupancy
    return self.with_column(_values_to_expr(self, frac_occupancy, polars.Float64).alias('frac_occupancy'))

apply_wobble

apply_wobble(
    rng: Union[Generator, int, None] = 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/atoms.py

def apply_wobble(self, rng: t.Union[numpy.random.Generator, int, None] = 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.
    """
    if 'wobble' not in self.columns:
        return self
    rng = numpy.random.default_rng(seed=rng)

    stddev = self.select((polars.col('wobble') / 3.).sqrt()).to_series().to_numpy()
    coords = self.coords()
    coords += stddev[:, None] * rng.standard_normal(coords.shape)
    return self.with_coords(coords)

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) -> 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/atoms.py

def with_type(self, types: t.Optional[AtomValues] = 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]`
    """
    if types is not None:
        return self.with_columns(type=_values_to_expr(self, types, polars.Int32))
    if 'type' in self.columns:
        return self

    unique = Atoms(self._get_frame().unique(maintain_order=False, subset=['elem', 'symbol']).sort(['elem', 'symbol']), _unchecked=True)
    new = self.with_column(polars.Series('type', values=numpy.zeros(len(self)), dtype=polars.Int32))

    logging.warning("Auto-assigning element types")
    for (i, (elem, sym)) in enumerate(unique.select(('elem', 'symbol')).rows()):
        print(f"Assigning type {i+1} to element '{sym}'")
        new = new.with_column(polars.when((polars.col('elem') == elem) & (polars.col('symbol') == sym))
                                    .then(polars.lit(i+1))
                                    .otherwise(polars.col('type'))
                                    .alias('type'))

    assert (new.get_column('type') == 0).sum() == 0
    return new

with_mass

with_mass(mass: Optional[ArrayLike] = 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/atoms.py

def with_mass(self, mass: t.Optional[ArrayLike] = 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.
    """
    if mass is not None:
        return self.with_column(_values_to_expr(self, mass, polars.Float32).alias('mass'))
    if 'mass' in self.columns:
        return self

    unique_elems = self.get_column('elem').unique()
    new = self.with_column(polars.Series('mass', values=numpy.zeros(len(self)), dtype=polars.Float32))

    logging.warning("Auto-assigning element masses")
    for elem in unique_elems:
        new = new.with_column(polars.when(polars.col('elem') == elem)
                                    .then(polars.lit(get_mass(elem)))
                                    .otherwise(polars.col('mass'))
                                    .alias('mass'))

    assert (new.get_column('mass').abs() < 1e-10).sum() == 0
    return new

with_symbol

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

Return self with the given atomic symbols.

Source code in atomlib/atoms.py

def with_symbol(self, symbols: ArrayLike, selection: t.Optional[AtomSelection] = None) -> Self:
    """
    Return `self` with the given atomic symbols.
    """
    if selection is not None:
        selection = _selection_to_numpy(self, selection)
        new_symbols = self.get_column('symbol')
        new_symbols[selection] = polars.Series(list(numpy.broadcast_to(symbols, len(selection))), dtype=polars.Utf8)
        symbols = new_symbols

    # TODO better cast here
    symbols = polars.Series('symbol', list(numpy.broadcast_to(symbols, len(self))), dtype=polars.Utf8)
    return self.with_columns((symbols, get_elem(symbols)))

with_coords

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

Return self replaced with the given atomic positions.

Source code in atomlib/atoms.py

def with_coords(self, pts: ArrayLike, selection: t.Optional[AtomSelection] = None, *, frame: t.Literal['local'] = 'local') -> Self:
    """
    Return `self` replaced with the given atomic positions.
    """
    if selection is not None:
        selection = _selection_to_numpy(self, selection)
        new_pts = self.coords()
        pts = numpy.atleast_2d(pts)
        assert pts.shape[-1] == 3
        new_pts[selection] = pts
        pts = new_pts

    # https://github.com/pola-rs/polars/issues/18369
    pts = numpy.broadcast_to(pts, (len(self), 3)) if len(self) else []
    return self.with_columns(polars.Series('coords', pts, polars.Array(polars.Float64, 3)))

with_velocity

with_velocity(
    pts: Optional[ArrayLike] = None,
    selection: Optional[AtomSelection] = 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/atoms.py

def with_velocity(self, pts: t.Optional[ArrayLike] = None,
                  selection: t.Optional[AtomSelection] = None) -> Self:
    """
    Return `self` replaced with the given atomic velocities.
    If `pts` is not specified, use the already existing velocities or zero.
    """
    if pts is None:
        if 'velocity' in self:
            return self
        all_pts = numpy.zeros((len(self), 3))
    else:
        all_pts = self['velocity'].to_numpy()

    if selection is None:
        all_pts = pts or all_pts
    elif pts is not None:
        selection = _selection_to_numpy(self, selection)
        all_pts = numpy.require(all_pts, requirements=['WRITEABLE'])
        pts = numpy.atleast_2d(pts)
        assert pts.shape[-1] == 3
        all_pts[selection] = pts

    all_pts = numpy.broadcast_to(all_pts, (len(self), 3))
    return self.with_columns(polars.Series('velocity', all_pts, polars.Array(polars.Float64, 3)))

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

empty staticmethod

empty() -> Atoms

Return an empty Atoms with only the mandatory columns.

Source code in atomlib/atoms.py

@staticmethod
def empty() -> Atoms:
    """
    Return an empty Atoms with only the mandatory columns.
    """
    return Atoms()

get_atoms

get_atoms(frame: Literal['local'] = 'local') -> Atoms

Source code in atomlib/atoms.py

def get_atoms(self, frame: t.Literal['local'] = 'local') -> Atoms:
    if frame != 'local':
        raise ValueError(f"Atoms without a cell only support the 'local' coordinate frame, not '{frame}'.")
    return self

with_atoms

with_atoms(
    atoms: HasAtoms, frame: Literal["local"] = "local"
) -> Atoms

Source code in atomlib/atoms.py

def with_atoms(self, atoms: HasAtoms, frame: t.Literal['local'] = 'local') -> Atoms:
    if frame != 'local':
        raise ValueError(f"Atoms without a cell only support the 'local' coordinate frame, not '{frame}'.")
    return atoms.get_atoms()

bbox

bbox() -> BBox3D

Return the bounding box of all the points in self.

Source code in atomlib/atoms.py

def bbox(self) -> BBox3D:
    """Return the bounding box of all the points in `self`."""
    if self._bbox is None:
        self._bbox = BBox3D.from_pts(self.coords())

    return self._bbox