Added the gravity-specific tasks to the weight calculation.

src/partition.c

@@ -577,9 +577,11 @@ static void repart_edge_metis(int partweights, int bothweights, int timebins,
     int cid = ci - cells;
 
     /* Different weights for different tasks. */
-    if (t->type == task_type_ghost || t->type == task_type_kick1 ||
-        t->type == task_type_kick2 || t->type == task_type_timestep ||
-        t->type == task_type_drift_part || t->type == task_type_drift_gpart) {
+    if (t->type == task_type_drift_part || t->type == task_type_drift_gpart ||
+	t->type == task_type_ghost || t->type == task_type_extra_ghost ||
+	t->type == task_type_kick1 || t->type == task_type_kick2 || 
+	t->type == task_type_timestep || t->type == task_type_init_grav ||
+	t->type == task_type_grav_down || t->type == task_type_grav_long_range) {
 
       /* Particle updates add only to vertex weight. */
       if (taskvweights) weights_v[cid] += w;

src/scheduler.c

@@ -1143,20 +1143,32 @@ void scheduler_reweight(struct scheduler *s, int verbose) {
       case task_type_ghost:
         if (t->ci == t->ci->super) cost = wscale * t->ci->count;
         break;
+      case task_type_extra_ghost:
+        if (t->ci == t->ci->super) cost = wscale * t->ci->count;
+        break;
       case task_type_drift_part:
         cost = wscale * t->ci->count;
         break;
       case task_type_drift_gpart:
         cost = wscale * t->ci->gcount;
         break;
+      case task_type_init_grav:
+	cost = wscale * t->ci->gcount;
+	break;
+      case task_type_grav_down:
+	cost = wscale * t->ci->gcount;
+	break;
+      case task_type_grav_long_range:
+	cost = wscale * t->ci->gcount;
+	break;
       case task_type_kick1:
-        cost = wscale * t->ci->count;
+        cost = wscale * t->ci->count + wscale * t->ci->gcount;
         break;
       case task_type_kick2:
-        cost = wscale * t->ci->count;
+        cost = wscale * t->ci->count + wscale * t->ci->gcount;
         break;
       case task_type_timestep:
-        cost = wscale * t->ci->count;
+        cost = wscale * t->ci->count + wscale * t->ci->gcount;
         break;
       case task_type_send:
         cost = 10 * wscale * t->ci->count * t->ci->count;