diff --git a/examples/01-charges-example.py b/examples/01-charges-example.py index 92d82949..53d6b964 100644 --- a/examples/01-charges-example.py +++ b/examples/01-charges-example.py @@ -42,12 +42,12 @@ # %% # # Create the properties CsCl unit cell - +dtype = torch.float64 symbols = ("Cs", "Cl") types = torch.tensor([55, 17]) -charges = torch.tensor([[1.0], [-1.0]], dtype=torch.float64) -positions = torch.tensor([(0, 0, 0), (0.5, 0.5, 0.5)], dtype=torch.float64) -cell = torch.eye(3, dtype=torch.float64) +charges = torch.tensor([[1.0], [-1.0]], dtype=dtype) +positions = torch.tensor([(0, 0, 0), (0.5, 0.5, 0.5)], dtype=dtype) +cell = torch.eye(3, dtype=dtype) pbc = torch.tensor([True, True, True]) @@ -72,6 +72,7 @@ cutoff=cutoff, neighbor_indices=neighbor_indices, neighbor_distances=neighbor_distances, + dtype=dtype, ) # %% @@ -101,7 +102,7 @@ # will be used to *compute* the potential energy of the system. calculator = torchpme.PMECalculator( - torchpme.CoulombPotential(smearing=smearing), **pme_params + torchpme.CoulombPotential(smearing=smearing, dtype=dtype), dtype=dtype, **pme_params ) # %% @@ -112,7 +113,7 @@ # As a first application of multiple charge channels, we start simply by using the # classic definition of one charge channel per atom. -charges = torch.tensor([[1.0], [-1.0]], dtype=torch.float64) +charges = torch.tensor([[1.0], [-1.0]], dtype=dtype) # %% # @@ -160,7 +161,7 @@ # species-specific potentials and facilitating the learning process for machine learning # algorithms. -charges_one_hot = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=torch.float64) +charges_one_hot = torch.tensor([[1.0, 0.0], [0.0, 1.0]], dtype=dtype) # %% # @@ -205,7 +206,7 @@ # creating a new calculator with the metatensor interface. calculator_metatensor = torchpme.metatensor.PMECalculator( - torchpme.CoulombPotential(smearing=smearing), **pme_params + torchpme.CoulombPotential(smearing=smearing, dtype=dtype), dtype=dtype, **pme_params ) # %% diff --git a/examples/02-neighbor-lists-usage.py b/examples/02-neighbor-lists-usage.py index 322e6a69..e72689f0 100644 --- a/examples/02-neighbor-lists-usage.py +++ b/examples/02-neighbor-lists-usage.py @@ -55,6 +55,7 @@ # # As a test system, we use a 2x2x2 supercell of an CsCl crystal in a cubic cell. +dtype = torch.float64 atoms_unitcell = ase.Atoms( symbols=["Cs", "Cl"], positions=np.array([(0, 0, 0), (0.5, 0.5, 0.5)]), @@ -97,7 +98,7 @@ nl = vesin.torch.NeighborList(cutoff=cutoff, full_list=False) neighbor_indices, neighbor_distances = nl.compute( points=positions.to(dtype=torch.float64, device="cpu"), - box=cell.to(dtype=torch.float64, device="cpu"), + box=cell.to(dtype=dtype, device="cpu"), periodic=True, quantities="Pd", ) @@ -109,6 +110,7 @@ cutoff=cutoff, neighbor_indices=neighbor_indices, neighbor_distances=neighbor_distances, + dtype=dtype, ) # %% @@ -193,7 +195,9 @@ def distances( # compute the potential. pme = torchpme.PMECalculator( - potential=torchpme.CoulombPotential(smearing=smearing), **pme_params + potential=torchpme.CoulombPotential(smearing=smearing, dtype=dtype), + dtype=dtype, + **pme_params, ) potential = pme( charges=charges, diff --git a/examples/08-combined-potential.py b/examples/08-combined-potential.py index df455e98..8d0667bb 100644 --- a/examples/08-combined-potential.py +++ b/examples/08-combined-potential.py @@ -28,6 +28,8 @@ from torchpme import CombinedPotential, EwaldCalculator, InversePowerLawPotential from torchpme.prefactors import eV_A +dtype = torch.float64 + # %% # Combined potentials # ------------------- @@ -65,10 +67,10 @@ # evaluation, and so one has to set it also for the combined potential, even if it is # not used explicitly in the evaluation of the combination. -pot_1 = InversePowerLawPotential(exponent=1, smearing=smearing) -pot_2 = InversePowerLawPotential(exponent=2, smearing=smearing) +pot_1 = InversePowerLawPotential(exponent=1, smearing=smearing, dtype=dtype) +pot_2 = InversePowerLawPotential(exponent=2, smearing=smearing, dtype=dtype) -potential = CombinedPotential(potentials=[pot_1, pot_2], smearing=smearing) +potential = CombinedPotential(potentials=[pot_1, pot_2], smearing=smearing, dtype=dtype) # Note also that :class:`CombinedPotential` can be used with any combination of # potentials, as long they are all either direct or range separated. For instance, one @@ -80,7 +82,7 @@ # We now plot of the individual and combined ``potential`` functions together with an # explicit sum of the two potentials. -dist = torch.logspace(-3, 2, 1000) +dist = torch.logspace(-3, 2, 1000, dtype=dtype) fig, ax = plt.subplots() @@ -115,7 +117,7 @@ # combines all terms in a range-separated potential, including the k-space # kernel. -k = torch.logspace(-2, 2, 1000) +k = torch.logspace(-2, 2, 1000, dtype=dtype) fig, ax = plt.subplots() @@ -154,9 +156,8 @@ # much bigger system. calculator = EwaldCalculator( - potential=potential, lr_wavelength=lr_wavelength, prefactor=eV_A + potential=potential, lr_wavelength=lr_wavelength, prefactor=eV_A, dtype=dtype ) -calculator.to(dtype=torch.float64) # %% diff --git a/examples/10-tuning.py b/examples/10-tuning.py index c2d61881..183b0f45 100644 --- a/examples/10-tuning.py +++ b/examples/10-tuning.py @@ -120,7 +120,9 @@ pme_params = {"mesh_spacing": 1.0, "interpolation_nodes": 4} pme = torchpme.PMECalculator( - potential=torchpme.CoulombPotential(smearing=smearing), + potential=torchpme.CoulombPotential(smearing=smearing, device=device, dtype=dtype), + device=device, + dtype=dtype, **pme_params, # type: ignore[arg-type] ) @@ -168,6 +170,8 @@ neighbor_indices=neighbor_indices, neighbor_distances=neighbor_distances, run_backward=True, + device=device, + dtype=dtype, ) estimated_timing = timings(pme) @@ -210,15 +214,19 @@ def filter_neighbors(cutoff, neighbor_indices, neighbor_distances): return neighbor_indices[filter_idx], neighbor_distances[filter_idx] -def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes): +def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes, device, dtype): filter_indices, filter_distances = filter_neighbors( cutoff, neighbor_indices, neighbor_distances ) pme = torchpme.PMECalculator( - potential=torchpme.CoulombPotential(smearing=smearing), + potential=torchpme.CoulombPotential( + smearing=smearing, device=device, dtype=dtype + ), mesh_spacing=mesh_spacing, interpolation_nodes=interpolation_nodes, + device=device, + dtype=dtype, ) potential = pme( charges=charges, @@ -239,6 +247,8 @@ def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes): run_backward=True, n_warmup=1, n_repeat=4, + device=device, + dtype=dtype, ) estimated_timing = timings(pme) return madelung, estimated_timing @@ -251,7 +261,9 @@ def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes): bounds = np.zeros((len(smearing_grid), len(spacing_grid))) for ism, smearing in enumerate(smearing_grid): for isp, spacing in enumerate(spacing_grid): - results[ism, isp], timings[ism, isp] = timed_madelung(8.0, smearing, spacing, 4) + results[ism, isp], timings[ism, isp] = timed_madelung( + 8.0, smearing, spacing, 4, device, dtype + ) bounds[ism, isp] = error_bounds(8.0, smearing, spacing, 4) # %% @@ -374,7 +386,7 @@ def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes): for inint, nint in enumerate(nint_grid): for isp, spacing in enumerate(spacing_grid): results[inint, isp], timings[inint, isp] = timed_madelung( - 5.0, 1.0, spacing, nint + 5.0, 1.0, spacing, nint, device=device, dtype=dtype ) bounds[inint, isp] = error_bounds(5.0, 1.0, spacing, nint) @@ -445,15 +457,19 @@ def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes): cutoff=5.0, neighbor_indices=neighbor_indices, neighbor_distances=neighbor_distances, + device=device, + dtype=dtype, ) -print(f""" +print( + f""" Estimated PME parameters (cutoff={5.0} Å): Smearing: {smearing} Å Mesh spacing: {parameters["mesh_spacing"]} Å Interpolation order: {parameters["interpolation_nodes"]} Estimated time per step: {timing} s -""") +""" +) # %% # What is the best cutoff? @@ -476,6 +492,8 @@ def timed_madelung(cutoff, smearing, mesh_spacing, interpolation_nodes): cutoff=cutoff, neighbor_indices=filter_indices, neighbor_distances=filter_distances, + device=device, + dtype=dtype, ) timings_grid.append(timing) diff --git a/examples/basic-usage.py b/examples/basic-usage.py index d726a18f..47bc1506 100644 --- a/examples/basic-usage.py +++ b/examples/basic-usage.py @@ -146,7 +146,7 @@ # contains all the necessary functions (such as those defining the short-range and # long-range splits) for this potential and makes them useable in the rest of the code. -potential = CoulombPotential(smearing=smearing) +potential = CoulombPotential(smearing=smearing, device=device, dtype=dtype) # %% # @@ -193,7 +193,9 @@ # Since our structure is relatively small, we use the :class:`EwaldCalculator`. # We start by the initialization of the class. -calculator = EwaldCalculator(potential=potential, lr_wavelength=lr_wavelength) +calculator = EwaldCalculator( + potential=potential, lr_wavelength=lr_wavelength, device=device, dtype=dtype +) # %% # diff --git a/tests/helpers.py b/tests/helpers.py index d38e9f9b..bca5feb7 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -256,7 +256,10 @@ def neighbor_list( nl = NeighborList(cutoff=cutoff, full_list=full_neighbor_list) neighbor_indices, d, S = nl.compute( - points=positions, box=box, periodic=periodic, quantities="PdS" + points=positions.to(dtype=torch.float64, device="cpu"), + box=box.to(dtype=torch.float64, device="cpu"), + periodic=periodic, + quantities="PdS", ) neighbor_indices = torch.from_numpy(neighbor_indices.astype(int)).to(