smyalygames/FiniteVolumeGPU
1
0
Fork 0
mirror of https://github.com/smyalygames/FiniteVolumeGPU.git synced 2026-01-13 23:28:41 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

feat(mpi): add parallel netCDF

Browse Source
This commit is contained in:
Anthony Berg
2025-07-24 13:42:03 +02:00
parent 320daff96d
commit d99ec9420b
10 changed files with 172 additions and 123 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
EulerTesting.ipynb
View File

@@ -101,22 +101,22 @@
"\n",
" #Create a netcdf file and simulate\n",
" with DataDumper(outfile, mode='w', clobber=False) as outdata:\n",
" outdata.ncfile.createDimension('time', None)\n",
" outdata.ncfile.createDimension('x', sim.nx)\n",
" outdata.ncfile.createDimension('y', sim.ny)\n",
" outdata.nc.createDimension('time', None)\n",
" outdata.nc.createDimension('x', sim.nx)\n",
" outdata.nc.createDimension('y', sim.ny)\n",
"\n",
" #Create variables\n",
" outdata.time_var = outdata.ncfile.createVariable('time', np.dtype('float32').char, 'time')\n",
" outdata.x_var = outdata.ncfile.createVariable('x', np.dtype('float32').char, 'x')\n",
" outdata.y_var = outdata.ncfile.createVariable('y', np.dtype('float32').char, 'y')\n",
" outdata.rho_var = outdata.ncfile.createVariable('rho', np.dtype('float32').char, ('time', 'y', 'x'), zlib=True,\n",
" least_significant_digit=3)\n",
" outdata.rho_u_var = outdata.ncfile.createVariable('rho_u', np.dtype('float32').char, ('time', 'y', 'x'),\n",
" zlib=True, least_significant_digit=3)\n",
" outdata.rho_v_var = outdata.ncfile.createVariable('rho_v', np.dtype('float32').char, ('time', 'y', 'x'),\n",
" zlib=True, least_significant_digit=3)\n",
" outdata.E_var = outdata.ncfile.createVariable('E', np.dtype('float32').char, ('time', 'y', 'x'), zlib=True,\n",
" least_significant_digit=3)\n",
" outdata.time_var = outdata.nc.createVariable('time', np.dtype('float32').char, 'time')\n",
" outdata.x_var = outdata.nc.createVariable('x', np.dtype('float32').char, 'x')\n",
" outdata.y_var = outdata.nc.createVariable('y', np.dtype('float32').char, 'y')\n",
" outdata.rho_var = outdata.nc.createVariable('rho', np.dtype('float32').char, ('time', 'y', 'x'), zlib=True,\n",
" least_significant_digit=3)\n",
" outdata.rho_u_var = outdata.nc.createVariable('rho_u', np.dtype('float32').char, ('time', 'y', 'x'),\n",
" zlib=True, least_significant_digit=3)\n",
" outdata.rho_v_var = outdata.nc.createVariable('rho_v', np.dtype('float32').char, ('time', 'y', 'x'),\n",
" zlib=True, least_significant_digit=3)\n",
" outdata.E_var = outdata.nc.createVariable('E', np.dtype('float32').char, ('time', 'y', 'x'), zlib=True,\n",
" least_significant_digit=3)\n",
"\n",
" #Create attributes\n",
" def to_json(in_dict):\n",
@@ -133,8 +133,8 @@
"\n",
" return json.dumps(out_dict)\n",
"\n",
" outdata.ncfile.created = time.ctime(time.time())\n",
" outdata.ncfile.sim_args = to_json(sim_args)\n",
" outdata.nc.created = time.ctime(time.time())\n",
" outdata.nc.sim_args = to_json(sim_args)\n",
"\n",
" outdata.x_var[:] = np.linspace(0, sim.nx * sim.dx, sim.nx)\n",
" outdata.y_var[:] = np.linspace(0, sim.ny * sim.dy, sim.ny)\n",
@@ -185,15 +185,15 @@
" fig = plt.figure(**fig_args)\n",
"\n",
" with DataDumper(infile, 'r') as indata:\n",
" time = indata.ncfile.variables['time'][:]\n",
" x = indata.ncfile.variables['x'][:]\n",
" y = indata.ncfile.variables['y'][:]\n",
" rho = indata.ncfile.variables['rho'][0]\n",
" rho_u = indata.ncfile.variables['rho_u'][0]\n",
" rho_v = indata.ncfile.variables['rho_v'][0]\n",
" time = indata.nc.variables['time'][:]\n",
" x = indata.nc.variables['x'][:]\n",
" y = indata.nc.variables['y'][:]\n",
" rho = indata.nc.variables['rho'][0]\n",
" rho_u = indata.nc.variables['rho_u'][0]\n",
" rho_v = indata.nc.variables['rho_v'][0]\n",
"\n",
" created = indata.ncfile.created\n",
" sim_args = json.loads(indata.ncfile.sim_args)\n",
" created = indata.nc.created\n",
" sim_args = json.loads(indata.nc.sim_args)\n",
" for key in sim_args:\n",
" if isinstance(sim_args[key], list):\n",
" sim_args[key] = \"[...]\"\n",
@@ -232,9 +232,9 @@
" progress_printer = ProgressPrinter(num_frames, print_every=5)\n",
"\n",
" def animate(i):\n",
" rho = indata.ncfile.variables['rho'][i]\n",
" rho_u = indata.ncfile.variables['rho_u'][i]\n",
" rho_v = indata.ncfile.variables['rho_v'][i]\n",
" rho = indata.nc.variables['rho'][i]\n",
" rho_u = indata.nc.variables['rho_u'][i]\n",
" rho_v = indata.nc.variables['rho_v'][i]\n",
"\n",
" if vis_type == VisType.Schlieren:\n",
" im.set_data(Visualization.gen_colors(rho, rho_u, rho_v, cmap, vmax, vmin))\n",

4
GPUSimulators/SHMEMSimulator.py
View File

@@ -103,10 +103,10 @@ class SHMEMSimulator(BaseSimulator):
if (gj == 0 and boundary_conditions.south != BoundaryCondition.Type.Periodic):
self.south = None
new_boundary_conditions.south = boundary_conditions.south;
if (gi == grid.grid[0]-1 and boundary_conditions.east != BoundaryCondition.Type.Periodic):
if (gi == grid.x-1 and boundary_conditions.east != BoundaryCondition.Type.Periodic):
self.east = None
new_boundary_conditions.east = boundary_conditions.east;
if (gj == grid.grid[1]-1 and boundary_conditions.north != BoundaryCondition.Type.Periodic):
if (gj == grid.y-1 and boundary_conditions.north != BoundaryCondition.Type.Periodic):
self.north = None
new_boundary_conditions.north = boundary_conditions.north;
sim.set_boundary_conditions(new_boundary_conditions)

4
GPUSimulators/SHMEMSimulatorGroup.py
View File

@@ -237,10 +237,10 @@ class SHMEMSimulatorGroup(object):
if (gj == 0 and boundary_conditions.south != BoundaryCondition.Type.Periodic):
self.south = None
new_boundary_conditions.south = boundary_conditions.south;
if (gi == grid.grid[0]-1 and boundary_conditions.east != BoundaryCondition.Type.Periodic):
if (gi == grid.x-1 and boundary_conditions.east != BoundaryCondition.Type.Periodic):
self.east = None
new_boundary_conditions.east = boundary_conditions.east;
if (gj == grid.grid[1]-1 and boundary_conditions.north != BoundaryCondition.Type.Periodic):
if (gj == grid.y-1 and boundary_conditions.north != BoundaryCondition.Type.Periodic):
self.north = None
new_boundary_conditions.north = boundary_conditions.north;
sim.set_boundary_conditions(new_boundary_conditions)

102
GPUSimulators/common/common.py
View File

@@ -28,7 +28,9 @@ import json
import numpy as np
from tqdm.auto import tqdm
from mpi4py import MPI
import GPUSimulators.mpi
from GPUSimulators.common.data_dumper import DataDumper
from GPUSimulators.common.timer import Timer
@@ -36,12 +38,12 @@ from GPUSimulators.common.timer import Timer
def safe_call(cmd):
logger = logging.getLogger(__name__)
try:
#git rev-parse HEAD
# git rev-parse HEAD
current_dir = os.path.dirname(os.path.realpath(__file__))
params = dict()
params['stderr'] = subprocess.STDOUT
params['cwd'] = current_dir
params['universal_newlines'] = True #text=True in more recent python
params['universal_newlines'] = True # text=True in more recent python
params['shell'] = False
if os.name == 'nt':
params['creationflags'] = subprocess.CREATE_NEW_PROCESS_GROUP
@@ -89,7 +91,7 @@ def run_simulation(simulator, simulator_args, outfile, save_times, save_var_name
save_var_names can be set to None if you do not want to save them
"""
profiling_data_sim_runner = { 'start': {}, 'end': {} }
profiling_data_sim_runner = {'start': {}, 'end': {}}
profiling_data_sim_runner["start"]["t_sim_init"] = 0
profiling_data_sim_runner["end"]["t_sim_init"] = 0
profiling_data_sim_runner["start"]["t_nc_write"] = 0
@@ -108,52 +110,71 @@ def run_simulation(simulator, simulator_args, outfile, save_times, save_var_name
sim = simulator(**simulator_args)
logger.info(f"Constructed in {str(t.secs)} seconds")
#Create a netcdf file and simulate
with DataDumper(outfile, mode='w', clobber=False) as outdata:
#Create attributes (metadata)
outdata.ncfile.created = time.ctime(time.time())
outdata.ncfile.git_hash = get_git_hash()
outdata.ncfile.git_status = get_git_status()
outdata.ncfile.simulator = str(simulator)
# Create a netcdf file and simulate
with DataDumper(outfile, mode='w', parallel=True, comm=MPI.COMM_WORLD, info=MPI.Info()) as outdata:
logger.info("Created NetCDF4 file.")
# Create attributes (metadata)
outdata.nc.created = time.ctime(time.time())
outdata.nc.git_hash = get_git_hash()
outdata.nc.git_status = get_git_status()
outdata.nc.simulator = str(simulator)
# do not write fields to attributes (they are to large)
simulator_args_for_ncfile = simulator_args.copy()
del simulator_args_for_ncfile["rho"]
del simulator_args_for_ncfile["rho_u"]
del simulator_args_for_ncfile["rho_v"]
del simulator_args_for_ncfile["E"]
outdata.ncfile.sim_args = to_json(simulator_args_for_ncfile)
#Create dimensions
outdata.ncfile.createDimension('time', len(save_times))
outdata.ncfile.createDimension('x', simulator_args['nx'])
outdata.ncfile.createDimension('y', simulator_args['ny'])
outdata.nc.sim_args = to_json(simulator_args_for_ncfile)
#Create variables for dimensions
ncvars = {'time': outdata.ncfile.createVariable('time', np.dtype('float32').char, 'time'),
'x': outdata.ncfile.createVariable('x', np.dtype('float32').char, 'x'),
'y': outdata.ncfile.createVariable('y', np.dtype('float32').char, 'y')}
# Create dimensions
if isinstance(sim, GPUSimulators.mpi.MPISimulator):
x_size = sim.grid.x
y_size = sim.grid.y
logger.info(f"Grid is - x: {x_size}, y: {y_size}")
else:
x_size = 1
y_size = 1
#Fill variables with proper values
grid_x0 = sim.grid.x0
grid_x1 = sim.grid.x1
grid_y0 = sim.grid.y0
grid_y1 = sim.grid.y1
x = simulator_args['nx'] * x_size
y = simulator_args['ny'] * y_size
outdata.nc.createDimension('time', len(save_times))
outdata.nc.createDimension('x', x)
outdata.nc.createDimension('y', y)
# Create variables for dimensions
ncvars = {'time': outdata.nc.createVariable('time', 'f4', ('time',)),
'x': outdata.nc.createVariable('x', 'f4', ('x',)),
'y': outdata.nc.createVariable('y', 'f4', ('y',))}
# Fill variables with proper values
ncvars['time'][:] = save_times
extent = sim.get_extent()
ncvars['x'][:] = np.linspace(extent[0], extent[1], simulator_args['nx'])
ncvars['y'][:] = np.linspace(extent[2], extent[3], simulator_args['ny'])
#Choose which variables to download (prune None from the list, but keep the index)
ncvars['time'].units = "s"
x0, x1, y0, y1 = sim.get_extent()
ncvars['x'][grid_x0:grid_x1] = np.linspace(x0, x1, simulator_args['nx'])
ncvars['y'][grid_y0:grid_y1] = np.linspace(y0, y1, simulator_args['ny'])
# Choose which variables to download (prune None from the list, but keep the index)
download_vars = []
for i, var_name in enumerate(save_var_names):
if var_name is not None:
download_vars += [i]
save_var_names = list(save_var_names[i] for i in download_vars)
#Create variables
for var_name in save_var_names:
ncvars[var_name] = outdata.ncfile.createVariable(
var_name, np.dtype('float32').char, ('time', 'y', 'x'), zlib=True, least_significant_digit=3)
#Create step sizes between each save
# Create variables
for var_name in save_var_names:
ncvars[var_name] = outdata.nc.createVariable(
var_name, 'f4', ('time', 'y', 'x'), zlib=True, least_significant_digit=3)
ncvars[var_name].set_collective(True)
# Create step sizes between each save
t_steps = np.empty_like(save_times)
t_steps[0] = save_times[0]
t_steps[1:] = save_times[1:] - save_times[0:-1]
@@ -162,8 +183,8 @@ def run_simulation(simulator, simulator_args, outfile, save_times, save_var_name
with tqdm(total=save_times[-1], desc="Simulation progress", unit="sim s") as pbar:
# Start simulation loop
for k, t_step in enumerate(t_steps):
t_end = k
for save_step, t_step in enumerate(t_steps):
t_end = save_step
# Sanity check simulator
try:
@@ -182,17 +203,16 @@ def run_simulation(simulator, simulator_args, outfile, save_times, save_var_name
profiling_data_sim_runner["start"]["t_nc_write"] += time.time()
#Download
# Download
save_vars = sim.download(download_vars)
#Save to file
# Save to file
for i, var_name in enumerate(save_var_names):
ncvars[var_name][k, :] = save_vars[i]
ncvars[var_name][save_step, grid_y0:grid_y1] = save_vars[i]
profiling_data_sim_runner["end"]["t_nc_write"] += time.time()
logger.debug(f"Simulated to t={t_end} in "
+ f"{sim.sim_steps()} timesteps (average dt={sim.sim_time() / sim.sim_steps()})")
return outdata.filename, profiling_data_sim_runner, sim.profiling_data_mpi

18
GPUSimulators/common/data_dumper.py
View File

@@ -2,8 +2,8 @@ import json
import logging
import os
import netCDF4
import numpy as np
from netCDF4 import Dataset
def to_json(in_dict):
@@ -48,13 +48,6 @@ class DataDumper(object):
mode = kwargs['mode']
# Create a new unique file if writing
if mode:
if ("w" in mode) or ("+" in mode) or ("a" in mode):
i = 0
stem, ext = os.path.splitext(filename)
while os.path.isfile(filename):
filename = f"{stem}_{str(i).zfill(4)}{ext}"
i = i + 1
self.filename = os.path.abspath(filename)
# Save arguments
@@ -70,10 +63,15 @@ class DataDumper(object):
self.logger.info("Arguments: " + str(self.args))
if self.kwargs:
self.logger.info("Keyword arguments: " + str(self.kwargs))
self.ncfile = netCDF4.Dataset(self.filename, *self.args, **self.kwargs)
self.nc = Dataset(
self.filename,
*self.args,
**self.kwargs
)
return self
def __exit__(self, *args):
self.logger.info("Closing " + self.filename)
self.ncfile.close()
self.nc.close()

4
GPUSimulators/helpers/initial_conditions.py
View File

@@ -28,8 +28,8 @@ from GPUSimulators.simulator.boundary import BoundaryCondition
def get_extent(width, height, nx, ny, grid, index=None):
if grid is not None:
gx = grid.grid[0]
gy = grid.grid[1]
gx = grid.x
gy = grid.y
if index is not None:
i, j = grid.get_coordinate(index)
else:

36
GPUSimulators/mpi/grid.py
View File

@@ -1,4 +1,5 @@
import logging
from enum import IntEnum
import numpy as np
from mpi4py import MPI
@@ -28,7 +29,7 @@ def get_grid(num_nodes, num_dims):
# Spent all factors: return number itself
if left == 1:
return (n, [n])
return n, [n]
# Find a new factor
i = 2
@@ -79,7 +80,7 @@ class MPIGrid(object):
neighboring nodes
"""
def __init__(self, comm, ndims=2):
def __init__(self, comm, nx, ny, ndims=2):
self.logger = logging.getLogger(__name__)
if ndims != 2:
@@ -87,40 +88,53 @@ class MPIGrid(object):
if comm.size < 1:
raise ValueError("Must have at least one node")
self.grid = get_grid(comm.size, ndims)
grid = get_grid(comm.size, ndims)
self.x = grid[0]
self.y = grid[1]
self.x0 = nx * (self.x-1)
self.x1 = self.x0 + nx
self.y0 = ny * (self.y-1)
self.y1 = self.y0 + ny
self.comm = comm
self.logger.debug(
f"Created MPI grid: {self.grid}. Rank {self.comm.rank} has coordinate {self.get_coordinate()}")
f"Created MPI grid: {grid}. Rank {self.comm.rank} has coordinate {self.get_coordinate()}")
def get_coordinate(self, rank=None):
if rank is None:
rank = self.comm.rank
i = (rank % self.grid[0])
j = (rank // self.grid[0])
i = (rank % self.x)
j = (rank // self.x)
return i, j
def get_grid_coordinate(self, rank=None):
"""
Gets the coordinate of the top left position of the grid in relation
to the entire grid.
"""
def get_rank(self, i, j):
return j * self.grid[0] + i
return j * self.x + i
def get_east(self):
i, j = self.get_coordinate(self.comm.rank)
i = (i + 1) % self.grid[0]
i = (i + 1) % self.x
return self.get_rank(i, j)
def get_west(self):
i, j = self.get_coordinate(self.comm.rank)
i = (i + self.grid[0] - 1) % self.grid[0]
i = (i + self.x - 1) % self.x
return self.get_rank(i, j)
def get_north(self):
i, j = self.get_coordinate(self.comm.rank)
j = (j + 1) % self.grid[1]
j = (j + 1) % self.y
return self.get_rank(i, j)
def get_south(self):
i, j = self.get_coordinate(self.comm.rank)
j = (j + self.grid[1] - 1) % self.grid[1]
j = (j + self.y - 1) % self.y
return self.get_rank(i, j)
def gather(self, data, root=0):

11
GPUSimulators/mpi/simulator.py
View File

@@ -86,10 +86,10 @@ class BaseMPISimulator(BaseSimulator):
if gj == 0 and boundary_conditions.south != BoundaryCondition.Type.Periodic:
self.south = None
new_boundary_conditions.south = boundary_conditions.south
if gi == grid.grid[0] - 1 and boundary_conditions.east != BoundaryCondition.Type.Periodic:
if gi == grid.x - 1 and boundary_conditions.east != BoundaryCondition.Type.Periodic:
self.east = None
new_boundary_conditions.east = boundary_conditions.east
if gj == grid.grid[1] - 1 and boundary_conditions.north != BoundaryCondition.Type.Periodic:
if gj == grid.y - 1 and boundary_conditions.north != BoundaryCondition.Type.Periodic:
self.north = None
new_boundary_conditions.north = boundary_conditions.north
sim.set_boundary_conditions(new_boundary_conditions)
@@ -188,12 +188,11 @@ class BaseMPISimulator(BaseSimulator):
width = self.sim.nx * self.sim.dx
height = self.sim.ny * self.sim.dy
i, j = self.grid.get_coordinate()
x0 = i * width
y0 = j * height
x0 = self.grid.x * width
y0 = self.grid.y * height
x1 = x0 + width
y1 = y0 + height
return [x0, x1, y0, y1]
return x0, x1, y0, y1
def full_exchange(self):
####

24
MPITest.ipynb
View File

@@ -217,15 +217,15 @@
"%%px\n",
"\n",
"with DataDumper(outfile, 'r') as indata:\n",
" t = indata.ncfile.variables['time'][:]\n",
" x = indata.ncfile.variables['x'][:]\n",
" y = indata.ncfile.variables['y'][:]\n",
" t = indata.nc.variables['time'][:]\n",
" x = indata.nc.variables['x'][:]\n",
" y = indata.nc.variables['y'][:]\n",
"\n",
" x_data = grid.gather(x)\n",
" y_data = grid.gather(y)\n",
"\n",
" created = indata.ncfile.created\n",
" sim_args = json.loads(indata.ncfile.sim_args)\n",
" created = indata.nc.created\n",
" sim_args = json.loads(indata.nc.sim_args)\n",
" for key in sim_args:\n",
" if isinstance(sim_args[key], list):\n",
" sim_args[key] = \"[...]\"\n",
@@ -234,7 +234,7 @@
" print(\"Simulator arguments: \\n\", sim_args)\n",
"\n",
" for i in range(num_steps):\n",
" rho = indata.ncfile.variables['rho'][i]\n",
" rho = indata.nc.variables['rho'][i]\n",
" rho_data = grid.gather(rho)\n",
"\n",
" #Plot on rank 0\n",
@@ -311,15 +311,15 @@
"%%px\n",
"\n",
"with DataDumper(outfile, 'r') as indata:\n",
" t = indata.ncfile.variables['time'][:]\n",
" x = indata.ncfile.variables['x'][:]\n",
" y = indata.ncfile.variables['y'][:]\n",
" t = indata.nc.variables['time'][:]\n",
" x = indata.nc.variables['x'][:]\n",
" y = indata.nc.variables['y'][:]\n",
"\n",
" x_data = grid.gather(x)\n",
" y_data = grid.gather(y)\n",
"\n",
" created = indata.ncfile.created\n",
" sim_args = json.loads(indata.ncfile.sim_args)\n",
" created = indata.nc.created\n",
" sim_args = json.loads(indata.nc.sim_args)\n",
" for key in sim_args:\n",
" if isinstance(sim_args[key], list):\n",
" sim_args[key] = \"[...]\"\n",
@@ -328,7 +328,7 @@
" print(\"Simulator arguments: \\n\", sim_args)\n",
"\n",
" for i in range(num_steps):\n",
" h = indata.ncfile.variables['h'][i]\n",
" h = indata.nc.variables['h'][i]\n",
" #hu = indata.ncfile.variables['hu'][i]\n",
" #hv = indata.ncfile.variables['hv'][i]\n",
" h_data = grid.gather(h)\n",

38
mpi_testing_hip.py
View File

@@ -31,15 +31,20 @@ from mpi4py import MPI
from hip import hip
from GPUSimulators.mpi import MPISimulator, MPIGrid
from GPUSimulators.common import run_simulation, get_git_hash, get_git_status, hip_check
from GPUSimulators.common import run_simulation, get_git_hash, get_git_status, hip_check, utils
from GPUSimulators.gpu import KernelContext
from GPUSimulators.model import EE2DKP07Dimsplit
from GPUSimulators.helpers import initial_conditions as IC
# Purely for local debugging
# import pydevd_pycharm
# pydevd_pycharm.settrace('localhost', port=24785, stdoutToServer=True, stderrToServer=True)
parser = argparse.ArgumentParser(description='Strong and weak scaling experiments.')
parser.add_argument('-nx', type=int, default=128)
parser.add_argument('-ny', type=int, default=128)
parser.add_argument('--profile', action='store_true') # default: False
parser.add_argument('--compile_opts', type=str, help="Compiler options for HIP code.")
args = parser.parse_args()
@@ -49,15 +54,19 @@ if args.profile:
t_total_start = time.time()
t_init_start = time.time()
nx = args.nx
ny = args.ny
# Get MPI COMM to use
comm = MPI.COMM_WORLD
rank = comm.rank
####
# Initialize logging
####
log_level_console = 20
log_level_file = 10
log_filename = 'mpi_' + str(comm.rank) + '.log'
log_filename = 'mpi_' + str(rank) + '.log'
logger = logging.getLogger('GPUSimulators')
logger.setLevel(min(log_level_console, log_level_file))
@@ -78,7 +87,7 @@ logger.info(f"File logger using level {logging.getLevelName(log_level_file)} to
# Initialize MPI grid etc
####
logger.info("Creating MPI grid")
grid = MPIGrid(MPI.COMM_WORLD)
grid = MPIGrid(comm, nx, ny)
####
# Initialize HIP
@@ -98,8 +107,6 @@ context = KernelContext(device=hip_device, autotuning=False)
np.random.seed(42)
logger.info("Generating initial conditions")
nx = args.nx
ny = args.ny
dt = 0.001
@@ -107,14 +114,20 @@ gamma = 1.4
# save_times = np.linspace(0, 0.000009, 2)
# save_times = np.linspace(0, 0.000099, 11)
# save_times = np.linspace(0, 0.000099, 2)
save_times = np.linspace(0, 20, 21)
outfile = "mpi_out_" + str(MPI.COMM_WORLD.rank) + ".nc"
save_times = np.linspace(0, 5, 21)
outfile = "mpi_out.nc4"
save_var_names = ['rho', 'rho_u', 'rho_v', 'E']
arguments = IC.gen_kelvin_helmholtz(nx, ny, gamma, grid=grid)
arguments['context'] = context
arguments['theta'] = 1.2
arguments['grid'] = grid
arguments['compile_opts'] = ['-g', '-g3', '-ggdb', '-gdwarf-4', '-O0']
compile_opts = args.compile_opts
if compile_opts is not None:
arguments['compile_opts'] += compile_opts
if args.profile:
t_init_end = time.time()
@@ -126,7 +139,6 @@ if args.profile:
####
logger.info("Running simulation")
# Helper function to create MPI simulator
@@ -139,14 +151,20 @@ def gen_sim(grid, **kwargs):
outfile, sim_runner_profiling_data, sim_profiling_data = run_simulation(
gen_sim, arguments, outfile, save_times, save_var_names, dt)
# Move NetCDF4 file to a unique file, for the next run.
if rank == 0:
new_filename = utils.unique_file(outfile)
os.rename(outfile, new_filename)
##### Profiling ######
if args.profile:
t_total_end = time.time()
t_total = t_total_end - t_total_start
profiling_data["t_total"] = t_total
print(f"Total run time on rank {str(MPI.COMM_WORLD.rank)} is {str(t_total)} s")
print(f"Total run time on rank {str(rank)} is {str(t_total)} s")
# write profiling to JSON file
if args.profile and MPI.COMM_WORLD.rank == 0:
if args.profile and rank == 0:
job_id = ""
if "SLURM_JOB_ID" in os.environ:
job_id = int(os.environ["SLURM_JOB_ID"])