From 497d6422dbf2bbdfca587e14702ee12329510b2a Mon Sep 17 00:00:00 2001
From: olender <olender@usp.br>
Date: Fri, 26 Jul 2024 14:03:34 -0300
Subject: [PATCH] added parallel load and save test

---
 spyro/examples/immersed_polygon.py |  2 +-
 test_polygon_fwi.py                | 60 ++++++++++++++++++++++++++++++
 2 files changed, 61 insertions(+), 1 deletion(-)
 create mode 100644 test_polygon_fwi.py

diff --git a/spyro/examples/immersed_polygon.py b/spyro/examples/immersed_polygon.py
index 8c6e9257..1796bca2 100644
--- a/spyro/examples/immersed_polygon.py
+++ b/spyro/examples/immersed_polygon.py
@@ -144,7 +144,7 @@ def _polygon_velocity_model(self):
         cond = fire.conditional(z <= -0.5 - 0.2*x, vl, cond)
 
         cond = fire.conditional(300*((x-0.5)*(-z-0.5))**2 + ((x-0.5)+(-z-0.5))**2 <= 0.300**2, v2+dv, cond)
-        
+
         if self.abc_pad_length > 0.0:
             middle_of_pad = -self.length_z - self.abc_pad_length*0.5
             cond = fire.conditional(z <= middle_of_pad, v0, cond)
diff --git a/test_polygon_fwi.py b/test_polygon_fwi.py
new file mode 100644
index 00000000..1ebbbddc
--- /dev/null
+++ b/test_polygon_fwi.py
@@ -0,0 +1,60 @@
+# from mpi4py.MPI import COMM_WORLD
+# import debugpy
+# debugpy.listen(3000 + COMM_WORLD.rank)
+# debugpy.wait_for_client()
+import spyro
+import numpy as np
+
+
+def test_real_shot_record_generation_parallel():
+    dictionary = {}
+    dictionary["absorving_boundary_conditions"] = {
+        "pad_length": 2.0,  # True or false
+    }
+    dictionary["mesh"] = {
+        "h": 0.05,  # mesh size in km
+    }
+    dictionary["polygon_options"] = {
+        "water_layer_is_present": True,
+        "upper_layer": 2.0,
+        "middle_layer": 2.5,
+        "lower_layer": 3.0,
+        "polygon_layer_perturbation": 0.3,
+    }
+    dictionary["acquisition"] = {
+        "source_locations": spyro.create_transect((-0.1, 0.1), (-0.1, 0.9), 4),
+        "frequency": 5.0,
+        "receiver_locations": spyro.create_transect((-0.16, 0.1), (-0.16, 0.9), 100),
+    }
+    dictionary["visualization"] = {
+        "debug_output": True,
+    }
+    dictionary["time_axis"] = {
+        "final_time": 1.0,  # Final time for event
+        "dt": 0.0005,  # timestep size
+        "amplitude": 1,  # the Ricker has an amplitude of 1.
+        "output_frequency": 500,  # how frequently to output solution to pvds
+        # how frequently to save solution to RAM
+        "gradient_sampling_frequency": 1,
+    }
+    fwi = spyro.examples.Polygon_acoustic_FWI(dictionary=dictionary, periodic=True)
+    fwi.generate_real_shot_record(plot_model=True, save_shot_record=True)
+
+    dictionary["inversion"] = {
+        "real_shot_record_files": "shots/shot_record_",
+    }
+    fwi2 = spyro.examples.Polygon_acoustic_FWI(dictionary=dictionary, periodic=True)
+
+    max_value = np.max(fwi2.real_shot_record[:, fwi2.comm.ensemble_comm.rank*33])
+    test_core = np.isclose(max_value, 0.184, atol=1e-2)
+
+    test_core_all = fwi2.comm.ensemble_comm.allgather(test_core)
+    test = all(test_core_all)
+
+    print(f"Correctly loaded shots: {test}", flush=True)
+
+    assert test
+
+
+if __name__ == "__main__":
+    test_real_shot_record_generation_parallel()