If task cells are not immediate neighbours, we cannot use these as edge weights, so ignore

These are expected to be gravity tasks

src/partition.c

@@ -617,41 +617,46 @@ static void repart_edge_metis(int partweights, int bothweights, int timebins,
           if (cj->nodeID == nodeID) weights_v[cjd] += 0.5 * w;
         }
 
-        if (timebins) {
-          /* Add weights to edge for all cells based on the expected
-           * interaction time (calculated as the time to the last expected
-           * time) as we want to avoid having active cells on the edges, so we
-           * cut for that. Note that weight is added to the local and remote
-           * cells, as we want to keep both away from any cuts, this can
-           * overflow int, so take care. */
-          int dti = num_time_bins - get_time_bin(ci->ti_end_min);
-          int dtj = num_time_bins - get_time_bin(cj->ti_end_min);
-          double dt = (double)(1 << dti) + (double)(1 << dtj);
-
-          /* ci */
-          int kk;
-          for (kk = 26 * cid; inds[kk] != cjd; kk++)
-            ;
-          weights_e[kk] += dt;
-
-          /* cj */
-          for (kk = 26 * cjd; inds[kk] != cid; kk++)
-            ;
-          weights_e[kk] += dt;
-        } else {
-
-          /* Add weights from task costs to the edge. */
-
-          /* ci */
-          int kk;
-          for (kk = 26 * cid; inds[kk] != cjd; kk++)
-            ;
-          weights_e[kk] += w;
-
-          /* cj */
-          for (kk = 26 * cjd; inds[kk] != cid; kk++)
-            ;
-          weights_e[kk] += w;
+        /* Find indices of ci/cj neighbours. Note with gravity these cells may
+         * not be neighbours, in that case we ignore any edge weight for that
+         * pair. */
+        int ik = -1;
+        for (int k = 26 * cid; k < 26 * nr_cells; k++) {
+          if (inds[k] == cjd) {
+            ik = k;
+            break;
+          }
+        }
+
+        /* cj */
+        int jk = -1;
+        for (int k = 26 * cjd; k < 26 * nr_cells; k++) {
+          if (inds[k] == cid) {
+            jk = k;
+            break;
+          }
+        }
+        if (ik != -1 && jk != -1) {
+
+          if (timebins) {
+            /* Add weights to edge for all cells based on the expected
+             * interaction time (calculated as the time to the last expected
+             * time) as we want to avoid having active cells on the edges, so
+             * we cut for that. Note that weight is added to the local and
+             * remote cells, as we want to keep both away from any cuts, this
+             * can overflow int, so take care. */
+            int dti = num_time_bins - get_time_bin(ci->ti_end_min);
+            int dtj = num_time_bins - get_time_bin(cj->ti_end_min);
+            double dt = (double)(1 << dti) + (double)(1 << dtj);
+            weights_e[ik] += dt;
+            weights_e[jk] += dt;
+
+          } else {
+
+            /* Add weights from task costs to the edge. */
+            weights_e[ik] += w;
+            weights_e[jk] += w;
+          }
         }
       }
     }