def test_distances(self):
self.set_positions_small_set()
longDiag = np.sqrt(2 * 0.8 * 0.8)
longLine = 0.8
shortDiag = np.sqrt(0.2 * 0.2 + 0.2 * 0.2)
shortLine = 0.2
semiDiag = np.sqrt(0.2 * 0.2 + 0.8 * 0.8)
# smallDrifterSet is initially with periodic boundary conditions
assertListAlmostEqual(self, self.smallDrifterSet.getDistances().tolist(), \
[shortDiag, shortLine, shortLine], 6,
'distance with periodic boundaries')
self.smallDrifterSet.setBoundaryConditions(
Common.BoundaryConditions(1, 1, 1, 1))
assertListAlmostEqual(self, self.smallDrifterSet.getDistances().tolist(), \
[longDiag, longLine, longLine], 6,
'distances with non-periodic boundaries')
self.smallDrifterSet.setBoundaryConditions(
Common.BoundaryConditions(1, 2, 1, 2))
assertListAlmostEqual(self, self.smallDrifterSet.getDistances().tolist(), \
[semiDiag, shortLine, longLine], 6,
'distances with periodic boundaries in east-west')
self.smallDrifterSet.setBoundaryConditions(
Common.BoundaryConditions(2, 1, 2, 1))
assertListAlmostEqual(self, self.smallDrifterSet.getDistances().tolist(), \
[semiDiag, longLine, shortLine], 6,
'distances with periodic boundaries in north-south')
def test_innovations(self):
self.set_positions_small_set()
zero = 0.0
close = 0.2
far = -0.8
fasit = [[close, close], [close, zero], [zero, close]]
# smallDrifterSet is initially with periodic boundary conditions
assert2DListAlmostEqual(self, self.smallDrifterSet.getInnovations().tolist(), \
fasit, 6,
'innovations with periodic boundaries')
fasit = [[far, far], [far, zero], [zero, far]]
self.smallDrifterSet.setBoundaryConditions(
Common.BoundaryConditions(1, 1, 1, 1))
assert2DListAlmostEqual(self, self.smallDrifterSet.getInnovations().tolist(), \
fasit, 6,
'innovations with non-periodic boundaries')
fasit = [[close, far], [close, zero], [zero, far]]
self.smallDrifterSet.setBoundaryConditions(
Common.BoundaryConditions(1, 2, 1, 2))
assert2DListAlmostEqual(self, self.smallDrifterSet.getInnovations().tolist(), \
fasit, 6,
'innovations with periodic boundaries in east-west')
fasit = [[far, close], [far, zero], [zero, close]]
self.smallDrifterSet.setBoundaryConditions(
Common.BoundaryConditions(2, 1, 2, 1))
assert2DListAlmostEqual(self, self.smallDrifterSet.getInnovations().tolist(), \
fasit, 6,
'innovations with periodic boundaries in north-south')
def setBoundaryConditions(self, bcSettings=1):
if (bcSettings == 1):
self.boundaryConditions = Common.BoundaryConditions()
elif (bcSettings == 2):
self.boundaryConditions = Common.BoundaryConditions(2, 2, 2, 2)
elif bcSettings == 3:
self.boundaryConditions = Common.BoundaryConditions(2, 1, 2, 1)
else:
self.boundaryConditions = Common.BoundaryConditions(1, 2, 1, 2)
def fromfilename(cls, gpu_ctx, filename, cont_write_netcdf=True):
"""
Initialize and hotstart simulation from nc-file.
cont_write_netcdf: Continue to write the results after each superstep to a new netCDF file
filename: Continue simulation based on parameters and last timestep in this file
"""
# open nc-file
sim_reader = SimReader.SimNetCDFReader(filename,
ignore_ghostcells=False)
sim_name = str(sim_reader.get('simulator_short'))
assert sim_name == cls.__name__, \
"Trying to initialize a " + \
cls.__name__ + " simulator with netCDF file based on " \
+ sim_name + " results."
# read parameters
nx = sim_reader.get("nx")
ny = sim_reader.get("ny")
dx = sim_reader.get("dx")
dy = sim_reader.get("dy")
width = nx * dx
height = ny * dy
dt = sim_reader.get("dt")
g = sim_reader.get("g")
r = sim_reader.get("bottom_friction_r")
f = sim_reader.get("coriolis_force")
beta = sim_reader.get("coriolis_beta")
minmodTheta = sim_reader.get("minmod_theta")
timeIntegrator = sim_reader.get("time_integrator")
y_zero_reference_cell = sim_reader.get("y_zero_reference_cell")
try:
wind_stress_type = sim_reader.get("wind_stress_type")
wind = Common.WindStressParams(type=wind_stress_type)
except:
wind = WindStress.WindStress()
boundaryConditions = Common.BoundaryConditions( \
sim_reader.getBC()[0], sim_reader.getBC()[1], \
sim_reader.getBC()[2], sim_reader.getBC()[3], \
sim_reader.getBCSpongeCells())
h0 = sim_reader.getH()
# get last timestep (including simulation time of last timestep)
eta0, hu0, hv0, time0 = sim_reader.getLastTimeStep()
return cls(gpu_ctx, \
h0, eta0, hu0, hv0, \
nx, ny, \
dx, dy, dt, \
g, f, r, \
t=time0, \
wind_stress=wind, \
boundary_conditions=boundaryConditions, \
write_netcdf=cont_write_netcdf)
def setUp(self):
self.nx = 1
self.ny = 1
self.dx = 1.0
self.dy = 1.0
self.dt = 1.0
self.numParticles = 3
self.observationVariance = 0.5
self.boundaryCondition = Common.BoundaryConditions(2, 2, 2, 2)
self.smallParticleSet = None
# to be initialized by child class with above values
# create_small_particle_set:
self.cl_ctx = None
self.smallPositionSetHost = np.array([[0.9, 0.9], [0.9, 0.1],
[0.1, 0.9], [0.1, 0.1]])
self.resampleNumParticles = 6
self.resamplingParticleArray = np.zeros((7, 2))
self.resamplingObservationVariance = 0.1
for i in range(2):
self.resamplingParticleArray[3 * i + 0, :] = [0.25, 0.35 + i * 0.3]
self.resamplingParticleArray[3 * i + 1, :] = [0.4, 0.35 + i * 0.3]
self.resamplingParticleArray[3 * i + 2, :] = [0.65, 0.35 + i * 0.3]
self.resamplingParticleArray[6, :] = [0.25, 0.5]
self.resamplingParticleSet = None
# to be initialized by child class wit resampleNumParticles only.
self.resamplingVar = 1e-8
def setUp(self):
self.gpu_ctx = None
self.numDrifters = 3
self.observationVariance = 0.25
self.boundaryCondition = Common.BoundaryConditions(2, 2, 2, 2)
self.smallDrifterSet = None
# to be initialized by child class with above values
self.smallPositionSetHost = np.array([[0.9, 0.9], [0.9, 0.1],
[0.1, 0.9], [0.1, 0.1]])
self.resampleNumDrifters = 6
self.resamplingDriftersArray = np.zeros((7, 2))
for i in range(2):
self.resamplingDriftersArray[3 * i + 0, :] = [0.25, 0.35 + i * 0.3]
self.resamplingDriftersArray[3 * i + 1, :] = [0.4, 0.35 + i * 0.3]
self.resamplingDriftersArray[3 * i + 2, :] = [0.65, 0.35 + i * 0.3]
self.resamplingDriftersArray[6, :] = [0.25, 0.5]
self.resamplingDrifterSet = None
# to be initialized by child class wit resampleNumDrifters only.
self.resamplingVar = 1e-8
self.largeDrifterSet = None
def __init__(self,
numDrifters,
observation_variance=0.01,
boundaryConditions=Common.BoundaryConditions(),
initialization_cov_drifters=None,
domain_size_x=1.0,
domain_size_y=1.0):
"""
Creates a GlobalParticles object for drift trajectory ensemble.
numDrifters: number of drifters in the collection, not included the observation
observation_variance: uncertainty of observation position
boundaryConditions: BoundaryConditions object, relevant during re-initialization of particles.
"""
# Call parent constructor
super(CPUDrifterCollection,
self).__init__(numDrifters,
observation_variance=observation_variance,
boundaryConditions=boundaryConditions,
domain_size_x=domain_size_x,
domain_size_y=domain_size_y)
# One position for every particle plus observation
self.positions = np.zeros((self.numDrifters + 1, 2))
# Initialize drifters:
self.uniformly_distribute_drifters(
initialization_cov_drifters=initialization_cov_drifters)
def __init__(self, numDrifters, observation_variance=0.1, \
boundaryConditions=Common.BoundaryConditions(), \
domain_size_x=1.0, domain_size_y=1.0):
"""
Creates a GlobalParticles object for drift trajectory ensemble.
numDrifters: number of drifers in the collection, not included the observation
observation_variance: uncertainty of observation position
boundaryConditions: BoundaryConditions object, relevant during re-initialization of particles.
"""
self.numDrifters = numDrifters
# Observation index is the last particle
self.obs_index = self.numDrifters
self.observation_variance = observation_variance
self.positions = None # Needs to be allocated in the child class
# Should represent all particles plus observation
self.domain_size_x = domain_size_x
self.domain_size_y = domain_size_y
# Boundary conditions are read from a BoundaryConditions object
self.boundaryConditions = boundaryConditions
def setBoundaryConditions(self, bcSettings=1):
if (bcSettings == 1):
self.boundaryConditions = Common.BoundaryConditions()
#self.ghosts = [0,0,0,0] # north, east, south, west
self.arrayRange = [None, None, 0, 0]
elif (bcSettings == 2):
self.boundaryConditions = Common.BoundaryConditions(2, 2, 2, 2)
#self.ghosts = [1,1,0,0] # Both periodic
self.arrayRange = [-1, -1, 0, 0]
elif bcSettings == 3:
self.boundaryConditions = Common.BoundaryConditions(2, 1, 2, 1)
#self.ghosts = [1,0,0,0] # periodic north-south
self.arrayRange = [-1, None, 0, 0]
else:
self.boundaryConditions = Common.BoundaryConditions(1, 2, 1, 2)
#self.ghosts = [0,1,0,0] # periodic east-west
self.arrayRange = [None, -1, 0, 0]
def setGridInfo(self,
nx,
ny,
dx,
dy,
dt,
boundaryConditions=Common.BoundaryConditions(),
eta=None,
hu=None,
hv=None,
H=None):
self.nx = nx
self.ny = ny
self.dx = dx
self.dy = dy
self.dt = dt
# Default values for now:
self.initialization_variance = 10 * dx
self.midPoint = 0.5 * np.array([self.nx * self.dx, self.ny * self.dy])
self.initialization_cov = np.eye(2) * self.initialization_variance
self.boundaryConditions = boundaryConditions
assert (self.simType == 'CDKLM16'
), 'CDKLM16 is currently the only supported scheme'
#if self.simType == 'CDKLM16':
self.ghostCells = np.array([2, 2, 2, 2])
if self.boundaryConditions.isSponge():
sponge = self.boundaryConditions.getSponge()
for i in range(4):
if sponge[i] > 0:
self.ghostCells[i] = sponge[i]
dataShape = (ny + self.ghostCells[0] + self.ghostCells[2],
nx + self.ghostCells[1] + self.ghostCells[3])
self.base_eta = eta
self.base_hu = hu
self.base_hv = hv
self.base_H = H
# Create base initial data:
if self.base_eta is None:
self.base_eta = np.zeros(dataShape, dtype=np.float32, order='C')
if self.base_hu is None:
self.base_hu = np.zeros(dataShape, dtype=np.float32, order='C')
if self.base_hv is None:
self.base_hv = np.zeros(dataShape, dtype=np.float32, order='C')
# Bathymetry:
if self.base_H is None:
waterDepth = 10
self.base_H = np.ones((dataShape[0] + 1, dataShape[1] + 1),
dtype=np.float32,
order='C') * waterDepth
self.setParameters()
def test_collection_mean(self):
self.set_positions_small_set()
periodicMean = [1-0.1/3, 1-0.1/3]
nonPeriodicMean = [(0.9 + 0.9 + 0.1)/3, (0.9 + 0.9 + 0.1)/3]
semiPeriodicMean = [nonPeriodicMean[0], periodicMean[1]]
assertListAlmostEqual(self, self.smallDrifterSet.getCollectionMean().tolist(),
periodicMean, 6,
'periodic mean')
self.smallDrifterSet.setBoundaryConditions(Common.BoundaryConditions(1,1,1,1))
assertListAlmostEqual(self, self.smallDrifterSet.getCollectionMean().tolist(),
nonPeriodicMean, 6,
'non-periodic mean')
self.smallDrifterSet.setBoundaryConditions(Common.BoundaryConditions(2,1,2,1))
assertListAlmostEqual(self, self.smallDrifterSet.getCollectionMean().tolist(),
semiPeriodicMean, 6,
'north-south-periodic mean')
def _setGridInfo(self,
nx,
ny,
dx,
dy,
dt,
boundaryConditions=Common.BoundaryConditions(),
eta=None,
hu=None,
hv=None,
H=None):
"""
Declaring grid-related member variables
"""
self.nx = nx
self.ny = ny
self.dx = dx
self.dy = dy
self.dt = dt
self.boundaryConditions = boundaryConditions
assert (self.simType == 'CDKLM16'
), 'CDKLM16 is currently the only supported scheme'
self.ghostCells = np.array([2, 2, 2, 2])
if self.boundaryConditions.isSponge():
sponge = self.boundaryConditions.getSponge()
for i in range(4):
if sponge[i] > 0:
self.ghostCells[i] = sponge[i]
dataShape = (ny + self.ghostCells[0] + self.ghostCells[2],
nx + self.ghostCells[1] + self.ghostCells[3])
self.base_eta = eta
self.base_hu = hu
self.base_hv = hv
self.base_H = H
# Create base initial data:
if self.base_eta is None:
self.base_eta = np.zeros(dataShape, dtype=np.float32, order='C')
if self.base_hu is None:
self.base_hu = np.zeros(dataShape, dtype=np.float32, order='C')
if self.base_hv is None:
self.base_hv = np.zeros(dataShape, dtype=np.float32, order='C')
# Bathymetry:
if self.base_H is None:
waterDepth = 10
self.base_H = np.ones((dataShape[0] + 1, dataShape[1] + 1),
dtype=np.float32,
order='C') * waterDepth
def create_noise(self):
n, e, s, w = 1, 1, 1, 1
if self.periodicNS:
n, s = 2, 2
if self.periodicEW:
e, w = 2, 2
self.noise = OceanStateNoise(self.gpu_ctx,
self.gpu_stream,
self.nx,
self.ny,
self.dx,
self.dy,
Common.BoundaryConditions(n, e, s, w),
staggered=self.staggered)
def __init__(self, gpu_ctx, numDrifters, \
observation_variance=0.01, \
boundaryConditions=Common.BoundaryConditions(), \
initialization_cov_drifters=None, \
domain_size_x=1.0, domain_size_y=1.0, \
gpu_stream=None, \
block_width = 64):
super(GPUDrifterCollection, self).__init__(numDrifters,
observation_variance=observation_variance,
boundaryConditions=boundaryConditions,
domain_size_x=domain_size_x,
domain_size_y=domain_size_y)
# Define CUDA environment:
self.gpu_ctx = gpu_ctx
self.block_width = block_width
self.block_height = 1
# TODO: Where should the cl_queue come from?
# For sure, the drifter and the ocean simulator should use
# the same queue...
self.gpu_stream = gpu_stream
if self.gpu_stream is None:
self.gpu_stream = cuda.Stream()
self.sensitivity = 1.0
self.driftersHost = np.zeros((self.getNumDrifters() + 1, 2)).astype(np.float32, order='C')
self.driftersDevice = Common.CUDAArray2D(self.gpu_stream, \
2, self.getNumDrifters()+1, 0, 0, \
self.driftersHost)
self.drift_kernels = gpu_ctx.get_kernel("driftKernels.cu", \
defines={'block_width': self.block_width, 'block_height': self.block_height})
# Get CUDA functions and define data types for prepared_{async_}call()
self.passiveDrifterKernel = self.drift_kernels.get_function("passiveDrifterKernel")
self.passiveDrifterKernel.prepare("iifffiiPiPiPifiiiPif")
self.enforceBoundaryConditionsKernel = self.drift_kernels.get_function("enforceBoundaryConditions")
self.enforceBoundaryConditionsKernel.prepare("ffiiiPi")
self.local_size = (self.block_width, self.block_height, 1)
self.global_size = (\
int(np.ceil((self.getNumDrifters() + 2)/float(self.block_width))), \
1)
# Initialize drifters:
self.uniformly_distribute_drifters(initialization_cov_drifters=initialization_cov_drifters)
def create_noise(self, factor=1):
n, e, s, w = 1, 1, 1, 1
if self.periodicNS:
n, s = 2, 2
if self.periodicEW:
e, w = 2, 2
self.noise = OceanStateNoise(self.gpu_ctx,
self.gpu_stream,
self.nx,
self.ny,
self.dx,
self.dy,
Common.BoundaryConditions(n, e, s, w),
staggered=self.staggered,
interpolation_factor=factor,
use_lcg=self.useLCG())
def __init__(self, \
gpu_ctx, \
H, eta0, hu0, hv0, \
nx, ny, \
dx, dy, dt, \
g, f, r, \
t=0.0, \
coriolis_beta=0.0, \
y_zero_reference_cell = 1, \
wind_stress=WindStress.WindStress(), \
boundary_conditions=Common.BoundaryConditions(), \
write_netcdf=False, \
comm=None, \
ignore_ghostcells=False, \
offset_x=0, offset_y=0, \
block_width=16, block_height=16):
"""
Initialization routine
H: Water depth incl ghost cells, (nx+2)*(ny+2) cells
eta0: Initial deviation from mean sea level incl ghost cells, (nx+2)*(ny+2) cells
hu0: Initial momentum along x-axis incl ghost cells, (nx+1)*(ny+2) cells
hv0: Initial momentum along y-axis incl ghost cells, (nx+2)*(ny+3) cells
nx: Number of cells along x-axis
ny: Number of cells along y-axis
dx: Grid cell spacing along x-axis (20 000 m)
dy: Grid cell spacing along y-axis (20 000 m)
dt: Size of each timestep (90 s)
g: Gravitational accelleration (9.81 m/s^2)
f: Coriolis parameter (1.2e-4 s^1), effectively as f = f + beta*y
r: Bottom friction coefficient (2.4e-3 m/s)
coriolis_beta: Coriolis linear factor -> f = f + beta*y
y_zero_reference_cell: The cell representing y_0 in the above, defined as the lower face of the cell .
wind_stress: Wind stress parameters
boundary_conditions: Boundary condition object
write_netcdf: Write the results after each superstep to a netCDF file
comm: MPI communicator
"""
#### THIS ALLOWS MAKES IT POSSIBLE TO GIVE THE OLD INPUT SHAPES TO NEW GHOST CELL REGIME: Only valid for benchmarking!
if (eta0.shape == (ny, nx)):
new_eta = np.zeros((ny+2, nx+2), dtype=np.float32)
new_eta[:ny, :nx] = eta0.copy()
eta0 = new_eta.copy()
if (H.shape == (ny, nx)):
new_H = np.ones((ny+2, nx+2), dtype=np.float32)*np.max(H)
new_H[:ny,:nx] = H.copy()
H = new_H.copy()
if (hu0.shape == (ny, nx+1)):
new_hu = np.zeros((ny+2, nx+1), dtype=np.float32)
new_hu[:ny, :nx+1] = hu0.copy()
hu0 = new_hu.copy()
if (hv0.shape == (ny+1, nx)):
new_hv = np.zeros((ny+3, nx+2), dtype=np.float32)
new_hv[:ny+1,:nx] = hv0.copy()
hv0 = new_hv.copy()
#Create data by uploading to device
ghost_cells_x = 1
ghost_cells_y = 1
y_zero_reference_cell = y_zero_reference_cell
# Index range for interior domain (north, east, south, west)
# so that interior domain of eta is
# eta[self.interior_domain_indices[2]:self.interior_domain_indices[0], \
# self.interior_domain_indices[3]:self.interior_domain_indices[1] ]
self.interior_domain_indices = np.array([-1, -1, 1, 1])
self.boundary_conditions = boundary_conditions
if boundary_conditions.isSponge():
nx = nx - 2 + boundary_conditions.spongeCells[1] + boundary_conditions.spongeCells[3]
ny = ny - 2 + boundary_conditions.spongeCells[0] + boundary_conditions.spongeCells[2]
y_zero_reference_cell = y_zero_reference_cell + boundary_conditions.spongeCells[2]
rk_order = None
theta = None
A = None
super(FBL, self).__init__(gpu_ctx, \
nx, ny, \
ghost_cells_x, \
ghost_cells_y, \
dx, dy, dt, \
g, f, r, A, \
t, \
theta, rk_order, \
coriolis_beta, \
y_zero_reference_cell, \
wind_stress, \
write_netcdf, \
ignore_ghostcells, \
offset_x, offset_y, \
comm, \
block_width, block_height)
self._set_interior_domain_from_sponge_cells()
#Get kernels
self.step_kernel = gpu_ctx.get_kernel("FBL_step_kernel.cu",
defines={'block_width': block_width, 'block_height': block_height},
compile_args={
'no_extern_c': True,
'options': ["--use_fast_math"],
#'options': ["--generate-line-info"],
#'options': ["--maxrregcount=32"]
#'arch': "compute_50",
#'code': "sm_50"
},
jit_compile_args={
#jit_options=[(cuda.jit_option.MAX_REGISTERS, 39)]
}
)
# Get CUDA functions
self.fblStepKernel = self.step_kernel.get_function("fblStepKernel")
# Prepare kernel lauches
self.fblStepKernel.prepare("iiffffffffPiPiPiPiif")
# Set up textures
self.update_wind_stress(self.step_kernel, self.fblStepKernel)
self.H = Common.CUDAArray2D(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, H)
self.gpu_data = Common.SWEDataArakawaC(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, eta0, hu0, hv0, fbl=True)
# Domain including ghost cells
self.nx_halo = np.int32(nx + 2)
self.ny_halo = np.int32(ny + 2)
self.bc_kernel = FBL_boundary_conditions(self.gpu_ctx, \
self.nx, \
self.ny, \
self.boundary_conditions
)
# Bit-wise boolean for wall boundary conditions
self.wall_bc = np.int32(0)
if (self.boundary_conditions.north == 1):
self.wall_bc = self.wall_bc | 0x01
if (self.boundary_conditions.east == 1):
self.wall_bc = self.wall_bc | 0x02
if (self.boundary_conditions.south == 1):
self.wall_bc = self.wall_bc | 0x04
if (self.boundary_conditions.west == 1):
self.wall_bc = self.wall_bc | 0x08
if self.write_netcdf:
self.sim_writer = SimWriter.SimNetCDFWriter(self, ignore_ghostcells=self.ignore_ghostcells, \
staggered_grid=True, \
offset_x=self.offset_x, offset_y=self.offset_y)
def __init__(self, \
gpu_ctx, \
eta0, hu0, hv0, Hi, \
nx, ny, \
dx, dy, dt, \
g, f, r, \
t=0.0, \
theta=1.3, rk_order=2, \
coriolis_beta=0.0, \
y_zero_reference_cell = 0, \
max_wind_direction_perturbation = 0, \
wind_stress=WindStress.WindStress(), \
boundary_conditions=Common.BoundaryConditions(), \
small_scale_perturbation=False, \
small_scale_perturbation_amplitude=None, \
h0AsWaterElevation=False, \
reportGeostrophicEquilibrium=False, \
write_netcdf=False, \
ignore_ghostcells=False, \
offset_x=0, offset_y=0, \
block_width=32, block_height=4):
"""
Initialization routine
eta0: Initial deviation from mean sea level incl ghost cells, (nx+2)*(ny+2) cells
hu0: Initial momentum along x-axis incl ghost cells, (nx+1)*(ny+2) cells
hv0: Initial momentum along y-axis incl ghost cells, (nx+2)*(ny+1) cells
Hi: Depth from equilibrium defined on cell corners, (nx+5)*(ny+5) corners
nx: Number of cells along x-axis
ny: Number of cells along y-axis
dx: Grid cell spacing along x-axis (20 000 m)
dy: Grid cell spacing along y-axis (20 000 m)
dt: Size of each timestep (90 s)
g: Gravitational accelleration (9.81 m/s^2)
f: Coriolis parameter (1.2e-4 s^1), effectively as f = f + beta*y
r: Bottom friction coefficient (2.4e-3 m/s)
t: Start simulation at time t
theta: MINMOD theta used the reconstructions of the derivatives in the numerical scheme
rk_order: Order of Runge Kutta method {1,2*,3}
coriolis_beta: Coriolis linear factor -> f = f + beta*(y-y_0)
y_zero_reference_cell: The cell representing y_0 in the above, defined as the lower face of the cell .
max_wind_direction_perturbation: Large-scale model error emulation by per-time-step perturbation of wind direction by +/- max_wind_direction_perturbation (degrees)
wind_stress: Wind stress parameters
boundary_conditions: Boundary condition object
h0AsWaterElevation: True if h0 is described by the surface elevation, and false if h0 is described by water depth
reportGeostrophicEquilibrium: Calculate the Geostrophic Equilibrium variables for each superstep
write_netcdf: Write the results after each superstep to a netCDF file
"""
## After changing from (h, B) to (eta, H), several of the simulator settings used are wrong. This check will help detect that.
if ( np.sum(eta0 - Hi[:-1, :-1] > 0) > nx):
assert(False), "It seems you are using water depth/elevation h and bottom topography B, while you should use water level eta and equillibrium depth H."
assert( rk_order < 4 or rk_order > 0 ), "Only 1st, 2nd and 3rd order Runge Kutta supported"
if (rk_order == 3):
assert(r == 0.0), "3rd order Runge Kutta supported only without friction"
# Sort out internally represented ghost_cells in the presence of given
# boundary conditions
ghost_cells_x = 2
ghost_cells_y = 2
y_zero_reference_cell = 2 + y_zero_reference_cell
# Boundary conditions
self.boundary_conditions = boundary_conditions
if (boundary_conditions.isSponge()):
nx = nx + boundary_conditions.spongeCells[1] + boundary_conditions.spongeCells[3] - 2*ghost_cells_x
ny = ny + boundary_conditions.spongeCells[0] + boundary_conditions.spongeCells[2] - 2*ghost_cells_y
y_zero_reference_cell = boundary_conditions.spongeCells[2] + y_zero_reference_cell
A = None
self.max_wind_direction_perturbation = max_wind_direction_perturbation
super(CDKLM16, self).__init__(gpu_ctx, \
nx, ny, \
ghost_cells_x, \
ghost_cells_y, \
dx, dy, dt, \
g, f, r, A, \
t, \
theta, rk_order, \
coriolis_beta, \
y_zero_reference_cell, \
wind_stress, \
write_netcdf, \
ignore_ghostcells, \
offset_x, offset_y, \
block_width, block_height)
# Index range for interior domain (north, east, south, west)
# so that interior domain of eta is
# eta[self.interior_domain_indices[2]:self.interior_domain_indices[0], \
# self.interior_domain_indices[3]:self.interior_domain_indices[1] ]
self.interior_domain_indices = np.array([-2,-2,2,2])
self._set_interior_domain_from_sponge_cells()
#Get kernels
self.kernel = gpu_ctx.get_kernel("CDKLM16_kernel.cu", defines={'block_width': block_width, 'block_height': block_height})
# Get CUDA functions and define data types for prepared_{async_}call()
self.swe_2D = self.kernel.get_function("swe_2D")
self.swe_2D.prepare("iifffffffffiiPiPiPiPiPiPiPiPifiiiiiPiPiPi")
self.update_wind_stress(self.kernel, self.swe_2D)
#Create data by uploading to device
self.gpu_data = Common.SWEDataArakawaA(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, eta0, hu0, hv0)
## Allocating memory for geostrophical equilibrium variables
self.reportGeostrophicEquilibrium = np.int32(reportGeostrophicEquilibrium)
dummy_zero_array = np.zeros((ny+2*ghost_cells_y, nx+2*ghost_cells_x), dtype=np.float32, order='C')
self.geoEq_uxpvy = Common.CUDAArray2D(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, dummy_zero_array)
self.geoEq_Kx = Common.CUDAArray2D(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, dummy_zero_array)
self.geoEq_Ly = Common.CUDAArray2D(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, dummy_zero_array)
#Bathymetry
self.bathymetry = Common.Bathymetry(gpu_ctx, self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, Hi, boundary_conditions)
self.h0AsWaterElevation = h0AsWaterElevation
if self.h0AsWaterElevation:
self.bathymetry.waterElevationToDepth(self.gpu_data.h0)
self.constant_equilibrium_depth = np.max(Hi)
self.bc_kernel = Common.BoundaryConditionsArakawaA(gpu_ctx, \
self.nx, \
self.ny, \
ghost_cells_x, \
ghost_cells_y, \
self.boundary_conditions, \
)
# Small scale perturbation:
self.small_scale_perturbation = small_scale_perturbation
self.small_scale_model_error = None
if small_scale_perturbation:
if small_scale_perturbation_amplitude is None:
self.small_scale_model_error = OceanStateNoise.OceanStateNoise.fromsim(self)
else:
self.small_scale_model_error = OceanStateNoise.OceanStateNoise.fromsim(self, soar_q0=small_scale_perturbation_amplitude)
if self.write_netcdf:
self.sim_writer = SimWriter.SimNetCDFWriter(self, ignore_ghostcells=self.ignore_ghostcells, \
offset_x=self.offset_x, offset_y=self.offset_y)
def test_set_boundary_condition(self):
self.set_positions_small_set()
self.smallDrifterSet.setBoundaryConditions(
Common.BoundaryConditions(2, 1, 2, 1))
self.assertEqual(self.smallDrifterSet.getBoundaryConditions().get(),
[2, 1, 2, 1])
def __init__(self,
gpu_ctx,
perturbation_type=DoubleJetPerturbationType.SteadyState,
model_error=True,
commonSpinUpTime=200000):
"""
Class that generates initial conditions for a double jet case (both perturbed and unperturbed).
The use of initial perturbations/spin up periods are given by the perturbation_type argument,
which should be a DoubleJetPerturbationType instance.
"""
# The following parameters are the standard choices we have made for our double jet case.
# If any of them are to be altered, they should be made optional input parameters to the
# constructor, with the values below given as default parameters.
# Check that the provided perturbation type is valid
DoubleJetPerturbationType._assert_valid(perturbation_type)
self.perturbation_type = perturbation_type
# Domain-related parameters
self.phi_0 = 72 * np.pi / 180.0
self.phi_05 = 75 * np.pi / 180.0
self.phi_1 = 78 * np.pi / 180.0
self.midpoint_phi_pos = 73.5 * np.pi / 180
self.midpoint_phi_neg = 76.5 * np.pi / 180
self.phi_delta = 5.5 * np.pi / 180
self.phi_pos_min = self.midpoint_phi_pos - self.phi_delta
self.phi_pos_max = self.midpoint_phi_pos + self.phi_delta
self.phi_neg_min = self.midpoint_phi_neg - self.phi_delta
self.phi_neg_max = self.midpoint_phi_neg + self.phi_delta
self.e_n = np.exp(-4 / (self.phi_delta * 2)**2)
distance_between_latitudes = 111e3 # m
degrees_0 = self.phi_0 * 180 / np.pi
degrees_1 = self.phi_1 * 180 / np.pi
y_south = degrees_0 * distance_between_latitudes
y_north = degrees_1 * distance_between_latitudes
degrees_mid = self.phi_05 * 180 / np.pi
self.ny = 300
self.dy = (y_north - y_south) / self.ny
self.dx = self.dy
self.nx = 500
self.ghosts = np.array([2, 2, 2, 2]) # north, east, south, west
self.dataShape = (self.ny + self.ghosts[0] + self.ghosts[2],
self.nx + self.ghosts[1] + self.ghosts[3])
# Physical parameters
self.g = 9.80616 # m/s^2 - gravitational acceleration
omega = 7.2722e-5 # 1/s - Angular rotation speed of the earth
self.earth_radius = 6.37122e6 # m - radius of the Earth
self.u_max = 3 # m/s - Gulf stream has "maximum speed typically about 2.5 m/s"
self.h_0 = 230 # m - It was found to be 230.03, but with a dobious calculation.
# - Better then just to set the depth to a constant :)
self.commonSpinUpTime = commonSpinUpTime # s - Because it just seems like a good measure.
self.individualSpinUpTime = 100000 # s - Because it just seems like a good measure.
self.f = 2 * omega * np.sin(self.phi_05)
self.tan = np.tan(self.phi_05)
# Initial data
sim_h_init, redef_hu_init = self._initSteadyState()
sim_h_init_mean = sim_h_init.mean()
self.delta_eta = np.max(sim_h_init) - np.min(sim_h_init)
max_dt = 0.25 * self.dx / (np.max(redef_hu_init / sim_h_init +
np.sqrt(self.g * sim_h_init)))
dt = 0.8 * max_dt
self.base_cpu_Hi = np.ones(
(self.dataShape[0] + 1, self.dataShape[1] + 1),
dtype=np.float32) * sim_h_init_mean
self.base_cpu_eta = -np.ones(self.dataShape,
dtype=np.float32) * sim_h_init_mean
self.base_cpu_hu = np.zeros(self.dataShape, dtype=np.float32)
self.base_cpu_hv = np.zeros(self.dataShape, dtype=np.float32)
for i in range(self.dataShape[1]):
self.base_cpu_eta[:, i] += sim_h_init
self.base_cpu_hu[:, i] = redef_hu_init
self.sim_args = {
"gpu_ctx": gpu_ctx,
"nx": self.nx,
"ny": self.ny,
"dx": self.dy,
"dy": self.dy,
"dt": dt,
"g": self.g,
"f": self.f,
"coriolis_beta": 0.0,
"r": 0.0,
"H": self.base_cpu_Hi,
"t": 0.0,
"rk_order": 2,
"boundary_conditions": Common.BoundaryConditions(2, 2, 2, 2),
"small_scale_perturbation": model_error,
"small_scale_perturbation_amplitude": 0.0003,
"small_scale_perturbation_interpolation_factor": 5,
}
self.base_init = {
"eta0": self.base_cpu_eta,
"hu0": self.base_cpu_hu,
"hv0": self.base_cpu_hv
}
if self.perturbation_type == DoubleJetPerturbationType.SpinUp or \
self.perturbation_type == DoubleJetPerturbationType.LowFrequencySpinUp or \
self.perturbation_type == DoubleJetPerturbationType.LowFrequencyStandardSpinUp:
if self.perturbation_type == DoubleJetPerturbationType.LowFrequencySpinUp:
self.commonSpinUpTime = self.commonSpinUpTime
self.individualSpinUpTime = self.individualSpinUpTime * 1.5
elif self.perturbation_type == DoubleJetPerturbationType.LowFrequencyStandardSpinUp:
self.sim_args, self.base_init = self.getStandardPerturbedInitConditions(
)
self.commonSpinUpTime = self.commonSpinUpTime * 2
tmp_sim = CDKLM16.CDKLM16(**self.sim_args, **self.base_init)
tmp_t = tmp_sim.step(self.commonSpinUpTime)
tmp_eta, tmp_hu, tmp_hv = tmp_sim.download(
interior_domain_only=False)
self.base_init['eta0'] = tmp_eta
self.base_init['hu0'] = tmp_hu
self.base_init['hv0'] = tmp_hv
self.sim_args['t'] = tmp_sim.t
tmp_sim.cleanUp()
# The IEWPFPaperCase - isolated to give a better overview
if self.perturbation_type == DoubleJetPerturbationType.IEWPFPaperCase:
self.sim_args["small_scale_perturbation_amplitude"] = 0.00025
self.sim_args["model_time_step"] = 60 # sec
tmp_sim = CDKLM16.CDKLM16(**self.sim_args, **self.base_init)
tmp_sim.updateDt()
three_days = 3 * 24 * 60 * 60
tmp_t = tmp_sim.dataAssimilationStep(three_days)
tmp_eta, tmp_hu, tmp_hv = tmp_sim.download(
interior_domain_only=False)
self.base_init['eta0'] = tmp_eta
self.base_init['hu0'] = tmp_hu
self.base_init['hv0'] = tmp_hv
self.sim_args['t'] = tmp_sim.t
tmp_sim.cleanUp()
def getInitialConditions(source_url_list, x0, x1, y0, y1, \
timestep_indices=None, \
land_value=5.0, \
iterations=10, \
sponge_cells={'north':20, 'south': 20, 'east': 20, 'west': 20}, \
erode_land=0,
download_data=True):
ic = {}
if type(source_url_list) is not list:
source_url_list = [source_url_list]
num_files = len(source_url_list)
for i in range(len(source_url_list)):
source_url_list[i] = checkCachedNetCDF(source_url_list[i],
download_data=download_data)
# Read constants and initial values from the first source url
source_url = source_url_list[0]
try:
ncfile = Dataset(source_url)
H_m = ncfile.variables['h'][y0 - 1:y1 + 1, x0 - 1:x1 + 1]
eta0 = ncfile.variables['zeta'][0, y0 - 1:y1 + 1, x0 - 1:x1 + 1]
u0 = ncfile.variables['ubar'][0, y0:y1, x0:x1]
v0 = ncfile.variables['vbar'][0, y0:y1, x0:x1]
angle = ncfile.variables['angle'][y0:y1, x0:x1]
latitude = ncfile.variables['lat'][y0:y1, x0:x1]
x = ncfile.variables['X'][x0:x1]
y = ncfile.variables['Y'][y0:y1]
except Exception as e:
raise e
finally:
ncfile.close()
# Get time steps:
if timestep_indices is None:
timestep_indices = [None] * num_files
elif type(timestep_indices) is not list:
timestep_indices_tmp = [None] * num_files
for i in range(num_files):
timestep_indices_tmp[i] = timestep_indices
timestep_indices = timestep_indices_tmp
timesteps = [None] * num_files
for i in range(num_files):
try:
ncfile = Dataset(source_url_list[i])
if (timestep_indices[i] is not None):
timesteps[i] = ncfile.variables['time'][timestep_indices[i][:]]
else:
timesteps[i] = ncfile.variables['time'][:]
timestep_indices[i] = range(len(timesteps[i]))
except Exception as e:
print('exception in obtaining timestep for file ' + str(i))
raise e
finally:
ncfile.close()
#Generate timesteps in reference to t0
t0 = timesteps[0][0]
for ts in timesteps:
t0 = min(t0, min(ts))
assert (np.all(np.diff(timesteps) >= 0))
for i in range(num_files):
timesteps[i] = timesteps[i] - t0
#Generate intersections bathymetry
H_m_mask = eta0.mask.copy()
H_m = np.ma.array(H_m, mask=H_m_mask)
for i in range(erode_land):
new_water = H_m.mask ^ binary_erosion(H_m.mask)
eps = 1.0e-5 #Make new Hm slighlyt different from land_value
eta0_dil = grey_dilation(eta0.filled(0.0), size=(3, 3))
H_m[new_water] = land_value + eps
eta0[new_water] = eta0_dil[new_water]
H_i, _ = OceanographicUtilities.midpointsToIntersections(
H_m, land_value=land_value, iterations=iterations)
eta0 = eta0[1:-1, 1:-1]
h0 = OceanographicUtilities.intersectionsToMidpoints(H_i).filled(
land_value) + eta0.filled(0.0)
#Generate physical variables
eta0 = np.ma.array(eta0.filled(0), mask=eta0.mask.copy())
hu0 = np.ma.array(h0 * u0.filled(0), mask=eta0.mask.copy())
hv0 = np.ma.array(h0 * v0.filled(0), mask=eta0.mask.copy())
#Spong cells for e.g., flow relaxation boundary conditions
ic['sponge_cells'] = sponge_cells
#Number of cells
ic['NX'] = x1 - x0
ic['NY'] = y1 - y0
# Domain size without ghost cells
ic['nx'] = ic['NX'] - 4
ic['ny'] = ic['NY'] - 4
#Dx and dy
#FIXME: Assumes equal for all.. .should check
ic['dx'] = np.average(x[1:] - x[:-1])
ic['dy'] = np.average(y[1:] - y[:-1])
#Gravity and friction
#FIXME: Friction coeff from netcdf?
ic['g'] = 9.81
ic['r'] = 3.0e-3
#Physical variables
ic['H'] = H_i
ic['eta0'] = eta0
ic['hu0'] = hu0
ic['hv0'] = hv0
#Coriolis angle and beta
ic['angle'] = angle
ic['latitude'] = OceanographicUtilities.degToRad(latitude)
ic['f'] = 0.0 #Set using latitude instead
# The beta plane of doing it:
# ic['f'], ic['coriolis_beta'] = OceanographicUtilities.calcCoriolisParams(OceanographicUtilities.degToRad(latitude[0, 0]))
#Boundary conditions
ic['boundary_conditions_data'] = getBoundaryConditionsData(
source_url_list, timestep_indices, timesteps, x0, x1, y0, y1)
ic['boundary_conditions'] = Common.BoundaryConditions(
north=3, south=3, east=3, west=3, spongeCells=sponge_cells)
#Wind stress (shear stress acting on the ocean surface)
ic['wind_stress'] = getWindSourceterm(source_url_list, timestep_indices,
timesteps, x0, x1, y0, y1)
#Note
ic['note'] = datetime.datetime.now().isoformat(
) + ": Generated from " + str(source_url_list)
#Initial reference time and all timesteps
ic['t0'] = t0
ic['timesteps'] = np.ravel(timesteps)
#wind (wind speed in m/s used for forcing on drifter)
ic['wind'] = getWind(source_url_list, timestep_indices, timesteps, x0, x1,
y0, y1)
return ic
def setUpAndStartEnsemble(self):
self.nx = 40
self.ny = 40
self.dx = 4.0
self.dy = 4.0
self.dt = 0.05
self.g = 9.81
self.r = 0.0
self.f = 0.05
self.beta = 0.0
self.waterDepth = 10.0
self.ensembleSize = 3
self.driftersPerOceanModel = 3
ghosts = np.array([2, 2, 2, 2]) # north, east, south, west
validDomain = np.array([2, 2, 2, 2])
self.boundaryConditions = Common.BoundaryConditions(2, 2, 2, 2)
# Define which cell index which has lower left corner as position (0,0)
x_zero_ref = 2
y_zero_ref = 2
dataShape = (self.ny + ghosts[0] + ghosts[2],
self.nx + ghosts[1] + ghosts[3])
dataShapeHi = (self.ny + ghosts[0] + ghosts[2] + 1,
self.nx + ghosts[1] + ghosts[3] + 1)
eta0 = np.zeros(dataShape, dtype=np.float32, order='C')
eta0_extra = np.zeros(dataShape, dtype=np.float32, order='C')
hv0 = np.zeros(dataShape, dtype=np.float32, order='C')
hu0 = np.zeros(dataShape, dtype=np.float32, order='C')
Hi = np.ones(dataShapeHi, dtype=np.float32,
order='C') * self.waterDepth
# Add disturbance:
rel_grid_size = self.nx * 1.0 / self.dx
BC.addBump(eta0, self.nx, self.ny, self.dx, self.dy, 0.3, 0.5,
0.05 * rel_grid_size, validDomain)
eta0 = eta0 * 0.3
BC.addBump(eta0, self.nx, self.ny, self.dx, self.dy, 0.7, 0.3,
0.10 * rel_grid_size, validDomain)
eta0 = eta0 * (-1.3)
BC.addBump(eta0, self.nx, self.ny, self.dx, self.dy, 0.15, 0.8,
0.03 * rel_grid_size, validDomain)
eta0 = eta0 * 1.0
BC.addBump(eta0, self.nx, self.ny, self.dx, self.dy, 0.6, 0.75,
0.06 * rel_grid_size, validDomain)
BC.addBump(eta0, self.nx, self.ny, self.dx, self.dy, 0.2, 0.2,
0.01 * rel_grid_size, validDomain)
eta0 = eta0 * (-0.03)
BC.addBump(eta0_extra, self.nx, self.ny, self.dx, self.dy, 0.5, 0.5,
0.4 * rel_grid_size, validDomain)
eta0 = eta0 + 0.02 * eta0_extra
BC.initializeBalancedVelocityField(eta0, Hi, hu0, hv0, \
self.f, self.beta, self.g, \
self.nx, self.ny, self.dx ,self.dy, ghosts)
eta0 = eta0 * 0.5
self.q0 = 0.5 * self.dt * self.f / (self.g * self.waterDepth)
self.sim = CDKLM16.CDKLM16(self.gpu_ctx, eta0, hu0, hv0, Hi, \
self.nx, self.ny, self.dx, self.dy, self.dt, \
self.g, self.f, self.r, \
boundary_conditions=self.boundaryConditions, \
write_netcdf=False, \
small_scale_perturbation=True, \
small_scale_perturbation_amplitude=self.q0)
self.ensemble = OceanNoiseEnsemble.OceanNoiseEnsemble(
self.gpu_ctx,
self.ensembleSize,
self.sim,
num_drifters=self.driftersPerOceanModel,
observation_type=dautils.ObservationType.DirectUnderlyingFlow,
observation_variance=0.01**2)
self.iewpf = IEWPFOcean.IEWPFOcean(self.ensemble)
gpu_ctx = Common.CUDAContext()
toc = time.time()
print("{:02.4f} s: ".format(toc - tic) + "Created context on " +
gpu_ctx.cuda_device.name())
# Set benchmark sizes
dx = 200.0
dy = 200.0
dt = 0.95 / 100
g = 9.81
f = 0.00
r = 0.0
boundaryConditions = Common.BoundaryConditions()
# Generate initial conditions
waterHeight = 60
x_center = dx * args.nx / 2.0
y_center = dy * args.ny / 2.0
size = 0.4 * min(args.nx * dx, args.ny * dy)
def my_exp(i, j):
x = dx * i - x_center
y = dy * j - y_center
return np.exp(-10 * (x * x / (size * size) + y * y /
(size * size))) * (np.sqrt(x**2 + y**2) < size)
def __init__(self, \
gpu_ctx, \
H, eta0, hu0, hv0, \
nx, ny, \
dx, dy, dt, \
g, f, r, \
t=0.0, \
coriolis_beta=0.0, \
y_zero_reference_cell = 0, \
wind_stress=WindStress.WindStress(), \
boundary_conditions=Common.BoundaryConditions(), \
write_netcdf=False, \
ignore_ghostcells=False, \
offset_x=0, offset_y=0, \
block_width=16, block_height=16):
"""
Initialization routine
H: Water depth incl ghost cells, (nx+2)*(ny+2) cells
eta0: Initial deviation from mean sea level incl ghost cells, (nx+2)*(ny+2) cells
hu0: Initial momentum along x-axis incl ghost cells, (nx+1)*(ny+2) cells
hv0: Initial momentum along y-axis incl ghost cells, (nx+2)*(ny+1) cells
nx: Number of cells along x-axis
ny: Number of cells along y-axis
dx: Grid cell spacing along x-axis (20 000 m)
dy: Grid cell spacing along y-axis (20 000 m)
dt: Size of each timestep (90 s)
g: Gravitational accelleration (9.81 m/s^2)
f: Coriolis parameter (1.2e-4 s^1), effectively as f = f + beta*y
r: Bottom friction coefficient (2.4e-3 m/s)
coriolis_beta: Coriolis linear factor -> f = f + beta*y
y_zero_reference_cell: The cell representing y_0 in the above, defined as the lower face of the cell .
wind_stress: Wind stress parameters
boundary_conditions: Boundary condition object
write_netcdf: Write the results after each superstep to a netCDF file
"""
#Create data by uploading to device
ghost_cells_x = 0
ghost_cells_y = 0
y_zero_reference_cell = y_zero_reference_cell
self.asym_ghost_cells = [0, 0, 0, 0] # [N, E, S, W]
# Index range for interior domain (north, east, south, west)
# so that interior domain of eta is
# eta[self.interior_domain_indices[2]:self.interior_domain_indices[0], \
# self.interior_domain_indices[3]:self.interior_domain_indices[1] ]
self.interior_domain_indices = np.array([None, None, 0, 0])
self.boundary_conditions = boundary_conditions
# Add asym ghost cell if periodic boundary condition:
if (self.boundary_conditions.north == 2) or \
(self.boundary_conditions.south == 2):
self.asym_ghost_cells[0] = 1
self.interior_domain_indices[0] = -1
if (self.boundary_conditions.east == 2) or \
(self.boundary_conditions.west == 2):
self.asym_ghost_cells[1] = 1
self.interior_domain_indices[1] = -1
if boundary_conditions.isSponge():
nx = nx + boundary_conditions.spongeCells[
1] + boundary_conditions.spongeCells[
3] # - self.asym_ghost_cells[1] - self.asym_ghost_cells[3]
ny = ny + boundary_conditions.spongeCells[
0] + boundary_conditions.spongeCells[
2] # - self.asym_ghost_cells[0] - self.asym_ghost_cells[2]
y_zero_reference_cell = y_zero_reference_cell + boundary_conditions.spongeCells[
2]
rk_order = None
theta = None
A = None
super(FBL, self).__init__(gpu_ctx, \
nx, ny, \
ghost_cells_x, \
ghost_cells_y, \
dx, dy, dt, \
g, f, r, A, \
t, \
theta, rk_order, \
coriolis_beta, \
y_zero_reference_cell, \
wind_stress, \
write_netcdf, \
ignore_ghostcells, \
offset_x, offset_y, \
block_width, block_height)
self._set_interior_domain_from_sponge_cells()
#Get kernels
self.u_kernel = gpu_ctx.get_kernel("FBL_U_kernel.cu",
defines={
'block_width': block_width,
'block_height': block_height
})
self.v_kernel = gpu_ctx.get_kernel("FBL_V_kernel.cu",
defines={
'block_width': block_width,
'block_height': block_height
})
self.eta_kernel = gpu_ctx.get_kernel("FBL_eta_kernel.cu",
defines={
'block_width': block_width,
'block_height': block_height
})
# Get CUDA functions
self.computeUKernel = self.u_kernel.get_function("computeUKernel")
self.computeVKernel = self.v_kernel.get_function("computeVKernel")
self.computeEtaKernel = self.eta_kernel.get_function(
"computeEtaKernel")
# Prepare kernel lauches
self.computeUKernel.prepare("iiffffffffPiPiPiPif")
self.computeVKernel.prepare("iiffffffffPiPiPiPif")
self.computeEtaKernel.prepare("iiffffffffPiPiPiPi")
# Set up textures
self.update_wind_stress(self.u_kernel, self.computeUKernel)
self.update_wind_stress(self.v_kernel, self.computeVKernel)
self.H = Common.CUDAArray2D(self.gpu_stream, nx, ny, ghost_cells_x,
ghost_cells_y, H, self.asym_ghost_cells)
self.gpu_data = Common.SWEDataArakawaC(self.gpu_stream, nx, ny,
ghost_cells_x, ghost_cells_y,
eta0, hu0, hv0,
self.asym_ghost_cells)
# Overwrite halo with asymetric ghost cells
self.nx_halo = np.int32(nx + self.asym_ghost_cells[1] +
self.asym_ghost_cells[3])
self.ny_halo = np.int32(ny + self.asym_ghost_cells[0] +
self.asym_ghost_cells[2])
self.bc_kernel = FBL_periodic_boundary(self.gpu_ctx, \
self.nx, \
self.ny, \
self.boundary_conditions, \
self.asym_ghost_cells
)
self.totalNumIterations = 0
if self.write_netcdf:
self.sim_writer = SimWriter.SimNetCDFWriter(self, ignore_ghostcells=self.ignore_ghostcells, \
staggered_grid=True, offset_x=self.offset_x, offset_y=self.offset_y)
def fromfilename(cls, gpu_ctx, filename, cont_write_netcdf=True, use_lcg=False, new_netcdf_filename=None):
"""
Initialize and hotstart simulation from nc-file.
cont_write_netcdf: Continue to write the results after each superstep to a new netCDF file
filename: Continue simulation based on parameters and last timestep in this file
new_netcdf_filename: If we want to continue to write netcdf, we should use this filename. Automatically generated if None.
"""
# open nc-file
sim_reader = SimReader.SimNetCDFReader(filename, ignore_ghostcells=False)
sim_name = str(sim_reader.get('simulator_short'))
assert sim_name == cls.__name__, \
"Trying to initialize a " + \
cls.__name__ + " simulator with netCDF file based on " \
+ sim_name + " results."
# read the most recent state
H = sim_reader.getH();
# get last timestep (including simulation time of last timestep)
eta0, hu0, hv0, time0 = sim_reader.getLastTimeStep()
# For some reason, some old netcdf had 3-dimensional bathymetry.
# This fix ensures that we only use a valid H
if len(H.shape) == 3:
print("norm diff H: ", np.linalg.norm(H[0,:,:] - H[1,:,:]))
H = H[0,:,:]
# Set simulation parameters
sim_params = {
'gpu_ctx': gpu_ctx,
'eta0': eta0,
'hu0': hu0,
'hv0': hv0,
'H': H,
'nx': sim_reader.get("nx"),
'ny': sim_reader.get("ny"),
'dx': sim_reader.get("dx"),
'dy': sim_reader.get("dy"),
'dt': sim_reader.get("dt"),
'g': sim_reader.get("g"),
'f': sim_reader.get("coriolis_force"),
'r': sim_reader.get("bottom_friction_r"),
't': time0,
'theta': sim_reader.get("minmod_theta"),
'rk_order': sim_reader.get("time_integrator"),
'coriolis_beta': sim_reader.get("coriolis_beta"),
'y_zero_reference_cell': sim_reader.get("y_zero_reference_cell"),
'write_netcdf': cont_write_netcdf,
'use_lcg': use_lcg,
'netcdf_filename': new_netcdf_filename
}
# Wind stress
try:
wind_stress_type = sim_reader.get("wind_stress_type")
wind = Common.WindStressParams(type=wind_stress_type)
except:
wind = WindStress.WindStress()
sim_params['wind_stress'] = wind
# Boundary conditions
sim_params['boundary_conditions'] = Common.BoundaryConditions( \
sim_reader.getBC()[0], sim_reader.getBC()[1], \
sim_reader.getBC()[2], sim_reader.getBC()[3], \
sim_reader.getBCSpongeCells())
# Model errors
if sim_reader.has('small_scale_perturbation'):
sim_params['small_scale_perturbation'] = sim_reader.get('small_scale_perturbation') == 'True'
if sim_params['small_scale_perturbation']:
sim_params['small_scale_perturbation_amplitude'] = sim_reader.get('small_scale_perturbation_amplitude')
sim_params['small_scale_perturbation_interpolation_factor'] = sim_reader.get('small_scale_perturbation_interpolation_factor')
# Data assimilation parameters:
if sim_reader.has('model_time_step'):
sim_params['model_time_step'] = sim_reader.get('model_time_step')
return cls(**sim_params)
def __init__(self, \
gpu_ctx, \
eta0, hu0, hv0, H, \
nx, ny, \
dx, dy, dt, \
g, f, r, \
angle=np.array([[0]], dtype=np.float32), \
t=0.0, \
theta=1.3, rk_order=2, \
coriolis_beta=0.0, \
max_wind_direction_perturbation = 0, \
wind_stress=WindStress.WindStress(), \
boundary_conditions=Common.BoundaryConditions(), \
boundary_conditions_data=Common.BoundaryConditionsData(), \
small_scale_perturbation=False, \
small_scale_perturbation_amplitude=None, \
small_scale_perturbation_interpolation_factor = 1, \
model_time_step=None,
reportGeostrophicEquilibrium=False, \
use_lcg=False, \
write_netcdf=False, \
comm=None, \
netcdf_filename=None, \
ignore_ghostcells=False, \
courant_number=0.8, \
offset_x=0, offset_y=0, \
flux_slope_eps = 1.0e-1, \
desingularization_eps = 1.0e-1, \
depth_cutoff = 1.0e-5, \
block_width=32, block_height=8, num_threads_dt=256,
block_width_model_error=16, block_height_model_error=16):
"""
Initialization routine
eta0: Initial deviation from mean sea level incl ghost cells, (nx+2)*(ny+2) cells
hu0: Initial momentum along x-axis incl ghost cells, (nx+1)*(ny+2) cells
hv0: Initial momentum along y-axis incl ghost cells, (nx+2)*(ny+1) cells
H: Depth from equilibrium defined on cell corners, (nx+5)*(ny+5) corners
nx: Number of cells along x-axis
ny: Number of cells along y-axis
dx: Grid cell spacing along x-axis (20 000 m)
dy: Grid cell spacing along y-axis (20 000 m)
dt: Size of each timestep (90 s)
g: Gravitational accelleration (9.81 m/s^2)
f: Coriolis parameter (1.2e-4 s^1), effectively as f = f + beta*y
r: Bottom friction coefficient (2.4e-3 m/s)
angle: Angle of rotation from North to y-axis
t: Start simulation at time t
theta: MINMOD theta used the reconstructions of the derivatives in the numerical scheme
rk_order: Order of Runge Kutta method {1,2*,3}
coriolis_beta: Coriolis linear factor -> f = f + beta*(y-y_0)
max_wind_direction_perturbation: Large-scale model error emulation by per-time-step perturbation of wind direction by +/- max_wind_direction_perturbation (degrees)
wind_stress: Wind stress parameters
boundary_conditions: Boundary condition object
small_scale_perturbation: Boolean value for applying a stochastic model error
small_scale_perturbation_amplitude: Amplitude (q0 coefficient) for model error
small_scale_perturbation_interpolation_factor: Width factor for correlation in model error
model_time_step: The size of a data assimilation model step (default same as dt)
reportGeostrophicEquilibrium: Calculate the Geostrophic Equilibrium variables for each superstep
use_lcg: Use LCG as the random number generator. Default is False, which means using curand.
write_netcdf: Write the results after each superstep to a netCDF file
comm: MPI communicator
desingularization_eps: Used for desingularizing hu/h
flux_slope_eps: Used for setting zero flux for symmetric Riemann fan
depth_cutoff: Used for defining dry cells
netcdf_filename: Use this filename. (If not defined, a filename will be generated by SimWriter.)
"""
self.logger = logging.getLogger(__name__)
assert( rk_order < 4 or rk_order > 0 ), "Only 1st, 2nd and 3rd order Runge Kutta supported"
if (rk_order == 3):
assert(r == 0.0), "3rd order Runge Kutta supported only without friction"
# Sort out internally represented ghost_cells in the presence of given
# boundary conditions
ghost_cells_x = 2
ghost_cells_y = 2
#Coriolis at "first" cell
x_zero_reference_cell = ghost_cells_x
y_zero_reference_cell = ghost_cells_y # In order to pass it to the super constructor
# Boundary conditions
self.boundary_conditions = boundary_conditions
if (boundary_conditions.isSponge()):
nx = nx + boundary_conditions.spongeCells[1] + boundary_conditions.spongeCells[3] - 2*ghost_cells_x
ny = ny + boundary_conditions.spongeCells[0] + boundary_conditions.spongeCells[2] - 2*ghost_cells_y
x_zero_reference_cell += boundary_conditions.spongeCells[3]
y_zero_reference_cell += boundary_conditions.spongeCells[2]
#Compensate f for reference cell (first cell in internal of domain)
north = np.array([np.sin(angle[0,0]), np.cos(angle[0,0])])
f = f - coriolis_beta * (x_zero_reference_cell*dx*north[0] + y_zero_reference_cell*dy*north[1])
x_zero_reference_cell = 0
y_zero_reference_cell = 0
A = None
self.max_wind_direction_perturbation = max_wind_direction_perturbation
super(CDKLM16, self).__init__(gpu_ctx, \
nx, ny, \
ghost_cells_x, \
ghost_cells_y, \
dx, dy, dt, \
g, f, r, A, \
t, \
theta, rk_order, \
coriolis_beta, \
y_zero_reference_cell, \
wind_stress, \
write_netcdf, \
ignore_ghostcells, \
offset_x, offset_y, \
comm, \
block_width, block_height)
# Index range for interior domain (north, east, south, west)
# so that interior domain of eta is
# eta[self.interior_domain_indices[2]:self.interior_domain_indices[0], \
# self.interior_domain_indices[3]:self.interior_domain_indices[1] ]
self.interior_domain_indices = np.array([-2,-2,2,2])
self._set_interior_domain_from_sponge_cells()
defines={'block_width': block_width, 'block_height': block_height,
'KPSIMULATOR_DESING_EPS': str(desingularization_eps)+'f',
'KPSIMULATOR_FLUX_SLOPE_EPS': str(flux_slope_eps)+'f',
'KPSIMULATOR_DEPTH_CUTOFF': str(depth_cutoff)+'f'}
#Get kernels
self.kernel = gpu_ctx.get_kernel("CDKLM16_kernel.cu",
defines=defines,
compile_args={ # default, fast_math, optimal
'options' : ["--ftz=true", # false, true, true
"--prec-div=false", # true, false, false,
"--prec-sqrt=false", # true, false, false
"--fmad=false"] # true, true, false
#'options': ["--use_fast_math"]
#'options': ["--generate-line-info"],
#nvcc_options=["--maxrregcount=39"],
#'arch': "compute_50",
#'code': "sm_50"
},
jit_compile_args={
#jit_options=[(cuda.jit_option.MAX_REGISTERS, 39)]
}
)
# Get CUDA functions and define data types for prepared_{async_}call()
self.cdklm_swe_2D = self.kernel.get_function("cdklm_swe_2D")
self.cdklm_swe_2D.prepare("iiffffffffiiPiPiPiPiPiPiPiPiffi")
self.update_wind_stress(self.kernel, self.cdklm_swe_2D)
# CUDA functions for finding max time step size:
self.num_threads_dt = num_threads_dt
self.num_blocks_dt = np.int32(self.global_size[0]*self.global_size[1])
self.update_dt_kernels = gpu_ctx.get_kernel("max_dt.cu",
defines={'block_width': block_width,
'block_height': block_height,
'NUM_THREADS': self.num_threads_dt})
self.per_block_max_dt_kernel = self.update_dt_kernels.get_function("per_block_max_dt")
self.per_block_max_dt_kernel.prepare("iifffPiPiPiPifPi")
self.max_dt_reduction_kernel = self.update_dt_kernels.get_function("max_dt_reduction")
self.max_dt_reduction_kernel.prepare("iPP")
# Bathymetry
self.bathymetry = Common.Bathymetry(gpu_ctx, self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, H, boundary_conditions)
# Adjust eta for possible dry states
Hm = self.downloadBathymetry()[1]
eta0 = np.maximum(eta0, -Hm)
# Create data by uploading to device
self.gpu_data = Common.SWEDataArakawaA(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, eta0, hu0, hv0)
# Allocate memory for calculating maximum timestep
host_dt = np.zeros((self.global_size[1], self.global_size[0]), dtype=np.float32)
self.device_dt = Common.CUDAArray2D(self.gpu_stream, self.global_size[0], self.global_size[1],
0, 0, host_dt)
host_max_dt_buffer = np.zeros((1,1), dtype=np.float32)
self.max_dt_buffer = Common.CUDAArray2D(self.gpu_stream, 1, 1, 0, 0, host_max_dt_buffer)
self.courant_number = courant_number
## Allocating memory for geostrophical equilibrium variables
self.reportGeostrophicEquilibrium = np.int32(reportGeostrophicEquilibrium)
self.geoEq_uxpvy = None
self.geoEq_Kx = None
self.geoEq_Ly = None
if self.reportGeostrophicEquilibrium:
dummy_zero_array = np.zeros((ny+2*ghost_cells_y, nx+2*ghost_cells_x), dtype=np.float32, order='C')
self.geoEq_uxpvy = Common.CUDAArray2D(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, dummy_zero_array)
self.geoEq_Kx = Common.CUDAArray2D(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, dummy_zero_array)
self.geoEq_Ly = Common.CUDAArray2D(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, dummy_zero_array)
self.constant_equilibrium_depth = np.max(H)
self.bc_kernel = Common.BoundaryConditionsArakawaA(gpu_ctx, \
self.nx, \
self.ny, \
ghost_cells_x, \
ghost_cells_y, \
self.boundary_conditions, \
boundary_conditions_data, \
)
# Small scale perturbation:
self.small_scale_perturbation = small_scale_perturbation
self.small_scale_model_error = None
self.small_scale_perturbation_interpolation_factor = small_scale_perturbation_interpolation_factor
if small_scale_perturbation:
if small_scale_perturbation_amplitude is None:
self.small_scale_model_error = OceanStateNoise.OceanStateNoise.fromsim(self,
interpolation_factor=small_scale_perturbation_interpolation_factor,
use_lcg=use_lcg,
block_width=block_width_model_error,
block_height=block_height_model_error)
else:
self.small_scale_model_error = OceanStateNoise.OceanStateNoise.fromsim(self,
soar_q0=small_scale_perturbation_amplitude,
interpolation_factor=small_scale_perturbation_interpolation_factor,
use_lcg=use_lcg,
block_width=block_width_model_error,
block_height=block_height_model_error)
# Data assimilation model step size
self.model_time_step = model_time_step
if model_time_step is None:
self.model_time_step = self.dt
self.total_time_steps = 0
if self.write_netcdf:
self.sim_writer = SimWriter.SimNetCDFWriter(self, filename=netcdf_filename, ignore_ghostcells=self.ignore_ghostcells, \
offset_x=self.offset_x, offset_y=self.offset_y)
#Upload data to GPU and bind to texture reference
self.angle_texref = self.kernel.get_texref("angle_tex")
self.angle_texref.set_array(cuda.np_to_array(np.ascontiguousarray(angle, dtype=np.float32), order="C"))
# Set texture parameters
self.angle_texref.set_filter_mode(cuda.filter_mode.LINEAR) #bilinear interpolation
self.angle_texref.set_address_mode(0, cuda.address_mode.CLAMP) #no indexing outside domain
self.angle_texref.set_address_mode(1, cuda.address_mode.CLAMP)
self.angle_texref.set_flags(cuda.TRSF_NORMALIZED_COORDINATES) #Use [0, 1] indexing
def __init__(self, \
gpu_ctx, \
eta0, Hi, hu0, hv0, \
nx, ny, \
dx, dy, dt, \
g, f=0.0, r=0.0, \
t=0.0, \
theta=1.3, use_rk2=True,
coriolis_beta=0.0, \
y_zero_reference_cell = 0, \
wind_stress=WindStress.WindStress(), \
boundary_conditions=Common.BoundaryConditions(), \
write_netcdf=False, \
ignore_ghostcells=False, \
offset_x=0, offset_y=0, \
block_width=32, block_height=16):
"""
Initialization routine
eta0: Initial deviation from mean sea level incl ghost cells, (nx+2)*(ny+2) cells
hu0: Initial momentum along x-axis incl ghost cells, (nx+1)*(ny+2) cells
hv0: Initial momentum along y-axis incl ghost cells, (nx+2)*(ny+1) cells
Hi: Depth from equilibrium defined on cell corners, (nx+5)*(ny+5) corners
nx: Number of cells along x-axis
ny: Number of cells along y-axis
dx: Grid cell spacing along x-axis (20 000 m)
dy: Grid cell spacing along y-axis (20 000 m)
dt: Size of each timestep (90 s)
g: Gravitational accelleration (9.81 m/s^2)
f: Coriolis parameter (1.2e-4 s^1), effectively as f = f + beta*y
r: Bottom friction coefficient (2.4e-3 m/s)
t: Start simulation at time t
theta: MINMOD theta used the reconstructions of the derivatives in the numerical scheme
use_rk2: Boolean if to use 2nd order Runge-Kutta (false -> 1st order forward Euler)
coriolis_beta: Coriolis linear factor -> f = f + beta*(y-y_0)
y_zero_reference_cell: The cell representing y_0 in the above, defined as the lower face of the cell .
wind_stress: Wind stress parameters
boundary_conditions: Boundary condition object
write_netcdf: Write the results after each superstep to a netCDF file
"""
## After changing from (h, B) to (eta, H), several of the simulator settings used are wrong. This check will help detect that.
if ( np.sum(eta0 - Hi[:-1, :-1] > 0) > nx):
assert(False), "It seems you are using water depth/elevation h and bottom topography B, while you should use water level eta and equillibrium depth H."
ghost_cells_x = 2
ghost_cells_y = 2
y_zero_reference_cell = 2.0 + y_zero_reference_cell
# Boundary conditions
self.boundary_conditions = boundary_conditions
# Extend the computational domain if the boundary conditions
# require it
if (boundary_conditions.isSponge()):
nx = nx + boundary_conditions.spongeCells[1] + boundary_conditions.spongeCells[3] - 2*ghost_cells_x
ny = ny + boundary_conditions.spongeCells[0] + boundary_conditions.spongeCells[2] - 2*ghost_cells_y
y_zero_reference_cell = boundary_conditions.spongeCells[2] + y_zero_reference_cell
self.use_rk2 = use_rk2
rk_order = np.int32(use_rk2 + 1)
A = None
super(KP07, self).__init__(gpu_ctx, \
nx, ny, \
ghost_cells_x, \
ghost_cells_y, \
dx, dy, dt, \
g, f, r, A, \
t, \
theta, rk_order, \
coriolis_beta, \
y_zero_reference_cell, \
wind_stress, \
write_netcdf, \
ignore_ghostcells, \
offset_x, offset_y, \
block_width, block_height)
# Index range for interior domain (north, east, south, west)
# so that interior domain of eta is
# eta[self.interior_domain_indices[2]:self.interior_domain_indices[0], \
# self.interior_domain_indices[3]:self.interior_domain_indices[1] ]
self.interior_domain_indices = np.array([-2,-2,2,2])
self._set_interior_domain_from_sponge_cells()
#Get kernels
self.kp07_kernel = gpu_ctx.get_kernel("KP07_kernel.cu", defines={'block_width': block_width, 'block_height': block_height})
# Get CUDA functions and define data types for prepared_{async_}call()
self.swe_2D = self.kp07_kernel.get_function("swe_2D")
self.swe_2D.prepare("iifffffffffiPiPiPiPiPiPiPiPiiiiif")
self.update_wind_stress(self.kp07_kernel, self.swe_2D)
#Create data by uploading to device
self.gpu_data = Common.SWEDataArakawaA(self.gpu_stream, nx, ny, ghost_cells_x, ghost_cells_y, eta0, hu0, hv0)
#Bathymetry
self.bathymetry = Common.Bathymetry(self.gpu_ctx, self.gpu_stream, \
nx, ny, ghost_cells_x, ghost_cells_y, Hi, boundary_conditions)
self.bc_kernel = Common.BoundaryConditionsArakawaA(gpu_ctx, \
self.nx, \
self.ny, \
ghost_cells_x, \
ghost_cells_y, \
self.boundary_conditions)
if self.write_netcdf:
self.sim_writer = SimWriter.SimNetCDFWriter(self, ignore_ghostcells=self.ignore_ghostcells, \
offset_x=self.offset_x, offset_y=self.offset_y)
def __init__(self, \
gpu_ctx, \
eta0, H, hu0, hv0, \
nx, ny, \
dx, dy, dt, \
g, f=0.0, r=0.0, \
t=0.0, \
theta=1.3, use_rk2=True,
coriolis_beta=0.0, \
y_zero_reference_cell = 0, \
wind_stress=WindStress.WindStress(), \
boundary_conditions=Common.BoundaryConditions(), \
write_netcdf=False, \
comm=None, \
ignore_ghostcells=False, \
offset_x=0, offset_y=0, \
flux_slope_eps = 1.0e-1, \
depth_cutoff = 1.0e-5, \
block_width=32, block_height=16):
"""
Initialization routine
eta0: Initial deviation from mean sea level incl ghost cells, (nx+2)*(ny+2) cells
hu0: Initial momentum along x-axis incl ghost cells, (nx+1)*(ny+2) cells
hv0: Initial momentum along y-axis incl ghost cells, (nx+2)*(ny+1) cells
H: Depth from equilibrium defined on cell corners, (nx+5)*(ny+5) corners
nx: Number of cells along x-axis
ny: Number of cells along y-axis
dx: Grid cell spacing along x-axis (20 000 m)
dy: Grid cell spacing along y-axis (20 000 m)
dt: Size of each timestep (90 s)
g: Gravitational accelleration (9.81 m/s^2)
f: Coriolis parameter (1.2e-4 s^1), effectively as f = f + beta*y
r: Bottom friction coefficient (2.4e-3 m/s)
t: Start simulation at time t
theta: MINMOD theta used the reconstructions of the derivatives in the numerical scheme
use_rk2: Boolean if to use 2nd order Runge-Kutta (false -> 1st order forward Euler)
coriolis_beta: Coriolis linear factor -> f = f + beta*(y-y_0)
y_zero_reference_cell: The cell representing y_0 in the above, defined as the lower face of the cell .
wind_stress: Wind stress parameters
boundary_conditions: Boundary condition object
write_netcdf: Write the results after each superstep to a netCDF file
comm: MPI communicator
depth_cutoff: Used for defining dry cells
flux_slope_eps: Used for desingularization with dry cells
"""
ghost_cells_x = 2
ghost_cells_y = 2
y_zero_reference_cell = 2.0 + y_zero_reference_cell
# Boundary conditions
self.boundary_conditions = boundary_conditions
# Extend the computational domain if the boundary conditions
# require it
if (boundary_conditions.isSponge()):
nx = nx + boundary_conditions.spongeCells[
1] + boundary_conditions.spongeCells[3] - 2 * ghost_cells_x
ny = ny + boundary_conditions.spongeCells[
0] + boundary_conditions.spongeCells[2] - 2 * ghost_cells_y
y_zero_reference_cell = boundary_conditions.spongeCells[
2] + y_zero_reference_cell
self.use_rk2 = use_rk2
rk_order = np.int32(use_rk2 + 1)
A = None
super(KP07, self).__init__(gpu_ctx, \
nx, ny, \
ghost_cells_x, \
ghost_cells_y, \
dx, dy, dt, \
g, f, r, A, \
t, \
theta, rk_order, \
coriolis_beta, \
y_zero_reference_cell, \
wind_stress, \
write_netcdf, \
ignore_ghostcells, \
offset_x, offset_y, \
comm, \
block_width, block_height)
# Index range for interior domain (north, east, south, west)
# so that interior domain of eta is
# eta[self.interior_domain_indices[2]:self.interior_domain_indices[0], \
# self.interior_domain_indices[3]:self.interior_domain_indices[1] ]
self.interior_domain_indices = np.array([-2, -2, 2, 2])
self._set_interior_domain_from_sponge_cells()
# The ocean simulators and the swashes cases are defined on
# completely different scales. We therefore specify a different
# desingularization parameter if we run a swashes case.
# Typical values:
#ifndef SWASHES
#define KPSIMULATOR_FLUX_SLOPE_EPS 1e-1f
#define KPSIMULATOR_FLUX_SLOPE_EPS_4 1.0e-4f
#else
#define KPSIMULATOR_FLUX_SLOPE_EPS 1.0e-4f
#define KPSIMULATOR_FLUX_SLOPE_EPS_4 1.0e-16f
#endif
defines = {
'block_width': block_width,
'block_height': block_height,
'KPSIMULATOR_FLUX_SLOPE_EPS': str(flux_slope_eps) + 'f',
'KPSIMULATOR_FLUX_SLOPE_EPS_4': str(flux_slope_eps**4) + 'f',
'KPSIMULATOR_DEPTH_CUTOFF': str(depth_cutoff) + 'f'
}
#Get kernels
self.kp07_kernel = gpu_ctx.get_kernel(
"KP07_kernel.cu",
defines=defines,
compile_args={ # default, fast_math, optimal
'options': [
"--ftz=true", # false, true, true
"--prec-div=false", # true, false, false,
"--prec-sqrt=false", # true, false, false
"--fmad=false"
] # true, true, false
#'options': ["--use_fast_math"]
#'options': ["--generate-line-info"],
#nvcc_options=["--maxrregcount=39"],
#'arch': "compute_50",
#'code': "sm_50"
},
jit_compile_args={
#jit_options=[(cuda.jit_option.MAX_REGISTERS, 39)]
})
# Get CUDA functions and define data types for prepared_{async_}call()
self.swe_2D = self.kp07_kernel.get_function("swe_2D")
self.swe_2D.prepare("iifffffffffiPiPiPiPiPiPiPiPiiiiif")
self.update_wind_stress(self.kp07_kernel, self.swe_2D)
# Upload Bathymetry
self.bathymetry = Common.Bathymetry(self.gpu_ctx, self.gpu_stream, \
nx, ny, ghost_cells_x, ghost_cells_y, H, boundary_conditions)
# Adjust eta for possible dry states
Hm = self.downloadBathymetry()[1]
eta0 = np.maximum(eta0, -Hm)
#Create data by uploading to device
self.gpu_data = Common.SWEDataArakawaA(self.gpu_stream, nx, ny,
ghost_cells_x, ghost_cells_y,
eta0, hu0, hv0)
self.bc_kernel = Common.BoundaryConditionsArakawaA(gpu_ctx, \
self.nx, \
self.ny, \
ghost_cells_x, \
ghost_cells_y, \
self.boundary_conditions)
if self.write_netcdf:
self.sim_writer = SimWriter.SimNetCDFWriter(self, ignore_ghostcells=self.ignore_ghostcells, \
offset_x=self.offset_x, offset_y=self.offset_y)
def __init__(self, \
gpu_ctx, \
H, eta0, hu0, hv0, \
nx, ny, \
dx, dy, dt, \
g, f, r, A=0.0, \
t=0.0, \
coriolis_beta=0.0, \
y_zero_reference_cell = 0, \
wind_stress=WindStress.WindStress(), \
boundary_conditions=Common.BoundaryConditions(), \
write_netcdf=False, \
ignore_ghostcells=False, \
offset_x=0, offset_y=0, \
block_width=16, block_height=16):
"""
Initialization routine
H: Water depth incl ghost cells, (nx+2)*(ny+2) cells
eta0: Initial deviation from mean sea level incl ghost cells, (nx+2)*(ny+2) cells
hu0: Initial momentum along x-axis incl ghost cells, (nx+1)*(ny+2) cells
hv0: Initial momentum along y-axis incl ghost cells, (nx+2)*(ny+1) cells
nx: Number of cells along x-axis
ny: Number of cells along y-axis
dx: Grid cell spacing along x-axis (20 000 m)
dy: Grid cell spacing along y-axis (20 000 m)
dt: Size of each timestep (90 s)
g: Gravitational accelleration (9.81 m/s^2)
f: Coriolis parameter (1.2e-4 s^1), effectively as f = f + beta*y
r: Bottom friction coefficient (2.4e-3 m/s)
A: Eddy viscosity coefficient (O(dx))
t: Start simulation at time t
coriolis_beta: Coriolis linear factor -> f = f + beta*(y-y_0)
y_zero_reference_cell: The cell representing y_0 in the above, defined as the lower face of the cell .
wind_stress: Wind stress parameters
boundary_conditions: Boundary condition object
write_netcdf: Write the results after each superstep to a netCDF file
"""
# Sort out internally represented ghost_cells in the presence of given
# boundary conditions
halo_x = 1
halo_y = 1
ghost_cells_x = 1
ghost_cells_y = 1
y_zero_reference_cell = y_zero_reference_cell + 1
self.boundary_conditions = boundary_conditions
if boundary_conditions.isSponge():
nx = nx + boundary_conditions.spongeCells[
1] + boundary_conditions.spongeCells[3] - 2 * ghost_cells_x
ny = ny + boundary_conditions.spongeCells[
0] + boundary_conditions.spongeCells[2] - 2 * ghost_cells_y
y_zero_reference_cell = y_zero_reference_cell + boundary_conditions.spongeCells[
2]
# self.<parameters> are sat in parent constructor:
rk_order = None
theta = None
super(CTCS, self).__init__(gpu_ctx, \
nx, ny, \
ghost_cells_x, \
ghost_cells_y, \
dx, dy, dt, \
g, f, r, A, \
t, \
theta, rk_order, \
coriolis_beta, \
y_zero_reference_cell, \
wind_stress, \
write_netcdf, \
ignore_ghostcells, \
offset_x, offset_y, \
block_width, block_height)
# Index range for interior domain (north, east, south, west)
# so that interior domain of eta is
# eta[self.interior_domain_indices[2]:self.interior_domain_indices[0], \
# self.interior_domain_indices[3]:self.interior_domain_indices[1] ]
self.interior_domain_indices = np.array([-1, -1, 1, 1])
self._set_interior_domain_from_sponge_cells()
#Get kernels
self.u_kernel = gpu_ctx.get_kernel("CTCS_U_kernel.cu",
defines={
'block_width': block_width,
'block_height': block_height
})
self.v_kernel = gpu_ctx.get_kernel("CTCS_V_kernel.cu",
defines={
'block_width': block_width,
'block_height': block_height
})
self.eta_kernel = gpu_ctx.get_kernel("CTCS_eta_kernel.cu",
defines={
'block_width': block_width,
'block_height': block_height
})
# Get CUDA functions
self.computeUKernel = self.u_kernel.get_function("computeUKernel")
self.computeVKernel = self.v_kernel.get_function("computeVKernel")
self.computeEtaKernel = self.eta_kernel.get_function(
"computeEtaKernel")
# Prepare kernel lauches
self.computeUKernel.prepare("iiiifffffffffPiPiPiPiPif")
self.computeVKernel.prepare("iiiifffffffffPiPiPiPiPif")
self.computeEtaKernel.prepare("iiffffffffPiPiPi")
# Set up textures
self.update_wind_stress(self.u_kernel, self.computeUKernel)
self.update_wind_stress(self.v_kernel, self.computeVKernel)
#Create data by uploading to device
self.H = Common.CUDAArray2D(self.gpu_stream, nx, ny, halo_x, halo_y, H)
self.gpu_data = Common.SWEDataArakawaC(self.gpu_stream, nx, ny, halo_x,
halo_y, eta0, hu0, hv0)
# Global size needs to be larger than the default from parent.__init__
self.global_size = ( \
int(np.ceil((self.nx+2*halo_x) / float(self.local_size[0]))), \
int(np.ceil((self.ny+2*halo_y) / float(self.local_size[1]))) \
)
self.bc_kernel = CTCS_boundary_condition(gpu_ctx, \
self.nx, \
self.ny, \
self.boundary_conditions, \
halo_x, halo_y \
)
if self.write_netcdf:
self.sim_writer = SimWriter.SimNetCDFWriter(self, ignore_ghostcells=self.ignore_ghostcells, \
staggered_grid=True, offset_x=self.offset_x, offset_y=self.offset_y)
def __init__(self, \
gpu_ctx, \
eta0, hu0, hv0, H, \
nx, ny, \
dx, dy, dt, \
g, f, r, \
subsample_f=10, \
angle=np.array([[0]], dtype=np.float32), \
subsample_angle=10, \
latitude=None, \
t=0.0, \
theta=1.3, rk_order=2, \
coriolis_beta=0.0, \
max_wind_direction_perturbation = 0, \
wind_stress=WindStress.WindStress(), \
boundary_conditions=Common.BoundaryConditions(), \
boundary_conditions_data=Common.BoundaryConditionsData(), \
small_scale_perturbation=False, \
small_scale_perturbation_amplitude=None, \
small_scale_perturbation_interpolation_factor = 1, \
model_time_step=None,
reportGeostrophicEquilibrium=False, \
use_lcg=False, \
write_netcdf=False, \
comm=None, \
local_particle_id=0, \
super_dir_name=None, \
netcdf_filename=None, \
ignore_ghostcells=False, \
courant_number=0.8, \
offset_x=0, offset_y=0, \
flux_slope_eps = 1.0e-1, \
desingularization_eps = 1.0e-1, \
depth_cutoff = 1.0e-5, \
block_width=12, block_height=32, num_threads_dt=256,
block_width_model_error=16, block_height_model_error=16):
"""
Initialization routine
eta0: Initial deviation from mean sea level incl ghost cells, (nx+2)*(ny+2) cells
hu0: Initial momentum along x-axis incl ghost cells, (nx+1)*(ny+2) cells
hv0: Initial momentum along y-axis incl ghost cells, (nx+2)*(ny+1) cells
H: Depth from equilibrium defined on cell corners, (nx+5)*(ny+5) corners
nx: Number of cells along x-axis
ny: Number of cells along y-axis
dx: Grid cell spacing along x-axis (20 000 m)
dy: Grid cell spacing along y-axis (20 000 m)
dt: Size of each timestep (90 s)
g: Gravitational accelleration (9.81 m/s^2)
f: Coriolis parameter (1.2e-4 s^1), effectively as f = f + beta*y
r: Bottom friction coefficient (2.4e-3 m/s)
subsample_f: Subsample the coriolis f when creating texture by factor
angle: Angle of rotation from North to y-axis as a texture (cuda.Array) or numpy array (in radians)
subsample_angle: Subsample the angles given as input when creating texture by factor
latitude: Specify latitude. This will override any f and beta plane already set (in radians)
t: Start simulation at time t
theta: MINMOD theta used the reconstructions of the derivatives in the numerical scheme
rk_order: Order of Runge Kutta method {1,2*,3}
coriolis_beta: Coriolis linear factor -> f = f + beta*(y-y_0)
max_wind_direction_perturbation: Large-scale model error emulation by per-time-step perturbation of wind direction by +/- max_wind_direction_perturbation (degrees)
wind_stress: Wind stress parameters
boundary_conditions: Boundary condition object
small_scale_perturbation: Boolean value for applying a stochastic model error
small_scale_perturbation_amplitude: Amplitude (q0 coefficient) for model error
small_scale_perturbation_interpolation_factor: Width factor for correlation in model error
model_time_step: The size of a data assimilation model step (default same as dt)
reportGeostrophicEquilibrium: Calculate the Geostrophic Equilibrium variables for each superstep
use_lcg: Use LCG as the random number generator. Default is False, which means using curand.
write_netcdf: Write the results after each superstep to a netCDF file
comm: MPI communicator
local_particle_id: Local (for each MPI process) particle id
desingularization_eps: Used for desingularizing hu/h
flux_slope_eps: Used for setting zero flux for symmetric Riemann fan
depth_cutoff: Used for defining dry cells
super_dir_name: Directory to write netcdf files to
netcdf_filename: Use this filename. (If not defined, a filename will be generated by SimWriter.)
"""
self.logger = logging.getLogger(__name__)
assert (rk_order < 4 or rk_order > 0
), "Only 1st, 2nd and 3rd order Runge Kutta supported"
if (rk_order == 3):
assert (r == 0.0
), "3rd order Runge Kutta supported only without friction"
# Sort out internally represented ghost_cells in the presence of given
# boundary conditions
ghost_cells_x = 2
ghost_cells_y = 2
#Coriolis at "first" cell
x_zero_reference_cell = ghost_cells_x
y_zero_reference_cell = ghost_cells_y # In order to pass it to the super constructor
# Boundary conditions
self.boundary_conditions = boundary_conditions
#Compensate f for reference cell (first cell in internal of domain)
north = np.array([np.sin(angle[0, 0]), np.cos(angle[0, 0])])
f = f - coriolis_beta * (x_zero_reference_cell * dx * north[0] +
y_zero_reference_cell * dy * north[1])
x_zero_reference_cell = 0
y_zero_reference_cell = 0
A = None
self.max_wind_direction_perturbation = max_wind_direction_perturbation
super(CDKLM16, self).__init__(gpu_ctx, \
nx, ny, \
ghost_cells_x, \
ghost_cells_y, \
dx, dy, dt, \
g, f, r, A, \
t, \
theta, rk_order, \
coriolis_beta, \
y_zero_reference_cell, \
wind_stress, \
write_netcdf, \
ignore_ghostcells, \
offset_x, offset_y, \
comm, \
block_width, block_height,
local_particle_id=local_particle_id)
# Index range for interior domain (north, east, south, west)
# so that interior domain of eta is
# eta[self.interior_domain_indices[2]:self.interior_domain_indices[0], \
# self.interior_domain_indices[3]:self.interior_domain_indices[1] ]
self.interior_domain_indices = np.array([-2, -2, 2, 2])
defines = {
'block_width': block_width,
'block_height': block_height,
'KPSIMULATOR_DESING_EPS': "{:.12f}f".format(desingularization_eps),
'KPSIMULATOR_FLUX_SLOPE_EPS': "{:.12f}f".format(flux_slope_eps),
'KPSIMULATOR_DEPTH_CUTOFF': "{:.12f}f".format(depth_cutoff),
'THETA': "{:.12f}f".format(self.theta),
'RK_ORDER': int(self.rk_order),
'NX': int(self.nx),
'NY': int(self.ny),
'DX': "{:.12f}f".format(self.dx),
'DY': "{:.12f}f".format(self.dy),
'GRAV': "{:.12f}f".format(self.g),
'FRIC': "{:.12f}f".format(self.r)
}
#Get kernels
self.kernel = gpu_ctx.get_kernel(
"CDKLM16_kernel.cu",
defines=defines,
compile_args={ # default, fast_math, optimal
'options': [
"--ftz=true", # false, true, true
"--prec-div=false", # true, false, false,
"--prec-sqrt=false", # true, false, false
"--fmad=false"
] # true, true, false
#'options': ["--use_fast_math"]
#'options': ["--generate-line-info"],
#nvcc_options=["--maxrregcount=39"],
#'arch': "compute_50",
#'code': "sm_50"
},
jit_compile_args={
#jit_options=[(cuda.jit_option.MAX_REGISTERS, 39)]
})
# Get CUDA functions and define data types for prepared_{async_}call()
self.cdklm_swe_2D = self.kernel.get_function("cdklm_swe_2D")
self.cdklm_swe_2D.prepare("fiPiPiPiPiPiPiPiPiffi")
self.update_wind_stress(self.kernel, self.cdklm_swe_2D)
# CUDA functions for finding max time step size:
self.num_threads_dt = num_threads_dt
self.num_blocks_dt = np.int32(self.global_size[0] *
self.global_size[1])
self.update_dt_kernels = gpu_ctx.get_kernel("max_dt.cu",
defines={
'block_width':
block_width,
'block_height':
block_height,
'NUM_THREADS':
self.num_threads_dt
})
self.per_block_max_dt_kernel = self.update_dt_kernels.get_function(
"per_block_max_dt")
self.per_block_max_dt_kernel.prepare("iifffPiPiPiPifPi")
self.max_dt_reduction_kernel = self.update_dt_kernels.get_function(
"max_dt_reduction")
self.max_dt_reduction_kernel.prepare("iPP")
# Bathymetry
self.bathymetry = Common.Bathymetry(gpu_ctx, self.gpu_stream, nx, ny,
ghost_cells_x, ghost_cells_y, H,
boundary_conditions)
# Adjust eta for possible dry states
Hm = self.downloadBathymetry()[1]
eta0 = np.maximum(eta0, -Hm)
# Create data by uploading to device
self.gpu_data = Common.SWEDataArakawaA(self.gpu_stream, nx, ny,
ghost_cells_x, ghost_cells_y,
eta0, hu0, hv0)
# Allocate memory for calculating maximum timestep
host_dt = np.zeros((self.global_size[1], self.global_size[0]),
dtype=np.float32)
self.device_dt = Common.CUDAArray2D(self.gpu_stream,
self.global_size[0],
self.global_size[1], 0, 0, host_dt)
host_max_dt_buffer = np.zeros((1, 1), dtype=np.float32)
self.max_dt_buffer = Common.CUDAArray2D(self.gpu_stream, 1, 1, 0, 0,
host_max_dt_buffer)
self.courant_number = courant_number
## Allocating memory for geostrophical equilibrium variables
self.reportGeostrophicEquilibrium = np.int32(
reportGeostrophicEquilibrium)
self.geoEq_uxpvy = None
self.geoEq_Kx = None
self.geoEq_Ly = None
if self.reportGeostrophicEquilibrium:
dummy_zero_array = np.zeros(
(ny + 2 * ghost_cells_y, nx + 2 * ghost_cells_x),
dtype=np.float32,
order='C')
self.geoEq_uxpvy = Common.CUDAArray2D(self.gpu_stream, nx, ny,
ghost_cells_x, ghost_cells_y,
dummy_zero_array)
self.geoEq_Kx = Common.CUDAArray2D(self.gpu_stream, nx, ny,
ghost_cells_x, ghost_cells_y,
dummy_zero_array)
self.geoEq_Ly = Common.CUDAArray2D(self.gpu_stream, nx, ny,
ghost_cells_x, ghost_cells_y,
dummy_zero_array)
self.constant_equilibrium_depth = np.max(H)
self.bc_kernel = Common.BoundaryConditionsArakawaA(gpu_ctx, \
self.nx, \
self.ny, \
ghost_cells_x, \
ghost_cells_y, \
self.boundary_conditions, \
boundary_conditions_data, \
)
def subsample_texture(data, factor):
ny, nx = data.shape
dx, dy = 1 / nx, 1 / ny
I = interp2d(np.linspace(0.5 * dx, 1 - 0.5 * dx, nx),
np.linspace(0.5 * dy, 1 - 0.5 * dy, ny),
data,
kind='linear')
new_nx, new_ny = max(2, nx // factor), max(2, ny // factor)
new_dx, new_dy = 1 / new_nx, 1 / new_ny
x_new = np.linspace(0.5 * new_dx, 1 - 0.5 * new_dx, new_nx)
y_new = np.linspace(0.5 * new_dy, 1 - 0.5 * new_dy, new_ny)
return I(x_new, y_new)
# Texture for angle
self.angle_texref = self.kernel.get_texref("angle_tex")
if isinstance(angle, cuda.Array):
# angle is already a texture, so we just set the texture reference
self.angle_texref.set_array(angle)
else:
#Upload data to GPU and bind to texture reference
if (subsample_angle and angle.size >= eta0.size):
self.logger.info("Subsampling angle texture by factor " +
str(subsample_angle))
self.logger.warning(
"This will give inaccurate angle along the border!")
angle = subsample_texture(angle, subsample_angle)
self.angle_texref.set_array(
cuda.np_to_array(np.ascontiguousarray(angle, dtype=np.float32),
order="C"))
# Set texture parameters
self.angle_texref.set_filter_mode(
cuda.filter_mode.LINEAR) #bilinear interpolation
self.angle_texref.set_address_mode(
0, cuda.address_mode.CLAMP) #no indexing outside domain
self.angle_texref.set_address_mode(1, cuda.address_mode.CLAMP)
self.angle_texref.set_flags(
cuda.TRSF_NORMALIZED_COORDINATES) #Use [0, 1] indexing
# Texture for coriolis f
self.coriolis_texref = self.kernel.get_texref("coriolis_f_tex")
# Create the CPU coriolis
if (latitude is not None):
if (self.f != 0.0):
raise RuntimeError(
"Cannot specify both latitude and f. Make your mind up.")
coriolis_f, _ = OceanographicUtilities.calcCoriolisParams(latitude)
coriolis_f = coriolis_f.astype(np.float32)
else:
if (self.coriolis_beta != 0.0):
if (angle.size != 1):
raise RuntimeError(
"non-constant angle cannot be combined with beta plane model (makes no sense)"
)
#Generate coordinates for all cells, including ghost cells from center to center
# [-3/2dx, nx+3/2dx] for ghost_cells_x == 2
x = np.linspace((-self.ghost_cells_x + 0.5) * self.dx,
(self.nx + self.ghost_cells_x - 0.5) * self.dx,
self.nx + 2 * self.ghost_cells_x)
y = np.linspace((-self.ghost_cells_y + 0.5) * self.dy,
(self.ny + self.ghost_cells_y - 0.5) * self.dy,
self.ny + 2 * self.ghost_cells_x)
self.logger.info(
"Using latitude to create Coriolis f texture ({:f}x{:f} cells)"
.format(x.size, y.size))
x, y = np.meshgrid(x, y)
n = x * np.sin(angle[0, 0]) + y * np.cos(
angle[0, 0]) #North vector
coriolis_f = self.f + self.coriolis_beta * n
else:
if (self.f.size == 1):
coriolis_f = np.array([[self.f]], dtype=np.float32)
elif (self.f.shape == eta0.shape):
coriolis_f = np.array(self.f, dtype=np.float32)
else:
raise RuntimeError(
"The shape of f should match up with eta or be scalar."
)
if (subsample_f and coriolis_f.size >= eta0.size):
self.logger.info("Subsampling coriolis texture by factor " +
str(subsample_f))
self.logger.warning(
"This will give inaccurate coriolis along the border!")
coriolis_f = subsample_texture(coriolis_f, subsample_f)
#Upload data to GPU and bind to texture reference
self.coriolis_texref.set_array(
cuda.np_to_array(np.ascontiguousarray(coriolis_f,
dtype=np.float32),
order="C"))
# Set texture parameters
self.coriolis_texref.set_filter_mode(
cuda.filter_mode.LINEAR) #bilinear interpolation
self.coriolis_texref.set_address_mode(
0, cuda.address_mode.CLAMP) #no indexing outside domain
self.coriolis_texref.set_address_mode(1, cuda.address_mode.CLAMP)
self.coriolis_texref.set_flags(
cuda.TRSF_NORMALIZED_COORDINATES) #Use [0, 1] indexing
# Small scale perturbation:
self.small_scale_perturbation = small_scale_perturbation
self.small_scale_model_error = None
self.small_scale_perturbation_interpolation_factor = small_scale_perturbation_interpolation_factor
if small_scale_perturbation:
self.small_scale_model_error = OceanStateNoise.OceanStateNoise.fromsim(
self,
soar_q0=small_scale_perturbation_amplitude,
interpolation_factor=
small_scale_perturbation_interpolation_factor,
use_lcg=use_lcg,
block_width=block_width_model_error,
block_height=block_height_model_error)
# Data assimilation model step size
self.model_time_step = model_time_step
self.total_time_steps = 0
if model_time_step is None:
self.model_time_step = self.dt
if self.write_netcdf:
self.sim_writer = SimWriter.SimNetCDFWriter(self, super_dir_name=super_dir_name, filename=netcdf_filename, \
ignore_ghostcells=self.ignore_ghostcells, offset_x=self.offset_x, offset_y=self.offset_y)
# Update timestep if dt is given as zero
if self.dt <= 0:
self.updateDt()
def __init__(self,
comm,
local_ensemble_size=None,
drifter_positions=[],
sim_args={},
sim_ic={},
sim_bc_args={},
ensemble_args={}):
"""
Initialize the ensemble. Only rank 0 should receive the optional arguments.
The constructor handles initialization across nodes
"""
self.logger = logging.getLogger(__name__ + "_rank=" + str(comm.rank))
self.logger.debug("Initializing")
#Broadcast general information about ensemble
##########################
self.comm = comm
self.num_nodes = self.comm.size - 1 #Root does not participate
assert self.comm.size >= 2, "You appear to not be using enough MPI nodes (at least two required)"
self.local_ensemble_size = local_ensemble_size
self.local_ensemble_size = self.comm.bcast(self.local_ensemble_size,
root=0)
self.num_drifters = len(drifter_positions)
self.num_drifters = self.comm.bcast(self.num_drifters, root=0)
#Broadcast initial conditions for simulator
##########################
self.sim_args = sim_args
self.sim_args = self.comm.bcast(self.sim_args, root=0)
self.data_shape = (self.sim_args['ny'] + 4, self.sim_args['nx'] + 4)
if (self.comm.rank == 0):
assert (sim_ic['H'].dtype == np.float32)
assert (sim_ic['eta0'].dtype == np.float32)
assert (sim_ic['hu0'].dtype == np.float32)
assert (sim_ic['hv0'].dtype == np.float32)
assert (sim_ic['H'].shape == (self.data_shape[0] + 1,
self.data_shape[1] + 1))
assert (sim_ic['eta0'].shape == self.data_shape)
assert (sim_ic['hu0'].shape == self.data_shape)
assert (sim_ic['hv0'].shape == self.data_shape)
else:
#FIXME: hardcoded for CDKLM four ghost cells
sim_ic['H'] = np.empty(
(self.data_shape[0] + 1, self.data_shape[1] + 1),
dtype=np.float32)
sim_ic['eta0'] = np.empty(self.data_shape, dtype=np.float32)
sim_ic['hu0'] = np.empty(self.data_shape, dtype=np.float32)
sim_ic['hv0'] = np.empty(self.data_shape, dtype=np.float32)
#FIXME: Optimize this to one transfer by packing arrays?
self.comm.Bcast(sim_ic['H'], root=0)
self.comm.Bcast(sim_ic['eta0'], root=0)
self.comm.Bcast(sim_ic['hu0'], root=0)
self.comm.Bcast(sim_ic['hv0'], root=0)
self.logger.debug("eta0 is %s", str(sim_ic['eta0']))
#Broadcast arguments that we do not store in self
##############################
ensemble_args = self.comm.bcast(ensemble_args, root=0)
sim_bc_args = self.comm.bcast(sim_bc_args, root=0)
sim_ic['boundary_conditions'] = Common.BoundaryConditions(
**sim_bc_args)
#Create ensemble on local node
##############################
self.logger.info("Creating ensemble with %d members",
self.local_ensemble_size)
self.gpu_ctx = Common.CUDAContext()
if (self.comm.rank == 0):
num_ensemble_members = 1
else:
num_ensemble_members = self.local_ensemble_size
self.ensemble = OceanModelEnsemble.OceanModelEnsemble(
self.gpu_ctx,
self.sim_args,
sim_ic,
num_ensemble_members,
drifter_positions=drifter_positions,
**ensemble_args)
