Merge pull request #17 from equinor/vectorized-mesh-updates

Vectorized mesh updates
equinor · Jan 17, 2024 · 6fbf345 · 6fbf345
2 parents 44291a0 + cac5fbd
commit 6fbf345
Show file tree

Hide file tree

Showing 3 changed files with 343 additions and 412 deletions.
diff --git a/tests/warmth3d/test_3d.py b/tests/warmth3d/test_3d.py
@@ -5,10 +5,15 @@
 import numpy as np
 from mpi4py import MPI
 import pickle
+import time
 
 def test_3d_compare():
     comm = MPI.COMM_WORLD
-    if comm.rank == 0:
+    inc = 1000
+    model_pickled = f"model-out-inc_{inc}.p"
+    if comm.rank == 0 and not os.path.isfile(model_pickled):
+        global runtime_1D_sim
+        st = time.time()
         maps_dir = Path("./docs/notebooks/data/")
         model = warmth.Model()
 
@@ -22,8 +27,6 @@ def test_3d_compare():
         inputs.loc[6]=[182,"182.gri",None,"Erosive"]
         model.builder.input_horizons=inputs
 
-
-        inc = 1000
         model.builder.define_geometry(maps_dir/"0.gri",xinc=inc,yinc=inc,fformat="irap_binary")
 
         model.builder.extract_nodes(4,maps_dir)
@@ -50,12 +53,13 @@ def test_3d_compare():
         model.parameters.tetha = 0
         model.parameters.alphav = 0
 
-        model.simulator.run(save=True,purge=True, parallel=False)
+        model.simulator.run(save=True, purge=True, parallel=True, use_mpi=(comm.size>1))
 
-        import pickle
-        pickle.dump( model, open( "model-out.p", "wb" ) )
-        model = pickle.load( open( "model-out.p", "rb" ) )
+        runtime_1D_sim = time.time() - st
+        print("Total time 1D simulations:", runtime_1D_sim)
 
+        pickle.dump( model, open( model_pickled, "wb" ) )
+        model = pickle.load( open( model_pickled, "rb" ) )
         try:
             os.mkdir('mesh')
         except FileExistsError:
@@ -66,41 +70,43 @@ def test_3d_compare():
             pass
 
     comm.Barrier()
-    import pickle
-    model = pickle.load( open( "model-out.p", "rb" ) )
+    model = pickle.load( open( model_pickled, "rb" ) )
     comm.Barrier()
-    # if comm.rank == 0:            
-    # mm2 = 
-    mm2, posarr, Tarr = run_3d(model.builder,model.parameters,start_time=model.parameters.time_start,end_time=0, pad_num_nodes=2,writeout=False, base_flux=None)
-
+    mm2 = run_3d(model.builder,model.parameters,start_time=model.parameters.time_start,end_time=0, pad_num_nodes=2,writeout=False, base_flux=None)
 
-    nnx = (model.builder.grid.num_nodes_x+2*mm2.padX) 
-    nny = (model.builder.grid.num_nodes_y+2*mm2.padX) 
-    # hx = model.builder.grid.num_nodes_x // 2
-    # hy = model.builder.grid.num_nodes_y // 2
-    hx = nnx // 2
-    hy = nny // 2
+    comm.Barrier()
+    if comm.rank == 0:            
+        nnx = (model.builder.grid.num_nodes_x+2*mm2.padX) 
+        nny = (model.builder.grid.num_nodes_y+2*mm2.padX) 
+        hx = nnx // 2
+        hy = nny // 2
+
+        nn = model.builder.nodes[hy-mm2.padX][hx-mm2.padX]
+        dd = nn._depth_out[:,0]
+
+        mm2_pos, mm2_temp = mm2.get_node_pos_and_temp(-1)
 
-    nn = model.builder.nodes[hy-mm2.padX][hx-mm2.padX]
-    dd = nn._depth_out[:,0]
+        node_ind = hy*nnx + hx
+        v_per_n = int(mm2_pos.shape[0]/(mm2.num_nodes))
 
-    node_ind = hy*nnx + hx
-    # v_per_n = int(mm2.mesh_vertices.shape[0]/(model.builder.grid.num_nodes_y*model.builder.grid.num_nodes_x))
-    v_per_n = int(mm2.mesh_vertices.shape[0]/(mm2.num_nodes))
+        temp_1d = np.nan_to_num(nn.temperature_out[:,0], nan=5.0)
+
+        # temp_3d_ind = np.array([ np.where([mm2.mesh_reindex==i])[1][0] for i in range(node_ind*v_per_n, (node_ind+1)*v_per_n) ] )
+        temp_3d_ind = np.array([ i for i in range(node_ind*v_per_n, (node_ind+1)*v_per_n) ] )
 
-    temp_1d = np.nan_to_num(nn.temperature_out[:,0], nan=5.0)
-    temp_3d_ind = np.array([ np.where([mm2.mesh_reindex==i])[1][0] for i in range(node_ind*v_per_n, (node_ind+1)*v_per_n) ] )
-    dd_mesh = mm2.mesh.geometry.x[temp_3d_ind,2] - mm2.sed_diff_z[range(node_ind*v_per_n, (node_ind+1)*v_per_n)]
-    temp_3d_mesh = mm2.u_n.x.array[temp_3d_ind]
+        # dd_mesh = mm2.mesh.geometry.x[temp_3d_ind,2] - mm2.sed_diff_z[range(node_ind*v_per_n, (node_ind+1)*v_per_n)]
+        dd_mesh = mm2_pos[temp_3d_ind,2] - mm2.sed_diff_z[range(node_ind*v_per_n, (node_ind+1)*v_per_n)]
+        # temp_3d_mesh = mm2.u_n.x.array[temp_3d_ind]
+        temp_3d_mesh = mm2_temp[temp_3d_ind]
 
-    temp_1d_at_mesh_pos = np.interp(dd_mesh, dd, temp_1d)
-    dd_subset = np.where(dd_mesh<5000)
+        temp_1d_at_mesh_pos = np.interp(dd_mesh, dd, temp_1d)
+        dd_subset = np.where(dd_mesh<5000)
 
-    max_abs_error = np.amax(np.abs(temp_1d_at_mesh_pos-temp_3d_mesh))
-    max_abs_error_shallow = np.amax(np.abs(temp_1d_at_mesh_pos[dd_subset]-temp_3d_mesh[dd_subset]))    
+        max_abs_error = np.amax(np.abs(temp_1d_at_mesh_pos-temp_3d_mesh))
+        max_abs_error_shallow = np.amax(np.abs(temp_1d_at_mesh_pos[dd_subset]-temp_3d_mesh[dd_subset]))    
 
-    print("Max error overall:", max_abs_error)
-    print("Max error <5000m:", max_abs_error_shallow)
+        print("Max error overall:", max_abs_error)
+        print("Max error <5000m:", max_abs_error_shallow)
 
     # assert (max_abs_error<25.0), "Temperature difference between 3D and 1D simulations is >25"
     # assert (max_abs_error_shallow<5.0),  "Temperature difference between 3D and 1D simulations is >5 in the sediments."
@@ -109,4 +115,6 @@ def test_3d_compare():
 
 
 if __name__ == "__main__":
-    test_3d_compare()
+    test_3d_compare()
+    if 'runtime_1D_sim' in globals():
+        print("Total time 1D simulations:", runtime_1D_sim)