Restructure: put radiation flux update into a function (#750)

### Description

This PR wraps the radiation flux update in source_terms_multi_group.hpp
into a single function `UpdateFlux`. A new struct `FluxUpdateResult` is
created to store the results from `UpdateFlux`.

### Related issues
None. 

### Checklist
_Before this pull request can be reviewed, all of these tasks should be
completed. Denote completed tasks with an `x` inside the square brackets
`[ ]` in the Markdown source below:_
- [x] I have added a description (see above).
- [ ] I have added a link to any related issues see (see above).
- [x] I have read the [Contributing
Guide](https://github.com/quokka-astro/quokka/blob/development/CONTRIBUTING.md).
- [ ] I have added tests for any new physics that this PR adds to the
code.
- [x] I have tested this PR on my local computer and all tests pass.
- [x] I have manually triggered the GPU tests with the magic comment
`/azp run`.
- [x] I have requested a reviewer for this PR.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/radiation/radiation_system.hpp

@@ -13,15 +13,12 @@
 
 #include <array>
 #include <cmath>
-#include <functional>
 
 // library headers
 #include "AMReX.H" // IWYU pragma: keep
 #include "AMReX_Array.H"
 #include "AMReX_BLassert.H"
 #include "AMReX_GpuQualifiers.H"
-#include "AMReX_IParser_Y.H"
-#include "AMReX_IntVect.H"
 #include "AMReX_REAL.H"
 
 // internal headers
@@ -121,6 +118,14 @@ template <typename problem_t> struct JacobianResult {
 	quokka::valarray<double, Physics_Traits<problem_t>::nGroups> Fg;  // (g) components of the residual, g = 1, 2, ..., nGroups
 };
 
+// A struct to hold the results of UpdateFlux(), containing the following elements:
+// Erad, gasMomentum, Frad
+template <typename problem_t> struct FluxUpdateResult {
+	quokka::valarray<double, Physics_Traits<problem_t>::nGroups> Erad;			   // radiation energy density
+	amrex::GpuArray<double, 3> gasMomentum;							   // gas momentum
+	amrex::GpuArray<amrex::GpuArray<amrex::Real, Physics_Traits<problem_t>::nGroups>, 3> Frad; // radiation flux
+};
+
 [[nodiscard]] AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE static auto minmod_func(double a, double b) -> double
 {
 	return 0.5 * (sgn(a) + sgn(b)) * std::min(std::abs(a), std::abs(b));
@@ -234,6 +239,9 @@ template <typename problem_t> class RadSystem : public HyperbolicSystem<problem_
 				       amrex::GpuArray<amrex::Real, AMREX_SPACEDIM> const &prob_lo, amrex::GpuArray<amrex::Real, AMREX_SPACEDIM> const &prob_hi,
 				       amrex::Real time);
 
+	AMREX_GPU_DEVICE static auto UpdateFlux(int i, int j, int k, arrayconst_t const &consPrev, NewtonIterationResult<problem_t> &energy, double dt,
+						double gas_update_factor, double Ekin0) -> FluxUpdateResult<problem_t>;
+
 	static void AddSourceTermsMultiGroup(array_t &consVar, arrayconst_t &radEnergySource, amrex::Box const &indexRange, amrex::Real dt, int stage,
 					     double dustGasCoeff, int *p_iteration_counter, int *p_iteration_failure_counter);