Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
north_diag = vs.hur[2:-2, 3:-1] / vs.dyu[np.newaxis, 2:-2] / \
vs.dyt[np.newaxis, 3:-1] * vs.cost[np.newaxis, 3:-1] / vs.cosu[np.newaxis, 2:-2]
south_diag = vs.hur[2:-2, 2:-2] / vs.dyu[np.newaxis, 2:-2] / \
vs.dyt[np.newaxis, 2:-2] * vs.cost[np.newaxis, 2:-2] / vs.cosu[np.newaxis, 2:-2]
# construct sparse matrix
cf = tuple(diag for diag in (
boundary_mask * main_diag + (1 - boundary_mask),
boundary_mask * east_diag,
boundary_mask * west_diag,
boundary_mask * north_diag,
boundary_mask * south_diag
))
row = PETSc.Mat.Stencil()
col = PETSc.Mat.Stencil()
ij_offsets = [(0, 0), (1, 0), (-1, 0), (0, 1), (0, -1)]
(i0, i1), (j0, j1) = self._da.getRanges()
for j in range(j0, j1):
for i in range(i0, i1):
iloc, jloc = i % (vs.nx // rs.num_proc[0]), j % (vs.ny // rs.num_proc[1])
row.index = (i, j)
for diag, offset in zip(cf, ij_offsets):
io, jo = (i + offset[0], j + offset[1])
col.index = (io, jo)
matrix.setValueStencil(row, col, diag[iloc, jloc])
matrix.assemble()
vs.dxt[2:-2, np.newaxis] / vs.cosu[np.newaxis, 2:-2]**2
north_diag = vs.hur[2:-2, 3:-1] / vs.dyu[np.newaxis, 2:-2] / \
vs.dyt[np.newaxis, 3:-1] * vs.cost[np.newaxis, 3:-1] / vs.cosu[np.newaxis, 2:-2]
south_diag = vs.hur[2:-2, 2:-2] / vs.dyu[np.newaxis, 2:-2] / \
vs.dyt[np.newaxis, 2:-2] * vs.cost[np.newaxis, 2:-2] / vs.cosu[np.newaxis, 2:-2]
# construct sparse matrix
cf = tuple(diag for diag in (
boundary_mask * main_diag + (1 - boundary_mask),
boundary_mask * east_diag,
boundary_mask * west_diag,
boundary_mask * north_diag,
boundary_mask * south_diag
))
row = PETSc.Mat.Stencil()
col = PETSc.Mat.Stencil()
ij_offsets = [(0, 0), (1, 0), (-1, 0), (0, 1), (0, -1)]
(i0, i1), (j0, j1) = self._da.getRanges()
for j in range(j0, j1):
for i in range(i0, i1):
iloc, jloc = i % (vs.nx // rs.num_proc[0]), j % (vs.ny // rs.num_proc[1])
row.index = (i, j)
for diag, offset in zip(cf, ij_offsets):
io, jo = (i + offset[0], j + offset[1])
col.index = (io, jo)
matrix.setValueStencil(row, col, diag[iloc, jloc])
matrix.assemble()
col.field = 1
A.setValueStencil(row, col, value)
A.assemble()
A.view()
exit()
Id = da.createMatrix()
Id.setType('aij') # sparse
Id.setFromOptions()
Id.setPreallocationNNZ((5, 5))
Id.setUp()
Id.zeroEntries()
row = PETSc.Mat.Stencil()
col = PETSc.Mat.Stencil()
mx, my = da.getSizes()
(xs, xe), (ys, ye) = da.getRanges()
for j in range(ys, ye):
for i in range(xs, xe):
row.index = (i, j)
row.field = 0
col.index = (i, j)
col.field = 0
Id.setValueStencil(row, row, 1.0)
row.field = 1
col.field = 1
Id.setValueStencil(row, col, 1.0)
Id.assemble()
# (xs, xe), (ys, ye) = da.getRanges()
# print(A.getValues(range(n*n), range(n*n)))
Helper function to assemble PETSc matrix A
Returns:
PETSc matrix object
"""
# create matrix and set basic options
A = self.init.createMatrix()
A.setType('aij') # sparse
A.setFromOptions()
A.setPreallocationNNZ((5, 5))
A.setUp()
# fill matrix
A.zeroEntries()
row = PETSc.Mat.Stencil()
col = PETSc.Mat.Stencil()
mx, my = self.init.getSizes()
(xs, xe), (ys, ye) = self.init.getRanges()
for j in range(ys, ye):
for i in range(xs, xe):
row.index = (i, j)
row.field = 0
if i == 0 or j == 0 or i == mx - 1 or j == my - 1:
A.setValueStencil(row, row, 1.0)
else:
diag = self.params.nu * (-2.0 / self.dx ** 2 - 2.0 / self.dy ** 2)
for index, value in [
((i, j - 1), self.params.nu / self.dy ** 2),
((i - 1, j), self.params.nu / self.dx ** 2),
((i, j), diag),
((i + 1, j), self.params.nu / self.dx ** 2),
((i, j + 1), self.params.nu / self.dy ** 2),
"""
Function to return the Jacobian matrix
Args:
snes: nonlinear solver object
X: input vector
J: Jacobian matrix (current?)
P: Jacobian matrix (new)
Returns:
matrix status
"""
self.da.globalToLocal(X, self.localX)
x = self.da.getVecArray(self.localX)
P.zeroEntries()
row = PETSc.Mat.Stencil()
col = PETSc.Mat.Stencil()
(xs, xe), (ys, ye) = self.da.getRanges()
for j in range(ys, ye):
for i in range(xs, xe):
# diagnoal 2-by-2 block (for u and v)
row.index = (i, j)
col.index = (i, j)
row.field = 0
col.field = 0
val = 1.0 - self.factor * (self.params.Du * (-2.0 / self.dx ** 2 - 2.0 / self.dy ** 2) -
x[i, j, 1] ** 2 - self.params.A)
P.setValueStencil(row, col, val)
row.field = 0
col.field = 1
boundary_type=('ghosted', 'ghosted'))
# Create the rhs vector based on the DA.
b = da.createGlobalVec()
b_val = da.getVecArray(b) # Obtain access to elements of b.
b_val[50, 50] = 1; # Set central value to 1.
# Create (a vector to store) the solution vector.
x = da.createGlobalVec()
# Create the matrix.
A = da.getMatrix('aij')
# Stencil objects make it easy to set the values of the matrix elements.
row = PETSc.Mat.Stencil()
col = PETSc.Mat.Stencil()
# Set matrix elements to correct values.
(i0, i1), (j0, j1) = da.getRanges()
for j in range(j0, j1):
for i in range(i0, i1):
row.index = (i, j)
for index, value in [((i, j), -4 + w**2),
((i-1, j), 1),
((i+1, j), 1),
((i, j-1), 1),
((i, j+1), 1)]:
col.index = index
A.setValueStencil(row, col, value) # Sets a single matrix element.
A.assemblyBegin() # Make matrices useable.
A.assemblyEnd()
Helper function to assemble PETSc identity matrix
Returns:
PETSc matrix object
"""
# create matrix and set basic options
Id = self.init.createMatrix()
Id.setType('aij') # sparse
Id.setFromOptions()
Id.setPreallocationNNZ((1, 1))
Id.setUp()
# fill matrix
Id.zeroEntries()
row = PETSc.Mat.Stencil()
(xs, xe), (ys, ye) = self.init.getRanges()
for j in range(ys, ye):
for i in range(xs, xe):
row.index = (i, j)
row.field = 0
Id.setValueStencil(row, row, 1.0)
Id.assemble()
return Id
Helper function to assemble PETSc matrix A
Returns:
PETSc matrix object
"""
# create matrix and set basic options
A = self.init.createMatrix()
A.setType('aij') # sparse
A.setFromOptions()
A.setPreallocationNNZ((5, 5))
A.setUp()
# fill matrix
A.zeroEntries()
row = PETSc.Mat.Stencil()
col = PETSc.Mat.Stencil()
mx, my = self.init.getSizes()
(xs, xe), (ys, ye) = self.init.getRanges()
for j in range(ys, ye):
for i in range(xs, xe):
row.index = (i, j)
row.field = 0
A.setValueStencil(row, row, self.params.Du * (-2.0 / self.dx ** 2 - 2.0 / self.dy ** 2))
row.field = 1
A.setValueStencil(row, row, self.params.Dv * (-2.0 / self.dx ** 2 - 2.0 / self.dy ** 2))
col.index = (i, j - 1)
col.field = 0
row.field = 0
A.setValueStencil(row, col, self.params.Du / self.dy ** 2)
col.field = 1
row.field = 1
A.setValueStencil(row, col, self.params.Dv / self.dy ** 2)
Helper function to assemble PETSc matrix A
Returns:
PETSc matrix object
"""
# create matrix and set basic options
A = self.init.createMatrix()
A.setType('aij') # sparse
A.setFromOptions()
A.setPreallocationNNZ((3, 3))
A.setUp()
# fill matrix
A.zeroEntries()
row = PETSc.Mat.Stencil()
col = PETSc.Mat.Stencil()
mx = self.init.getSizes()[0]
(xs, xe) = self.init.getRanges()[0]
for i in range(xs, xe):
row.i = i
row.field = 0
if i == 0 or i == mx - 1:
A.setValueStencil(row, row, 1.0)
else:
diag = -2.0 / self.dx ** 2
for index, value in [
(i - 1, 1.0 / self.dx ** 2),
(i, diag),
(i + 1, 1.0 / self.dx ** 2),
]:
col.i = index
col.field = 0
params: problem parameters
factor: temporal factor (dt*Qd)
dx: grid spacing in x direction
"""
assert da.getDim() == 1
self.da = da
self.mx = self.da.getSizes()[0]
(self.xs, self.xe) = self.da.getRanges()[0]
self.dx = 100 / (self.mx - 1)
print(self.mx, self.dx, self.xs, self.xe)
self.lambda0 = 2.0
self.nu = 1
self.localX = self.da.createLocalVec()
self.row = PETSc.Mat.Stencil()
self.col = PETSc.Mat.Stencil()
self.mat = self.da.createMatrix()
self.mat.setType('aij') # sparse
self.mat.setFromOptions()
self.mat.setPreallocationNNZ((3, 3))
self.mat.setUp()
self.gvec = self.da.createGlobalVec()
self.rhs = self.da.createGlobalVec()