Index

Top-level exports:

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.

CoordinateFrame module-attribute

CoordinateFrame: TypeAlias = Literal[
    "cell",
    "cell_frac",
    "cell_box",
    "ortho",
    "ortho_frac",
    "ortho_box",
    "linear",
    "local",
    "global",
]

A coordinate frame to use.

• cell: Real-space units along crystal axes
• cell_frac: Fraction of unit cells
• cell_box: Fraction of cell box
• ortho: Real-space units along orthogonal cell
• ortho_frac: Fraction of orthogonal cell
• ortho_box: Fraction of orthogonal box
• linear: Angstroms in local coordinate system (without affine transformation)
• local: Angstroms in local coordinate system (with affine transformation)
• global: Angstroms in global coordinate system (with all transformations)

For more information, see the documentation at Coordinate systems, or the example notebook at examples/coords.ipynb.

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

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

Cell dataclass

Bases: HasCell

Internal class for representing the coordinate systems of a crystal.

The overall transformation from crystal coordinates to real-space coordinates is is split into four transformations, applied from bottom to top. First is n_cells, which scales from fractions of a unit cell to fractions of a supercell. Next is cell_size, which scales to real-space units. ortho is an orthogonalization matrix, a det = 1 upper-triangular matrix which transforms crystal axes to an orthogonal coordinate system. Finally, affine contains any remaining transformations to the local coordinate system, which atoms are stored in.

Source code in atomlib/cell.py

@dataclass(frozen=True, init=False)
class Cell(HasCell):
    """
    Internal class for representing the coordinate systems of a crystal.

    The overall transformation from crystal coordinates to real-space coordinates is
    is split into four transformations, applied from bottom to top. First is `n_cells`,
    which scales from fractions of a unit cell to fractions of a supercell. Next is
    `cell_size`, which scales to real-space units. `ortho` is an orthogonalization
    matrix, a det = 1 upper-triangular matrix which transforms crystal axes to
    an orthogonal coordinate system. Finally, `affine` contains any remaining
    transformations to the local coordinate system, which atoms are stored in.
    """

    def get_cell(self) -> Cell:
        return self

    def with_cell(self: Cell, cell: Cell) -> Cell:
        return cell

    _affine: AffineTransform3D = AffineTransform3D()
    """
    Affine transformation. Holds transformation from `'ortho'` to `'local'` coordinates,
    including rotation away from the standard crystal orientation.
    """

    _ortho: LinearTransform3D = LinearTransform3D()
    """
    Orthogonalization transformation. Skews but does not scale the crystal axes to cartesian axes.
    """

    _cell_size: NDArray[numpy.float64]
    """Unit cell size."""
    _cell_angle: NDArray[numpy.float64] = field(default_factory=lambda: numpy.full(3, numpy.pi/2.))
    """Unit cell angles, in radians."""
    _n_cells: NDArray[numpy.int64] = field(default_factory=lambda: numpy.ones(3, numpy.int64))
    """Number of unit cells."""
    _pbc: NDArray[numpy.bool_] = field(default_factory=lambda: numpy.ones(3, numpy.bool_))
    """Flags indicating the presence of periodic boundary conditions along each axis."""

    def __init__(self, *,
        affine: t.Optional[AffineTransform3D] = None, ortho: t.Optional[LinearTransform3D] = None,
        cell_size: VecLike, cell_angle: t.Optional[VecLike] = None,
        n_cells: t.Optional[VecLike] = None, pbc: t.Optional[VecLike] = None):

        object.__setattr__(self, '_affine', AffineTransform3D() if affine is None else affine)
        object.__setattr__(self, '_ortho', LinearTransform3D() if ortho is None else ortho)
        object.__setattr__(self, '_cell_size', to_vec3(cell_size))
        object.__setattr__(self, '_cell_angle', numpy.full(3, numpy.pi/2.) if cell_angle is None else to_vec3(cell_angle))
        object.__setattr__(self, '_n_cells', numpy.ones(3, numpy.int_) if n_cells is None else to_vec3(n_cells, numpy.int64))
        object.__setattr__(self, '_pbc', numpy.ones(3, numpy.bool_) if pbc is None else to_vec3(pbc, numpy.bool_))

    @staticmethod
    def from_unit_cell(cell_size: VecLike, cell_angle: t.Optional[VecLike] = None, n_cells: t.Optional[VecLike] = None,
                       pbc: t.Optional[VecLike] = None):
        return Cell(
            ortho=cell_to_ortho([1.]*3, cell_angle),
            n_cells=to_vec3([1]*3 if n_cells is None else n_cells, numpy.int_),
            cell_size=to_vec3(cell_size),
            cell_angle=to_vec3([numpy.pi/2.]*3 if cell_angle is None else cell_angle),
            pbc=pbc
        )

    @staticmethod
    def from_ortho(ortho: AffineTransform3D, n_cells: t.Optional[VecLike] = None, pbc: t.Optional[VecLike] = None):
        lin = ortho.to_linear()
        # decompose into orthogonal and upper triangular
        q, r = numpy.linalg.qr(lin.inner)

        # flip QR decomposition so R has positive diagonals
        signs = numpy.sign(numpy.diagonal(r))
        # multiply flips to columns of Q, rows of R
        q = q * signs
        r = r * signs[:, None]
        #numpy.testing.assert_allclose(q @ r, lin.inner)
        if numpy.linalg.det(q) < 0:
            warn("Crystal is left-handed. This is currently unsupported, and may cause errors.")
            # currently, behavior is to leave `ortho` proper, and move the inversion into the affine transform

        cell_size, cell_angle = ortho_to_cell(lin)
        return Cell(
            affine=LinearTransform3D(q).translate(ortho.translation()),
            ortho=LinearTransform3D(r / cell_size).round_near_zero(),
            cell_size=cell_size, cell_angle=cell_angle,
            n_cells=to_vec3([1]*3 if n_cells is None else n_cells, numpy.int_),
            pbc=pbc,
        )

    def __str__(self) -> str:
        return "\n".join((
            self.__class__.__name__,
            f"Cell size: {self.cell_size!r}",
            f"Cell angle: {self.cell_angle!r}",
            f"# cells: {self.n_cells!r}",
            f"pbc: {self.pbc!r}",
        ))

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}("
            f"ortho={self.ortho}, affine={self.affine}, cell_size={self.cell_size}, "
            f"cell_angle={self.cell_angle}, n_cells={self.n_cells}, pbc={self.pbc})"
        )

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

affine property

affine: AffineTransform3D

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

ortho property

ortho: LinearTransform3D

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

metric property

metric: LinearTransform3D

Cell metric tensor

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

cell_size property

cell_size: NDArray[float64]

Unit cell size.

cell_angle property

cell_angle: NDArray[float64]

Unit cell angles, in radians.

n_cells property

n_cells: NDArray[int_]

Number of unit cells.

pbc property

pbc: NDArray[bool_]

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

ortho_size property

ortho_size: NDArray[float64]

Return size of orthogonal unit cell.

Equivalent to the diagonal of the orthogonalization matrix.

box_size property

box_size: NDArray[float64]

Return size of the cell box.

Equivalent to self.n_cells * self.cell_size.

bbox class-attribute instance-attribute

bbox = bbox_cell

get_transform

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

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

Source code in atomlib/cell.py

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

corners

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

Source code in atomlib/cell.py

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

bbox_cell

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

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

Source code in atomlib/cell.py

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

is_orthogonal

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

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

Source code in atomlib/cell.py

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

is_orthogonal_in_local

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

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

Source code in atomlib/cell.py

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

to_ortho

to_ortho() -> AffineTransform3D

Source code in atomlib/cell.py

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

transform_cell

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

Apply the given transform to the unit cell, and return a new Cell. The transform is applied in coordinate frame 'frame'. Orthogonal and affine transformations are applied to the affine matrix component, while skew and scaling is applied to the orthogonalization matrix/cell_size.

Source code in atomlib/cell.py

def transform_cell(self: HasCellT, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> HasCellT:
    """
    Apply the given transform to the unit cell, and return a new `Cell`.
    The transform is applied in coordinate frame 'frame'.
    Orthogonal and affine transformations are applied to the affine matrix component,
    while skew and scaling is applied to the orthogonalization matrix/cell_size.
    """
    transform = t.cast(AffineTransform3D, self.change_transform(transform, 'local', frame))
    if not transform.to_linear().is_orthogonal():
        raise NotImplementedError()
    return self.with_cell(Cell(
        affine=transform @ self.affine,
        ortho=self.ortho,
        cell_size=self.cell_size,
        cell_angle=self.cell_angle,
        n_cells=self.n_cells,
        pbc=self.pbc,
    ))

strain_orthogonal

strain_orthogonal() -> HasCellT

Orthogonalize using strain.

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

Source code in atomlib/cell.py

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

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

repeat

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

Tile the cell by n in each dimension.

Source code in atomlib/cell.py

def repeat(self: HasCellT, n: t.Union[int, VecLike]) -> HasCellT:
    """Tile the cell by `n` in each dimension."""
    ns = numpy.broadcast_to(n, 3)
    if not numpy.issubdtype(ns.dtype, numpy.integer):
        raise ValueError("repeat() argument must be an integer or integer array.")
    return self.with_cell(Cell(
        affine=self.affine,
        ortho=self.ortho,
        cell_size=self.cell_size,
        cell_angle=self.cell_angle,
        n_cells=self.n_cells * numpy.broadcast_to(n, 3),
        pbc = self.pbc | (ns > 1)  # assume periodic along repeated directions
    ))

explode

explode() -> HasCellT

Materialize repeated cells as one supercell.

Source code in atomlib/cell.py

def explode(self: HasCellT) -> HasCellT:
    """Materialize repeated cells as one supercell."""
    return self.with_cell(Cell(
        affine=self.affine,
        ortho=self.ortho,
        cell_size=self.cell_size*self.n_cells,
        cell_angle=self.cell_angle,
        pbc=self.pbc,
    ))

explode_z

explode_z() -> HasCellT

Materialize repeated cells as one supercell in z.

Source code in atomlib/cell.py

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

crop

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

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

Source code in atomlib/cell.py

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

    if not frame.lower().startswith('cell'):
        if not self.is_orthogonal():
            raise ValueError("Cannot crop a non-orthogonal cell in orthogonal coordinates. Use crop_atoms instead.")

    min = to_vec3([x_min, y_min, z_min])
    max = to_vec3([x_max, y_max, z_max])
    (min, max) = self.get_transform('cell_box', frame).transform([min, max])
    new_box = BBox3D(min, max) & BBox3D.unit()
    cropped = (new_box.min > 0.) | (new_box.max < 1.)

    return self.with_cell(Cell(
        affine=self.affine @ AffineTransform3D.translate(-new_box.min),
        ortho=self.ortho,
        cell_size=new_box.size * self.cell_size * numpy.where(cropped, self.n_cells, 1),
        n_cells=numpy.where(cropped, 1, self.n_cells),
        cell_angle=self.cell_angle,
        pbc=self.pbc & ~cropped  # remove periodicity along cropped directions
    ))

change_transform

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

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

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

Coordinate-change a transformation from frame_from into frame_to.

Source code in atomlib/cell.py

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

assert_equal

assert_equal(other: Any)

Source code in atomlib/cell.py

def assert_equal(self, other: t.Any):
    assert isinstance(other, HasCell) and type(self) is type(other)
    numpy.testing.assert_array_almost_equal(self.affine.inner, other.affine.inner, 6)
    numpy.testing.assert_array_almost_equal(self.ortho.inner, other.ortho.inner, 6)
    numpy.testing.assert_array_almost_equal(self.cell_size, other.cell_size, 6)
    numpy.testing.assert_array_equal(self.n_cells, other.n_cells)
    numpy.testing.assert_array_equal(self.pbc, other.pbc)

get_cell

get_cell() -> Cell

Source code in atomlib/cell.py

def get_cell(self) -> Cell:
    return self

with_cell

with_cell(cell: Cell) -> Cell

Source code in atomlib/cell.py

def with_cell(self: Cell, cell: Cell) -> Cell:
    return cell

from_unit_cell staticmethod

from_unit_cell(
    cell_size: VecLike,
    cell_angle: Optional[VecLike] = None,
    n_cells: Optional[VecLike] = None,
    pbc: Optional[VecLike] = None,
)

Source code in atomlib/cell.py

@staticmethod
def from_unit_cell(cell_size: VecLike, cell_angle: t.Optional[VecLike] = None, n_cells: t.Optional[VecLike] = None,
                   pbc: t.Optional[VecLike] = None):
    return Cell(
        ortho=cell_to_ortho([1.]*3, cell_angle),
        n_cells=to_vec3([1]*3 if n_cells is None else n_cells, numpy.int_),
        cell_size=to_vec3(cell_size),
        cell_angle=to_vec3([numpy.pi/2.]*3 if cell_angle is None else cell_angle),
        pbc=pbc
    )

from_ortho staticmethod

from_ortho(
    ortho: AffineTransform3D,
    n_cells: Optional[VecLike] = None,
    pbc: Optional[VecLike] = None,
)

Source code in atomlib/cell.py

@staticmethod
def from_ortho(ortho: AffineTransform3D, n_cells: t.Optional[VecLike] = None, pbc: t.Optional[VecLike] = None):
    lin = ortho.to_linear()
    # decompose into orthogonal and upper triangular
    q, r = numpy.linalg.qr(lin.inner)

    # flip QR decomposition so R has positive diagonals
    signs = numpy.sign(numpy.diagonal(r))
    # multiply flips to columns of Q, rows of R
    q = q * signs
    r = r * signs[:, None]
    #numpy.testing.assert_allclose(q @ r, lin.inner)
    if numpy.linalg.det(q) < 0:
        warn("Crystal is left-handed. This is currently unsupported, and may cause errors.")
        # currently, behavior is to leave `ortho` proper, and move the inversion into the affine transform

    cell_size, cell_angle = ortho_to_cell(lin)
    return Cell(
        affine=LinearTransform3D(q).translate(ortho.translation()),
        ortho=LinearTransform3D(r / cell_size).round_near_zero(),
        cell_size=cell_size, cell_angle=cell_angle,
        n_cells=to_vec3([1]*3 if n_cells is None else n_cells, numpy.int_),
        pbc=pbc,
    )

HasCell

Source code in atomlib/cell.py

class HasCell:
    # abstract methods

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

    @abc.abstractmethod
    def with_cell(self: HasCellT, cell: Cell) -> HasCellT:
        """Replace the cell in ``self`` with ``cell``."""
        ...

    # getters

    @property
    def affine(self) -> AffineTransform3D:
        """
        Affine transformation. Holds transformation from 'ortho' to 'local' coordinates,
        including rotation away from the standard crystal orientation.
        """
        return self.get_cell()._affine

    @property
    def ortho(self) -> LinearTransform3D:
        """
        Orthogonalization transformation. Skews but does not scale the crystal axes to cartesian axes.
        """
        return self.get_cell()._ortho

    @property
    def metric(self) -> LinearTransform3D:
        r"""
        Cell metric tensor

        Returns the dot product between every combination of basis vectors.
        :math:`\mathbf{a} \cdot \mathbf{b} = a_i M_ij b_j`
        """
        ortho = self.get_cell()._ortho.scale(self.cell_size)
        return ortho.T @ ortho

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

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

    @property
    def n_cells(self) -> NDArray[numpy.int_]:
        """Number of unit cells."""
        return self.get_cell()._n_cells

    @property
    def pbc(self) -> NDArray[numpy.bool_]:
        """Flags indicating the presence of periodic boundary conditions along each axis."""
        return self.get_cell()._pbc

    @property
    def ortho_size(self) -> NDArray[numpy.float64]:
        """
        Return size of orthogonal unit cell.

        Equivalent to the diagonal of the orthogonalization matrix.
        """
        return self.cell_size * numpy.diag(self.ortho.inner)

    @property
    def box_size(self) -> NDArray[numpy.float64]:
        """
        Return size of the cell box.

        Equivalent to ``self.n_cells * self.cell_size``.
        """
        return self.n_cells * self.cell_size

    # get transforms

    def _get_transform_to_local(self, frame: CoordinateFrame) -> AffineTransform3D:
        """Get the transform from `frame` to local coordinates."""
        frame = t.cast(CoordinateFrame, frame.lower())

        if frame == 'local' or frame == 'global':
            return LinearTransform3D()

        if frame == 'linear':
            return self.affine.to_translation()

        if frame.startswith('cell'):
            transform = self.affine @ self.ortho
            cell_size = self.cell_size
        elif frame.startswith('ortho'):
            transform = self.affine
            cell_size = self.ortho_size
        else:
            raise ValueError(f"Unknown coordinate frame '{frame}'")

        if '_' not in frame:
            return transform
        end = frame.split('_', 2)[1]
        if end == 'frac':
            return transform @ LinearTransform3D.scale(cell_size)
        if end == 'box':
            return transform @ LinearTransform3D.scale(cell_size * self.n_cells)
        raise ValueError(f"Unknown coordinate frame '{frame}'")

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

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

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

    bbox = bbox_cell

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

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

    def _cell_size_in_local(self) -> Vec3:
        """Calculate cell_size in the local coordinate system. Assumes `self.is_orthogonal_in_local()`."""
        return numpy.abs(self.get_transform('local', 'ortho').transform_vec(self.cell_size))

    def _box_size_in_local(self) -> Vec3:
        """Calculate box_size in the local coordinate system. Assumes `self.is_orthogonal_in_local()`."""
        return numpy.abs(self.get_transform('local', 'ortho').transform_vec(self.box_size))

    def _n_cells_in_local(self) -> NDArray[numpy.int_]:
        """Calculate n_cells after any local rotation. Assumes `self.is_orthogonal_in_local()`."""
        return numpy.abs(numpy.round(self.get_transform('local', 'ortho').transform_vec(self.n_cells)).astype(int))

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

    def transform_cell(self: HasCellT, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> HasCellT:
        """
        Apply the given transform to the unit cell, and return a new `Cell`.
        The transform is applied in coordinate frame 'frame'.
        Orthogonal and affine transformations are applied to the affine matrix component,
        while skew and scaling is applied to the orthogonalization matrix/cell_size.
        """
        transform = t.cast(AffineTransform3D, self.change_transform(transform, 'local', frame))
        if not transform.to_linear().is_orthogonal():
            raise NotImplementedError()
        return self.with_cell(Cell(
            affine=transform @ self.affine,
            ortho=self.ortho,
            cell_size=self.cell_size,
            cell_angle=self.cell_angle,
            n_cells=self.n_cells,
            pbc=self.pbc,
        ))

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

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

    def repeat(self: HasCellT, n: t.Union[int, VecLike]) -> HasCellT:
        """Tile the cell by `n` in each dimension."""
        ns = numpy.broadcast_to(n, 3)
        if not numpy.issubdtype(ns.dtype, numpy.integer):
            raise ValueError("repeat() argument must be an integer or integer array.")
        return self.with_cell(Cell(
            affine=self.affine,
            ortho=self.ortho,
            cell_size=self.cell_size,
            cell_angle=self.cell_angle,
            n_cells=self.n_cells * numpy.broadcast_to(n, 3),
            pbc = self.pbc | (ns > 1)  # assume periodic along repeated directions
        ))

    def explode(self: HasCellT) -> HasCellT:
        """Materialize repeated cells as one supercell."""
        return self.with_cell(Cell(
            affine=self.affine,
            ortho=self.ortho,
            cell_size=self.cell_size*self.n_cells,
            cell_angle=self.cell_angle,
            pbc=self.pbc,
        ))

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

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

        if not frame.lower().startswith('cell'):
            if not self.is_orthogonal():
                raise ValueError("Cannot crop a non-orthogonal cell in orthogonal coordinates. Use crop_atoms instead.")

        min = to_vec3([x_min, y_min, z_min])
        max = to_vec3([x_max, y_max, z_max])
        (min, max) = self.get_transform('cell_box', frame).transform([min, max])
        new_box = BBox3D(min, max) & BBox3D.unit()
        cropped = (new_box.min > 0.) | (new_box.max < 1.)

        return self.with_cell(Cell(
            affine=self.affine @ AffineTransform3D.translate(-new_box.min),
            ortho=self.ortho,
            cell_size=new_box.size * self.cell_size * numpy.where(cropped, self.n_cells, 1),
            n_cells=numpy.where(cropped, 1, self.n_cells),
            cell_angle=self.cell_angle,
            pbc=self.pbc & ~cropped  # remove periodicity along cropped directions
        ))

    @t.overload
    def change_transform(self, transform: AffineTransform3D,
                         frame_to: t.Optional[CoordinateFrame] = None,
                         frame_from: t.Optional[CoordinateFrame] = None) -> AffineTransform3D:
        ...

    @t.overload
    def change_transform(self, transform: Transform3D,
                         frame_to: t.Optional[CoordinateFrame] = None,
                         frame_from: t.Optional[CoordinateFrame] = None) -> Transform3D:
        ...

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

    def assert_equal(self, other: t.Any):
        assert isinstance(other, HasCell) and type(self) is type(other)
        numpy.testing.assert_array_almost_equal(self.affine.inner, other.affine.inner, 6)
        numpy.testing.assert_array_almost_equal(self.ortho.inner, other.ortho.inner, 6)
        numpy.testing.assert_array_almost_equal(self.cell_size, other.cell_size, 6)
        numpy.testing.assert_array_equal(self.n_cells, other.n_cells)
        numpy.testing.assert_array_equal(self.pbc, other.pbc)

affine property

affine: AffineTransform3D

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

ortho property

ortho: LinearTransform3D

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

metric property

metric: LinearTransform3D

Cell metric tensor

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

cell_size property

cell_size: NDArray[float64]

Unit cell size.

cell_angle property

cell_angle: NDArray[float64]

Unit cell angles, in radians.

n_cells property

n_cells: NDArray[int_]

Number of unit cells.

pbc property

pbc: NDArray[bool_]

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

ortho_size property

ortho_size: NDArray[float64]

Return size of orthogonal unit cell.

Equivalent to the diagonal of the orthogonalization matrix.

box_size property

box_size: NDArray[float64]

Return size of the cell box.

Equivalent to self.n_cells * self.cell_size.

bbox class-attribute instance-attribute

bbox = bbox_cell

get_cell abstractmethod

get_cell() -> Cell

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

Source code in atomlib/cell.py

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

with_cell abstractmethod

with_cell(cell: Cell) -> HasCellT

Replace the cell in self with cell.

Source code in atomlib/cell.py

@abc.abstractmethod
def with_cell(self: HasCellT, cell: Cell) -> HasCellT:
    """Replace the cell in ``self`` with ``cell``."""
    ...

get_transform

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

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

Source code in atomlib/cell.py

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

corners

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

Source code in atomlib/cell.py

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

bbox_cell

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

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

Source code in atomlib/cell.py

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

is_orthogonal

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

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

Source code in atomlib/cell.py

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

is_orthogonal_in_local

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

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

Source code in atomlib/cell.py

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

to_ortho

to_ortho() -> AffineTransform3D

Source code in atomlib/cell.py

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

transform_cell

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

Apply the given transform to the unit cell, and return a new Cell. The transform is applied in coordinate frame 'frame'. Orthogonal and affine transformations are applied to the affine matrix component, while skew and scaling is applied to the orthogonalization matrix/cell_size.

Source code in atomlib/cell.py

def transform_cell(self: HasCellT, transform: AffineTransform3D, frame: CoordinateFrame = 'local') -> HasCellT:
    """
    Apply the given transform to the unit cell, and return a new `Cell`.
    The transform is applied in coordinate frame 'frame'.
    Orthogonal and affine transformations are applied to the affine matrix component,
    while skew and scaling is applied to the orthogonalization matrix/cell_size.
    """
    transform = t.cast(AffineTransform3D, self.change_transform(transform, 'local', frame))
    if not transform.to_linear().is_orthogonal():
        raise NotImplementedError()
    return self.with_cell(Cell(
        affine=transform @ self.affine,
        ortho=self.ortho,
        cell_size=self.cell_size,
        cell_angle=self.cell_angle,
        n_cells=self.n_cells,
        pbc=self.pbc,
    ))

strain_orthogonal

strain_orthogonal() -> HasCellT

Orthogonalize using strain.

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

Source code in atomlib/cell.py

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

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

repeat

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

Tile the cell by n in each dimension.

Source code in atomlib/cell.py

def repeat(self: HasCellT, n: t.Union[int, VecLike]) -> HasCellT:
    """Tile the cell by `n` in each dimension."""
    ns = numpy.broadcast_to(n, 3)
    if not numpy.issubdtype(ns.dtype, numpy.integer):
        raise ValueError("repeat() argument must be an integer or integer array.")
    return self.with_cell(Cell(
        affine=self.affine,
        ortho=self.ortho,
        cell_size=self.cell_size,
        cell_angle=self.cell_angle,
        n_cells=self.n_cells * numpy.broadcast_to(n, 3),
        pbc = self.pbc | (ns > 1)  # assume periodic along repeated directions
    ))

explode

explode() -> HasCellT

Materialize repeated cells as one supercell.

Source code in atomlib/cell.py

def explode(self: HasCellT) -> HasCellT:
    """Materialize repeated cells as one supercell."""
    return self.with_cell(Cell(
        affine=self.affine,
        ortho=self.ortho,
        cell_size=self.cell_size*self.n_cells,
        cell_angle=self.cell_angle,
        pbc=self.pbc,
    ))

explode_z

explode_z() -> HasCellT

Materialize repeated cells as one supercell in z.

Source code in atomlib/cell.py

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

crop

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

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

Source code in atomlib/cell.py

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

    if not frame.lower().startswith('cell'):
        if not self.is_orthogonal():
            raise ValueError("Cannot crop a non-orthogonal cell in orthogonal coordinates. Use crop_atoms instead.")

    min = to_vec3([x_min, y_min, z_min])
    max = to_vec3([x_max, y_max, z_max])
    (min, max) = self.get_transform('cell_box', frame).transform([min, max])
    new_box = BBox3D(min, max) & BBox3D.unit()
    cropped = (new_box.min > 0.) | (new_box.max < 1.)

    return self.with_cell(Cell(
        affine=self.affine @ AffineTransform3D.translate(-new_box.min),
        ortho=self.ortho,
        cell_size=new_box.size * self.cell_size * numpy.where(cropped, self.n_cells, 1),
        n_cells=numpy.where(cropped, 1, self.n_cells),
        cell_angle=self.cell_angle,
        pbc=self.pbc & ~cropped  # remove periodicity along cropped directions
    ))

change_transform

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

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

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

Coordinate-change a transformation from frame_from into frame_to.

Source code in atomlib/cell.py

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

assert_equal

assert_equal(other: Any)

Source code in atomlib/cell.py

def assert_equal(self, other: t.Any):
    assert isinstance(other, HasCell) and type(self) is type(other)
    numpy.testing.assert_array_almost_equal(self.affine.inner, other.affine.inner, 6)
    numpy.testing.assert_array_almost_equal(self.ortho.inner, other.ortho.inner, 6)
    numpy.testing.assert_array_almost_equal(self.cell_size, other.cell_size, 6)
    numpy.testing.assert_array_equal(self.n_cells, other.n_cells)
    numpy.testing.assert_array_equal(self.pbc, other.pbc)

AtomCell dataclass

Bases: AtomCellIOMixin, HasAtomCell

Cell of atoms with known size and periodic boundary conditions.

Source code in atomlib/atomcell.py

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        self.__post_init__()

    def __post_init__(self):
        pass

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

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

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

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

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

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

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

affine property

affine: AffineTransform3D

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

ortho property

ortho: LinearTransform3D

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

metric property

metric: LinearTransform3D

Cell metric tensor

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

cell_size property

cell_size: NDArray[float64]

Unit cell size.

cell_angle property

cell_angle: NDArray[float64]

Unit cell angles, in radians.

n_cells property

n_cells: NDArray[int_]

Number of unit cells.

pbc property

pbc: NDArray[bool_]

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

ortho_size property

ortho_size: NDArray[float64]

Return size of orthogonal unit cell.

Equivalent to the diagonal of the orthogonalization matrix.

box_size property

box_size: NDArray[float64]

Return size of the cell box.

Equivalent to self.n_cells * self.cell_size.

columns property

columns: List[str]

Return the column names in self.

RETURNS	DESCRIPTION
List[str]	A sequence of column names

dtypes property

dtypes: List[DataType]

Return the datatypes in self.

RETURNS	DESCRIPTION
List[DataType]	A sequence of column DataTypes

schema property

schema: Schema

Return the schema of self.

RETURNS	DESCRIPTION
Schema	A dictionary of column names and DataTypes

with_column class-attribute instance-attribute

with_column = with_columns

unique class-attribute instance-attribute

unique = deduplicate

atoms instance-attribute

atoms: Atoms

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

cell instance-attribute

cell: Cell

Cell coordinate system.

frame class-attribute instance-attribute

frame: CoordinateFrame = 'local'

Coordinate frame 'atoms' are stored in.

get_transform

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

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

Source code in atomlib/cell.py

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

corners

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

Source code in atomlib/cell.py

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

bbox_cell

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

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

Source code in atomlib/cell.py

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

bbox

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

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

Source code in atomlib/atomcell.py

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

is_orthogonal

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

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

Source code in atomlib/cell.py

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

is_orthogonal_in_local

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

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

Source code in atomlib/cell.py

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

to_ortho

to_ortho() -> AffineTransform3D

Source code in atomlib/cell.py

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

transform_cell

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

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

Source code in atomlib/atomcell.py

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

strain_orthogonal

strain_orthogonal() -> HasCellT

Orthogonalize using strain.

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

Source code in atomlib/cell.py

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

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

repeat

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

Tile the cell

Source code in atomlib/atomcell.py

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

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

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

explode

explode() -> Self

Materialize repeated cells as one supercell.

Source code in atomlib/atomcell.py

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

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

explode_z

explode_z() -> HasCellT

Materialize repeated cells as one supercell in z.

Source code in atomlib/cell.py

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

crop

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

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

To crop atoms only, use crop_atoms instead.

Source code in atomlib/atomcell.py

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

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

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

change_transform

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

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

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

Coordinate-change a transformation from frame_from into frame_to.

Source code in atomlib/cell.py

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

describe

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

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

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

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

Source code in atomlib/atomcell.py

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

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

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

with_columns

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

Return a copy of self with the given columns added.

Source code in atomlib/atomcell.py

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

insert_column

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

Source code in atomlib/atoms.py

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

get_column

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

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

Source code in atomlib/atomcell.py

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

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

get_columns

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

Return all columns from self as a list of Series.

Source code in atomlib/atomcell.py

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

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

get_column_index

get_column_index(name: str) -> int

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

Source code in atomlib/atoms.py

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

group_by

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

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

Source code in atomlib/atomcell.py

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

pipe

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

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

Source code in atomlib/atomcell.py

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

drop

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

Return self with the specified columns removed.

Source code in atomlib/atoms.py

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

filter

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

Filter self, removing rows which evaluate to False.

Source code in atomlib/atomcell.py

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

sort

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

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

Source code in atomlib/atomcell.py

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

slice

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

Return a slice of the rows in self.

Source code in atomlib/atomcell.py

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

head

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

Return the first n rows of self.

Source code in atomlib/atomcell.py

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

tail

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

Return the last n rows of self.

Source code in atomlib/atomcell.py

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

drop_nulls

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

Drop rows that contain nulls in any of columns subset.

Source code in atomlib/atoms.py

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

fill_null

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

Source code in atomlib/atomcell.py

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

fill_nan

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

Source code in atomlib/atomcell.py

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

concat classmethod

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

Concatenate multiple Atoms together, handling metadata appropriately.

Source code in atomlib/atoms.py

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

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

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

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

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

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

partition_by

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

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

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

Group by the given columns and partition into separate dataframes.

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

Source code in atomlib/atoms.py

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

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

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

select

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

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

Expressions may either be columns or expressions of columns.

Source code in atomlib/atomcell.py

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

    Expressions may either be columns or expressions of columns.

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

select_schema

select_schema(schema: SchemaDict) -> DataFrame

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

Source code in atomlib/atoms.py

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

select_props

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

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

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

Source code in atomlib/atomcell.py

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

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

try_select

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

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

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

Source code in atomlib/atomcell.py

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

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

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

try_get_column

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

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

Source code in atomlib/atoms.py

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

bbox_atoms

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

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

Source code in atomlib/atomcell.py

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

transform_atoms

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

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

Source code in atomlib/atomcell.py

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

transform

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

Source code in atomlib/atomcell.py

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

round_near_zero

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

Round atom position values near zero to zero.

Source code in atomlib/atomcell.py

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

crop_atoms

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

Source code in atomlib/atomcell.py

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

deduplicate

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

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

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

Source code in atomlib/atoms.py

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

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

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

    indices = numpy.arange(len(self))

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

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

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

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

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

with_bounds

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

Return a periodic cell with the given orthogonal cell dimensions.

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

Source code in atomlib/atoms.py

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

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

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

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

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

coords

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

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

Source code in atomlib/atomcell.py

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

x

x() -> Expr

Source code in atomlib/atoms.py

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

y

y() -> Expr

Source code in atomlib/atoms.py

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

z

z() -> Expr

Source code in atomlib/atoms.py

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

velocities

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

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

Source code in atomlib/atomcell.py

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

types

types() -> Optional[Series]

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

Source code in atomlib/atoms.py

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

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

masses

masses() -> Optional[Series]

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

Source code in atomlib/atoms.py

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

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

add_atom

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

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

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

Return a copy of self with an extra atom.

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

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

Source code in atomlib/atomcell.py

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

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

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

pos

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

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

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

Select all atoms at a given position.

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

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

Source code in atomlib/atoms.py

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

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

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

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

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

    return selection

with_index

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

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

Source code in atomlib/atomcell.py

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

with_wobble

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

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

Source code in atomlib/atomcell.py

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

with_occupancy

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

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

Source code in atomlib/atomcell.py

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

apply_wobble

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

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

Source code in atomlib/atomcell.py

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

apply_occupancy

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

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

Source code in atomlib/atoms.py

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

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

with_type

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

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

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

Source code in atomlib/atomcell.py

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

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

with_mass

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

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

Source code in atomlib/atomcell.py

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

with_symbol

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

Return self with the given atomic symbols.

Source code in atomlib/atomcell.py

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

with_coords

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

Return self replaced with the given atomic positions.

Source code in atomlib/atomcell.py

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

with_velocity

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

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

Source code in atomlib/atomcell.py

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

get_atomcell

get_atomcell() -> AtomCell

Source code in atomlib/atomcell.py

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

to_frame

to_frame(frame: CoordinateFrame) -> Self

Convert the stored Atoms to the given coordinate frame.

Source code in atomlib/atomcell.py

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

crop_to_box

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

Source code in atomlib/atomcell.py

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

wrap

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

Wrap atoms around the cell boundaries.

Source code in atomlib/atomcell.py

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

repeat_to

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

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

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

Source code in atomlib/atomcell.py

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

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

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

    return atom_cell

repeat_x

repeat_x(n: int) -> Self

Tile the cell in the x axis.

Source code in atomlib/atomcell.py

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

repeat_y

repeat_y(n: int) -> Self

Tile the cell in the y axis.

Source code in atomlib/atomcell.py

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

repeat_z

repeat_z(n: int) -> Self

Tile the cell in the z axis.

Source code in atomlib/atomcell.py

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

repeat_to_x

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

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

Source code in atomlib/atomcell.py

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

repeat_to_y

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

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

Source code in atomlib/atomcell.py

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

repeat_to_z

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

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

Source code in atomlib/atomcell.py

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

repeat_to_aspect

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

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

Source code in atomlib/atomcell.py

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

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

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

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

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

periodic_duplicate

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

Add duplicate copies of atoms near periodic boundaries.

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

Source code in atomlib/atomcell.py

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

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

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

    return self.to_frame(frame_save)

read classmethod

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

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

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

Read a structure from a file.

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

Source code in atomlib/mixins.py

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

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

read_cif classmethod

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

Read a structure from a CIF file.

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

Source code in atomlib/mixins.py

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

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

read_xyz classmethod

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

Read a structure from an XYZ file.

Source code in atomlib/mixins.py

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

read_xsf classmethod

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

Read a structure from an XSF file.

Source code in atomlib/mixins.py

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

read_cfg classmethod

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

Read a structure from a CFG file.

Source code in atomlib/mixins.py

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

read_lmp classmethod

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

Read a structure from a LAAMPS data file.

Source code in atomlib/mixins.py

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

write_cif

write_cif(f: FileOrPath)

Source code in atomlib/mixins.py

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

write_xyz

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

Source code in atomlib/mixins.py

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

write_xsf

write_xsf(f: FileOrPath)

Source code in atomlib/mixins.py

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

write_cfg

write_cfg(f: FileOrPath)

Source code in atomlib/mixins.py

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

write_lmp

write_lmp(f: FileOrPath)

Source code in atomlib/mixins.py

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

write

write(path: FileOrPath, ty: FileType)

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

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

Write this structure to a file.

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

Source code in atomlib/mixins.py

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

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

write_mslice

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

Write a structure to an mslice file.

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

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

Source code in atomlib/mixins.py

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

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

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

write_qe

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

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

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

Source code in atomlib/mixins.py

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

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

get_cell

get_cell() -> Cell

Source code in atomlib/atomcell.py

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

with_cell

with_cell(cell: Cell) -> Self

Source code in atomlib/atomcell.py

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

get_atoms

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

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

Source code in atomlib/atomcell.py

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

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

with_atoms

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

Source code in atomlib/atomcell.py

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

get_frame

get_frame() -> CoordinateFrame

Get the coordinate frame atoms are stored in.

Source code in atomlib/atomcell.py

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

from_ortho classmethod

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

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

Source code in atomlib/atomcell.py

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

from_unit_cell classmethod

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

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

Source code in atomlib/atomcell.py

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

orthogonalize

orthogonalize() -> OrthoCell

Source code in atomlib/atomcell.py

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

clone

clone() -> AtomCellT

Make a deep copy of self.

Source code in atomlib/atomcell.py

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

assert_equal

assert_equal(other: Any)

Assert this structure is equal to other

Source code in atomlib/atomcell.py

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

HasAtomCell

Bases: HasAtoms, HasCell, ABC

Source code in atomlib/atomcell.py

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

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

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

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

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

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

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

    # transformation

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

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

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

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

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

    # crop methods

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return atom_cell

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

        return self.to_frame(frame_save)

    # add frame to some HasAtoms methods

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

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

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

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

    with_column = with_columns

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

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

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

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

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

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

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

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

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

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

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

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

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

    # TODO: partition_by

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

        Expressions may either be columns or expressions of columns.

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

affine property

affine: AffineTransform3D

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

ortho property

ortho: LinearTransform3D

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

metric property

metric: LinearTransform3D

Cell metric tensor

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

cell_size property

cell_size: NDArray[float64]

Unit cell size.

cell_angle property

cell_angle: NDArray[float64]

Unit cell angles, in radians.

n_cells property

n_cells: NDArray[int_]

Number of unit cells.

pbc property

pbc: NDArray[bool_]

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

ortho_size property

ortho_size: NDArray[float64]

Return size of orthogonal unit cell.

Equivalent to the diagonal of the orthogonalization matrix.

box_size property

box_size: NDArray[float64]

Return size of the cell box.

Equivalent to self.n_cells * self.cell_size.

columns property

columns: List[str]

Return the column names in self.

RETURNS	DESCRIPTION
List[str]	A sequence of column names

dtypes property

dtypes: List[DataType]

Return the datatypes in self.

RETURNS	DESCRIPTION
List[DataType]	A sequence of column DataTypes

schema property

schema: Schema

Return the schema of self.

RETURNS	DESCRIPTION
Schema	A dictionary of column names and DataTypes

unique class-attribute instance-attribute

unique = deduplicate

with_column class-attribute instance-attribute

with_column = with_columns

get_cell abstractmethod

get_cell() -> Cell

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

Source code in atomlib/cell.py

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

get_transform

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

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

Source code in atomlib/cell.py

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

corners

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

Source code in atomlib/cell.py

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

bbox_cell

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

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

Source code in atomlib/cell.py

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

is_orthogonal

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

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

Source code in atomlib/cell.py

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

is_orthogonal_in_local

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

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

Source code in atomlib/cell.py

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

to_ortho

to_ortho() -> AffineTransform3D

Source code in atomlib/cell.py

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

strain_orthogonal

strain_orthogonal() -> HasCellT

Orthogonalize using strain.

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

Source code in atomlib/cell.py

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

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

explode_z

explode_z() -> HasCellT

Materialize repeated cells as one supercell in z.

Source code in atomlib/cell.py

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

change_transform

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

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

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

Coordinate-change a transformation from frame_from into frame_to.

Source code in atomlib/cell.py

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

assert_equal

assert_equal(other: Any)

Source code in atomlib/atoms.py

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

insert_column

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

Source code in atomlib/atoms.py

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

get_column_index

get_column_index(name: str) -> int

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

Source code in atomlib/atoms.py

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

clone

clone() -> DataFrame

Return a copy of self.

Source code in atomlib/atoms.py

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

drop

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

Return self with the specified columns removed.

Source code in atomlib/atoms.py

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

drop_nulls

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

Drop rows that contain nulls in any of columns subset.

Source code in atomlib/atoms.py

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

concat classmethod

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

Concatenate multiple Atoms together, handling metadata appropriately.

Source code in atomlib/atoms.py

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

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

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

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

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

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

partition_by

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

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

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

Group by the given columns and partition into separate dataframes.

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

Source code in atomlib/atoms.py

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

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

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

select_schema

select_schema(schema: SchemaDict) -> DataFrame

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

Source code in atomlib/atoms.py

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

try_get_column

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

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

Source code in atomlib/atoms.py

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

deduplicate

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

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

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

Source code in atomlib/atoms.py

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

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

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

    indices = numpy.arange(len(self))

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

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

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

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

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

with_bounds

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

Return a periodic cell with the given orthogonal cell dimensions.

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

Source code in atomlib/atoms.py

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

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

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

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

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

x

x() -> Expr

Source code in atomlib/atoms.py

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

y

y() -> Expr

Source code in atomlib/atoms.py

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

z

z() -> Expr

Source code in atomlib/atoms.py

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

types

types() -> Optional[Series]

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

Source code in atomlib/atoms.py

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

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

masses

masses() -> Optional[Series]

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

Source code in atomlib/atoms.py

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

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

pos

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

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

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

Select all atoms at a given position.

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

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

Source code in atomlib/atoms.py

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

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

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

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

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

    return selection

apply_occupancy

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

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

Source code in atomlib/atoms.py

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

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

get_frame abstractmethod

get_frame() -> CoordinateFrame

Get the coordinate frame atoms are stored in.

Source code in atomlib/atomcell.py

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

with_atoms abstractmethod

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

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

Source code in atomlib/atomcell.py

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

with_cell

with_cell(cell: Cell) -> Self

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

Source code in atomlib/atomcell.py

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

get_atomcell

get_atomcell() -> AtomCell

Source code in atomlib/atomcell.py

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

get_atoms abstractmethod

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

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

Source code in atomlib/atomcell.py

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

bbox_atoms

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

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

Source code in atomlib/atomcell.py

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

bbox

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

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

Source code in atomlib/atomcell.py

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

to_frame

to_frame(frame: CoordinateFrame) -> Self

Convert the stored Atoms to the given coordinate frame.

Source code in atomlib/atomcell.py

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

transform_atoms

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

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

Source code in atomlib/atomcell.py

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

transform_cell

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

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

Source code in atomlib/atomcell.py

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

transform

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

Source code in atomlib/atomcell.py

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

crop

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

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

To crop atoms only, use crop_atoms instead.

Source code in atomlib/atomcell.py

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

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

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

crop_atoms

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

Source code in atomlib/atomcell.py

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

crop_to_box

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

Source code in atomlib/atomcell.py

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

wrap

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

Wrap atoms around the cell boundaries.

Source code in atomlib/atomcell.py

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

repeat

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

Tile the cell

Source code in atomlib/atomcell.py

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

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

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

repeat_to

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

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

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

Source code in atomlib/atomcell.py

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

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

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

    return atom_cell

repeat_x

repeat_x(n: int) -> Self

Tile the cell in the x axis.

Source code in atomlib/atomcell.py

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

repeat_y

repeat_y(n: int) -> Self

Tile the cell in the y axis.

Source code in atomlib/atomcell.py

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

repeat_z

repeat_z(n: int) -> Self

Tile the cell in the z axis.

Source code in atomlib/atomcell.py

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

repeat_to_x

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

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

Source code in atomlib/atomcell.py

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

repeat_to_y

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

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

Source code in atomlib/atomcell.py

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

repeat_to_z

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

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

Source code in atomlib/atomcell.py

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

repeat_to_aspect

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

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

Source code in atomlib/atomcell.py

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

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

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

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

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

explode

explode() -> Self

Materialize repeated cells as one supercell.

Source code in atomlib/atomcell.py

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

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

periodic_duplicate

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

Add duplicate copies of atoms near periodic boundaries.

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

Source code in atomlib/atomcell.py

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

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

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

    return self.to_frame(frame_save)

describe

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

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

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

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

Source code in atomlib/atomcell.py

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

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

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

with_columns

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

Return a copy of self with the given columns added.

Source code in atomlib/atomcell.py

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

get_column

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

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

Source code in atomlib/atomcell.py

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

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

get_columns

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

Return all columns from self as a list of Series.

Source code in atomlib/atomcell.py

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

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

group_by

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

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

Source code in atomlib/atomcell.py

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

pipe

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

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

Source code in atomlib/atomcell.py

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

filter

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

Filter self, removing rows which evaluate to False.

Source code in atomlib/atomcell.py

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

sort

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

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

Source code in atomlib/atomcell.py

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

slice

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

Return a slice of the rows in self.

Source code in atomlib/atomcell.py

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

head

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

Return the first n rows of self.

Source code in atomlib/atomcell.py

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

tail

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

Return the last n rows of self.

Source code in atomlib/atomcell.py

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

fill_null

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

Source code in atomlib/atomcell.py

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

fill_nan

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

Source code in atomlib/atomcell.py

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

select

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

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

Expressions may either be columns or expressions of columns.

Source code in atomlib/atomcell.py

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

    Expressions may either be columns or expressions of columns.

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

select_props

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

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

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

Source code in atomlib/atomcell.py

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

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

try_select

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

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

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

Source code in atomlib/atomcell.py

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

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

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

round_near_zero

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

Round atom position values near zero to zero.

Source code in atomlib/atomcell.py

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

coords

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

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

Source code in atomlib/atomcell.py

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

velocities

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

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

Source code in atomlib/atomcell.py

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

add_atom

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

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

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

Return a copy of self with an extra atom.

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

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

Source code in atomlib/atomcell.py

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

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

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

with_index

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

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

Source code in atomlib/atomcell.py

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

with_wobble

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

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

Source code in atomlib/atomcell.py

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

with_occupancy

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

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

Source code in atomlib/atomcell.py

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

apply_wobble

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

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

Source code in atomlib/atomcell.py

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

with_type

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

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

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

Source code in atomlib/atomcell.py

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

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

with_mass

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

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

Source code in atomlib/atomcell.py

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

with_symbol

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

Return self with the given atomic symbols.

Source code in atomlib/atomcell.py

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

with_coords

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

Return self replaced with the given atomic positions.

Source code in atomlib/atomcell.py

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

with_velocity

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

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

Source code in atomlib/atomcell.py

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