diff --git a/methods/conservedquantitycomputation.py b/methods/conservedquantitycomputation.py index 53958b269b7f8c1ab86c47f6862e7e64adafd708..9f5fb0b47a1518cf242d56c2ed35e6196abd22b9 100644 --- a/methods/conservedquantitycomputation.py +++ b/methods/conservedquantitycomputation.py @@ -177,7 +177,7 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation): if 'density' in output_quantity_names_to_symbols: density_output_symbol = output_quantity_names_to_symbols['density'] if isinstance(density_output_symbol, Field): - density_output_symbol = density_output_symbol() + density_output_symbol = density_output_symbol.center if density_output_symbol != self._symbolOrder0: main_assignments.append(Assignment(density_output_symbol, self._symbolOrder0)) else: @@ -186,8 +186,7 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation): if 'velocity' in output_quantity_names_to_symbols: vel_output_symbols = output_quantity_names_to_symbols['velocity'] if isinstance(vel_output_symbols, Field): - field = vel_output_symbols - vel_output_symbols = [field(i) for i in range(len(self._symbolsOrder1))] + vel_output_symbols = vel_output_symbols.center_vector if tuple(vel_output_symbols) != tuple(self._symbolsOrder1): main_assignments += [Assignment(a, b) for a, b in zip(vel_output_symbols, self._symbolsOrder1)] else: @@ -208,6 +207,17 @@ class DensityVelocityComputation(AbstractConservedQuantityComputation): self._forceModel, self._compressible) for sym, val in zip(momentum_density_output_symbols, mom_density_eq_coll.main_assignments[1:]): main_assignments.append(Assignment(sym, val.rhs)) + if 'moment0' in output_quantity_names_to_symbols: + moment0_output_symbol = output_quantity_names_to_symbols['moment0'] + if isinstance(moment0_output_symbol, Field): + moment0_output_symbol = moment0_output_symbol.center + main_assignments.append(Assignment(moment0_output_symbol, sum(pdfs))) + if 'moment1' in output_quantity_names_to_symbols: + moment1_output_symbol = output_quantity_names_to_symbols['moment1'] + if isinstance(moment1_output_symbol, Field): + moment1_output_symbol = moment1_output_symbol.center_vector + main_assignments.extend([Assignment(lhs, sum(d[i] * pdf for d, pdf in zip(self._stencil, pdfs))) + for i, lhs in enumerate(moment1_output_symbol)]) ac = ac.copy(main_assignments, [Assignment(a, b) for a, b in eqs.items()]) return ac.new_without_unused_subexpressions() diff --git a/methods/entropic_eq_srt.py b/methods/entropic_eq_srt.py index 7a7b039cd146dc54b0150943d9b7199072a7ef27..07afe6488d4f8bde27bd12faeea3ff51328d47db 100644 --- a/methods/entropic_eq_srt.py +++ b/methods/entropic_eq_srt.py @@ -39,6 +39,10 @@ class EntropicEquilibriumSRT(AbstractLbMethod): conserved_quantity_equations=conserved_quantity_equations, include_force_terms=include_force_terms) + def get_equilibrium_terms(self): + equilibrium = self.get_equilibrium() + return sp.Matrix([eq.rhs for eq in equilibrium.main_assignments]) + def _get_collision_rule_with_relaxation_rate(self, relaxation_rate, include_force_terms=True, conserved_quantity_equations=None): f = sp.Matrix(self.pre_collision_pdf_symbols) @@ -72,8 +76,9 @@ class EntropicEquilibriumSRT(AbstractLbMethod): cr.simplification_hints['relaxation_rates'] = [] return cr - def get_collision_rule(self): - return self._get_collision_rule_with_relaxation_rate(self._relaxationRate) + def get_collision_rule(self, conserved_quantity_equations=None): + return self._get_collision_rule_with_relaxation_rate(self._relaxationRate, + conserved_quantity_equations=conserved_quantity_equations) def create_srt_entropic(stencil, relaxation_rate, force_model, compressible):