schrodinger.application.matsci.flywheel.ldmol.generator module

High-level LDMol generator for text-to-molecule generation.

Integrates MolT5 text encoder, VAE autoencoder, DiT, and DDIM sampling.

class schrodinger.application.matsci.flywheel.ldmol.generator.SearchScoreResult(scores, scorable_mask=None, molecule_keys=None, optimal_mask=None, stop_search=False)

Bases: object

Scores returned by a DiffGraph Search scoring callback.

Parameters:

  • scores (array-like) – Candidate scores in decoded-SMILES order

  • scorable_mask (array-like or None) – Candidates with valid property evaluations. If not provided, candidates with finite scores are considered scorable.

  • molecule_keys (list(str) or None) – Canonical identities in decoded-SMILES order

  • optimal_mask (array-like or None) – Candidates satisfying the strict completion criterion

  • stop_search (bool) – Whether the caller has enough complete candidates

__init__(scores, scorable_mask=None, molecule_keys=None, optimal_mask=None, stop_search=False)
class schrodinger.application.matsci.flywheel.ldmol.generator.DiffGraphSearchSeedPool(latents: object = None, scores: object = None, offset: int = 0, smiles: object = None, molecule_keys: object = None)

Bases: object

Optional molecular seeds used to initialize DiffGraph Search.

Parameters:

  • latents (numpy.ndarray or None) – Clean molecular latents

  • scores (array-like or None) – Search scores aligned with latents

  • offset (int) – Starting position in the ranked seed pool

  • smiles (list(str) or None) – Decoded molecules aligned with latents

  • molecule_keys (list(str) or None) – Canonical identities aligned with latents

latents: object
scores: object
offset: int
smiles: object
molecule_keys: object
__init__(latents: object = None, scores: object = None, offset: int = 0, smiles: object = None, molecule_keys: object = None) → None
class schrodinger.application.matsci.flywheel.ldmol.generator.DiffGraphSearchElite(clean: object, score: float, smiles: str = None, molecule_key: str = None)

Bases: object

Best scored molecule and its aligned search data.

Parameters:

  • clean (torch.Tensor) – Predicted clean latent for the molecule

  • score (float) – Search score for the molecule

  • smiles (str or None) – Decoded SMILES for the molecule

  • molecule_key (str or None) – Canonical molecular identity

clean: object
score: float
smiles: str
molecule_key: str
__init__(clean: object, score: float, smiles: str = None, molecule_key: str = None) → None
class schrodinger.application.matsci.flywheel.ldmol.generator.PromptBatch(text_embs: object, attn_mask: object, batch_size: int)

Bases: object

Encoded prompts and their expanded batch size.

Parameters:

  • text_embs (torch.Tensor) – Conditional prompt embeddings

  • attn_mask (torch.Tensor) – Conditional prompt attention mask

  • batch_size (int) – Expanded prompt batch size

text_embs: object
attn_mask: object
batch_size: int
__init__(text_embs: object, attn_mask: object, batch_size: int) → None
class schrodinger.application.matsci.flywheel.ldmol.generator.DiffGraphSearchConfig(score_fn: object, branch_factor: int = 1, max_length: int = 150, mc_samples: int = 1)

Bases: object

Fixed configuration for DiffGraph Search sampling.

Parameters:

  • score_fn (callable) – Maps decoded SMILES to SearchScoreResult

  • branch_factor (int) – Transitions generated per active state

  • max_length (int) – Maximum autoregressive SMILES decode length

  • mc_samples (int) – Stochastic decoder samples per transition

score_fn: object
branch_factor: int
max_length: int
mc_samples: int
__init__(score_fn: object, branch_factor: int = 1, max_length: int = 150, mc_samples: int = 1) → None
class schrodinger.application.matsci.flywheel.ldmol.generator.DiffGraphSearchActive(states: object, model_kwargs: dict, model_output: dict = None, seed_latents: object = None, scores: object = None, smiles: object = None, molecule_keys: object = None)

Bases: object

Aligned data for the active DiffGraph Search branches.

Parameters:

  • states (torch.Tensor) – Active noisy diffusion states

  • model_kwargs (dict) – Model inputs aligned with states

  • model_output (dict or None) – Model outputs aligned with states

  • seed_latents (torch.Tensor or None) – Initial clean latents used to create states

  • scores (numpy.ndarray or None) – Scores aligned with states

  • smiles (list(str) or None) – Decoded molecules aligned with states

  • molecule_keys (list(str) or None) – Canonical identities aligned with states

states: object
model_kwargs: dict
model_output: dict
seed_latents: object
scores: object
smiles: object
molecule_keys: object
__init__(states: object, model_kwargs: dict, model_output: dict = None, seed_latents: object = None, scores: object = None, smiles: object = None, molecule_keys: object = None) → None
class schrodinger.application.matsci.flywheel.ldmol.generator.DiffGraphSearchCandidates(states: object, model_kwargs: dict, model_output: dict, clean_latents: object, smiles: object = None, score_result: object = None)

Bases: object

Aligned data for one set of proposed search branches.

Parameters:

  • states (torch.Tensor) – Proposed noisy branch states

  • model_kwargs (dict) – Model inputs aligned with states

  • model_output (dict) – Model outputs aligned with states

  • clean_latents (torch.Tensor) – Predicted clean latents aligned with states

  • smiles (list(str) or None) – Decoded candidate molecules

  • score_result (SearchScoreResult or None) – Scores and identities aligned with smiles

states: object
model_kwargs: dict
model_output: dict
clean_latents: object
smiles: object
score_result: object
__init__(states: object, model_kwargs: dict, model_output: dict, clean_latents: object, smiles: object = None, score_result: object = None) → None
class schrodinger.application.matsci.flywheel.ldmol.generator.DiffGraphSearchSelection(clean_latents: object, smiles: object)

Bases: object

Selected clean latents and their decoded molecules.

Parameters:

  • clean_latents (torch.Tensor) – Selected predicted clean latents

  • smiles (list(str)) – Decoded molecules aligned with clean_latents

clean_latents: object
smiles: object
__init__(clean_latents: object, smiles: object) → None
class schrodinger.application.matsci.flywheel.ldmol.generator.DiffGraphSearchResult(latents: object, elite_smiles: str = None, elite_index: object = None, candidate_smiles: object = None)

Bases: object

Fixed result returned by DiffGraph Search sampling.

Parameters:

  • latents (torch.Tensor) – Final denoised molecular latents

  • elite_smiles (str or None) – Best decoded molecule found by the search

  • elite_index (int or None) – Position of the elite molecule in latents

  • candidate_smiles (list(str) or None) – Scored molecules selected at early completion

latents: object
elite_smiles: str
elite_index: object
candidate_smiles: object
__init__(latents: object, elite_smiles: str = None, elite_index: object = None, candidate_smiles: object = None) → None
class schrodinger.application.matsci.flywheel.ldmol.generator.DiffGraphSearchState(text_embs, attn_mask, batch_size, diffusion, seed_pool)

Bases: object

Inputs and initialized active state for one DiffGraph Search run.

Parameters:

  • text_embs (torch.Tensor) – Conditional prompt embeddings

  • attn_mask (torch.Tensor) – Conditional prompt attention mask

  • batch_size (int) – Active search-state count

  • diffusion (SpacedDiffusion) – Diffusion sampling process

  • seed_pool (DiffGraphSearchSeedPool) – Optional ranked molecular seeds

__init__(text_embs, attn_mask, batch_size, diffusion, seed_pool)
property active_count

Number of initialized active states.

Return type:

int

Returns:

Number of active diffusion states

Whether candidate scoring requested early search completion.

configure(config)

Set configuration used by one search run.

Parameters:

config (DiffGraphSearchConfig) – Fixed search scoring and branching configuration

validate()

Validate optional seed-pool data before initialization.

setInitialState(active, initial_pool_text)

Store initialized active states and aligned metadata.

Parameters:

  • active (DiffGraphSearchActive) – Initialized active branches and aligned metadata

  • initial_pool_text (str) – Initial-pool summary for logging

setBranchCandidates(candidates)

Store branch transitions and their predicted clean latents.

Parameters:

candidates (DiffGraphSearchCandidates) – Proposed branches and aligned model data

setCandidateScores(smiles_list, score_result)

Store validated decoded-candidate scores and identities.

Parameters:

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

  • score_result (SearchScoreResult) – Scores and identities aligned with smiles_list

static repeatBatch(batch, repeat_count)

Repeat every batched tensor for independent search branches.

Parameters:

  • batch (dict) – Model inputs containing batched tensors

  • repeat_count (int) – Repetitions per tensor row

Return type:

dict

Returns:

Model inputs with repeated tensor rows

static selectBatch(batch, selected)

Select retained rows from batched tensors.

Parameters:

  • batch (dict) – Model inputs containing batched tensors

  • selected (torch.Tensor) – Row indices to retain

Return type:

dict

Returns:

Model inputs with selected tensor rows

retainCandidateBranches()

Retain ranked candidate branches and update the active state.

Return type:

bool

Returns:

Whether at least one scored candidate was retained

initializeElite()

Initialize the global elite from the optional seed pool.

finalize()

Return final latents and optional global-elite metadata.

Return type:

DiffGraphSearchResult

Returns:

Final latents with elite and candidate metadata

schrodinger.application.matsci.flywheel.ldmol.generator.select_diff_graph_search_candidates(scores, num_select, scorable_mask=None, smiles_list=None, molecule_keys=None, priority_mask=None)

Return global highest-scoring search candidate indices.

Parameters:

  • scores (numpy.ndarray) – Flat candidate scores

  • num_select (int) – Number of candidates to retain

  • scorable_mask (numpy.ndarray or None) – Candidates with valid property scores

  • smiles_list (list(str) or None) – Optional decoded molecules used for deduplication

  • molecule_keys (list(str) or None) – Canonical identities used instead of raw SMILES

  • priority_mask (array-like or None) – Rank these candidates before other rankable entries

Return type:

numpy.ndarray

Returns:

Selected indices in descending score order

class schrodinger.application.matsci.flywheel.ldmol.generator.LDMolGenerator(device='cpu', logger=None)

Bases: BaseModel

Text-to-molecule generator using LDMol.

Generates molecules from text descriptions via latent diffusion. All models are loaded lazily on first use.

__init__(device='cpu', logger=None)
Parameters:

  • device (str) – Compute device (‘cpu’, ‘cuda’, ‘mps’)

  • logger (logging.Logger or None) – Logger instance for debug output

ensureModelsLoaded()

Load all models on first use (lazy loading).

Loads MolT5 text encoder, VAE autoencoder, and DiT in eval mode.

loadDiT()

Download (if needed), create, and load the DiT checkpoint.

Weights are loaded with strict=False to allow EMA checkpoints that may contain extra keys not in the model definition.

encodePrompts(text_prompts, num_samples)

Encode text prompts to MolT5 embeddings and tile for batch.

Parameters:

  • text_prompts (list(str)) – Text descriptions

  • num_samples (int) – Number of molecules per prompt

Return type:

PromptBatch

Returns:

Encoded prompt tensors and expanded batch size

setGenerationSeed(seed)

Seed all random generators used by molecule generation.

Parameters:

seed (int or None) – Random seed, or None to preserve the global RNG state

runDiffusionSampling(text_embs, attn_mask, batch_size, diffusion, cfg_scale, seed=None)

Run DDIM reverse diffusion with optional classifier-free guidance.

Parameters:

  • text_embs (torch.Tensor) – Conditioned text embeddings (batch, seq, dim)

  • attn_mask (torch.Tensor) – Text attention mask (batch, seq)

  • batch_size (int) – Number of molecules to generate

  • diffusion (SpacedDiffusion) – Diffusion sampling process

  • cfg_scale (float) – Guidance scale (>= 1.0)

  • seed (int or None) – Random seed set immediately before noise sampling for reproducibility. Seeds both torch (noise tensor) and numpy (consumed by downstream alignment code). None leaves the global RNG state unchanged.

Return type:

torch.Tensor

Returns:

Denoised latents (batch, seq_len, channels)

decodeDiffGraphSearchCandidates(clean_latents, max_length, mc_samples)

Decode independent molecules for each search clean latent.

Parameters:

  • clean_latents (torch.Tensor) – Latents with shape (batch, channels, seq, 1)

  • max_length (int) – Maximum autoregressive SMILES length

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

Return type:

list(str)

Returns:

Flattened decoded SMILES in latent/sample order

initializeDiffGraphSearchState(search_state)

Populate one validated DiffGraph Search state.

Valid initial molecules are ranked and noised at the highest diffusion timestep. If no initial pool is supplied, random Gaussian states are created instead.

Parameters:

search_state (DiffGraphSearchState) – Search inputs and output state

advanceDiffGraphSearchState()

Take one deterministic transition without branch scoring.

The current search state is held by the generator while a search run is active. This path is used for the final timestep after branch selection, and preserves the current model output when another reverse-diffusion step remains.

branchDiffGraphSearchState()

Propose, score, and retain the best search branch states.

Each active state gets branch_factor continuations: one deterministic elite continuation and branch_factor - 1 stochastic continuations. Each transition is decoded once ordinarily and mc_samples additional times. The cached global elite is retained after branch selection and is not rescored here.

Return type:

bool

Returns:

False when no candidate can be scored and the deterministic parent paths were retained; True after selecting scored branches

getDiffGraphSearchTimesteps(count)

Create a timestep tensor for the active search step.

Parameters:

count (int) – Number of timestep entries

Return type:

torch.Tensor

Returns:

Current timestep repeated count times

predictDiffGraphSearchStates(states, timesteps, model_kwargs)

Predict clean latents for a batch at the requested timesteps.

Parameters:

  • states (torch.Tensor) – Noisy diffusion states

  • timesteps (torch.Tensor) – Diffusion timestep for each state

  • model_kwargs (dict) – Model inputs aligned with states

Return type:

dict

Returns:

Diffusion model outputs

buildDiffGraphSearchBranches()

Build deterministic and stochastic branch transitions.

scoreDiffGraphSearchBranches()

Decode and score the current branch candidates.

selectDiffGraphSearchBranches()

Select scored branches and update the active search state.

Return type:

bool

Returns:

Whether at least one scored branch was selected

reinsertDiffGraphSearchElite()

Keep the global elite in the active diffusion state.

runDiffGraphSearchLoop()

Run the stateful reverse-diffusion and branch-selection loop.

runDiffGraphSearchSampling(prompt_batch, diffusion, config, seed_pool=None, seed=None)

Run DiffGraph Search with direct clean-molecule scoring.

This search is a variation of TreeG-SC that ranks decoded molecular property scores directly and does not average Monte Carlo samples.

Each active state takes config.branch_factor reverse transitions: one deterministic transition and config.branch_factor - 1 stochastic transitions. Each transition is decoded once ordinarily plus config.mc_samples stochastic decoder samples. The cached global elite is retained without rescoring. The best clean latent is retained as an elite state.

Parameters:

  • prompt_batch (PromptBatch) – Encoded prompts and active-set size

  • diffusion (SpacedDiffusion) – Diffusion sampling process

  • config (DiffGraphSearchConfig) – Search scoring and branching configuration

  • seed_pool (DiffGraphSearchSeedPool or None) – Optional ranked molecular seeds

  • seed (int or None) – Random seed set before initial and branch noise sampling

Return type:

DiffGraphSearchResult

Returns:

Denoised latents and aligned search metadata

generate(text_prompts, num_samples=1, num_sampling_steps=100, cfg_scale=7.5, max_length=150, seed=None, return_latents=False)

Generate molecules from text descriptions.

Pipeline: text → MolT5 embeddings → DiT DDIM denoising → VAE decode → SMILES.

Parameters:

  • text_prompts (str or list(str)) – Text description(s) to condition generation

  • num_samples (int) – Number of molecules to generate per prompt

  • num_sampling_steps (int) – DDIM reverse diffusion steps

  • cfg_scale (float) – Classifier-free guidance scale (>= 1.0)

  • max_length (int) – Maximum autoregressive SMILES decode length

  • seed (int or None) – Random seed set immediately before noise sampling. None leaves the global RNG state unchanged.

  • return_latents (bool) – If True, return (smiles, latents) tuple

Return type:

list(str) or tuple(list(str), torch.Tensor)

Returns:

Generated SMILES strings, or (smiles, latents) if return_latents

generateDiffGraphSearch(text_prompts, config, num_samples=1, num_sampling_steps=100, seed=None, return_latents=False, seed_pool=None)

Generate molecules using DiffGraph Search.

Parameters:

  • text_prompts (str or list(str)) – Text description(s) to condition generation

  • config (DiffGraphSearchConfig) – Search scoring and branching configuration

  • num_samples (int) – Number of molecules to generate per prompt

  • num_sampling_steps (int) – DDIM reverse diffusion steps

  • seed (int or None) – Random seed set immediately before sampling

  • return_latents (bool) – If True, return (smiles, latents) tuple

  • seed_pool (DiffGraphSearchSeedPool or None) – Optional ranked molecular seeds

Return type:

list(str) or tuple(list(str), torch.Tensor)

Returns:

Generated SMILES strings, or (smiles, latents) if requested

decodeLatents(latents, max_length=150)

Decode latent tensors to SMILES strings.

Parameters:

  • latents (torch.Tensor) – Latents (batch, seq_len, channels)

  • max_length (int) – Maximum autoregressive decode length

Return type:

list(str)

Returns:

Decoded SMILES strings

encodeSmiles(smiles_list)

Encode SMILES strings into the molecular VAE latent space.

Parameters:

smiles_list (list(str)) – SMILES strings to encode

Return type:

torch.Tensor

Returns:

Latents with shape (batch, sequence length, channels)