def multTranspose(self, mat, x, y): sizes = self.mat.getSizes() for i in range(self.dimension): start = i stride = self.dimension xa = x.array_r[start::stride] ya = y.array_r[start::stride] xi = PETSc.Vec().createWithArray(xa, size=sizes[0], comm=x.comm) yi = PETSc.Vec().createWithArray(ya, size=sizes[1], comm=y.comm) self.mat.multTranspose(xi, yi) y.array[start::stride] = yi.array_r
def gather(self, global_indices=None): """Gather a :class:`Vector` to all processes :arg global_indices: the globally numbered indices to gather (should be the same on all processes). If `None`, gather the entire :class:`Vector`.""" if global_indices is None: N = self.size() v = PETSc.Vec().createSeq(N, comm=PETSc.COMM_SELF) is_ = PETSc.IS().createStride(N, 0, 1, comm=PETSc.COMM_SELF) else: global_indices = np.asarray(global_indices, dtype=np.int32) N = len(global_indices) v = PETSc.Vec().createSeq(N, comm=PETSc.COMM_SELF) is_ = PETSc.IS().createGeneral(global_indices, comm=PETSc.COMM_SELF) with self.dat.vec_ro as vec: vscat = PETSc.Scatter().create(vec, is_, v, None) vscat.scatterBegin(vec, v, addv=PETSc.InsertMode.INSERT_VALUES) vscat.scatterEnd(vec, v, addv=PETSc.InsertMode.INSERT_VALUES) return v.array
def __init__(self, function_space, coords): "create and assemble interpolation matrix" if not isinstance(function_space, WithGeometry): raise TypeError( "bad input type for function_space: must be a FunctionSpace") self.coords = np.copy(coords) self.function_space = function_space self.comm = function_space.comm self.n_data = coords.shape[0] assert (coords.shape[1] == self.function_space.mesh().cell_dimension() ), "shape of coordinates does not match mesh dimension" # allocate working vectors to handle parallel matrix operations and data transfer # dataspace_vector is a distributed PETSc vector in the data space self.dataspace_distrib = PETSc.Vec().create(comm=self.comm) self.dataspace_distrib.setSizes((-1, self.n_data)) self.dataspace_distrib.setFromOptions() self.n_data_local = self.dataspace_distrib.getSizes()[0] # all data computations are done on root process, so create gather method to # facilitate this data transfer self.petsc_scatter, self.dataspace_gathered = PETSc.Scatter.toZero( self.dataspace_distrib) self.meshspace_vector = Function(self.function_space).vector() self.n_mesh_local = self.meshspace_vector.local_size() self.n_mesh = self.meshspace_vector.size() nnz = len(self.function_space.cell_node_list[0]) self.interp = PETSc.Mat().create(comm=self.comm) self.interp.setSizes( ((self.n_mesh_local, -1), (self.n_data_local, -1))) self.interp.setPreallocationNNZ(nnz) self.interp.setFromOptions() self.interp.setUp() self.is_assembled = False
def __init__(self, function_space, sigma, l, cutoff=1.e-3, regularization=1.e-8, cov=sqexp): """ Create new forcing covariance Creates a new ForcingCovariance object from a function space, parameters, and covariance function. Required parameters are the function space and sigma and correlation length parameters needed to compute the covariance matrix. Note that this just initializes the object, and does not compute the matrix entries or assemble the final PETSc matrix. This is done using the ``assemble`` method, though if you attempt to use an unassembled matrix assembly will automatically be done. However the domain decomposition is done here to determine the number of DOFs handled by each process. """ # need to investigate parallelization here, load balancing likely to be uneven # if we just use the local ownership from the distributed matrix # since each row has an uneven amount of work # know that we have reduced bandwidth (though unclear if this translates to a low # bandwidth of the assembled covariance matrix) if not isinstance(function_space, WithGeometry): raise TypeError( "bad input type for function_space: must be a FunctionSpace") self.function_space = function_space self.comm = function_space.comm # extract mesh and process local information self.nx = Function(self.function_space).vector().size() self.nx_local = Function(self.function_space).vector().local_size() # set parameters and covariance assert regularization >= 0., "regularization parameter must be non-negative" self.sigma = sigma self.l = l self.cutoff = cutoff self.regularization = regularization self.cov = cov # get local ownership information of distributed matrix vtemp = PETSc.Vec().create(comm=self.comm) vtemp.setSizes((self.nx_local, -1)) vtemp.setFromOptions() vtemp.setUp() self.local_startind, self.local_endind = vtemp.getOwnershipRange() vtemp.destroy() self.is_assembled = False self.G = None
def interp_covariance_to_data(im_left, G, ls, im_right, ensemble_comm=COMM_SELF): """ Solve for the interpolated covariance matrix Solve for the covariance matrix interpolated to the sensor data locations. Note that the arguments allow for two different interpolation matrices to be used, in the event that we wish to compute the covariance matrix for other locations to make predictions. Note that since the Covariance Matrix is symmetric, it is advantageous to put the matrix with fewer spatial locations on the right as it will lead to fewer FEM solves (simply take the transpose if the reverse order is desired) This function solves :math:`\Phi_l^T A^{-1}GA^{-1}\Phi_r`, returning it as a 2D numpy array on the root process. This requires doing two FEM solves for each location provided in :math:`\Phi_r` plus some sparse matrix multiplication. Non-root processes will return an empty 2D numpy array (shape ``(0,0)``). The solves can be done independently, so optionally a Firedrake Ensemble communicator can be provided for dividing up the solves. The solves are simply divided up among the processes, so it is up to the user to determine an appropriate number of processes given the number of sensor locations that are needed. If not provided, the solves will be done serially. :param im_left: Left side InterpolationMatrix :type im_left: InterpolationMatrix :param G: Forcing covariance matrix to be used in the solve. :type G: ForcingCovariance :param ls: Firedrake LinearSolver to be used for the FEM solution :type ls: Firedrake LinearSolver :param im_right: Right side Interpolation Matrix :type im_right: InterpolationMatrix :param ensemble_comm: MPI Communicator over which to parallelize the FEM solves (optional, default is solve in series) :type ensemble_comm: MPI Communicator :returns: Covariance matrix interpolated to sensor locations. If run in parallel, this is returned on the root process as a 2D numpy array, while all other processes return an empty array (shape ``(0,0)``) :rtype: ndarray """ if not isinstance(im_left, InterpolationMatrix): raise TypeError("first argument to interp_covariance_to_data must be an InterpolationMatrix") if not isinstance(ls, LinearSolver): raise TypeError("ls must be a firedrake LinearSolver") if not isinstance(G, ForcingCovariance): raise TypeError("G must be a ForcingCovariance class") if not isinstance(im_right, InterpolationMatrix): raise TypeError("fourth argument to interp_covariance_to_data must be an InterpolationMatrix") if not isinstance(ensemble_comm, type(COMM_SELF)): raise TypeError("ensemble_comm must be an MPI communicator created from a firedrake Ensemble") # use ensemble comm to split up solves across ensemble processes v_tmp = PETSc.Vec().create(comm=ensemble_comm) v_tmp.setSizes((-1, im_right.n_data)) v_tmp.setFromOptions() imin, imax = v_tmp.getOwnershipRange() v_tmp.destroy() # create array for holding results # if root on base comm, will have data at the end of the solve/interpolation # otherwise, size will be zero if im_left.comm.rank == 0: n_local = im_left.n_data else: n_local = 0 # additional index is for the column vectors that this process owns in the # ensemble, which has length imax - imin result_tmparray = np.zeros((imax - imin, n_local)) for i in range(imin, imax): rhs = im_right.get_meshspace_column_vector(i) tmp = solve_forcing_covariance(G, ls, rhs) result_tmparray[i - imin] = im_left.interp_mesh_to_data(tmp) # create distributed vector for gathering results at root cov_distrib = PETSc.Vec().create(comm=ensemble_comm) cov_distrib.setSizes((n_local*(imax - imin), -1)) cov_distrib.setFromOptions() cov_distrib.array = result_tmparray.flatten() scatterfunc, cov_gathered = PETSc.Scatter.toZero(cov_distrib) scatterfunc.scatter(cov_distrib, cov_gathered, mode=PETSc.ScatterMode.SCATTER_FORWARD) out_array = np.copy(cov_gathered.array) cov_distrib.destroy() cov_gathered.destroy() # reshape output -- if I am root on both the main comm and ensemble comm then # I have the whole array. Other processes have nothing if im_left.comm.rank == 0 and ensemble_comm.rank == 0: outsize = (im_left.n_data, im_right.n_data) else: outsize = (0,0) return np.reshape(out_array, outsize)