gt4sd.frameworks.gflownet.tests.qm9 module

Summary

Classes:

QM9Dataset	QM9 dataset compatible with gflownet.
QM9GapTask	QM9 task compatible with gflownet.

Functions:

thermometer

Compute a thermometer reward using gap.

Reference

class QM9Dataset(h5_file, target='gap', properties=['rA', 'rB', 'rC', 'mu', 'alpha', 'homo', 'lumo', 'gap', 'r2', 'zpve', 'U0', 'U', 'H', 'G', 'Cv'])

Bases: GFlowNetDataset

QM9 dataset compatible with gflownet.

__init__(h5_file, target='gap', properties=['rA', 'rB', 'rC', 'mu', 'alpha', 'homo', 'lumo', 'gap', 'r2', 'zpve', 'U0', 'U', 'H', 'G', 'Cv'])

Initialize QM9 dataset.

Parameters

• h5_file (str) – path to the h5 file containing the dataset.

• target (str) – target property to optimize and build the reward.

• properties (List[str]) – list of properties to use as features.

__annotations__ = {}

__doc__ = 'QM9 dataset compatible with gflownet.'

__module__ = 'gt4sd.frameworks.gflownet.tests.qm9'

__parameters__ = ()

thermometer(v, n_bins=50, vmin=0, vmax=1)

Compute a thermometer reward using gap.

Parameters

• v (Tensor) – tensor of values to compute the reward.

• n_bins (int) – number of bins to use.

• vmin (int) – minimum value of the range.

• vmax (int) – maximum value of the range.

Return type

Tensor

Returns

tensor of the reward.

class QM9GapTask(configuration, dataset, reward_model=None, wrap_model=None)

Bases: GFlowNetTask

QM9 task compatible with gflownet.

__init__(configuration, dataset, reward_model=None, wrap_model=None)

Initialize QM9 task.

Code adapted from: https://github.com/recursionpharma/gflownet/blob/trunk/src/gflownet/tasks/qm9/qm9.py.

Parameters

• configuration (Dict[str, Any]) – configuration of the task.

• dataset (GFlowNetDataset) – dataset to use for the task.

• reward_model (Optional[Module, None]) – model to use for the reward.

• wrap_model (Optional[Callable[[Module], Module], None]) – function to wrap the model.

flat_reward_transform(_y)

Transforms a target quantity y (e.g. the LUMO energy in QM9) to a positive reward scalar.

Parameters

_y (Union[float, Tensor]) – target quantity to transform.

Return type

FlatRewards

Returns

reward scalar.

inverse_flat_reward_transform(rp)

Inverse transform a reward scalar to a target quantity y (e.g. the LUMO energy in QM9).

Parameters

rp – reward scalar to transform.

Returns

target quantity.

load_task_models()

Loads the models for the task.

Return type

Dict[str, Module]

Returns

dictionary of models.

sample_conditional_information(n)

Sample conditional information for the task.

Parameters

n – number of samples to sample.

Returns

dictionary of conditional information.

cond_info_to_reward(cond_info, _flat_reward)

Compute the reward for a given conditional information.

Parameters

• cond_info (Dict[str, Tensor]) – dictionary of conditional information.

• _flat_reward (FlatRewards) – flat reward.

Return type

RewardScalar

Returns

reward scalar.

compute_flat_rewards(mols)

Computes the flat rewards for a list of molecules.

Parameters

mols (List[Mol]) – list of molecules.

Return type

Tuple[RewardScalar, Tensor]

Returns

reward scalar and validity.

__annotations__ = {}

__doc__ = 'QM9 task compatible with gflownet.'

__module__ = 'gt4sd.frameworks.gflownet.tests.qm9'