Sampling

Once a query step has filtered the query database down to the rows you care about — see Query Examples for the SQL filter syntax — its sampling block lets you pick a subset of those rows to keep. In basic usage, method is set to random or all; for more advanced use cases, it can instead point at a Python file to run your own selection strategy — for example, greedy diversity sampling based on Tanimoto distance, shown below as a worked example.

Basic Usage

The sampling block controls both how many of the filtered rows are kept and how they are chosen. method selects the strategy — either a built-in name or a path to a plugin file — and the rest of the block supplies whatever that strategy needs.

For the built-in strategies, that’s just n_samples and an optional seed:

sampling:
  n_samples: 100000
  method: random
  seed: 42 # for reproducibility

method: random draws a uniform random subset of size n_samples; method: all skips sampling entirely and keeps every filtered row. To use your own selection logic instead, point method at a Python file and pass it whatever columns and keyword arguments it needs:

sampling:
  method: /path/to/custom_sampling.py
  required_columns: ["num_atoms", "energy"]
  kwargs:
    some_kwarg: 42

• method — random, all, or a path to a Python file defining a custom_sampling_function function.

• n_samples — number of rows to keep. Applies to random and to custom samplers that respect it; omit it (or set it to null) to let a custom sampler decide how many rows to keep on its own.

• seed — random seed used only for method: random.

• required_columns — extra column name(s) to load and pass to a custom sampler alongside entry_key, in this example, number of atoms and energy. Ignored by the built-in methods.

• kwargs — any additional keyword arguments forwarded to a custom sampler’s custom_sampling_function function.

Writing a Sampling Plugin

Create a standard Python file (e.g. my_sampler.py) that defines a top-level custom_sampling_function function:

# my_sampler.py
import polars as pl


def custom_sampling_function(
    frame: pl.LazyFrame | pl.DataFrame,
    n_samples: int | None,
    some_kwarg: int,
) -> list[str]:
    """Return the entry_key of every selected row.

    `some_kwarg` is an example for a keyword argument (here, not actually used below).
    """
    collected = frame.collect() if isinstance(frame, pl.LazyFrame) else frame
    return collected["entry_key"].to_list()[:n_samples]

frame is the filtered query result (with entry_key plus whatever columns you requested via required_columns); custom_sampling_function must return the list of selected entry_key values. The CLI loads the file with runpy.run_path and raises FileNotFoundError/ValueError if the file or the custom_sampling_function function is missing. The function name and signature shown above must be followed exactly — the CLI calls custom_sampling_function(frame, n_samples, **kwargs) positionally, so renaming it or reordering its parameters will break the plugin.

Note

The plugin mechanism mirrors the “Custom I/O Plugins” section of the CLI guide: a plain Python file loaded dynamically, no package installation required.

Example: Greedy Tanimoto Sampling

greedy_sampling.py is primarily meant as a worked example of a custom sampling plugin: it shows how to write a custom_sampling_function that streams through a large filtered result set in chunks and selects rows based on a pairwise distance computed between them. That said, it is also directly usable if greedy diversity sampling fits your use case, not just a reference to copy from.

The plugin greedily selects a chemically diverse subset based on molecular fingerprint bit vectors: to compute fingerprint-based similarity, we use the Tanimoto index[1]: starting from one seed structure, it repeatedly adds the next structure whose Tanimoto distance to every structure already selected exceeds a threshold, streaming through the filtered rows in chunks so it scales to large databases. It relies on JAX for the distance computation, so make sure jax is installed in your environment; it is not a dependency of ChemReporter itself.

Here is an example sampling config using it:

sampling:
  method: /path/to/greedy_sampling.py
  required_columns: fingerprint_bits
  n_samples: null
  kwargs:
    seed: 42
    chunk_size: 25000
    min_distance_threshold: 0.65

• n_samples: null keeps adding structures until no remaining candidate clears the threshold, rather than stopping at a fixed count. Set it to an integer to cap the subset size at that count instead.

• required_columns: fingerprint_bits — fingerprint_bits is not itself a query database column; it’s a special preset name that expands to every individual fingerprint_* bit column (fingerprint_1, fingerprint_2, …) needed to compute Tanimoto distances, so you don’t have to list them all yourself.

• min_distance_threshold — minimum Tanimoto distance (1 - similarity) a candidate must have to every structure already selected. Raise it for a more diverse (and smaller) subset, lower it to keep more structures.

• chunk_size — number of rows processed per streaming batch. Lower it if you are memory-constrained.

• seed — used to pick the random starting structure when you don’t specify a start_idx yourself.

---

[1]

Why Tanimoto? Bajusz et al. demonstrated that the Tanimoto index, alongside Dice and Cosine metrics, consistently yields optimal similarity rankings for binary molecular fingerprints across unbiased compound datasets (Bajusz et al., 2015).