Actual source code: petsclandau.h
1: #pragma once
3: #include <petscdmplex.h>
4: #include <petscts.h>
6: PETSC_EXTERN PetscErrorCode DMPlexLandauPrintNorms(Vec, PetscInt);
7: PETSC_EXTERN PetscErrorCode DMPlexLandauCreateVelocitySpace(MPI_Comm, PetscInt, const char[], Vec *, Mat *, DM *);
8: PETSC_EXTERN PetscErrorCode DMPlexLandauDestroyVelocitySpace(DM *);
9: PETSC_EXTERN PetscErrorCode DMPlexLandauAccess(DM, Vec, PetscErrorCode (*)(DM, Vec, PetscInt, PetscInt, PetscInt, void *), void *);
10: PETSC_EXTERN PetscErrorCode DMPlexLandauAddMaxwellians(DM, Vec, PetscReal, PetscReal[], PetscReal[], PetscInt, PetscInt, PetscInt, void *);
11: PETSC_EXTERN PetscErrorCode DMPlexLandauCreateMassMatrix(DM dm, Mat *Amat);
12: PETSC_EXTERN PetscErrorCode DMPlexLandauIFunction(TS, PetscReal, Vec, Vec, Vec, void *);
13: PETSC_EXTERN PetscErrorCode DMPlexLandauIJacobian(TS, PetscReal, Vec, Vec, PetscReal, Mat, Mat, void *);
15: typedef PetscInt LandauIdx;
17: /* the Fokker-Planck-Landau context */
18: #if !defined(LANDAU_MAX_SPECIES)
19: #if defined(PETSC_USE_DMLANDAU_2D)
20: #define LANDAU_MAX_SPECIES 10
21: #define LANDAU_MAX_GRIDS 3
22: #else
23: #define LANDAU_MAX_SPECIES 10
24: #define LANDAU_MAX_GRIDS 3
25: #endif
26: #else
27: #define LANDAU_MAX_GRIDS 3
28: #endif
30: #if !defined(LANDAU_MAX_Q)
31: #if defined(LANDAU_MAX_NQND)
32: #error "LANDAU_MAX_NQND but not LANDAU_MAX_Q. Use -DLANDAU_MAX_Q=4 for Q3 elements"
33: #endif
34: #if defined(PETSC_USE_DMLANDAU_2D)
35: #define LANDAU_MAX_Q 6
36: #else
37: #define LANDAU_MAX_Q 5
38: #endif
39: #else
40: #undef LANDAU_MAX_NQND
41: #endif
43: #if defined(PETSC_USE_DMLANDAU_2D)
44: #define LANDAU_MAX_Q_FACE LANDAU_MAX_Q
45: #define LANDAU_MAX_NQND (LANDAU_MAX_Q * LANDAU_MAX_Q)
46: #define LANDAU_MAX_BATCH_SZ 1024
47: #define LANDAU_DIM 2
48: #else
49: #define LANDAU_MAX_Q_FACE (LANDAU_MAX_Q * LANDAU_MAX_Q)
50: #define LANDAU_MAX_NQND (LANDAU_MAX_Q * LANDAU_MAX_Q * LANDAU_MAX_Q)
51: #define LANDAU_MAX_BATCH_SZ 64
52: #define LANDAU_DIM 3
53: #endif
55: typedef enum {
56: LANDAU_KOKKOS,
57: LANDAU_CPU
58: } LandauDeviceType;
60: // static data - will be "device" data
61: typedef struct {
62: void *invJ; // nip*dim*dim
63: void *D; // nq*nb*dim
64: void *B; // nq*nb
65: void *alpha; // ns
66: void *beta; // ns
67: void *invMass; // ns
68: void *w; // nip
69: void *x; // nip
70: void *y; // nip
71: void *z; // nip
72: void *Eq_m; // ns - dynamic
73: void *f; // nip*Nf - dynamic (IP)
74: void *dfdx; // nip*Nf - dynamic (IP)
75: void *dfdy; // nip*Nf - dynamic (IP)
76: void *dfdz; // nip*Nf - dynamic (IP)
77: int dim_, ns_, nip_, nq_, nb_;
78: void *NCells; // remove and use elem_offset - TODO
79: void *species_offset; // for each grid, but same for all batched vertices
80: void *mat_offset; // for each grid, but same for all batched vertices
81: void *elem_offset; // for each grid, but same for all batched vertices
82: void *ip_offset; // for each grid, but same for all batched vertices
83: void *ipf_offset; // for each grid, but same for all batched vertices
84: void *ipfdf_data; // for each grid, but same for all batched vertices
85: void *maps; // for each grid, but same for all batched vertices
86: // COO
87: void *coo_elem_offsets;
88: void *coo_elem_point_offsets;
89: void *coo_elem_fullNb;
90: void *coo_vals;
91: void *lambdas;
92: LandauIdx coo_n_cellsTot;
93: LandauIdx coo_size;
94: LandauIdx coo_max_fullnb;
95: } LandauStaticData;
97: typedef enum {
98: LANDAU_EX2_TSSOLVE,
99: LANDAU_MATRIX_TOTAL,
100: LANDAU_OPERATOR,
101: LANDAU_JACOBIAN_COUNT,
102: LANDAU_JACOBIAN,
103: LANDAU_MASS,
104: LANDAU_F_DF,
105: LANDAU_KERNEL,
106: KSP_FACTOR,
107: KSP_SOLVE,
108: LANDAU_NUM_TIMERS
109: } LandauOMPTimers;
111: typedef struct {
112: PetscBool interpolate; /* Generate intermediate mesh elements */
113: PetscBool gpu_assembly;
114: MPI_Comm comm; /* global communicator to use for errors and diagnostics */
115: double times[LANDAU_NUM_TIMERS];
116: PetscBool use_matrix_mass;
117: /* FE */
118: PetscFE fe[LANDAU_MAX_SPECIES];
119: /* geometry */
120: PetscReal radius[LANDAU_MAX_GRIDS];
121: PetscReal radius_par[LANDAU_MAX_GRIDS];
122: PetscReal radius_perp[LANDAU_MAX_GRIDS];
123: PetscReal re_radius; /* RE: radius of refinement along v_perp=0, z>0 */
124: PetscReal vperp0_radius1; /* RE: radius of refinement along v_perp=0 */
125: PetscReal vperp0_radius2; /* RE: radius of refinement along v_perp=0 after origin AMR refinement */
126: PetscBool sphere;
127: PetscReal sphere_inner_radius_90degree[LANDAU_MAX_GRIDS];
128: PetscReal sphere_inner_radius_45degree[LANDAU_MAX_GRIDS];
129: PetscInt cells0[3];
130: /* AMR */
131: PetscBool use_p4est;
132: PetscInt numRERefine; /* RE: refinement along v_perp=0, z > 0 */
133: PetscInt nZRefine1; /* RE: origin refinement after v_perp=0 refinement */
134: PetscInt nZRefine2; /* RE: origin refinement after origin AMR refinement */
135: PetscInt numAMRRefine[LANDAU_MAX_GRIDS]; /* normal AMR - refine from origin */
136: PetscInt postAMRRefine[LANDAU_MAX_GRIDS]; /* uniform refinement of AMR */
137: PetscBool simplex;
138: char filename[PETSC_MAX_PATH_LEN];
139: PetscReal thermal_speed[LANDAU_MAX_GRIDS];
140: PetscBool sphere_uniform_normal;
141: /* relativistic */
142: PetscBool use_energy_tensor_trick;
143: PetscBool use_relativistic_corrections;
144: /* physics */
145: PetscReal thermal_temps[LANDAU_MAX_SPECIES];
146: PetscReal masses[LANDAU_MAX_SPECIES]; /* mass of each species */
147: PetscReal charges[LANDAU_MAX_SPECIES]; /* charge of each species */
148: PetscReal n[LANDAU_MAX_SPECIES]; /* number density of each species */
149: PetscReal m_0; /* reference mass */
150: PetscReal v_0; /* reference velocity */
151: PetscReal n_0; /* reference number density */
152: PetscReal t_0; /* reference time */
153: PetscReal Ez;
154: PetscReal epsilon0;
155: PetscReal k;
156: PetscReal lambdas[LANDAU_MAX_GRIDS][LANDAU_MAX_GRIDS];
157: PetscReal electronShift;
158: PetscInt num_species;
159: PetscInt num_grids;
160: PetscInt species_offset[LANDAU_MAX_GRIDS + 1]; // for each grid, but same for all batched vertices
161: PetscInt mat_offset[LANDAU_MAX_GRIDS + 1]; // for each grid, but same for all batched vertices
162: // batching
163: PetscBool jacobian_field_major_order; // this could be a type but lets not get pedantic
164: VecScatter plex_batch;
165: Vec work_vec;
166: IS batch_is;
167: PetscErrorCode (*seqaij_mult)(Mat, Vec, Vec);
168: PetscErrorCode (*seqaij_multtranspose)(Mat, Vec, Vec);
169: PetscErrorCode (*seqaij_solve)(Mat, Vec, Vec);
170: PetscErrorCode (*seqaij_getdiagonal)(Mat, Vec);
171: /* COO */
172: Mat J;
173: Mat M;
174: Vec X;
175: /* derived type */
176: void *data;
177: /* computing */
178: LandauDeviceType deviceType;
179: DM pack;
180: DM plex[LANDAU_MAX_GRIDS];
181: LandauStaticData SData_d; /* static geometric data on device */
182: /* diagnostics */
183: PetscInt verbose;
184: PetscLogEvent events[20];
185: PetscLogStage stage;
186: PetscObjectState norm_state;
187: PetscInt batch_sz;
188: PetscInt batch_view_idx;
189: } LandauCtx;
191: #define LANDAU_SPECIES_MAJOR
192: #if !defined(LANDAU_SPECIES_MAJOR)
193: #define LAND_PACK_IDX(_b, _g) (_b * ctx->num_grids + _g)
194: #define LAND_MOFFSET(_b, _g, _nbch, _ngrid, _mat_off) (_b * _mat_off[_ngrid] + _mat_off[_g])
195: #else
196: #define LAND_PACK_IDX(_b, _g) (_g * ctx->batch_sz + _b)
197: #define LAND_MOFFSET(_b, _g, _nbch, _ngrid, _mat_off) (_nbch * _mat_off[_g] + _b * (_mat_off[_g + 1] - _mat_off[_g]))
198: #endif
200: typedef struct {
201: PetscReal scale;
202: LandauIdx gid; // Landau matrix index (<10,000)
203: } pointInterpolationP4est;
204: typedef struct _lP4estVertexMaps {
205: LandauIdx (*gIdx)[LANDAU_MAX_SPECIES][LANDAU_MAX_NQND]; // #elems * LANDAU_MAX_NQND
206: LandauIdx num_elements;
207: LandauIdx num_reduced;
208: LandauIdx num_face; // (Q or Q^2 for 3D)
209: LandauDeviceType deviceType;
210: PetscInt Nf;
211: pointInterpolationP4est (*c_maps)[LANDAU_MAX_Q_FACE];
212: struct _lP4estVertexMaps *d_self;
213: void *vp1, *vp2, *vp3;
214: PetscInt numgrids;
215: } P4estVertexMaps;
217: #if defined(PETSC_HAVE_KOKKOS)
218: PETSC_EXTERN PetscErrorCode LandauKokkosJacobian(DM[], const PetscInt, const PetscInt, const PetscInt, const PetscInt, const PetscInt[], PetscReal[], PetscScalar[], const PetscScalar[], const LandauStaticData *, const PetscReal, const PetscLogEvent[], const PetscInt[], const PetscInt[], Mat[], Mat);
219: PETSC_EXTERN PetscErrorCode LandauKokkosCreateMatMaps(P4estVertexMaps *, pointInterpolationP4est (*)[LANDAU_MAX_Q_FACE], PetscInt[], PetscInt);
220: PETSC_EXTERN PetscErrorCode LandauKokkosDestroyMatMaps(P4estVertexMaps *, PetscInt);
221: PETSC_EXTERN PetscErrorCode LandauKokkosStaticDataSet(DM, const PetscInt, const PetscInt, const PetscInt, const PetscInt, PetscInt[], PetscInt[], PetscInt[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], PetscReal[], LandauStaticData *);
222: PETSC_EXTERN PetscErrorCode LandauKokkosStaticDataClear(LandauStaticData *);
223: #endif