diff --git a/src/seismic_hazard_forecasting.py b/src/seismic_hazard_forecasting.py
index c3b0426..1b7880e 100644
--- a/src/seismic_hazard_forecasting.py
+++ b/src/seismic_hazard_forecasting.py
@@ -803,6 +803,19 @@ verbose: {verbose}")
         iml_grid_prep = iml_grid.copy()  # temp ground motion grids
 
         if use_AOI:
+
+            for j in range(0, len(products)):
+                # Reassemble the grid cleanly using the original shape
+                # Initialize a flat array filled entirely with NaNs
+                iml_grid_flat = np.full(num_points, np.nan, dtype=np.float64)
+                
+                # Assign the computed values to their exact original 1D index positions
+                iml_grid_flat[indices_filtered] = iml_grid_raw[j]
+                
+                # Reshape back using the exact shape of your original xx/yy grids
+                iml_grid_prep[j] = iml_grid_flat.reshape(grid_shape)
+
+            
             for i in indices:
                 if i in indices_filtered:
                     for j in range(0, len(products)):