© The Authors 2025

1 Introduction

The tiny dust grains originating from the interstellar medium play a crucial role in its evolution at all scales (e.g. Draine 2003; Draine et al. 2007). They are dynamically coupled to the gas via a drag term and can thus exchange momentum and affect the gas dynamics. They act as charge reservoirs, allowing them to interact with the local magnetic field. They also interact with radiation, which affects the light that reaches Earth, and they play a key role in the local thermal balance by significantly contributing to opacity. Additionally, they regulate local chemistry and serve as fundamental building blocks for the formation of planetary embryos. Consequently, to achieve quantitative accuracy, numerical simulations of structure formation increasingly incorporate the most realistic models for dust processing (e.g. the discussions by Bai & Stone 2010; Laibe & Price 2012b; Lombart & Laibe 2021; Moseley & Teyssier 2025).

A major challenge arises from the fact that, locally, dust distributions can vary significantly, particularly in terms of size, which ranges from submicron scales to millimeter-sized particles, as seen in protoplanetary disks (e.g. Bae et al. 2023; Lesur et al. 2023a). This variability is referred to as the polydisperse distribution of dust. This polydispersity adds complexity to the evolution of the system by generating dynamics that do not appear when all grains are of a single size (Bai & Stone 2010; Laibe & Price 2014). For instance, in a protoplanetary disk, larger grains tend to drift inward, which in turn induces a radial outward drift of the gas as a result of feedback effects (Dipierro et al. 2018). This outward gas motion then carries along part of the dust distribution. The consequences of this polydispersity can be critical by favouring or preventing the development of instabilities that are necessary to concentrate dust and form the first planetary embryos (e.g. Krapp et al. 2019; Paardekooper et al. 2020; Krapp et al. 2020; Paardekooper et al. 2021; Wu et al. 2024; Cañas et al. 2024; Paardekooper & Aly 2025a,b; Matthijsse et al. 2025). Last but not least, the dust grain size distribution evolves during the formation of a star and planets via coagulation and fragmentation processes. The latter are regulated by the coupling of the dust grains with the gas and the magnetic fields, for which the dust grain size is key (e.g., Guillet et al. 2020; Lombart & Laibe 2021). In these scenarios, the behaviour of the gas-dust mixture is highly sensitive to the nature and distribution of the dust grains. Precise and rigorous numerical modelling is required to accurately capture these effects.

Numerically, addressing the polydispersity of strong constraints in hydrodynamic solvers presents several challenges (e.g. Bai & Stone 2010; Hutchison et al. 2018; Benítez-Llambay et al. 2019; Mentiplay et al. 2020). On one hand, increasing the number of fields required to evolve in order to simulate mixtures with different dust sizes imposes significant constraints on the computation time, memory access, and data management. In this regard, the solution lies in algorithmic and software advancements. In particular, the advent of exascale computing offers the possibility of overcoming some of these challenges, albeit at the cost of significantly higher energy and environmental demands (e.g. Wibking & Krumholz 2022; Grete et al. 2023; Lesur et al. 2023b; David-Cléris et al. 2025). On the other hand, intense friction regimes associated with small grains impose strict constraints on time-integration schemes. One way to circumvent this difficulty is to reformulate the gas-grain mixture equations in a limit where its evolution can be treated as an equivalent diffusive process (Hutchison et al. 2018; Lebreuilly et al. 2019; Mentiplay et al. 2020). This method is particularly effective when all grains are well coupled with the gas, but it becomes inadequate when part of the distribution is not (e.g. Tricco et al. 2017; Commerçon et al. 2023). Another approach involves using unconditionally stable implicit numerical schemes (Bai & Stone 2010; Benítez-Llambay et al. 2019; Krapp & Benítez-Llambay 2020; Huang & Bai 2022; Krapp et al. 2024). These schemes typically require significantly more computational resources than explicit schemes, however, and they often lead to a trade-off in accuracy for the sake of practical feasibility. This is contrary to the goal of a high-precision numerical treatment that is necessary for accurately simulating complex physical processes. For these reasons, we aim to design a dust algorithm compatible with new multi-GPU architectures. By properly handling the friction matrix decomposition, we ensure maximum precision of the drag calculation without significantly impacting the overall computation time.

The manuscript is organised as follows. Section 2 presents the equations and the most recent description for dust- and gas-mixture hydrodynamics and introduces our novel exponential solver. In Sect. 3 we test our method with respect to other drag solvers in different regimes using standard gas- and dust-dynamics test. We also assess and discuss its performance on GPUs. Section 4 concludes the paper.

2 Numerical method

2.1 Equations of the evolution of a dust and gas mixture

We considered an astrophysical dust-gas mixture and modelled N pressure-less dust fluids. Each dust fluid i was coupled to the gas via a drag term parametrised by T_s,i, the stopping time of the i th dust species. In most astrophysical conditions, the mean free path of the gas exceeds the size of the grain by far. In this case, the grain stopping time of dust species i is given by Epstein (1924) $$T_{\mathrm{s},i}=\sqrt{\frac{\pi \gamma }{8}}\frac{{\rho }_{\text{grain},i}}{\rho }\frac{S_{\text{grain},i}}{c_{\mathrm{s}}},$$(1) where γ is the adiabatic index of the gas, c_s is the sound speed of the gas, and S_{grain, i} and ρ_{grain, i} are the radius and intrinsic density of the grain. This intrinsic grain density is different from the dust density ρ_d (see below).

Conservation of momentum was ensured by the backreaction term onto the gas momentum. The equations of evolution for the mixture are given by $$\frac{\partial {\rho }_{\mathrm{g}}}{\partial t}+\nabla ·\left[{\rho }_{\mathrm{g}}{\mathit{v}}_{\mathrm{g}}\right]=0$$(2) $$\frac{\partial \left({\rho }_{\mathrm{g}}{\mathit{v}}_{\mathrm{g}}\right)}{\partial t}+\nabla ·[{\rho }_{\mathrm{g}}{\mathit{v}}_{\mathrm{g}}\otimes {\mathit{v}}_{\mathrm{g}}+P_{\mathrm{g}}\mathrm{I}]-{\rho }_{\mathrm{g}}{\mathit{f}}_{\mathrm{g}}={\sum }_{i=1}^N{\rho }_{\mathrm{d},i}\frac{{\mathit{v}}_{\mathrm{d},i}-{\mathit{v}}_{\mathrm{g}}}{T_{\mathrm{s},i}}$$(3) $$\begin{array}{rl}\frac{\mathrm{\partial }{E}_{\mathrm{g}}}{\mathrm{\partial }t}+\mathrm{\nabla }\cdot \left[(E_{\mathrm{g}}+P_{\mathrm{g}}){\mathit{v}}_{\mathrm{g}}\right]-{\rho }_{\mathrm{g}}{\mathit{f}}_{\mathrm{g}}\cdot {\mathit{v}}_{\mathrm{g}}& =\sum _{i=1}^N{\rho }_{\mathrm{d},i}\frac{{\mathit{v}}_{\mathrm{d},i}-{\mathit{v}}_{\mathrm{g}}}{T_{\mathrm{s},i}}\cdot {\mathit{v}}_{\mathrm{g}}\\ & +w\sum _{i=1}^N{\rho }_{\mathrm{d},i}\frac{{({\mathit{v}}_{\mathrm{d},i}-{\mathit{v}}_{\mathrm{g}})}^2}{T_{\mathrm{s},i}}\end{array}$$(4) $$\frac{\partial {\rho }_{\mathrm{d},i}}{\partial t}+\nabla ·\left[{\rho }_{\mathrm{d},i}{\mathit{v}}_{\mathrm{d},i}\right]=0,$$(5) $$\frac{\partial \left({\rho }_{\mathrm{d},i}{\mathit{v}}_{\mathrm{d},i}\right)}{\partial t}+\nabla ·[{\rho }_{\mathrm{d},i}{\mathit{v}}_{\mathrm{d},i}\otimes {\mathit{v}}_{\mathrm{d},i}]-{\rho }_{\mathrm{d},i}{\mathit{f}}_{\mathrm{d},i}={\rho }_{\mathrm{d},i}\frac{{\mathit{v}}_{\mathrm{g}}-{\mathit{v}}_{\mathrm{d},i}}{T_{\mathrm{s},i}},$$(6) where ρ_g, v_g, P_g, and E_g denote the gas density, the gas velocity and gas pressure, and the total gas energy $E_{\mathrm{g}}=\frac{1}{2}{\rho }_{\mathrm{g}}{\mathit{v}}_{\mathrm{g}}^2+e$, respectively, with e=P_g(γ−1) the gas internal energy. The densities and velocities of the dust fluid i are denoted ρ_d,i and v_{d, i}, respectively. I is the matrix identity. The external forces specific to each phase are denoted f_g and f_{d, i}. w is a control variable that parametrises frictional heating: w=0 turns the frictional heating off, and w=1 means that the entire dissipation is deposited into the gas. Eqs. (2)–(6) can be written in the following compact form: $$\frac{\partial \mathit{U}}{\partial t}+\nabla ·F\left(\mathit{U}\right)+{\mathit{f}}_{\mathrm{src},\mathrm{ext}}={\mathit{f}}_{\mathrm{drag}},$$(7) where $$\mathit{U}\equiv [{\rho }_{\mathrm{g}},{\rho }_{\mathrm{g}}{\mathit{v}}_{\mathrm{g}},E_{\mathrm{g}},{\rho }_{\mathrm{d},i},{\rho }_{\mathrm{d},i}{\mathit{v}}_{\mathrm{d},i}],$$(8) $$F\left(\mathit{U}\right)\equiv [{\rho }_{\mathrm{g}}{\mathit{v}}_{\mathrm{g}},{\rho }_{\mathrm{g}}{\mathit{v}}_{\mathrm{g}}\otimes {\mathit{v}}_{\mathrm{g}}+P_{\mathrm{g}}\mathit{I},(E_{\mathrm{g}}+P_{\mathrm{g}}){\mathit{v}}_{\mathrm{g}},{\rho }_{\mathrm{d},i}{\mathit{v}}_{\mathrm{d},i},{\rho }_{\mathrm{d},i}{\mathit{v}}_{\mathrm{d},i}\otimes {\mathit{v}}_{\mathrm{d},i}],$$(9) $${\mathit{f}}_{\text{src,ext}}\equiv [0,{\rho }_{\mathrm{g}}{\mathit{f}}_{\mathrm{g}},{\rho }_{\mathrm{g}}{\mathit{f}}_{\mathrm{g}}·{\mathit{v}}_{\mathrm{g}},0,{\rho }_{\mathrm{d},i}{\mathit{f}}_{\mathrm{d},i}],$$(10) and $$\begin{array}{rl}{\mathit{f}}_{\mathrm{d}\mathrm{r}\mathrm{a}\mathrm{g}}\equiv [0,\sum _{i=1}^N{\rho }_{\mathrm{d},i}& \frac{{\mathit{v}}_{\mathrm{d},i}-{\mathit{v}}_{\mathrm{g}}}{T_{\mathrm{s},i}},\sum _{i=1}^N{\rho }_{\mathrm{d},i}\frac{{\mathit{v}}_{\mathrm{d},i}-{\mathit{v}}_{\mathrm{g}}}{T_{\mathrm{d},i}}\cdot {\mathit{v}}_{\mathrm{g}}\\ & +w\sum _{i=1}^N{\rho }_{\mathrm{d},i}\frac{{({\mathit{v}}_{\mathrm{d},i}-{\mathit{v}}_{\mathrm{g}})}^2}{T_{\mathrm{s},i}},0,{\rho }_{\mathrm{d},i}\frac{{\mathit{v}}_{\mathrm{g}}-{\mathit{v}}_{\mathrm{d},i}}{T_{\mathrm{s},i}}].\end{array}$$(11)

In Eqs. (7)-(11), U denotes the state vector of conservative variables, and F(U) is a rank-2 tensor, the rows of which are the fluxes of each component of U. f_src,ext and f_drag denote external source terms and drag source terms.

2.2 Review of current dust-gas integrators

Equation (7) is usually solved by splitting the time step to separately integrate the transport term, the source terms, and the drag terms. Gas and dust are treated in a similar way, except for the energy equation. Drag integration requires solving the following equation: $$\frac{\partial \tilde{\mathit{U}}}{\partial t}=\left[\begin{array}{c}-{\sum }_{i=1}^N{\alpha }_i({ϵ}_i{\mathit{M}}_{\mathrm{g}}-{\mathit{M}}_{\mathrm{d},i})\\ {\alpha }_1({ϵ}_1{\mathit{M}}_{\mathrm{g}}-{\mathit{M}}_{\mathrm{d},1})\\ {\alpha }_2({ϵ}_2{\mathit{M}}_{\mathrm{g}}-{\mathit{M}}_{\mathrm{d},2})\\ ⋮\\ {\alpha }_N({ϵ}_N{\mathit{M}}_{\mathrm{g}}-{\mathit{M}}_{\mathrm{d},N})\end{array}\right]\equiv {\tilde{\mathit{f}}}_{\text{drag}},$$(12) where $$\tilde{\mathit{U}}\equiv {[{\rho }_{\mathrm{g}}{\mathit{v}}_{\mathrm{g}},{\rho }_{\mathrm{d},1}{\mathit{v}}_{\mathrm{d},1},...,{\rho }_{\mathrm{d},N}{\mathit{v}}_{\mathrm{d},N}]}^T={[{\mathit{M}}_{\mathrm{g}},{\mathit{M}}_{\mathrm{d},1},...,{\mathit{M}}_{\mathrm{d},N}]}^T,$$(13) $${ϵ}_i\equiv \frac{{\rho }_{\mathrm{d},i}}{{\rho }_{\mathrm{g}}},$$(14) and $${\alpha }_i\equiv \frac{1}{T_{\mathrm{s},i}}.$$(15)

Assuming linear drag terms, Eq. (12) can be set in the form of a system of first-order ordinary differential equations, $$\frac{\partial \tilde{\mathit{U}}}{\partial t}=\mathit{\Omega }\tilde{U},$$(16) where Ω is the (N+1) x(N+1) sparse matrix, defined by $$\mathit{\Omega }\equiv \left[\begin{array}{ccccc}-{\sum }_{i=1}^N{ϵ}_i{\alpha }_i& {\alpha }_1& {\alpha }_2& ...& {\alpha }_N\\ {ϵ}_1{\alpha }_1& -{\alpha }_1& 0& ...& 0\\ {ϵ}_2{\alpha }_2& 0& -{\alpha }_2& ...& 0\\ ⋮& ⋮& ⋮& \ddots & ⋮\\ {ϵ}_N{\alpha }_N& 0& 0& ...& -{\alpha }_N\end{array}\right].$$(17)

Several methods have been proposed to solve Eq. (16). Benítez-Llambay et al. (2019) and Krapp & Benítez-Llambay (2020) proposed a fully implicit first-order method that involves a similar sparse matrix that can be inverted analytically with the complexity of O(N). Huang & Bai (2022) proposed several second-order methods that are explicit or implicit, where the implicit schemes involve dense matrices that are computationally expensive to invert. The diagonally implicit Runge-Kutta (DIRK) method of Krapp et al. (2024) combines the advantages of the first- and second-order fully implicit methods.

2.3 Exponential solver

2.3.1 Series methods for the matrix exponential

We chose to use the analytical solution of Eq. (16), $$\tilde{\mathit{U}}\left(t\right)={\mathrm{e}}^{t\mathit{\Omega }}\tilde{\mathit{U}}\left(t_0\right).$$(18) to build an implicit solver that preserves the accuracy of the scheme. The key difficulty was to maintain reasonable computational costs. Several algorithms have been proposed for the computation of the matrix exponential (Moler & Van Loan 2003). For a given matrix A, the naive approach consists of performing a Taylor expansion in the truncation order m according to $${\mathrm{e}}^{\mathit{A}}={\sum }_{i=0}^{\infty }\frac{{\mathit{A}}^i}{i!}\simeq {\sum }_{i=0}^m\frac{{\mathit{A}}^i}{i!}\equiv T_m\left(\mathit{A}\right),$$(19) which is usually not accurate enough because it entails truncation errors (both arithmetic and series truncation) and round-off errors (Moler & Van Loan 2003). An alternative method consists of using the (p, q) Padé approximant of e^A, defined by $${\mathit{R}}_{pq}\left(\mathit{A}\right)=N_{pq}\left(\mathit{A}\right){\left({\mathit{D}}_{pq}\right)}^{-1},$$(20) where $${\mathit{N}}_{pq}\left(\mathit{A}\right)={\sum }_{j=0}^p\frac{(p+q-j)!p!}{(p+q)!j!(p-j)!}{\mathit{A}}^j,$$(21) and $${\mathit{D}}_{pq}\left(\mathit{A}\right)={\sum }_{j=0}^q\frac{(p+q-j)!q!}{(p+q)!j!(q-j)!}{(-\mathit{A})}^j.$$(22)

This approximation holds when D_{p q}(A) is not singular. This property is a priori not satisfied for all pairs (p, q). The Taylor series of order m is just a (m, 0) Padé approximant. This method provides a much better approximation for a given order than the Taylor series (see Table 1 in Moler & Van Loan 2003). The Padé approximation is also limited by its own truncation and round-off errors, however. Both algorithms can be used when ‖A‖<1 (Sastre et al. 2015).

2.3.2 Scaling and squaring method

The most widely used method for computing the matrix exponential is the so-called scaling and squaring method (Al-Mohy & Higham 2010). The method relies on the following identity: $${\mathrm{e}}^A={\left({\mathrm{e}}^{A/n}\right)}^n,$$(23) where n is a non-negative integer. It consists of computing an approximation for the exponential of the scaled matrix e^A/n, and then squaring the result n times to have the approximation of e^A. The scaling step is performed with a Padé or Taylor approximation, which requires that ‖A‖ is sufficiently small (i.e. ‖A‖<1) to control round-off errors.

The scaling step is commonly performed using a Padé approximation to achieve double precision in the institute of electrical and electronics engineers (IEEE) 754 standard at lowest computational cost (Skaflestad & Wright 2009; Al-Mohy & Higham 2009; Higham 2009). This requires a matrix inversion, as shown by Eq. (20). It is possible to achieve almost similar precision without paying the cost of this matrix inversion by the Taylor series, however (Ruiz et al. 2016). The scale factor s is defined as $$2^s\equiv n.$$(24)

Importantly, Fasi & Higham (2019) extended these methods for arbitrary precision, which is of prime interest for our drag algorithm because the aim is to achieve the required accuracy at minimum computational cost.

Sastre et al. (2015) proposed a modified version of the Taylor-based scaling and squaring algorithm. The first modification is relative to the choice of parameters $m_k^{*}$ (optimal Taylor polynomial order), k^* (optimal matrix product required for the Taylor polynomial evaluation), and s^* (optimal squaring parameter), such that the total number of matrix products k^*+s^* is smaller than or equal to the number for the standard Taylor evaluation. The second modification they proposed is to neglect high-order terms of the Taylor polynomials when for the given unit round-off u the accuracy of the final result is not affected.

With these two modifications, they proposed four new Taylor-based scaling and squaring algorithms. Although two of the algorithms they proposed gave a comparable error to the expm routine of MATLAB (matrix exponential routine of MATLAB based on a Padé-approximation), however, we need to point out that all their algorithms are memory-bound. The reason is the need of storing a sequence of the matrix power {A, A², \ldots, A^q}, with q a given integer known at runtime. The memory required to store the sequence of matrices for every cell in our simulation renders this method unfeasible on both GPUs and CPUs. In Sastre et al. (2015), the Taylor algorithm consisted of two parts:

1. A Paterson-Stockmeyer coefficient evaluation (see ALGORITHM PS-COEFF in Sastre et al. 2015), and

2. a Horner-Paterson-Stockmeyer evaluation of the order- m Taylor polynomial (see ALGORITHM HPS-EVAL in Sastre et al. 2015).

We merged the two algorithms to reduce the memory footprint. The algorithmic details are presented in Appendix B.

2.4 Implementation in Dyablo and Shamrock

The different drag solvers we described above are all implemented in the codes Dyablo (Delorme et al. 2025) and Shamrock (David-Cléris et al. 2025). These are two multi-GPU codes for astrophysical hydrodynamics simulations with the aim of targeting exascale supercomputers. They are both written in C++ (currently, C++17) and use MPI for the domain decomposition. For the shared memory parallelism, Dyablo uses Kokkos (Trott et al. 2022), and Shamrock uses Sycl (Group 2021). Dyablo supports Eulerian AMR grid-based method. Shamrock supports both Lagrangian smoothed particle hydrodynamics (SPH) and Eulerian adaptive-mesh-refinement(AMR) grid-based methods.

As pointed out in Sect. 2.2, Eq. (7)–(11) was solved by the splitting technique. To evolve Uⁿ from time t=tⁿ to a new value Uⁿ⁺¹ at t=tⁿ⁺¹, we first solved $$\frac{\partial \mathit{U}}{\partial t}+\nabla ·F\left(\mathit{U}\right)+{\mathit{f}}_{\mathrm{src},\mathrm{ext}}=0,$$(25) with U(tⁿ)=Uⁿ as initial condition, to obtain an intermediate state value Ūⁿ⁺¹. Then, using Ūⁿ⁺¹ as initial condition, we solved $$\frac{\partial \mathit{U}}{\partial t}={\mathit{f}}_{\text{drag}}.$$(26)

We integrated Eq. (25) using a second-order Godunov method, whereas Eq. (26) was solved by one of the different drag solvers we described in Sect. 2.2.

3 Tests and performance

In this section, we present the numerical tests we performed to validate the method. We also present a performance analysis of our exponential drag solver within both Dyablo and Shamrock.

3.1 DUSTYBOX test

In this first test, we focused on the accuracy and the robustness of our matrix exponential solver alone (i.e. without the hydrodynamical solver step). The 1D dust-gas collisions test (Benítez-Llambay et al. 2019) is a derivative of the DUSTYBOX test (Laibe & Price 2011), which consists of a system of gas and an arbitrary number of dust fluids with constant stopping times. The system evolves only under the momentum exchange between gas and dust fluids. In other words, we integrated Eq. (16) for a given initial condition Ũ₀. Benítez-Llambay et al. (2019) showed that the velocity of all fluids (gas and dust) is expected to converge to an asymptotic steady-state velocity (the velocity of the centre of mass), given by $$\mathit{V}\left(t\right)={\mathit{V}}_{\mathrm{CM}}+{\sum }_{i=1}^N{\mathit{c}}_i{e}^{{\lambda }_{i}t},$$(27) where V_CM is the velocity of the centre of mass, defined as $${\mathit{V}}_{\mathrm{CM}}=\frac{{\mathit{v}}_{\mathrm{g}}+{ϵ}_i{\mathit{v}}_{\mathrm{d},i}}{1+{\sum }_{i=1}^N{ϵ}_i},$$(28) and λ_i<0 is obtained by solving the eigenvalue problem, $$\mathit{\Omega }\overline{\mathit{x}}=\lambda \overline{\mathit{x}}.$$(29) The associate vector c_i is then obtained by expressing the initial conditions into the eigenvectors.

We performed tests B (stiff case with short stopping times) and C (stiff case with high dust-to-gas ratios) described by Huang & Bai (2022) for our EXP¹ and different drag solvers of reference. In Table 1 we recall the initial conditions, the eigenvalues, and the coefficients for the analytical solution we used.

Figure 1 shows the evolution of velocities for Collisions tests B and C at a fixed time step of 0.005 and 0.05, respectively, for our EXP integrator (and others, presented in Sect. 2.2). In both cases, the numerical velocity converges to the asymptotic steady-state velocity, and it matches the analytical solution well.

To quantify how well our matrix exponential solver performed in comparison with the drag solvers used in the literature, we performed a convergence analysis in time where the error was obtained following Eq.(40) of Krapp et al. (2024), $$E{r}_1\equiv \frac{1}{M}{\sum }_{k=0}^2{\sum }_{i=0}^{M-1}\frac{|{\tilde{v}}_k\left(\Delta t_i\right)-v_k\left(\Delta t_i\right)|}{v_k\left(\Delta t_i\right)}$$(30) with v and ṽ denoting exact and numerical quantities, respectively. Here, Δ t_i is the time step, M is the number of time steps during numerical integration, and k=0, 1, 2 is the index for the gas and dust fluids. Starting at Δ t=10⁻⁴(10⁻³) for test B (for test C), we doubled the time step until we reached a final time of t_final=10⁻¹ (2) for test B and (test C). Here and in the remainder of the paper, we refer to RK1 as the first-order implicit drag solver, to RK2 as the second-order implicit Runge-Kutta solver, to MDIRK- and MDIRK+ as the versions of the multi-fluid second-order diagonally implicit Runge-Kutta solver, and to EXP as the exponential solver.

Figure 2 shows the evolution of error E r 1 as a function of the time step Δ t_i. The RK1 and MDIRK+ solvers are first order, and the RK2 and MDIRK- solvers are second order. On the other hand, the error of the EXP matrix exponential solver remains almost constant around 10⁻⁸, regardless of Δ t. This property arises because the choice of the optimal Taylor polynomial order $m_k^{*}$ and the optimal squaring parameter was made such that the forward error was smaller than the unit round-off u (we set its value to u=2⁻⁵³). This indicates that the EXP method is not affected by truncation errors and provides an accurate solution for the drag even at large time steps.

To test the robustness of the matrix exponential solver, we also performed a collision test with 20 dust fluids. The stopping time T_s,i was uniformly spaced in log scale between T_{s, min}=10⁻³ and T_s,max=10. The dust-to-gas ratios were given by $${ϵ}_i=\frac{ϵ(\sqrt{T_{\mathrm{s},i+1}}-\sqrt{T_{\mathrm{s},i}})}{(\sqrt{T_{\mathrm{s},max}}-\sqrt{T_{\mathrm{s},min}})},$$(31) such that $ϵ={\sum }_{i=1}^{N=20}{ϵ}_i=1.0$. Figure 3 shows the evolution of the gas and each of the 20 fluid velocities. Even in stiff regimes (short stopping times), the matrix exponential solver still performed well and provided errors of ∼ 10⁻⁷ for all dust fluids. Suppose for the RK1 scheme that E r₁ ∝ Δ t and the simulation time denotes $T_1\propto \frac{1}{\Delta t}$, then to achieve an error of ∼ 10⁻⁷, the RK1 scheme requires ∼ 10⁴(∼ 55) more time on a Collisions test with two (20) dust fluids.

Fig. 1 Collisions test. Evolution of the gas and dust velocities (from top to bottom) for tests B and C (left and right columns, respectively). The dot-dashed black line represents the analytical solution. The other lines compare the numerical solutions obtained with the EXP integrator (orange) with other integrators, namely the first-order implicit RK1 (blue), the second-order fully implicit RK2 (purple), the MDIRK method with $\gamma =1-1/\sqrt{2}$ (green), and the MDIRK method with $\gamma =1+1/\sqrt{2}$ (red).

Fig. 2 Collisions tests B (left) and C (right). The evolution of error E r1 is defined by Eq. (30) for the EXP (orange) and other integrators: RK1 (blue), RK2 (purple), MKDIRK- (green), and MKDIRK+ (red).

3.2 DUSTYWAVE test

The Dustywave test consisted in following the evolution of linear sound waves in a mixture of gas and dust. Because the dust phase is assumed to be pressure-less, it cannot support sound waves. The gas sound wave propagates with a modified sound speed in the mixture, however, because of gas pressure and momentum exchange. This test served as a benchmark for the coupling between the transport and drag steps by solving Eqs. (2)-(6). The analytical solution for one gas and one dust fluid was proposed by Laibe & Price (2012a) and was extended to an arbitrary number of dust fluids by Laibe & Price (2014).

We reproduced the 1D test for one gas fluid with one and four dust fluids. The setup consisted of an isothermal sound wave ($P_{\mathrm{g}}=c_{\mathrm{s}}^2{\rho }_{\mathrm{g}}$, with c_s=1) in a periodic box of size L=1 and zero background velocity. The perturbation equation writes $$\delta f=A[Re\left(\delta \widehat{f}\right)cos\left(kx\right)-Im\left(\delta \widehat{f}\right)sin\left(kx\right)],$$(32) with A a small amplitude to ensure linearity (10⁻⁴ c_s and ${10}^{-4}{\rho }_{\mathrm{g}}^0$ for velocity and density perturbations), and with the wavenumber $k=\frac{2\pi }{L}$. The test was performed on a uniform grid of N_cells= 128³.

The initial background densities, perturbation amplitudes, and stopping were taken from Table 2 of Benítez-Llambay et al. (2019).

Figure 4 shows the evolution of the normalised velocities $\left(\frac{\delta \widehat{v}}{c_{\mathrm{s}}A}\right)$ and densities $\left(\frac{\delta \widehat{\rho }}{{\rho }_{0}A}\right)$ for two and five dust fluids, respectively. The numerical solutions recover the analytical solutions well (a few 0.1−1%), which indicates a good coupling of the drag solver and the transport step. We additionally studied the convergence in space of the scheme (transport and drag together) by comparing the results with the RK1 solver, using a small time step Δ t=10⁻⁴ such that the influence of the time error was minimised. For the transport step, we used a second-order MUSCL-Hancock scheme with the Minmod slope limiter. To compute the error, we used Eq. (46) from Krapp et al. (2024) as follows: $$E{r}_2\equiv \frac{1}{N_{\text{cells}}}{\sum }_{k=1}^2{\sum }_{i=1}^2|\delta f_{i,\text{numeric},k}-\delta f_{i,\text{analytic},k}|,$$(33) with k and i the indices of the fluid and the fields (velocity and density), respectively. The lower left panel of Fig. 4 shows that the scheme has second-order accuracy in space, which is consistent with the order of the transport scheme we used. As the resolution increases, the EXP solver gives a better accuracy than the RK1 solver by a moderate factor.

Fig. 3 Evolution of the gas and dust velocities for a Collisions test with 20 dust fluids obtained with the EXP solver and compared to the analytical solution given by Eq. (27) (dashed black lines). For the first five dust species, their stopping time is short, such that they are strongly coupled to the gas. We therefore chose to plot the first and the fifth species. In addition to the 6th dust species, we chose to skip every second fluid from the 10th to the 20th dust species.

Fig. 4 Dustywave test. Evolution of normalised densities (upper panels) and velocities (lower panels) for one (left panel) and four dust fluids (right panel). The dashed black lines correspond to the analytical solution. The other markers and colours represent the numerical results obtained with the EXP algorithm. We additionally plot the L1 error given by Eq. (33) for the one-dust fluid case as an inset in the bottom left panel using a fixed time step of Δ t=10−4 for EXP and RK1.

3.3 DUSTYSHOCK test

The dust-gas mixture shock test is a canonical benchmarking problem to test the robustness and accuracy of numerical codes for dusty gas. The steady-state J-type shock solution (when the shock speed exceeds the signal speed) in the case of a one-dust fluid, proposed by Lehmann & Wardle (2018), was extended to an arbitrary number of dust fluids by Benítez-Llambay et al. (2019). The left state is given by $$\begin{array}{rl}{\rho }_{\mathrm{g}}^{-}& ={\rho }_{\mathrm{g}0}\\ {\rho }_{\mathrm{d},i}^{-}& ={ϵ}_i{\rho }_{\mathrm{g}0}\\ w_{\mathrm{g}}^{-}& =w_{\mathrm{d},i}=1.0\end{array}$$(34) and the asymptotic right states are given by $$\begin{array}{rl}{\rho }_{\mathrm{g}/\mathrm{d},i}^{+}& =\frac{{\rho }_{\mathrm{g}/\mathrm{d},i}^{-}{w}_{\mathrm{g}/\mathrm{d},i}^{-}}{w_{\mathrm{g}/\mathrm{d},i}^{+}}{\rho }_{\mathrm{g}0}\\ w_{\mathrm{g}}^{+}& =w_{\mathrm{d},i}^{+}={(1+\sum _{i=1}^N{ϵ}_i)}^{-1}{\mathcal{M}}^{-2}\end{array}$$(35) where w_g/d, i=v_g/d, i/v_s, v_s is the shock velocity, and ℳ=v_s/c_s is the Mach number. Following the same setup as in Sect. 3.3.2 from Benítez-Llambay et al. (2019), we performed the two-fluid and five-fluid tests. We list the initial conditions for the two-fluid and four-fluid tests below, $$\begin{array}{rl}(K,{\rho }^{-},\rho +,w^{-},w^{+})& =(1.0,1.0,8.0,1.0,0.125)\\ (K_1,K_2,K_3,{\rho }^{-},{\rho }^{+},w^{-},w^{+})& =(1.0,3.0,5.0,1.0,16.0,1.0,0.0625)\\ (c_{\mathrm{s}},v_{\mathrm{s}},ϵ,x_0,x_{\text{start\hspace{0.17em}}},x_{\text{end\hspace{0.17em}}},N_{\text{cells\hspace{0.17em}}})& =(1,2,1,4,0,40,512)\end{array}$$(36) x₀ is the initial jump position, and we used the zero-gradient boundary condition. We used an HLLE Riemann solver for the gas and the upwind solver of Huang & Bai (2022) for the dust. We used the van Leer slope limiter.

Figure 5 shows the normalised velocities and densities for gas and for one- (first two panels) and three- (last two panels) dust fluids at t=500.

The good agreement between the analytical and numerical solutions (∼ 0.1−1%) shows that our drag solver is correctly coupled to the transport step in a challenging regime.

3.4 Performance

We report the performance of our EXP algorithm on GPUs by performing the Dustywave test with various numbers of dust fluids. Following the same procedure as in Sect. (3.2), we generated uniform stopping times, computed the dust-to-gas ratio with Eq. (31), solved the dispersion relation, and computed the associated eigenvectors (see Eq. (45)–(48) in Benítez-Llambay et al. 2019). The tests were performed on a NVIDIA DGX station with four A100 GPUs and an AMD 7742 CPU (64 cores, 2.25−3.4 GHz). By default, we used N_cells=64³ and integrated until t=2. Profiling consisted in measuring the execution time of the relevant kernels using the tools Nsight compute and Nsight system of the NVIDIA Nsight developer tools suit. Although absolute performance is expected to improve with future hardware advancements and potential optimisation strategies, these tests provide useful estimates of the current tractability of the EXP algorithm to astrophysical simulations.

In the Dyablo code, the matrices required for the Taylor evaluation step are handled with the Kokkos array. These are static arrays, which implies that the code must be recompiled whenever the number of dust fluids changes. Because each cell needs to solve the drag equation, the cells of the overall grid are divided into blocks. Within each block, a parallel_for execution policy is applied in conjunction with a RangePolicy with a length equal to the block size, which allows each cell to integrate the drag equation. In the Shamrock code, the matrices are handled using mdspan, which are dynamic arrays that allow adjusting the number of dust fluids at runtime without recompilation. Within Shamrock, we experimented with two implementation strategies. In the first implementation, mdspan arrays were attached to the device or global memory, and we used a basic data-parallel kernel. In the second implementation, we used ND-Range data-parallel kernels, where the number of work-items² per work-group³ is computed with respect to the size of the shared memory available on the device and the number of dust fluids. In this case, mdspan arrays were directly attached to the local memory of the work-group.

Figure 6 shows the time per cell (top panels) and the relative cost of the drag kernels per time step (bottom panels) for both codes as function of the number of dust fluids N for a resolution of N_cells=64³. For reference, the computational cost of the RK1 solver scales as O(N), as expected, regardless of the implementation. For the EXP solver, the computational cost within Dyablo initially scales as O(N²) up to 8 fluids and then becomes dominated by matrix-matrix operations. This results in a time complexity of O(N³). Within Shamrock, the local memory implementation of the EXP solver offers a slight performance improvement over the global memory implementation for small numbers of fluids (e.g. a factor of 10 for N=4). As the number of fluids and thus the matrix sizes increases, the size of the work-groups decreases. Hence, the local and global memory implementations both have nearly identical execution times. One possible path for improvement might be to use hierarchical dataparallel kernel in which a work-group will compute the matrix exponential instead of assigning this task to a single thread, as is currently done. Furthermore, within Dyablo, the EXP solver takes 55% of the overall time steps for up to 16 fluids, while with Shamrock, the drag time step with the same number of fluids is about the same as the hydro time step in both implementations. For instance, for up to 8 fluids, the EXP solver in Dyablo takes less than 8%, whereas the local memory version in Shamrock is slightly lower than 50%.

Additionally, we analysed the effect of the resolution on the measured performances. Figure 7 shows the time per cell (top panels) and the relative cost of the drag solver per time step (bottom panels) as a function of the number of dust fluids for N_cells=256³ (left panels) and N_cells=64³ (right panels) cells (Dyablo code). We observe the same performance profile as for N_cells=64³, where for EXP, we have a time complexity of 𝒪(N²) and ℴO(N³) for N ≤ 8 and N ≥ 9, respectively. The computational cost with N=8(∼eq 9.7 × 10⁻⁹ s) multiplied by about 4.5 when moving to N=9(∼eq 4.4 × 10⁻⁸ s) clearly shows the change in order in the time complexity. With the Nsight Compute profiler, both kernels have a large number of registers (the fastest memory storage area with a streaming multiprocessor) per thread, 256 each. There are 0.18 and 0.09 eligible warps (groups of threads that are scheduled together and are executed in parallel) per cycle for N=8 and N=9 fluids, respectively, on average, however. For each warp, the latency (the amount of time wasted doing nothing) is 12.20 and 25.02 cycles on average between two consecutive executed instructions for 8 and 9 fluids, respectively. This shows that the actual implementation of the exponential solver is over using registers per thread, and as a result, it causes a low occupancy (a measure of the efficiency of the GPU computing resources). We provide the performances for the RK1 solver, which relies on an analytical prescription, for reference.

The above analyses suggest optimisation paths to make the EXP algorithm more competitive. A reduced pressure on the registers might lead to a better occupancy. Using several threads (e.g. a thread block) by matrix operation instead of a single thread, as is done currently, is worth considering. This might be done, for example, by using hierarchical parallelism. Finally, future GPUs with larger bandwidth and shared memory capacity (notably for the implementation in local memory) are also expected to reduce the latency of the algorithm.

Fig. 5 Dustyshock test. Normalised velocities (first and third panels for one- and three-dust fluid(s), respectively) and density (second and fourth panels for one- and three-dust fluid(s), respectively). The dashed lines show the analytical solutions, and the circle markers show the numerical results we obtained with the EXP algorithm. With the EXP solver, the shock is captured within one to two cells.

Fig. 6 Performance of the EXP drag solver as function of the number of fluids (gas and dust fluids) on one A100 GPU obtained with an implementation with a moderate level of optimisation in the Dyablo (left) and the Shamrock (right) codes. We provide the results obtained with the RK1 algorithm for comparison. The tests with Dyablo used static arrays and global memory. The tests with Shamrock used dynamical arrays and either the local (EXP-LOC) or the global (EXP-GLB) memory. For 8-16 dust fluids, the computational cost of the drag step is ∼ 10−100% of the hydro step, which makes it usable for astrophysical applications.

Fig. 7 Performances of the EXP drag solver on a Dustywave test performed on a A100 GPU using Dyablo for Ncells=643 (left) and Ncells=2563 cells for 1 to 15 dust fluids, respectively. The top panels show the time per cell for the drag step. The bottom panels show the relative cost of the drag step with respect to the total time step (hydro+drag) per time step.

4 Conclusion

We have presented and studied the EXP algorithm that we designed to compute the momentum exchange for dust-gas polydisperse mixtures very accurately on GPU architectures. The algorithm is based on a Taylor-based scaling and squaring matrix exponential computation. It was implemented in two different codes, namely Dyablo and Shamrock, using standard moderate optimisation strategies, and it was benchmarked on standard drag-dominated and hydro-dominated tests. The performance measurements showed that the cost of the drag time step does not significantly affect the code performance for up to ∼ 8−16 dust fluids.

The method is therefore suitable for astrophysical simulations even with growth and fragmentation as long as high-order solvers are used (Lombart & Laibe 2021; Lombart et al. 2022, 2024). The algorithm can also be applied for linear or linearised source terms, such as in radiative transfer and chemistry.

Data availability

The Dyablo code and the relevant sources are publicly distributed on Github (https://github.com/Dyablo-HPC/Dyablo). The Shamrock code and the relevant sources are publicly distributed on Github under the open source CeCILL v2.1 license (https://github.com/Shamrock-code/Shamrock).

Acknowledgements

We thank T. David--Cléris, A. Durocher, M. Delorme, G. Lesur, U. Lebreuilly, G. Verrier, P. Hennebelle, F. Lovascio for useful comments and discussions, P. Huang for helping with the RK2 test, E. Quemener and the Centre Blaise Pascal de Simulation et de Modélisation Numérique for support on the benchmarks, the ENS de Lyon for partly funding the DGX Nvidia Workstation and the electric costs associated with local simulations. We acknowledge support from the WP 5.1 MHD@Exascale (ANR-22-EXOR-0015) of the Programme et équipements prioritaires de recherche (PEPR) Origins (PI: A. Morbidelli). The SHAMROCK code follows the developments of the WP 5.1 of PEPR Origins. We acknowledge funding from the ERC CoG project PODCAST No. 864965. BC also acknowledges the funding from program ANR-20-CE49-0006 (ANR DISKBUILD) and funding from the European High Performance Computing Joint Undertaking (JU) under grant agreement No. 101093441, supported by contributions from multiple member states, including ANR (ANR-23-EHPC-0007-02). We thank the referee whose comments, questions and suggestions help to improve the quality of the paper.

Appendix A Drag integrators

To solve for Eqs. (7)–(11), we first integrated the transport and the source terms between the steps (n) and $(n+\frac{1}{2})$. We then solved for $$\frac{\partial \tilde{\mathit{U}}}{\partial t}={\tilde{\mathit{f}}}_{\text{drag}},$$(A.1) where $$\tilde{\mathit{U}}\equiv {[{\rho }_{\mathrm{g}}{\mathit{v}}_{\mathrm{g}},{\rho }_{\mathrm{d},1}{\mathit{v}}_{\mathrm{d},1},...,{\rho }_{\mathrm{d},N}{\mathit{v}}_{\mathrm{d},N}]}^T={[{\mathit{M}}_{\mathrm{g}},{\mathit{M}}_{\mathrm{d},1},...,{\mathit{M}}_{\mathrm{d},N}]}^T,$$(A.2) and f̃_drag is defined by Eq. (12), between steps $(n+\frac{1}{2})$ and (n+1). For IMEX (Implicit-Explicit)-like schemes, momenta terms should be modified accordingly for consistency. In the following, we review the different algorithms used to solve Eqs. (A.1)–(A.2).

Appendix A.1 General first-order implicit RK1

Denoting Δ Ũ=Ũ⁽ⁿ⁺¹⁾−Ũ⁽ⁿ⁾, and integrating Eq. (A.1) from t=tⁿ to t=tⁿ⁺¹ provides the first-order implicit RK1 method and $$\begin{array}{rl}\frac{\mathrm{\Delta }\tilde{U}}{\mathrm{\Delta }t}& ={\tilde{\mathit{f}}}_{\text{drag\hspace{0.17em}}}\left({\tilde{\mathit{U}}}^{(n+1)}\right)\\ & ={\tilde{\mathit{f}}}_{\text{drag\hspace{0.17em}}}({\tilde{\mathit{U}}}^{(n)}+({\tilde{\mathit{U}}}^{(n+1)}-{\tilde{\mathit{U}}}^{(n)}))\\ & \approx {\tilde{\mathit{f}}}_{\text{drag\hspace{0.17em}}}\left({\tilde{\mathit{U}}}^{n)}\right)+\mathbf{\Omega }\mathrm{\Delta }\tilde{\mathit{U}}\end{array}$$(A.3) where the Jacobian matrix Ω is given by Eq. (17). Corrections for IMEX-like schemes are taken into account by replacing f̃_drag in Eq. (A.1) by f̃ given by $$\tilde{\mathit{f}}={\tilde{\mathit{f}}}_{\text{drag}}+\frac{({\tilde{\mathit{U}}}^{(n+\frac{1}{2})}-{\tilde{\mathit{U}}}^{\left(n\right)})}{\Delta t},$$(A.4) where Ũ⁽ⁿ⁾, Ũ^(n+½) denote the momenta vectors before and after the hydrodynamic steps respectively. As such, one solves for $${\tilde{\mathit{U}}}^{(n+1)}-{\tilde{\mathit{U}}}^{\left(n\right)}={(\mathit{I}-\Delta t\mathit{\Omega })}^{-1}[\Delta t{\tilde{\mathit{f}}}_{\text{drag}}\left({\tilde{\mathit{U}}}^{\left(n\right)}\right)+({\tilde{\mathit{U}}}^{(n+\frac{1}{2})}-{\tilde{\mathit{U}}}^{\left(n\right)})].$$(A.5)

For N dust fluids, Eq. (A.5) can be solved using direct methods with a time complexity of O(N³) such as PLU decomposition. This algorithm was used by Benítez-Llambay et al. (2019), only for testing purposes.

Appendix A.2 Implicit RK1 with analytical solution

A more effective alternative of complexity O(N) consists of using an analytical solution of the sparse matrix Eq. (A.5). For a given vector b=(b₀, b₁, …, b_N), we aim to find the vector q=(q₀, q₁, …, q_N) such that $$(\mathit{I}-\Delta t\mathit{\Omega })\mathit{q}=\mathit{b}.$$(A.6) With some algebra $$(A.6)⟺\{\begin{array}{c}(1+{\sum }_{i=1}^N{ϵ}_i{\alpha }_i\Delta t)q_0-{\sum }_{i=1}^N{\alpha }_i\Delta t{q}_i=b_0,\hfill \\ -{ϵ}_1{\alpha }_1\Delta t{q}_0+(1+{\alpha }_1\Delta t)q_1=b_1,\hfill \\ ⋮\hfill \\ -{ϵ}_N{\alpha }_N\Delta t{q}_0+(1+{\alpha }_N\Delta t)q_N=b_N.\hfill \end{array}(\mathrm{A}.7)$$(A.7)

From lines 2−(N+1) of Eq. (A.7) we have for i=1, …, N $$q_i=\frac{1}{(1+{\alpha }_i\Delta t)}[b_i+\left({ϵ}_i{\alpha }_i\Delta t\right)q_0].$$(A.8) Replacing Eq. (A.8) back into the first line of Eq. (A.7) gives $$q_0=\frac{b_0+\Delta t\mathcal{B}}{\mathcal{A}},$$(A.9) with $\mathcal{A}=1+{\sum }_{i=1}^N(1-\frac{1}{1+{\alpha }_i\Delta t}){ϵ}_i$ and $\mathcal{B}={\sum }_{i=1}^N\left(\frac{{\alpha }_i}{1+{\alpha }_i\Delta t}\right)b_i$.

Appendix A.3 Second-order implicit RK2

The second-order implicit RK2 drag solver proposed by Huang & Bai (2022) is consistent with the two-stage second-order Total Variation Diminishing (TVD) Runge-Kutta scheme of Gottlieb & Shu (1998). The first stage of the drag integration step consists of a first-order implicit RK1. The second stage of the drag integration step is described as follows $${\tilde{\mathit{U}}}^{(n+1)}-{\tilde{\mathit{U}}}^{\left(n\right)}=\frac{\Delta t}{2}\tilde{\mathit{f}}({\tilde{\mathit{U}}}^{\left(n\right)},{\tilde{\mathit{V}}}^{\left(n\right)})+\frac{\Delta t}{2}\tilde{\mathit{f}}({\tilde{\mathit{U}}}^{(n+1)},{\tilde{\mathit{V}}}^{(n+\frac{1}{2})}),$$(A.10) where $$\tilde{\mathit{V}}\equiv (ϵ,\alpha ),$$(A.11) with ε=[ε₁, ·s, ε_n] and α=[α₁, ·s, α_n].

From a Taylor expansion in ‖Ω‖ Δ t, we have $$\begin{array}{rl}\tilde{\mathit{f}}({\tilde{\mathit{U}}}^{(n)},{\tilde{\mathit{V}}}^{(n)})& =\tilde{\mathit{f}}({\tilde{\mathit{U}}}^{(n+1)}-\mathrm{\Delta }t\tilde{\mathit{f}}({\tilde{\mathit{U}}}^{(n+1)},{\tilde{\mathit{V}}}^{(n)}),{\tilde{\mathit{V}}}^{(n)})\\ & =(\mathit{I}-\mathrm{\Delta }t\mathbf{\Omega }({\tilde{\mathit{U}}}^{(n+1)},{\tilde{\mathit{V}}}^{(n)}))\tilde{\mathit{f}}({\tilde{\mathit{U}}}^{(n+1)},{\tilde{\mathit{V}}}^{(n)})\\ & =(\mathit{I}-\mathrm{\Delta }t\mathbf{\Omega }({\tilde{\mathit{U}}}^{(n+1)},{\tilde{\mathit{V}}}^{(n)}))(\tilde{\mathit{f}}({\tilde{\mathit{U}}}^{(n)},{\tilde{\mathit{V}}}^{(n)})+\\ \mathbf{\Omega }({\tilde{\mathit{U}}}^{(n)},{\tilde{\mathit{V}}}^{(n)})\mathrm{\Delta }\tilde{\mathit{U}}).& \end{array}$$(A.12) Similarly, $$\tilde{\mathit{f}}({\tilde{\mathit{U}}}^{(n+1)},{\tilde{\mathit{V}}}^{(n+\frac{1}{2})})=\tilde{\mathit{f}}({\tilde{\mathit{U}}}^{\left(n\right)},{\tilde{\mathit{V}}}^{(n+\frac{1}{2})})+\mathit{\Omega }({\tilde{\mathit{U}}}^{\left(n\right)},{\tilde{\mathit{V}}}^{(n+\frac{1}{2})})\Delta \tilde{\mathit{U}}.$$(A.13)

Replacing Eqs. (A.12)–(A.13) into Eq. (A.10) gives $$\begin{array}{c}\hfill \Delta \tilde{\mathit{U}}=\frac{\Delta t}{2}{\mathit{\Lambda }}^{-1}[\tilde{\mathit{f}}({\tilde{\mathit{U}}}^{\left(n\right)},{\tilde{\mathit{V}}}^{(n+\frac{1}{2})})+(\mathit{I}-\Delta t\mathit{\Omega }({\tilde{\mathit{U}}}^{(n+1)},{\tilde{\mathit{V}}}^{\left(n\right)}))\\ \hfill \tilde{\mathit{f}}({\tilde{\mathit{U}}}^{\left(n\right)},{\tilde{\mathit{V}}}^{\left(n\right)})]\end{array}(\mathrm{A}.14)$$(A.14) where $$\mathit{\Lambda }\equiv \mathit{I}-\frac{\Delta t}{2}(\mathit{I}-\Delta t\mathit{\Omega }({\tilde{\mathit{U}}}^{(n+1)},{\tilde{\mathit{V}}}^{\left(n\right)}))\mathit{\Omega }({\tilde{\mathit{U}}}^{\left(n\right)},{\tilde{\mathit{V}}}^{\left(n\right)})-\mathit{\Omega }({\tilde{\mathit{U}}}^{\left(n\right)},{\tilde{\mathit{V}}}^{(n+\frac{1}{2})}).$$(A.15) The inverse of Λ is obtained by a PLU decomposition.

Appendix A.4 Second-order diagonally implicit Runge-Kutta MDIRK

As an alternative to operator splitting, Krapp et al. (2024) suggest combining the explicit and implicit terms in a single time step when integrating from a step n to a step n+1, as $${\mathit{U}}^{(*)}={\mathit{U}}^{\left(n\right)}+\tilde{\gamma }\Delta t\mathcal{L}\left({\mathit{U}}^{\left(n\right)}\right),$$(A.16) $${\mathit{U}}^{(**)}={\mathit{U}}^{\left(n\right)}+(1-\delta )\Delta t\mathcal{L}({\mathit{U}}^{(*)}+\gamma \Delta t{\mathit{K}}^{(*)}),$$(A.17) $${\mathit{U}}^{\left(\mathit{p}\right)}={\mathit{U}}^{(**)}+\frac{\delta }{\tilde{\gamma }}({\mathit{U}}^{(*)}-{\mathit{U}}^{\left(n\right)})+b_2\Delta t{\mathbf{K}}^{\left(p\right)},$$(A.18) $${\mathit{U}}^{(\mathit{n}+\mathbf{1})}={\mathit{U}}^{\left(p\right)}+(b_1-\beta )\Delta t{\mathit{K}}^{(*)}+b_2\Delta t{\mathit{K}}^{\left(p\right)},$$(A.19) where $$\mathit{U}\equiv$$(A.20) is defined for both gas and dust. The operator $$\mathcal{L}\left({\mathit{U}}_i^{\left(n\right)}\right)=-\frac{1}{dx}({\mathit{F}}_{i+1/2}\left({\mathit{U}}^{\left(\mathit{n}\right)}\right)-{\mathit{F}}_{i-1/2}\left({\mathit{U}}^{\left(\mathit{n}\right)}\right)),$$(A.21) corresponds to the transport step. It returns flux divergence as described by Eq. (A.21) and $\tilde{\gamma }=\gamma ,\delta =1-1/\left(2\gamma \right),\beta =1-\gamma ,b_1=1-\gamma ,b_2=\gamma .{\mathit{K}}^{\left(p\right)}$ and K^(*) are obtained by solving the system $$(\mathit{I}-\gamma \Delta t\mathit{\Omega }){\mathit{K}}^{\left(p\right)}=\mathit{\Omega }{M}^{\left(p\right)},$$(A.22) $$(\mathit{I}-\gamma \Delta t\mathit{\Omega }){\mathit{K}}^{(*)}=\mathit{\Omega }{M}^{(*)}.$$(A.23) where $$\mathit{K}\equiv (k_{\mathrm{g}}^n,k_{\mathrm{d},1}^n,\cdots ,k_{\mathrm{d},N-1}^n).$$(A.24)

The resolution of Eqs. (A.22)–(A.24) can be performed analytically, similarly to Appendix A.2.

Appendix B Modified Taylor-based scaling and squaring algorithm

We explain the different steps involved in the modified Taylor-based scaling and squaring algorithm of Sastre et al. (2015) described in Sect. 2.3.1.

Appendix B.1 Backward error analysis

The choice of polynomial order m and the scale factor s is based on a backward error analysis, given by the following theorem (Sastre et al. 2015).

Theorem 1. Let the matrix exponential Taylor approximation T_m(A) of a given matrix A as defined by Eq. (19), satisfy $${\mathrm{e}}^{-2^{-s}\mathit{A}}{T}_m\left(2^{-s}\mathit{A}\right)=\mathit{I}+\mathit{G},$$(B.1) with ‖G‖<1. Then $${\left[T_m\left(2^{-s}\mathit{A}\right)\right]}^{2s}={\mathrm{e}}^{\mathit{A}+\mathit{E}},$$(B.2) where E commutes with A and $$\frac{\parallel \mathit{E}\parallel }{\parallel \mathit{A}\parallel }\le \frac{-log(1-\parallel \mathit{G}\parallel )}{\parallel 2^{-s}\mathit{A}\parallel }.$$(B.3)

The residue E quantifies the error between the exact exponential of the matrix and its numerical approximation. To improve the estimate, one should minimize the right-hand side of Eq. (B.3). This term can be expressed as a function of θ ≡ ‖2^−s A‖. For this purpose, define the function $$\rho \left(x\right)={\mathrm{e}}^{-x}{T}_m\left(x\right)-1.$$(B.4) One has $$\begin{array}{rl}\rho (x)& =-{\mathrm{e}}^{-x}[{\mathrm{e}}^x-T_m(x)]\\ & =-{\mathrm{e}}^{-x}\left(\sum _{j\ge m+1}\frac{x^j}{j!}\right)\\ & =\left(\sum _{j\ge 0}\frac{(-1)x^j}{j!}\right)\left(\sum _{j\ge m+1}\frac{x^j}{j!}\right)\\ & =\sum _{j\ge m+1}{c}_j{x}^j\end{array}$$(B.5) From Eq. (B.1), $$\parallel \mathit{G}\parallel =\parallel \rho \left(2^{-s}\mathit{A}\right)\parallel \le {\sum }_{j\ge m+1}\left|c_j\right|{\theta }^j\equiv f_m\left(\theta \right).$$(B.6)

Combining Eqs. (B.3) and (B.3) gives the bound $$\frac{\parallel \mathit{E}\parallel }{\parallel \mathit{A}\parallel }\le \frac{-log(1-f_m\left(\theta \right))}{\theta }.$$(B.7) Using a zero-finder tool (symbolic Math Toolbox of MATLAB), we determined the highest value of θ ≡‖2^−s A‖ for a range of m such that the bound is at most equal to u=2⁻⁵³ (the unit round-off in IEEE double precision arithmetic). $${\theta }_m\equiv max\{\theta :\frac{-log(1-f_m\left(\theta \right))}{\theta }\le u\},$$(B.8) u=2⁻⁵³ being the unit round-off in IEEE double precision arithmetic. One can take another value depending on the desired accuracy : 2⁻¹¹, 2⁻²⁴, 2⁻¹¹³ and 2⁻²³⁷ for half, single, quadruple and octuple precision arithmetic respectively.

Appendix B.2 Numerical cost

In the evaluation step, the matrix T_m(2^−s A) is calculated from Eq. (B.1). The matrix polynomial is computed with the Horner-Paterson-Stockmeyer method to minimize the number of matrix products involved (Higham 2005). In the rest of the paper we used $${\Pi }_m\equiv \lfloor \sqrt{m}\rfloor +\lfloor m/\sqrt{m}\rfloor ,$$(B.9) the optimal number of matrix multiplications in the scaling phase with a Taylor approximation of order m, where ⌊ x⌋ denotes the integer part of the real number x. One defines the scaling factor s by $$s\equiv \{\begin{array}{cc}\lfloor {log}_2(\parallel \mathit{A}\parallel /{\theta }_m)\rfloor ,\hfill & \parallel \mathit{A}\parallel \ge {\theta }_m\hfill \\ 0\hfill & \parallel \mathit{A}\parallel <{\theta }_m\hfill \end{array}.$$(B.10)

The total cost in terms of matrix multiplications to compute the order- m Taylor-based matrix exponential approximant is (Higham 2005) $${\Pi }_m+s={\Pi }_m+max(\lfloor {log}_2(\parallel \mathit{A}\parallel /{\theta }_m)\rfloor ,0).$$(B.11) The optimal truncate order of the Taylor polynomial is the one that minimizes the quantity Π_m+s. For ‖A‖ ≥ θ_m, the optimal m is such that Γ_m=Π_m−log₂(θ_m) is the minimum cost, since ‖A‖ does not depend on m.

Appendix B.3 First modification for choosing m and s

For a maximum number K of matrix product in the scaling step, we are looking for k^* (optimal number of matrix product), m^* (optimal Taylor polynomial order associated with k^*) and s^* (optimal scale factor). The total cost is then k^*+s^*.

The algorithm 1 presents the different strategy to choose the optimal order m^* and the optimal scale factor s^*. The algorithm 1 with option 0 is the standard algorithm to choose m^* and s^*. The first modification proposed by Sastre et al. is to reduce the total cost (k^*+s^*) in some cases (option 1 and option 2 in algorithm 1). Option 1 involves decreasing the number of matrix products by one, while option 2 involves the cost of additional squaring to select k^*<K so that the overall cost is less than one given by option 1.

The high number of branching logic in this algorithm is subject to a high rate of branch divergence on GPUs. So, we rewrote algorithm 1 for option 2 to reduce the number of if-else statements, and then reduce branch divergence. The pseudo-code for the new algorithm is presented in algorithm. 2.

Appendix B.4 Second modification

To evaluate the matrix polynomial of m order in Eq. (19), the Paterson-Stockmeyer method (Paterson & Stockmeyer 1973) consists of choosing an integer q<m and writing the matrix polynomial as a polynomial of degree $r=\lfloor \frac{m}{q}\rfloor$ in A^q, such that $$T_m\left(\mathit{A}\right)={\sum }_{i=0}^r{\mathit{B}}_i{\left({\mathit{A}}^q\right)}^i,$$(B.12) where B_i, i=0, …, r are matrix polynomials of degrees q $${\mathit{B}}_i\equiv {\sum }_{j=0}^{q-1}{b}_{qi+j}{\mathit{A}}^j$$(B.13) and $${\mathit{B}}_r\equiv b_{rq}\mathit{I}+...+b_m{\mathit{A}}_{m-qr}.$$(B.14)

The algorithm 3 presents the procedure for evaluating the Paterson-Stockmeyer coefficients r{B_k}.

Evaluating the matrix polynomial from Eq. (B.12) using the Paterson-Stockmeyer algorithm in combination with the Horner

technique reduces the number of matrix products (Higham 2005). It is called the Horner-Paterson-Stockmeyer method. Equation (B.14) shows that when q divides m, one has B_r=b_m I. The values of q are given in the third column of table B.1. And m_k+1=m_k+q_k. Thus, a reduction of the number of matrix products by 1 reduces the order of the Taylor polynomial by q_k. Based on this remark, when the inequality $$\frac{\parallel T_{m_k}\left(\mathit{A}\right)-T_{m_{k-1}}\left(\mathit{A}\right)\parallel }{\parallel {\mathrm{e}}^{\mathit{A}}\parallel }\le u,$$(B.15) holds, Sastre et al. proposed to neglect the q highest terms in T_mk(A). This condition is verified by performing a numerical

test, which consists of checking if the inequality $$\parallel {\mathrm{e}}^{-\mathit{A}}\parallel \parallel {\overline{B}}_{r-1}\parallel {\parallel {\mathit{A}}^q\parallel }^{r-1}\le u,$$(B.16) is satisfied. To do so, one estimates a bound for ‖e^−A‖ according to $$\parallel {\mathrm{e}}^{-\mathit{A}}\parallel \approx \parallel T_{m_k}(-\mathit{A})\parallel \le b_{exp}=\parallel b_0\mathit{I}+{\widehat{\mathit{B}}}_0\parallel +{\sum }_{l=1}^{r-1}\parallel {\widehat{\mathit{B}}}_l\parallel {\parallel {\mathit{A}}^q\parallel }^l,$$(B.17) where ${\widehat{\mathit{B}}}_l={\sum }_{j=1}^q{(-1)}^{ql+j}{b}_{ql+j}{\mathit{A}}^j,l=0,...,r-1$. The condition given by Eq. (B.16) can therefore be replaced in practice by $$\parallel T_{m_k}\left(\mathit{A}\right)-T_{m_{k-1}}\left(\mathit{A}\right)\parallel \le \parallel {\overline{B}}_{r-1}\parallel {\parallel {\mathit{A}}^q\parallel }^{r-1},$$(B.18) from Eq. (B.16). The numerical procedure to get b_exp the bound of ‖e^−A‖ is detailed in algorithm 8.

Algorithm 5 presents the Horner-Paterson-Stockmeyer procedure to evaluate the Taylor polynomial or order m during the scaling phase with the bound performance test to reduce the number of matrix products.

Together, the two modifications are resumed in the algorithm 6.

As mentioned in Sect. 2.3.2, algorithm 9 shows the new algorithm we proposed, which combines algorithms 3 and 5 (without bound tests) together.

References