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