Line data Source code
1 : /**********************************************************************/
2 : /* DO NOT MODIFY THIS HEADER */
3 : /* Swift, a Fourier spectral solver for MOOSE */
4 : /* */
5 : /* Copyright 2024 Battelle Energy Alliance, LLC */
6 : /* ALL RIGHTS RESERVED */
7 : /**********************************************************************/
8 :
9 : #include "LBMCollisionDynamics.h"
10 :
11 : registerMooseObject("SwiftApp", LBMBGKCollision);
12 : registerMooseObject("SwiftApp", LBMMRTCollision);
13 : registerMooseObject("SwiftApp", LBMSmagorinskyCollision);
14 : registerMooseObject("SwiftApp", LBMSmagorinskyMRTCollision);
15 :
16 : template <int coll_dyn>
17 : InputParameters
18 0 : LBMCollisionDynamicsTempl<coll_dyn>::validParams()
19 : {
20 0 : InputParameters params = LatticeBoltzmannOperator::validParams();
21 :
22 0 : params.addClassDescription("Template object for LBM collision dynamics");
23 0 : params.addRequiredParam<TensorInputBufferName>("f", "Input buffer distribution function");
24 0 : params.addRequiredParam<TensorInputBufferName>("feq",
25 : "Input buffer equilibrium distribution function");
26 0 : params.addRequiredParam<std::string>("tau0", "Relaxation parameter");
27 0 : params.addParam<std::string>("Cs", "0.1", "Relaxation parameter");
28 0 : params.addParam<bool>(
29 0 : "projection", false, "Whether or not to project non-equilibrium onto Hermite space.");
30 :
31 0 : return params;
32 0 : }
33 :
34 : template <int coll_dyn>
35 0 : LBMCollisionDynamicsTempl<coll_dyn>::LBMCollisionDynamicsTempl(const InputParameters & parameters)
36 : : LatticeBoltzmannOperator(parameters),
37 0 : _f(getInputBuffer("f")),
38 0 : _feq(getInputBuffer("feq")),
39 0 : _shape(_lb_problem.getGridSize()),
40 0 : _tau_0(_lb_problem.getConstant<Real>(getParam<std::string>("tau0"))),
41 0 : _C_s(_lb_problem.getConstant<Real>(getParam<std::string>("Cs"))),
42 0 : _delta_x(1.0),
43 0 : _projection(getParam<bool>("projection"))
44 : {
45 : //
46 0 : _fneq = torch::zeros({_shape[0], _shape[1], _shape[2], _stencil._q},
47 : MooseTensor::floatTensorOptions());
48 0 : }
49 :
50 : template <int coll_dyn>
51 : void
52 0 : LBMCollisionDynamicsTempl<coll_dyn>::HermiteRegularization()
53 : {
54 : /**
55 : * Regularization procedure projects non-equilibrium (fneq) distribution
56 : * onto the second order Hermite space.
57 : */
58 :
59 : using torch::indexing::Slice;
60 :
61 0 : int64_t nx = _shape[0];
62 0 : int64_t ny = _shape[1];
63 0 : int64_t nz = _shape[2];
64 :
65 0 : auto f_flat = _f.view({nx * ny * nz, _stencil._q});
66 0 : auto feq_flat = _feq.view({nx * ny * nz, _stencil._q});
67 0 : auto f_neq_hat = _fneq.view({nx * ny * nz, _stencil._q});
68 :
69 0 : torch::Tensor fneq = f_flat - feq_flat;
70 0 : torch::Tensor fneqtimescc = torch::zeros({nx * ny * nz, 9}, MooseTensor::floatTensorOptions());
71 0 : torch::Tensor e_xyz = torch::stack({_stencil._ex, _stencil._ey, _stencil._ez}, 0);
72 :
73 0 : for (int ic = 0; ic < _stencil._q; ic++)
74 : {
75 0 : auto exyz_ic = e_xyz.index({Slice(), ic}).flatten();
76 0 : torch::Tensor ccr = torch::outer(exyz_ic, exyz_ic).flatten();
77 0 : fneqtimescc += (fneq.select(1, ic).view({nx * ny * nz, 1}) * ccr.view({1, 9}));
78 : }
79 :
80 : // Compute Hermite tensor
81 0 : torch::Tensor H2 = torch::zeros({1, 9}, MooseTensor::floatTensorOptions());
82 0 : for (int ic = 0; ic < _stencil._q; ic++)
83 : {
84 0 : auto exyz_ic = e_xyz.index({Slice(), ic}).flatten();
85 0 : torch::Tensor ccr = torch::outer(exyz_ic, exyz_ic) / _lb_problem._cs2 -
86 : torch::eye(3, MooseTensor::floatTensorOptions());
87 0 : H2 = ccr.flatten().unsqueeze(0).expand({nx * ny * nz, 9});
88 :
89 : // Compute regularized non-equilibrium distribution
90 0 : f_neq_hat.index_put_(
91 0 : {Slice(), ic},
92 0 : (_stencil._weights[ic] * (1.0 / (2.0 * _lb_problem._cs2)) * (fneqtimescc * H2).sum(1)));
93 : }
94 :
95 0 : _fneq = f_neq_hat.view({nx, ny, nz, _stencil._q});
96 0 : }
97 :
98 : template <int coll_dyn>
99 : void
100 0 : LBMCollisionDynamicsTempl<coll_dyn>::computeRelaxationParameter()
101 : {
102 0 : int64_t nx = _shape[0];
103 0 : int64_t ny = _shape[1];
104 0 : int64_t nz = _shape[2];
105 :
106 0 : auto f_neq_hat = _fneq.view({nx * ny * nz, _stencil._q, 1, 1, 1});
107 :
108 : torch::Tensor S;
109 : {
110 : torch::Tensor outer_products;
111 : {
112 : torch::Tensor ex_2d;
113 : torch::Tensor ey_2d;
114 : torch::Tensor ez_2d;
115 :
116 : {
117 0 : auto zeros = torch::zeros({_stencil._q}, MooseTensor::intTensorOptions());
118 0 : auto ones = torch::ones({_stencil._q}, MooseTensor::intTensorOptions());
119 :
120 0 : ex_2d = torch::stack({_stencil._ex, zeros, zeros});
121 0 : ey_2d = torch::stack({zeros, _stencil._ey, zeros});
122 0 : ez_2d = torch::stack({zeros, zeros, _stencil._ez});
123 :
124 0 : if (nz == 1)
125 0 : ez_2d = torch::stack({ones, zeros, _stencil._ez});
126 : } // zeros and ones are out of scope
127 :
128 : // outer product
129 : // expected shape: _q, 3, 3, 3
130 0 : outer_products = torch::zeros({_stencil._q, 3, 3, 3}, MooseTensor::floatTensorOptions());
131 :
132 0 : for (int i = 0; i < _stencil._q; i++)
133 : {
134 0 : auto ex_col = ex_2d.index({Slice(), i});
135 0 : auto ey_col = ey_2d.index({Slice(), i});
136 0 : auto ez_col = ez_2d.index({Slice(), i});
137 0 : auto outer_product = torch::einsum("i,j,k->kij", {ex_col, ey_col, ez_col});
138 0 : outer_products[i] = outer_product;
139 : }
140 0 : outer_products = outer_products.view({1, _stencil._q, 3, 3, 3});
141 : } // ex_2d ey_2d ez_2d are out of scope
142 :
143 : torch::Tensor Q_mean;
144 : {
145 : // momentum flux
146 : torch::Tensor Q;
147 : {
148 : // auto f_neq_outer_prod = torch::einsum("anijk,mnxyz->mijk", {outer_products, f_neq_hat});
149 : // until we figure out a better way to optimzie memory consumption of above commented line
150 : // we will do this in smaller batches
151 : // this will take slightly longer .... ??
152 :
153 : const int64_t M = nx * ny * nz;
154 0 : const int64_t batch_size = M / 20; // just pulled out of my ***
155 0 : torch::Tensor f_neq_outer_prod_batched =
156 : torch::zeros({M, 3, 3, 3}, MooseTensor::floatTensorOptions());
157 :
158 0 : for (int64_t i = 0; i < M; i += batch_size)
159 : {
160 0 : int64_t current_batch_size = std::min(batch_size, M - i);
161 0 : torch::Tensor f_neq_hat_batch = f_neq_hat.slice(0, i, i + current_batch_size);
162 :
163 0 : torch::Tensor batch_result =
164 0 : torch::einsum("anijk,mnxyz->mijk", {outer_products, f_neq_hat_batch});
165 :
166 0 : f_neq_outer_prod_batched.slice(0, i, i + current_batch_size).copy_(batch_result);
167 : }
168 0 : Q = f_neq_outer_prod_batched.view({nx, ny, nz, 3, 3, 3});
169 : }
170 :
171 : // mean density
172 0 : _mean_density = torch::mean(torch::sum(_f, 3)).item<double>();
173 :
174 : // Frobenius norm
175 0 : Q_mean = torch::norm(Q, 2, {3, 4, 5}) / (_mean_density * _lb_problem._cs2);
176 : } // auto Q goes of scopoe
177 :
178 : // subgrid time scale factor
179 0 : auto t_sgs = sqrt(_C_s) * _delta_x / _lb_problem._cs;
180 0 : auto eta = _tau_0 / t_sgs;
181 :
182 : // mean strain rate
183 0 : auto Q_mean_sqrt = torch::sqrt(eta * eta + 4.0 * Q_mean);
184 0 : auto eta_Q_mean_sqrt = (-1.0 * eta + Q_mean_sqrt);
185 0 : S = eta_Q_mean_sqrt / (2.0 * t_sgs);
186 : } // outer_products is out of scope
187 :
188 : // relaxation parameter
189 0 : _relaxation_parameter = _tau_0 + _C_s * _delta_x * _delta_x * S / _lb_problem._cs2;
190 0 : _relaxation_parameter = _relaxation_parameter.view({nx, ny, nz, 1});
191 0 : }
192 :
193 : template <int coll_dyn>
194 : void
195 0 : LBMCollisionDynamicsTempl<coll_dyn>::computeLocalRelaxationMatrix()
196 : {
197 0 : if (_lb_problem.getTotalSteps() == 0)
198 : {
199 0 : std::array<int64_t, 5> local_relaxation_mrt_shape = {
200 0 : _shape[0], _shape[1], _shape[2], _stencil._q, _stencil._q};
201 :
202 0 : _local_relaxation_matrix =
203 : torch::zeros(local_relaxation_mrt_shape, MooseTensor::floatTensorOptions());
204 0 : torch::Tensor stencil_S_expanded = _stencil._S.view({_stencil._q, _stencil._q}).clone();
205 :
206 0 : stencil_S_expanded = stencil_S_expanded.unsqueeze(0).unsqueeze(0).unsqueeze(0);
207 0 : _local_relaxation_matrix = stencil_S_expanded.expand(local_relaxation_mrt_shape).clone();
208 : }
209 :
210 0 : for (int64_t sh_id = 0; sh_id < _stencil._id_kinematic_visc.size(0); sh_id++)
211 0 : _local_relaxation_matrix.index_put_({Slice(),
212 0 : Slice(),
213 0 : Slice(),
214 0 : _stencil._id_kinematic_visc[sh_id],
215 0 : _stencil._id_kinematic_visc[sh_id]},
216 : 1.0 / _relaxation_parameter.squeeze(-1));
217 0 : }
218 :
219 : template <int coll_dyn>
220 : void
221 0 : LBMCollisionDynamicsTempl<coll_dyn>::computeGlobalRelaxationMatrix()
222 : {
223 0 : if (_lb_problem.getTotalSteps() == 0)
224 : {
225 0 : _global_relaxation_matrix = _stencil._S.clone();
226 :
227 0 : _global_relaxation_matrix.index_put_({_stencil._id_kinematic_visc, _stencil._id_kinematic_visc},
228 0 : 1.0 / _tau_0);
229 : }
230 0 : }
231 :
232 : template <>
233 : void
234 0 : LBMCollisionDynamicsTempl<0>::BGKDynamics()
235 : {
236 : /* LBM BGK collision */
237 0 : _u = _feq + _fneq - 1.0 / _tau_0 * _fneq;
238 0 : _lb_problem.maskedFillSolids(_u, 0);
239 0 : }
240 :
241 : template <>
242 : void
243 0 : LBMCollisionDynamicsTempl<1>::MRTDynamics()
244 : {
245 0 : computeGlobalRelaxationMatrix();
246 :
247 : /* LBM MRT collision */
248 : // f = M^{-1} x S x M x (f - feq)
249 0 : auto m_neq = torch::einsum("ab,ijkb->ijka", {_stencil._M, _fneq});
250 0 : auto m_neq_relaxed = torch::einsum("ab,ijkb->ijka", {_global_relaxation_matrix, m_neq});
251 0 : auto f = torch::einsum("ab,ijkb->ijka", {_stencil._M_inv, m_neq_relaxed});
252 :
253 0 : _u = _feq + _fneq - f;
254 :
255 0 : _lb_problem.maskedFillSolids(_u, 0);
256 0 : }
257 :
258 : template <>
259 : void
260 0 : LBMCollisionDynamicsTempl<2>::SmagorinskyDynamics()
261 : {
262 0 : computeRelaxationParameter();
263 :
264 : // BGK collision
265 0 : _u = _feq + _fneq - 1.0 / _relaxation_parameter * _fneq;
266 :
267 0 : _lb_problem.maskedFillSolids(_u, 0);
268 0 : }
269 :
270 : template <>
271 : void
272 0 : LBMCollisionDynamicsTempl<3>::SmagorinskyMRTDynamics()
273 : {
274 0 : computeRelaxationParameter();
275 0 : computeLocalRelaxationMatrix();
276 :
277 : /* LBM MRT collision */
278 :
279 0 : auto m_neq = torch::einsum("ab,ijkb->ijka", {_stencil._M, _fneq});
280 0 : auto m_neq_relaxed = torch::einsum("ijklm,ijkm->ijkl", {_local_relaxation_matrix, m_neq});
281 0 : auto f = torch::einsum("ab,ijkb->ijka", {_stencil._M_inv, m_neq_relaxed});
282 :
283 : // f = M^{-1} x S x M x (f - feq)
284 0 : _u = _feq + _fneq - f;
285 0 : }
286 :
287 : template <int coll_dyn>
288 : void
289 0 : LBMCollisionDynamicsTempl<coll_dyn>::computeBuffer()
290 : {
291 0 : if (_projection)
292 0 : HermiteRegularization();
293 : else
294 0 : _fneq = _f - _feq;
295 :
296 : switch (coll_dyn)
297 : {
298 0 : case 0:
299 0 : BGKDynamics();
300 : break;
301 0 : case 1:
302 0 : MRTDynamics();
303 : break;
304 0 : case 2:
305 0 : SmagorinskyDynamics();
306 : break;
307 0 : case 3:
308 0 : SmagorinskyMRTDynamics();
309 : break;
310 : default:
311 : mooseError("Undefined template value");
312 : }
313 0 : _lb_problem.maskedFillSolids(_u, 0);
314 0 : }
315 :
316 : template class LBMCollisionDynamicsTempl<0>;
317 : template class LBMCollisionDynamicsTempl<1>;
318 : template class LBMCollisionDynamicsTempl<2>;
319 : template class LBMCollisionDynamicsTempl<3>;
|
