schrodinger.application.matsci.flywheel.ldmol.diffusion.respace module

Timestep re-spacing for DDIM and strided diffusion sampling.

schrodinger.application.matsci.flywheel.ldmol.diffusion.respace.space_timesteps(num_timesteps, section_counts)

Select a subset of timesteps from the full diffusion schedule.

Divides the schedule into equal sections and strides each to the requested count using fractional indexing, always anchoring at both endpoints of each section.

Parameters:

  • num_timesteps (int) – Total steps in the original schedule

  • section_counts (list(int) or str) – Step counts per section as a list or comma-separated string (e.g. ‘100’ or [100])

Return type:

set(int)

Returns:

Selected timestep indices

class schrodinger.application.matsci.flywheel.ldmol.diffusion.respace.SpacedDiffusion(use_timesteps, **kwargs)

Bases: GaussianDiffusion

Diffusion process that runs a subset of the original timesteps.

Selects evenly-spaced steps from the full diffusion schedule and remaps timestep indices so the wrapped model always receives values from the original 0..T-1 range.

__init__(use_timesteps, **kwargs)
Parameters:

  • use_timesteps (set(int) or list(int)) – Timestep indices to retain from the full schedule

  • kwargs – Arguments forwarded to GaussianDiffusion

pMeanVariance(model, *args, **kwargs)

Delegate to parent with timestep-remapped model.

trainingLosses(model, *args, **kwargs)

Delegate to parent with timestep-remapped model.

conditionMean(cond_fn, *args, **kwargs)

Delegate to parent with timestep-remapped condition function.

conditionScore(cond_fn, *args, **kwargs)

Delegate to parent with timestep-remapped condition function.

wrapModel(model)

Wrap model to remap spaced indices to original timestep values.

Parameters:

model – Model or condition function to wrap

Returns:

Wrapped model

Return type:

WrappedModel

scaleTimesteps(timesteps)

No-op: scaling is handled by the wrapped model.

Parameters:

timesteps (torch.Tensor) – Spaced timestep indices

Return type:

torch.Tensor

Returns:

Unchanged timesteps

class schrodinger.application.matsci.flywheel.ldmol.diffusion.respace.WrappedModel(model, timestep_map, original_num_steps)

Bases: object

Remaps spaced timestep indices to their original diffusion values.

SpacedDiffusion passes 0..N-1 spaced indices to the model. This wrapper translates them back to the original 0..T-1 schedule so the DiT’s timestep embeddings remain aligned with training.

__init__(model, timestep_map, original_num_steps)
Parameters:

  • model – Underlying model or condition function

  • timestep_map (list(int)) – Maps spaced index → original index

  • original_num_steps (int) – Total steps in the original schedule