schrodinger.application.matsci.flywheel.diff_graph_search_trainer module

DiffGraph Search for molecular alignment.

This algorithm is a property-guided variation of TreeG-SC. It directly scores decoded molecules at each reverse-diffusion step and ranks every decoder sample individually instead of averaging Monte Carlo scores.

Copyright Schrodinger, LLC. All rights reserved.

class schrodinger.application.matsci.flywheel.diff_graph_search_trainer.DiffGraphSearchTrainer(generator, screener, text_prompt, num_sampling_steps, branch_factor=4, mc_samples=1, elite_reseed_fraction=0.2, **kwargs)

Bases: BaseTrainer

DiffGraph Search with direct ML/QSAR branch scoring.

NAME = 'DiffGraph Search'
__init__(generator, screener, text_prompt, num_sampling_steps, branch_factor=4, mc_samples=1, elite_reseed_fraction=0.2, **kwargs)
Parameters:

  • generator (LDMolGenerator) – Frozen LDMol molecule generator

  • screener (ScreenerList) – Direct molecular property evaluator

  • text_prompt (str) – LDMol text condition

  • num_sampling_steps (int) – Number of reverse diffusion steps

  • branch_factor (int) – Total transitions per active state; one is deterministic and the remainder are stochastic

  • mc_samples (int) – Additional stochastic decoder samples per transition; one ordinary decode is always included

  • elite_reseed_fraction (float) – Active-state fraction reserved for non-feasible refined seeds

  • kwargs (dict) – Base trainer keyword arguments

buildTraversalSeedPool(active_count)

Build one bounded batch of new and refined molecular seeds.

updateTopSearchScores(score_items)

Retain the best score for the top unique molecules.

Parameters:

score_items (iterable(tuple(str, float))) – Molecular identity and score pairs

accumulateInitialBatch(smiles_list, latents)

Cache valid initial molecules for guided-diffusion seeding.

Parameters:

  • smiles_list (list(str)) – Initial decoded molecules

  • latents (numpy.ndarray) – Latents aligned with smiles_list

scoreCandidates(smiles_list)

Evaluate decoded molecules for property-guided branch selection.

Parameters:

smiles_list (list(str)) – Decoded candidates in search order

Return type:

ldmol_generator.SearchScoreResult

Returns:

Position-aligned scores and molecular identities

logRefinementStep(decode_evaluations)

Log metadata for the current property-guided refinement step.

Parameters:

decode_evaluations (int) – Number of decoded candidates in this step

getCandidateScoreResult(smiles_list, candidate_records)

Convert screened candidates into a position-aligned score result.

Parameters:

  • smiles_list (list(str)) – Decoded candidates in search order

  • candidate_records (dict(str, CandidateRecord)) – Screened records keyed by decoded SMILES

Return type:

ldmol_generator.SearchScoreResult

Returns:

Scores, masks, and canonical identities for each candidate

getCompleteCandidateCount()

Count complete candidates in the screener’s bounded top archive.

Return type:

int

Returns:

Number of retained candidates satisfying all constraints

accumulateCompleteCandidates()

Commit the screener’s complete top candidates as final results.

logScoreSummary(score_result, decode_evaluations, complete_count)

Update score history and log aggregate refinement statistics.

Parameters:

  • score_result (ldmol_generator.SearchScoreResult) – Current decoded-candidate scores and masks

  • decode_evaluations (int) – Number of decoded candidates in this step

  • complete_count (int) – Distinct complete candidates found so far

Return type:

float

Returns:

Best rankable score from this refinement step

getGenerationSeed()

Return a distinct deterministic seed for the next generation attempt.

Return type:

int or None

Returns:

Attempt seed, or None when the trainer is unseeded

evaluateCandidatesUntilValid()

Generate and screen one search batch in bounded-size chunks.