mesh_logging

class LogMolecule(max_molecule_resolution: int = 10, max_isosurface_resolution: float = 0.2, basis_transformation: str | Callable | None = 'auto', log_initial_guess: bool = True, log_gradient: bool = True, log_random_basis_functions: bool = False, **super_kwargs)

Logs the molecule to tensorboard as a mesh.

__init__(max_molecule_resolution: int = 10, max_isosurface_resolution: float = 0.2, basis_transformation: str | Callable | None = 'auto', log_initial_guess: bool = True, log_gradient: bool = True, log_random_basis_functions: bool = False, **super_kwargs) → None

Parameters:

• max_molecule_resolution – The maximum “resolution” of the sticks representation of the molecule.

• max_isosurface_resolution – The maximum grid resolution used to compute the isosurface, in Bohr.

• basis_transformation – The basis transformation to be applied to the molecule. Can be either None, ‘auto’, or a callable, e.g. a value of BASIS_TRANSFORMATIONS. If ‘auto’, the basis transformation is chosen based on the label_subdir of the datamodule. None means no transformation is applied. Defaults to ‘auto’.

• log_initial_guess – Whether to log the initial guess. Defaults to True.

• log_gradient – Whether to log the gradient. Defaults to True.

• log_random_basis_functions – Whether to log some random basis functions. Defaults to False. See LogMolecule.get_mesh_kwargs() for details.

• **super_kwargs – Keyword arguments passed to OnStepCallbackWithTiming.

execute(trainer: Trainer, pl_module: MLDFTLitModule, outputs: Tensor | Mapping[str, Any] | None, batch: Any, split: str) → None

Logs a mesh to tensorboard at the end of a training batch, if the timing matches.

get_mesh_kwargs(trainer: Trainer, pl_module: MLDFTLitModule, batch: Any, outputs: Tensor | Mapping[str, Any] | None, basis_info: BasisInfo) → dict

Create a dictionary of meshes to be logged.

Parameters:

• trainer – The lightning trainer.

• pl_module – The lightning module.

• batch – The batch.

• outputs – The outputs of the lightning module.

• basis_info – The BasisInfo object.

Returns:

A dict of dicts where each value is a dict with arguments to SummaryWriter.add_mesh.

Return type:

dict

get_molecule_dict(mol: Any) → dict

Get a sticks representation of the molecule as a mesh dict.

combine_mesh_dicts(mesh_dicts: list[dict]) → dict

Combines multiple mesh dicts into one.

Parameters:

mesh_dicts – A list of mesh dicts. Each dict must have exactly the keys vertices, faces, colors.

Returns:

The combined mesh dict.

Return type:

dict

hex_to_rgb(hex_color: str) → tuple

Converts a hex color string to an rgb tuple.

Parameters:

hex_color – The hex color string, e.g. #FF0000.

Returns:

The rgb tuple, e.g. (255, 0, 0).

Return type:

tuple