Unique variable names in madelung example

examples/madelung.py

@@ -49,7 +49,7 @@
     interpolation_order=interpolation_order,
     subtract_self=True,
 )
-potentials_mesh = MP.compute(frame)
+potentials_torch = MP.compute(frame)
 
 # %%
 # The "potentials" that have been computed so far are not the actual electrostatic
@@ -58,27 +58,28 @@
 # separately. Thus, to get the Madelung constant, we need to take a linear combination
 # of these "potentials" weighted by the charges of the atoms.
 
-atomic_energies = torch.zeros((n_atoms, 1))
+atomic_energies_torch = torch.zeros((n_atoms, 1))
 for idx_c in range(n_atoms):
     for idx_n in range(n_atoms):
         # The coulomb potential between atoms i and j is charge_i * charge_j / d_ij
         # The features are simply computing a pure 1/r potential with no prefactors.
         # Thus, to compute the energy between atoms of species i and j, we need to
         # multiply by the charges of i and j.
-        atomic_energies[idx_c] += (
-            charges[idx_c] * charges[idx_n] * potentials_mesh[idx_c, idx_n]
+        print(charges[idx_c] * charges[idx_n], potentials_torch[idx_n, idx_c])
+        atomic_energies_torch[idx_c] += (
+            charges[idx_c] * charges[idx_n] * potentials_torch[idx_c, idx_n]
         )
 
 # %%
 # The total energy is just the sum of all atomic energies
-total_energy = torch.sum(atomic_energies)
+total_energy_torch = torch.sum(atomic_energies_torch)
 
 # %%
 # Compare against reference Madelung constant and reference energy:
 print("Using the torch version")
-print(f"Computed energies on each atom = {atomic_energies.tolist()}")
+print(f"Computed energies on each atom = {atomic_energies_torch.tolist()}")
 print(f"Reference Madelung constant = {madelung:.3f}")
-print(f"Total energy = {total_energy:.3f}\n")
+print(f"Total energy = {total_energy_torch:.3f}\n")
 
 
 # %%
@@ -93,34 +94,37 @@
     interpolation_order=interpolation_order,
     subtract_self=True,
 )
-potentials_mesh = MP.compute(frame)
+potential_metatensor = MP.compute(frame)
 
 
 # %%
 # To get the Madelung constant, we again need to take a linear combination
 # of the "potentials" weighted by the charges of the atoms.
 
-atomic_energies = torch.zeros((n_atoms, 1))
+atomic_energies_metatensor = torch.zeros((n_atoms, 1))
 for idx_c, c in enumerate(atomic_numbers):
     for idx_n, n in enumerate(atomic_numbers):
         # Take the coefficients with the correct center atom and neighbor atom species
-        block = potentials_mesh.block(
+        block = potential_metatensor.block(
             {"species_center": int(c), "species_neighbor": int(n)}
         )
 
         # The coulomb potential between atoms i and j is charge_i * charge_j / d_ij
         # The features are simply computing a pure 1/r potential with no prefactors.
         # Thus, to compute the energy between atoms of species i and j, we need to
         # multiply by the charges of i and j.
-        atomic_energies[idx_c] += charges[idx_c] * charges[idx_n] * block.values[0, 0]
+        print(c, n, charges[idx_c] * charges[idx_n], block.values[0, 0])
+        atomic_energies_metatensor[idx_c] += (
+            charges[idx_c] * charges[idx_n] * block.values[0, 0]
+        )
 
 # %%
 # The total energy is just the sum of all atomic energies
-total_energy = torch.sum(atomic_energies)
+total_energy_metatensor = torch.sum(atomic_energies_metatensor)
 
 # %%
 # Compare against reference Madelung constant and reference energy:
 print("Using the metatensor version")
-print(f"Computed energies on each atom = {atomic_energies.tolist()}")
+print(f"Computed energies on each atom = {atomic_energies_metatensor.tolist()}")
 print(f"Reference Madelung constant = {madelung:.3f}")
-print(f"Total energy = {total_energy:.3f}")
+print(f"Total energy = {total_energy_metatensor:.3f}")

src/meshlode/calculators/meshpotential.py

@@ -111,15 +111,8 @@ def compute(
         if not isinstance(systems, list):
             systems = [systems]
 
-        # Generate a dictionary to map atomic species to array indices
-        # In general, the species are sorted according to atomic number
-        # and assigned the array indices 0, 1, 2,...
-        # Example: for H2O: H is mapped to 0 and O is mapped to 1.
-        all_species = []
-        n_atoms_tot = 0
-        for system in systems:
-            n_atoms_tot += len(system)
-            all_species.append(system.species)
+        # Extract all species/atomic_numbers from the list of systems
+        all_species = [system.species for system in systems]
         all_species = torch.hstack(all_species)
         atomic_numbers = _my_1d_tolist(torch.unique(all_species))
         n_species = len(atomic_numbers)