Base Classes

Current scheme behind all tuning functions is grid-searching based, focusing on the Fourier space parameters like lr_wavelength, mesh_spacing and interpolation_nodes. For real space parameter cutoff, it is treated as a hyperparameter here, which should be manually specified by the user. The parameter smearing is determined by the real space error formula and is set to achieve a real space error of desired_accuracy / 4.

The Fourier space parameters are all discrete, so it’s convenient to do the grid-search. Default searching-ranges are provided for those parameters. For lr_wavelength, the values are chosen to be with a minimum of 1 and a maximum of 13 mesh points in each spatial direction (x, y, z). For mesh_spacing, the values are set to have minimally 2 and maximally 7 mesh points in each spatial direction, for both the P3M and PME method. The values of interpolation_nodes are the same as those supported in torchpme.lib.MeshInterpolator.

In the grid-searching, all possible parameter combinations are evaluated. The error associated with the parameter is estimated by the error formulas implemented in the subclasses of torchpme.tuning.tuner.TuningErrorBounds. Parameter with the error within the desired accuracy are benchmarked for computational time by torchpme.tuning.tuner.TuningTimings The timing of the other parameters are not tested and set to infinity.

The return of these tuning functions contains the smearing and a dictionary, in which there is parameter for the Fourier space. The parameter is that of the desired accuracy and the shortest timing. The parameter of the smallest error will be returned in the case that no parameter can fulfill the accuracy requirement.

class torchpme.tuning.tuner.TunerBase(charges: Tensor, cell: Tensor, positions: Tensor, cutoff: float, calculator: type[Calculator], exponent: int = 1, full_neighbor_list: bool = False, prefactor: float = 1.0)[source]

Base class defining the interface for a parameter tuner.

This class provides a framework for tuning the parameters of a calculator. The class itself supports estimating the smearing from the real space cutoff based on the real space error formula. The TunerBase.tune() defines the interface for a sophisticated tuning process, which takes a value of the desired accuracy.

Parameters:
  • charges (Tensor) – torch.tensor of shape (len(positions, 1)) containing the atomic (pseudo-)charges

  • cell (Tensor) – torch.tensor of shape (3, 3), where cell[i] is the i-th basis vector of the unit cell

  • positions (Tensor) – torch.tensor of shape (N, 3) containing the Cartesian coordinates of the N particles within the supercell.

  • cutoff (float) – real space cutoff, serves as a hyperparameter here.

  • calculator (type[Calculator]) – the calculator to be tuned

  • exponent (int) – exponent of the potential, only exponent = 1 is supported

  • full_neighbor_list (bool) – If set to True, a “full” neighbor list is expected as input. This means that each atom pair appears twice. If set to False, a “half” neighbor list is expected.

  • prefactor (float) – electrostatics prefactor; see Prefactors for details and common values.

Example

>>> import torch
>>> import torchpme
>>> positions = torch.tensor([[0.0, 0.0, 0.0], [0.4, 0.4, 0.4]])
>>> charges = torch.tensor([[1.0], [-1.0]])
>>> cell = torch.eye(3)
>>> tuner = TunerBase(charges, cell, positions, 4.4, torchpme.EwaldCalculator)
>>> smearing = tuner.estimate_smearing(1e-3)
>>> print(smearing)
1.1069526756106463
estimate_smearing(accuracy: float) float[source]

Estimate the smearing based on the error formula of the real space. The smearing is set as leading to a real space error of accuracy/4.

Parameters:

accuracy (float) – a float, the desired accuracy

Returns:

a float, the estimated smearing

Return type:

float

static filter_neighbors(cutoff: float, neighbor_indices: Tensor, neighbor_distances: Tensor)[source]

Filter neighbor indices and distances based on a user given cutoff. This allows users pre-computing the neighbor list with a larger cutoff and then filtering the neighbors based on a smaller cutoff, leading to a faster tuning on the cutoff.

Parameters:
  • cutoff (float) – real space cutoff

  • neighbor_indices (Tensor) – torch.tensor with the i,j indices of neighbors for which the potential should be computed in real space.

  • neighbor_distances (Tensor) – torch.tensor with the pair distances of the neighbors for which the potential should be computed in real space.

class torchpme.tuning.tuner.GridSearchTuner(charges: Tensor, cell: Tensor, positions: Tensor, cutoff: float, calculator: type[Calculator], error_bounds: type[TuningErrorBounds], params: list[dict], neighbor_indices: Tensor, neighbor_distances: Tensor, full_neighbor_list: bool = False, prefactor: float = 1.0, exponent: int = 1)[source]

Tuner using grid search.

The tuner uses the error formula to estimate the error of a given parameter set. If the error is smaller than the accuracy, the timing is measured and returned. If the error is larger than the accuracy, the timing is set to infinity and the parameter is skipped.

Note

The cutoff is treated as a hyperparameter here. In case one wants to tune the cutoff, one could instantiate the tuner with different cutoff values and manually pick the best from the tuning results.

Parameters:
  • charges (Tensor) – torch.tensor of shape (len(positions, 1)) containing the atomic (pseudo-)charges

  • cell (Tensor) – torch.tensor of shape (3, 3), where cell[i] is the i-th basis vector of the unit cell

  • positions (Tensor) – torch.tensor of shape (N, 3) containing the Cartesian coordinates of the N particles within the supercell.

  • cutoff (float) – real space cutoff, serves as a hyperparameter here.

  • calculator (type[Calculator]) – the calculator to be tuned

  • error_bounds (type[TuningErrorBounds]) – error bounds for the calculator

  • params (list[dict]) – list of Fourier space parameter sets for which the error is estimated

  • neighbor_indices (Tensor) – torch.tensor with the i,j indices of neighbors for which the potential should be computed in real space.

  • neighbor_distances (Tensor) – torch.tensor with the pair distances of the neighbors for which the potential should be computed in real space.

  • full_neighbor_list (bool) – If set to True, a “full” neighbor list is expected as input. This means that each atom pair appears twice. If set to False, a “half” neighbor list is expected.

  • prefactor (float) – electrostatics prefactor; see Prefactors for details and common values.

  • exponent (int) – exponent of the potential, only exponent = 1 is supported

tune(accuracy: float = 0.001) tuple[list[float], list[float]][source]

Estimate the error and timing for each parameter set. Only parameters for which the error is smaller than the accuracy are timed, the others’ timing is set to infinity.

Parameters:

accuracy (float) – a float, the desired accuracy

Returns:

a list of errors and a list of timings

Return type:

tuple[list[float], list[float]]

class torchpme.tuning.tuner.TuningTimings(charges: Tensor, cell: Tensor, positions: Tensor, neighbor_indices: Tensor, neighbor_distances: Tensor, n_repeat: int = 4, n_warmup: int = 4, run_backward: bool | None = True)[source]

Class for timing a calculator.

The class estimates the average execution time of a given calculater after several warmup runs. The class takes the information of the structure that one wants to benchmark on, and the configuration of the timing process as inputs.

Parameters:
  • charges (Tensor) – torch.tensor of shape (len(positions, 1)) containing the atomic (pseudo-)charges

  • cell (Tensor) – torch.tensor of shape (3, 3), where cell[i] is the i-th basis vector of the unit cell

  • positions (Tensor) – torch.tensor of shape (N, 3) containing the Cartesian coordinates of the N particles within the supercell.

  • cutoff – real space cutoff, serves as a hyperparameter here.

  • neighbor_indices (Tensor) – torch.tensor with the i,j indices of neighbors for which the potential should be computed in real space.

  • neighbor_distances (Tensor) – torch.tensor with the pair distances of the neighbors for which the potential should be computed in real space.

  • n_repeat (int) – number of times to repeat to estimate the average timing

  • n_warmup (int) – number of warmup runs, recommended to be at least 4

  • run_backward (bool | None) – whether to run the backward pass

forward(calculator: Module)[source]

Estimate the execution time of a given calculator for the structure to be used as benchmark.

Parameters:

calculator (Module) – the calculator to be tuned

Returns:

a float, the average execution time

class torchpme.tuning.tuner.TuningErrorBounds(charges: Tensor, cell: Tensor, positions: Tensor)[source]

Base class for error bounds. This class calculates the real space error and the Fourier space error based on the error formula. This class is used in the tuning process. It can also be used with the torchpme.tuning.tuner.TunerBase to build up a custom parameter tuner.

Parameters:
  • charges (Tensor) – torch.tensor of shape (len(positions, 1)) containing the atomic (pseudo-)charges

  • cell (Tensor) – torch.tensor of shape (3, 3), where cell[i] is the i-th basis vector of the unit cell

  • positions (Tensor) – torch.tensor of shape (N, 3) containing the Cartesian coordinates of the N particles within the supercell.

forward(*args, **kwargs)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Examples using Tuning Classes