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