Query & Selection

See Atom Selection Syntax for a narrative guide with rendered examples.

class atomworks.io.utils.query.AtomSelection(chain_id: str = '*', res_name: str = '*', res_id: int | str = '*', atom_name: str = '*', transformation_id: int | str = '*')

Bases: object

A single-valued selection of atoms in a molecular structure.

A selection is specified by chain_id, res_name, res_id, atom_name, and (optionally) transformation_id. Each field is either an exact value or the wildcard "*" (match anything). Fields are combined with logical AND.

For example:

• specifying only chain_id selects all atoms in that chain

• specifying chain_id and res_name selects all atoms of that residue type in that chain

• specifying only atom_name selects all atoms with that name, in any chain/residue

For multi-valued selections (lists, ranges, unions), use AtomSelectionStack.

classmethod from_pymol_str(pymol_string: str) → AtomSelection

Create a selection from a PyMOL atom label CHAIN/RES_NAME`RES_ID/ATOM.

Such strings are produced by clicking an atom/residue in PyMOL, e.g. "A/ASP`37/OD2". "*" may be used as a wildcard. transformation_id is not supported by PyMOL strings.

classmethod from_selection_str(selection_string: str) → AtomSelection

Create a selection from CHAIN/RES_NAME/RES_ID/ATOM/TRANSFORM syntax.

"*" acts as a wildcard for any field; trailing fields may be omitted and default to "*". Bracket-list syntax ([...]) is not accepted here - use AtomSelectionStack.from_query() or the query sel('...') function for multi-valued selections.

Examples

>>> AtomSelection.from_selection_str("A/ALA/1/CA")
A/ALA/1/CA
>>> AtomSelection.from_selection_str("*/ALA/*/CB")
*/ALA/*/CB

get_idxs(atom_array: AtomArray) → ndarray

Get the indices of atoms selected by this AtomSelection.

get_mask(atom_array: AtomArray, raise_on_empty: bool = True) → ndarray

Create a boolean mask for this selection on atom_array.

Parameters:

raise_on_empty – If True (default), raise ValueError when the selection matches no atoms. Set False for union/composition contexts (used by sel('...')).

class atomworks.io.utils.query.AtomSelectionStack(selections: list[AtomSelection])

Bases: object

A union (logical OR) of AtomSelection objects.

Enables a single string to select multiple atoms/segments via from_query() (extended syntax with [...] lists and ranges) or from_contig() (contiguous residue ranges).

classmethod from_contig(contig: str) → AtomSelectionStack

Create a stack from contiguous residue ranges like "A1-2, B3-10".

classmethod from_query(query: str | list[str]) → AtomSelectionStack

Create a stack from the extended path-selection syntax.

Grammar (fields in order CHAIN/RES_NAME/RES_ID/ATOM/TRANSFORM, trailing fields default to "*"):

• Each field is a scalar (A), the wildcard *, or a bracket list [a, b, ...].

• Inside a bracket list, each item is a scalar or - for res_id only - an inclusive range lo-hi (e.g. [1-5, 9, 12-14]). Ranges must be bracketed; a bare 5-10 is not a range.

• Multiple whole tokens, separated by top-level commas (or given as a list[str]), are unioned. Commas inside [...] are part of the list, not token separators.

Examples

>>> AtomSelectionStack.from_query("[A,B]/ALA/1/[CA,CB]")  # 4 selections, unioned
>>> AtomSelectionStack.from_query("A/*/[5-10]")  # residues 5..10 in chain A
>>> AtomSelectionStack.from_query("A/*/[5-10], B/*/[3-8]")  # different range per chain

get_center_of_mass(atom_array: AtomArray | AtomArrayStack) → ndarray

Return the center of mass of the selected atoms.

get_mask(atom_array: AtomArray | AtomArrayStack, raise_on_empty: bool = True) → ndarray

Create a boolean mask by unioning (logical OR) all member selections.

Parameters:

raise_on_empty – Passed through to each member selection. If True (default), a member matching no atoms raises ValueError. sel('...') sets this to False so unions degrade gracefully.

get_principal_components(atom_array: AtomArray | AtomArrayStack) → ndarray

Return principal axes (eigenvectors) of the selected atoms via SVD.

Returns:

(3, 3) array for AtomArray. (n_models, 3, 3) array for AtomArrayStack.

class atomworks.io.utils.query.QueryExpression(expr: str)

Bases: object

Query evaluator for biotite AtomArrays using pandas-like syntax.

Examples

Select all CA atoms in chain A:

>>> expr = QueryExpression("(chain_id == 'A') & (atom_name == 'CA')")
>>> ca_atoms = expr.query(atom_array)

Select atoms without NaN coordinates:

>>> expr = QueryExpression("~has_nan_coord()")
>>> valid_atoms = expr.query(atom_array)

Select bonded atoms in specific residues:

>>> expr = QueryExpression("has_bonds() & (res_name in ['ALA', 'GLY', 'VAL'])")

Combine path-selection syntax with predicates via sel('...'):

>>> expr = QueryExpression("sel('[A,B]/ALA') & (x > 0)")

OPS = mappingproxy({<class 'ast.Eq'>: <built-in function eq>, <class 'ast.NotEq'>: <built-in function ne>, <class 'ast.Lt'>: <built-in function lt>, <class 'ast.LtE'>: <built-in function le>, <class 'ast.Gt'>: <built-in function gt>, <class 'ast.GtE'>: <built-in function ge>, <class 'ast.In'>: None, <class 'ast.NotIn'>: None, <class 'ast.And'>: <ufunc 'logical_and'>, <class 'ast.Or'>: <ufunc 'logical_or'>, <class 'ast.Not'>: <ufunc 'logical_not'>, <class 'ast.BitAnd'>: <ufunc 'bitwise_and'>, <class 'ast.BitOr'>: <ufunc 'bitwise_or'>, <class 'ast.Invert'>: <ufunc 'invert'>, <class 'ast.UAdd'>: <built-in function pos>, <class 'ast.USub'>: <built-in function neg>})

STRING_ARG_FUNCTIONS = frozenset({'sel'})

idxs(atom_array: AtomArray | AtomArrayStack) → ndarray

Apply the query expression to an AtomArray and return the indices of the matching atoms.

Parameters:

atom_array – The atom array to query.

Returns:

Numpy array of indices for atoms that match the query expression.

mask(atom_array: AtomArray | AtomArrayStack) → ndarray

Apply the query expression to an AtomArray and return a boolean mask.

Parameters:

atom_array – The atom array to query.

Returns:

Boolean numpy array indicating which atoms match the query.

query(atom_array: AtomArray | AtomArrayStack) → AtomArray | AtomArrayStack

Apply the query expression to an AtomArray and return a filtered AtomArray.

Parameters:

atom_array – The atom array to query.

Returns:

Filtered atom array containing only atoms that match the query expression.

atomworks.io.utils.query.get_mask_from_atom_selection(atom_array: AtomArray, atom_selection: AtomSelection, raise_on_empty: bool = True) → ndarray

Create a boolean mask from an AtomSelection.

Parameters:

raise_on_empty – If True (default), raise ValueError when no atoms match.

atomworks.io.utils.query.get_mask_from_selection_string(atom_array: AtomArray, selection_string: str) → ndarray

Create a boolean mask from a CHAIN/RES_NAME/RES_ID/ATOM/TRANSFORM selection string.

"*" acts as a wildcard for any field.

atomworks.io.utils.query.idxs(atom_array: AtomArray | AtomArrayStack, expr: str) → ndarray

Query the AtomArray using pandas-like syntax and return the indices of the matching atoms.

atomworks.io.utils.query.mask(atom_array: AtomArray | AtomArrayStack, expr: str) → ndarray

Query the AtomArray using pandas-like syntax and return a boolean mask.

atomworks.io.utils.query.parse_pymol_string(pymol_string: str) → AtomSelection

Parse a PyMOL CHAIN/RES_NAME`RES_ID/ATOM string into an AtomSelection.

Wildcards ("*") are supported; transformation_id is not.

atomworks.io.utils.query.parse_selection_string(selection_string: str) → AtomSelection

Parse a CHAIN/RES_NAME/RES_ID/ATOM/TRANSFORM string into an AtomSelection.

"*" acts as a wildcard for any field. Trailing fields may be omitted and default to "*".

Raises:

ValueError – If bracket-list syntax [...] is used (multi-valued selections must go through AtomSelectionStack.from_query()).

atomworks.io.utils.query.query(atom_array: AtomArray | AtomArrayStack, expr: str) → AtomArray | AtomArrayStack

Query the AtomArray using pandas-like syntax. :param atom_array: The atom array to query. :param expr: Query expression in pandas-like syntax.

Returns:

Filtered atom array containing only atoms that match the query expression.

Examples

>>> # Select all CA atoms in chain A
>>> ca_atoms = query(atom_array, "(chain_id == 'A') & (atom_name == 'CA')")

>>> # Select atoms without NaN coordinates
>>> valid_atoms = query(atom_array, "~has_nan_coord()")

>>> # Select bonded atoms in specific residues
>>> bonded = query(atom_array, "has_bonds() & (res_name in ['ALA', 'GLY', 'VAL'])")