Skip to content

Commit

Permalink
fixup
Browse files Browse the repository at this point in the history
  • Loading branch information
hexane360 committed Apr 20, 2023
1 parent bccdc25 commit 0d21f28
Showing 1 changed file with 13 additions and 3 deletions.
16 changes: 13 additions & 3 deletions structlib/atoms.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,16 @@ def _selection_to_expr(selection: AtomSelection) -> polars.Expr:
return _values_to_expr(selection, ty=polars.Boolean)


def _selection_to_numpy(df: t.Union[polars.DataFrame, HasAtoms], selection: AtomSelection) -> NDArray[numpy.bool_]:
series = _selection_to_series(df, selection)
try:
import pyarrow
except ImportError:
# workaround without pyarrow
return numpy.array(series.to_list(), dtype=numpy.bool_)
return series.to_numpy(zero_copy_only=False)


def _select_schema(df: t.Union[polars.DataFrame, HasAtoms], schema: SchemaDict) -> polars.DataFrame:
"""
Select columns from ``self`` and cast to the given schema.
Expand Down Expand Up @@ -593,7 +603,7 @@ def with_symbol(self: HasAtomsT, symbols: ArrayLike, selection: t.Optional[AtomS
Return `self` with the given atomic symbols.
"""
if selection is not None:
selection = _selection_to_series(self, selection).to_numpy()
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
Expand All @@ -607,7 +617,7 @@ def with_coords(self: HasAtomsT, pts: ArrayLike, selection: t.Optional[AtomSelec
Return `self` replaced with the given atomic positions.
"""
if selection is not None:
selection = _selection_to_series(self, selection).to_numpy()
selection = _selection_to_numpy(self, selection)
new_pts = self.coords()
pts = numpy.atleast_2d(pts)
assert pts.shape[-1] == 3
Expand Down Expand Up @@ -778,4 +788,4 @@ def _repr_pretty_(self, p, cycle: bool) -> None:

__all__ = [
'Atoms', 'HasAtoms', 'IntoAtoms', 'AtomSelection', 'AtomValues',
]
]

0 comments on commit 0d21f28

Please sign in to comment.