MolecularDiffusion.modules.models.jodo.utils¶

Masking / centering / symmetric-noise helpers, ported verbatim from JODO models/utils.py (the model registry and DataParallel factory are dropped – Hydra _target_ is this platform’s registry).

Functions¶

`assert_correctly_masked`(variable, node_mask)
`assert_mean_zero_with_mask`(x, node_mask[, eps])
`check_mask_correct`(variables, node_mask)
`coord2diff`(x, edge_index[, norm_constant])
`coord2diff_adj`(x, edge_index[, spatial_th])
`coord2dist`(x, edge_index)
`get_rw_feat`(k_step, dense_adj)	Compute k_step Random Walk for given dense adjacency matrix.
`remove_mean`(x)
`remove_mean_with_mask`(x, node_mask)
`sample_center_gravity_zero_gaussian_with_mask`(size, ...)
`sample_combined_position_feature_noise`(n_samples, ...)	Sample mean-centered normal noise for z_x, and standard normal noise for z_h.
`sample_gaussian_with_mask`(size, device, node_mask)
`sample_symmetric_edge_feature_noise`(n_samples, ...)	sample symmetric normal noise for edge feature.
`to_dense_edge_attr`(edge_index, edge_attr, edge_final, ...)

Module Contents¶

MolecularDiffusion.modules.models.jodo.utils.assert_correctly_masked(variable, node_mask)¶

MolecularDiffusion.modules.models.jodo.utils.assert_mean_zero_with_mask(x, node_mask, eps=1e-10)¶

MolecularDiffusion.modules.models.jodo.utils.check_mask_correct(variables, node_mask)¶

MolecularDiffusion.modules.models.jodo.utils.coord2diff(x, edge_index, norm_constant=1)¶

MolecularDiffusion.modules.models.jodo.utils.coord2diff_adj(x, edge_index, spatial_th=2.0)¶

MolecularDiffusion.modules.models.jodo.utils.coord2dist(x, edge_index)¶

MolecularDiffusion.modules.models.jodo.utils.get_rw_feat(k_step, dense_adj)¶: Compute k_step Random Walk for given dense adjacency matrix.

MolecularDiffusion.modules.models.jodo.utils.remove_mean(x)¶

MolecularDiffusion.modules.models.jodo.utils.remove_mean_with_mask(x, node_mask)¶

MolecularDiffusion.modules.models.jodo.utils.sample_center_gravity_zero_gaussian_with_mask(size, device, node_mask)¶

MolecularDiffusion.modules.models.jodo.utils.sample_combined_position_feature_noise(n_samples, n_nodes, in_node_nf, node_mask)¶: Sample mean-centered normal noise for z_x, and standard normal noise for z_h.

MolecularDiffusion.modules.models.jodo.utils.sample_gaussian_with_mask(size, device, node_mask)¶

MolecularDiffusion.modules.models.jodo.utils.sample_symmetric_edge_feature_noise(n_samples, n_nodes, edge_ch, edge_mask)¶: sample symmetric normal noise for edge feature.

MolecularDiffusion.modules.models.jodo.utils.to_dense_edge_attr(edge_index, edge_attr, edge_final, bs, n_nodes)¶