image_logging

class LogDistanceEmbeddings(max_distance: float = 5.0, n_distances: int = 1000, **super_kwargs)

Logs a line plot of the distance embeddings for a range of distances.

__init__(max_distance: float = 5.0, n_distances: int = 1000, **super_kwargs)

Plot the distance embeddings for a range of distances.

Parameters:

• max_distance (float) – The maximum distance to consider.

• n_distances (int) – The number of distances to consider.

Returns:

The plot.

Return type:

plt.Figure

get_figure(pl_module: MLDFTLitModule, batch: Any, outputs: Tensor | Mapping[str, Any] | None, basis_info: BasisInfo) → Figure

Create the figure to be plotted: A line plot of the distance embeddings

class LogGradientScatter(train_timing: CallbackTiming = None, val_timing: CallbackTiming = None, name: str = 'auto')

Logs a scatter plot of the target and predicted gradients per basis function.

get_figure(pl_module: MLDFTLitModule, batch: Any, outputs: Tensor | Mapping[str, Any] | None, basis_info: BasisInfo) → Figure

Create the figure to be plotted: A scatter plot of the target and predicted gradients per basis function, as well as their projected difference.

Parameters:

• pl_module – The lightning module.

• batch – The batch.

• outputs – The outputs of the lightning module.

• basis_info – The BasisInfo object.

Returns:

The figure to be plotted.

Return type:

Figure

class LogMatplotlibToTensorboard(train_timing: CallbackTiming = None, val_timing: CallbackTiming = None, name: str = 'auto')

Base Class to log matplotlib figures to tensorboard.

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

Generates a figure using get_figure() and logs it to tensorboard.

get_figure(pl_module: MLDFTLitModule, batch: Any, outputs: Tensor | Mapping[str, Any] | None, basis_info: BasisInfo) → Figure

Create the figure to be plotted.

Parameters:

• pl_module – The lightning module.

• batch – The batch.

• outputs – The outputs of the lightning module.

• basis_info – The BasisInfo object.

Returns:

The figure to be plotted.

Return type:

Figure

class LogTargetPredScatters(with_atom_ref: bool | str = 'auto', **super_kwargs)

Logs scatter plots of the target and predicted energies, gradients and initial guess deltas.

__init__(with_atom_ref: bool | str = 'auto', **super_kwargs)

Parameters:

• with_atom_ref (bool | str) – Whether to add two additional plots of energy / gradient minus their AtomRef values. If ‘auto’, it is checked whether the model has a property atom_ref_module. Defaults to ‘auto’.

• **super_kwargs – Additional kwargs for the superclass.

get_figure(pl_module: MLDFTLitModule, batch: Any, outputs: Tensor | Mapping[str, Any] | None, basis_info: BasisInfo) → Figure

Create two scatter plots: One for the energies, and one for the gradients.

Parameters:

• pl_module – The lightning module.

• batch – The batch.

• outputs – The outputs of the lightning module.

• basis_info – The BasisInfo object.

Returns:

The figure to be plotted.

Return type:

Figure