diff --git a/exercise4/wave_2d_parallel.c b/exercise4/wave_2d_parallel.c
index 4b196ac..b09c464 100644
--- a/exercise4/wave_2d_parallel.c
+++ b/exercise4/wave_2d_parallel.c
@@ -31,9 +31,9 @@ real_t
 // TASK: T1b
 // Declare variables each MPI process will need
 // BEGIN: T1b
-#define U_prv(i, j) buffers[0][((i) + 1) * (N + 2) + (j) + 1]
-#define U(i, j) buffers[1][((i) + 1) * (N + 2) + (j) + 1]
-#define U_nxt(i, j) buffers[2][((i) + 1) * (N + 2) + (j) + 1]
+#define U_prv(i, j) buffers[0][((i) + 1) * (local_N + 2) + (j) + 1]
+#define U(i, j) buffers[1][((i) + 1) * (local_N + 2) + (j) + 1]
+#define U_nxt(i, j) buffers[2][((i) + 1) * (local_N + 2) + (j) + 1]
 
 int world_size;
 int world_rank;
@@ -137,9 +137,9 @@ void boundary_condition(void) {
 // and set the time step.
 void domain_initialize(void) {
   // BEGIN: T4
-  buffers[0] = malloc((M + 2) * (N + 2) * sizeof(real_t));
-  buffers[1] = malloc((M + 2) * (N + 2) * sizeof(real_t));
-  buffers[2] = malloc((M + 2) * (N + 2) * sizeof(real_t));
+  buffers[0] = malloc((local_M + 2) * (local_N + 2) * sizeof(real_t));
+  buffers[1] = malloc((local_M + 2) * (local_N + 2) * sizeof(real_t));
+  buffers[2] = malloc((local_M + 2) * (local_N + 2) * sizeof(real_t));
 
   if (!buffers[0] || !buffers[1] || !buffers[2]) {
     perror("malloc");
@@ -147,10 +147,12 @@ void domain_initialize(void) {
   }
 
   // initialize interior (physical cells)
-  for (int_t i = 0; i < M; i++) {
-    for (int_t j = 0; j < N; j++) {
-      real_t dx_i = (i - M / 2.0);
-      real_t dy_j = (j - N / 2.0);
+  int rows_offset = coordinates[0] * local_M;
+  int cols_offset = coordinates[1] * local_N;
+  for (int_t i = 0; i < local_M; i++) {
+    for (int_t j = 0; j < local_N; j++) {
+      real_t dx_i = (i + rows_offset - M / 2.0);
+      real_t dy_j = (j + cols_offset - N / 2.0);
       real_t delta = sqrt((dx_i * dx_i) / (real_t)M + (dy_j * dy_j) / (real_t)N);
       U_prv(i, j) = U(i, j) = exp(-4.0 * delta * delta);
     }