schrodinger.seam.transforms.grouping module¶

class schrodinger.seam.transforms.grouping.GroupIntoBins(bin_key_func: Callable[[V], float], bin_width: int)¶

Bases: PTransform

Bins (key, value)’s into ((key, bin_index), [values]) tuples based on a given bin_width.

The “bin index” is calculated by 1) First passing the value into bin_key_fun to get a bin value, 2) Then, dividing the bin value by the bin width and rounding down to the nearest integer.

Parameters:

bin_key_func – A function that takes a value and returns a value to bin by. If not specified, the value itself is used.
bin_width – The width of each bin.

Example usage:

>>> import apache_beam as beam
>>> from schrodinger.application.transforms.dock import Dock
>>> from schrodinger.seam.io import chemio
>>> from schrodinger.test import mmshare_data_file
>>> with beam.Pipeline() as p:
...    input_data = [('a', 1), ('a', 2), ('a', 6), ('b', 5), ('b', 9)]
...    result = (
...        p
...        | beam.Create(input_data)
...        | GroupIntoBins(bin_key_func=lambda x: x, bin_width=3)
...        | beam.Map(print))
(('a', 0), [1, 2])
(('a', 2), [6])
(('b', 1), [5])
(('b', 3), [9])

__init__(bin_key_func: Callable[[V], float], bin_width: int)¶

expand(pcoll: PCollection[Tuple[K, V]]) → PCollection[Tuple[Tuple[K, int], Iterable[V]]]¶

class schrodinger.seam.transforms.grouping.GroupIntoBinsByMemoryConfig(*, target_bin_bytes: float, bytes_func: Callable[[Any], float], key_func: Callable[[Any], Any], max_bin_count: Optional[int] = None, min_bin_count: int = 1)¶

Bases: BaseModel

Validated configuration for GroupIntoBinsByMemory. See that transform for what each field means.

model_config: ClassVar[ConfigDict] = {'frozen': True}¶: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_bin_bytes: float¶

bytes_func: Callable[[Any], float]¶

key_func: Callable[[Any], Any]¶

max_bin_count: Optional[int]¶

min_bin_count: int¶

class schrodinger.seam.transforms.grouping.GroupIntoBinsByMemory(*, _config: pydantic.main.BaseModel | None = None, **kwargs)¶

Bases: PTransformWithConfig

Groups (key, value) pairs into ((key, bin_index), [values]) based on an estimate of how much memory the values in a bin will need.

A bin is cut when the estimated memory of its values reaches target_bin_bytes, so a bin holds many small values or few large ones, which keeps the memory needed to process one bin roughly constant regardless of how large the values are.

max_bin_count additionally caps the number of values in a bin. It is useful as a guard on how long a single bin takes to process, which target_bin_bytes alone does not bound. min_bin_count is the floor: a bin holds at least that many values even if they overrun target_bin_bytes, except for the last bin of a key group, which holds whatever values are left over. With the default of 1, a bin overruns the budget only when a single value exceeds it on its own, which no binning can avoid. A larger floor gives up that memory bound in exchange for bins big enough to be worth processing, so it should be set no higher than the available memory allows. Setting min_bin_count equal to max_bin_count gives fixed-size bins, ignoring the byte budget entirely.

This transform: 1) Adds a random UUID to each value’s sort key to break ties 2) Extracts sort keys and sizes from all values 3) Sorts the extracted keys to compute bin boundaries 4) Assigns each value to a bin based on these boundaries

This approach avoids loading all values into memory for sorting - only the extracted keys and sizes need to be sorted, which is much more memory-efficient for large objects like Structure instances. The UUID also means that the limits are strictly enforced even when there are ties in key_func.

Parameters:

target_bin_bytes – Estimated memory a bin’s values should add up to. This is a target rather than a cap, since min_bin_count may push a bin over it.
bytes_func – A function that takes a value and returns its estimated memory in bytes.
key_func – A function that takes a value and returns a sortable key.
max_bin_count – Maximum number of values per bin, if any.
min_bin_count – Minimum number of values per bin.

Example usage:

>>> import apache_beam as beam
>>> with beam.Pipeline() as p:
...     # Each value is (energy, estimated_bytes).
...     input_data = [
...         ('grp', (1.0, 40)),
...         ('grp', (2.0, 40)),
...         ('grp', (3.0, 30)),
...         ('grp', (4.0, 150)),
...     ]
...     result = (
...         p
...         | beam.Create(input_data)
...         | GroupIntoBinsByMemory(
...             target_bin_bytes=100,
...             bytes_func=lambda x: x[1],
...             key_func=lambda x: x[0])
...         | beam.Map(print))
(('grp', 0), [(1.0, 40), (2.0, 40)])
(('grp', 1), [(3.0, 30)])
(('grp', 2), [(4.0, 150)])  # oversized, so it overruns the budget

config_class¶: alias of GroupIntoBinsByMemoryConfig

expand(pcoll)¶

class schrodinger.seam.transforms.grouping.UnbinAndPartitionBoundaryValues(k: int, key_func: Optional[Callable[[V], float]] = None)¶

Bases: PTransform

Unbins and partitions the input into two pcollections of boundary and non-boundary values.

The boundary values are the first and last k values in each bin after sorting. The non-boundary values are the rest of the values in each bin.

Example usage:

>>> import apache_beam as beam
>>> from schrodinger.seam.transforms.grouping import UnbinAndPartitionBoundaryValues
>>> with beam.Pipeline() as p:
...     bins = [
...         ('a', [1, 2, 3])
...     ]
...     pcoll = p | beam.Create(bins)
...     split = pcoll | UnbinAndPartitionBoundaryValues(k=1)
...     _ = split.boundary_values | beam.Map(lambda x: print("Boundary", x))
...     _ = split.nonboundary_values | beam.Map(lambda x: print("NonBoundary", x))
Boundary ('a', 1)
NonBoundary ('a', 2)
Boundary ('a', 3)

BOUNDARY_TAG = 'boundary_values'¶

NON_BOUNDARY_TAG = 'nonboundary_values'¶

__init__(k: int, key_func: Optional[Callable[[V], float]] = None)¶

Parameters:

k – Used to define the first and last k elements as boundary values
key_func – Sorting key used to sort a list of values

expand(pcoll)¶

class schrodinger.seam.transforms.grouping.AssignNearestBinKey(bin_key_func: Callable[[V], float], bin_width: float)¶

Bases: PTransform

Assigns the nearest bin key to each element in the input PCollection.

The bin key is calculated by dividing the value by the bin width and rounding to the nearest integer.

Parameters:

bin_key_func – A function that takes a value and returns a value to bin by. If not specified, the value itself is used.
bin_width – The width of each bin.

Example usage:

>>> import apache_beam as beam
>>> from schrodinger.seam.transforms.grouping import AssignNearestBinKey
>>> with beam.Pipeline() as p:
...     input_data = [('a', 7), ('a', 12), ('a', 22), ('b', 4), ('b', 25)]
...     result = (
...         p
...         | beam.Create(input_data)
...         | AssignNearestBinKey(bin_key_func=lambda x: x, bin_width=10)
...     )
...     _ = result | beam.Map(print)
(('a', 1), 7)
(('a', 1), 12)
(('a', 2), 22)
(('b', 0), 4)
(('b', 2), 25)

__init__(bin_key_func: Callable[[V], float], bin_width: float)¶

expand(pcoll: PCollection[Tuple[K, V]]) → PCollection[Tuple[Tuple[K, int], V]]¶

class schrodinger.seam.transforms.grouping.UnbinAndPartitionBoundaryValuesByMemoryConfig(*, target_boundary_bytes: float, bytes_func: ~typing.Callable[[~typing.Any], float], key_func: ~typing.Callable[[~typing.Any], ~typing.Any] = <function UnbinAndPartitionBoundaryValuesByMemoryConfig.<lambda>>, max_boundary_count: ~typing.Optional[int] = None, min_boundary_count: int = 1)¶

Bases: BaseModel

Validated configuration for UnbinAndPartitionBoundaryValuesByMemory. See that transform for what each field means.

model_config: ClassVar[ConfigDict] = {'frozen': True}¶: Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

target_boundary_bytes: float¶

bytes_func: Callable[[Any], float]¶

key_func: Callable[[Any], Any]¶

max_boundary_count: Optional[int]¶

min_boundary_count: int¶

class schrodinger.seam.transforms.grouping.UnbinAndPartitionBoundaryValuesByMemory(*, _config: pydantic.main.BaseModel | None = None, **kwargs)¶

Bases: PTransformWithConfig

Unbins and partitions the input into boundary and non-boundary values, sizing each boundary by the estimated memory of its values.

Each bin’s values are sorted by key_func, and its lowest and highest values become boundary values. Indices are remapped so that bins and boundaries never share one:

bin N, sorted:  [ low boundary | middle | high boundary ]
output index:         2N         2N + 1       2N + 2

so boundary index 2N + 2 collects the high boundary of bin N and the low boundary of bin N + 1. Boundary values go to boundary_values and middle values to nonboundary_values.

A boundary takes values until target_boundary_bytes or max_boundary_count would be exceeded, but always at least min_boundary_count. When a bin’s two boundaries would overlap, the overlapping values go only to the low boundary, so a bin that small sends nothing to index 2N + 2 (see DESMOND-17832).

Parameters:

target_boundary_bytes – Estimated memory one boundary’s values should add up to; a boundary index collects two boundaries. This is a target rather than a cap, since min_boundary_count may push a boundary over it.
bytes_func – A function that takes a value and returns its estimated memory in bytes.
key_func – Sorting key used to sort values within each bin.
max_boundary_count – Maximum number of values at one boundary, if any.
min_boundary_count – Minimum number of values at one boundary.

Example usage:

>>> import apache_beam as beam
>>> with beam.Pipeline() as p:
...     # Each value is (energy, estimated_bytes).
...     bins = [
...         (('a', 0), [(1.0, 30), (2.0, 30), (3.0, 30), (4.0, 30)]),
...     ]
...     pcoll = p | beam.Create(bins)
...     split = pcoll | UnbinAndPartitionBoundaryValuesByMemory(
...         target_boundary_bytes=60,
...         bytes_func=lambda x: x[1],
...         key_func=lambda x: x[0])
...     _ = split.boundary_values | beam.Map(lambda x: print("Boundary", x))
...     _ = split.nonboundary_values | beam.Map(lambda x: print("NonBoundary", x))
Boundary (('a', 0), (1.0, 30))
Boundary (('a', 0), (2.0, 30))
Boundary (('a', 2), (3.0, 30))
Boundary (('a', 2), (4.0, 30))

BOUNDARY_TAG = 'boundary_values'¶

NON_BOUNDARY_TAG = 'nonboundary_values'¶

config_class¶: alias of UnbinAndPartitionBoundaryValuesByMemoryConfig

expand(pcoll)¶