SciPost Submission Page
Python-JAX-based Fast Stokesian Dynamics
by Kim William Torre, Raoul D. Schram, Joost de Graaf
Submission summary
| Authors (as registered SciPost users): | Kim William Torre |
| Submission information | |
|---|---|
| Preprint Link: | https://arxiv.org/abs/2503.07847v1 (pdf) |
| Code repository: | https://github.com/torrewk/Python-Jax-Fast-Stokesian-Dynamics |
| Code version: | v0.2.0 |
| Code license: | Apache-2.0 |
| Date submitted: | 2025-03-12 14:15 |
| Submitted by: | Torre, Kim William |
| Submitted to: | SciPost Physics Codebases |
| Ontological classification | |
|---|---|
| Academic field: | Physics |
| Specialties: |
|
| Approaches: | Theoretical, Computational |
Abstract
Stokesian Dynamics (SD) is a powerful computational framework for simulating the motion of particles in a viscous Newtonian fluid under Stokes-flow conditions. Traditional SD implementations can be computationally expensive as they rely on the inversion of large mobility matrices to determine hydrodynamic interactions. Recently, however, the simulation of thermalized systems with large numbers of particles has become feasible [Fiore and Swan, J. Fluid. Mech. $\textbf{878}$, 544 (2019)]. Their ``fast Stokesian dynamics'' (FSD) method leverages a saddle-point formulation to ensure overall scaling of the algorithm that is linear in the number of particles $\mathcal{O}(N)$; performance relies on dedicated graphics-processing-unit computing. Here, we present a different route toward implementing FSD, which instead leverages the Just-in-Time (JIT) compilation capabilities of Google JAX. We refer to this implementation as JFSD and perform benchmarks on it to verify that it has the right scaling and is sufficiently fast by the standards of modern computational physics. In addition, we provide a series of physical test cases that help ensure accuracy and robustness, as the code undergoes further development. Thus, JFSD is ready to facilitate the study of hydrodynamic effects in particle suspensions across the domains of soft, active, and granular matter.
Current status:
Reports on this Submission
Strengths
1- Provides a new, faster implementation of Stokesian Dynamics for low Reynolds number flows of suspensions
2- Thorough validation of correct performance against existing literature
3- Good benchmarking for speed against current state of the art
4- Clever use of JIT compilation which scales well for large systems (though it is poor for small systems because of the slow initial compilation step)
Weaknesses
1- Small grammatical errors - a few sentences that are not sentences (e.g. page 1, line 5; page 5, line 3)
2- Could add plans for extending to planar extensional flow (using Kraynik-Reinelt boundary conditions) which would strengthen the software considerably
3- In section 3, I would welcome more discussion of why box size affects the performance of the code
4- Formatting of references is poor (many missing capital letters in article titles)
Report
This is a lovely paper. It presents very clearly why Stokesian Dynamics is still a method worth using, and provides a new, faster implementation. Everything is properly validated and benchmarked, and example applications provided in detail. As far as I can tell this is being made available in a coherent and easy-to-access way, though that is not my area of expertise.
Requested changes
1- Correct small grammatical errors at page 1, line 5; page 5, line 3
2- In section 3, explain why box size affects the performance of the code
3- Correct formatting of references
Recommendation
Publish (easily meets expectations and criteria for this Journal; among top 50%)