Python-JAX-based Fast Stokesian Dynamics
Extracted source text · arXiv 2503.07847
% =========================================================================
% SciPost LaTeX template
% Version 2024-07
%
% Submissions to SciPost Journals should make use of this template.
%
% INSTRUCTIONS: simply look for the `TODO:' tokens and adapt your file.
% ========================================================================
\documentclass{SciPost}
%\documentclass[nolinelabels]{SciPost}
% Prevent all line breaks in inline equations.
\binoppenalty=10000
\relpenalty=10000
\hypersetup{
colorlinks,
linkcolor={red!50!black},
citecolor={blue!50!black},
urlcolor={blue!80!black}
}
\usepackage[bitstream-charter]{mathdesign}
\urlstyle{same}
% Fix \cal and \mathcal characters look (so it's not the same as \mathscr)
\DeclareSymbolFont{usualmathcal}{OMS}{cmsy}{m}{n}
\DeclareSymbolFontAlphabet{\mathcal}{usualmathcal}
\fancypagestyle{SPstyle}{
\fancyhf{}
\lhead{\colorbox{scipostblue}{\bf \color{white} ~SciPost Physics Codebases }}
\rhead{{\bf \color{scipostdeepblue} ~Submission }}
\renewcommand{\headrulewidth}{1pt}
\fancyfoot[C]{\textbf{\thepage}}
}
\begin{document}
\pagestyle{SPstyle}
\begin{center}{\Large \textbf{\color{scipostdeepblue}{
\textbf{\texttt{Python-JAX}}-based Fast Stokesian Dynamics\\
}}}\end{center}
\begin{center}\textbf{
%%%%%%%%%% TODO: AUTHORS
% Write the author list here.
% Use (full) first name (+ middle name initials) + surname format.
% Separate subsequent authors by a comma, omit comma and use "and" for the last author.
% Mark the corresponding author(s) with a superscript symbol in this order
% \star, \dagger, \ddagger, \circ, \S, \P, \parallel, ...
Kim William Torre\textsuperscript{1$\star$},
Raoul D. Schram\textsuperscript{2$\ddagger$}, and
Joost de Graaf\textsuperscript{1$\dagger$}
%%%%%%%%%% END TODO: AUTHORS
}\end{center}
\begin{center}
{\bf 1} Institute for Theoretical Physics, Center for Extreme Matter and Emergent Phenomena, Utrecht University, Princetonplein 5, 3584 CC Utrecht, The Netherlands
\\
{\bf 2} Research and Data Management Services, Information and Technology Services, Utrecht University, Heidelberglaan 8, 3584 CS Utrecht, The Netherlands
\\[\baselineskip]
$\star$ \href{mailto:email1}{\small k.w.torre@uu.nl}\,,\quad
$\ddagger$ \href{mailto:email2}{\small r.d.schram@uu.nl}\,,\quad
$\dagger$ \href{mailto:email3}{\small j.degraaf@uu.nl}
\end{center}
\section*{\color{scipostdeepblue}{Abstract}}
\textbf{\boldmath{%
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)]\nocite{fiore2019fast}. 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 \textbf{\texttt{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.
}}
\vspace{\baselineskip}
%%%%%%%%%% BLOCK: Copyright information
% This block will be filled during the proof stage, and finilized just before publication.
% It exists here only as a placeholder, and should not be modified by authors.
\noindent\textcolor{white!90!black}{%
\fbox{\parbox{0.975\linewidth}{%
\textcolor{white!40!black}{\begin{tabular}{lr}%
\begin{minipage}{0.6\textwidth}%
{\small Copyright attribution to authors. \newline
This work is a submission to SciPost Physics Codebases. \newline
License information to appear upon publication. \newline
Publication information to appear upon publication.}
\end{minipage} & \begin{minipage}{0.4\textwidth}
{\small Received Date \newline Accepted Date \newline Published Date}%
\end{minipage}
\end{tabular}}
}}
}
%%%%%%%%%% BLOCK: Copyright information
%%%%%%%%%% TODO: LINENO
% For convenience during refereeing we turn on line numbers:
%\linenumbers
% You should run LaTeX twice in order for the line numbers to appear.
%%%%%%%%%% END TODO: LINENO
%%%%%%%%%% TODO: TOC
% Guideline: if your paper is longer that 6 pages, include a TOC
% To remove the TOC, simply cut the following block
\vspace{10pt}
\noindent\rule{\textwidth}{1pt}
\tableofcontents
\noindent\rule{\textwidth}{1pt}
\vspace{10pt}
%%%%%%%%%% END TODO: TOC
%%%%%%%%% TODO: CONTENTS
% Write your article contents here, starting from first \section.
% An example structure is given below.
\section{Introduction}
\label{sec:intro}
The study of hydrodynamic interactions in particle suspensions has attracted considerable interest in colloid science and fluid mechanics over the past decades~\cite{batchelor1977effect,happel1983mechanics,kim2013microhydrodynamics}. The development of Stokesian Dynamics (SD) by Brady and Bossis~\cite{brady1988stokesian} in the late 1980s provided an accurate framework for simulating the many-body hydrodynamic interactions between spherical particles at zero (low) Reynolds number. That is, in the regime where inertia is completely dominated by friction. The Stokes equations that describe the fluid dynamics of such systems are linear. This is leveraged in SD to divide the mobility problem into a far- and near-field component. The split enables the precise modeling of long-range, many-body hydrodynamic interactions and short-range, lubrication forces~\cite{jeffrey1984calculation, kim1985resistance, jeffrey1992calculation}, both of which can strongly influence the behavior of colloidal suspensions~\cite{eckstein1977self, furukawa2010key,whitmer2011influence,royall2015probing,varga2018modelling,swan-furst2019,turetta2022role, torre2023structuring, de2023hydrodynamic, torre2023hydrodynamic, TorredeGraaf2024}.
%
However, despite its accuracy, SD did not become as widely adopted as other hydrodynamics methods for particle suspensions~\cite{bolintineanu2014particle}. Among these, lattice-based methods for computational fluid dynamics gained considerable traction, including lattice-Boltzmann (LB)~\cite{chen1998lattice,ladd2001lattice,dunweg2009lattice,aidun2010lattice,rohm2012lattice}, multi-particle collision dynamics (MPCD)~\cite{gompper2009multi, howard2018efficient} --- or stochastic rotation dynamics (SRD)~\cite{malevanets1999mesoscopic} --- and Fluid Particle Dynamics (FPD)~\cite{tanaka2000simulation}.
%
Contrasting the features of these methods, we see the following advantages for SD. The algorithm relies on multipole expansion and pairwise approximations to compute hydrodynamic interactions without directly solving for the fluid flow. This means that at low dilution, there is an intrinsic advantage to using SD. However, LB and MPCD are particularly effective for determining fluid-structure interactions involving complex geometries~\cite{chen1998lattice, ladd2001lattice, dunweg2009lattice, aidun2010lattice, rohm2012lattice}. Yet, achieving the same level of accuracy in both many-body and lubrication interactions as SD requires the use of prohibitively fine resolution in LB, leading to higher computational demands~\cite{nguyen2002lubrication, aidun2010lattice}. MPCD and SRD operate on similar principles~\cite{malevanets1999mesoscopic, gompper2009multi, howard2018efficient}; however, they are intrinsically stochastic, with control parameters such as viscosity emerging as collective properties of the system. SD instead provides explicit control over thermal and athermal motion, like LB. Lastly, FPD, as introduced by Tanaka and Araki~\cite{tanaka2000simulation}, directly solves the Navier-Stokes equation on a lattice. It treats colloidal particles as regions of higher viscosity fluid with smooth interfacial profiles. This eliminates the need to explicitly resolve solid-fluid boundaries, but it introduces approximations. These depend on the viscosity ratio between the particles and surrounding fluid, as well as on the lattice resolution.
%
The major bottlenecks for the use of SD were: (i)~the computational cost involved in propagating the dynamics and (ii) that it is challenging to implement the algorithm correctly. Advancements following the original implementation~\cite{brady1988stokesian} have therefore focused on improving SD's computational performance~\cite{sierou2001accelerated,banchio2003accelerated,wang2016spectral}. These include the incorporation of iterative solvers and preconditioning strategies. Fiore and Swan overcame some of the last remaining hurdles in making SD truly efficient, through their Fast Stokesian Dynamics (FSD) approach~\cite{fiore2019fast}. They combined matrix-free techniques with spectral Ewald methods, which allowed them to further reduce computational costs, improve the scaling to $\mathcal{O}(N)$ overall, and maintain accuracy. They also made their algorithm widely available through integration in the \textbf{\texttt{HOOMD}}~\cite{ANDERSON2020109363} simulation package. This means that, provided that complex geometries are not critical to the analysis of a system, FSD provides a competitive alternative to more widely used lattice-based approaches. However, the efficiency obtained in FSD was limited to \textbf{\texttt{NVIDIA}}\textsuperscript{\textregistered} graphics processing units (GPUs), through an internal reliance on \textbf{\texttt{CUDA}}\textsuperscript{\textregistered} programming.
%
In this paper, we present an alternative route toward making FSD more widely available and cross-platform compatible. We have ported the FSD algorithm to a \textbf{\texttt{Python}}-based framework that relies on the \textbf{\texttt{Google JAX}} library~\cite{jax2018github}. We name this version of SD, \textbf{\texttt{Python-JAX}} Fast Stokesian Dynamics, or \textbf{\texttt{Python}}-JFSD for short. Our implementation ensures compatibility with contemporary \textbf{\texttt{CUDA}}\textsuperscript{\textregistered} and \textbf{\texttt{CuDNN}}\textsuperscript{\textregistered} libraries. It also addresses limitations in the original code caused by updates to the \textbf{\texttt{HOOMD}} application programming interface (API). \textbf{\texttt{Python}}-JFSD further offers a user-friendly interface and provides flexibility in choosing between periodic and open boundary conditions for hydrodynamic interactions. By leveraging \textbf{\texttt{Google JAX}}’s Just-In-Time (JIT) compilation~\cite{jax2018github}, the framework supports the efficient addition of new features, ensuring modularity and extensibility.
%
The remainder of this paper is organized as follows. In Section~\ref{sec:meth}, we first provide a brief overview of the FSD method. Section~\ref{sec:test} contains the results of a set of physical unit test that probe the implementation accuracy and benchmark its performances. In Section~\ref{sec:disc}, we discuss the applicability range of JFSD, as well as its strong and weak points. We close with a summary and outlook in Section~\ref{sec:concl}.
%%%%%%%%
\section{\label{sec:meth}Method}
%%%%%%%%
Consider a system comprising $N$ spherical particles of diameter $\sigma$, suspended in a Newtonian solvent with viscosity $\eta$. We have a prescribed background flow $\mathcal{U}^{\infty}$, which is a $6N$-dimensional vector containing the linear and angular fluid velocities at the position of each particle. We assume viscous forces dominate inertial ones, such that the Reynolds number $\mathrm{Re} \ll 1$, and the particles undergo over-damped motion. Under these conditions, the particle dynamics are prescribed by a linear system of equations that couple the dynamical degrees of freedom (forces and torques $\mathcal{F}$, and stresses $\mathcal{S}$) to kinematic degrees of freedom (linear and angular velocities $\mathcal{U}$ and strain rates $\mathcal{E}$) via the grand-resistance matrix $\mathcal{R}$. The relation is given by
\begin{align}
\label{eqn:granres_def}
\begin{pmatrix}
\mathcal{F}^H \\
\mathcal{S}^H
\end{pmatrix}
= - {\mathcal{R}} \cdot
\begin{pmatrix}
\mathcal{U} - \mathcal{U}^{\infty} \\
-\mathcal{E} ^{\infty}
\end{pmatrix},
\end{align}
where $\mathcal{U}$ is a $6N$-dimensional vector containing the linear and angular velocities of all particles, $\mathcal{E}^{\infty}$ is the strain rate of the background flow, $\mathcal{F}^H$ denotes the hydrodynamic drag forces and torques, and $\mathcal{S}^H$ is a $5N$-dimensional vector of the independent components of the hydrodynamic stresslets.
%
The generalized velocities of all $N$ particles can then be expressed as
\begin{align} \label{eq:particle_disp}
\mathcal{U} - \mathcal{U}^{\infty} &= \mathbf{R}_{\mathcal{FU}}^{-1} \cdot (\mathcal{F^P} + \mathbf{R}_{\mathcal{FE}} \cdot \mathcal{E}^{\infty}) + \sqrt{\frac{2k_{\mathrm{B}} T}{\Delta t}} \mathbf{R}_{\mathcal{FU}}^{-1/2} \cdot \mathcal{\psi} + k_{\mathrm{B}} T~\nabla \cdot \mathbf{R}_{\mathcal{FU}}^{-1}.
\end{align}
Here, $\mathcal{F^P}$ is a $6N$-dimensional vector whose first $3N$ components are the forces applied to the particles, and the last $3N$ components are the applied torques. The vector $\mathcal{\psi}$ is a $6N$-dimensional vector containing random variables normally distributed with zero mean and unit variance, and $\Delta t$ is the simulation timestep. The term $k_{\mathrm{B}} T$ represents the thermal energy; $k_{\mathrm{B}}$ is the Boltzmann constant and $T$ the temperature. The tensor $\mathbf{R}_{\mathcal{FU}}$ is the upper-left subblock of $\mathcal{R}$, which maps velocities to forces, while $\mathbf{R}_{\mathcal{FE}}$ is the upper-right subblock that couples strain rates to forces in $\mathcal{R}$, respectively. In Eq.~\eqref{eq:particle_disp}, the first two terms represent deterministic contributions to the particle dynamics coming from the applied forces and torques, as well as the background flow, while the last two terms account for thermal fluctuations.
%
In SD~\cite{brady1988stokesian}, hydrodynamic interactions are incorporated through a combination of the mobility and resistance frameworks. The construction of the grand resistance matrix $\mathcal{R}$ begins with the far-field mobility matrix $\mathcal{M}^{\text{ff}}$. Short-range interactions are resolved in a pairwise manner by adding near-field resistance contributions $\mathcal{R}^{\text{nf}}$~\cite{jeffrey1984calculation,jeffrey1992calculation}. To avoid double-counting, the far-field pairwise contributions $\mathcal{R}^{\text{ff}}_{\text{2B}}$ are subtracted
\begin{align} \label{eq:total_grandres_SD}
\mathcal{R} = (\mathcal{M}^{\text{ff}})^{-1} + \mathcal{R}^{\text{nf}} - \mathcal{R}^{\text{ff}}_{\text{2B}}.
\end{align}
%
The FSD method~\cite{fiore2019fast} extends the SD framework by formulating the particle dynamics as a linear system of $17N$ equations using a saddle-point matrix that integrates near-field (short-ranged, pairwise additive) and far-field (long-ranged, many-body) interactions. This formulation enables iterative solvers with specialized preconditioners to efficiently handle the resulting linear systems. The approach avoids the explicit inversion of ill-conditioned hydrodynamic operators and drastically reducing computational costs. Brownian forces are computed by combining the positively-split Ewald method~\cite{fiore2017rapid} and an iterative Krylov subspace method~\cite{chow2014preconditioned}, both seamlessly integrated into the saddle-point formulation. We refer the interested reader to Ref.~\cite{fiore2019fast} for the full details.
%
Our implementation of the method, JFSD, is structured as a modular Python package optimized for scalability and high-performance hydrodynamic simulations. In particular, the configuration management module leverages a user-friendly TOML file that specifies simulation parameters --- such as the number of steps, particle count, time step, initialization settings, and output options --- providing a transparent and easily modifiable interface for setting up simulations. The modular design allows for easy extension of pairwise interaction models, enabling users to implement custom hydrodynamic kernels or introduce \textit{ad hoc} force laws with minimal modifications. The framework also provides built-in support for steady and oscillatory shear flows, with a straightforward pathway toward incorporating more complex shear protocols by modifying existing velocity gradients. Boundary conditions can be adjusted efficiently by altering the scalar mobility functions, as demonstrated in the already implemented transition between open and periodic boundary conditions. Lastly, our neighbor list management was derived from the one used in JAX-MD~\cite{schoenholz2021jax}, a widely used molecular dynamics framework. Our implementation ensures FSD-compatible handling of particle interactions in both dense and dilute systems. Full details of our implementation can be found in the documentation accompanying the JFSD release~\cite{jfsdgithub}.
%
A key advantage of JFSD is its integration with \textbf{\texttt{Google JAX}}~\cite{jax2018github}, which enables Just-In-Time (JIT) compilation and GPU acceleration for significant performance improvements. JAX compiles entire computation graphs into highly optimized routines that eliminate Python overhead by fusing multiple operations into a single efficient kernel. Moreover, its vectorized operations and automatic backend dispatch ensure that computations run optimally on both GPU and CPU, making the framework highly scalable across diverse hardware configurations.
%%%%%%%%
\section{\label{sec:test}Method Validations}
%%%%%%%%
Following a general discussion of (F)SD, we now show results obtained using JFSD. These are contrasted against two-body analytical expressions and/or independently obtained numerical data. We also showcase the efficiency of the JFSD software by benchmarking the code for the case of a hard-sphere system across a range of volume fractions $\phi$.
%%%%%%%%
\subsection{Sedimenting Particles}
%%%%%%%%
The deterministic mapping from forces to velocities was validated by simulating the sedimentation of three particles aligned along the $x$-axis. Our results are compared to the original Stokesian Dynamics (SD) trajectories by Durlofsky~\textit{et al.}~\cite{durlofsky1987dynamic} for an unbounded fluid solvent in Fig.~\ref{fig:dancing_spheres}. This shows the sedimentation trajectories of the three particles. Here, JFSD demonstrates excellent agreement with the SD reference result. The trajectories aligning closely over the full time interval considered. With JFSD we also considered periodic boundaries, for which discrepancies become pronounced as time progresses. These deviations are expected and arise due to interaction with periodic images. Despite these differences, JFSD faithfully reproduces the initial dynamics. The combined result is sufficient to conclude that JFSD is consistent with the original SD framework.
%
\begin{figure}
\centerline{\includegraphics{Figure1_jfsd.pdf}}
\caption{\label{fig:dancing_spheres}
Trajectories of three equal spheres sedimenting vertically, starting from initial positions at $\boldsymbol{r}_1 = (-2.5\sigma,0,0)$ (green), $\boldsymbol{r}_2 = (0,0,0)$ (red), and $\boldsymbol{r}_3 = (3.5\sigma,0,0)$ (gray). Results obtained using the JFSD method under (a) open boundary conditions (circles) and (b) periodic boundary conditions (rhombi), in a box of size $L=25 \sigma$. Both panels include reference data from Ref.~\cite{durlofsky1987dynamic} (dashed curves), where the original data points have been smoothed using spline interpolation.}
\end{figure}
%
Next, we computed the sedimentation velocity $U_{s}$ of a periodic cubic array of $N=$~2,197 hard spheres across a range of $\phi$, see Fig.~\ref{fig:sedim_mobility}. These results were compared against the benchmark data from Sierou and Brady~\cite{sierou2001accelerated}, which revealed that JFSD accurately captures the non-linear decay of $U_{s}$ with increasing $\phi$. Our implementation gives consistent results even in dense regimes, which implies that it can reliably propagate suspension dynamics across a wide range of concentrations.
%
\begin{figure}
\centerline{\includegraphics{Figure2_jfsd.pdf}}
\caption{\label{fig:sedim_mobility}Comparison of sedimenting velocity $U_{s}$ (blue dots) for a simple cubic array of $N=$~2,197 particles (insets show examples) as a function of the volume fraction $\phi$. The sedimentation velocity is normalized by the infinite dilution limit $U_{0}$. Results from the JFSD method under periodic boundary conditions are contrasted with those obtained using Sierou's accelerated Stokesian dynamics~\cite{sierou2001accelerated} (red rhombi) and Brady~\textit{et al.}'s original Stokesian dynamics~\cite{brady1988dynamic} (gray squares). The curves connecting the two reference data sets, obtained through spline interpolation, serve to guide the eye.}
\end{figure}
%%%%%%%%
\subsection{Hydrodynamic Effects under Shear Flow}
%%%%%%%%
Now that we have examined the quality of the force-torque and (angular) velocity coupling, we turn to shear. We investigated the behavior of particle pairs in simple shear flow to benchmark the accuracy of JFSD under these conditions, see Fig.~\ref{fig:pairshear}, which shows three example trajectories. The results align closely with analytical predictions~\cite{da1996shear,kim2013microhydrodynamics}, demonstrating JFSD’s performance.
%
\begin{figure}
\centerline{\includegraphics{Figure3_jfsd.pdf}}
\caption{\label{fig:pairshear}Trajectories of a particle pair in simple shear flow under open boundary conditions. Results from JFSD (symbols) are compared against analytical predictions~\cite{da1996shear, kim2013microhydrodynamics} (dotted lines) for three initial center-to-center separations $\Delta r$. These trajectories are represented in the distance along the $x$- and $y$-axes, respectively, normalized by the particle diameter $\sigma$. The insets show three snapshots of a trajectory for spheres spaced $\Delta r = 1.01 \sigma$ apart orthogonal to the shear plane.
}
\end{figure}
%
In addition to pair interactions, we computed the shear viscosity $\eta_{\mathrm{eff}}$ of a periodic simple cubic array of (again) $N =$~2,197 hard spheres as a function of $\phi$. Figure~\ref{fig:shear_viscosity} compares our result with data from the original FSD implementation by Fiore and Swan~\cite{fiore2019fast}. The accurate agreement shows that both algorithms produce the same results and are capable of computing rheological properties even at high $\phi$.
%
\begin{figure}
\centerline{\includegraphics{Figure4_jfsd.pdf}}
\caption{\label{fig:shear_viscosity}Shear viscosity $\eta_{\mathrm{eff}}$ of a simple cubic array of $N =$~2,197 particles as a function of volume fraction $\phi$. The results are normalized by the bare viscosity of the suspending medium $\eta_{0}$. \textbf{\texttt{Python}}-JFSD results (blue circles) are compared with data from Fiore and Swan~\cite{fiore2019fast} (open squares). The dotted connecting curve serves to guide the eye.}
\end{figure}
%
Our combined shear results allow us to conclude that this part of the algorithm is correctly implemented. Of course, additional testing was performed on the behavior of the $\mathcal{R}$ matrix, but we chose only to report on physical observables here.
%%%%%%%%
\subsection{Thermal Motion}
%%%%%%%%
Having verified the deterministic part of the algorithm, we turned our attention to thermal fluctuations. The mean squared displacement (MSD) $\langle \boldsymbol{r}^{2} \rangle$ of a single particle in 3D was computed under periodic boundary conditions and benchmarked against Hasimoto's corrected expression for periodic boundaries~\cite{hasimoto1959periodic}. To generate the MSD curves in Fig.~\ref{fig:msd_periodic}, we simulated for $100$ Brownian times and repeated the process for $10$ independent samples for each box size $L$, averaging the results across these runs. The figure reveals that the simulated MSD aligns closely with theoretical predictions, confirming the accurate implementation of periodic images and the validity of thermal diffusion in the JFSD framework.
%
\begin{figure}
\centerline{\includegraphics{Figure5_jfsd.pdf}}
\caption{\label{fig:msd_periodic}Mean squared displacement $\langle \boldsymbol{r}^{2} \rangle$ of a single particle as a function of time $t$ divided by the Brownian time $t_{\mathrm{B}}$. The JFSD data is given by open circles with error bars representing the standard error of the mean for two box sizes $L$ as labelled. The data for the $L = 25 \sigma$ ($\sigma$ is the sphere diameter) box is shifted by the constant offset $+3$ in the vertical direction to help aid the presentation. The dotted line shows the analytic $\langle \boldsymbol{r}^{2} \rangle = 6Dt$ result for an unbounded fluid, while the solid line applies Hasimoto's correction to obtain the appropriate expression for a periodic domain~\cite{hasimoto1959periodic}.}
\end{figure}
%
However, MSD analysis alone is not sufficient to validate the thermal part of the algorithm. This is because in SD, there is a `Brownian drift' term that is required to ensure stationarity under the Gibbs–Boltzmann distribution and generate particle configurations with the correct statistics at equilibrium~\cite{fiore2018rapid,fiore2019fast}. To demonstrate that we have correctly implemented this aspect, we draw inspiration from a test originally considered by Fiore and Swan~\cite{fiore2018rapid,fiore2019fast}. That is, we simulated a pair of particles, in an unbounded three-dimensional (3D) fluid, interacting \textit{via} a linear potential
\begin{align} \label{eq:linear_potential}
V( r) = k | r - r_0|.
\end{align}
Here, $k$ is the potential strength, $ r$ is the radial distance between the sphere centers, and $ r_0$ is the equilibrium distance. Figure~\ref{fig:dimer_thermal} compares the average distribution of radial distances $p(r)$ derived from 5 independent simulations --- with duration of $5$ Brownian times each --- against the theoretical Boltzmann probability density function~\cite{fiore2019fast}. The good agreement demonstrates \textbf{\texttt{Python}}-JFSD's ability to reproduce expected equilibrium distributions. Small deviations in the tail of the distribution arise primarily from the accumulation of numerical errors in long trajectories, due to the use of a first-order Euler integration scheme~\cite{ermak1978brownian}.
\begin{figure}
\centerline{\includegraphics{Figure6_jfsd.pdf}}
\caption{\label{fig:dimer_thermal}Radial distribution $p(r) \propto r^2 \exp(-V(r)/k_{\mathrm{B}}T)$ of a pair of particles interacting \textit{via} a linear potential (see Eq.~\eqref{eq:linear_potential}) under open boundary conditions, with $r$ the the center-to-center distance between the two spheres. The JFSD data is shown using open circles, with the vertical error bars showing the standard error of the mean (computed from 5 independent samples). The horizontal error bars follow from the fact that the data was obtained by a binning procedure.}
\end{figure}
%
We also considered large particle number systems ($N =$~2,500) with periodic boundary conditions. For these, we computed the short-time self-diffusion coefficient $D_{s}$ as a function of $\phi$ for hard-sphere suspensions, see Fig.~\ref{fig:diffusion_random}. Our results are compared therein with two independently derived datasets. The first is by Ladd~\cite{ladd1990hydrodynamic} and was obtained from a multipole‐moment expansion, combined with pairwise lubrication interactions, that contains a higher number of moments than conventional SD. The second follows from Sierou's and Brady's work~\cite{sierou2001accelerated} and represents the Accellerated Stokesian Dynamics method. Both datasets were obtained for open systems, hence we have corrected for the periodicity effects using Hasimoto's scaling factor~\cite{hasimoto1959periodic,ladd1990hydrodynamic,sierou2001accelerated}. The close agreement between our results and the (scaled) literature values confirms JFSD's ability to resolve thermal fluctuations in systems with many particles in periodic boundary conditions.
%
\begin{figure}
\centerline{\includegraphics{Figure7_jfsd.pdf}}
\caption{\label{fig:diffusion_random}Short-time self-diffusion coefficient $D_{s}$ for hard-sphere suspensions under periodic boundary conditions as a function of volume fraction $\phi$. The diffusion coefficient is normalized by the bulk value $D_{0}$. Results from \textbf{\texttt{Python}}-JFSD (open circles) have error bars that show the standard error of the mean. These are compared to the work by Ladd~\cite{ladd1990hydrodynamic} (open circle) and Sierou~\cite{sierou2001accelerated} (red crosses). Both datasets are corrected for finite-size effects \textit{via} the Hasimoto factor~\cite{hasimoto1959periodic} and the dashed curves help guide the eye.}
\end{figure}
%%%%%%%%
\subsection{Algorithmic Performance and Scaling}
%%%%%%%%
We assessed the computational performance of \textbf{\texttt{Python}}-JFSD using an NVIDIA GeForce RTX 4060 Ti with 8GB of VRAM. As a standard performance metric, we define the particle time steps per second (PTPS) as $\text{PTPS} = N/ T_{\mathrm{step}}$, where $T_{\mathrm{step}}$ is the clock time in seconds required to perform a single simulation step as a function of the number of particles $N$. The setup consists of $N$ hard spheres arranged in a cubic array, interacting only via excluded-volume interactions to prevent overlap. The particles undergo thermal motion, and a simple shear flow is applied in the $xy$-plane (with vorticity in the $z$-direction, flow in the $x$-direction, and flow gradient in the $y$-direction). Periodic boundary conditions are imposed. To ensure robust benchmarking, the system is simulated for 50 steps, and the average $T_{\mathrm{step}}$ is computed after excluding the first step, which involves JIT compilation and can take up to 1–2 minutes.
%
Since computational performance depends on the volume fraction $\phi$, we considered four representative values, $\phi \in \{ 0.05, 0.15, 0.4, 0.5 \}$, to assess scaling behavior. This dependence arises from the iterative solver used in the FSD algorithm~\cite{fiore2019fast}, where the number of iterations required for convergence increases due to lubrication interactions, which become more significant at higher $\phi$ as particles are closer together. The results, shown in Fig.~\ref{fig:performance_scaling}, present PTPS as a function of $N$.
\begin{figure}
\centerline{\includegraphics{Figure8_jfsd.pdf}}
\caption{\label{fig:performance_scaling}Particle time steps per second (PTPS) as a function of the number of particles $N$ in a periodic simulation volume. Four volume fractions were considered: $\phi \in \{ 0.05, 0.15, 0.4, 0.5 \}$, as labeled. At low $\phi$, PTPS plateaus for large $N$, while for high $\phi$, PTPS decreases with $N$ due to system-size dependence. The inset shows the normalized step time per particle per unit box size, $T_{\mathrm{step}} / (N L)$, which remains constant for large $\phi$, confirming that the algorithm maintains linear scaling despite the decrease in PTPS in the main plot.}
\end{figure}
%
The PTPS trends indicate that at low $\phi$, performance reaches a steady plateau, as expected for sufficiently large $N$. Conversely, at higher volume fractions, PTPS continues to decrease with increasing $N$. This behavior suggests that system size influences solver convergence, which introduces additional computational overhead. The inset in Fig.~\ref{fig:performance_scaling} clarifies this effect: for large $\phi$, the normalized time step per particle per unit box size remains unchanged, verifying that the scaling of the FSD algorithm is preserved. This shows that although PTPS decreases at high $\phi$, the core computational efficiency of the method remains intact, consistent with theoretical predictions~\cite{fiore2019fast}.
%%%%%%%%
\section{\label{sec:disc}Discussion}
%%%%%%%%
SD is a well-established framework for simulating particle suspensions~\cite{durlofsky1987dynamic,brady1988stokesian,sierou2001accelerated,banchio2003accelerated,wang2016spectral,fiore2019fast}, capable of accurately resolving both many-body hydrodynamic interactions and near-field lubrication forces. It is particularly effective for systems requiring precise particle-level hydrodynamic modeling without complicated boundary conditions. SD relies on direct calculations of forces and torques, which allows for computational efficiency without the need for fine spatial discretization. Recent applications include sedimentation in colloidal suspensions~\cite{varga2018modelling}, colloidal gelation~\cite{torre2023hydrodynamic}, and the dynamics of active matter~\cite{ge2025hydrodynamic} in confining geometries~\cite{elfring2022active,shaik2023confined}. These examples showcase SD's versatility across a broad range of (colloidal) particulate systems.
%
Implementing SD methods, however, presents significant challenges due to the algorithm's reliance on complex mathematical algorithms~\cite{townsend2023generating} and the precision required to accurately model lubrication forces and many-body interactions. Ensuring correctness and robustness needs meticulous validation. To address this, our JFSD implementation incorporates continuous integration workflows, automatically checking each pull request -- originating from Git-based version control -- against an extensive suite of benchmarks~\cite{jfsdgithub}. This rigorous testing, combined with a modular design, allows users to choose between periodic or open boundary conditions and various hydrodynamic models, including FSD, which has been the focus of this paper. In addition to these features, our software suite is also capable of modeling hydrodynamic interactions at the Rotne-Prager-Yamakawa level~\cite{rotne1969variational,fiore2017rapid, pelaez2022complex}, which is suitable for capturing pair-wise, far-field coupling. We have further implemented a simple Brownian Dynamics code, which focuses on the hydrodynamic self-interaction only. However, this last algorithm has not been optimized to the same degree as is available in dedicated software packages, as our primary focus has been on realizing a new version of FSD.
%
JFSD was created to provide greater accessibility to SD by leveraging \textbf{\texttt{Python}}’s simplicity and interoperability, while maintaining the competitive computational performance of dedicated implementations. Thanks to \textbf{\texttt{Google JAX}}'s JIT compilation and reliance on GPU acceleration, our software achieves speeds comparable to previous \textbf{\texttt{CUDA}}\textsuperscript{\textregistered} implementations~\cite{fiore2019fast}.
Presently, JFSD is capable of simulating approximately 30,000–50,000 particles within the memory constraints of mid-range GPUs. While this imposes a practical upper limit, ongoing developments in \textbf{\texttt{Google JAX}}, particularly the support for multi-core CPU and distributed computing, could further expand the range of the software in the future.
%%%%%%%%
\section{\label{sec:concl}Conclusion and Outlook}
%%%%%%%%
In summary, we have presented
a \textbf{\texttt{Python}}-based implementation of fast Stokesian Dynamics that relies on the \textbf{\texttt{Google JAX}}, which we refer to as JFSD. Our work addresses the use issues posed by the currently outdated \textbf{\texttt{CUDA}}\textsuperscript{\textregistered} code base~\cite{fiore2019fast} originally developed by Fiore and Swan. By utilizing Just-In-Time compilation~\cite{jax2018github}, JFSD achieves identical scaling, while maintaining performance and offering a user-friendly and modular interface. Our software is capable of simulations involving tens of thousands of particles on a modern desktop GPU.
% Stokesian Dynamics is a powerful method for efficiently resolving many-body hydrodynamic interactions and short-range lubrication forces in particle suspensions. JFSD brings this framework into a modern \textbf{\texttt{Python-JAX}} implementation, Building on the strengths of Stokesian Dynamics, the implementation delivers accurate and efficient simulations across a range of scenarios.
%
We built upon the extensive literature for SD and hydrodynamic interactions in suspensions to validate the code. Key benchmarks performed in this paper include: sedimentation, shear rheology, and thermal motion. We found strong agreement with the literature, demonstrating the quality of our implementation. The inclusion of both open and periodic boundary conditions, along with alternative hydrodynamic models, ensures flexibility and broad applicability for diverse physical systems. JFSD is thus ready to enable new research into the effect of hydrodynamic interactions in particle suspensions.
\section*{Acknowledgements}
We developed JFSD with the aim of giving back to the fluid-dynamics community, whose efforts and developments~\cite{durlofsky1987dynamic,fiore2019fast,fiore2017rapid,fiore2018rapid,ANDERSON2020109363,pelaez2022complex,townsend2024stokesian} helped support our own research~\cite{de2019hydrodynamics, torre2023hydrodynamic}. We hope that our implementation of FSD will allow researchers to continue to make use of the excellent algorithms developed by Fiore and Swan~\cite{fiore2019fast} for many years to come. Thank you Jim for the kindness that you showed us. We are also grateful to Andrew M. Fiore, Zhouyang Ge, and Gwynn Elfring; Luca Leone and Henri Menke; and Athanasios Machas and Dimos Aslanis for useful discussions, which helped us significantly improve our understanding of FSD, \textbf{\texttt{Google JAX}}, and the user/developer perspective, respectively.
\paragraph{Data availability statement}
An open data package containing the means to reproduce the results of the simulations is available at: DOI. The specific version of JFSD used to generate the results in this paper is tagged as \texttt{v0.2.0}, which will aid with reproducibility.
\paragraph{Author contributions}
Conceptualization, J.d.G. \& K.W.T.; Methodology, K.W.T. \& R.S.; Numerical calculations, K.W.T., Validation, K.W.T. \& R.S.; Investigation,
K.W.T.; Writing --- Original Draft, K.W.T.; Writing --- Review \& Editing, J.d.G.; Funding Acquisition, J.d.G.;
Resources, J.d.G.; Supervision, J.d.G.
\paragraph{Funding information}
The authors acknowledge the Dutch Research Council NWO for funding through OCENW.KLEIN.354, as well as the International Fine Particle Research Institute for funding through collaboration grant CRR-118-01.
\bibliography{SciPost_Example_BiBTeX_File.bib}
\end{document}
\begin{thebibliography}{10}
\providecommand{\url}[1]{\texttt{#1}}
\providecommand{\urlprefix}{URL }
\expandafter\ifx\csname urlstyle\endcsname\relax
\providecommand{\doi}[1]{doi:\discretionary{}{}{}#1}\else
\providecommand{\doi}{doi:\discretionary{}{}{}\begingroup \urlstyle{rm}\Url}\fi
\providecommand{\eprint}[2][]{\url{#2}}
\bibitem{fiore2019fast}
A.~M. Fiore and J.~W. Swan,
\newblock \emph{Fast stokesian dynamics},
\newblock Journal of Fluid Mechanics \textbf{878}, 544 (2019).
\bibitem{batchelor1977effect}
G.~K. Batchelor,
\newblock \emph{The effect of brownian motion on the bulk stress in a suspension of spherical particles},
\newblock Journal of fluid mechanics \textbf{83}(1), 97 (1977).
\bibitem{happel1983mechanics}
J.~Happel, H.~Brenner and L.~R.~N. Hydrodynamics,
\newblock \emph{Mechanics of fluids and transport processes},
\newblock Low reynolds number hydrodynamics. With special applications to particulate media. M. Nijhoff, Leiden (1983).
\bibitem{kim2013microhydrodynamics}
S.~Kim and S.~J. Karrila,
\newblock \emph{Microhydrodynamics: principles and selected applications},
\newblock Courier Corporation (2013).
\bibitem{brady1988stokesian}
J.~F. Brady, G.~Bossis \emph{et~al.},
\newblock \emph{Stokesian dynamics},
\newblock Annual review of fluid mechanics \textbf{20}(1), 111 (1988).
\bibitem{jeffrey1984calculation}
D.~Jeffrey and Y.~Onishi,
\newblock \emph{Calculation of the resistance and mobility functions for two unequal rigid spheres in low-reynolds-number flow},
\newblock Journal of Fluid Mechanics \textbf{139}, 261 (1984).
\bibitem{kim1985resistance}
S.~Kim and R.~T. Mifflin,
\newblock \emph{The resistance and mobility functions of two equal spheres in low-reynolds-number flow},
\newblock The Physics of fluids \textbf{28}(7), 2033 (1985).
\bibitem{jeffrey1992calculation}
D.~Jeffrey,
\newblock \emph{The calculation of the low reynolds number resistance functions for two unequal spheres},
\newblock Physics of Fluids A: Fluid Dynamics \textbf{4}(1), 16 (1992).
\bibitem{eckstein1977self}
E.~C. Eckstein, D.~G. Bailey and A.~H. Shapiro,
\newblock \emph{Self-diffusion of particles in shear flow of a suspension},
\newblock Journal of Fluid Mechanics \textbf{79}(1), 191 (1977).
\bibitem{furukawa2010key}
A.~Furukawa and H.~Tanaka,
\newblock \emph{Key role of hydrodynamic interactions in colloidal gelation},
\newblock Phys. Rev. Lett. \textbf{104}(24), 245702 (2010).
\bibitem{whitmer2011influence}
J.~K. Whitmer and E.~Luijten,
\newblock \emph{Influence of hydrodynamics on cluster formation in colloid- polymer mixtures},
\newblock The Journal of Physical Chemistry B \textbf{115}(22), 7294 (2011).
\bibitem{royall2015probing}
C.~P. Royall, J.~Eggers, A.~Furukawa and H.~Tanaka,
\newblock \emph{Probing colloidal gels at multiple length scales: The role of hydrodynamics},
\newblock Physical review letters \textbf{114}(25), 258302 (2015).
\bibitem{varga2018modelling}
Z.~Varga, J.~L. Hofmann and J.~W. Swan,
\newblock \emph{Modelling a hydrodynamic instability in freely settling colloidal gels},
\newblock Journal of Fluid Mechanics \textbf{856}, 1014 (2018).
\bibitem{swan-furst2019}
K.~A. Whitaker, Z.~Varga, L.~C. Hsiao, M.~J. Solomon, J.~W. Swan and E.~M. Furst,
\newblock \emph{Colloidal gel elasticity arises from the packing of locally glassy clusters},
\newblock Nat. Commun. \textbf{10}(1), 1 (2019).
\bibitem{turetta2022role}
L.~Turetta and M.~Lattuada,
\newblock \emph{The role of hydrodynamic interactions on the aggregation kinetics of sedimenting colloidal particles},
\newblock Soft Matter \textbf{18}(8), 1715 (2022).
\bibitem{torre2023structuring}
K.~W. Torre and J.~de~Graaf,
\newblock \emph{Structuring colloidal gels via micro-bubble oscillations},
\newblock Soft Matter \textbf{19}(15), 2771 (2023).
\bibitem{de2023hydrodynamic}
J.~De~Graaf, K.~W. Torre, W.~C. Poon and M.~Hermes,
\newblock \emph{Hydrodynamic stability criterion for colloidal gelation under gravity},
\newblock Physical Review E \textbf{107}(3), 034608 (2023).
\bibitem{torre2023hydrodynamic}
K.~W. Torre and J.~de~Graaf,
\newblock \emph{Hydrodynamic lubrication in colloidal gels},
\newblock Soft Matter \textbf{19}(38), 7388 (2023).
\bibitem{TorredeGraaf2024}
K.~Torre and J.~de~Graaf,
\newblock \emph{Delayed gravitational collapse of attractive colloidal suspensions},
\newblock Journal of Fluid Mechanics \textbf{1000}, A73 (2024),
\newblock \doi{10.1017/jfm.2024.1076}.
\bibitem{bolintineanu2014particle}
D.~S. Bolintineanu, G.~S. Grest, J.~B. Lechman, F.~Pierce, S.~J. Plimpton and P.~R. Schunk,
\newblock \emph{Particle dynamics modeling methods for colloid suspensions},
\newblock Computational Particle Mechanics \textbf{1}, 321 (2014).
\bibitem{chen1998lattice}
S.~Chen and G.~D. Doolen,
\newblock \emph{Lattice boltzmann method for fluid flows},
\newblock Annual review of fluid mechanics \textbf{30}(1), 329 (1998).
\bibitem{ladd2001lattice}
A.~J. Ladd and R.~Verberg,
\newblock \emph{Lattice-boltzmann simulations of particle-fluid suspensions},
\newblock Journal of statistical physics \textbf{104}, 1191 (2001).
\bibitem{dunweg2009lattice}
B.~D{\"u}nweg and A.~J. Ladd,
\newblock \emph{Lattice boltzmann simulations of soft matter systems},
\newblock Advanced computer simulation approaches for soft matter sciences III pp. 89--166 (2009).
\bibitem{aidun2010lattice}
C.~K. Aidun and J.~R. Clausen,
\newblock \emph{Lattice-boltzmann method for complex flows},
\newblock Annual review of fluid mechanics \textbf{42}(1), 439 (2010).
\bibitem{rohm2012lattice}
D.~R{\"o}hm and A.~Arnold,
\newblock \emph{Lattice boltzmann simulations on gpus with espresso},
\newblock The European Physical Journal Special Topics \textbf{210}(1), 89 (2012).
\bibitem{gompper2009multi}
G.~Gompper, T.~Ihle, D.~Kroll and R.~Winkler,
\newblock \emph{Multi-particle collision dynamics: A particle-based mesoscale simulation approach to the hydrodynamics of complex fluids},
\newblock Advanced computer simulation approaches for soft matter sciences III pp. 1--87 (2009).
\bibitem{howard2018efficient}
M.~P. Howard, A.~Z. Panagiotopoulos and A.~Nikoubashman,
\newblock \emph{Efficient mesoscale hydrodynamics: Multiparticle collision dynamics with massively parallel gpu acceleration},
\newblock Computer Physics Communications \textbf{230}, 10 (2018).
\bibitem{malevanets1999mesoscopic}
A.~Malevanets and R.~Kapral,
\newblock \emph{Mesoscopic model for solvent dynamics},
\newblock The Journal of chemical physics \textbf{110}(17), 8605 (1999).
\bibitem{tanaka2000simulation}
H.~Tanaka and T.~Araki,
\newblock \emph{Simulation method of colloidal suspensions with hydrodynamic interactions: Fluid particle dynamics},
\newblock Physical review letters \textbf{85}(6), 1338 (2000).
\bibitem{nguyen2002lubrication}
N.-Q. Nguyen and A.~Ladd,
\newblock \emph{Lubrication corrections for lattice-boltzmann simulations of particle suspensions},
\newblock Physical Review E \textbf{66}(4), 046708 (2002).
\bibitem{sierou2001accelerated}
A.~Sierou and J.~F. Brady,
\newblock \emph{Accelerated stokesian dynamics simulations},
\newblock Journal of fluid mechanics \textbf{448}, 115 (2001).
\bibitem{banchio2003accelerated}
A.~J. Banchio and J.~F. Brady,
\newblock \emph{Accelerated stokesian dynamics: Brownian motion},
\newblock The Journal of chemical physics \textbf{118}(22), 10323 (2003).
\bibitem{wang2016spectral}
M.~Wang and J.~F. Brady,
\newblock \emph{Spectral ewald acceleration of stokesian dynamics for polydisperse suspensions},
\newblock Journal of Computational Physics \textbf{306}, 443 (2016).
\bibitem{ANDERSON2020109363}
J.~A. Anderson, J.~Glaser and S.~C. Glotzer,
\newblock \emph{Hoomd-blue: A python package for high-performance molecular dynamics and hard particle monte carlo simulations},
\newblock Comput. Mater. Sci. \textbf{173}, 109363 (2020),
\newblock \doi{https://doi.org/10.1016/j.commatsci.2019.109363}.
\bibitem{jax2018github}
J.~Bradbury, R.~Frostig, P.~Hawkins, M.~J. Johnson, C.~Leary, D.~Maclaurin, G.~Necula, A.~Paszke, J.~Vander{P}las, S.~Wanderman-{M}ilne and Q.~Zhang,
\newblock \emph{{JAX}: composable transformations of {P}ython+{N}um{P}y programs} (2018).
\bibitem{fiore2017rapid}
A.~M. Fiore, F.~Balboa~Usabiaga, A.~Donev and J.~W. Swan,
\newblock \emph{Rapid sampling of stochastic displacements in brownian dynamics simulations},
\newblock The Journal of chemical physics \textbf{146}(12) (2017).
\bibitem{chow2014preconditioned}
E.~Chow and Y.~Saad,
\newblock \emph{Preconditioned krylov subspace methods for sampling multivariate gaussian distributions},
\newblock SIAM Journal on Scientific Computing \textbf{36}(2), A588 (2014).
\bibitem{schoenholz2021jax}
S.~S. Schoenholz and E.~D. Cubuk,
\newblock \emph{Jax, md a framework for differentiable physics},
\newblock Journal of Statistical Mechanics: Theory and Experiment \textbf{2021}(12), 124016 (2021).
\bibitem{jfsdgithub}
K.~W. Torre, J.~De~Graaf and R.~Schram,
\newblock \emph{Python-jax-fast-stokesian-dynamics},
\newblock Available at: \url{https://github.com/torrewk/Python-Jax-Fast-Stokesian-Dynamics}.
\bibitem{durlofsky1987dynamic}
L.~Durlofsky, J.~F. Brady and G.~Bossis,
\newblock \emph{Dynamic simulation of hydrodynamically interacting particles},
\newblock Journal of fluid mechanics \textbf{180}, 21 (1987).
\bibitem{brady1988dynamic}
J.~F. Brady, R.~J. Phillips, J.~C. Lester and G.~Bossis,
\newblock \emph{Dynamic simulation of hydrodynamically interacting suspensions},
\newblock Journal of Fluid Mechanics \textbf{195}, 257 (1988).
\bibitem{da1996shear}
F.~Da~Cunha and E.~Hinch,
\newblock \emph{Shear-induced dispersion in a dilute suspension of rough spheres},
\newblock Journal of fluid mechanics \textbf{309}, 211 (1996).
\bibitem{hasimoto1959periodic}
H.~Hasimoto,
\newblock \emph{On the periodic fundamental solutions of the stokes equations and their application to viscous flow past a cubic array of spheres},
\newblock Journal of Fluid Mechanics \textbf{5}(2), 317 (1959).
\bibitem{fiore2018rapid}
A.~M. Fiore and J.~W. Swan,
\newblock \emph{Rapid sampling of stochastic displacements in brownian dynamics simulations with stresslet constraints},
\newblock The Journal of chemical physics \textbf{148}(4) (2018).
\bibitem{ermak1978brownian}
D.~L. Ermak and J.~A. McCammon,
\newblock \emph{Brownian dynamics with hydrodynamic interactions},
\newblock The Journal of chemical physics \textbf{69}(4), 1352 (1978).
\bibitem{ladd1990hydrodynamic}
A.~J. Ladd,
\newblock \emph{Hydrodynamic transport coefficients of random dispersions of hard spheres},
\newblock The Journal of chemical physics \textbf{93}(5), 3484 (1990).
\bibitem{ge2025hydrodynamic}
Z.~Ge and G.~J. Elfring,
\newblock \emph{Hydrodynamic diffusion in apolar active suspensions of squirmers},
\newblock Journal of Fluid Mechanics \textbf{1003}, A17 (2025).
\bibitem{elfring2022active}
G.~J. Elfring and J.~F. Brady,
\newblock \emph{Active stokesian dynamics},
\newblock Journal of Fluid Mechanics \textbf{952}, A19 (2022).
\bibitem{shaik2023confined}
V.~A. Shaik, Z.~Peng, J.~F. Brady and G.~J. Elfring,
\newblock \emph{Confined active matter in external fields},
\newblock Soft Matter \textbf{19}(7), 1384 (2023).
\bibitem{townsend2023generating}
A.~K. Townsend,
\newblock \emph{Generating, from scratch, the near-field asymptotic forms of scalar resistance functions for two unequal rigid spheres in low reynolds number flow},
\newblock Physics of Fluids \textbf{35}(12) (2023).
\bibitem{rotne1969variational}
J.~Rotne and S.~Prager,
\newblock \emph{Variational treatment of hydrodynamic interaction in polymers},
\newblock The Journal of Chemical Physics \textbf{50}(11), 4831 (1969).
\bibitem{pelaez2022complex}
R.~P. Pel{\'a}ez,
\newblock \emph{Complex fluids in the GPU era: Algorithms and simulations},
\newblock Ph.D. thesis, Universidad Aut{\'o}noma de Madrid (2022).
\bibitem{townsend2024stokesian}
A.~K. Townsend,
\newblock \emph{Stokesian dynamics in python},
\newblock Journal of Open Source Software \textbf{9}(94), 6011 (2024).
\bibitem{de2019hydrodynamics}
J.~De~Graaf, W.~C. Poon, M.~J. Haughey and M.~Hermes,
\newblock \emph{Hydrodynamics strongly affect the dynamics of colloidal gelation but not gel structure},
\newblock Soft Matter \textbf{15}(1), 10 (2019).
\end{thebibliography}