1 #include <../src/mat/impls/aij/seq/aij.h> 2 #include <petsc/private/dmpleximpl.h> /*I "petscdmplex.h" I*/ 3 #include <petsclandau.h> /*I "petsclandau.h" I*/ 4 #include <petscts.h> 5 #include <petscdmforest.h> 6 #include <petscdmcomposite.h> 7 8 /* Landau collision operator */ 9 10 /* relativistic terms */ 11 #if defined(PETSC_USE_REAL_SINGLE) 12 #define SPEED_OF_LIGHT 2.99792458e8F 13 #define C_0(v0) (SPEED_OF_LIGHT / v0) /* needed for relativistic tensor on all architectures */ 14 #else 15 #define SPEED_OF_LIGHT 2.99792458e8 16 #define C_0(v0) (SPEED_OF_LIGHT / v0) /* needed for relativistic tensor on all architectures */ 17 #endif 18 19 #include "land_tensors.h" 20 21 #if defined(PETSC_HAVE_OPENMP) 22 #include <omp.h> 23 #endif 24 25 static PetscErrorCode LandauGPUMapsDestroy(void **ptr) 26 { 27 P4estVertexMaps *maps = (P4estVertexMaps *)*ptr; 28 29 PetscFunctionBegin; 30 // free device data 31 if (maps[0].deviceType != LANDAU_CPU) { 32 #if defined(PETSC_HAVE_KOKKOS) 33 if (maps[0].deviceType == LANDAU_KOKKOS) { 34 PetscCall(LandauKokkosDestroyMatMaps(maps, maps[0].numgrids)); // implies Kokkos does 35 } 36 #endif 37 } 38 // free host data 39 for (PetscInt grid = 0; grid < maps[0].numgrids; grid++) { 40 PetscCall(PetscFree(maps[grid].c_maps)); 41 PetscCall(PetscFree(maps[grid].gIdx)); 42 } 43 PetscCall(PetscFree(maps)); 44 PetscFunctionReturn(PETSC_SUCCESS); 45 } 46 static PetscErrorCode energy_f(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf_dummy, PetscScalar *u, void *actx) 47 { 48 PetscReal v2 = 0; 49 50 PetscFunctionBegin; 51 /* compute v^2 / 2 */ 52 for (PetscInt i = 0; i < dim; ++i) v2 += x[i] * x[i]; 53 /* evaluate the Maxwellian */ 54 u[0] = v2 / 2; 55 PetscFunctionReturn(PETSC_SUCCESS); 56 } 57 58 /* needs double */ 59 static PetscErrorCode gamma_m1_f(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf_dummy, PetscScalar *u, void *actx) 60 { 61 PetscReal *c2_0_arr = ((PetscReal *)actx); 62 double u2 = 0, c02 = (double)*c2_0_arr, xx; 63 64 PetscFunctionBegin; 65 /* compute u^2 / 2 */ 66 for (PetscInt i = 0; i < dim; ++i) u2 += x[i] * x[i]; 67 /* gamma - 1 = g_eps, for conditioning and we only take derivatives */ 68 xx = u2 / c02; 69 #if defined(PETSC_USE_DEBUG) 70 u[0] = PetscSqrtReal(1. + xx); 71 #else 72 u[0] = xx / (PetscSqrtReal(1. + xx) + 1.) - 1.; // better conditioned. -1 might help condition and only used for derivative 73 #endif 74 PetscFunctionReturn(PETSC_SUCCESS); 75 } 76 77 /* 78 LandauFormJacobian_Internal - Evaluates Jacobian matrix. 79 80 Input Parameters: 81 . globX - input vector 82 . actx - optional user-defined context 83 . dim - dimension 84 85 Output Parameter: 86 . J0acP - Jacobian matrix filled, not created 87 */ 88 static PetscErrorCode LandauFormJacobian_Internal(Vec a_X, Mat JacP, const PetscInt dim, PetscReal shift, void *a_ctx) 89 { 90 LandauCtx *ctx = (LandauCtx *)a_ctx; 91 PetscInt numCells[LANDAU_MAX_GRIDS], Nq, Nb; 92 PetscQuadrature quad; 93 PetscReal Eq_m[LANDAU_MAX_SPECIES]; // could be static data w/o quench (ex2) 94 PetscScalar *cellClosure = NULL; 95 const PetscScalar *xdata = NULL; 96 PetscDS prob; 97 PetscContainer container; 98 P4estVertexMaps *maps; 99 Mat subJ[LANDAU_MAX_GRIDS * LANDAU_MAX_BATCH_SZ]; 100 101 PetscFunctionBegin; 102 PetscValidHeaderSpecific(a_X, VEC_CLASSID, 1); 103 PetscValidHeaderSpecific(JacP, MAT_CLASSID, 2); 104 PetscAssertPointer(ctx, 5); 105 /* check for matrix container for GPU assembly. Support CPU assembly for debugging */ 106 PetscCheck(ctx->plex[0] != NULL, ctx->comm, PETSC_ERR_ARG_WRONG, "Plex not created"); 107 PetscCall(PetscLogEventBegin(ctx->events[10], 0, 0, 0, 0)); 108 PetscCall(DMGetDS(ctx->plex[0], &prob)); // same DS for all grids 109 PetscCall(PetscObjectQuery((PetscObject)JacP, "assembly_maps", (PetscObject *)&container)); 110 if (container) { 111 PetscCheck(ctx->gpu_assembly, ctx->comm, PETSC_ERR_ARG_WRONG, "maps but no GPU assembly"); 112 PetscCall(PetscContainerGetPointer(container, (void **)&maps)); 113 PetscCheck(maps, ctx->comm, PETSC_ERR_ARG_WRONG, "empty GPU matrix container"); 114 for (PetscInt i = 0; i < ctx->num_grids * ctx->batch_sz; i++) subJ[i] = NULL; 115 } else { 116 PetscCheck(!ctx->gpu_assembly, ctx->comm, PETSC_ERR_ARG_WRONG, "No maps but GPU assembly"); 117 for (PetscInt tid = 0; tid < ctx->batch_sz; tid++) { 118 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) PetscCall(DMCreateMatrix(ctx->plex[grid], &subJ[LAND_PACK_IDX(tid, grid)])); 119 } 120 maps = NULL; 121 } 122 // get dynamic data (Eq is odd, for quench and Spitzer test) for CPU assembly and raw data for Jacobian GPU assembly. Get host numCells[], Nq (yuck) 123 PetscCall(PetscFEGetQuadrature(ctx->fe[0], &quad)); 124 PetscCall(PetscQuadratureGetData(quad, NULL, NULL, &Nq, NULL, NULL)); 125 PetscCall(PetscFEGetDimension(ctx->fe[0], &Nb)); 126 PetscCheck(Nq <= LANDAU_MAX_NQND, ctx->comm, PETSC_ERR_ARG_WRONG, "Order too high. Nq = %" PetscInt_FMT " > LANDAU_MAX_NQND (%d)", Nq, LANDAU_MAX_NQND); 127 PetscCheck(Nb <= LANDAU_MAX_NQND, ctx->comm, PETSC_ERR_ARG_WRONG, "Order too high. Nb = %" PetscInt_FMT " > LANDAU_MAX_NQND (%d)", Nb, LANDAU_MAX_NQND); 128 // get metadata for collecting dynamic data 129 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 130 PetscInt cStart, cEnd; 131 PetscCheck(ctx->plex[grid] != NULL, ctx->comm, PETSC_ERR_ARG_WRONG, "Plex not created"); 132 PetscCall(DMPlexGetHeightStratum(ctx->plex[grid], 0, &cStart, &cEnd)); 133 numCells[grid] = cEnd - cStart; // grids can have different topology 134 } 135 PetscCall(PetscLogEventEnd(ctx->events[10], 0, 0, 0, 0)); 136 if (shift == 0) { /* create dynamic point data: f_alpha for closure of each cell (cellClosure[nbatch,ngrids,ncells[g],f[Nb,ns[g]]]) or xdata */ 137 DM pack; 138 PetscCall(VecGetDM(a_X, &pack)); 139 PetscCheck(pack, PETSC_COMM_SELF, PETSC_ERR_PLIB, "pack has no DM"); 140 PetscCall(PetscLogEventBegin(ctx->events[1], 0, 0, 0, 0)); 141 for (PetscInt fieldA = 0; fieldA < ctx->num_species; fieldA++) { 142 Eq_m[fieldA] = ctx->Ez * ctx->t_0 * ctx->charges[fieldA] / (ctx->v_0 * ctx->masses[fieldA]); /* normalize dimensionless */ 143 if (dim == 2) Eq_m[fieldA] *= 2 * PETSC_PI; /* add the 2pi term that is not in Landau */ 144 } 145 if (!ctx->gpu_assembly) { 146 Vec *locXArray, *globXArray; 147 PetscScalar *cellClosure_it; 148 PetscInt cellClosure_sz = 0, nDMs, Nf[LANDAU_MAX_GRIDS]; 149 PetscSection section[LANDAU_MAX_GRIDS], globsection[LANDAU_MAX_GRIDS]; 150 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 151 PetscCall(DMGetLocalSection(ctx->plex[grid], §ion[grid])); 152 PetscCall(DMGetGlobalSection(ctx->plex[grid], &globsection[grid])); 153 PetscCall(PetscSectionGetNumFields(section[grid], &Nf[grid])); 154 } 155 /* count cellClosure size */ 156 PetscCall(DMCompositeGetNumberDM(pack, &nDMs)); 157 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) cellClosure_sz += Nb * Nf[grid] * numCells[grid]; 158 PetscCall(PetscMalloc1(cellClosure_sz * ctx->batch_sz, &cellClosure)); 159 cellClosure_it = cellClosure; 160 PetscCall(PetscMalloc(sizeof(*locXArray) * nDMs, &locXArray)); 161 PetscCall(PetscMalloc(sizeof(*globXArray) * nDMs, &globXArray)); 162 PetscCall(DMCompositeGetLocalAccessArray(pack, a_X, nDMs, NULL, locXArray)); 163 PetscCall(DMCompositeGetAccessArray(pack, a_X, nDMs, NULL, globXArray)); 164 for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) { // OpenMP (once) 165 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 166 Vec locX = locXArray[LAND_PACK_IDX(b_id, grid)], globX = globXArray[LAND_PACK_IDX(b_id, grid)], locX2; 167 PetscInt cStart, cEnd, ei; 168 PetscCall(VecDuplicate(locX, &locX2)); 169 PetscCall(DMGlobalToLocalBegin(ctx->plex[grid], globX, INSERT_VALUES, locX2)); 170 PetscCall(DMGlobalToLocalEnd(ctx->plex[grid], globX, INSERT_VALUES, locX2)); 171 PetscCall(DMPlexGetHeightStratum(ctx->plex[grid], 0, &cStart, &cEnd)); 172 for (ei = cStart; ei < cEnd; ++ei) { 173 PetscScalar *coef = NULL; 174 PetscCall(DMPlexVecGetClosure(ctx->plex[grid], section[grid], locX2, ei, NULL, &coef)); 175 PetscCall(PetscMemcpy(cellClosure_it, coef, Nb * Nf[grid] * sizeof(*cellClosure_it))); /* change if LandauIPReal != PetscScalar */ 176 PetscCall(DMPlexVecRestoreClosure(ctx->plex[grid], section[grid], locX2, ei, NULL, &coef)); 177 cellClosure_it += Nb * Nf[grid]; 178 } 179 PetscCall(VecDestroy(&locX2)); 180 } 181 } 182 PetscCheck(cellClosure_it - cellClosure == cellClosure_sz * ctx->batch_sz, PETSC_COMM_SELF, PETSC_ERR_PLIB, "iteration wrong %" PetscCount_FMT " != cellClosure_sz = %" PetscInt_FMT, cellClosure_it - cellClosure, cellClosure_sz * ctx->batch_sz); 183 PetscCall(DMCompositeRestoreLocalAccessArray(pack, a_X, nDMs, NULL, locXArray)); 184 PetscCall(DMCompositeRestoreAccessArray(pack, a_X, nDMs, NULL, globXArray)); 185 PetscCall(PetscFree(locXArray)); 186 PetscCall(PetscFree(globXArray)); 187 xdata = NULL; 188 } else { 189 PetscMemType mtype; 190 if (ctx->jacobian_field_major_order) { // get data in batch ordering 191 PetscCall(VecScatterBegin(ctx->plex_batch, a_X, ctx->work_vec, INSERT_VALUES, SCATTER_FORWARD)); 192 PetscCall(VecScatterEnd(ctx->plex_batch, a_X, ctx->work_vec, INSERT_VALUES, SCATTER_FORWARD)); 193 PetscCall(VecGetArrayReadAndMemType(ctx->work_vec, &xdata, &mtype)); 194 } else { 195 PetscCall(VecGetArrayReadAndMemType(a_X, &xdata, &mtype)); 196 } 197 PetscCheck(mtype == PETSC_MEMTYPE_HOST || ctx->deviceType != LANDAU_CPU, ctx->comm, PETSC_ERR_ARG_WRONG, "CPU run with device data: use -mat_type aij"); 198 cellClosure = NULL; 199 } 200 PetscCall(PetscLogEventEnd(ctx->events[1], 0, 0, 0, 0)); 201 } else xdata = cellClosure = NULL; 202 203 /* do it */ 204 if (ctx->deviceType == LANDAU_KOKKOS) { 205 #if defined(PETSC_HAVE_KOKKOS) 206 PetscCall(LandauKokkosJacobian(ctx->plex, Nq, Nb, ctx->batch_sz, ctx->num_grids, numCells, Eq_m, cellClosure, xdata, &ctx->SData_d, shift, ctx->events, ctx->mat_offset, ctx->species_offset, subJ, JacP)); 207 #else 208 SETERRQ(ctx->comm, PETSC_ERR_ARG_WRONG, "-landau_device_type %s not built", "kokkos"); 209 #endif 210 } else { /* CPU version */ 211 PetscTabulation *Tf; // used for CPU and print info. Same on all grids and all species 212 PetscInt ip_offset[LANDAU_MAX_GRIDS + 1], ipf_offset[LANDAU_MAX_GRIDS + 1], elem_offset[LANDAU_MAX_GRIDS + 1], IPf_sz_glb, IPf_sz_tot, num_grids = ctx->num_grids, Nf[LANDAU_MAX_GRIDS]; 213 PetscReal *ff, *dudx, *dudy, *dudz, *invJ_a = (PetscReal *)ctx->SData_d.invJ, *xx = (PetscReal *)ctx->SData_d.x, *yy = (PetscReal *)ctx->SData_d.y, *zz = (PetscReal *)ctx->SData_d.z, *ww = (PetscReal *)ctx->SData_d.w; 214 PetscReal *nu_alpha = (PetscReal *)ctx->SData_d.alpha, *nu_beta = (PetscReal *)ctx->SData_d.beta, *invMass = (PetscReal *)ctx->SData_d.invMass; 215 PetscReal(*lambdas)[LANDAU_MAX_GRIDS][LANDAU_MAX_GRIDS] = (PetscReal(*)[LANDAU_MAX_GRIDS][LANDAU_MAX_GRIDS])ctx->SData_d.lambdas; 216 PetscSection section[LANDAU_MAX_GRIDS], globsection[LANDAU_MAX_GRIDS]; 217 PetscScalar *coo_vals = NULL; 218 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 219 PetscCall(DMGetLocalSection(ctx->plex[grid], §ion[grid])); 220 PetscCall(DMGetGlobalSection(ctx->plex[grid], &globsection[grid])); 221 PetscCall(PetscSectionGetNumFields(section[grid], &Nf[grid])); 222 } 223 /* count IPf size, etc */ 224 PetscCall(PetscDSGetTabulation(prob, &Tf)); // Bf, &Df same for all grids 225 const PetscReal *const BB = Tf[0]->T[0], *const DD = Tf[0]->T[1]; 226 ip_offset[0] = ipf_offset[0] = elem_offset[0] = 0; 227 for (PetscInt grid = 0; grid < num_grids; grid++) { 228 PetscInt nfloc = ctx->species_offset[grid + 1] - ctx->species_offset[grid]; 229 elem_offset[grid + 1] = elem_offset[grid] + numCells[grid]; 230 ip_offset[grid + 1] = ip_offset[grid] + numCells[grid] * Nq; 231 ipf_offset[grid + 1] = ipf_offset[grid] + Nq * nfloc * numCells[grid]; 232 } 233 IPf_sz_glb = ipf_offset[num_grids]; 234 IPf_sz_tot = IPf_sz_glb * ctx->batch_sz; 235 // prep COO 236 PetscCall(PetscMalloc1(ctx->SData_d.coo_size, &coo_vals)); // allocate every time? 237 if (shift == 0.0) { /* compute dynamic data f and df and init data for Jacobian */ 238 #if defined(PETSC_HAVE_THREADSAFETY) 239 double starttime, endtime; 240 starttime = MPI_Wtime(); 241 #endif 242 PetscCall(PetscLogEventBegin(ctx->events[8], 0, 0, 0, 0)); 243 PetscCall(PetscMalloc4(IPf_sz_tot, &ff, IPf_sz_tot, &dudx, IPf_sz_tot, &dudy, (dim == 3 ? IPf_sz_tot : 0), &dudz)); 244 // F df/dx 245 for (PetscInt tid = 0; tid < ctx->batch_sz * elem_offset[num_grids]; tid++) { // for each element 246 const PetscInt b_Nelem = elem_offset[num_grids], b_elem_idx = tid % b_Nelem, b_id = tid / b_Nelem; // b_id == OMP thd_id in batch 247 // find my grid: 248 PetscInt grid = 0; 249 while (b_elem_idx >= elem_offset[grid + 1]) grid++; // yuck search for grid 250 { 251 const PetscInt loc_nip = numCells[grid] * Nq, loc_Nf = ctx->species_offset[grid + 1] - ctx->species_offset[grid], loc_elem = b_elem_idx - elem_offset[grid]; 252 const PetscInt moffset = LAND_MOFFSET(b_id, grid, ctx->batch_sz, ctx->num_grids, ctx->mat_offset); //b_id*b_N + ctx->mat_offset[grid]; 253 PetscScalar *coef, coef_buff[LANDAU_MAX_SPECIES * LANDAU_MAX_NQND]; 254 PetscReal *invJe = &invJ_a[(ip_offset[grid] + loc_elem * Nq) * dim * dim]; // ingJ is static data on batch 0 255 PetscInt b, f, q; 256 if (cellClosure) { 257 coef = &cellClosure[b_id * IPf_sz_glb + ipf_offset[grid] + loc_elem * Nb * loc_Nf]; // this is const 258 } else { 259 coef = coef_buff; 260 for (f = 0; f < loc_Nf; ++f) { 261 LandauIdx *const Idxs = &maps[grid].gIdx[loc_elem][f][0]; 262 for (b = 0; b < Nb; ++b) { 263 PetscInt idx = Idxs[b]; 264 if (idx >= 0) { 265 coef[f * Nb + b] = xdata[idx + moffset]; 266 } else { 267 idx = -idx - 1; 268 coef[f * Nb + b] = 0; 269 for (q = 0; q < maps[grid].num_face; q++) { 270 PetscInt id = maps[grid].c_maps[idx][q].gid; 271 PetscScalar scale = maps[grid].c_maps[idx][q].scale; 272 coef[f * Nb + b] += scale * xdata[id + moffset]; 273 } 274 } 275 } 276 } 277 } 278 /* get f and df */ 279 for (PetscInt qi = 0; qi < Nq; qi++) { 280 const PetscReal *invJ = &invJe[qi * dim * dim]; 281 const PetscReal *Bq = &BB[qi * Nb]; 282 const PetscReal *Dq = &DD[qi * Nb * dim]; 283 PetscReal u_x[LANDAU_DIM]; 284 /* get f & df */ 285 for (f = 0; f < loc_Nf; ++f) { 286 const PetscInt idx = b_id * IPf_sz_glb + ipf_offset[grid] + f * loc_nip + loc_elem * Nq + qi; 287 PetscInt b, e; 288 PetscReal refSpaceDer[LANDAU_DIM]; 289 ff[idx] = 0.0; 290 for (PetscInt d = 0; d < LANDAU_DIM; ++d) refSpaceDer[d] = 0.0; 291 for (b = 0; b < Nb; ++b) { 292 const PetscInt cidx = b; 293 ff[idx] += Bq[cidx] * PetscRealPart(coef[f * Nb + cidx]); 294 for (PetscInt d = 0; d < dim; ++d) refSpaceDer[d] += Dq[cidx * dim + d] * PetscRealPart(coef[f * Nb + cidx]); 295 } 296 for (PetscInt d = 0; d < LANDAU_DIM; ++d) { 297 for (e = 0, u_x[d] = 0.0; e < LANDAU_DIM; ++e) u_x[d] += invJ[e * dim + d] * refSpaceDer[e]; 298 } 299 dudx[idx] = u_x[0]; 300 dudy[idx] = u_x[1]; 301 #if LANDAU_DIM == 3 302 dudz[idx] = u_x[2]; 303 #endif 304 } 305 } // q 306 } // grid 307 } // grid*batch 308 PetscCall(PetscLogEventEnd(ctx->events[8], 0, 0, 0, 0)); 309 #if defined(PETSC_HAVE_THREADSAFETY) 310 endtime = MPI_Wtime(); 311 if (ctx->stage) ctx->times[LANDAU_F_DF] += (endtime - starttime); 312 #endif 313 } // Jacobian setup 314 // assemble Jacobian (or mass) 315 for (PetscInt tid = 0; tid < ctx->batch_sz * elem_offset[num_grids]; tid++) { // for each element 316 const PetscInt b_Nelem = elem_offset[num_grids]; 317 const PetscInt glb_elem_idx = tid % b_Nelem, b_id = tid / b_Nelem; 318 PetscInt grid = 0; 319 #if defined(PETSC_HAVE_THREADSAFETY) 320 double starttime, endtime; 321 starttime = MPI_Wtime(); 322 #endif 323 while (glb_elem_idx >= elem_offset[grid + 1]) grid++; 324 { 325 const PetscInt loc_Nf = ctx->species_offset[grid + 1] - ctx->species_offset[grid], loc_elem = glb_elem_idx - elem_offset[grid]; 326 const PetscInt moffset = LAND_MOFFSET(b_id, grid, ctx->batch_sz, ctx->num_grids, ctx->mat_offset), totDim = loc_Nf * Nq, elemMatSize = totDim * totDim; 327 PetscScalar *elemMat; 328 const PetscReal *invJe = &invJ_a[(ip_offset[grid] + loc_elem * Nq) * dim * dim]; 329 PetscCall(PetscMalloc1(elemMatSize, &elemMat)); 330 PetscCall(PetscMemzero(elemMat, elemMatSize * sizeof(*elemMat))); 331 if (shift == 0.0) { // Jacobian 332 PetscCall(PetscLogEventBegin(ctx->events[4], 0, 0, 0, 0)); 333 } else { // mass 334 PetscCall(PetscLogEventBegin(ctx->events[16], 0, 0, 0, 0)); 335 } 336 for (PetscInt qj = 0; qj < Nq; ++qj) { 337 const PetscInt jpidx_glb = ip_offset[grid] + qj + loc_elem * Nq; 338 PetscReal g0[LANDAU_MAX_SPECIES], g2[LANDAU_MAX_SPECIES][LANDAU_DIM], g3[LANDAU_MAX_SPECIES][LANDAU_DIM][LANDAU_DIM]; // could make a LANDAU_MAX_SPECIES_GRID ~ number of ions - 1 339 PetscInt d, d2, dp, d3, IPf_idx; 340 if (shift == 0.0) { // Jacobian 341 const PetscReal *const invJj = &invJe[qj * dim * dim]; 342 PetscReal gg2[LANDAU_MAX_SPECIES][LANDAU_DIM], gg3[LANDAU_MAX_SPECIES][LANDAU_DIM][LANDAU_DIM], gg2_temp[LANDAU_DIM], gg3_temp[LANDAU_DIM][LANDAU_DIM]; 343 const PetscReal vj[3] = {xx[jpidx_glb], yy[jpidx_glb], zz ? zz[jpidx_glb] : 0}, wj = ww[jpidx_glb]; 344 // create g2 & g3 345 for (d = 0; d < LANDAU_DIM; d++) { // clear accumulation data D & K 346 gg2_temp[d] = 0; 347 for (d2 = 0; d2 < LANDAU_DIM; d2++) gg3_temp[d][d2] = 0; 348 } 349 /* inner beta reduction */ 350 IPf_idx = 0; 351 for (PetscInt grid_r = 0, f_off = 0, ipidx = 0; grid_r < ctx->num_grids; grid_r++, f_off = ctx->species_offset[grid_r]) { // IPf_idx += nip_loc_r*Nfloc_r 352 PetscInt nip_loc_r = numCells[grid_r] * Nq, Nfloc_r = Nf[grid_r]; 353 for (PetscInt ei_r = 0, loc_fdf_idx = 0; ei_r < numCells[grid_r]; ++ei_r) { 354 for (PetscInt qi = 0; qi < Nq; qi++, ipidx++, loc_fdf_idx++) { 355 const PetscReal wi = ww[ipidx], x = xx[ipidx], y = yy[ipidx]; 356 PetscReal temp1[3] = {0, 0, 0}, temp2 = 0; 357 #if LANDAU_DIM == 2 358 PetscReal Ud[2][2], Uk[2][2], mask = (PetscAbs(vj[0] - x) < 100 * PETSC_SQRT_MACHINE_EPSILON && PetscAbs(vj[1] - y) < 100 * PETSC_SQRT_MACHINE_EPSILON) ? 0. : 1.; 359 LandauTensor2D(vj, x, y, Ud, Uk, mask); 360 #else 361 PetscReal U[3][3], z = zz[ipidx], mask = (PetscAbs(vj[0] - x) < 100 * PETSC_SQRT_MACHINE_EPSILON && PetscAbs(vj[1] - y) < 100 * PETSC_SQRT_MACHINE_EPSILON && PetscAbs(vj[2] - z) < 100 * PETSC_SQRT_MACHINE_EPSILON) ? 0. : 1.; 362 if (ctx->use_relativistic_corrections) { 363 LandauTensor3DRelativistic(vj, x, y, z, U, mask, C_0(ctx->v_0)); 364 } else { 365 LandauTensor3D(vj, x, y, z, U, mask); 366 } 367 #endif 368 for (PetscInt f = 0; f < Nfloc_r; ++f) { 369 const PetscInt idx = b_id * IPf_sz_glb + ipf_offset[grid_r] + f * nip_loc_r + ei_r * Nq + qi; // IPf_idx + f*nip_loc_r + loc_fdf_idx; 370 temp1[0] += dudx[idx] * nu_beta[f + f_off] * invMass[f + f_off] * (*lambdas)[grid][grid_r]; 371 temp1[1] += dudy[idx] * nu_beta[f + f_off] * invMass[f + f_off] * (*lambdas)[grid][grid_r]; 372 #if LANDAU_DIM == 3 373 temp1[2] += dudz[idx] * nu_beta[f + f_off] * invMass[f + f_off] * (*lambdas)[grid][grid_r]; 374 #endif 375 temp2 += ff[idx] * nu_beta[f + f_off] * (*lambdas)[grid][grid_r]; 376 } 377 temp1[0] *= wi; 378 temp1[1] *= wi; 379 #if LANDAU_DIM == 3 380 temp1[2] *= wi; 381 #endif 382 temp2 *= wi; 383 #if LANDAU_DIM == 2 384 for (d2 = 0; d2 < 2; d2++) { 385 for (d3 = 0; d3 < 2; ++d3) { 386 /* K = U * grad(f): g2=e: i,A */ 387 gg2_temp[d2] += Uk[d2][d3] * temp1[d3]; 388 /* D = -U * (I \kron (fx)): g3=f: i,j,A */ 389 gg3_temp[d2][d3] += Ud[d2][d3] * temp2; 390 } 391 } 392 #else 393 for (d2 = 0; d2 < 3; ++d2) { 394 for (d3 = 0; d3 < 3; ++d3) { 395 /* K = U * grad(f): g2 = e: i,A */ 396 gg2_temp[d2] += U[d2][d3] * temp1[d3]; 397 /* D = -U * (I \kron (fx)): g3 = f: i,j,A */ 398 gg3_temp[d2][d3] += U[d2][d3] * temp2; 399 } 400 } 401 #endif 402 } // qi 403 } // ei_r 404 IPf_idx += nip_loc_r * Nfloc_r; 405 } /* grid_r - IPs */ 406 PetscCheck(IPf_idx == IPf_sz_glb, PETSC_COMM_SELF, PETSC_ERR_PLIB, "IPf_idx != IPf_sz %" PetscInt_FMT " %" PetscInt_FMT, IPf_idx, IPf_sz_glb); 407 // add alpha and put in gg2/3 408 for (PetscInt fieldA = 0, f_off = ctx->species_offset[grid]; fieldA < loc_Nf; ++fieldA) { 409 for (d2 = 0; d2 < LANDAU_DIM; d2++) { 410 gg2[fieldA][d2] = gg2_temp[d2] * nu_alpha[fieldA + f_off]; 411 for (d3 = 0; d3 < LANDAU_DIM; d3++) gg3[fieldA][d2][d3] = -gg3_temp[d2][d3] * nu_alpha[fieldA + f_off] * invMass[fieldA + f_off]; 412 } 413 } 414 /* add electric field term once per IP */ 415 for (PetscInt fieldA = 0, f_off = ctx->species_offset[grid]; fieldA < loc_Nf; ++fieldA) gg2[fieldA][LANDAU_DIM - 1] += Eq_m[fieldA + f_off]; 416 /* Jacobian transform - g2, g3 */ 417 for (PetscInt fieldA = 0; fieldA < loc_Nf; ++fieldA) { 418 for (d = 0; d < dim; ++d) { 419 g2[fieldA][d] = 0.0; 420 for (d2 = 0; d2 < dim; ++d2) { 421 g2[fieldA][d] += invJj[d * dim + d2] * gg2[fieldA][d2]; 422 g3[fieldA][d][d2] = 0.0; 423 for (d3 = 0; d3 < dim; ++d3) { 424 for (dp = 0; dp < dim; ++dp) g3[fieldA][d][d2] += invJj[d * dim + d3] * gg3[fieldA][d3][dp] * invJj[d2 * dim + dp]; 425 } 426 g3[fieldA][d][d2] *= wj; 427 } 428 g2[fieldA][d] *= wj; 429 } 430 } 431 } else { // mass 432 PetscReal wj = ww[jpidx_glb]; 433 /* Jacobian transform - g0 */ 434 for (PetscInt fieldA = 0; fieldA < loc_Nf; ++fieldA) { 435 if (dim == 2) { 436 g0[fieldA] = wj * shift * 2. * PETSC_PI; // move this to below and remove g0 437 } else { 438 g0[fieldA] = wj * shift; // move this to below and remove g0 439 } 440 } 441 } 442 /* FE matrix construction */ 443 { 444 PetscInt fieldA, d, f, d2, g; 445 const PetscReal *BJq = &BB[qj * Nb], *DIq = &DD[qj * Nb * dim]; 446 /* assemble - on the diagonal (I,I) */ 447 for (fieldA = 0; fieldA < loc_Nf; fieldA++) { 448 for (f = 0; f < Nb; f++) { 449 const PetscInt i = fieldA * Nb + f; /* Element matrix row */ 450 for (g = 0; g < Nb; ++g) { 451 const PetscInt j = fieldA * Nb + g; /* Element matrix column */ 452 const PetscInt fOff = i * totDim + j; 453 if (shift == 0.0) { 454 for (d = 0; d < dim; ++d) { 455 elemMat[fOff] += DIq[f * dim + d] * g2[fieldA][d] * BJq[g]; 456 for (d2 = 0; d2 < dim; ++d2) elemMat[fOff] += DIq[f * dim + d] * g3[fieldA][d][d2] * DIq[g * dim + d2]; 457 } 458 } else { // mass 459 elemMat[fOff] += BJq[f] * g0[fieldA] * BJq[g]; 460 } 461 } 462 } 463 } 464 } 465 } /* qj loop */ 466 if (shift == 0.0) { // Jacobian 467 PetscCall(PetscLogEventEnd(ctx->events[4], 0, 0, 0, 0)); 468 } else { 469 PetscCall(PetscLogEventEnd(ctx->events[16], 0, 0, 0, 0)); 470 } 471 #if defined(PETSC_HAVE_THREADSAFETY) 472 endtime = MPI_Wtime(); 473 if (ctx->stage) ctx->times[LANDAU_KERNEL] += (endtime - starttime); 474 #endif 475 /* assemble matrix */ 476 if (!container) { 477 PetscInt cStart; 478 PetscCall(PetscLogEventBegin(ctx->events[6], 0, 0, 0, 0)); 479 PetscCall(DMPlexGetHeightStratum(ctx->plex[grid], 0, &cStart, NULL)); 480 PetscCall(DMPlexMatSetClosure(ctx->plex[grid], section[grid], globsection[grid], subJ[LAND_PACK_IDX(b_id, grid)], loc_elem + cStart, elemMat, ADD_VALUES)); 481 PetscCall(PetscLogEventEnd(ctx->events[6], 0, 0, 0, 0)); 482 } else { // GPU like assembly for debugging 483 PetscInt fieldA, q, f, g, d, nr, nc, rows0[LANDAU_MAX_Q_FACE] = {0}, cols0[LANDAU_MAX_Q_FACE] = {0}, rows[LANDAU_MAX_Q_FACE], cols[LANDAU_MAX_Q_FACE]; 484 PetscScalar vals[LANDAU_MAX_Q_FACE * LANDAU_MAX_Q_FACE] = {0}, row_scale[LANDAU_MAX_Q_FACE] = {0}, col_scale[LANDAU_MAX_Q_FACE] = {0}; 485 LandauIdx *coo_elem_offsets = (LandauIdx *)ctx->SData_d.coo_elem_offsets, *coo_elem_fullNb = (LandauIdx *)ctx->SData_d.coo_elem_fullNb, (*coo_elem_point_offsets)[LANDAU_MAX_NQND + 1] = (LandauIdx(*)[LANDAU_MAX_NQND + 1]) ctx->SData_d.coo_elem_point_offsets; 486 /* assemble - from the diagonal (I,I) in this format for DMPlexMatSetClosure */ 487 for (fieldA = 0; fieldA < loc_Nf; fieldA++) { 488 LandauIdx *const Idxs = &maps[grid].gIdx[loc_elem][fieldA][0]; 489 for (f = 0; f < Nb; f++) { 490 PetscInt idx = Idxs[f]; 491 if (idx >= 0) { 492 nr = 1; 493 rows0[0] = idx; 494 row_scale[0] = 1.; 495 } else { 496 idx = -idx - 1; 497 for (q = 0, nr = 0; q < maps[grid].num_face; q++, nr++) { 498 if (maps[grid].c_maps[idx][q].gid < 0) break; 499 rows0[q] = maps[grid].c_maps[idx][q].gid; 500 row_scale[q] = maps[grid].c_maps[idx][q].scale; 501 } 502 } 503 for (g = 0; g < Nb; ++g) { 504 idx = Idxs[g]; 505 if (idx >= 0) { 506 nc = 1; 507 cols0[0] = idx; 508 col_scale[0] = 1.; 509 } else { 510 idx = -idx - 1; 511 nc = maps[grid].num_face; 512 for (q = 0, nc = 0; q < maps[grid].num_face; q++, nc++) { 513 if (maps[grid].c_maps[idx][q].gid < 0) break; 514 cols0[q] = maps[grid].c_maps[idx][q].gid; 515 col_scale[q] = maps[grid].c_maps[idx][q].scale; 516 } 517 } 518 const PetscInt i = fieldA * Nb + f; /* Element matrix row */ 519 const PetscInt j = fieldA * Nb + g; /* Element matrix column */ 520 const PetscScalar Aij = elemMat[i * totDim + j]; 521 if (coo_vals) { // mirror (i,j) in CreateStaticGPUData 522 const PetscInt fullNb = coo_elem_fullNb[glb_elem_idx], fullNb2 = fullNb * fullNb; 523 const PetscInt idx0 = b_id * coo_elem_offsets[elem_offset[num_grids]] + coo_elem_offsets[glb_elem_idx] + fieldA * fullNb2 + fullNb * coo_elem_point_offsets[glb_elem_idx][f] + nr * coo_elem_point_offsets[glb_elem_idx][g]; 524 for (PetscInt q = 0, idx2 = idx0; q < nr; q++) { 525 for (PetscInt d = 0; d < nc; d++, idx2++) coo_vals[idx2] = row_scale[q] * col_scale[d] * Aij; 526 } 527 } else { 528 for (q = 0; q < nr; q++) rows[q] = rows0[q] + moffset; 529 for (d = 0; d < nc; d++) cols[d] = cols0[d] + moffset; 530 for (q = 0; q < nr; q++) { 531 for (d = 0; d < nc; d++) vals[q * nc + d] = row_scale[q] * col_scale[d] * Aij; 532 } 533 PetscCall(MatSetValues(JacP, nr, rows, nc, cols, vals, ADD_VALUES)); 534 } 535 } 536 } 537 } 538 } 539 if (loc_elem == -1) { 540 PetscCall(PetscPrintf(ctx->comm, "CPU Element matrix\n")); 541 for (PetscInt d = 0; d < totDim; ++d) { 542 for (PetscInt f = 0; f < totDim; ++f) PetscCall(PetscPrintf(ctx->comm, " %12.5e", (double)PetscRealPart(elemMat[d * totDim + f]))); 543 PetscCall(PetscPrintf(ctx->comm, "\n")); 544 } 545 exit(12); 546 } 547 PetscCall(PetscFree(elemMat)); 548 } /* grid */ 549 } /* outer element & batch loop */ 550 if (shift == 0.0) { // mass 551 PetscCall(PetscFree4(ff, dudx, dudy, dudz)); 552 } 553 if (!container) { // 'CPU' assembly move nest matrix to global JacP 554 for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) { // OpenMP 555 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 556 const PetscInt moffset = LAND_MOFFSET(b_id, grid, ctx->batch_sz, ctx->num_grids, ctx->mat_offset); // b_id*b_N + ctx->mat_offset[grid]; 557 PetscInt nloc, nzl, colbuf[1024], row; 558 const PetscInt *cols; 559 const PetscScalar *vals; 560 Mat B = subJ[LAND_PACK_IDX(b_id, grid)]; 561 PetscCall(MatAssemblyBegin(B, MAT_FINAL_ASSEMBLY)); 562 PetscCall(MatAssemblyEnd(B, MAT_FINAL_ASSEMBLY)); 563 PetscCall(MatGetSize(B, &nloc, NULL)); 564 for (PetscInt i = 0; i < nloc; i++) { 565 PetscCall(MatGetRow(B, i, &nzl, &cols, &vals)); 566 PetscCheck(nzl <= 1024, PetscObjectComm((PetscObject)B), PETSC_ERR_PLIB, "Row too big: %" PetscInt_FMT, nzl); 567 for (PetscInt j = 0; j < nzl; j++) colbuf[j] = moffset + cols[j]; 568 row = moffset + i; 569 PetscCall(MatSetValues(JacP, 1, &row, nzl, colbuf, vals, ADD_VALUES)); 570 PetscCall(MatRestoreRow(B, i, &nzl, &cols, &vals)); 571 } 572 PetscCall(MatDestroy(&B)); 573 } 574 } 575 } 576 if (coo_vals) { 577 PetscCall(MatSetValuesCOO(JacP, coo_vals, ADD_VALUES)); 578 PetscCall(PetscFree(coo_vals)); 579 } 580 } /* CPU version */ 581 PetscCall(MatAssemblyBegin(JacP, MAT_FINAL_ASSEMBLY)); 582 PetscCall(MatAssemblyEnd(JacP, MAT_FINAL_ASSEMBLY)); 583 /* clean up */ 584 if (cellClosure) PetscCall(PetscFree(cellClosure)); 585 if (xdata) PetscCall(VecRestoreArrayReadAndMemType(a_X, &xdata)); 586 PetscFunctionReturn(PETSC_SUCCESS); 587 } 588 589 static PetscErrorCode GeometryDMLandau(DM base, PetscInt point, PetscInt dim, const PetscReal abc[], PetscReal xyz[], void *a_ctx) 590 { 591 PetscReal r = abc[0], z = abc[1]; 592 LandauCtx *ctx = (LandauCtx *)a_ctx; 593 594 PetscFunctionBegin; 595 if (ctx->sphere && dim == 3) { // make sphere: works for one AMR and Q2 596 PetscInt nzero = 0, idx = 0; 597 xyz[0] = r; 598 xyz[1] = z; 599 xyz[2] = abc[2]; 600 for (PetscInt i = 0; i < 3; i++) { 601 if (PetscAbs(xyz[i]) < PETSC_SQRT_MACHINE_EPSILON) nzero++; 602 else idx = i; 603 } 604 if (nzero == 2) xyz[idx] *= 1.732050807568877; // sqrt(3) 605 else if (nzero == 1) { 606 for (PetscInt i = 0; i < 3; i++) xyz[i] *= 1.224744871391589; // sqrt(3/2) 607 } 608 } else { 609 xyz[0] = r; 610 xyz[1] = z; 611 if (dim == 3) xyz[2] = abc[2]; 612 } 613 PetscFunctionReturn(PETSC_SUCCESS); 614 } 615 616 /* create DMComposite of meshes for each species group */ 617 static PetscErrorCode LandauDMCreateVMeshes(MPI_Comm comm_self, const PetscInt dim, const char prefix[], LandauCtx *ctx, DM pack) 618 { 619 PetscFunctionBegin; 620 { /* p4est, quads */ 621 /* Create plex mesh of Landau domain */ 622 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 623 PetscReal par_radius = ctx->radius_par[grid], perp_radius = ctx->radius_perp[grid]; 624 if (!ctx->sphere && !ctx->simplex) { // 2 or 3D (only 3D option) 625 PetscReal lo[] = {-perp_radius, -par_radius, -par_radius}, hi[] = {perp_radius, par_radius, par_radius}; 626 DMBoundaryType periodicity[3] = {DM_BOUNDARY_NONE, dim == 2 ? DM_BOUNDARY_NONE : DM_BOUNDARY_NONE, DM_BOUNDARY_NONE}; 627 if (dim == 2) lo[0] = 0; 628 else { 629 lo[1] = -perp_radius; 630 hi[1] = perp_radius; // 3D y is a perp 631 } 632 PetscCall(DMPlexCreateBoxMesh(comm_self, dim, PETSC_FALSE, ctx->cells0, lo, hi, periodicity, PETSC_TRUE, 0, PETSC_TRUE, &ctx->plex[grid])); // TODO: make composite and create dm[grid] here 633 PetscCall(DMLocalizeCoordinates(ctx->plex[grid])); /* needed for periodic */ 634 if (dim == 3) PetscCall(PetscObjectSetName((PetscObject)ctx->plex[grid], "cube")); 635 else PetscCall(PetscObjectSetName((PetscObject)ctx->plex[grid], "half-plane")); 636 } else if (dim == 2) { 637 size_t len; 638 PetscCall(PetscStrlen(ctx->filename, &len)); 639 if (len) { 640 Vec coords; 641 PetscScalar *x; 642 PetscInt N; 643 char str[] = "-dm_landau_view_file_0"; 644 str[21] += grid; 645 PetscCall(DMPlexCreateFromFile(comm_self, ctx->filename, "plexland.c", PETSC_TRUE, &ctx->plex[grid])); 646 PetscCall(DMPlexOrient(ctx->plex[grid])); 647 PetscCall(DMGetCoordinatesLocal(ctx->plex[grid], &coords)); 648 PetscCall(VecGetSize(coords, &N)); 649 PetscCall(VecGetArray(coords, &x)); 650 /* scale by domain size */ 651 for (PetscInt i = 0; i < N; i += 2) { 652 x[i + 0] *= ctx->radius_perp[grid]; 653 x[i + 1] *= ctx->radius_par[grid]; 654 } 655 PetscCall(VecRestoreArray(coords, &x)); 656 PetscCall(PetscObjectSetName((PetscObject)ctx->plex[grid], ctx->filename)); 657 PetscCall(PetscInfo(ctx->plex[grid], "%" PetscInt_FMT ") Read %s mesh file (%s)\n", grid, ctx->filename, str)); 658 PetscCall(DMViewFromOptions(ctx->plex[grid], NULL, str)); 659 } else { // simplex forces a sphere 660 PetscInt numCells = ctx->simplex ? 12 : 6, cell_size = ctx->simplex ? 3 : 4, j; 661 const PetscInt numVerts = 11; 662 PetscInt cellsT[][4] = { 663 {0, 1, 6, 5 }, 664 {1, 2, 7, 6 }, 665 {2, 3, 8, 7 }, 666 {3, 4, 9, 8 }, 667 {5, 6, 7, 10}, 668 {10, 7, 8, 9 } 669 }; 670 PetscInt cellsS[][3] = { 671 {0, 1, 6 }, 672 {1, 2, 6 }, 673 {6, 2, 7 }, 674 {7, 2, 8 }, 675 {8, 2, 3 }, 676 {8, 3, 4 }, 677 {0, 6, 5 }, 678 {5, 6, 7 }, 679 {5, 7, 10}, 680 {10, 7, 9 }, 681 {9, 7, 8 }, 682 {9, 8, 4 } 683 }; 684 const PetscInt *pcell = (const PetscInt *)(ctx->simplex ? &cellsS[0][0] : &cellsT[0][0]); 685 PetscReal coords[11][2], *flatCoords = &coords[0][0]; 686 PetscReal rad = ctx->radius[grid]; 687 for (j = 0; j < 5; j++) { // outside edge 688 PetscReal z, r, theta = -PETSC_PI / 2 + (j % 5) * PETSC_PI / 4; 689 r = rad * PetscCosReal(theta); 690 coords[j][0] = r; 691 z = rad * PetscSinReal(theta); 692 coords[j][1] = z; 693 } 694 coords[j][0] = 0; 695 coords[j++][1] = -rad * ctx->sphere_inner_radius_90degree[grid]; 696 coords[j][0] = rad * ctx->sphere_inner_radius_45degree[grid] * 0.707106781186548; 697 coords[j++][1] = -rad * ctx->sphere_inner_radius_45degree[grid] * 0.707106781186548; 698 coords[j][0] = rad * ctx->sphere_inner_radius_90degree[grid]; 699 coords[j++][1] = 0; 700 coords[j][0] = rad * ctx->sphere_inner_radius_45degree[grid] * 0.707106781186548; 701 coords[j++][1] = rad * ctx->sphere_inner_radius_45degree[grid] * 0.707106781186548; 702 coords[j][0] = 0; 703 coords[j++][1] = rad * ctx->sphere_inner_radius_90degree[grid]; 704 coords[j][0] = 0; 705 coords[j++][1] = 0; 706 PetscCall(DMPlexCreateFromCellListPetsc(comm_self, 2, numCells, numVerts, cell_size, ctx->interpolate, pcell, 2, flatCoords, &ctx->plex[grid])); 707 PetscCall(PetscObjectSetName((PetscObject)ctx->plex[grid], "semi-circle")); 708 PetscCall(PetscInfo(ctx->plex[grid], "\t%" PetscInt_FMT ") Make circle %s mesh\n", grid, ctx->simplex ? "simplex" : "tensor")); 709 } 710 } else { 711 PetscCheck(dim == 3 && ctx->sphere && !ctx->simplex, ctx->comm, PETSC_ERR_ARG_WRONG, "not: dim == 3 && ctx->sphere && !ctx->simplex"); 712 PetscReal rad = ctx->radius[grid] / 1.732050807568877, inner_rad = rad * ctx->sphere_inner_radius_45degree[grid], outer_rad = rad; 713 const PetscInt numCells = 7, cell_size = 8, numVerts = 16; 714 const PetscInt cells[][8] = { 715 {0, 3, 2, 1, 4, 5, 6, 7 }, 716 {0, 4, 5, 1, 8, 9, 13, 12}, 717 {1, 5, 6, 2, 9, 10, 14, 13}, 718 {2, 6, 7, 3, 10, 11, 15, 14}, 719 {0, 3, 7, 4, 8, 12, 15, 11}, 720 {0, 1, 2, 3, 8, 11, 10, 9 }, 721 {4, 7, 6, 5, 12, 13, 14, 15} 722 }; 723 PetscReal coords[16 /* numVerts */][3]; 724 for (PetscInt j = 0; j < 4; j++) { // inner edge, low 725 coords[j][0] = inner_rad * (j == 0 || j == 3 ? 1 : -1); 726 coords[j][1] = inner_rad * (j / 2 < 1 ? 1 : -1); 727 coords[j][2] = inner_rad * -1; 728 } 729 for (PetscInt j = 0, jj = 4; j < 4; j++, jj++) { // inner edge, hi 730 coords[jj][0] = inner_rad * (j == 0 || j == 3 ? 1 : -1); 731 coords[jj][1] = inner_rad * (j / 2 < 1 ? 1 : -1); 732 coords[jj][2] = inner_rad * 1; 733 } 734 for (PetscInt j = 0, jj = 8; j < 4; j++, jj++) { // outer edge, low 735 coords[jj][0] = outer_rad * (j == 0 || j == 3 ? 1 : -1); 736 coords[jj][1] = outer_rad * (j / 2 < 1 ? 1 : -1); 737 coords[jj][2] = outer_rad * -1; 738 } 739 for (PetscInt j = 0, jj = 12; j < 4; j++, jj++) { // outer edge, hi 740 coords[jj][0] = outer_rad * (j == 0 || j == 3 ? 1 : -1); 741 coords[jj][1] = outer_rad * (j / 2 < 1 ? 1 : -1); 742 coords[jj][2] = outer_rad * 1; 743 } 744 PetscCall(DMPlexCreateFromCellListPetsc(comm_self, 3, numCells, numVerts, cell_size, ctx->interpolate, (const PetscInt *)cells, 3, (const PetscReal *)coords, &ctx->plex[grid])); 745 PetscCall(PetscObjectSetName((PetscObject)ctx->plex[grid], "cubed sphere")); 746 PetscCall(PetscInfo(ctx->plex[grid], "\t%" PetscInt_FMT ") Make cubed sphere %s mesh\n", grid, ctx->simplex ? "simplex" : "tensor")); 747 } 748 PetscCall(DMSetFromOptions(ctx->plex[grid])); 749 } // grid loop 750 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)pack, prefix)); 751 { /* convert to p4est (or whatever), wait for discretization to create pack */ 752 char convType[256]; 753 PetscBool flg; 754 755 PetscOptionsBegin(ctx->comm, prefix, "Mesh conversion options", "DMPLEX"); 756 PetscCall(PetscOptionsFList("-dm_landau_type", "Convert DMPlex to another format (p4est)", "plexland.c", DMList, DMPLEX, convType, 256, &flg)); 757 PetscOptionsEnd(); 758 if (flg) { 759 ctx->use_p4est = PETSC_TRUE; /* flag for Forest */ 760 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 761 DM dmforest; 762 PetscCall(DMConvert(ctx->plex[grid], convType, &dmforest)); 763 if (dmforest) { 764 PetscBool isForest; 765 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)dmforest, prefix)); 766 PetscCall(DMIsForest(dmforest, &isForest)); 767 if (isForest) { 768 if (ctx->sphere) PetscCall(DMForestSetBaseCoordinateMapping(dmforest, GeometryDMLandau, ctx)); 769 PetscCall(DMDestroy(&ctx->plex[grid])); 770 ctx->plex[grid] = dmforest; // Forest for adaptivity 771 } else SETERRQ(ctx->comm, PETSC_ERR_PLIB, "Converted to non Forest?"); 772 } else SETERRQ(ctx->comm, PETSC_ERR_PLIB, "Convert failed?"); 773 } 774 } else ctx->use_p4est = PETSC_FALSE; /* flag for Forest */ 775 } 776 } /* non-file */ 777 PetscCall(DMSetDimension(pack, dim)); 778 PetscCall(PetscObjectSetName((PetscObject)pack, "Mesh")); 779 PetscCall(DMSetApplicationContext(pack, ctx)); 780 PetscFunctionReturn(PETSC_SUCCESS); 781 } 782 783 static PetscErrorCode SetupDS(DM pack, PetscInt dim, PetscInt grid, LandauCtx *ctx) 784 { 785 PetscInt ii, i0; 786 char buf[256]; 787 PetscSection section; 788 789 PetscFunctionBegin; 790 for (ii = ctx->species_offset[grid], i0 = 0; ii < ctx->species_offset[grid + 1]; ii++, i0++) { 791 if (ii == 0) PetscCall(PetscSNPrintf(buf, sizeof(buf), "e")); 792 else PetscCall(PetscSNPrintf(buf, sizeof(buf), "i%" PetscInt_FMT, ii)); 793 /* Setup Discretization - FEM */ 794 PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, ctx->simplex, NULL, PETSC_DECIDE, &ctx->fe[ii])); 795 PetscCall(PetscObjectSetName((PetscObject)ctx->fe[ii], buf)); 796 PetscCall(DMSetField(ctx->plex[grid], i0, NULL, (PetscObject)ctx->fe[ii])); 797 } 798 PetscCall(DMCreateDS(ctx->plex[grid])); 799 PetscCall(DMGetLocalSection(ctx->plex[grid], §ion)); 800 for (PetscInt ii = ctx->species_offset[grid], i0 = 0; ii < ctx->species_offset[grid + 1]; ii++, i0++) { 801 if (ii == 0) PetscCall(PetscSNPrintf(buf, sizeof(buf), "se")); 802 else PetscCall(PetscSNPrintf(buf, sizeof(buf), "si%" PetscInt_FMT, ii)); 803 PetscCall(PetscSectionSetComponentName(section, i0, 0, buf)); 804 } 805 PetscFunctionReturn(PETSC_SUCCESS); 806 } 807 808 /* Define a Maxwellian function for testing out the operator. */ 809 810 /* Using cartesian velocity space coordinates, the particle */ 811 /* density, [1/m^3], is defined according to */ 812 813 /* $$ n=\int_{R^3} dv^3 \left(\frac{m}{2\pi T}\right)^{3/2}\exp [- mv^2/(2T)] $$ */ 814 815 /* Using some constant, c, we normalize the velocity vector into a */ 816 /* dimensionless variable according to v=c*x. Thus the density, $n$, becomes */ 817 818 /* $$ n=\int_{R^3} dx^3 \left(\frac{mc^2}{2\pi T}\right)^{3/2}\exp [- mc^2/(2T)*x^2] $$ */ 819 820 /* Defining $\theta=2T/mc^2$, we thus find that the probability density */ 821 /* for finding the particle within the interval in a box dx^3 around x is */ 822 823 /* f(x;\theta)=\left(\frac{1}{\pi\theta}\right)^{3/2} \exp [ -x^2/\theta ] */ 824 825 typedef struct { 826 PetscReal v_0; 827 PetscReal kT_m; 828 PetscReal n; 829 PetscReal shift; 830 } MaxwellianCtx; 831 832 static PetscErrorCode maxwellian(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf_dummy, PetscScalar *u, void *actx) 833 { 834 MaxwellianCtx *mctx = (MaxwellianCtx *)actx; 835 PetscInt i; 836 PetscReal v2 = 0, theta = 2 * mctx->kT_m / (mctx->v_0 * mctx->v_0), shift; /* theta = 2kT/mc^2 */ 837 838 PetscFunctionBegin; 839 /* compute the exponents, v^2 */ 840 for (i = 0; i < dim; ++i) v2 += x[i] * x[i]; 841 /* evaluate the Maxwellian */ 842 if (mctx->shift < 0) shift = -mctx->shift; 843 else { 844 u[0] = mctx->n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta)); 845 shift = mctx->shift; 846 } 847 if (shift != 0.) { 848 v2 = 0; 849 for (i = 0; i < dim - 1; ++i) v2 += x[i] * x[i]; 850 v2 += (x[dim - 1] - shift) * (x[dim - 1] - shift); 851 /* evaluate the shifted Maxwellian */ 852 u[0] += mctx->n * PetscPowReal(PETSC_PI * theta, -1.5) * (PetscExpReal(-v2 / theta)); 853 } 854 PetscFunctionReturn(PETSC_SUCCESS); 855 } 856 857 /*@ 858 DMPlexLandauAddMaxwellians - Add a Maxwellian distribution to a state 859 860 Collective 861 862 Input Parameters: 863 + dm - The mesh (local) 864 . time - Current time 865 . temps - Temperatures of each species (global) 866 . ns - Number density of each species (global) 867 . grid - index into current grid - just used for offset into `temp` and `ns` 868 . b_id - batch index 869 . n_batch - number of batches 870 - actx - Landau context 871 872 Output Parameter: 873 . X - The state (local to this grid) 874 875 Level: beginner 876 877 .seealso: `DMPlexLandauCreateVelocitySpace()` 878 @*/ 879 PetscErrorCode DMPlexLandauAddMaxwellians(DM dm, Vec X, PetscReal time, PetscReal temps[], PetscReal ns[], PetscInt grid, PetscInt b_id, PetscInt n_batch, void *actx) 880 { 881 LandauCtx *ctx = (LandauCtx *)actx; 882 PetscErrorCode (*initu[LANDAU_MAX_SPECIES])(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *); 883 PetscInt dim; 884 MaxwellianCtx *mctxs[LANDAU_MAX_SPECIES], data[LANDAU_MAX_SPECIES]; 885 886 PetscFunctionBegin; 887 PetscCall(DMGetDimension(dm, &dim)); 888 if (!ctx) PetscCall(DMGetApplicationContext(dm, &ctx)); 889 for (PetscInt ii = ctx->species_offset[grid], i0 = 0; ii < ctx->species_offset[grid + 1]; ii++, i0++) { 890 mctxs[i0] = &data[i0]; 891 data[i0].v_0 = ctx->v_0; // v_0 same for all grids 892 data[i0].kT_m = ctx->k * temps[ii] / ctx->masses[ii]; /* kT/m */ 893 data[i0].n = ns[ii]; 894 initu[i0] = maxwellian; 895 data[i0].shift = 0; 896 } 897 data[0].shift = ctx->electronShift; 898 /* need to make ADD_ALL_VALUES work - TODO */ 899 PetscCall(DMProjectFunction(dm, time, initu, (void **)mctxs, INSERT_ALL_VALUES, X)); 900 PetscFunctionReturn(PETSC_SUCCESS); 901 } 902 903 /* 904 LandauSetInitialCondition - Adds Maxwellians with context 905 906 Collective 907 908 Input Parameters: 909 . dm - The mesh 910 - grid - index into current grid - just used for offset into temp and ns 911 . b_id - batch index 912 - n_batch - number of batches 913 + actx - Landau context with T and n 914 915 Output Parameter: 916 . X - The state 917 918 Level: beginner 919 920 .seealso: `DMPlexLandauCreateVelocitySpace()`, `DMPlexLandauAddMaxwellians()` 921 */ 922 static PetscErrorCode LandauSetInitialCondition(DM dm, Vec X, PetscInt grid, PetscInt b_id, PetscInt n_batch, void *actx) 923 { 924 LandauCtx *ctx = (LandauCtx *)actx; 925 926 PetscFunctionBegin; 927 if (!ctx) PetscCall(DMGetApplicationContext(dm, &ctx)); 928 PetscCall(VecZeroEntries(X)); 929 PetscCall(DMPlexLandauAddMaxwellians(dm, X, 0.0, ctx->thermal_temps, ctx->n, grid, b_id, n_batch, ctx)); 930 PetscFunctionReturn(PETSC_SUCCESS); 931 } 932 933 // adapt a level once. Forest in/out 934 #if defined(PETSC_USE_INFO) 935 static const char *s_refine_names[] = {"RE", "Z1", "Origin", "Z2", "Uniform"}; 936 #endif 937 static PetscErrorCode adaptToleranceFEM(PetscFE fem, Vec sol, PetscInt type, PetscInt grid, LandauCtx *ctx, DM *newForest) 938 { 939 DM forest, plex, adaptedDM = NULL; 940 PetscDS prob; 941 PetscBool isForest; 942 PetscQuadrature quad; 943 PetscInt Nq, Nb, *Nb2, cStart, cEnd, c, dim, qj, k; 944 DMLabel adaptLabel = NULL; 945 946 PetscFunctionBegin; 947 forest = ctx->plex[grid]; 948 PetscCall(DMCreateDS(forest)); 949 PetscCall(DMGetDS(forest, &prob)); 950 PetscCall(DMGetDimension(forest, &dim)); 951 PetscCall(DMIsForest(forest, &isForest)); 952 PetscCheck(isForest, ctx->comm, PETSC_ERR_ARG_WRONG, "! Forest"); 953 PetscCall(DMConvert(forest, DMPLEX, &plex)); 954 PetscCall(DMPlexGetHeightStratum(plex, 0, &cStart, &cEnd)); 955 PetscCall(DMLabelCreate(PETSC_COMM_SELF, "adapt", &adaptLabel)); 956 PetscCall(PetscFEGetQuadrature(fem, &quad)); 957 PetscCall(PetscQuadratureGetData(quad, NULL, NULL, &Nq, NULL, NULL)); 958 PetscCheck(Nq <= LANDAU_MAX_NQND, ctx->comm, PETSC_ERR_ARG_WRONG, "Order too high. Nq = %" PetscInt_FMT " > LANDAU_MAX_NQND (%d)", Nq, LANDAU_MAX_NQND); 959 PetscCall(PetscFEGetDimension(ctx->fe[0], &Nb)); 960 PetscCall(PetscDSGetDimensions(prob, &Nb2)); 961 PetscCheck(Nb2[0] == Nb, ctx->comm, PETSC_ERR_ARG_WRONG, " Nb = %" PetscInt_FMT " != Nb (%" PetscInt_FMT ")", Nb, Nb2[0]); 962 PetscCheck(Nb <= LANDAU_MAX_NQND, ctx->comm, PETSC_ERR_ARG_WRONG, "Order too high. Nb = %" PetscInt_FMT " > LANDAU_MAX_NQND (%d)", Nb, LANDAU_MAX_NQND); 963 PetscCall(PetscInfo(sol, "%" PetscInt_FMT ") Refine phase: %s\n", grid, s_refine_names[type])); 964 if (type == 4) { 965 for (c = cStart; c < cEnd; c++) PetscCall(DMLabelSetValue(adaptLabel, c, DM_ADAPT_REFINE)); 966 } else if (type == 2) { 967 PetscInt rCellIdx[8], nr = 0, nrmax = (dim == 3) ? 8 : 2; 968 PetscReal minRad = PETSC_INFINITY, r; 969 for (c = cStart; c < cEnd; c++) { 970 PetscReal tt, v0[LANDAU_MAX_NQND * 3], J[LANDAU_MAX_NQND * 9], invJ[LANDAU_MAX_NQND * 9], detJ[LANDAU_MAX_NQND]; 971 PetscCall(DMPlexComputeCellGeometryFEM(plex, c, quad, v0, J, invJ, detJ)); 972 (void)J; 973 (void)invJ; 974 for (qj = 0; qj < Nq; ++qj) { 975 tt = PetscSqr(v0[dim * qj + 0]) + PetscSqr(v0[dim * qj + 1]) + PetscSqr((dim == 3) ? v0[dim * qj + 2] : 0); 976 r = PetscSqrtReal(tt); 977 if (r < minRad - PETSC_SQRT_MACHINE_EPSILON * 10.) { 978 minRad = r; 979 nr = 0; 980 rCellIdx[nr++] = c; 981 PetscCall(PetscInfo(sol, "\t\t%" PetscInt_FMT ") Found first inner r=%e, cell %" PetscInt_FMT ", qp %" PetscInt_FMT "/%" PetscInt_FMT "\n", grid, (double)r, c, qj + 1, Nq)); 982 } else if ((r - minRad) < PETSC_SQRT_MACHINE_EPSILON * 100. && nr < nrmax) { 983 for (k = 0; k < nr; k++) 984 if (c == rCellIdx[k]) break; 985 if (k == nr) { 986 rCellIdx[nr++] = c; 987 PetscCall(PetscInfo(sol, "\t\t\t%" PetscInt_FMT ") Found another inner r=%e, cell %" PetscInt_FMT ", qp %" PetscInt_FMT "/%" PetscInt_FMT ", d=%e\n", grid, (double)r, c, qj + 1, Nq, (double)(r - minRad))); 988 } 989 } 990 } 991 } 992 for (k = 0; k < nr; k++) PetscCall(DMLabelSetValue(adaptLabel, rCellIdx[k], DM_ADAPT_REFINE)); 993 PetscCall(PetscInfo(sol, "\t\t\t%" PetscInt_FMT ") Refined %" PetscInt_FMT " origin cells %" PetscInt_FMT ",%" PetscInt_FMT " r=%g\n", grid, nr, rCellIdx[0], rCellIdx[1], (double)minRad)); 994 } else if (type == 0 || type == 1 || type == 3) { /* refine along r=0 axis */ 995 PetscScalar *coef = NULL; 996 Vec coords; 997 PetscInt csize, Nv, d, nz, nrefined = 0; 998 DM cdm; 999 PetscSection cs; 1000 PetscCall(DMGetCoordinatesLocal(forest, &coords)); 1001 PetscCall(DMGetCoordinateDM(forest, &cdm)); 1002 PetscCall(DMGetLocalSection(cdm, &cs)); 1003 for (c = cStart; c < cEnd; c++) { 1004 PetscInt doit = 0, outside = 0; 1005 PetscCall(DMPlexVecGetClosure(cdm, cs, coords, c, &csize, &coef)); 1006 Nv = csize / dim; 1007 for (nz = d = 0; d < Nv; d++) { 1008 PetscReal z = PetscRealPart(coef[d * dim + (dim - 1)]), x = PetscSqr(PetscRealPart(coef[d * dim + 0])) + ((dim == 3) ? PetscSqr(PetscRealPart(coef[d * dim + 1])) : 0); 1009 x = PetscSqrtReal(x); 1010 if (type == 0) { 1011 if (ctx->re_radius > PETSC_SQRT_MACHINE_EPSILON && (z < -PETSC_MACHINE_EPSILON * 10. || z > ctx->re_radius + PETSC_MACHINE_EPSILON * 10.)) outside++; /* first pass don't refine bottom */ 1012 } else if (type == 1 && (z > ctx->vperp0_radius1 || z < -ctx->vperp0_radius1)) { 1013 outside++; /* don't refine outside electron refine radius */ 1014 PetscCall(PetscInfo(sol, "\t%" PetscInt_FMT ") (debug) found %s cells\n", grid, s_refine_names[type])); 1015 } else if (type == 3 && (z > ctx->vperp0_radius2 || z < -ctx->vperp0_radius2)) { 1016 outside++; /* refine r=0 cells on refinement front */ 1017 PetscCall(PetscInfo(sol, "\t%" PetscInt_FMT ") (debug) found %s cells\n", grid, s_refine_names[type])); 1018 } 1019 if (x < PETSC_MACHINE_EPSILON * 10. && (type != 0 || ctx->re_radius > PETSC_SQRT_MACHINE_EPSILON)) nz++; 1020 } 1021 PetscCall(DMPlexVecRestoreClosure(cdm, cs, coords, c, &csize, &coef)); 1022 if (doit || (outside < Nv && nz)) { 1023 PetscCall(DMLabelSetValue(adaptLabel, c, DM_ADAPT_REFINE)); 1024 nrefined++; 1025 } 1026 } 1027 PetscCall(PetscInfo(sol, "\t%" PetscInt_FMT ") Refined %" PetscInt_FMT " cells\n", grid, nrefined)); 1028 } 1029 PetscCall(DMDestroy(&plex)); 1030 PetscCall(DMAdaptLabel(forest, adaptLabel, &adaptedDM)); 1031 PetscCall(DMLabelDestroy(&adaptLabel)); 1032 *newForest = adaptedDM; 1033 if (adaptedDM) { 1034 if (isForest) { 1035 PetscCall(DMForestSetAdaptivityForest(adaptedDM, NULL)); // ???? 1036 } 1037 PetscCall(DMConvert(adaptedDM, DMPLEX, &plex)); 1038 PetscCall(DMPlexGetHeightStratum(plex, 0, &cStart, &cEnd)); 1039 PetscCall(PetscInfo(sol, "\t\t\t\t%" PetscInt_FMT ") %" PetscInt_FMT " cells, %" PetscInt_FMT " total quadrature points\n", grid, cEnd - cStart, Nq * (cEnd - cStart))); 1040 PetscCall(DMDestroy(&plex)); 1041 } else *newForest = NULL; 1042 PetscFunctionReturn(PETSC_SUCCESS); 1043 } 1044 1045 // forest goes in (ctx->plex[grid]), plex comes out 1046 static PetscErrorCode adapt(PetscInt grid, LandauCtx *ctx, Vec *uu) 1047 { 1048 PetscInt adaptIter; 1049 1050 PetscFunctionBegin; 1051 PetscInt type, limits[5] = {(grid == 0) ? ctx->numRERefine : 0, (grid == 0) ? ctx->nZRefine1 : 0, ctx->numAMRRefine[grid], (grid == 0) ? ctx->nZRefine2 : 0, ctx->postAMRRefine[grid]}; 1052 for (type = 0; type < 5; type++) { 1053 for (adaptIter = 0; adaptIter < limits[type]; adaptIter++) { 1054 DM newForest = NULL; 1055 PetscCall(adaptToleranceFEM(ctx->fe[0], *uu, type, grid, ctx, &newForest)); 1056 if (newForest) { 1057 PetscCall(DMDestroy(&ctx->plex[grid])); 1058 PetscCall(VecDestroy(uu)); 1059 PetscCall(DMCreateGlobalVector(newForest, uu)); 1060 PetscCall(LandauSetInitialCondition(newForest, *uu, grid, 0, 1, ctx)); 1061 ctx->plex[grid] = newForest; 1062 } else { 1063 PetscCall(PetscInfo(*uu, "No refinement\n")); 1064 } 1065 } 1066 } 1067 PetscCall(PetscObjectSetName((PetscObject)*uu, "uAMR")); 1068 PetscFunctionReturn(PETSC_SUCCESS); 1069 } 1070 1071 // make log(Lambdas) from NRL Plasma formulary 1072 static PetscErrorCode makeLambdas(LandauCtx *ctx) 1073 { 1074 PetscFunctionBegin; 1075 for (PetscInt gridi = 0; gridi < ctx->num_grids; gridi++) { 1076 PetscInt iii = ctx->species_offset[gridi]; 1077 PetscReal Ti_ev = (ctx->thermal_temps[iii] / 1.1604525e7) * 1000; // convert (back) to eV 1078 PetscReal ni = ctx->n[iii] * ctx->n_0; 1079 for (PetscInt gridj = gridi; gridj < ctx->num_grids; gridj++) { 1080 PetscInt jjj = ctx->species_offset[gridj]; 1081 PetscReal Zj = ctx->charges[jjj] / 1.6022e-19; 1082 if (gridi == 0) { 1083 if (gridj == 0) { // lam_ee 1084 ctx->lambdas[gridi][gridj] = 23.5 - PetscLogReal(PetscSqrtReal(ni) * PetscPowReal(Ti_ev, -1.25)) - PetscSqrtReal(1e-5 + PetscSqr(PetscLogReal(Ti_ev) - 2) / 16); 1085 } else { // lam_ei == lam_ie 1086 if (10 * Zj * Zj > Ti_ev) { 1087 ctx->lambdas[gridi][gridj] = ctx->lambdas[gridj][gridi] = 23 - PetscLogReal(PetscSqrtReal(ni) * Zj * PetscPowReal(Ti_ev, -1.5)); 1088 } else { 1089 ctx->lambdas[gridi][gridj] = ctx->lambdas[gridj][gridi] = 24 - PetscLogReal(PetscSqrtReal(ni) / Ti_ev); 1090 } 1091 } 1092 } else { // lam_ii' 1093 PetscReal mui = ctx->masses[iii] / 1.6720e-27, Zi = ctx->charges[iii] / 1.6022e-19; 1094 PetscReal Tj_ev = (ctx->thermal_temps[jjj] / 1.1604525e7) * 1000; // convert (back) to eV 1095 PetscReal muj = ctx->masses[jjj] / 1.6720e-27; 1096 PetscReal nj = ctx->n[jjj] * ctx->n_0; 1097 ctx->lambdas[gridi][gridj] = ctx->lambdas[gridj][gridi] = 23 - PetscLogReal(Zi * Zj * (mui + muj) / (mui * Tj_ev + muj * Ti_ev) * PetscSqrtReal(ni * Zi * Zi / Ti_ev + nj * Zj * Zj / Tj_ev)); 1098 } 1099 } 1100 } 1101 //PetscReal v0 = PetscSqrtReal(ctx->k * ctx->thermal_temps[iii] / ctx->masses[iii]); /* arbitrary units for non-dimensionalization: plasma formulary def */ 1102 PetscFunctionReturn(PETSC_SUCCESS); 1103 } 1104 1105 static PetscErrorCode ProcessOptions(LandauCtx *ctx, const char prefix[]) 1106 { 1107 PetscBool flg, fileflg; 1108 PetscInt ii, nt, nm, nc, num_species_grid[LANDAU_MAX_GRIDS], non_dim_grid; 1109 PetscReal lnLam = 10; 1110 DM dummy; 1111 1112 PetscFunctionBegin; 1113 PetscCall(DMCreate(ctx->comm, &dummy)); 1114 /* get options - initialize context */ 1115 ctx->verbose = 1; // should be 0 for silent compliance 1116 ctx->batch_sz = 1; 1117 ctx->batch_view_idx = 0; 1118 ctx->interpolate = PETSC_TRUE; 1119 ctx->gpu_assembly = PETSC_TRUE; 1120 ctx->norm_state = 0; 1121 ctx->electronShift = 0; 1122 ctx->M = NULL; 1123 ctx->J = NULL; 1124 /* geometry and grids */ 1125 ctx->sphere = PETSC_FALSE; 1126 ctx->use_p4est = PETSC_FALSE; 1127 ctx->simplex = PETSC_FALSE; 1128 for (PetscInt grid = 0; grid < LANDAU_MAX_GRIDS; grid++) { 1129 ctx->radius[grid] = 5.; /* thermal radius (velocity) */ 1130 ctx->radius_perp[grid] = 5.; /* thermal radius (velocity) */ 1131 ctx->radius_par[grid] = 5.; /* thermal radius (velocity) */ 1132 ctx->numAMRRefine[grid] = 0; 1133 ctx->postAMRRefine[grid] = 0; 1134 ctx->species_offset[grid + 1] = 1; // one species default 1135 num_species_grid[grid] = 0; 1136 ctx->plex[grid] = NULL; /* cache as expensive to Convert */ 1137 } 1138 ctx->species_offset[0] = 0; 1139 ctx->re_radius = 0.; 1140 ctx->vperp0_radius1 = 0; 1141 ctx->vperp0_radius2 = 0; 1142 ctx->nZRefine1 = 0; 1143 ctx->nZRefine2 = 0; 1144 ctx->numRERefine = 0; 1145 num_species_grid[0] = 1; // one species default 1146 /* species - [0] electrons, [1] one ion species eg, duetarium, [2] heavy impurity ion, ... */ 1147 ctx->charges[0] = -1; /* electron charge (MKS) */ 1148 ctx->masses[0] = 1 / 1835.469965278441013; /* temporary value in proton mass */ 1149 ctx->n[0] = 1; 1150 ctx->v_0 = 1; /* thermal velocity, we could start with a scale != 1 */ 1151 ctx->thermal_temps[0] = 1; 1152 /* constants, etc. */ 1153 ctx->epsilon0 = 8.8542e-12; /* permittivity of free space (MKS) F/m */ 1154 ctx->k = 1.38064852e-23; /* Boltzmann constant (MKS) J/K */ 1155 ctx->n_0 = 1.e20; /* typical plasma n, but could set it to 1 */ 1156 ctx->Ez = 0; 1157 for (PetscInt grid = 0; grid < LANDAU_NUM_TIMERS; grid++) ctx->times[grid] = 0; 1158 for (PetscInt ii = 0; ii < LANDAU_DIM; ii++) ctx->cells0[ii] = 2; 1159 if (LANDAU_DIM == 2) ctx->cells0[0] = 1; 1160 ctx->use_matrix_mass = PETSC_FALSE; 1161 ctx->use_relativistic_corrections = PETSC_FALSE; 1162 ctx->use_energy_tensor_trick = PETSC_FALSE; /* Use Eero's trick for energy conservation v --> grad(v^2/2) */ 1163 ctx->SData_d.w = NULL; 1164 ctx->SData_d.x = NULL; 1165 ctx->SData_d.y = NULL; 1166 ctx->SData_d.z = NULL; 1167 ctx->SData_d.invJ = NULL; 1168 ctx->jacobian_field_major_order = PETSC_FALSE; 1169 ctx->SData_d.coo_elem_offsets = NULL; 1170 ctx->SData_d.coo_elem_point_offsets = NULL; 1171 ctx->SData_d.coo_elem_fullNb = NULL; 1172 ctx->SData_d.coo_size = 0; 1173 PetscOptionsBegin(ctx->comm, prefix, "Options for Fokker-Plank-Landau collision operator", "none"); 1174 #if defined(PETSC_HAVE_KOKKOS) 1175 ctx->deviceType = LANDAU_KOKKOS; 1176 PetscCall(PetscStrncpy(ctx->filename, "kokkos", sizeof(ctx->filename))); 1177 #else 1178 ctx->deviceType = LANDAU_CPU; 1179 PetscCall(PetscStrncpy(ctx->filename, "cpu", sizeof(ctx->filename))); 1180 #endif 1181 PetscCall(PetscOptionsString("-dm_landau_device_type", "Use kernels on 'cpu' 'kokkos'", "plexland.c", ctx->filename, ctx->filename, sizeof(ctx->filename), NULL)); 1182 PetscCall(PetscStrcmp("cpu", ctx->filename, &flg)); 1183 if (flg) { 1184 ctx->deviceType = LANDAU_CPU; 1185 } else { 1186 PetscCall(PetscStrcmp("kokkos", ctx->filename, &flg)); 1187 if (flg) ctx->deviceType = LANDAU_KOKKOS; 1188 else SETERRQ(ctx->comm, PETSC_ERR_ARG_WRONG, "-dm_landau_device_type %s", ctx->filename); 1189 } 1190 ctx->filename[0] = '\0'; 1191 PetscCall(PetscOptionsString("-dm_landau_filename", "file to read mesh from", "plexland.c", ctx->filename, ctx->filename, sizeof(ctx->filename), &fileflg)); 1192 PetscCall(PetscOptionsReal("-dm_landau_electron_shift", "Shift in thermal velocity of electrons", "none", ctx->electronShift, &ctx->electronShift, NULL)); 1193 PetscCall(PetscOptionsInt("-dm_landau_verbose", "Level of verbosity output", "plexland.c", ctx->verbose, &ctx->verbose, NULL)); 1194 PetscCall(PetscOptionsInt("-dm_landau_batch_size", "Number of 'vertices' to batch", "ex2.c", ctx->batch_sz, &ctx->batch_sz, NULL)); 1195 PetscCheck(LANDAU_MAX_BATCH_SZ >= ctx->batch_sz, ctx->comm, PETSC_ERR_ARG_WRONG, "LANDAU_MAX_BATCH_SZ %d < ctx->batch_sz %" PetscInt_FMT, LANDAU_MAX_BATCH_SZ, ctx->batch_sz); 1196 PetscCall(PetscOptionsInt("-dm_landau_batch_view_idx", "Index of batch for diagnostics like plotting", "ex2.c", ctx->batch_view_idx, &ctx->batch_view_idx, NULL)); 1197 PetscCheck(ctx->batch_view_idx < ctx->batch_sz, ctx->comm, PETSC_ERR_ARG_WRONG, "-ctx->batch_view_idx %" PetscInt_FMT " > ctx->batch_sz %" PetscInt_FMT, ctx->batch_view_idx, ctx->batch_sz); 1198 PetscCall(PetscOptionsReal("-dm_landau_Ez", "Initial parallel electric field in unites of Conner-Hastie critical field", "plexland.c", ctx->Ez, &ctx->Ez, NULL)); 1199 PetscCall(PetscOptionsReal("-dm_landau_n_0", "Normalization constant for number density", "plexland.c", ctx->n_0, &ctx->n_0, NULL)); 1200 PetscCall(PetscOptionsBool("-dm_landau_use_mataxpy_mass", "Use fast but slightly fragile MATAXPY to add mass term", "plexland.c", ctx->use_matrix_mass, &ctx->use_matrix_mass, NULL)); 1201 PetscCall(PetscOptionsBool("-dm_landau_use_relativistic_corrections", "Use relativistic corrections", "plexland.c", ctx->use_relativistic_corrections, &ctx->use_relativistic_corrections, NULL)); 1202 PetscCall(PetscOptionsBool("-dm_landau_simplex", "Use simplex elements", "plexland.c", ctx->simplex, &ctx->simplex, NULL)); 1203 PetscCall(PetscOptionsBool("-dm_landau_sphere", "use sphere/semi-circle domain instead of rectangle", "plexland.c", ctx->sphere, &ctx->sphere, NULL)); 1204 if (LANDAU_DIM == 2 && ctx->use_relativistic_corrections) ctx->use_relativistic_corrections = PETSC_FALSE; // should warn 1205 PetscCall(PetscOptionsBool("-dm_landau_use_energy_tensor_trick", "Use Eero's trick of using grad(v^2/2) instead of v as args to Landau tensor to conserve energy with relativistic corrections and Q1 elements", "plexland.c", ctx->use_energy_tensor_trick, 1206 &ctx->use_energy_tensor_trick, NULL)); 1207 1208 /* get num species with temperature, set defaults */ 1209 for (ii = 1; ii < LANDAU_MAX_SPECIES; ii++) { 1210 ctx->thermal_temps[ii] = 1; 1211 ctx->charges[ii] = 1; 1212 ctx->masses[ii] = 1; 1213 ctx->n[ii] = 1; 1214 } 1215 nt = LANDAU_MAX_SPECIES; 1216 PetscCall(PetscOptionsRealArray("-dm_landau_thermal_temps", "Temperature of each species [e,i_0,i_1,...] in keV (must be set to set number of species)", "plexland.c", ctx->thermal_temps, &nt, &flg)); 1217 if (flg) { 1218 PetscCall(PetscInfo(dummy, "num_species set to number of thermal temps provided (%" PetscInt_FMT ")\n", nt)); 1219 ctx->num_species = nt; 1220 } else SETERRQ(ctx->comm, PETSC_ERR_ARG_WRONG, "-dm_landau_thermal_temps ,t1,t2,.. must be provided to set the number of species"); 1221 for (ii = 0; ii < ctx->num_species; ii++) ctx->thermal_temps[ii] *= 1.1604525e7; /* convert to Kelvin */ 1222 nm = LANDAU_MAX_SPECIES - 1; 1223 PetscCall(PetscOptionsRealArray("-dm_landau_ion_masses", "Mass of each species in units of proton mass [i_0=2,i_1=40...]", "plexland.c", &ctx->masses[1], &nm, &flg)); 1224 PetscCheck(!flg || nm == ctx->num_species - 1, ctx->comm, PETSC_ERR_ARG_WRONG, "num ion masses %" PetscInt_FMT " != num species %" PetscInt_FMT, nm, ctx->num_species - 1); 1225 nm = LANDAU_MAX_SPECIES; 1226 PetscCall(PetscOptionsRealArray("-dm_landau_n", "Number density of each species = n_s * n_0", "plexland.c", ctx->n, &nm, &flg)); 1227 PetscCheck(!flg || nm == ctx->num_species, ctx->comm, PETSC_ERR_ARG_WRONG, "wrong num n: %" PetscInt_FMT " != num species %" PetscInt_FMT, nm, ctx->num_species); 1228 for (ii = 0; ii < LANDAU_MAX_SPECIES; ii++) ctx->masses[ii] *= 1.6720e-27; /* scale by proton mass kg */ 1229 ctx->masses[0] = 9.10938356e-31; /* electron mass kg (should be about right already) */ 1230 nc = LANDAU_MAX_SPECIES - 1; 1231 PetscCall(PetscOptionsRealArray("-dm_landau_ion_charges", "Charge of each species in units of proton charge [i_0=2,i_1=18,...]", "plexland.c", &ctx->charges[1], &nc, &flg)); 1232 if (flg) PetscCheck(nc == ctx->num_species - 1, ctx->comm, PETSC_ERR_ARG_WRONG, "num charges %" PetscInt_FMT " != num species %" PetscInt_FMT, nc, ctx->num_species - 1); 1233 for (ii = 0; ii < LANDAU_MAX_SPECIES; ii++) ctx->charges[ii] *= 1.6022e-19; /* electron/proton charge (MKS) */ 1234 /* geometry and grids */ 1235 nt = LANDAU_MAX_GRIDS; 1236 PetscCall(PetscOptionsIntArray("-dm_landau_num_species_grid", "Number of species on each grid: [ 1, ....] or [S, 0 ....] for single grid", "plexland.c", num_species_grid, &nt, &flg)); 1237 if (flg) { 1238 ctx->num_grids = nt; 1239 for (ii = nt = 0; ii < ctx->num_grids; ii++) nt += num_species_grid[ii]; 1240 PetscCheck(ctx->num_species == nt, ctx->comm, PETSC_ERR_ARG_WRONG, "-dm_landau_num_species_grid: sum %" PetscInt_FMT " != num_species = %" PetscInt_FMT ". %" PetscInt_FMT " grids (check that number of grids <= LANDAU_MAX_GRIDS = %d)", nt, ctx->num_species, 1241 ctx->num_grids, LANDAU_MAX_GRIDS); 1242 } else { 1243 if (ctx->num_species > LANDAU_MAX_GRIDS) { 1244 num_species_grid[0] = 1; 1245 num_species_grid[1] = ctx->num_species - 1; 1246 ctx->num_grids = 2; 1247 } else { 1248 ctx->num_grids = ctx->num_species; 1249 for (ii = 0; ii < ctx->num_grids; ii++) num_species_grid[ii] = 1; 1250 } 1251 } 1252 for (ctx->species_offset[0] = ii = 0; ii < ctx->num_grids; ii++) ctx->species_offset[ii + 1] = ctx->species_offset[ii] + num_species_grid[ii]; 1253 PetscCheck(ctx->species_offset[ctx->num_grids] == ctx->num_species, ctx->comm, PETSC_ERR_ARG_WRONG, "ctx->species_offset[ctx->num_grids] %" PetscInt_FMT " != ctx->num_species = %" PetscInt_FMT " ???????????", ctx->species_offset[ctx->num_grids], 1254 ctx->num_species); 1255 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1256 PetscInt iii = ctx->species_offset[grid]; // normalize with first (arbitrary) species on grid 1257 ctx->thermal_speed[grid] = PetscSqrtReal(ctx->k * ctx->thermal_temps[iii] / ctx->masses[iii]); /* arbitrary units for non-dimensionalization: plasma formulary def */ 1258 } 1259 // get lambdas here because we need them for t_0 etc 1260 PetscCall(PetscOptionsReal("-dm_landau_ln_lambda", "Universal cross section parameter. Default uses NRL formulas", "plexland.c", lnLam, &lnLam, &flg)); 1261 if (flg) { 1262 for (PetscInt grid = 0; grid < LANDAU_MAX_GRIDS; grid++) { 1263 for (PetscInt gridj = 0; gridj < LANDAU_MAX_GRIDS; gridj++) ctx->lambdas[gridj][grid] = lnLam; /* cross section ratio large - small angle collisions */ 1264 } 1265 } else { 1266 PetscCall(makeLambdas(ctx)); 1267 } 1268 non_dim_grid = 0; 1269 PetscCall(PetscOptionsInt("-dm_landau_normalization_grid", "Index of grid to use for setting v_0, m_0, t_0. (Not recommended)", "plexland.c", non_dim_grid, &non_dim_grid, &flg)); 1270 if (non_dim_grid != 0) PetscCall(PetscInfo(dummy, "Normalization grid set to %" PetscInt_FMT ", but non-default not well verified\n", non_dim_grid)); 1271 PetscCheck(non_dim_grid >= 0 && non_dim_grid < ctx->num_species, ctx->comm, PETSC_ERR_ARG_WRONG, "Normalization grid wrong: %" PetscInt_FMT, non_dim_grid); 1272 ctx->v_0 = ctx->thermal_speed[non_dim_grid]; /* arbitrary units for non dimensionalization: global mean velocity in 1D of electrons */ 1273 ctx->m_0 = ctx->masses[non_dim_grid]; /* arbitrary reference mass, electrons */ 1274 ctx->t_0 = 8 * PETSC_PI * PetscSqr(ctx->epsilon0 * ctx->m_0 / PetscSqr(ctx->charges[non_dim_grid])) / ctx->lambdas[non_dim_grid][non_dim_grid] / ctx->n_0 * PetscPowReal(ctx->v_0, 3); /* note, this t_0 makes nu[non_dim_grid,non_dim_grid]=1 */ 1275 /* domain */ 1276 nt = LANDAU_MAX_GRIDS; 1277 PetscCall(PetscOptionsRealArray("-dm_landau_domain_radius", "Phase space size in units of thermal velocity of grid", "plexland.c", ctx->radius, &nt, &flg)); 1278 if (flg) { 1279 PetscCheck(nt >= ctx->num_grids, ctx->comm, PETSC_ERR_ARG_WRONG, "-dm_landau_domain_radius: given %" PetscInt_FMT " radius != number grids %" PetscInt_FMT, nt, ctx->num_grids); 1280 while (nt--) ctx->radius_par[nt] = ctx->radius_perp[nt] = ctx->radius[nt]; 1281 } else { 1282 nt = LANDAU_MAX_GRIDS; 1283 PetscCall(PetscOptionsRealArray("-dm_landau_domain_max_par", "Parallel velocity domain size in units of thermal velocity of grid", "plexland.c", ctx->radius_par, &nt, &flg)); 1284 if (flg) PetscCheck(nt >= ctx->num_grids, ctx->comm, PETSC_ERR_ARG_WRONG, "-dm_landau_domain_max_par: given %" PetscInt_FMT " radius != number grids %" PetscInt_FMT, nt, ctx->num_grids); 1285 PetscCall(PetscOptionsRealArray("-dm_landau_domain_max_perp", "Perpendicular velocity domain size in units of thermal velocity of grid", "plexland.c", ctx->radius_perp, &nt, &flg)); 1286 if (flg) PetscCheck(nt >= ctx->num_grids, ctx->comm, PETSC_ERR_ARG_WRONG, "-dm_landau_domain_max_perp: given %" PetscInt_FMT " radius != number grids %" PetscInt_FMT, nt, ctx->num_grids); 1287 } 1288 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1289 if (flg && ctx->radius[grid] <= 0) { /* negative is ratio of c - need to set par and perp with this -- todo */ 1290 if (ctx->radius[grid] == 0) ctx->radius[grid] = 0.75; 1291 else ctx->radius[grid] = -ctx->radius[grid]; 1292 ctx->radius[grid] = ctx->radius[grid] * SPEED_OF_LIGHT / ctx->v_0; // use any species on grid to normalize (v_0 same for all on grid) 1293 PetscCall(PetscInfo(dummy, "Change domain radius to %g for grid %" PetscInt_FMT "\n", (double)ctx->radius[grid], grid)); 1294 } 1295 ctx->radius[grid] *= ctx->thermal_speed[grid] / ctx->v_0; // scale domain by thermal radius relative to v_0 1296 ctx->radius_perp[grid] *= ctx->thermal_speed[grid] / ctx->v_0; // scale domain by thermal radius relative to v_0 1297 ctx->radius_par[grid] *= ctx->thermal_speed[grid] / ctx->v_0; // scale domain by thermal radius relative to v_0 1298 } 1299 /* amr parameters */ 1300 if (!fileflg) { 1301 nt = LANDAU_MAX_GRIDS; 1302 PetscCall(PetscOptionsIntArray("-dm_landau_amr_levels_max", "Number of AMR levels of refinement around origin, after (RE) refinements along z", "plexland.c", ctx->numAMRRefine, &nt, &flg)); 1303 PetscCheck(!flg || nt >= ctx->num_grids, ctx->comm, PETSC_ERR_ARG_WRONG, "-dm_landau_amr_levels_max: given %" PetscInt_FMT " != number grids %" PetscInt_FMT, nt, ctx->num_grids); 1304 nt = LANDAU_MAX_GRIDS; 1305 PetscCall(PetscOptionsIntArray("-dm_landau_amr_post_refine", "Number of levels to uniformly refine after AMR", "plexland.c", ctx->postAMRRefine, &nt, &flg)); 1306 for (ii = 1; ii < ctx->num_grids; ii++) ctx->postAMRRefine[ii] = ctx->postAMRRefine[0]; // all grids the same now 1307 PetscCall(PetscOptionsInt("-dm_landau_amr_re_levels", "Number of levels to refine along v_perp=0, z>0", "plexland.c", ctx->numRERefine, &ctx->numRERefine, &flg)); 1308 PetscCall(PetscOptionsInt("-dm_landau_amr_z_refine_pre", "Number of levels to refine along v_perp=0 before origin refine", "plexland.c", ctx->nZRefine1, &ctx->nZRefine1, &flg)); 1309 PetscCall(PetscOptionsInt("-dm_landau_amr_z_refine_post", "Number of levels to refine along v_perp=0 after origin refine", "plexland.c", ctx->nZRefine2, &ctx->nZRefine2, &flg)); 1310 PetscCall(PetscOptionsReal("-dm_landau_re_radius", "velocity range to refine on positive (z>0) r=0 axis for runaways", "plexland.c", ctx->re_radius, &ctx->re_radius, &flg)); 1311 PetscCall(PetscOptionsReal("-dm_landau_z_radius_pre", "velocity range to refine r=0 axis (for electrons)", "plexland.c", ctx->vperp0_radius1, &ctx->vperp0_radius1, &flg)); 1312 PetscCall(PetscOptionsReal("-dm_landau_z_radius_post", "velocity range to refine r=0 axis (for electrons) after origin AMR", "plexland.c", ctx->vperp0_radius2, &ctx->vperp0_radius2, &flg)); 1313 /* spherical domain */ 1314 if (ctx->sphere || ctx->simplex) { 1315 ctx->sphere_uniform_normal = PETSC_FALSE; 1316 PetscCall(PetscOptionsBool("-dm_landau_sphere_uniform_normal", "Scaling of circle radius to get uniform particles per cell with Maxwellians (not used)", "plexland.c", ctx->sphere_uniform_normal, &ctx->sphere_uniform_normal, NULL)); 1317 if (!ctx->sphere_uniform_normal) { // true 1318 nt = LANDAU_MAX_GRIDS; 1319 PetscCall(PetscOptionsRealArray("-dm_landau_sphere_inner_radius_90degree_scale", "Scaling of radius for inner circle on 90 degree grid", "plexland.c", ctx->sphere_inner_radius_90degree, &nt, &flg)); 1320 if (flg && nt < ctx->num_grids) { 1321 for (PetscInt grid = nt; grid < ctx->num_grids; grid++) ctx->sphere_inner_radius_90degree[grid] = ctx->sphere_inner_radius_90degree[0]; 1322 } else if (!flg || nt == 0) { 1323 if (LANDAU_DIM == 2) { 1324 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) ctx->sphere_inner_radius_90degree[grid] = 0.4; // optimized for R=5, Q4, AMR=0 1325 } else { 1326 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) ctx->sphere_inner_radius_90degree[grid] = 0.577 * 0.40; 1327 } 1328 } 1329 nt = LANDAU_MAX_GRIDS; 1330 PetscCall(PetscOptionsRealArray("-dm_landau_sphere_inner_radius_45degree_scale", "Scaling of radius for inner circle on 45 degree grid", "plexland.c", ctx->sphere_inner_radius_45degree, &nt, &flg)); 1331 if (flg && nt < ctx->num_grids) { 1332 for (PetscInt grid = nt; grid < ctx->num_grids; grid++) ctx->sphere_inner_radius_45degree[grid] = ctx->sphere_inner_radius_45degree[0]; 1333 } else if (!flg || nt == 0) { 1334 if (LANDAU_DIM == 2) { 1335 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) ctx->sphere_inner_radius_45degree[grid] = 0.45; // optimized for R=5, Q4, AMR=0 1336 } else { 1337 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) ctx->sphere_inner_radius_45degree[grid] = 0.4; // 3D sphere 1338 } 1339 } 1340 if (ctx->sphere) PetscCall(PetscInfo(ctx->plex[0], "sphere : , 45 degree scaling = %g; 90 degree scaling = %g\n", (double)ctx->sphere_inner_radius_45degree[0], (double)ctx->sphere_inner_radius_90degree[0])); 1341 } else { 1342 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1343 switch (ctx->numAMRRefine[grid]) { 1344 case 0: 1345 case 1: 1346 case 2: 1347 case 3: 1348 default: 1349 if (LANDAU_DIM == 2) { 1350 ctx->sphere_inner_radius_90degree[grid] = 0.40; 1351 ctx->sphere_inner_radius_45degree[grid] = 0.45; 1352 } else { 1353 ctx->sphere_inner_radius_45degree[grid] = 0.25; 1354 } 1355 } 1356 } 1357 } 1358 } else { 1359 nt = LANDAU_DIM; 1360 PetscCall(PetscOptionsIntArray("-dm_landau_num_cells", "Number of cells in each dimension of base grid", "plexland.c", ctx->cells0, &nt, &flg)); 1361 } 1362 } 1363 /* processing options */ 1364 PetscCall(PetscOptionsBool("-dm_landau_gpu_assembly", "Assemble Jacobian on GPU", "plexland.c", ctx->gpu_assembly, &ctx->gpu_assembly, NULL)); 1365 PetscCall(PetscOptionsBool("-dm_landau_jacobian_field_major_order", "Reorder Jacobian for GPU assembly with field major, or block diagonal, ordering (DEPRECATED)", "plexland.c", ctx->jacobian_field_major_order, &ctx->jacobian_field_major_order, NULL)); 1366 if (ctx->jacobian_field_major_order) PetscCheck(ctx->gpu_assembly, ctx->comm, PETSC_ERR_ARG_WRONG, "-dm_landau_jacobian_field_major_order requires -dm_landau_gpu_assembly"); 1367 PetscCheck(!ctx->jacobian_field_major_order, ctx->comm, PETSC_ERR_ARG_WRONG, "-dm_landau_jacobian_field_major_order DEPRECATED"); 1368 PetscOptionsEnd(); 1369 1370 for (ii = ctx->num_species; ii < LANDAU_MAX_SPECIES; ii++) ctx->masses[ii] = ctx->thermal_temps[ii] = ctx->charges[ii] = 0; 1371 if (ctx->verbose != 0) { 1372 PetscReal pmassunit = PetscRealConstant(1.6720e-27); 1373 1374 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "masses: e=%10.3e; ions in proton mass units: %10.3e %10.3e ...\n", (double)ctx->masses[0], (double)(ctx->masses[1] / pmassunit), (double)(ctx->num_species > 2 ? ctx->masses[2] / pmassunit : 0))); 1375 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "charges: e=%10.3e; charges in elementary units: %10.3e %10.3e\n", (double)ctx->charges[0], (double)(-ctx->charges[1] / ctx->charges[0]), (double)(ctx->num_species > 2 ? -ctx->charges[2] / ctx->charges[0] : 0))); 1376 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "n: e: %10.3e i: %10.3e %10.3e\n", (double)ctx->n[0], (double)ctx->n[1], (double)(ctx->num_species > 2 ? ctx->n[2] : 0))); 1377 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "thermal T (K): e=%10.3e i=%10.3e %10.3e. Normalization grid %" PetscInt_FMT ": v_0=%10.3e (%10.3ec) n_0=%10.3e t_0=%10.3e %" PetscInt_FMT " batched, view batch %" PetscInt_FMT "\n", (double)ctx->thermal_temps[0], 1378 (double)ctx->thermal_temps[1], (double)((ctx->num_species > 2) ? ctx->thermal_temps[2] : 0), non_dim_grid, (double)ctx->v_0, (double)(ctx->v_0 / SPEED_OF_LIGHT), (double)ctx->n_0, (double)ctx->t_0, ctx->batch_sz, ctx->batch_view_idx)); 1379 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Domain radius (AMR levels) grid %d: par=%10.3e perp=%10.3e (%" PetscInt_FMT ") ", 0, (double)ctx->radius_par[0], (double)ctx->radius_perp[0], ctx->numAMRRefine[0])); 1380 for (ii = 1; ii < ctx->num_grids; ii++) PetscCall(PetscPrintf(PETSC_COMM_WORLD, ", %" PetscInt_FMT ": par=%10.3e perp=%10.3e (%" PetscInt_FMT ") ", ii, (double)ctx->radius_par[ii], (double)ctx->radius_perp[ii], ctx->numAMRRefine[ii])); 1381 if (ctx->use_relativistic_corrections) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nUse relativistic corrections\n")); 1382 else PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n")); 1383 } 1384 PetscCall(DMDestroy(&dummy)); 1385 { 1386 PetscMPIInt rank; 1387 PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank)); 1388 ctx->stage = 0; 1389 PetscCall(PetscLogEventRegister("Landau Create", DM_CLASSID, &ctx->events[13])); /* 13 */ 1390 PetscCall(PetscLogEventRegister(" GPU ass. setup", DM_CLASSID, &ctx->events[2])); /* 2 */ 1391 PetscCall(PetscLogEventRegister(" Build matrix", DM_CLASSID, &ctx->events[12])); /* 12 */ 1392 PetscCall(PetscLogEventRegister(" Assembly maps", DM_CLASSID, &ctx->events[15])); /* 15 */ 1393 PetscCall(PetscLogEventRegister("Landau Mass mat", DM_CLASSID, &ctx->events[14])); /* 14 */ 1394 PetscCall(PetscLogEventRegister("Landau Operator", DM_CLASSID, &ctx->events[11])); /* 11 */ 1395 PetscCall(PetscLogEventRegister("Landau Jacobian", DM_CLASSID, &ctx->events[0])); /* 0 */ 1396 PetscCall(PetscLogEventRegister("Landau Mass", DM_CLASSID, &ctx->events[9])); /* 9 */ 1397 PetscCall(PetscLogEventRegister(" Preamble", DM_CLASSID, &ctx->events[10])); /* 10 */ 1398 PetscCall(PetscLogEventRegister(" static IP Data", DM_CLASSID, &ctx->events[7])); /* 7 */ 1399 PetscCall(PetscLogEventRegister(" dynamic IP-Jac", DM_CLASSID, &ctx->events[1])); /* 1 */ 1400 PetscCall(PetscLogEventRegister(" Kernel-init", DM_CLASSID, &ctx->events[3])); /* 3 */ 1401 PetscCall(PetscLogEventRegister(" Jac-f-df (GPU)", DM_CLASSID, &ctx->events[8])); /* 8 */ 1402 PetscCall(PetscLogEventRegister(" J Kernel (GPU)", DM_CLASSID, &ctx->events[4])); /* 4 */ 1403 PetscCall(PetscLogEventRegister(" M Kernel (GPU)", DM_CLASSID, &ctx->events[16])); /* 16 */ 1404 PetscCall(PetscLogEventRegister(" Copy to CPU", DM_CLASSID, &ctx->events[5])); /* 5 */ 1405 PetscCall(PetscLogEventRegister(" CPU assemble", DM_CLASSID, &ctx->events[6])); /* 6 */ 1406 1407 if (rank) { /* turn off output stuff for duplicate runs - do we need to add the prefix to all this? */ 1408 PetscCall(PetscOptionsClearValue(NULL, "-snes_converged_reason")); 1409 PetscCall(PetscOptionsClearValue(NULL, "-ksp_converged_reason")); 1410 PetscCall(PetscOptionsClearValue(NULL, "-snes_monitor")); 1411 PetscCall(PetscOptionsClearValue(NULL, "-ksp_monitor")); 1412 PetscCall(PetscOptionsClearValue(NULL, "-ts_monitor")); 1413 PetscCall(PetscOptionsClearValue(NULL, "-ts_view")); 1414 PetscCall(PetscOptionsClearValue(NULL, "-ts_adapt_monitor")); 1415 PetscCall(PetscOptionsClearValue(NULL, "-dm_landau_amr_dm_view")); 1416 PetscCall(PetscOptionsClearValue(NULL, "-dm_landau_amr_vec_view")); 1417 PetscCall(PetscOptionsClearValue(NULL, "-dm_landau_mass_dm_view")); 1418 PetscCall(PetscOptionsClearValue(NULL, "-dm_landau_mass_view")); 1419 PetscCall(PetscOptionsClearValue(NULL, "-dm_landau_jacobian_view")); 1420 PetscCall(PetscOptionsClearValue(NULL, "-dm_landau_mat_view")); 1421 PetscCall(PetscOptionsClearValue(NULL, "-pc_bjkokkos_ksp_converged_reason")); 1422 PetscCall(PetscOptionsClearValue(NULL, "-pc_bjkokkos_ksp_monitor")); 1423 PetscCall(PetscOptionsClearValue(NULL, "-")); 1424 PetscCall(PetscOptionsClearValue(NULL, "-info")); 1425 } 1426 } 1427 PetscFunctionReturn(PETSC_SUCCESS); 1428 } 1429 1430 static PetscErrorCode CreateStaticData(PetscInt dim, IS grid_batch_is_inv[], LandauCtx *ctx) 1431 { 1432 PetscSection section[LANDAU_MAX_GRIDS], globsection[LANDAU_MAX_GRIDS]; 1433 PetscQuadrature quad; 1434 const PetscReal *quadWeights; 1435 PetscReal invMass[LANDAU_MAX_SPECIES], nu_alpha[LANDAU_MAX_SPECIES], nu_beta[LANDAU_MAX_SPECIES]; 1436 PetscInt numCells[LANDAU_MAX_GRIDS], Nq, Nb, Nf[LANDAU_MAX_GRIDS], ncellsTot = 0, MAP_BF_SIZE = 64 * LANDAU_DIM * LANDAU_DIM * LANDAU_MAX_Q_FACE * LANDAU_MAX_SPECIES; 1437 PetscTabulation *Tf; 1438 PetscDS prob; 1439 1440 PetscFunctionBegin; 1441 PetscCall(PetscFEGetDimension(ctx->fe[0], &Nb)); 1442 PetscCheck(Nb <= LANDAU_MAX_NQND, ctx->comm, PETSC_ERR_ARG_WRONG, "Order too high. Nb = %" PetscInt_FMT " > LANDAU_MAX_NQND (%d)", Nb, LANDAU_MAX_NQND); 1443 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1444 for (PetscInt ii = ctx->species_offset[grid]; ii < ctx->species_offset[grid + 1]; ii++) { 1445 invMass[ii] = ctx->m_0 / ctx->masses[ii]; 1446 nu_alpha[ii] = PetscSqr(ctx->charges[ii] / ctx->m_0) * ctx->m_0 / ctx->masses[ii]; 1447 nu_beta[ii] = PetscSqr(ctx->charges[ii] / ctx->epsilon0) / (8 * PETSC_PI) * ctx->t_0 * ctx->n_0 / PetscPowReal(ctx->v_0, 3); 1448 } 1449 } 1450 if (ctx->verbose == 4) { 1451 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "nu_alpha: ")); 1452 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1453 PetscInt iii = ctx->species_offset[grid]; 1454 for (PetscInt ii = iii; ii < ctx->species_offset[grid + 1]; ii++) PetscCall(PetscPrintf(PETSC_COMM_WORLD, " %e", (double)nu_alpha[ii])); 1455 } 1456 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nnu_beta: ")); 1457 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1458 PetscInt iii = ctx->species_offset[grid]; 1459 for (PetscInt ii = iii; ii < ctx->species_offset[grid + 1]; ii++) PetscCall(PetscPrintf(PETSC_COMM_WORLD, " %e", (double)nu_beta[ii])); 1460 } 1461 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\nnu_alpha[i]*nu_beta[j]*lambda[i][j]:\n")); 1462 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1463 PetscInt iii = ctx->species_offset[grid]; 1464 for (PetscInt ii = iii; ii < ctx->species_offset[grid + 1]; ii++) { 1465 for (PetscInt gridj = 0; gridj < ctx->num_grids; gridj++) { 1466 PetscInt jjj = ctx->species_offset[gridj]; 1467 for (PetscInt jj = jjj; jj < ctx->species_offset[gridj + 1]; jj++) PetscCall(PetscPrintf(PETSC_COMM_WORLD, " %14.9e", (double)(nu_alpha[ii] * nu_beta[jj] * ctx->lambdas[grid][gridj]))); 1468 } 1469 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n")); 1470 } 1471 } 1472 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "lambda[i][j]:\n")); 1473 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1474 PetscInt iii = ctx->species_offset[grid]; 1475 for (PetscInt ii = iii; ii < ctx->species_offset[grid + 1]; ii++) { 1476 for (PetscInt gridj = 0; gridj < ctx->num_grids; gridj++) { 1477 PetscInt jjj = ctx->species_offset[gridj]; 1478 for (PetscInt jj = jjj; jj < ctx->species_offset[gridj + 1]; jj++) PetscCall(PetscPrintf(PETSC_COMM_WORLD, " %14.9e", (double)ctx->lambdas[grid][gridj])); 1479 } 1480 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n")); 1481 } 1482 } 1483 } 1484 PetscCall(DMGetDS(ctx->plex[0], &prob)); // same DS for all grids 1485 PetscCall(PetscDSGetTabulation(prob, &Tf)); // Bf, &Df same for all grids 1486 /* DS, Tab and quad is same on all grids */ 1487 PetscCheck(ctx->plex[0], ctx->comm, PETSC_ERR_ARG_WRONG, "Plex not created"); 1488 PetscCall(PetscFEGetQuadrature(ctx->fe[0], &quad)); 1489 PetscCall(PetscQuadratureGetData(quad, NULL, NULL, &Nq, NULL, &quadWeights)); 1490 PetscCheck(Nq <= LANDAU_MAX_NQND, ctx->comm, PETSC_ERR_ARG_WRONG, "Order too high. Nq = %" PetscInt_FMT " > LANDAU_MAX_NQND (%d)", Nq, LANDAU_MAX_NQND); 1491 /* setup each grid */ 1492 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1493 PetscInt cStart, cEnd; 1494 PetscCheck(ctx->plex[grid] != NULL, ctx->comm, PETSC_ERR_ARG_WRONG, "Plex not created"); 1495 PetscCall(DMPlexGetHeightStratum(ctx->plex[grid], 0, &cStart, &cEnd)); 1496 numCells[grid] = cEnd - cStart; // grids can have different topology 1497 PetscCall(DMGetLocalSection(ctx->plex[grid], §ion[grid])); 1498 PetscCall(DMGetGlobalSection(ctx->plex[grid], &globsection[grid])); 1499 PetscCall(PetscSectionGetNumFields(section[grid], &Nf[grid])); 1500 ncellsTot += numCells[grid]; 1501 } 1502 /* create GPU assembly data */ 1503 if (ctx->gpu_assembly) { /* we need GPU object with GPU assembly */ 1504 PetscContainer container; 1505 PetscScalar *elemMatrix, *elMat; 1506 pointInterpolationP4est(*pointMaps)[LANDAU_MAX_Q_FACE]; 1507 P4estVertexMaps *maps; 1508 const PetscInt *plex_batch = NULL, elMatSz = Nb * Nb * ctx->num_species * ctx->num_species; 1509 LandauIdx *coo_elem_offsets = NULL, *coo_elem_fullNb = NULL, (*coo_elem_point_offsets)[LANDAU_MAX_NQND + 1] = NULL; 1510 /* create GPU assembly data */ 1511 PetscCall(PetscInfo(ctx->plex[0], "Make GPU maps %d\n", 1)); 1512 PetscCall(PetscLogEventBegin(ctx->events[2], 0, 0, 0, 0)); 1513 PetscCall(PetscMalloc(sizeof(*maps) * ctx->num_grids, &maps)); 1514 PetscCall(PetscMalloc(sizeof(*pointMaps) * MAP_BF_SIZE, &pointMaps)); 1515 PetscCall(PetscMalloc(sizeof(*elemMatrix) * elMatSz, &elemMatrix)); 1516 1517 { // setup COO assembly -- put COO metadata directly in ctx->SData_d 1518 PetscCall(PetscMalloc3(ncellsTot + 1, &coo_elem_offsets, ncellsTot, &coo_elem_fullNb, ncellsTot, &coo_elem_point_offsets)); // array of integer pointers 1519 coo_elem_offsets[0] = 0; // finish later 1520 PetscCall(PetscInfo(ctx->plex[0], "COO initialization, %" PetscInt_FMT " cells\n", ncellsTot)); 1521 ctx->SData_d.coo_n_cellsTot = ncellsTot; 1522 ctx->SData_d.coo_elem_offsets = (void *)coo_elem_offsets; 1523 ctx->SData_d.coo_elem_fullNb = (void *)coo_elem_fullNb; 1524 ctx->SData_d.coo_elem_point_offsets = (void *)coo_elem_point_offsets; 1525 } 1526 1527 ctx->SData_d.coo_max_fullnb = 0; 1528 for (PetscInt grid = 0, glb_elem_idx = 0; grid < ctx->num_grids; grid++) { 1529 PetscInt cStart, cEnd, Nfloc = Nf[grid], totDim = Nfloc * Nb; 1530 if (grid_batch_is_inv[grid]) PetscCall(ISGetIndices(grid_batch_is_inv[grid], &plex_batch)); 1531 PetscCheck(!plex_batch, ctx->comm, PETSC_ERR_ARG_WRONG, "-dm_landau_jacobian_field_major_order DEPRECATED"); 1532 PetscCall(DMPlexGetHeightStratum(ctx->plex[grid], 0, &cStart, &cEnd)); 1533 // make maps 1534 maps[grid].d_self = NULL; 1535 maps[grid].num_elements = numCells[grid]; 1536 maps[grid].num_face = (PetscInt)(pow(Nq, 1. / ((double)dim)) + .001); // Q 1537 maps[grid].num_face = (PetscInt)(pow(maps[grid].num_face, (double)(dim - 1)) + .001); // Q^2 1538 maps[grid].num_reduced = 0; 1539 maps[grid].deviceType = ctx->deviceType; 1540 maps[grid].numgrids = ctx->num_grids; 1541 // count reduced and get 1542 PetscCall(PetscMalloc(maps[grid].num_elements * sizeof(*maps[grid].gIdx), &maps[grid].gIdx)); 1543 for (PetscInt ej = cStart, eidx = 0; ej < cEnd; ++ej, ++eidx, glb_elem_idx++) { 1544 if (coo_elem_offsets) coo_elem_offsets[glb_elem_idx + 1] = coo_elem_offsets[glb_elem_idx]; // start with last one, then add 1545 for (PetscInt fieldA = 0; fieldA < Nf[grid]; fieldA++) { 1546 PetscInt fullNb = 0; 1547 for (PetscInt q = 0; q < Nb; ++q) { 1548 PetscInt numindices, *indices; 1549 PetscScalar *valuesOrig = elMat = elemMatrix; 1550 PetscCall(PetscArrayzero(elMat, totDim * totDim)); 1551 elMat[(fieldA * Nb + q) * totDim + fieldA * Nb + q] = 1; 1552 PetscCall(DMPlexGetClosureIndices(ctx->plex[grid], section[grid], globsection[grid], ej, PETSC_TRUE, &numindices, &indices, NULL, &elMat)); 1553 if (ctx->simplex) { 1554 PetscCheck(numindices == Nb, ctx->comm, PETSC_ERR_ARG_WRONG, "numindices != Nb numindices=%" PetscInt_FMT " Nb=%" PetscInt_FMT, numindices, Nb); 1555 for (PetscInt q = 0; q < numindices; ++q) { maps[grid].gIdx[eidx][fieldA][q] = indices[q]; } 1556 fullNb++; 1557 } else { 1558 for (PetscInt f = 0; f < numindices; ++f) { // look for a non-zero on the diagonal (is this too complicated for simplices?) 1559 if (PetscAbs(PetscRealPart(elMat[f * numindices + f])) > PETSC_MACHINE_EPSILON) { 1560 // found it 1561 if (PetscAbs(PetscRealPart(elMat[f * numindices + f] - 1.)) < PETSC_MACHINE_EPSILON) { // normal vertex 1.0 1562 if (plex_batch) { 1563 maps[grid].gIdx[eidx][fieldA][q] = plex_batch[indices[f]]; 1564 } else { 1565 maps[grid].gIdx[eidx][fieldA][q] = indices[f]; 1566 } 1567 fullNb++; 1568 } else { //found a constraint 1569 PetscInt jj = 0; 1570 PetscReal sum = 0; 1571 const PetscInt ff = f; 1572 maps[grid].gIdx[eidx][fieldA][q] = -maps[grid].num_reduced - 1; // store (-)index: id = -(idx+1): idx = -id - 1 1573 PetscCheck(!ctx->simplex, ctx->comm, PETSC_ERR_ARG_WRONG, "No constraints with simplex"); 1574 do { // constraints are continuous in Plex - exploit that here 1575 PetscInt ii; // get 'scale' 1576 for (ii = 0, pointMaps[maps[grid].num_reduced][jj].scale = 0; ii < maps[grid].num_face; ii++) { // sum row of outer product to recover vector value 1577 if (ff + ii < numindices) { // 3D has Q and Q^2 interps so might run off end. We could test that elMat[f*numindices + ff + ii] > 0, and break if not 1578 pointMaps[maps[grid].num_reduced][jj].scale += PetscRealPart(elMat[f * numindices + ff + ii]); 1579 } 1580 } 1581 sum += pointMaps[maps[grid].num_reduced][jj].scale; // diagnostic 1582 // get 'gid' 1583 if (pointMaps[maps[grid].num_reduced][jj].scale == 0) pointMaps[maps[grid].num_reduced][jj].gid = -1; // 3D has Q and Q^2 interps 1584 else { 1585 if (plex_batch) { 1586 pointMaps[maps[grid].num_reduced][jj].gid = plex_batch[indices[f]]; 1587 } else { 1588 pointMaps[maps[grid].num_reduced][jj].gid = indices[f]; 1589 } 1590 fullNb++; 1591 } 1592 } while (++jj < maps[grid].num_face && ++f < numindices); // jj is incremented if we hit the end 1593 while (jj < maps[grid].num_face) { 1594 pointMaps[maps[grid].num_reduced][jj].scale = 0; 1595 pointMaps[maps[grid].num_reduced][jj].gid = -1; 1596 jj++; 1597 } 1598 if (PetscAbs(sum - 1.0) > 10 * PETSC_MACHINE_EPSILON) { // debug 1599 PetscInt d, f; 1600 PetscReal tmp = 0; 1601 PetscCall( 1602 PetscPrintf(PETSC_COMM_SELF, "\t\t%" PetscInt_FMT ".%" PetscInt_FMT ".%" PetscInt_FMT ") ERROR total I = %22.16e (LANDAU_MAX_Q_FACE=%d, #face=%" PetscInt_FMT ")\n", eidx, q, fieldA, (double)sum, LANDAU_MAX_Q_FACE, maps[grid].num_face)); 1603 for (d = 0, tmp = 0; d < numindices; ++d) { 1604 if (tmp != 0 && PetscAbs(tmp - 1.0) > 10 * PETSC_MACHINE_EPSILON) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%3" PetscInt_FMT ") %3" PetscInt_FMT ": ", d, indices[d])); 1605 for (f = 0; f < numindices; ++f) tmp += PetscRealPart(elMat[d * numindices + f]); 1606 if (tmp != 0) PetscCall(PetscPrintf(ctx->comm, " | %22.16e\n", (double)tmp)); 1607 } 1608 } 1609 maps[grid].num_reduced++; 1610 PetscCheck(maps[grid].num_reduced < MAP_BF_SIZE, PETSC_COMM_SELF, PETSC_ERR_PLIB, "maps[grid].num_reduced %" PetscInt_FMT " > %" PetscInt_FMT, maps[grid].num_reduced, MAP_BF_SIZE); 1611 } 1612 break; 1613 } 1614 } 1615 } // !simplex 1616 // cleanup 1617 PetscCall(DMPlexRestoreClosureIndices(ctx->plex[grid], section[grid], globsection[grid], ej, PETSC_TRUE, &numindices, &indices, NULL, &elMat)); 1618 if (elMat != valuesOrig) PetscCall(DMRestoreWorkArray(ctx->plex[grid], numindices * numindices, MPIU_SCALAR, &elMat)); 1619 } 1620 { // setup COO assembly 1621 coo_elem_offsets[glb_elem_idx + 1] += fullNb * fullNb; // one species block, adds a block for each species, on this element in this grid 1622 if (fieldA == 0) { // cache full Nb for this element, on this grid per species 1623 coo_elem_fullNb[glb_elem_idx] = fullNb; 1624 if (fullNb > ctx->SData_d.coo_max_fullnb) ctx->SData_d.coo_max_fullnb = fullNb; 1625 } else PetscCheck(coo_elem_fullNb[glb_elem_idx] == fullNb, PETSC_COMM_SELF, PETSC_ERR_PLIB, "full element size change with species %" PetscInt_FMT " %" PetscInt_FMT, coo_elem_fullNb[glb_elem_idx], fullNb); 1626 } 1627 } // field 1628 } // cell 1629 // allocate and copy point data maps[grid].gIdx[eidx][field][q] 1630 PetscCall(PetscMalloc(maps[grid].num_reduced * sizeof(*maps[grid].c_maps), &maps[grid].c_maps)); 1631 for (PetscInt ej = 0; ej < maps[grid].num_reduced; ++ej) { 1632 for (PetscInt q = 0; q < maps[grid].num_face; ++q) { 1633 maps[grid].c_maps[ej][q].scale = pointMaps[ej][q].scale; 1634 maps[grid].c_maps[ej][q].gid = pointMaps[ej][q].gid; 1635 } 1636 } 1637 #if defined(PETSC_HAVE_KOKKOS) 1638 if (ctx->deviceType == LANDAU_KOKKOS) { 1639 PetscCall(LandauKokkosCreateMatMaps(maps, pointMaps, Nf, grid)); // implies Kokkos does 1640 } 1641 #endif 1642 if (plex_batch) { 1643 PetscCall(ISRestoreIndices(grid_batch_is_inv[grid], &plex_batch)); 1644 PetscCall(ISDestroy(&grid_batch_is_inv[grid])); // we are done with this 1645 } 1646 } /* grids */ 1647 // finish COO 1648 { // setup COO assembly 1649 PetscInt *oor, *ooc; 1650 ctx->SData_d.coo_size = coo_elem_offsets[ncellsTot] * ctx->batch_sz; 1651 PetscCall(PetscMalloc2(ctx->SData_d.coo_size, &oor, ctx->SData_d.coo_size, &ooc)); 1652 for (PetscInt i = 0; i < ctx->SData_d.coo_size; i++) oor[i] = ooc[i] = -1; 1653 // get 1654 for (PetscInt grid = 0, glb_elem_idx = 0; grid < ctx->num_grids; grid++) { 1655 for (PetscInt ej = 0; ej < numCells[grid]; ++ej, glb_elem_idx++) { 1656 const PetscInt fullNb = coo_elem_fullNb[glb_elem_idx]; 1657 const LandauIdx *const Idxs = &maps[grid].gIdx[ej][0][0]; // just use field-0 maps, They should be the same but this is just for COO storage 1658 coo_elem_point_offsets[glb_elem_idx][0] = 0; 1659 for (PetscInt f = 0, cnt2 = 0; f < Nb; f++) { 1660 PetscInt idx = Idxs[f]; 1661 coo_elem_point_offsets[glb_elem_idx][f + 1] = coo_elem_point_offsets[glb_elem_idx][f]; // start at last 1662 if (idx >= 0) { 1663 cnt2++; 1664 coo_elem_point_offsets[glb_elem_idx][f + 1]++; // inc 1665 } else { 1666 idx = -idx - 1; 1667 for (PetscInt q = 0; q < maps[grid].num_face; q++) { 1668 if (maps[grid].c_maps[idx][q].gid < 0) break; 1669 cnt2++; 1670 coo_elem_point_offsets[glb_elem_idx][f + 1]++; // inc 1671 } 1672 } 1673 PetscCheck(cnt2 <= fullNb, PETSC_COMM_SELF, PETSC_ERR_PLIB, "wrong count %" PetscInt_FMT " < %" PetscInt_FMT, fullNb, cnt2); 1674 } 1675 PetscCheck(coo_elem_point_offsets[glb_elem_idx][Nb] == fullNb, PETSC_COMM_SELF, PETSC_ERR_PLIB, "coo_elem_point_offsets size %" PetscInt_FMT " != fullNb=%" PetscInt_FMT, coo_elem_point_offsets[glb_elem_idx][Nb], fullNb); 1676 } 1677 } 1678 // set 1679 for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) { 1680 for (PetscInt grid = 0, glb_elem_idx = 0; grid < ctx->num_grids; grid++) { 1681 const PetscInt moffset = LAND_MOFFSET(b_id, grid, ctx->batch_sz, ctx->num_grids, ctx->mat_offset); 1682 for (PetscInt ej = 0; ej < numCells[grid]; ++ej, glb_elem_idx++) { 1683 const PetscInt fullNb = coo_elem_fullNb[glb_elem_idx], fullNb2 = fullNb * fullNb; 1684 // set (i,j) 1685 for (PetscInt fieldA = 0; fieldA < Nf[grid]; fieldA++) { 1686 const LandauIdx *const Idxs = &maps[grid].gIdx[ej][fieldA][0]; 1687 PetscInt rows[LANDAU_MAX_Q_FACE], cols[LANDAU_MAX_Q_FACE]; 1688 for (PetscInt f = 0; f < Nb; ++f) { 1689 const PetscInt nr = coo_elem_point_offsets[glb_elem_idx][f + 1] - coo_elem_point_offsets[glb_elem_idx][f]; 1690 if (nr == 1) rows[0] = Idxs[f]; 1691 else { 1692 const PetscInt idx = -Idxs[f] - 1; 1693 for (PetscInt q = 0; q < nr; q++) rows[q] = maps[grid].c_maps[idx][q].gid; 1694 } 1695 for (PetscInt g = 0; g < Nb; ++g) { 1696 const PetscInt nc = coo_elem_point_offsets[glb_elem_idx][g + 1] - coo_elem_point_offsets[glb_elem_idx][g]; 1697 if (nc == 1) cols[0] = Idxs[g]; 1698 else { 1699 const PetscInt idx = -Idxs[g] - 1; 1700 for (PetscInt q = 0; q < nc; q++) cols[q] = maps[grid].c_maps[idx][q].gid; 1701 } 1702 const PetscInt idx0 = b_id * coo_elem_offsets[ncellsTot] + coo_elem_offsets[glb_elem_idx] + fieldA * fullNb2 + fullNb * coo_elem_point_offsets[glb_elem_idx][f] + nr * coo_elem_point_offsets[glb_elem_idx][g]; 1703 for (PetscInt q = 0, idx = idx0; q < nr; q++) { 1704 for (PetscInt d = 0; d < nc; d++, idx++) { 1705 oor[idx] = rows[q] + moffset; 1706 ooc[idx] = cols[d] + moffset; 1707 } 1708 } 1709 } 1710 } 1711 } 1712 } // cell 1713 } // grid 1714 } // batch 1715 PetscCall(MatSetPreallocationCOO(ctx->J, ctx->SData_d.coo_size, oor, ooc)); 1716 PetscCall(PetscFree2(oor, ooc)); 1717 } 1718 PetscCall(PetscFree(pointMaps)); 1719 PetscCall(PetscFree(elemMatrix)); 1720 PetscCall(PetscContainerCreate(PETSC_COMM_SELF, &container)); 1721 PetscCall(PetscContainerSetPointer(container, (void *)maps)); 1722 PetscCall(PetscContainerSetCtxDestroy(container, LandauGPUMapsDestroy)); 1723 PetscCall(PetscObjectCompose((PetscObject)ctx->J, "assembly_maps", (PetscObject)container)); 1724 PetscCall(PetscContainerDestroy(&container)); 1725 PetscCall(PetscLogEventEnd(ctx->events[2], 0, 0, 0, 0)); 1726 } // end GPU assembly 1727 { /* create static point data, Jacobian called first, only one vertex copy */ 1728 PetscReal *invJe, *ww, *xx, *yy, *zz = NULL, *invJ_a; 1729 PetscInt outer_ipidx, outer_ej, grid, nip_glb = 0; 1730 PetscFE fe; 1731 PetscCall(PetscLogEventBegin(ctx->events[7], 0, 0, 0, 0)); 1732 PetscCall(PetscInfo(ctx->plex[0], "Initialize static data\n")); 1733 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) nip_glb += Nq * numCells[grid]; 1734 /* collect f data, first time is for Jacobian, but make mass now */ 1735 if (ctx->verbose != 0) { 1736 PetscInt ncells = 0, N; 1737 MatInfo info; 1738 PetscCall(MatGetInfo(ctx->J, MAT_LOCAL, &info)); 1739 PetscCall(MatGetSize(ctx->J, &N, NULL)); 1740 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) ncells += numCells[grid]; 1741 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%d) %s %" PetscInt_FMT " IPs, %" PetscInt_FMT " cells total, Nb=%" PetscInt_FMT ", Nq=%" PetscInt_FMT ", dim=%" PetscInt_FMT ", Tab: Nb=%" PetscInt_FMT " Nf=%" PetscInt_FMT " Np=%" PetscInt_FMT " cdim=%" PetscInt_FMT " N=%" PetscInt_FMT " nnz= %" PetscInt_FMT "\n", 0, "FormLandau", nip_glb, ncells, Nb, Nq, dim, Nb, 1742 ctx->num_species, Nb, dim, N, (PetscInt)info.nz_used)); 1743 } 1744 PetscCall(PetscMalloc4(nip_glb, &ww, nip_glb, &xx, nip_glb, &yy, nip_glb * dim * dim, &invJ_a)); 1745 if (dim == 3) PetscCall(PetscMalloc1(nip_glb, &zz)); 1746 if (ctx->use_energy_tensor_trick) { 1747 PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, ctx->simplex, NULL, PETSC_DECIDE, &fe)); 1748 PetscCall(PetscObjectSetName((PetscObject)fe, "energy")); 1749 } 1750 /* init each grids static data - no batch */ 1751 for (grid = 0, outer_ipidx = 0, outer_ej = 0; grid < ctx->num_grids; grid++) { // OpenMP (once) 1752 Vec v2_2 = NULL; // projected function: v^2/2 for non-relativistic, gamma... for relativistic 1753 PetscSection e_section; 1754 DM dmEnergy; 1755 PetscInt cStart, cEnd, ej; 1756 1757 PetscCall(DMPlexGetHeightStratum(ctx->plex[grid], 0, &cStart, &cEnd)); 1758 // prep energy trick, get v^2 / 2 vector 1759 if (ctx->use_energy_tensor_trick) { 1760 PetscErrorCode (*energyf[1])(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *) = {ctx->use_relativistic_corrections ? gamma_m1_f : energy_f}; 1761 Vec glob_v2; 1762 PetscReal *c2_0[1], data[1] = {PetscSqr(C_0(ctx->v_0))}; 1763 1764 PetscCall(DMClone(ctx->plex[grid], &dmEnergy)); 1765 PetscCall(PetscObjectSetName((PetscObject)dmEnergy, "energy")); 1766 PetscCall(DMSetField(dmEnergy, 0, NULL, (PetscObject)fe)); 1767 PetscCall(DMCreateDS(dmEnergy)); 1768 PetscCall(DMGetLocalSection(dmEnergy, &e_section)); 1769 PetscCall(DMGetGlobalVector(dmEnergy, &glob_v2)); 1770 PetscCall(PetscObjectSetName((PetscObject)glob_v2, "trick")); 1771 c2_0[0] = &data[0]; 1772 PetscCall(DMProjectFunction(dmEnergy, 0., energyf, (void **)c2_0, INSERT_ALL_VALUES, glob_v2)); 1773 PetscCall(DMGetLocalVector(dmEnergy, &v2_2)); 1774 PetscCall(VecZeroEntries(v2_2)); /* zero BCs so don't set */ 1775 PetscCall(DMGlobalToLocalBegin(dmEnergy, glob_v2, INSERT_VALUES, v2_2)); 1776 PetscCall(DMGlobalToLocalEnd(dmEnergy, glob_v2, INSERT_VALUES, v2_2)); 1777 PetscCall(DMViewFromOptions(dmEnergy, NULL, "-energy_dm_view")); 1778 PetscCall(VecViewFromOptions(glob_v2, NULL, "-energy_vec_view")); 1779 PetscCall(DMRestoreGlobalVector(dmEnergy, &glob_v2)); 1780 } 1781 /* append part of the IP data for each grid */ 1782 for (ej = 0; ej < numCells[grid]; ++ej, ++outer_ej) { 1783 PetscScalar *coefs = NULL; 1784 PetscReal vj[LANDAU_MAX_NQND * LANDAU_DIM], detJj[LANDAU_MAX_NQND], Jdummy[LANDAU_MAX_NQND * LANDAU_DIM * LANDAU_DIM], c0 = C_0(ctx->v_0), c02 = PetscSqr(c0); 1785 invJe = invJ_a + outer_ej * Nq * dim * dim; 1786 PetscCall(DMPlexComputeCellGeometryFEM(ctx->plex[grid], ej + cStart, quad, vj, Jdummy, invJe, detJj)); 1787 if (ctx->use_energy_tensor_trick) PetscCall(DMPlexVecGetClosure(dmEnergy, e_section, v2_2, ej + cStart, NULL, &coefs)); 1788 /* create static point data */ 1789 for (PetscInt qj = 0; qj < Nq; qj++, outer_ipidx++) { 1790 const PetscInt gidx = outer_ipidx; 1791 const PetscReal *invJ = &invJe[qj * dim * dim]; 1792 ww[gidx] = detJj[qj] * quadWeights[qj]; 1793 if (dim == 2) ww[gidx] *= vj[qj * dim + 0]; /* cylindrical coordinate, w/o 2pi */ 1794 // get xx, yy, zz 1795 if (ctx->use_energy_tensor_trick) { 1796 double refSpaceDer[3], eGradPhi[3]; 1797 const PetscReal *const DD = Tf[0]->T[1]; 1798 const PetscReal *Dq = &DD[qj * Nb * dim]; 1799 for (PetscInt d = 0; d < 3; ++d) refSpaceDer[d] = eGradPhi[d] = 0.0; 1800 for (PetscInt b = 0; b < Nb; ++b) { 1801 for (PetscInt d = 0; d < dim; ++d) refSpaceDer[d] += Dq[b * dim + d] * PetscRealPart(coefs[b]); 1802 } 1803 xx[gidx] = 1e10; 1804 if (ctx->use_relativistic_corrections) { 1805 double dg2_c2 = 0; 1806 //for (PetscInt d = 0; d < dim; ++d) refSpaceDer[d] *= c02; 1807 for (PetscInt d = 0; d < dim; ++d) dg2_c2 += PetscSqr(refSpaceDer[d]); 1808 dg2_c2 *= (double)c02; 1809 if (dg2_c2 >= .999) { 1810 xx[gidx] = vj[qj * dim + 0]; /* coordinate */ 1811 yy[gidx] = vj[qj * dim + 1]; 1812 if (dim == 3) zz[gidx] = vj[qj * dim + 2]; 1813 PetscCall(PetscPrintf(ctx->comm, "Error: %12.5e %" PetscInt_FMT ".%" PetscInt_FMT ") dg2/c02 = %12.5e x= %12.5e %12.5e %12.5e\n", (double)PetscSqrtReal(xx[gidx] * xx[gidx] + yy[gidx] * yy[gidx] + zz[gidx] * zz[gidx]), ej, qj, dg2_c2, (double)xx[gidx], (double)yy[gidx], (double)zz[gidx])); 1814 } else { 1815 PetscReal fact = c02 / PetscSqrtReal(1. - dg2_c2); 1816 for (PetscInt d = 0; d < dim; ++d) refSpaceDer[d] *= fact; 1817 // could test with other point u' that (grad - grad') * U (refSpaceDer, refSpaceDer') == 0 1818 } 1819 } 1820 if (xx[gidx] == 1e10) { 1821 for (PetscInt d = 0; d < dim; ++d) { 1822 for (PetscInt e = 0; e < dim; ++e) eGradPhi[d] += invJ[e * dim + d] * refSpaceDer[e]; 1823 } 1824 xx[gidx] = eGradPhi[0]; 1825 yy[gidx] = eGradPhi[1]; 1826 if (dim == 3) zz[gidx] = eGradPhi[2]; 1827 } 1828 } else { 1829 xx[gidx] = vj[qj * dim + 0]; /* coordinate */ 1830 yy[gidx] = vj[qj * dim + 1]; 1831 if (dim == 3) zz[gidx] = vj[qj * dim + 2]; 1832 } 1833 } /* q */ 1834 if (ctx->use_energy_tensor_trick) PetscCall(DMPlexVecRestoreClosure(dmEnergy, e_section, v2_2, ej + cStart, NULL, &coefs)); 1835 } /* ej */ 1836 if (ctx->use_energy_tensor_trick) { 1837 PetscCall(DMRestoreLocalVector(dmEnergy, &v2_2)); 1838 PetscCall(DMDestroy(&dmEnergy)); 1839 } 1840 } /* grid */ 1841 if (ctx->use_energy_tensor_trick) PetscCall(PetscFEDestroy(&fe)); 1842 /* cache static data */ 1843 if (ctx->deviceType == LANDAU_KOKKOS) { 1844 #if defined(PETSC_HAVE_KOKKOS) 1845 PetscCall(LandauKokkosStaticDataSet(ctx->plex[0], Nq, Nb, ctx->batch_sz, ctx->num_grids, numCells, ctx->species_offset, ctx->mat_offset, nu_alpha, nu_beta, invMass, (PetscReal *)ctx->lambdas, invJ_a, xx, yy, zz, ww, &ctx->SData_d)); 1846 /* free */ 1847 PetscCall(PetscFree4(ww, xx, yy, invJ_a)); 1848 if (dim == 3) PetscCall(PetscFree(zz)); 1849 #else 1850 SETERRQ(ctx->comm, PETSC_ERR_ARG_WRONG, "-landau_device_type kokkos not built"); 1851 #endif 1852 } else { /* CPU version, just copy in, only use part */ 1853 PetscReal *nu_alpha_p = (PetscReal *)ctx->SData_d.alpha, *nu_beta_p = (PetscReal *)ctx->SData_d.beta, *invMass_p = (PetscReal *)ctx->SData_d.invMass, *lambdas_p = NULL; // why set these ? 1854 ctx->SData_d.w = (void *)ww; 1855 ctx->SData_d.x = (void *)xx; 1856 ctx->SData_d.y = (void *)yy; 1857 ctx->SData_d.z = (void *)zz; 1858 ctx->SData_d.invJ = (void *)invJ_a; 1859 PetscCall(PetscMalloc4(ctx->num_species, &nu_alpha_p, ctx->num_species, &nu_beta_p, ctx->num_species, &invMass_p, LANDAU_MAX_GRIDS * LANDAU_MAX_GRIDS, &lambdas_p)); 1860 for (PetscInt ii = 0; ii < ctx->num_species; ii++) { 1861 nu_alpha_p[ii] = nu_alpha[ii]; 1862 nu_beta_p[ii] = nu_beta[ii]; 1863 invMass_p[ii] = invMass[ii]; 1864 } 1865 ctx->SData_d.alpha = (void *)nu_alpha_p; 1866 ctx->SData_d.beta = (void *)nu_beta_p; 1867 ctx->SData_d.invMass = (void *)invMass_p; 1868 ctx->SData_d.lambdas = (void *)lambdas_p; 1869 for (PetscInt grid = 0; grid < LANDAU_MAX_GRIDS; grid++) { 1870 PetscReal(*lambdas)[LANDAU_MAX_GRIDS][LANDAU_MAX_GRIDS] = (PetscReal(*)[LANDAU_MAX_GRIDS][LANDAU_MAX_GRIDS])ctx->SData_d.lambdas; 1871 for (PetscInt gridj = 0; gridj < LANDAU_MAX_GRIDS; gridj++) { (*lambdas)[grid][gridj] = ctx->lambdas[grid][gridj]; } 1872 } 1873 } 1874 PetscCall(PetscLogEventEnd(ctx->events[7], 0, 0, 0, 0)); 1875 } // initialize 1876 PetscFunctionReturn(PETSC_SUCCESS); 1877 } 1878 1879 /* < v, u > */ 1880 static void g0_1(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g0[]) 1881 { 1882 g0[0] = 1.; 1883 } 1884 1885 /* < v, u > */ 1886 static void g0_fake(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g0[]) 1887 { 1888 static double ttt = 1e-12; 1889 g0[0] = ttt++; 1890 } 1891 1892 /* < v, u > */ 1893 static void g0_r(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g0[]) 1894 { 1895 g0[0] = 2. * PETSC_PI * x[0]; 1896 } 1897 1898 /* 1899 LandauCreateJacobianMatrix - creates ctx->J with without real data. Hard to keep sparse. 1900 - Like DMPlexLandauCreateMassMatrix. Should remove one and combine 1901 - has old support for field major ordering 1902 */ 1903 static PetscErrorCode LandauCreateJacobianMatrix(MPI_Comm comm, Vec X, IS grid_batch_is_inv[LANDAU_MAX_GRIDS], LandauCtx *ctx) 1904 { 1905 PetscInt *idxs = NULL; 1906 Mat subM[LANDAU_MAX_GRIDS]; 1907 1908 PetscFunctionBegin; 1909 if (!ctx->gpu_assembly) { /* we need GPU object with GPU assembly */ 1910 PetscFunctionReturn(PETSC_SUCCESS); 1911 } 1912 // get the RCM for this grid to separate out species into blocks -- create 'idxs' & 'ctx->batch_is' -- not used 1913 if (ctx->gpu_assembly && ctx->jacobian_field_major_order) PetscCall(PetscMalloc1(ctx->mat_offset[ctx->num_grids] * ctx->batch_sz, &idxs)); 1914 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1915 const PetscInt *values, n = ctx->mat_offset[grid + 1] - ctx->mat_offset[grid]; 1916 Mat gMat; 1917 DM massDM; 1918 PetscDS prob; 1919 Vec tvec; 1920 // get "mass" matrix for reordering 1921 PetscCall(DMClone(ctx->plex[grid], &massDM)); 1922 PetscCall(DMCopyFields(ctx->plex[grid], PETSC_DETERMINE, PETSC_DETERMINE, massDM)); 1923 PetscCall(DMCreateDS(massDM)); 1924 PetscCall(DMGetDS(massDM, &prob)); 1925 for (PetscInt ix = 0, ii = ctx->species_offset[grid]; ii < ctx->species_offset[grid + 1]; ii++, ix++) PetscCall(PetscDSSetJacobian(prob, ix, ix, g0_fake, NULL, NULL, NULL)); 1926 PetscCall(PetscOptionsInsertString(NULL, "-dm_preallocate_only")); // this trick is need to both sparsify the matrix and avoid runtime error 1927 PetscCall(DMCreateMatrix(massDM, &gMat)); 1928 PetscCall(PetscOptionsInsertString(NULL, "-dm_preallocate_only false")); 1929 PetscCall(MatSetOption(gMat, MAT_STRUCTURALLY_SYMMETRIC, PETSC_TRUE)); 1930 PetscCall(MatSetOption(gMat, MAT_IGNORE_ZERO_ENTRIES, PETSC_TRUE)); 1931 PetscCall(DMCreateLocalVector(ctx->plex[grid], &tvec)); 1932 PetscCall(DMPlexSNESComputeJacobianFEM(massDM, tvec, gMat, gMat, ctx)); 1933 PetscCall(MatViewFromOptions(gMat, NULL, "-dm_landau_reorder_mat_view")); 1934 PetscCall(DMDestroy(&massDM)); 1935 PetscCall(VecDestroy(&tvec)); 1936 subM[grid] = gMat; 1937 if (ctx->gpu_assembly && ctx->jacobian_field_major_order) { 1938 MatOrderingType rtype = MATORDERINGRCM; 1939 IS isrow, isicol; 1940 PetscCall(MatGetOrdering(gMat, rtype, &isrow, &isicol)); 1941 PetscCall(ISInvertPermutation(isrow, PETSC_DECIDE, &grid_batch_is_inv[grid])); 1942 PetscCall(ISGetIndices(isrow, &values)); 1943 for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) { // add batch size DMs for this species grid 1944 #if !defined(LANDAU_SPECIES_MAJOR) 1945 PetscInt N = ctx->mat_offset[ctx->num_grids], n0 = ctx->mat_offset[grid] + b_id * N; 1946 for (PetscInt ii = 0; ii < n; ++ii) idxs[n0 + ii] = values[ii] + n0; 1947 #else 1948 PetscInt n0 = ctx->mat_offset[grid] * ctx->batch_sz + b_id * n; 1949 for (PetscInt ii = 0; ii < n; ++ii) idxs[n0 + ii] = values[ii] + n0; 1950 #endif 1951 } 1952 PetscCall(ISRestoreIndices(isrow, &values)); 1953 PetscCall(ISDestroy(&isrow)); 1954 PetscCall(ISDestroy(&isicol)); 1955 } 1956 } 1957 if (ctx->gpu_assembly && ctx->jacobian_field_major_order) PetscCall(ISCreateGeneral(comm, ctx->mat_offset[ctx->num_grids] * ctx->batch_sz, idxs, PETSC_OWN_POINTER, &ctx->batch_is)); 1958 // get a block matrix 1959 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 1960 Mat B = subM[grid]; 1961 PetscInt nloc, nzl, *colbuf, row, COL_BF_SIZE = 1024; 1962 PetscCall(PetscMalloc(sizeof(*colbuf) * COL_BF_SIZE, &colbuf)); 1963 PetscCall(MatGetSize(B, &nloc, NULL)); 1964 for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) { 1965 const PetscInt moffset = LAND_MOFFSET(b_id, grid, ctx->batch_sz, ctx->num_grids, ctx->mat_offset); 1966 const PetscInt *cols; 1967 const PetscScalar *vals; 1968 for (PetscInt i = 0; i < nloc; i++) { 1969 PetscCall(MatGetRow(B, i, &nzl, NULL, NULL)); 1970 if (nzl > COL_BF_SIZE) { 1971 PetscCall(PetscFree(colbuf)); 1972 PetscCall(PetscInfo(ctx->plex[grid], "Realloc buffer %" PetscInt_FMT " to %" PetscInt_FMT " (row size %" PetscInt_FMT ") \n", COL_BF_SIZE, 2 * COL_BF_SIZE, nzl)); 1973 COL_BF_SIZE = nzl; 1974 PetscCall(PetscMalloc(sizeof(*colbuf) * COL_BF_SIZE, &colbuf)); 1975 } 1976 PetscCall(MatGetRow(B, i, &nzl, &cols, &vals)); 1977 for (PetscInt j = 0; j < nzl; j++) colbuf[j] = cols[j] + moffset; 1978 row = i + moffset; 1979 PetscCall(MatSetValues(ctx->J, 1, &row, nzl, colbuf, vals, INSERT_VALUES)); 1980 PetscCall(MatRestoreRow(B, i, &nzl, &cols, &vals)); 1981 } 1982 } 1983 PetscCall(PetscFree(colbuf)); 1984 } 1985 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) PetscCall(MatDestroy(&subM[grid])); 1986 PetscCall(MatAssemblyBegin(ctx->J, MAT_FINAL_ASSEMBLY)); 1987 PetscCall(MatAssemblyEnd(ctx->J, MAT_FINAL_ASSEMBLY)); 1988 1989 // debug 1990 PetscCall(MatViewFromOptions(ctx->J, NULL, "-dm_landau_mat_view")); 1991 if (ctx->gpu_assembly && ctx->jacobian_field_major_order) { 1992 Mat mat_block_order; 1993 PetscCall(MatCreateSubMatrix(ctx->J, ctx->batch_is, ctx->batch_is, MAT_INITIAL_MATRIX, &mat_block_order)); // use MatPermute 1994 PetscCall(MatViewFromOptions(mat_block_order, NULL, "-dm_landau_mat_view")); 1995 PetscCall(MatDestroy(&mat_block_order)); 1996 PetscCall(VecScatterCreate(X, ctx->batch_is, X, NULL, &ctx->plex_batch)); 1997 PetscCall(VecDuplicate(X, &ctx->work_vec)); 1998 } 1999 PetscFunctionReturn(PETSC_SUCCESS); 2000 } 2001 2002 PetscErrorCode DMPlexLandauCreateMassMatrix(DM pack, Mat *Amat); 2003 /*@C 2004 DMPlexLandauCreateVelocitySpace - Create a `DMPLEX` velocity space mesh 2005 2006 Collective 2007 2008 Input Parameters: 2009 + comm - The MPI communicator 2010 . dim - velocity space dimension (2 for axisymmetric, 3 for full 3X + 3V solver) 2011 - prefix - prefix for options (not tested) 2012 2013 Output Parameters: 2014 + pack - The `DM` object representing the mesh 2015 . X - A vector (user destroys) 2016 - J - Optional matrix (object destroys) 2017 2018 Level: beginner 2019 2020 .seealso: `DMPlexCreate()`, `DMPlexLandauDestroyVelocitySpace()` 2021 @*/ 2022 PetscErrorCode DMPlexLandauCreateVelocitySpace(MPI_Comm comm, PetscInt dim, const char prefix[], Vec *X, Mat *J, DM *pack) 2023 { 2024 LandauCtx *ctx; 2025 Vec Xsub[LANDAU_MAX_GRIDS]; 2026 IS grid_batch_is_inv[LANDAU_MAX_GRIDS]; 2027 2028 PetscFunctionBegin; 2029 PetscCheck(dim == 2 || dim == 3, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Only 2D and 3D supported"); 2030 PetscCheck(LANDAU_DIM == dim, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dim %" PetscInt_FMT " != LANDAU_DIM %d", dim, LANDAU_DIM); 2031 PetscCall(PetscNew(&ctx)); 2032 ctx->comm = comm; /* used for diagnostics and global errors */ 2033 /* process options */ 2034 PetscCall(ProcessOptions(ctx, prefix)); 2035 if (dim == 2) ctx->use_relativistic_corrections = PETSC_FALSE; 2036 /* Create Mesh */ 2037 PetscCall(DMCompositeCreate(PETSC_COMM_SELF, pack)); 2038 PetscCall(PetscLogEventBegin(ctx->events[13], 0, 0, 0, 0)); 2039 PetscCall(PetscLogEventBegin(ctx->events[15], 0, 0, 0, 0)); 2040 PetscCall(LandauDMCreateVMeshes(PETSC_COMM_SELF, dim, prefix, ctx, *pack)); // creates grids (Forest of AMR) 2041 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 2042 /* create FEM */ 2043 PetscCall(SetupDS(ctx->plex[grid], dim, grid, ctx)); 2044 /* set initial state */ 2045 PetscCall(DMCreateGlobalVector(ctx->plex[grid], &Xsub[grid])); 2046 PetscCall(PetscObjectSetName((PetscObject)Xsub[grid], "u_orig")); 2047 /* initial static refinement, no solve */ 2048 PetscCall(LandauSetInitialCondition(ctx->plex[grid], Xsub[grid], grid, 0, 1, ctx)); 2049 /* forest refinement - forest goes in (if forest), plex comes out */ 2050 if (ctx->use_p4est) { 2051 DM plex; 2052 PetscCall(adapt(grid, ctx, &Xsub[grid])); // forest goes in, plex comes out 2053 if (grid == 0) { 2054 PetscCall(DMViewFromOptions(ctx->plex[grid], NULL, "-dm_landau_amr_dm_view")); // need to differentiate - todo 2055 PetscCall(VecViewFromOptions(Xsub[grid], NULL, "-dm_landau_amr_vec_view")); 2056 } 2057 // convert to plex, all done with this level 2058 PetscCall(DMConvert(ctx->plex[grid], DMPLEX, &plex)); 2059 PetscCall(DMDestroy(&ctx->plex[grid])); 2060 ctx->plex[grid] = plex; 2061 } 2062 #if !defined(LANDAU_SPECIES_MAJOR) 2063 PetscCall(DMCompositeAddDM(*pack, ctx->plex[grid])); 2064 #else 2065 for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) { // add batch size DMs for this species grid 2066 PetscCall(DMCompositeAddDM(*pack, ctx->plex[grid])); 2067 } 2068 #endif 2069 PetscCall(DMSetApplicationContext(ctx->plex[grid], ctx)); 2070 } 2071 #if !defined(LANDAU_SPECIES_MAJOR) 2072 // stack the batched DMs, could do it all here!!! b_id=0 2073 for (PetscInt b_id = 1; b_id < ctx->batch_sz; b_id++) { 2074 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) PetscCall(DMCompositeAddDM(*pack, ctx->plex[grid])); 2075 } 2076 #endif 2077 // create ctx->mat_offset 2078 ctx->mat_offset[0] = 0; 2079 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 2080 PetscInt n; 2081 PetscCall(VecGetLocalSize(Xsub[grid], &n)); 2082 ctx->mat_offset[grid + 1] = ctx->mat_offset[grid] + n; 2083 } 2084 // creat DM & Jac 2085 PetscCall(DMSetApplicationContext(*pack, ctx)); 2086 PetscCall(PetscOptionsInsertString(NULL, "-dm_preallocate_only")); 2087 PetscCall(DMCreateMatrix(*pack, &ctx->J)); 2088 PetscCall(PetscOptionsInsertString(NULL, "-dm_preallocate_only false")); 2089 PetscCall(MatSetOption(ctx->J, MAT_STRUCTURALLY_SYMMETRIC, PETSC_TRUE)); 2090 PetscCall(MatSetOption(ctx->J, MAT_IGNORE_ZERO_ENTRIES, PETSC_TRUE)); 2091 PetscCall(PetscObjectSetName((PetscObject)ctx->J, "Jac")); 2092 // construct initial conditions in X 2093 PetscCall(DMCreateGlobalVector(*pack, X)); 2094 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 2095 PetscInt n; 2096 PetscCall(VecGetLocalSize(Xsub[grid], &n)); 2097 for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) { 2098 PetscScalar const *values; 2099 const PetscInt moffset = LAND_MOFFSET(b_id, grid, ctx->batch_sz, ctx->num_grids, ctx->mat_offset); 2100 PetscCall(LandauSetInitialCondition(ctx->plex[grid], Xsub[grid], grid, b_id, ctx->batch_sz, ctx)); 2101 PetscCall(VecGetArrayRead(Xsub[grid], &values)); // Drop whole grid in Plex ordering 2102 for (PetscInt i = 0, idx = moffset; i < n; i++, idx++) PetscCall(VecSetValue(*X, idx, values[i], INSERT_VALUES)); 2103 PetscCall(VecRestoreArrayRead(Xsub[grid], &values)); 2104 } 2105 } 2106 // cleanup 2107 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) PetscCall(VecDestroy(&Xsub[grid])); 2108 /* check for correct matrix type */ 2109 if (ctx->gpu_assembly) { /* we need GPU object with GPU assembly */ 2110 PetscBool flg; 2111 if (ctx->deviceType == LANDAU_KOKKOS) { 2112 PetscCall(PetscObjectTypeCompareAny((PetscObject)ctx->J, &flg, MATSEQAIJKOKKOS, MATMPIAIJKOKKOS, MATAIJKOKKOS, "")); 2113 #if defined(PETSC_HAVE_KOKKOS) 2114 PetscCheck(flg, ctx->comm, PETSC_ERR_ARG_WRONG, "must use '-dm_mat_type aijkokkos -dm_vec_type kokkos' for GPU assembly and Kokkos or use '-dm_landau_device_type cpu'"); 2115 #else 2116 PetscCheck(flg, ctx->comm, PETSC_ERR_ARG_WRONG, "must configure with '--download-kokkos-kernels' for GPU assembly and Kokkos or use '-dm_landau_device_type cpu'"); 2117 #endif 2118 } 2119 } 2120 PetscCall(PetscLogEventEnd(ctx->events[15], 0, 0, 0, 0)); 2121 2122 // create field major ordering 2123 ctx->work_vec = NULL; 2124 ctx->plex_batch = NULL; 2125 ctx->batch_is = NULL; 2126 for (PetscInt i = 0; i < LANDAU_MAX_GRIDS; i++) grid_batch_is_inv[i] = NULL; 2127 PetscCall(PetscLogEventBegin(ctx->events[12], 0, 0, 0, 0)); 2128 PetscCall(LandauCreateJacobianMatrix(comm, *X, grid_batch_is_inv, ctx)); 2129 PetscCall(PetscLogEventEnd(ctx->events[12], 0, 0, 0, 0)); 2130 2131 // create AMR GPU assembly maps and static GPU data 2132 PetscCall(CreateStaticData(dim, grid_batch_is_inv, ctx)); 2133 2134 PetscCall(PetscLogEventEnd(ctx->events[13], 0, 0, 0, 0)); 2135 2136 // create mass matrix 2137 PetscCall(DMPlexLandauCreateMassMatrix(*pack, NULL)); 2138 2139 if (J) *J = ctx->J; 2140 2141 if (ctx->gpu_assembly && ctx->jacobian_field_major_order) { 2142 PetscContainer container; 2143 // cache ctx for KSP with batch/field major Jacobian ordering -ksp_type gmres/etc -dm_landau_jacobian_field_major_order 2144 PetscCall(PetscContainerCreate(PETSC_COMM_SELF, &container)); 2145 PetscCall(PetscContainerSetPointer(container, (void *)ctx)); 2146 PetscCall(PetscObjectCompose((PetscObject)ctx->J, "LandauCtx", (PetscObject)container)); 2147 PetscCall(PetscContainerDestroy(&container)); 2148 // batch solvers need to map -- can batch solvers work 2149 PetscCall(PetscContainerCreate(PETSC_COMM_SELF, &container)); 2150 PetscCall(PetscContainerSetPointer(container, (void *)ctx->plex_batch)); 2151 PetscCall(PetscObjectCompose((PetscObject)ctx->J, "plex_batch_is", (PetscObject)container)); 2152 PetscCall(PetscContainerDestroy(&container)); 2153 } 2154 // for batch solvers 2155 { 2156 PetscContainer container; 2157 PetscInt *pNf; 2158 PetscCall(PetscContainerCreate(PETSC_COMM_SELF, &container)); 2159 PetscCall(PetscMalloc1(sizeof(*pNf), &pNf)); 2160 *pNf = ctx->batch_sz; 2161 PetscCall(PetscContainerSetPointer(container, (void *)pNf)); 2162 PetscCall(PetscContainerSetCtxDestroy(container, PetscCtxDestroyDefault)); 2163 PetscCall(PetscObjectCompose((PetscObject)ctx->J, "batch size", (PetscObject)container)); 2164 PetscCall(PetscContainerDestroy(&container)); 2165 } 2166 PetscFunctionReturn(PETSC_SUCCESS); 2167 } 2168 2169 /*@C 2170 DMPlexLandauAccess - Access to the distribution function with user callback 2171 2172 Collective 2173 2174 Input Parameters: 2175 + pack - the `DMCOMPOSITE` 2176 . func - call back function 2177 - user_ctx - user context 2178 2179 Input/Output Parameter: 2180 . X - Vector to data to 2181 2182 Level: advanced 2183 2184 .seealso: `DMPlexLandauCreateVelocitySpace()` 2185 @*/ 2186 PetscErrorCode DMPlexLandauAccess(DM pack, Vec X, PetscErrorCode (*func)(DM, Vec, PetscInt, PetscInt, PetscInt, void *), void *user_ctx) 2187 { 2188 LandauCtx *ctx; 2189 2190 PetscFunctionBegin; 2191 PetscCall(DMGetApplicationContext(pack, &ctx)); // uses ctx->num_grids; ctx->plex[grid]; ctx->batch_sz; ctx->mat_offset 2192 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 2193 PetscInt dim, n; 2194 PetscCall(DMGetDimension(pack, &dim)); 2195 for (PetscInt sp = ctx->species_offset[grid], i0 = 0; sp < ctx->species_offset[grid + 1]; sp++, i0++) { 2196 Vec vec; 2197 PetscInt vf[1] = {i0}; 2198 IS vis; 2199 DM vdm; 2200 PetscCall(DMCreateSubDM(ctx->plex[grid], 1, vf, &vis, &vdm)); 2201 PetscCall(DMSetApplicationContext(vdm, ctx)); // the user might want this 2202 PetscCall(DMCreateGlobalVector(vdm, &vec)); 2203 PetscCall(VecGetSize(vec, &n)); 2204 for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) { 2205 const PetscInt moffset = LAND_MOFFSET(b_id, grid, ctx->batch_sz, ctx->num_grids, ctx->mat_offset); 2206 PetscCall(VecZeroEntries(vec)); 2207 /* Add your data with 'dm' for species 'sp' to 'vec' */ 2208 PetscCall(func(vdm, vec, i0, grid, b_id, user_ctx)); 2209 /* add to global */ 2210 PetscScalar const *values; 2211 const PetscInt *offsets; 2212 PetscCall(VecGetArrayRead(vec, &values)); 2213 PetscCall(ISGetIndices(vis, &offsets)); 2214 for (PetscInt i = 0; i < n; i++) PetscCall(VecSetValue(X, moffset + offsets[i], values[i], ADD_VALUES)); 2215 PetscCall(VecRestoreArrayRead(vec, &values)); 2216 PetscCall(ISRestoreIndices(vis, &offsets)); 2217 } // batch 2218 PetscCall(VecDestroy(&vec)); 2219 PetscCall(ISDestroy(&vis)); 2220 PetscCall(DMDestroy(&vdm)); 2221 } 2222 } // grid 2223 PetscFunctionReturn(PETSC_SUCCESS); 2224 } 2225 2226 /*@ 2227 DMPlexLandauDestroyVelocitySpace - Destroy a `DMPLEX` velocity space mesh 2228 2229 Collective 2230 2231 Input/Output Parameters: 2232 . dm - the `DM` to destroy 2233 2234 Level: beginner 2235 2236 .seealso: `DMPlexLandauCreateVelocitySpace()` 2237 @*/ 2238 PetscErrorCode DMPlexLandauDestroyVelocitySpace(DM *dm) 2239 { 2240 LandauCtx *ctx; 2241 2242 PetscFunctionBegin; 2243 PetscCall(DMGetApplicationContext(*dm, &ctx)); 2244 PetscCall(MatDestroy(&ctx->M)); 2245 PetscCall(MatDestroy(&ctx->J)); 2246 for (PetscInt ii = 0; ii < ctx->num_species; ii++) PetscCall(PetscFEDestroy(&ctx->fe[ii])); 2247 PetscCall(ISDestroy(&ctx->batch_is)); 2248 PetscCall(VecDestroy(&ctx->work_vec)); 2249 PetscCall(VecScatterDestroy(&ctx->plex_batch)); 2250 if (ctx->deviceType == LANDAU_KOKKOS) { 2251 #if defined(PETSC_HAVE_KOKKOS) 2252 PetscCall(LandauKokkosStaticDataClear(&ctx->SData_d)); 2253 #else 2254 SETERRQ(ctx->comm, PETSC_ERR_ARG_WRONG, "-landau_device_type %s not built", "kokkos"); 2255 #endif 2256 } else { 2257 if (ctx->SData_d.x) { /* in a CPU run */ 2258 PetscReal *invJ = (PetscReal *)ctx->SData_d.invJ, *xx = (PetscReal *)ctx->SData_d.x, *yy = (PetscReal *)ctx->SData_d.y, *zz = (PetscReal *)ctx->SData_d.z, *ww = (PetscReal *)ctx->SData_d.w; 2259 LandauIdx *coo_elem_offsets = (LandauIdx *)ctx->SData_d.coo_elem_offsets, *coo_elem_fullNb = (LandauIdx *)ctx->SData_d.coo_elem_fullNb, (*coo_elem_point_offsets)[LANDAU_MAX_NQND + 1] = (LandauIdx(*)[LANDAU_MAX_NQND + 1]) ctx->SData_d.coo_elem_point_offsets; 2260 PetscCall(PetscFree4(ww, xx, yy, invJ)); 2261 if (zz) PetscCall(PetscFree(zz)); 2262 if (coo_elem_offsets) { 2263 PetscCall(PetscFree3(coo_elem_offsets, coo_elem_fullNb, coo_elem_point_offsets)); // could be NULL 2264 } 2265 PetscCall(PetscFree4(ctx->SData_d.alpha, ctx->SData_d.beta, ctx->SData_d.invMass, ctx->SData_d.lambdas)); 2266 } 2267 } 2268 2269 if (ctx->times[LANDAU_MATRIX_TOTAL] > 0) { // OMP timings 2270 PetscCall(PetscPrintf(ctx->comm, "TSStep N 1.0 %10.3e\n", ctx->times[LANDAU_EX2_TSSOLVE])); 2271 PetscCall(PetscPrintf(ctx->comm, "2: Solve: %10.3e with %" PetscInt_FMT " threads\n", ctx->times[LANDAU_EX2_TSSOLVE] - ctx->times[LANDAU_MATRIX_TOTAL], ctx->batch_sz)); 2272 PetscCall(PetscPrintf(ctx->comm, "3: Landau: %10.3e\n", ctx->times[LANDAU_MATRIX_TOTAL])); 2273 PetscCall(PetscPrintf(ctx->comm, "Landau Jacobian %" PetscInt_FMT " 1.0 %10.3e\n", (PetscInt)ctx->times[LANDAU_JACOBIAN_COUNT], ctx->times[LANDAU_JACOBIAN])); 2274 PetscCall(PetscPrintf(ctx->comm, "Landau Operator N 1.0 %10.3e\n", ctx->times[LANDAU_OPERATOR])); 2275 PetscCall(PetscPrintf(ctx->comm, "Landau Mass N 1.0 %10.3e\n", ctx->times[LANDAU_MASS])); 2276 PetscCall(PetscPrintf(ctx->comm, " Jac-f-df (GPU) N 1.0 %10.3e\n", ctx->times[LANDAU_F_DF])); 2277 PetscCall(PetscPrintf(ctx->comm, " Kernel (GPU) N 1.0 %10.3e\n", ctx->times[LANDAU_KERNEL])); 2278 PetscCall(PetscPrintf(ctx->comm, "MatLUFactorNum X 1.0 %10.3e\n", ctx->times[KSP_FACTOR])); 2279 PetscCall(PetscPrintf(ctx->comm, "MatSolve X 1.0 %10.3e\n", ctx->times[KSP_SOLVE])); 2280 } 2281 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) PetscCall(DMDestroy(&ctx->plex[grid])); 2282 PetscCall(PetscFree(ctx)); 2283 PetscCall(DMDestroy(dm)); 2284 PetscFunctionReturn(PETSC_SUCCESS); 2285 } 2286 2287 /* < v, ru > */ 2288 static void f0_s_den(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0) 2289 { 2290 PetscInt ii = (PetscInt)PetscRealPart(constants[0]); 2291 f0[0] = u[ii]; 2292 } 2293 2294 /* < v, ru > */ 2295 static void f0_s_mom(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0) 2296 { 2297 PetscInt ii = (PetscInt)PetscRealPart(constants[0]), jj = (PetscInt)PetscRealPart(constants[1]); 2298 f0[0] = x[jj] * u[ii]; /* x momentum */ 2299 } 2300 2301 static void f0_s_v2(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0) 2302 { 2303 PetscInt i, ii = (PetscInt)PetscRealPart(constants[0]); 2304 double tmp1 = 0.; 2305 for (i = 0; i < dim; ++i) tmp1 += x[i] * x[i]; 2306 f0[0] = tmp1 * u[ii]; 2307 } 2308 2309 static PetscErrorCode gamma_n_f(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nf, PetscScalar *u, void *actx) 2310 { 2311 const PetscReal *c2_0_arr = ((PetscReal *)actx); 2312 const PetscReal c02 = c2_0_arr[0]; 2313 2314 PetscFunctionBegin; 2315 for (PetscInt s = 0; s < Nf; s++) { 2316 PetscReal tmp1 = 0.; 2317 for (PetscInt i = 0; i < dim; ++i) tmp1 += x[i] * x[i]; 2318 #if defined(PETSC_USE_DEBUG) 2319 u[s] = PetscSqrtReal(1. + tmp1 / c02); // u[0] = PetscSqrtReal(1. + xx); 2320 #else 2321 { 2322 PetscReal xx = tmp1 / c02; 2323 u[s] = xx / (PetscSqrtReal(1. + xx) + 1.); // better conditioned = xx/(PetscSqrtReal(1. + xx) + 1.) 2324 } 2325 #endif 2326 } 2327 PetscFunctionReturn(PETSC_SUCCESS); 2328 } 2329 2330 /* < v, ru > */ 2331 static void f0_s_rden(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0) 2332 { 2333 PetscInt ii = (PetscInt)PetscRealPart(constants[0]); 2334 f0[0] = 2. * PETSC_PI * x[0] * u[ii]; 2335 } 2336 2337 /* < v, ru > */ 2338 static void f0_s_rmom(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0) 2339 { 2340 PetscInt ii = (PetscInt)PetscRealPart(constants[0]); 2341 f0[0] = 2. * PETSC_PI * x[0] * x[1] * u[ii]; 2342 } 2343 2344 static void f0_s_rv2(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar *f0) 2345 { 2346 PetscInt ii = (PetscInt)PetscRealPart(constants[0]); 2347 f0[0] = 2. * PETSC_PI * x[0] * (x[0] * x[0] + x[1] * x[1]) * u[ii]; 2348 } 2349 2350 /*@ 2351 DMPlexLandauPrintNorms - collects moments and prints them 2352 2353 Collective 2354 2355 Input Parameters: 2356 + X - the state 2357 - stepi - current step to print 2358 2359 Level: beginner 2360 2361 .seealso: `DMPlexLandauCreateVelocitySpace()` 2362 @*/ 2363 PetscErrorCode DMPlexLandauPrintNorms(Vec X, PetscInt stepi) 2364 { 2365 LandauCtx *ctx; 2366 PetscDS prob; 2367 DM pack; 2368 PetscInt cStart, cEnd, dim, ii, i0, nDMs; 2369 PetscScalar xmomentumtot = 0, ymomentumtot = 0, zmomentumtot = 0, energytot = 0, densitytot = 0, tt[LANDAU_MAX_SPECIES]; 2370 PetscScalar xmomentum[LANDAU_MAX_SPECIES], ymomentum[LANDAU_MAX_SPECIES], zmomentum[LANDAU_MAX_SPECIES], energy[LANDAU_MAX_SPECIES], density[LANDAU_MAX_SPECIES]; 2371 Vec *globXArray; 2372 2373 PetscFunctionBegin; 2374 PetscCall(VecGetDM(X, &pack)); 2375 PetscCheck(pack, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Vector has no DM"); 2376 PetscCall(DMGetDimension(pack, &dim)); 2377 PetscCheck(dim == 2 || dim == 3, PETSC_COMM_SELF, PETSC_ERR_PLIB, "dim %" PetscInt_FMT " not in [2,3]", dim); 2378 PetscCall(DMGetApplicationContext(pack, &ctx)); 2379 PetscCheck(ctx, PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 2380 /* print momentum and energy */ 2381 PetscCall(DMCompositeGetNumberDM(pack, &nDMs)); 2382 PetscCheck(nDMs == ctx->num_grids * ctx->batch_sz, PETSC_COMM_WORLD, PETSC_ERR_PLIB, "#DM wrong %" PetscInt_FMT " %" PetscInt_FMT, nDMs, ctx->num_grids * ctx->batch_sz); 2383 PetscCall(PetscMalloc(sizeof(*globXArray) * nDMs, &globXArray)); 2384 PetscCall(DMCompositeGetAccessArray(pack, X, nDMs, NULL, globXArray)); 2385 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 2386 Vec Xloc = globXArray[LAND_PACK_IDX(ctx->batch_view_idx, grid)]; 2387 PetscCall(DMGetDS(ctx->plex[grid], &prob)); 2388 for (ii = ctx->species_offset[grid], i0 = 0; ii < ctx->species_offset[grid + 1]; ii++, i0++) { 2389 PetscScalar user[2] = {(PetscScalar)i0, ctx->charges[ii]}; 2390 PetscCall(PetscDSSetConstants(prob, 2, user)); 2391 if (dim == 2) { /* 2/3X + 3V (cylindrical coordinates) */ 2392 PetscCall(PetscDSSetObjective(prob, 0, &f0_s_rden)); 2393 PetscCall(DMPlexComputeIntegralFEM(ctx->plex[grid], Xloc, tt, ctx)); 2394 density[ii] = tt[0] * ctx->n_0 * ctx->charges[ii]; 2395 PetscCall(PetscDSSetObjective(prob, 0, &f0_s_rmom)); 2396 PetscCall(DMPlexComputeIntegralFEM(ctx->plex[grid], Xloc, tt, ctx)); 2397 zmomentum[ii] = tt[0] * ctx->n_0 * ctx->v_0 * ctx->masses[ii]; 2398 PetscCall(PetscDSSetObjective(prob, 0, &f0_s_rv2)); 2399 PetscCall(DMPlexComputeIntegralFEM(ctx->plex[grid], Xloc, tt, ctx)); 2400 energy[ii] = tt[0] * 0.5 * ctx->n_0 * ctx->v_0 * ctx->v_0 * ctx->masses[ii]; 2401 zmomentumtot += zmomentum[ii]; 2402 energytot += energy[ii]; 2403 densitytot += density[ii]; 2404 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%3" PetscInt_FMT ") species-%" PetscInt_FMT ": charge density= %20.13e z-momentum= %20.13e energy= %20.13e", stepi, ii, (double)PetscRealPart(density[ii]), (double)PetscRealPart(zmomentum[ii]), (double)PetscRealPart(energy[ii]))); 2405 } else { /* 2/3Xloc + 3V */ 2406 PetscCall(PetscDSSetObjective(prob, 0, &f0_s_den)); 2407 PetscCall(DMPlexComputeIntegralFEM(ctx->plex[grid], Xloc, tt, ctx)); 2408 density[ii] = tt[0] * ctx->n_0 * ctx->charges[ii]; 2409 PetscCall(PetscDSSetObjective(prob, 0, &f0_s_mom)); 2410 user[1] = 0; 2411 PetscCall(PetscDSSetConstants(prob, 2, user)); 2412 PetscCall(DMPlexComputeIntegralFEM(ctx->plex[grid], Xloc, tt, ctx)); 2413 xmomentum[ii] = tt[0] * ctx->n_0 * ctx->v_0 * ctx->masses[ii]; 2414 user[1] = 1; 2415 PetscCall(PetscDSSetConstants(prob, 2, user)); 2416 PetscCall(DMPlexComputeIntegralFEM(ctx->plex[grid], Xloc, tt, ctx)); 2417 ymomentum[ii] = tt[0] * ctx->n_0 * ctx->v_0 * ctx->masses[ii]; 2418 user[1] = 2; 2419 PetscCall(PetscDSSetConstants(prob, 2, user)); 2420 PetscCall(DMPlexComputeIntegralFEM(ctx->plex[grid], Xloc, tt, ctx)); 2421 zmomentum[ii] = tt[0] * ctx->n_0 * ctx->v_0 * ctx->masses[ii]; 2422 if (ctx->use_relativistic_corrections) { 2423 /* gamma * M * f */ 2424 if (ii == 0 && grid == 0) { // do all at once 2425 Vec Mf, globGamma, *globMfArray, *globGammaArray; 2426 PetscErrorCode (*gammaf[1])(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *) = {gamma_n_f}; 2427 PetscReal *c2_0[1], data[1]; 2428 2429 PetscCall(VecDuplicate(X, &globGamma)); 2430 PetscCall(VecDuplicate(X, &Mf)); 2431 PetscCall(PetscMalloc(sizeof(*globMfArray) * nDMs, &globMfArray)); 2432 PetscCall(PetscMalloc(sizeof(*globMfArray) * nDMs, &globGammaArray)); 2433 /* M * f */ 2434 PetscCall(MatMult(ctx->M, X, Mf)); 2435 /* gamma */ 2436 PetscCall(DMCompositeGetAccessArray(pack, globGamma, nDMs, NULL, globGammaArray)); 2437 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // yes a grid loop in a grid loop to print nice, need to fix for batching 2438 Vec v1 = globGammaArray[LAND_PACK_IDX(ctx->batch_view_idx, grid)]; 2439 data[0] = PetscSqr(C_0(ctx->v_0)); 2440 c2_0[0] = &data[0]; 2441 PetscCall(DMProjectFunction(ctx->plex[grid], 0., gammaf, (void **)c2_0, INSERT_ALL_VALUES, v1)); 2442 } 2443 PetscCall(DMCompositeRestoreAccessArray(pack, globGamma, nDMs, NULL, globGammaArray)); 2444 /* gamma * Mf */ 2445 PetscCall(DMCompositeGetAccessArray(pack, globGamma, nDMs, NULL, globGammaArray)); 2446 PetscCall(DMCompositeGetAccessArray(pack, Mf, nDMs, NULL, globMfArray)); 2447 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { // yes a grid loop in a grid loop to print nice 2448 PetscInt Nf = ctx->species_offset[grid + 1] - ctx->species_offset[grid], N, bs; 2449 Vec Mfsub = globMfArray[LAND_PACK_IDX(ctx->batch_view_idx, grid)], Gsub = globGammaArray[LAND_PACK_IDX(ctx->batch_view_idx, grid)], v1, v2; 2450 // get each component 2451 PetscCall(VecGetSize(Mfsub, &N)); 2452 PetscCall(VecCreate(ctx->comm, &v1)); 2453 PetscCall(VecSetSizes(v1, PETSC_DECIDE, N / Nf)); 2454 PetscCall(VecCreate(ctx->comm, &v2)); 2455 PetscCall(VecSetSizes(v2, PETSC_DECIDE, N / Nf)); 2456 PetscCall(VecSetFromOptions(v1)); // ??? 2457 PetscCall(VecSetFromOptions(v2)); 2458 // get each component 2459 PetscCall(VecGetBlockSize(Gsub, &bs)); 2460 PetscCheck(bs == Nf, PETSC_COMM_SELF, PETSC_ERR_PLIB, "bs %" PetscInt_FMT " != num_species %" PetscInt_FMT " in Gsub", bs, Nf); 2461 PetscCall(VecGetBlockSize(Mfsub, &bs)); 2462 PetscCheck(bs == Nf, PETSC_COMM_SELF, PETSC_ERR_PLIB, "bs %" PetscInt_FMT " != num_species %" PetscInt_FMT, bs, Nf); 2463 for (PetscInt i = 0, ix = ctx->species_offset[grid]; i < Nf; i++, ix++) { 2464 PetscScalar val; 2465 PetscCall(VecStrideGather(Gsub, i, v1, INSERT_VALUES)); // this is not right -- TODO 2466 PetscCall(VecStrideGather(Mfsub, i, v2, INSERT_VALUES)); 2467 PetscCall(VecDot(v1, v2, &val)); 2468 energy[ix] = PetscRealPart(val) * ctx->n_0 * ctx->v_0 * ctx->v_0 * ctx->masses[ix]; 2469 } 2470 PetscCall(VecDestroy(&v1)); 2471 PetscCall(VecDestroy(&v2)); 2472 } /* grids */ 2473 PetscCall(DMCompositeRestoreAccessArray(pack, globGamma, nDMs, NULL, globGammaArray)); 2474 PetscCall(DMCompositeRestoreAccessArray(pack, Mf, nDMs, NULL, globMfArray)); 2475 PetscCall(PetscFree(globGammaArray)); 2476 PetscCall(PetscFree(globMfArray)); 2477 PetscCall(VecDestroy(&globGamma)); 2478 PetscCall(VecDestroy(&Mf)); 2479 } 2480 } else { 2481 PetscCall(PetscDSSetObjective(prob, 0, &f0_s_v2)); 2482 PetscCall(DMPlexComputeIntegralFEM(ctx->plex[grid], Xloc, tt, ctx)); 2483 energy[ii] = 0.5 * tt[0] * ctx->n_0 * ctx->v_0 * ctx->v_0 * ctx->masses[ii]; 2484 } 2485 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%3" PetscInt_FMT ") species %" PetscInt_FMT ": density=%20.13e, x-momentum=%20.13e, y-momentum=%20.13e, z-momentum=%20.13e, energy=%21.13e", stepi, ii, (double)PetscRealPart(density[ii]), (double)PetscRealPart(xmomentum[ii]), (double)PetscRealPart(ymomentum[ii]), (double)PetscRealPart(zmomentum[ii]), (double)PetscRealPart(energy[ii]))); 2486 xmomentumtot += xmomentum[ii]; 2487 ymomentumtot += ymomentum[ii]; 2488 zmomentumtot += zmomentum[ii]; 2489 energytot += energy[ii]; 2490 densitytot += density[ii]; 2491 } 2492 if (ctx->num_species > 1) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n")); 2493 } 2494 } 2495 PetscCall(DMCompositeRestoreAccessArray(pack, X, nDMs, NULL, globXArray)); 2496 PetscCall(PetscFree(globXArray)); 2497 /* totals */ 2498 PetscCall(DMPlexGetHeightStratum(ctx->plex[0], 0, &cStart, &cEnd)); 2499 if (ctx->num_species > 1) { 2500 if (dim == 2) { 2501 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\t%3" PetscInt_FMT ") Total: charge density=%21.13e, momentum=%21.13e, energy=%21.13e (m_i[0]/m_e = %g, %" PetscInt_FMT " cells on electron grid)", stepi, (double)PetscRealPart(densitytot), (double)PetscRealPart(zmomentumtot), (double)PetscRealPart(energytot), 2502 (double)(ctx->masses[1] / ctx->masses[0]), cEnd - cStart)); 2503 } else { 2504 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\t%3" PetscInt_FMT ") Total: charge density=%21.13e, x-momentum=%21.13e, y-momentum=%21.13e, z-momentum=%21.13e, energy=%21.13e (m_i[0]/m_e = %g, %" PetscInt_FMT " cells)", stepi, (double)PetscRealPart(densitytot), (double)PetscRealPart(xmomentumtot), (double)PetscRealPart(ymomentumtot), (double)PetscRealPart(zmomentumtot), (double)PetscRealPart(energytot), 2505 (double)(ctx->masses[1] / ctx->masses[0]), cEnd - cStart)); 2506 } 2507 } else PetscCall(PetscPrintf(PETSC_COMM_WORLD, " -- %" PetscInt_FMT " cells", cEnd - cStart)); 2508 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "\n")); 2509 PetscFunctionReturn(PETSC_SUCCESS); 2510 } 2511 2512 /*@ 2513 DMPlexLandauCreateMassMatrix - Create mass matrix for Landau in Plex space (not field major order of Jacobian) 2514 - puts mass matrix into ctx->M 2515 2516 Collective 2517 2518 Input Parameter: 2519 . pack - the `DM` object. Puts matrix in Landau context M field 2520 2521 Output Parameter: 2522 . Amat - The mass matrix (optional), mass matrix is added to the `DM` context 2523 2524 Level: beginner 2525 2526 .seealso: `DMPlexLandauCreateVelocitySpace()` 2527 @*/ 2528 PetscErrorCode DMPlexLandauCreateMassMatrix(DM pack, Mat *Amat) 2529 { 2530 DM mass_pack, massDM[LANDAU_MAX_GRIDS]; 2531 PetscDS prob; 2532 PetscInt ii, dim, N1 = 1, N2; 2533 LandauCtx *ctx; 2534 Mat packM, subM[LANDAU_MAX_GRIDS]; 2535 2536 PetscFunctionBegin; 2537 PetscValidHeaderSpecific(pack, DM_CLASSID, 1); 2538 if (Amat) PetscAssertPointer(Amat, 2); 2539 PetscCall(DMGetApplicationContext(pack, &ctx)); 2540 PetscCheck(ctx, PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 2541 PetscCall(PetscLogEventBegin(ctx->events[14], 0, 0, 0, 0)); 2542 PetscCall(DMGetDimension(pack, &dim)); 2543 PetscCall(DMCompositeCreate(PetscObjectComm((PetscObject)pack), &mass_pack)); 2544 /* create pack mass matrix */ 2545 for (PetscInt grid = 0, ix = 0; grid < ctx->num_grids; grid++) { 2546 PetscCall(DMClone(ctx->plex[grid], &massDM[grid])); 2547 PetscCall(DMCopyFields(ctx->plex[grid], PETSC_DETERMINE, PETSC_DETERMINE, massDM[grid])); 2548 PetscCall(DMCreateDS(massDM[grid])); 2549 PetscCall(DMGetDS(massDM[grid], &prob)); 2550 for (ix = 0, ii = ctx->species_offset[grid]; ii < ctx->species_offset[grid + 1]; ii++, ix++) { 2551 if (dim == 3) PetscCall(PetscDSSetJacobian(prob, ix, ix, g0_1, NULL, NULL, NULL)); 2552 else PetscCall(PetscDSSetJacobian(prob, ix, ix, g0_r, NULL, NULL, NULL)); 2553 } 2554 #if !defined(LANDAU_SPECIES_MAJOR) 2555 PetscCall(DMCompositeAddDM(mass_pack, massDM[grid])); 2556 #else 2557 for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) { // add batch size DMs for this species grid 2558 PetscCall(DMCompositeAddDM(mass_pack, massDM[grid])); 2559 } 2560 #endif 2561 PetscCall(DMCreateMatrix(massDM[grid], &subM[grid])); 2562 } 2563 #if !defined(LANDAU_SPECIES_MAJOR) 2564 // stack the batched DMs 2565 for (PetscInt b_id = 1; b_id < ctx->batch_sz; b_id++) { 2566 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) PetscCall(DMCompositeAddDM(mass_pack, massDM[grid])); 2567 } 2568 #endif 2569 PetscCall(PetscOptionsInsertString(NULL, "-dm_preallocate_only")); 2570 PetscCall(DMCreateMatrix(mass_pack, &packM)); 2571 PetscCall(PetscOptionsInsertString(NULL, "-dm_preallocate_only false")); 2572 PetscCall(MatSetOption(packM, MAT_STRUCTURALLY_SYMMETRIC, PETSC_TRUE)); 2573 PetscCall(MatSetOption(packM, MAT_IGNORE_ZERO_ENTRIES, PETSC_TRUE)); 2574 PetscCall(DMDestroy(&mass_pack)); 2575 /* make mass matrix for each block */ 2576 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 2577 Vec locX; 2578 DM plex = massDM[grid]; 2579 PetscCall(DMGetLocalVector(plex, &locX)); 2580 /* Mass matrix is independent of the input, so no need to fill locX */ 2581 PetscCall(DMPlexSNESComputeJacobianFEM(plex, locX, subM[grid], subM[grid], ctx)); 2582 PetscCall(DMRestoreLocalVector(plex, &locX)); 2583 PetscCall(DMDestroy(&massDM[grid])); 2584 } 2585 PetscCall(MatGetSize(ctx->J, &N1, NULL)); 2586 PetscCall(MatGetSize(packM, &N2, NULL)); 2587 PetscCheck(N1 == N2, PetscObjectComm((PetscObject)pack), PETSC_ERR_PLIB, "Incorrect matrix sizes: |Jacobian| = %" PetscInt_FMT ", |Mass|=%" PetscInt_FMT, N1, N2); 2588 /* assemble block diagonals */ 2589 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) { 2590 Mat B = subM[grid]; 2591 PetscInt nloc, nzl, *colbuf, COL_BF_SIZE = 1024, row; 2592 PetscCall(PetscMalloc(sizeof(*colbuf) * COL_BF_SIZE, &colbuf)); 2593 PetscCall(MatGetSize(B, &nloc, NULL)); 2594 for (PetscInt b_id = 0; b_id < ctx->batch_sz; b_id++) { 2595 const PetscInt moffset = LAND_MOFFSET(b_id, grid, ctx->batch_sz, ctx->num_grids, ctx->mat_offset); 2596 const PetscInt *cols; 2597 const PetscScalar *vals; 2598 for (PetscInt i = 0; i < nloc; i++) { 2599 PetscCall(MatGetRow(B, i, &nzl, NULL, NULL)); 2600 if (nzl > COL_BF_SIZE) { 2601 PetscCall(PetscFree(colbuf)); 2602 PetscCall(PetscInfo(pack, "Realloc buffer %" PetscInt_FMT " to %" PetscInt_FMT " (row size %" PetscInt_FMT ") \n", COL_BF_SIZE, 2 * COL_BF_SIZE, nzl)); 2603 COL_BF_SIZE = nzl; 2604 PetscCall(PetscMalloc(sizeof(*colbuf) * COL_BF_SIZE, &colbuf)); 2605 } 2606 PetscCall(MatGetRow(B, i, &nzl, &cols, &vals)); 2607 for (PetscInt j = 0; j < nzl; j++) colbuf[j] = cols[j] + moffset; 2608 row = i + moffset; 2609 PetscCall(MatSetValues(packM, 1, &row, nzl, colbuf, vals, INSERT_VALUES)); 2610 PetscCall(MatRestoreRow(B, i, &nzl, &cols, &vals)); 2611 } 2612 } 2613 PetscCall(PetscFree(colbuf)); 2614 } 2615 // cleanup 2616 for (PetscInt grid = 0; grid < ctx->num_grids; grid++) PetscCall(MatDestroy(&subM[grid])); 2617 PetscCall(MatAssemblyBegin(packM, MAT_FINAL_ASSEMBLY)); 2618 PetscCall(MatAssemblyEnd(packM, MAT_FINAL_ASSEMBLY)); 2619 PetscCall(PetscObjectSetName((PetscObject)packM, "mass")); 2620 PetscCall(MatViewFromOptions(packM, NULL, "-dm_landau_mass_view")); 2621 ctx->M = packM; 2622 if (Amat) *Amat = packM; 2623 PetscCall(PetscLogEventEnd(ctx->events[14], 0, 0, 0, 0)); 2624 PetscFunctionReturn(PETSC_SUCCESS); 2625 } 2626 2627 /*@ 2628 DMPlexLandauIFunction - `TS` residual calculation, confusingly this computes the Jacobian w/o mass 2629 2630 Collective 2631 2632 Input Parameters: 2633 + ts - The time stepping context 2634 . time_dummy - current time (not used) 2635 . X - Current state 2636 . X_t - Time derivative of current state 2637 - actx - Landau context 2638 2639 Output Parameter: 2640 . F - The residual 2641 2642 Level: beginner 2643 2644 .seealso: `DMPlexLandauCreateVelocitySpace()`, `DMPlexLandauIJacobian()` 2645 @*/ 2646 PetscErrorCode DMPlexLandauIFunction(TS ts, PetscReal time_dummy, Vec X, Vec X_t, Vec F, void *actx) 2647 { 2648 LandauCtx *ctx = (LandauCtx *)actx; 2649 PetscInt dim; 2650 DM pack; 2651 #if defined(PETSC_HAVE_THREADSAFETY) 2652 double starttime, endtime; 2653 #endif 2654 PetscObjectState state; 2655 2656 PetscFunctionBegin; 2657 PetscCall(TSGetDM(ts, &pack)); 2658 PetscCall(DMGetApplicationContext(pack, &ctx)); 2659 PetscCheck(ctx, PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 2660 if (ctx->stage) PetscCall(PetscLogStagePush(ctx->stage)); 2661 PetscCall(PetscLogEventBegin(ctx->events[11], 0, 0, 0, 0)); 2662 PetscCall(PetscLogEventBegin(ctx->events[0], 0, 0, 0, 0)); 2663 #if defined(PETSC_HAVE_THREADSAFETY) 2664 starttime = MPI_Wtime(); 2665 #endif 2666 PetscCall(DMGetDimension(pack, &dim)); 2667 PetscCall(PetscObjectStateGet((PetscObject)ctx->J, &state)); 2668 if (state != ctx->norm_state) { 2669 PetscCall(MatZeroEntries(ctx->J)); 2670 PetscCall(LandauFormJacobian_Internal(X, ctx->J, dim, 0.0, (void *)ctx)); 2671 PetscCall(MatViewFromOptions(ctx->J, NULL, "-dm_landau_jacobian_view")); 2672 PetscCall(PetscObjectStateGet((PetscObject)ctx->J, &state)); 2673 ctx->norm_state = state; 2674 } else { 2675 PetscCall(PetscInfo(ts, "WARNING Skip forming Jacobian, has not changed %" PetscInt64_FMT "\n", state)); 2676 } 2677 /* mat vec for op */ 2678 PetscCall(MatMult(ctx->J, X, F)); /* C*f */ 2679 /* add time term */ 2680 if (X_t) PetscCall(MatMultAdd(ctx->M, X_t, F, F)); 2681 #if defined(PETSC_HAVE_THREADSAFETY) 2682 if (ctx->stage) { 2683 endtime = MPI_Wtime(); 2684 ctx->times[LANDAU_OPERATOR] += (endtime - starttime); 2685 ctx->times[LANDAU_JACOBIAN] += (endtime - starttime); 2686 ctx->times[LANDAU_MATRIX_TOTAL] += (endtime - starttime); 2687 ctx->times[LANDAU_JACOBIAN_COUNT] += 1; 2688 } 2689 #endif 2690 PetscCall(PetscLogEventEnd(ctx->events[0], 0, 0, 0, 0)); 2691 PetscCall(PetscLogEventEnd(ctx->events[11], 0, 0, 0, 0)); 2692 if (ctx->stage) PetscCall(PetscLogStagePop()); 2693 PetscFunctionReturn(PETSC_SUCCESS); 2694 } 2695 2696 /*@ 2697 DMPlexLandauIJacobian - `TS` Jacobian construction, confusingly this adds mass 2698 2699 Collective 2700 2701 Input Parameters: 2702 + ts - The time stepping context 2703 . time_dummy - current time (not used) 2704 . X - Current state 2705 . U_tdummy - Time derivative of current state (not used) 2706 . shift - shift for du/dt term 2707 - actx - Landau context 2708 2709 Output Parameters: 2710 + Amat - Jacobian 2711 - Pmat - same as Amat 2712 2713 Level: beginner 2714 2715 .seealso: `DMPlexLandauCreateVelocitySpace()`, `DMPlexLandauIFunction()` 2716 @*/ 2717 PetscErrorCode DMPlexLandauIJacobian(TS ts, PetscReal time_dummy, Vec X, Vec U_tdummy, PetscReal shift, Mat Amat, Mat Pmat, void *actx) 2718 { 2719 LandauCtx *ctx = NULL; 2720 PetscInt dim; 2721 DM pack; 2722 #if defined(PETSC_HAVE_THREADSAFETY) 2723 double starttime, endtime; 2724 #endif 2725 PetscObjectState state; 2726 2727 PetscFunctionBegin; 2728 PetscCall(TSGetDM(ts, &pack)); 2729 PetscCall(DMGetApplicationContext(pack, &ctx)); 2730 PetscCheck(ctx, PETSC_COMM_SELF, PETSC_ERR_PLIB, "no context"); 2731 PetscCheck(Amat == Pmat && Amat == ctx->J, ctx->comm, PETSC_ERR_PLIB, "Amat!=Pmat || Amat!=ctx->J"); 2732 PetscCall(DMGetDimension(pack, &dim)); 2733 /* get collision Jacobian into A */ 2734 if (ctx->stage) PetscCall(PetscLogStagePush(ctx->stage)); 2735 PetscCall(PetscLogEventBegin(ctx->events[11], 0, 0, 0, 0)); 2736 PetscCall(PetscLogEventBegin(ctx->events[9], 0, 0, 0, 0)); 2737 #if defined(PETSC_HAVE_THREADSAFETY) 2738 starttime = MPI_Wtime(); 2739 #endif 2740 PetscCheck(shift != 0.0, ctx->comm, PETSC_ERR_PLIB, "zero shift"); 2741 PetscCall(PetscObjectStateGet((PetscObject)ctx->J, &state)); 2742 PetscCheck(state == ctx->norm_state, ctx->comm, PETSC_ERR_PLIB, "wrong state, %" PetscInt64_FMT " %" PetscInt64_FMT, ctx->norm_state, state); 2743 if (!ctx->use_matrix_mass) { 2744 PetscCall(LandauFormJacobian_Internal(X, ctx->J, dim, shift, (void *)ctx)); 2745 } else { /* add mass */ 2746 PetscCall(MatAXPY(Pmat, shift, ctx->M, SAME_NONZERO_PATTERN)); 2747 } 2748 #if defined(PETSC_HAVE_THREADSAFETY) 2749 if (ctx->stage) { 2750 endtime = MPI_Wtime(); 2751 ctx->times[LANDAU_OPERATOR] += (endtime - starttime); 2752 ctx->times[LANDAU_MASS] += (endtime - starttime); 2753 ctx->times[LANDAU_MATRIX_TOTAL] += (endtime - starttime); 2754 } 2755 #endif 2756 PetscCall(PetscLogEventEnd(ctx->events[9], 0, 0, 0, 0)); 2757 PetscCall(PetscLogEventEnd(ctx->events[11], 0, 0, 0, 0)); 2758 if (ctx->stage) PetscCall(PetscLogStagePop()); 2759 PetscFunctionReturn(PETSC_SUCCESS); 2760 } 2761