lbmpy: output of raw zeroth and first moment

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()

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):