Selector Tutorial¶

Learn how to select atoms and filter data using MolPy's powerful selector system! Selectors let you create boolean masks and filter Block objects with ease.

What are Selectors?¶

Selectors are composable predicates that create boolean masks for filtering data:

• AtomTypeSelector: Select by atom type
• ElementSelector: Select by element symbol
• AtomIndexSelector: Select by atom indices
• CoordinateRangeSelector: Select by coordinate ranges
• DistanceSelector: Select by distance from a point

You can combine them with & (and), | (or), and ~ (not) operators!

In [1]:

import molpy as mp
from molpy.core.selector import (
    AtomIndexSelector,
    AtomTypeSelector,
    CoordinateRangeSelector,
    DistanceSelector,
    ElementSelector,
)

import molpy as mp from molpy.core.selector import ( AtomIndexSelector, AtomTypeSelector, CoordinateRangeSelector, DistanceSelector, ElementSelector, )

Basic Selectors¶

Let's start with simple selectors:

In [2]:

# Create a frame with atoms
frame = mp.Frame()
frame["atoms"] = mp.Block(
    {
        "x": [0.0, 1.0, 2.0, 3.0],
        "y": [0.0, 0.0, 0.0, 0.0],
        "z": [0.0, 0.0, 0.0, 0.0],
        "element": ["C", "C", "H", "H"],
        "type": [1, 1, 2, 2],
    }
)

atoms = frame["atoms"]
print(f"Total atoms: {atoms.nrows}")

# Create a frame with atoms frame = mp.Frame() frame["atoms"] = mp.Block( { "x": [0.0, 1.0, 2.0, 3.0], "y": [0.0, 0.0, 0.0, 0.0], "z": [0.0, 0.0, 0.0, 0.0], "element": ["C", "C", "H", "H"], "type": [1, 1, 2, 2], } ) atoms = frame["atoms"] print(f"Total atoms: {atoms.nrows}")

Total atoms: 4

Select by Element¶

In [3]:

# Select all carbon atoms
carbon_sel = ElementSelector("C")
carbon_mask = carbon_sel.mask(atoms)
carbon_atoms = carbon_sel(atoms)  # or atoms[carbon_mask]

print(f"Carbon mask: {carbon_mask}")
print(f"Carbon atoms: {carbon_atoms.nrows}")
print(f"Carbon elements: {carbon_atoms['element']}")

# Select all carbon atoms carbon_sel = ElementSelector("C") carbon_mask = carbon_sel.mask(atoms) carbon_atoms = carbon_sel(atoms) # or atoms[carbon_mask] print(f"Carbon mask: {carbon_mask}") print(f"Carbon atoms: {carbon_atoms.nrows}") print(f"Carbon elements: {carbon_atoms['element']}")

Carbon mask: [ True  True False False]
Carbon atoms: 2
Carbon elements: ['C' 'C']

Select by Atom Type¶

In [4]:

# Select atoms with type 1
type1_sel = AtomTypeSelector(1)
type1_atoms = type1_sel(atoms)

print(f"Type 1 atoms: {type1_atoms.nrows}")
print(f"Type 1 elements: {type1_atoms['element']}")

# Select atoms with type 1 type1_sel = AtomTypeSelector(1) type1_atoms = type1_sel(atoms) print(f"Type 1 atoms: {type1_atoms.nrows}") print(f"Type 1 elements: {type1_atoms['element']}")

Type 1 atoms: 2
Type 1 elements: ['C' 'C']

Select by Index¶

In [5]:

# Select atoms by indices
index_sel = AtomIndexSelector([0, 2])
selected_atoms = index_sel(atoms)

print(f"Selected atoms: {selected_atoms.nrows}")
print(f"Selected elements: {selected_atoms['element']}")

# Select atoms by indices index_sel = AtomIndexSelector([0, 2]) selected_atoms = index_sel(atoms) print(f"Selected atoms: {selected_atoms.nrows}") print(f"Selected elements: {selected_atoms['element']}")

Selected atoms: 0
Selected elements: []

Coordinate-Based Selectors¶

In [6]:

# Select atoms in a coordinate range (for x axis)
x_range_sel = CoordinateRangeSelector(axis="x", min_value=0.5, max_value=2.5)
range_atoms = x_range_sel(atoms)

print(f"Atoms in x range: {range_atoms.nrows}")
print(f"X coordinates: {range_atoms['x']}")

# Combine multiple axis selectors
y_range_sel = CoordinateRangeSelector(axis="y", min_value=-0.5, max_value=0.5)
combined_sel = x_range_sel & y_range_sel
combined_atoms = combined_sel(atoms)
print(f"Atoms in x and y range: {combined_atoms.nrows}")

# Select atoms in a coordinate range (for x axis) x_range_sel = CoordinateRangeSelector(axis="x", min_value=0.5, max_value=2.5) range_atoms = x_range_sel(atoms) print(f"Atoms in x range: {range_atoms.nrows}") print(f"X coordinates: {range_atoms['x']}") # Combine multiple axis selectors y_range_sel = CoordinateRangeSelector(axis="y", min_value=-0.5, max_value=0.5) combined_sel = x_range_sel & y_range_sel combined_atoms = combined_sel(atoms) print(f"Atoms in x and y range: {combined_atoms.nrows}")

Atoms in x range: 2
X coordinates: [1. 2.]
Atoms in x and y range: 2

Distance-Based Selection¶

In [7]:

# Select atoms within 1.5 Angstrom of origin
distance_sel = DistanceSelector(center=[0.0, 0.0, 0.0], max_distance=1.5)
nearby_atoms = distance_sel(atoms)

print(f"Atoms within 1.5 Å: {nearby_atoms.nrows}")

# Select atoms in a shell (between min and max distance)
shell_sel = DistanceSelector(center=[1.0, 0.0, 0.0], min_distance=0.5, max_distance=2.0)
shell_atoms = shell_sel(atoms)
print(f"Atoms in shell: {shell_atoms.nrows}")

# Select atoms within 1.5 Angstrom of origin distance_sel = DistanceSelector(center=[0.0, 0.0, 0.0], max_distance=1.5) nearby_atoms = distance_sel(atoms) print(f"Atoms within 1.5 Å: {nearby_atoms.nrows}") # Select atoms in a shell (between min and max distance) shell_sel = DistanceSelector(center=[1.0, 0.0, 0.0], min_distance=0.5, max_distance=2.0) shell_atoms = shell_sel(atoms) print(f"Atoms in shell: {shell_atoms.nrows}")

Atoms within 1.5 Å: 2
Atoms in shell: 3

Combining Selectors¶

The real power comes from combining selectors:

In [8]:

# Carbon atoms AND type 1
carbon_type1 = ElementSelector("C") & AtomTypeSelector(1)
result = carbon_type1(atoms)
print(f"Carbon with type 1: {result.nrows} atoms")

# Carbon OR hydrogen
c_or_h = ElementSelector("C") | ElementSelector("H")
result = c_or_h(atoms)
print(f"Carbon or hydrogen: {result.nrows} atoms")

# NOT carbon (everything except carbon)
not_carbon = ~ElementSelector("C")
result = not_carbon(atoms)
print(f"Not carbon: {result.nrows} atoms")
print(f"Elements: {result['element']}")

# Carbon atoms AND type 1 carbon_type1 = ElementSelector("C") & AtomTypeSelector(1) result = carbon_type1(atoms) print(f"Carbon with type 1: {result.nrows} atoms") # Carbon OR hydrogen c_or_h = ElementSelector("C") | ElementSelector("H") result = c_or_h(atoms) print(f"Carbon or hydrogen: {result.nrows} atoms") # NOT carbon (everything except carbon) not_carbon = ~ElementSelector("C") result = not_carbon(atoms) print(f"Not carbon: {result.nrows} atoms") print(f"Elements: {result['element']}")

Carbon with type 1: 2 atoms
Carbon or hydrogen: 4 atoms
Not carbon: 2 atoms
Elements: ['H' 'H']

Complex Combinations¶

You can build complex selection logic:

In [9]:

# Complex selection: (Carbon OR type 1) AND within distance
complex_sel = (ElementSelector("C") | AtomTypeSelector(1)) & DistanceSelector(
    center=[1.0, 0.0, 0.0], max_distance=2.0
)
result = complex_sel(atoms)
print(f"Complex selection: {result.nrows} atoms")

# Complex selection: (Carbon OR type 1) AND within distance complex_sel = (ElementSelector("C") | AtomTypeSelector(1)) & DistanceSelector( center=[1.0, 0.0, 0.0], max_distance=2.0 ) result = complex_sel(atoms) print(f"Complex selection: {result.nrows} atoms")

Complex selection: 2 atoms

Using with Frames¶

Selectors work directly with Frame blocks:

In [10]:

# Apply selector to frame atoms
carbon_atoms = ElementSelector("C")(frame["atoms"])

# Create a new frame with selected atoms
new_frame = mp.Frame()
new_frame["atoms"] = carbon_atoms
new_frame.metadata["box"] = frame.metadata.get("box")

print(f"New frame has {new_frame['atoms'].nrows} atoms")

# Apply selector to frame atoms carbon_atoms = ElementSelector("C")(frame["atoms"]) # Create a new frame with selected atoms new_frame = mp.Frame() new_frame["atoms"] = carbon_atoms new_frame.metadata["box"] = frame.metadata.get("box") print(f"New frame has {new_frame['atoms'].nrows} atoms")

New frame has 2 atoms