neural_lam.utils

neural_lam.utils#

Utility helpers shared across Neural-LAM training and evaluation.

Module Contents#

class neural_lam.utils.BufferList(buffer_tensors, persistent=True)#

Bases: torch.nn.Module

A list of torch buffer tensors that sit together as a Module with no parameters and only buffers.

This should be replaced by a native torch BufferList once implemented. See: pytorch/pytorch#37386

Register a collection of tensors as buffers inside a module.

Parameters:

  • buffer_tensors (Sequence[torch.Tensor]) – Buffers to register in the order they should be indexed.

  • persistent (bool, optional) – If True, buffers are saved in checkpoints. Default True.

n_buffers#
neural_lam.utils.fractional_plot_bundle(fraction)#

Return a tueplots bundle scaled to a fraction of the page width.

Parameters:

fraction (float) – Denominator applied to the default NeurIPS figure width.

Returns:

Matplotlib rcParams bundle with updated figure.figsize.

Return type:

dict

neural_lam.utils.get_integer_time(tdelta) → tuple[int, str]#

Express a datetime.timedelta as an integer number of time units.

Parameters:

tdelta (datetime.timedelta) – Time interval to convert.

Returns:

Integer value and the corresponding unit (e.g. "hours"). If no unit yields an integer count, (1, "unknown") is returned.

Return type:

tuple[int, str]

Examples

>>> from datetime import timedelta
>>> get_integer_time(timedelta(days=14))
(2, 'weeks')
>>> get_integer_time(timedelta(hours=5))
(5, 'hours')
>>> get_integer_time(timedelta(milliseconds=1000))
(1, 'seconds')
>>> get_integer_time(timedelta(days=0.001))
(1, 'unknown')
neural_lam.utils.has_working_latex()#

Check whether a LaTeX toolchain is available on the system.

Returns:

True if latex and the required auxiliary tools are callable.

Return type:

bool

neural_lam.utils.init_training_logger_metrics(training_logger, val_steps)#

Configure validation metric aggregation for the active training logger.

Parameters:

  • training_logger (pytorch_lightning.loggers.Logger) – Logger instance used during training.

  • val_steps (Iterable[int]) – Autoregressive rollout lengths to log as separate metrics.

neural_lam.utils.inverse_sigmoid(x)#

Compute the logit (inverse sigmoid) while clamping to (0, 1).

Parameters:

x (torch.Tensor) – Input tensor assumed to contain logits after a sigmoid.

Returns:

Tensor containing log(x / (1 - x)) after clamping away from the saturation limits.

Return type:

torch.Tensor

neural_lam.utils.inverse_softplus(x, beta=1, threshold=20)#

Approximate the inverse of torch.nn.functional.softplus().

Parameters:

  • x (torch.Tensor) – Input tensor whose softplus inverse should be computed.

  • beta (float, optional) – Softplus beta parameter that controls the sharpness. Default 1.

  • threshold (float, optional) – Threshold applied to the input for numerical stability. Default 20.

Returns:

Tensor containing the inverse-softplus values.

Return type:

torch.Tensor

neural_lam.utils.load_graph(graph_dir_path, device='cpu')#

Load all tensors representing the graph from graph_dir_path.

Needs the following files for all graphs: - m2m_edge_index.pt - g2m_edge_index.pt - m2g_edge_index.pt - m2m_features.pt - g2m_features.pt - m2g_features.pt - mesh_features.pt

And in addition for hierarchical graphs: - mesh_up_edge_index.pt - mesh_down_edge_index.pt - mesh_up_features.pt - mesh_down_features.pt

Parameters:

  • graph_dir_path (str) – Path to directory containing the graph files.

  • device (str) – Device to load tensors to.

Returns:

  • hierarchical (bool) – Whether the graph is hierarchical.

  • graph (dict) – Dictionary containing the graph tensors, with keys as follows: - g2m_edge_index - m2g_edge_index - m2m_edge_index - mesh_up_edge_index - mesh_down_edge_index - g2m_features - m2g_features - m2m_features - mesh_up_features - mesh_down_features - mesh_static_features

neural_lam.utils.make_mlp(blueprint, layer_norm=True)#

Construct a multilayer perceptron from a blueprint of layer widths.

Parameters:

  • blueprint (list[int]) – Sequence of layer dimensions where blueprint[0] is the input size and blueprint[-1] is the output size.

  • layer_norm (bool, optional) – If True, append a LayerNorm to the output as in GraphCast.

Returns:

Sequential module implementing the specified MLP.

Return type:

torch.nn.Sequential

neural_lam.utils.rank_zero_print(*args, **kwargs)#

Print arguments only from the rank-zero process in distributed runs.

neural_lam.utils.setup_training_logger(datastore, args, run_name)#

Instantiate the configured experiment logger.

Parameters:

  • datastore (BaseDatastore) – Datastore providing metadata for logging configuration.

  • args (argparse.Namespace) – Parsed training arguments controlling the logger backend.

  • run_name (str) – Name of the run.

Returns:

logger – Logger object.

Return type:

pytorch_lightning.loggers.base

Contents