diff --git a/examples/10-tuning.py b/examples/10-tuning.py index b147848b..f66ac3c6 100644 --- a/examples/10-tuning.py +++ b/examples/10-tuning.py @@ -88,15 +88,24 @@ # and this is how long it took to run with these parameters (est.) -timings = TuningTimings(charges, cell, positions, cutoff=max_cutoff, run_backward=True) +timings = TuningTimings( + charges, + cell, + positions, + neighbor_indices=neighbor_indices, + neighbor_distances=neighbor_distances, + run_backward=True, +) estimated_timing = timings(pme) -print(f""" +print( + f""" Computed madelung constant: {madelung} Actual error: {madelung - madelung_ref} Estimated error: {estimated_error} Timing: {estimated_timing} seconds -""") +""" +) # %% # now set up a testing framework diff --git a/src/torchpme/calculators/calculator.py b/src/torchpme/calculators/calculator.py index 627283e5..bb9e22d1 100644 --- a/src/torchpme/calculators/calculator.py +++ b/src/torchpme/calculators/calculator.py @@ -40,9 +40,10 @@ def __init__( ): super().__init__() - assert isinstance(potential, Potential), ( - f"Potential must be an instance of Potential, got {type(potential)}" - ) + if not isinstance(potential, Potential): + raise TypeError( + f"Potential must be an instance of Potential, got {type(potential)}" + ) self.device = "cpu" if device is None else device self.dtype = torch.get_default_dtype() if dtype is None else dtype diff --git a/src/torchpme/tuning/ewald.py b/src/torchpme/tuning/ewald.py index dbc70814..730fce02 100644 --- a/src/torchpme/tuning/ewald.py +++ b/src/torchpme/tuning/ewald.py @@ -2,6 +2,7 @@ from typing import Optional import torch +import vesin.torch from ..calculators import EwaldCalculator from ..utils import _validate_parameters @@ -19,22 +20,22 @@ def tune_ewald( ns_lo: int = 1, ns_hi: int = 14, accuracy: float = 1e-3, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, ) -> tuple[float, dict[str, float]]: r""" Find the optimal parameters for :class:`torchpme.EwaldCalculator`. - The error formulas are given `online - `_ - (now not available, need to be updated later). Note the difference notation between - the parameters in the reference and ours: + .. note:: - .. math:: - - \alpha &= \left( \sqrt{2}\,\mathrm{smearing} \right)^{-1} - - K &= \frac{2 \pi}{\mathrm{lr\_wavelength}} - - r_c &= \mathrm{cutoff} + The :func:`torchpme.tuning.ewald.EwaldErrorBounds.forward` method takes floats + as the input, in order to be in consistency with the rest of the package -- + these parameters are always ``float`` but not ``torch.Tensor``. This design, + however, prevents the utilization of ``torch.autograd`` and other ``torch`` + features. To take advantage of these features, one can use the + :func:`torchpme.tuning.ewald.EwaldErrorBounds.err_rspace` and + :func:`torchpme.tuning.ewald.EwaldErrorBounds.err_kspace`, which takes + ``torch.Tensor`` as parameters. :param charges: torch.Tensor, atomic (pseudo-)charges :param cell: torch.Tensor, periodic supercell for the system @@ -72,25 +73,41 @@ def tune_ewald( _validate_parameters(charges, cell, positions, exponent) min_dimension = float(torch.min(torch.linalg.norm(cell, dim=1))) params = [{"lr_wavelength": min_dimension / ns} for ns in range(ns_lo, ns_hi + 1)] + if neighbor_indices is None and neighbor_distances is None: + nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False) + i, j, neighbor_distances = nl.compute( + points=positions.to(dtype=torch.float64, device="cpu"), + box=cell.to(dtype=torch.float64, device="cpu"), + periodic=True, + quantities="ijd", + ) + neighbor_indices = torch.stack([i, j], dim=1) + elif neighbor_indices is None or neighbor_distances is None: + raise ValueError( + "If neighbor_indices or neighbor_distances are None, both must be None." + ) + tuner = GridSearchTuner( - charges, - cell, - positions, - cutoff, + charges=charges, + cell=cell, + positions=positions, + cutoff=cutoff, exponent=exponent, neighbor_indices=neighbor_indices, neighbor_distances=neighbor_distances, calculator=EwaldCalculator, error_bounds=EwaldErrorBounds(charges=charges, cell=cell, positions=positions), params=params, + dtype=dtype, + device=device, ) smearing = tuner.estimate_smearing(accuracy) errs, timings = tuner.tune(accuracy) + # There are multiple errors below the accuracy, return the one with the shortest + # calculation time. The timing of those parameters leading to an higher error than + # the accuracy are set to infinity if any(err < accuracy for err in errs): - # There are multiple errors below the accuracy, return the one with the shortest - # calculation time. The timing of those parameters leading to an higher error - # than the accuracy are set to infinity return smearing, params[timings.index(min(timings))] # No parameter meets the requirement, return the one with the smallest error return smearing, params[errs.index(min(errs))] @@ -100,18 +117,9 @@ class EwaldErrorBounds(TuningErrorBounds): r""" Error bounds for :class:`torchpme.calculators.ewald.EwaldCalculator`. - The error formulas are given `online - `_ - (now not available, need to be updated later). Note the difference notation between - the parameters in the reference and ours: - - .. math:: - - \alpha &= \left( \sqrt{2}\,\mathrm{smearing} \right)^{-1} - - K &= \frac{2 \pi}{\mathrm{lr\_wavelength}} - - r_c &= \mathrm{cutoff} + .. math:: + \text{Error}_{\text{total}} = \sqrt{\text{Error}_{\text{real space}}^2 + + \text{Error}_{\text{Fourier space}}^2 :param charges: atomic charges :param cell: single tensor of shape (3, 3), describing the bounding diff --git a/src/torchpme/tuning/p3m.py b/src/torchpme/tuning/p3m.py index 4f2eb4f7..9102d38d 100644 --- a/src/torchpme/tuning/p3m.py +++ b/src/torchpme/tuning/p3m.py @@ -3,6 +3,7 @@ from typing import Optional import torch +import vesin.torch from ..calculators import P3MCalculator from ..utils import _validate_parameters @@ -79,6 +80,8 @@ def tune_p3m( mesh_lo: int = 2, mesh_hi: int = 7, accuracy: float = 1e-3, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, ): r""" Find the optimal parameters for :class:`torchpme.calculators.pme.PMECalculator`. @@ -141,18 +144,33 @@ def tune_p3m( range(nodes_lo, nodes_hi + 1), range(mesh_lo, mesh_hi + 1) ) ] + if neighbor_indices is None and neighbor_distances is None: + nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False) + i, j, neighbor_distances = nl.compute( + points=positions.to(dtype=torch.float64, device="cpu"), + box=cell.to(dtype=torch.float64, device="cpu"), + periodic=True, + quantities="ijd", + ) + neighbor_indices = torch.stack([i, j], dim=1) + elif neighbor_indices is None or neighbor_distances is None: + raise ValueError( + "If neighbor_indices or neighbor_distances are None, both must be None." + ) tuner = GridSearchTuner( - charges, - cell, - positions, - cutoff, + charges=charges, + cell=cell, + positions=positions, + cutoff=cutoff, exponent=exponent, neighbor_indices=neighbor_indices, neighbor_distances=neighbor_distances, calculator=P3MCalculator, error_bounds=P3MErrorBounds(charges=charges, cell=cell, positions=positions), params=params, + dtype=dtype, + device=device, ) smearing = tuner.estimate_smearing(accuracy) errs, timings = tuner.tune(accuracy) @@ -168,15 +186,18 @@ def tune_p3m( class P3MErrorBounds(TuningErrorBounds): r""" - " - Error bounds for :class:`torchpme.calculators.pme.P3MCalculator`. - - For the error formulas are given `here `_. - Note the difference notation between the parameters in the reference and ours: + " Error bounds for :class:`torchpme.calculators.pme.P3MCalculator`. - .. math:: + .. note:: - \alpha = \left(\sqrt{2}\,\mathrm{smearing} \right)^{-1} + The :func:`torchpme.tuning.p3m.P3MErrorBounds.forward` method takes floats as + the input, in order to be in consistency with the rest of the package -- these + parameters are always ``float`` but not ``torch.Tensor``. This design, however, + prevents the utilization of ``torch.autograd`` and other ``torch`` features. To + take advantage of these features, one can use the + :func:`torchpme.tuning.p3m.P3MErrorBounds.err_rspace` and + :func:`torchpme.tuning.p3m.P3MErrorBounds.err_kspace`, which takes + ``torch.Tensor`` as parameters. :param charges: atomic charges :param cell: single tensor of shape (3, 3), describing the bounding @@ -270,6 +291,10 @@ def forward( r""" Calculate the error bound of P3M. + .. math:: + \text{Error}_{\text{total}} = \sqrt{\text{Error}_{\text{real space}}^2 + + \text{Error}_{\text{Fourier space}}^2 + :param smearing: see :class:`torchpme.P3MCalculator` for details :param mesh_spacing: see :class:`torchpme.P3MCalculator` for details :param cutoff: see :class:`torchpme.P3MCalculator` for details diff --git a/src/torchpme/tuning/pme.py b/src/torchpme/tuning/pme.py index dda3d210..a849d4c8 100644 --- a/src/torchpme/tuning/pme.py +++ b/src/torchpme/tuning/pme.py @@ -3,6 +3,7 @@ from typing import Optional import torch +import vesin.torch from ..calculators import PMECalculator from ..utils import _validate_parameters @@ -22,6 +23,8 @@ def tune_pme( mesh_lo: int = 2, mesh_hi: int = 7, accuracy: float = 1e-3, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, ): r""" Find the optimal parameters for :class:`torchpme.PMECalculator`. @@ -84,18 +87,33 @@ def tune_pme( range(nodes_lo, nodes_hi + 1), range(mesh_lo, mesh_hi + 1) ) ] + if neighbor_indices is None and neighbor_distances is None: + nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False) + i, j, neighbor_distances = nl.compute( + points=positions.to(dtype=torch.float64, device="cpu"), + box=cell.to(dtype=torch.float64, device="cpu"), + periodic=True, + quantities="ijd", + ) + neighbor_indices = torch.stack([i, j], dim=1) + elif neighbor_indices is None or neighbor_distances is None: + raise ValueError( + "If neighbor_indices or neighbor_distances are None, both must be None." + ) tuner = GridSearchTuner( - charges, - cell, - positions, - cutoff, + charges=charges, + cell=cell, + positions=positions, + cutoff=cutoff, exponent=exponent, neighbor_indices=neighbor_indices, neighbor_distances=neighbor_distances, calculator=PMECalculator, error_bounds=PMEErrorBounds(charges=charges, cell=cell, positions=positions), params=params, + dtype=dtype, + device=device, ) smearing = tuner.estimate_smearing(accuracy) errs, timings = tuner.tune(accuracy) @@ -111,13 +129,18 @@ def tune_pme( class PMEErrorBounds(TuningErrorBounds): r""" - Error bounds for :class:`torchpme.PMECalculator`. For the error formulas are given - `elsewhere `_. Note the difference notation - between the parameters in the reference and ours: + Error bounds for :class:`torchpme.PMECalculator`. - .. math:: + .. note:: - \alpha = \left(\sqrt{2}\,\mathrm{smearing} \right)^{-1} + The :func:`torchpme.tuning.pme.PMEErrorBounds.forward` method takes floats as + the input, in order to be in consistency with the rest of the package -- these + parameters are always ``float`` but not ``torch.Tensor``. This design, however, + prevents the utilization of ``torch.autograd`` and other ``torch`` features. To + take advantage of these features, one can use the + :func:`torchpme.tuning.pme.PMEErrorBounds.err_rspace` and + :func:`torchpme.tuning.pme.PMEErrorBounds.err_kspace`, which takes + ``torch.Tensor`` as parameters. :param charges: atomic charges :param cell: single tensor of shape (3, 3), describing the bounding @@ -211,6 +234,10 @@ def error( r""" Calculate the error bound of PME. + .. math:: + \text{Error}_{\text{total}} = \sqrt{\text{Error}_{\text{real space}}^2 + + \text{Error}_{\text{Fourier space}}^2 + :param smearing: if its value is given, it will not be tuned, see :class:`torchpme.PMECalculator` for details :param mesh_spacing: if its value is given, it will not be tuned, see diff --git a/src/torchpme/tuning/tuner.py b/src/torchpme/tuning/tuner.py index 6ac250f7..d792f33e 100644 --- a/src/torchpme/tuning/tuner.py +++ b/src/torchpme/tuning/tuner.py @@ -3,10 +3,9 @@ from typing import Optional import torch -import vesin.torch from ..calculators import Calculator -from ..potentials import CoulombPotential +from ..potentials import InversePowerLawPotential from ..utils import _validate_parameters @@ -37,6 +36,9 @@ def __init__( def forward(self, *args, **kwargs): return self.error(*args, **kwargs) + def error(self, *args, **kwargs): + raise NotImplementedError + class TunerBase: """ @@ -79,6 +81,8 @@ def __init__( cutoff: float, calculator: type[Calculator], exponent: int = 1, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, ): _validate_parameters(charges, cell, positions, exponent) self.charges = charges @@ -87,8 +91,8 @@ def __init__( self.cutoff = cutoff self.calculator = calculator self.exponent = exponent - self._dtype = cell.dtype - self._device = cell.device + self.device = "cpu" if device is None else device + self.dtype = torch.get_default_dtype() if dtype is None else dtype self._prefac = 2 * float((charges**2).sum()) / math.sqrt(len(positions)) @@ -138,11 +142,11 @@ class GridSearchTuner(TunerBase): :param cutoff: real space cutoff, serves as a hyperparameter here. :param calculator: the calculator to be tuned :param params: list of Fourier space parameter sets for which the error is estimated - :param exponent: exponent of the potential, only exponent = 1 is supported :param neighbor_indices: torch.Tensor with the ``i,j`` indices of neighbors for which the potential should be computed in real space. :param neighbor_distances: torch.Tensor with the pair distances of the neighbors for which the potential should be computed in real space. + :param exponent: exponent of the potential, only exponent = 1 is supported """ def __init__( @@ -154,17 +158,21 @@ def __init__( calculator: type[Calculator], error_bounds: type[TuningErrorBounds], params: list[dict], + neighbor_indices: torch.Tensor, + neighbor_distances: torch.Tensor, exponent: int = 1, - neighbor_indices: Optional[torch.Tensor] = None, - neighbor_distances: Optional[torch.Tensor] = None, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, ): super().__init__( - charges, - cell, - positions, - cutoff, - calculator, - exponent, + charges=charges, + cell=cell, + positions=positions, + cutoff=cutoff, + calculator=calculator, + exponent=exponent, + dtype=dtype, + device=device, ) self.error_bounds = error_bounds self.params = params @@ -172,10 +180,11 @@ def __init__( charges, cell, positions, - cutoff, neighbor_indices, neighbor_distances, True, + dtype=dtype, + device=device, ) def tune(self, accuracy: float = 1e-3) -> tuple[list[float], list[float]]: @@ -193,22 +202,24 @@ def tune(self, accuracy: float = 1e-3) -> tuple[list[float], list[float]]: param_errors = [] param_timings = [] for param in self.params: - error = self.error_bounds(smearing=smearing, cutoff=self.cutoff, **param) # type: ignore[call-arg] + error = self.error_bounds(smearing=smearing, cutoff=self.cutoff, **param) # type: ignore[call-arg] param_errors.append(float(error)) - if error > accuracy: - param_timings.append(float("inf")) - continue - - param_timings.append(self._timing(smearing, param)) + param_timings.append( + self._timing(smearing, param) if error <= accuracy else float("inf") + ) return param_errors, param_timings def _timing(self, smearing: float, k_space_params: dict): calculator = self.calculator( - potential=CoulombPotential( - # exponent=self.exponent, # but only exponent = 1 is supported + potential=InversePowerLawPotential( + exponent=self.exponent, # but only exponent = 1 is supported smearing=smearing, + device=self.device, + dtype=self.dtype, ), + device=self.device, + dtype=self.dtype, **k_space_params, ) @@ -242,39 +253,26 @@ def __init__( charges: torch.Tensor, cell: torch.Tensor, positions: torch.Tensor, - cutoff: float, - neighbor_indices: Optional[torch.Tensor] = None, - neighbor_distances: Optional[torch.Tensor] = None, + neighbor_indices: torch.Tensor, + neighbor_distances: torch.Tensor, n_repeat: int = 4, n_warmup: int = 2, run_backward: Optional[bool] = True, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, ): super().__init__() - self._charges = charges - self._cell = cell - self._positions = positions - self._dtype = charges.dtype - self._device = charges.device - self._n_repeat = n_repeat - self._n_warmup = n_warmup - self._run_backward = run_backward - - if neighbor_indices is None and neighbor_distances is None: - nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False) - i, j, neighbor_distances = nl.compute( - points=self._positions.to(dtype=torch.float64, device="cpu"), - box=self._cell.to(dtype=torch.float64, device="cpu"), - periodic=True, - quantities="ijd", - ) - neighbor_indices = torch.stack([i, j], dim=1) - elif neighbor_indices is None or neighbor_distances is None: - raise ValueError( - "If neighbor_indices or neighbor_distances are None, both must be None." - ) - self._neighbor_indices = neighbor_indices.to(device=self._device) - self._neighbor_distances = neighbor_distances.to( - dtype=self._dtype, device=self._device + self.charges = charges + self.cell = cell + self.positions = positions + self.dtype = dtype + self.device = device + self.n_repeat = n_repeat + self.n_warmup = n_warmup + self.run_backward = run_backward + self.neighbor_indices = neighbor_indices.to(device=self.device) + self.neighbor_distances = neighbor_distances.to( + dtype=self.dtype, device=self.device ) def forward(self, calculator: torch.nn.Module): @@ -285,24 +283,24 @@ def forward(self, calculator: torch.nn.Module): :param calculator: the calculator to be tuned :return: a float, the average execution time """ - for _ in range(self._n_warmup): + for _ in range(self.n_warmup): result = calculator.forward( - positions=self._positions, - charges=self._charges, - cell=self._cell, - neighbor_indices=self._neighbor_indices, - neighbor_distances=self._neighbor_distances, + positions=self.positions, + charges=self.charges, + cell=self.cell, + neighbor_indices=self.neighbor_indices, + neighbor_distances=self.neighbor_distances, ) # measure time execution_time = 0.0 - for _ in range(self._n_repeat): - positions = self._positions.clone() - cell = self._cell.clone() - charges = self._charges.clone() + for _ in range(self.n_repeat): + positions = self.positions.clone() + cell = self.cell.clone() + charges = self.charges.clone() # nb - this won't compute gradiens involving the distances - if self._run_backward: + if self.run_backward: positions.requires_grad_(True) cell.requires_grad_(True) charges.requires_grad_(True) @@ -311,15 +309,15 @@ def forward(self, calculator: torch.nn.Module): positions=positions, charges=charges, cell=cell, - neighbor_indices=self._neighbor_indices, - neighbor_distances=self._neighbor_distances, + neighbor_indices=self.neighbor_indices, + neighbor_distances=self.neighbor_distances, ) value = result.sum() - if self._run_backward: + if self.run_backward: value.backward(retain_graph=True) - if self._device is torch.device("cuda"): + if self.device is torch.device("cuda"): torch.cuda.synchronize() execution_time += time.time() - return execution_time / self._n_repeat + return execution_time / self.n_repeat