schrodinger.application.matsci.flywheel.base_trainer module

Shared infrastructure for property-guided molecular alignment trainers.

This module owns screening/history accumulation, convergence, checkpointing, and result packaging. Concrete trainers implement candidate generation.

Copyright Schrodinger, LLC. All rights reserved.

class schrodinger.application.matsci.flywheel.base_trainer.CandidateRecord(smiles: str, molecule_key: str = None, structure: object = None, latent: object = None, score: float = None, mol_detail: dict = None, rejection_reason: str = None, evaluated: bool = False)

Bases: object

Persistent screening state and evaluation outcome for one molecule.

smiles: str
molecule_key: str
structure: object
latent: object
score: float
mol_detail: dict
rejection_reason: str
evaluated: bool
__init__(smiles: str, molecule_key: str = None, structure: object = None, latent: object = None, score: float = None, mol_detail: dict = None, rejection_reason: str = None, evaluated: bool = False) → None
schrodinger.application.matsci.flywheel.base_trainer.build_candidate_trace(smiles_list, accepted_candidates, proposal_source=None)

Build evaluation-order records for accepted candidates in debug mode.

Parameters:

  • smiles_list (list(str)) – Decoded SMILES strings

  • accepted_candidates (dict(int, CandidateRecord)) – Accepted records by input position

  • proposal_source (str or None) – Candidate-generation source

Return type:

list(dict)

Returns:

Candidate records, or an empty list outside debug mode

schrodinger.application.matsci.flywheel.base_trainer.build_final_details(mol_details)

Build columnar screening details from valid candidate records.

Parameters:

mol_details (list(dict)) – Screening details for valid candidates

Return type:

dict

Returns:

Screening details keyed by module constants

class schrodinger.application.matsci.flywheel.base_trainer.BaseTrainer(generator, screener, batch_size=10, num_molecules=10, convergence_threshold=0.01, seed=None, logger=None, checkpoint_callback=None)

Bases: object

Shared screening and result workflow for alignment trainers.

__init__(generator, screener, batch_size=10, num_molecules=10, convergence_threshold=0.01, seed=None, logger=None, checkpoint_callback=None)
Parameters:

  • generator (LDMolGenerator) – LDMol molecule generator

  • screener (ScreenerList) – Property screener

  • batch_size (int) – Molecules to acquire per search batch

  • num_molecules (int) – Number of top molecules used for cumulative mean statistics and convergence

  • convergence_threshold (float) – Minimum fitness improvement to reset the patience counter

  • seed (int or None) – Global random seed for numpy and torch; set at the start of run() for reproducibility. None leaves the state unseeded.

  • logger (logging.Logger or None) – Logger for progress output, or None to print to stdout

  • checkpoint_callback (callable or None) – Function receiving the current SMILES, alignment fitness, and screening details after each batch

log(message)
Parameters:

message (str) – Message to log

debugLog(message)

Log a technical diagnostic only in debug mode.

Parameters:

message (str) – Diagnostic message

getSearchScore(mol_detail, score_key, default=0.0)

Return an alignment-search score with a strict-score fallback.

Parameters:

  • mol_detail (dict) – Per-molecule screening details

  • score_key (str) – Score key to retrieve

  • default (float) – Value used when the score is unavailable

Return type:

float

Returns:

Search-time score value

getActualScore(mol_detail)

Return the alignment-search score for one molecule.

Parameters:

mol_detail (dict) – Per-molecule screening details

Return type:

float

Returns:

Alignment-search score

getTopScoreIndices(scores)

Return indices of the highest-scoring output-sized subset.

Parameters:

scores (list(float)) – Scores for accumulated molecules

Return type:

numpy.ndarray

Returns:

Indices of up to num_molecules highest scores

getTopScoreMean(scores)

Return the mean over the highest-scoring output-sized subset.

Parameters:

scores (list(float)) – Scores for accumulated molecules

Return type:

float

Returns:

Mean of up to num_molecules highest scores

getCandidateData(data_key=None)

Return accepted records or one aligned numerical array.

Parameters:

data_key (str or None) – CANDIDATE_LATENTS_KEY, CANDIDATE_SEARCH_SCORES_KEY, or CANDIDATE_RECORDS_KEY. None returns records.

Return type:

tuple(CandidateRecord) or numpy.ndarray

Returns:

Requested data in candidate-record insertion order

getFeasibleIndices()

Return indices of molecules satisfying all property constraints.

Return type:

list(int)

Returns:

Indices into the accepted candidate records

getBestProgressIndex()

Return index of best result for user-facing progress.

Feasible molecules rank by strict fitness. Until one is found, search score identifies the closest candidate.

Return type:

int

Returns:

Index into the accepted candidate records

formatPropertyPredictionLines(mol_detail)

Format predicted properties as separate user-facing lines.

Parameters:

mol_detail (dict) – Per-molecule screening details

Return type:

list(str)

Returns:

One formatted line per predicted property

logProgressSummary()

Log cumulative progress in user-facing terms.

flattenLatents(latents)

Flatten 3-D latents for storage with candidate records.

Parameters:

latents (numpy.ndarray) – Latents (N, seq_len, channels)

Return type:

numpy.ndarray

Returns:

Flattened latents (N, seq_len * channels)

accumulateInitialBatch(smiles_list, latents)

Screen and accumulate a pre-generated initial batch.

Initializes global best and top-result mean fitness from accepted results.

Parameters:

  • smiles_list (list) – Pre-generated SMILES strings

  • latents (numpy.ndarray) – Corresponding latents (N, seq_len, channels)

writeCheckpoint()

Write current alignment history through caller callback.

run(initial_smiles, initial_latents)

Run alignment from a pre-generated initial batch.

The caller generates the initial molecules. This method screens them and evaluates method-specific candidate batches.

Parameters:

  • initial_smiles (list) – SMILES strings for the initial batch

  • initial_latents (numpy.ndarray) – Corresponding latents (N, seq_len, channels)

Return type:

dict(str, CandidateRecord)

Returns:

Accepted records keyed by canonical molecule identity

screenAndAccumulate(smiles_list, latents, proposal_source=None)

Screen candidates and accumulate accepted molecule records.

Parameters:

  • smiles_list (list) – Decoded SMILES for batch

  • latents (numpy.ndarray) – Flat latent vectors (N, D)

  • proposal_source (str or None) – Candidate-generation source

Return type:

int

Returns:

Number of accepted molecules

evaluateCandidatesUntilValid()

Generate and evaluate one method-specific candidate batch.

Implementations returning True must consume candidates so a later call makes progress toward exhausting the finite search space.

Return type:

bool

Returns:

True when another search batch can be evaluated, or False when the trainer has no further candidates.

updateImprovementState()

Update convergence state and report accumulated progress.

checkOptimalResults()

Return whether requested output count reached maximum strict fitness.

Return type:

bool

Returns:

True when num_molecules results have strict fitness 1.0

getFinalFitnessScores()

Return strict reporting scores.

Return type:

list(float)

Returns:

Final score for each accumulated molecule

getBestResults()

Return accumulated results.

Return type:

dict(str, CandidateRecord)

Returns:

Accepted records keyed by canonical molecule identity