Hessian utils

mlip.data.helpers.hessian_utils.pad_systems_hessians(systems: list[ChemicalSystem]) → list[ChemicalSystem]

Pad the Hessian of each system in the given systems list to (n,3,N,3), where n is the number of atoms in the system, and N is the number of atoms in the largest system in the list.

Parameters:

systems – List of chemical systems with Hessian labels of shape (n,3,n,3).

Returns:

List of chemical systems with Hessian labels padded to shape (n,3,N,3).

mlip.data.helpers.hessian_utils.process_graph_hessian(batched_graph: Graph, num_rows: int) → Graph

Process Hessian labels to match the format of Hessian predictions.

First, rows are subsampled from graph-wise Hessians into a batch of shape (R,G,N,3) where R = num_rows, G = n_graph and N = max_system_size. Hessians are processed, cropped then permuted into shape (n, R, 3), where n is maximum possible number of graph nodes in a batch, ensuring the processed Hessian label shape is static.

Parameters:

  • batched_graph – batch of graphs with full, padded Hessian matrices.

  • num_rows – number of Hessian rows to be subsampled.

Returns:

A batched graph with processed subsampled Hessian labels.

mlip.data.helpers.hessian_utils.skip_graph_hessian(graph: Graph) → Graph

Returns a graph that has sample_hessian_rows=np.array(False) in the globals, causing the predictor to skip Hessian prediction for this graph.

Required for Hessian training when interleaving Hessian-labeled and non-Hessian batches.

Parameters:

graph – A graph for which Hessian prediction should be skipped.

Returns:

The graph with sample_hessian_rows=np.array(False) in the globals.

mlip.data.helpers.hessian_utils.get_hessian_processing_functions() → tuple[Callable[[list[ChemicalSystem]], list[ChemicalSystem]], Callable[[Graph], Graph]]

Return preprocessing and postprocessing functions for Hessian labels.

The first function pad_systems_hessians() operates on chemical systems prior to graph construction (e.g., padding Hessians), while the second function, process_graph_hessian() processes Hessian labels after graph objects have been created.

Returns:

A tuple containing a systems-level preprocessing function (applied before graph creation) and a graph-level postprocessing function (applied after graph creation).