Fix neighbor list bug for non-orthogonal cells

Author: peterspackman

examples/simple_inference.cpp

@@ -188,27 +188,6 @@ int main(int argc, char **argv) {
     return 1;
   }
 
-  // Build neighbor list
-  NeighborListOptions nl_opts;
-  nl_opts.cutoff =
-      5.0f; // PET-MAD uses 5.0 A cutoff (match PyTorch for comparison)
-  nl_opts.full_list = true;
-
-  NeighborListBuilder builder(nl_opts);
-  auto nlist = builder.build(system);
-
-  log::debug("Built neighbor list with {} pairs", nlist.num_pairs());
-
-  // Print first 5 edges
-  for (int e = 0; e < std::min(5, nlist.num_pairs()); ++e) {
-    log::trace(
-        "  {}: ({}->{}) shift=[{},{},{}] D=[{:.3f},{:.3f},{:.3f}] d={:.3f}", e,
-        nlist.centers[e], nlist.neighbors[e], nlist.cell_shifts[e][0],
-        nlist.cell_shifts[e][1], nlist.cell_shifts[e][2],
-        nlist.edge_vectors[e][0], nlist.edge_vectors[e][1],
-        nlist.edge_vectors[e][2], nlist.distances[e]);
-  }
-
   // Load model and run inference
   try {
     log::info("Loading model from {}", model_path);
@@ -234,6 +213,15 @@ int main(int argc, char **argv) {
       log::info("Overriding cutoff to: {:.2f} A", cutoff_override);
     }
 
+    // Log neighbor count using model's cutoff
+    {
+      NeighborListBuilder nl_builder(
+          NeighborListOptions{pet_model.cutoff(), true, false});
+      auto nlist = nl_builder.build(system);
+      log::info("Neighbor pairs: {} (avg {:.1f} per atom)", nlist.num_pairs(),
+                static_cast<double>(nlist.num_pairs()) / system.num_atoms());
+    }
+
     static constexpr std::array backend_names = {"auto", "cpu", "cuda", "hip",
                                                   "metal", "vulkan", "sycl", "cann"};
     log::info("Backend preference: {}", backend_names[static_cast<size_t>(backend_pref)]);

src/core/neighbor_list.cpp

@@ -67,9 +67,12 @@ struct CellGrid {
 
   std::array<int, 3> position_to_bin(const Vec3 &pos) const {
     double frac[3];
+    // Compute fractional coordinates: frac = pos @ inv_cell (i.e., inv_cell.T @ pos)
+    // For the standard convention where rows of the cell matrix are lattice vectors,
+    // fractional coords are frac[i] = sum_j pos[j] * inv_cell[j,i]
     for (int i = 0; i < 3; ++i) {
-      frac[i] = inv_cell[i * 3 + 0] * pos.x + inv_cell[i * 3 + 1] * pos.y +
-                inv_cell[i * 3 + 2] * pos.z;
+      frac[i] = inv_cell[0 * 3 + i] * pos.x + inv_cell[1 * 3 + i] * pos.y +
+                inv_cell[2 * 3 + i] * pos.z;
     }
     return {static_cast<int>(std::floor(frac[0] * n_bins[0])),
             static_cast<int>(std::floor(frac[1] * n_bins[1])),