Pass fluxes to elliptic source computers

More involved systems such as XCTS need all variables to compute source terms.
We pass the computed volume fluxes to the source computers as an optimization
so they don't have to be re-computed.

src/Elliptic/DiscontinuousGalerkin/InitializeFirstOrderOperator.hpp

@@ -71,7 +71,7 @@ struct InitializeFirstOrderOperator {
                                               vars_tag, PrimalVariables,
                                               AuxiliaryVariables>;
   using sources_compute_tag = elliptic::Tags::FirstOrderSourcesCompute<
-      SourcesComputer, vars_tag, PrimalVariables, AuxiliaryVariables>;
+      Dim, SourcesComputer, vars_tag, PrimalVariables, AuxiliaryVariables>;
 
   using exterior_tags = tmpl::list<
       // On exterior (ghost) boundary faces we compute the fluxes from the

src/Elliptic/FirstOrderComputeTags.hpp

@@ -46,23 +46,25 @@ struct FirstOrderFluxesCompute
   }
 };
 
-template <typename SourcesComputer, typename VariablesTag,
+template <size_t Dim, typename SourcesComputer, typename VariablesTag,
           typename PrimalVariables, typename AuxiliaryVariables>
 struct FirstOrderSourcesCompute
     : db::add_tag_prefix<::Tags::Source, VariablesTag>,
       db::ComputeTag {
   using base = db::add_tag_prefix<::Tags::Source, VariablesTag>;
   using argument_tags =
-      tmpl::push_front<typename SourcesComputer::argument_tags, VariablesTag>;
+      tmpl::push_front<typename SourcesComputer::argument_tags, VariablesTag,
+                       db::add_tag_prefix<::Tags::Flux, VariablesTag,
+                                          tmpl::size_t<Dim>, Frame::Inertial>>;
   using volume_tags = get_volume_tags<SourcesComputer>;
   using return_type = typename base::type;
-  template <typename... SourcesArgs>
+  template <typename Vars, typename Fluxes, typename... SourcesArgs>
   static void function(const gsl::not_null<return_type*> sources,
-                       const typename VariablesTag::type& vars,
+                       const Vars& vars, const Fluxes& fluxes,
                        const SourcesArgs&... sources_args) noexcept {
     *sources = return_type{vars.number_of_grid_points()};
-    elliptic::first_order_sources<PrimalVariables, AuxiliaryVariables,
-                                  SourcesComputer>(sources, vars,
+    elliptic::first_order_sources<Dim, PrimalVariables, AuxiliaryVariables,
+                                  SourcesComputer>(sources, vars, fluxes,
                                                    sources_args...);
   }
 };

src/Elliptic/FirstOrderOperator.hpp

@@ -43,34 +43,44 @@ struct FirstOrderFluxesImpl<Dim, tmpl::list<PrimalFields...>,
   }
 };
 
-template <typename PrimalFields, typename AuxiliaryFields,
+template <size_t Dim, typename PrimalFields, typename AuxiliaryFields,
           typename SourcesComputer>
 struct FirstOrderSourcesImpl;
 
-template <typename... PrimalFields, typename... AuxiliaryFields,
+template <size_t Dim, typename... PrimalFields, typename... AuxiliaryFields,
           typename SourcesComputer>
-struct FirstOrderSourcesImpl<tmpl::list<PrimalFields...>,
+struct FirstOrderSourcesImpl<Dim, tmpl::list<PrimalFields...>,
                              tmpl::list<AuxiliaryFields...>, SourcesComputer> {
   template <typename VarsTags, typename... SourcesArgs>
   static constexpr void apply(
       const gsl::not_null<
           Variables<db::wrap_tags_in<::Tags::Source, VarsTags>>*>
           sources,
       const Variables<VarsTags>& vars,
+      const Variables<db::wrap_tags_in<
+          ::Tags::Flux, VarsTags, tmpl::size_t<Dim>, Frame::Inertial>>& fluxes,
       const SourcesArgs&... sources_args) noexcept {
-    // Compute sources for primal fields
-    SourcesComputer::apply(
-        make_not_null(&get<::Tags::Source<PrimalFields>>(*sources))...,
-        sources_args..., get<PrimalFields>(vars)...);
-    // Compute sources for auxiliary fields. They are just the auxiliary field
-    // values.
+    // Set auxiliary field sources to the auxiliary field values to begin with.
+    // This is the standard choice, since the auxiliary equations define the
+    // auxiliary variables. Source computers can adjust or add to these sources.
     tmpl::for_each<tmpl::list<AuxiliaryFields...>>(
-        [&sources, &vars ](const auto auxiliary_field_tag_v) noexcept {
+        [&sources, &vars](const auto auxiliary_field_tag_v) noexcept {
           using auxiliary_field_tag =
               tmpl::type_from<decltype(auxiliary_field_tag_v)>;
           get<::Tags::Source<auxiliary_field_tag>>(*sources) =
               get<auxiliary_field_tag>(vars);
         });
+    // Call into the sources computer to set primal field sources and possibly
+    // adjust auxiliary field sources. The sources depend on the primal and the
+    // auxiliary variables. However, we pass the volume fluxes instead of the
+    // auxiliary variables to the source computer as an optimization so they
+    // don't have to be re-computed.
+    SourcesComputer::apply(
+        make_not_null(&get<::Tags::Source<PrimalFields>>(*sources))...,
+        make_not_null(&get<::Tags::Source<AuxiliaryFields>>(*sources))...,
+        sources_args..., get<PrimalFields>(vars)...,
+        get<::Tags::Flux<PrimalFields, tmpl::size_t<Dim>, Frame::Inertial>>(
+            fluxes)...);
   }
 };
 
@@ -118,30 +128,41 @@ auto first_order_fluxes(const Variables<VarsTags>& vars,
  * \brief Compute the sources \f$S(u)\f$ for the first-order formulation of
  * elliptic systems.
  *
+ * This function takes the `fluxes` as an argument in addition to the variables
+ * as an optimization. The fluxes will generally be computed before the sources
+ * anyway, so we pass them to the source computers to avoid having to re-compute
+ * them for source-terms that have the same form as the fluxes.
+ *
  * \see `elliptic::first_order_operator`
  */
-template <typename PrimalFields, typename AuxiliaryFields,
+template <size_t Dim, typename PrimalFields, typename AuxiliaryFields,
           typename SourcesComputer, typename VarsTags, typename... SourcesArgs>
 void first_order_sources(
     gsl::not_null<Variables<db::wrap_tags_in<::Tags::Source, VarsTags>>*>
         sources,
     const Variables<VarsTags>& vars,
+    const Variables<db::wrap_tags_in<::Tags::Flux, VarsTags, tmpl::size_t<Dim>,
+                                     Frame::Inertial>>& fluxes,
     const SourcesArgs&... sources_args) noexcept {
   first_order_operator_detail::FirstOrderSourcesImpl<
-      PrimalFields, AuxiliaryFields, SourcesComputer>::apply(std::move(sources),
-                                                             vars,
-                                                             sources_args...);
+      Dim, PrimalFields, AuxiliaryFields,
+      SourcesComputer>::apply(std::move(sources), vars, fluxes,
+                              sources_args...);
 }
 
-template <typename PrimalFields, typename AuxiliaryFields,
+template <size_t Dim, typename PrimalFields, typename AuxiliaryFields,
           typename SourcesComputer, typename VarsTags, typename... SourcesArgs>
-auto first_order_sources(const Variables<VarsTags>& vars,
-                         const SourcesArgs&... sources_args) noexcept {
+auto first_order_sources(
+    const Variables<VarsTags>& vars,
+    const Variables<db::wrap_tags_in<::Tags::Flux, VarsTags, tmpl::size_t<Dim>,
+                                     Frame::Inertial>>& fluxes,
+    const SourcesArgs&... sources_args) noexcept {
   Variables<db::wrap_tags_in<::Tags::Source, VarsTags>> sources{
       vars.number_of_grid_points()};
   first_order_operator_detail::FirstOrderSourcesImpl<
-      PrimalFields, AuxiliaryFields,
-      SourcesComputer>::apply(make_not_null(&sources), vars, sources_args...);
+      Dim, PrimalFields, AuxiliaryFields,
+      SourcesComputer>::apply(make_not_null(&sources), vars, fluxes,
+                              sources_args...);
   return sources;
 }
 // @}

src/Elliptic/Systems/Elasticity/Equations.hpp

@@ -94,7 +94,9 @@ struct Sources {
   using argument_tags = tmpl::list<>;
   static void apply(
       const gsl::not_null<tnsr::I<DataVector, Dim>*> source_for_displacement,
-      const tnsr::I<DataVector, Dim>& /*displacement*/) noexcept {
+      const gsl::not_null<tnsr::ii<DataVector, Dim>*> /*source_for_strain*/,
+      const tnsr::I<DataVector, Dim>& /*displacement*/,
+      const tnsr::IJ<DataVector, Dim>& /*stress*/) noexcept {
     for (size_t d = 0; d < Dim; d++) {
       source_for_displacement->get(d) = 0.;
     }

src/Elliptic/Systems/Poisson/Equations.hpp

@@ -110,8 +110,13 @@ struct NonEuclideanFluxes {
  */
 struct Sources {
   using argument_tags = tmpl::list<>;
-  static void apply(const gsl::not_null<Scalar<DataVector>*> source_for_field,
-                    const Scalar<DataVector>& /*field*/) noexcept {
+  template <size_t Dim>
+  static void apply(
+      const gsl::not_null<Scalar<DataVector>*> source_for_field,
+      const gsl::not_null<
+          tnsr::i<DataVector, Dim>*> /*source_for_field_gradient*/,
+      const Scalar<DataVector>& /*field*/,
+      const tnsr::I<DataVector, Dim>& /*field_flux*/) noexcept {
     get(*source_for_field) = 0.;
   }
 };

tests/Unit/Elliptic/DiscontinuousGalerkin/Test_InitializeFirstOrderOperator.cpp

@@ -86,8 +86,12 @@ struct Fluxes {
 
 struct Sources {
   using argument_tags = tmpl::list<>;
-  static void apply(const gsl::not_null<Scalar<DataVector>*> source_for_field,
-                    const Scalar<DataVector>& field) {
+  template <size_t Dim>
+  static void apply(
+      const gsl::not_null<Scalar<DataVector>*> source_for_field,
+      const gsl::not_null<tnsr::i<DataVector, Dim>*> /*source_for_aux_field*/,
+      const Scalar<DataVector>& field,
+      const tnsr::I<DataVector, Dim>& /*field_flux*/) {
     get(*source_for_field) = get(field);
   }
 };

tests/Unit/Elliptic/Systems/Elasticity/Test_Equations.cpp

@@ -63,6 +63,7 @@ void test_computers(const DataVector& used_for_size) {
   using auxiliary_flux_tag =
       ::Tags::Flux<auxiliary_field_tag, tmpl::size_t<Dim>, Frame::Inertial>;
   using field_source_tag = ::Tags::Source<field_tag>;
+  using auxiliary_source_tag = ::Tags::Source<auxiliary_field_tag>;
   using constitutive_relation_tag = Elasticity::Tags::ConstitutiveRelation<
       Elasticity::ConstitutiveRelations::IsotropicHomogeneous<Dim>>;
   using coordinates_tag = domain::Tags::Coordinates<Dim, Frame::Inertial>;
@@ -103,17 +104,23 @@ void test_computers(const DataVector& used_for_size) {
   }
   {
     INFO("Sources" << Dim << "D");
-    auto box = db::create<db::AddSimpleTags<field_tag, field_source_tag>>(
-        tnsr::I<DataVector, Dim>{num_points,
-                                 std::numeric_limits<double>::signaling_NaN()},
-        tnsr::I<DataVector, Dim>{num_points,
-                                 std::numeric_limits<double>::signaling_NaN()});
+    auto box =
+        db::create<db::AddSimpleTags<field_source_tag, auxiliary_source_tag>>(
+            tnsr::I<DataVector, Dim>{
+                num_points, std::numeric_limits<double>::signaling_NaN()},
+            tnsr::ii<DataVector, Dim>{
+                num_points, std::numeric_limits<double>::signaling_NaN()});
 
     const Elasticity::Sources<Dim> sources_computer{};
     using argument_tags = typename Elasticity::Sources<Dim>::argument_tags;
 
-    db::mutate_apply<tmpl::list<field_source_tag>, argument_tags>(
-        sources_computer, make_not_null(&box), get<field_tag>(box));
+    db::mutate_apply<tmpl::list<field_source_tag, auxiliary_source_tag>,
+                     argument_tags>(
+        sources_computer, make_not_null(&box),
+        tnsr::I<DataVector, Dim>{num_points,
+                                 std::numeric_limits<double>::signaling_NaN()},
+        tnsr::IJ<DataVector, Dim>{
+            num_points, std::numeric_limits<double>::signaling_NaN()});
     auto expected_field_source = tnsr::I<DataVector, Dim>{num_points, 0.};
     CHECK(get<field_source_tag>(box) == expected_field_source);
   }

tests/Unit/Elliptic/Systems/Poisson/Test_Equations.cpp

@@ -51,6 +51,7 @@ void test_computers(const DataVector& used_for_size) {
   using auxiliary_flux_tag =
       ::Tags::Flux<auxiliary_field_tag, tmpl::size_t<Dim>, Frame::Inertial>;
   using field_source_tag = ::Tags::Source<field_tag>;
+  using auxiliary_source_tag = ::Tags::Source<auxiliary_field_tag>;
 
   const size_t num_points = used_for_size.size();
   {
@@ -133,14 +134,23 @@ void test_computers(const DataVector& used_for_size) {
   }
   {
     INFO("Sources" << Dim << "D");
-    auto box = db::create<db::AddSimpleTags<field_tag, field_source_tag>>(
-        Scalar<DataVector>{num_points, 2.}, Scalar<DataVector>{num_points, 0.});
+    auto box =
+        db::create<db::AddSimpleTags<field_source_tag, auxiliary_source_tag>>(
+            Scalar<DataVector>{num_points,
+                               std::numeric_limits<double>::signaling_NaN()},
+            tnsr::i<DataVector, Dim>{
+                num_points, std::numeric_limits<double>::signaling_NaN()});
 
     const Poisson::Sources sources_computer{};
     using argument_tags = typename Poisson::Sources::argument_tags;
 
-    db::mutate_apply<tmpl::list<field_source_tag>, argument_tags>(
-        sources_computer, make_not_null(&box), get<field_tag>(box));
+    db::mutate_apply<tmpl::list<field_source_tag, auxiliary_source_tag>,
+                     argument_tags>(
+        sources_computer, make_not_null(&box),
+        Scalar<DataVector>{num_points,
+                           std::numeric_limits<double>::signaling_NaN()},
+        tnsr::I<DataVector, Dim>{num_points,
+                                 std::numeric_limits<double>::signaling_NaN()});
     auto expected_field_source = Scalar<DataVector>{num_points, 0.};
     CHECK(get<field_source_tag>(box) == expected_field_source);
   }

tests/Unit/Elliptic/Test_FirstOrderComputeTags.cpp

@@ -63,9 +63,13 @@ struct Fluxes {
 
 struct Sources {
   using argument_tags = tmpl::list<AnArgument, BaseArgumentTag>;
-  static void apply(const gsl::not_null<Scalar<DataVector>*> source_for_field,
-                    const double an_argument, const double base_tag_argument,
-                    const Scalar<DataVector>& field) {
+  template <size_t Dim>
+  static void apply(
+      const gsl::not_null<Scalar<DataVector>*> source_for_field,
+      const gsl::not_null<tnsr::i<DataVector, Dim>*> /*source_for_aux_field*/,
+      const double an_argument, const double base_tag_argument,
+      const Scalar<DataVector>& field,
+      const tnsr::I<DataVector, Dim>& /*field_flux*/) {
     get(*source_for_field) =
         get(field) * square(an_argument) + base_tag_argument;
   }
@@ -82,8 +86,8 @@ void test_first_order_compute_tags() {
       elliptic::Tags::FirstOrderFluxesCompute<Dim, Fluxes<Dim>, vars_tag,
                                               primal_vars, auxiliary_vars>;
   using first_order_sources_compute_tag =
-      elliptic::Tags::FirstOrderSourcesCompute<Sources, vars_tag, primal_vars,
-                                               auxiliary_vars>;
+      elliptic::Tags::FirstOrderSourcesCompute<Dim, Sources, vars_tag,
+                                               primal_vars, auxiliary_vars>;
 
   TestHelpers::db::test_compute_tag<first_order_fluxes_compute_tag>(
       "Variables(Flux(FieldTag),Flux(AuxiliaryFieldTag))");

tests/Unit/Elliptic/Test_FirstOrderOperator.cpp

@@ -65,9 +65,12 @@ struct Fluxes {
 
 struct Sources {
   using argument_tags = tmpl::list<AnArgument>;
-  static void apply(const gsl::not_null<Scalar<DataVector>*> source_for_field,
-                    const double an_argument,
-                    const Scalar<DataVector>& field) {
+  template <size_t Dim>
+  static void apply(
+      const gsl::not_null<Scalar<DataVector>*> source_for_field,
+      const gsl::not_null<tnsr::i<DataVector, Dim>*> /*source_for_aux_field*/,
+      const double an_argument, const Scalar<DataVector>& field,
+      const tnsr::I<DataVector, Dim>& /*field_flux*/) {
     get(*source_for_field) = get(field) * square(an_argument);
   }
 };
@@ -115,12 +118,12 @@ void test_fluxes_and_sources() {
 
   // Test sources
   Variables<db::wrap_tags_in<::Tags::Source, all_fields>> sources{num_points};
-  elliptic::first_order_sources<primal_fields, auxiliary_fields, Sources>(
-      make_not_null(&sources), vars, 3.);
+  elliptic::first_order_sources<Dim, primal_fields, auxiliary_fields, Sources>(
+      make_not_null(&sources), vars, fluxes, 3.);
   // Check return-by-ref and return-by-value functions are equal
   CHECK(sources ==
-        elliptic::first_order_sources<primal_fields, auxiliary_fields, Sources>(
-            vars, 3.));
+        elliptic::first_order_sources<Dim, primal_fields, auxiliary_fields,
+                                      Sources>(vars, fluxes, 3.));
   // Check computed values
   CHECK(get(get<::Tags::Source<PrimalField>>(sources)) ==
         DataVector{num_points, 18.});

tests/Unit/Helpers/PointwiseFunctions/AnalyticSolutions/FirstOrderEllipticSolutionsTestHelpers.hpp

@@ -54,6 +54,8 @@ void verify_solution(
   using primal_fields = typename System::primal_fields;
   using auxiliary_fields = typename System::auxiliary_fields;
 
+  CAPTURE(mesh);
+
   const size_t num_points = mesh.number_of_grid_points();
   const auto logical_coords = logical_coordinates(mesh);
   const auto inertial_coords = coord_map(logical_coords);
@@ -72,10 +74,11 @@ void verify_solution(
   auto div_fluxes = divergence(std::move(fluxes), mesh,
                                coord_map.inv_jacobian(logical_coords));
   auto sources = std::apply(
-      [&solution_fields](const auto&... expanded_sources_args) {
-        return ::elliptic::first_order_sources<primal_fields, auxiliary_fields,
+      [&solution_fields, &fluxes](const auto&... expanded_sources_args) {
+        return ::elliptic::first_order_sources<System::volume_dim,
+                                               primal_fields, auxiliary_fields,
                                                typename System::sources>(
-            solution_fields, expanded_sources_args...);
+            solution_fields, fluxes, expanded_sources_args...);
       },
       sources_args);
   Variables<db::wrap_tags_in<Tags::OperatorAppliedTo, all_fields>>