jfsd 0.1.0

Python implementation of Fast Stokesian Dynamics methods.

jfsd: a Python implementation of Fast Stokesian Dynamics with Jax

This repository contains a Python implementation of the Fast Stokesian Dynamics methods (A. Fiore, J. Swan, 2019, https://www.cambridge.org/core/journals/journal-of-fluid-mechanics/article/abs/fast-stokesian-dynamics/970BD1B80B43E21CD355C7BAD4644D46) which leverages the library JAX to just-in-time compile and run on GPU the simulation. The accuracy of hydrodynamic interactions can be scaled down in order to obtain Rotne-Prager-Yamakawa (up to 2-body effects, and no lubrication) or Brownian Dynamics (only 1-body effects). Boundary conditions can be switch from periodic to open, and the average height of asperities on particles can be specified, enabling simulations of rough colloids.

Installation Guide

Pre-requisites:

• cuda (tested with version 11.8, https://developer.nvidia.com/cuda-11-8-0-download-archive)
• cuDNN (tested with 8.6 for cuda 11, https://developer.nvidia.com/rdp/cudnn-archive)
• Python >= 3.9

1 - Set up work (virtual) environment:

git clone https://github.com/torrewk/Python-Jax-Fast-Stokesian-Dynamics.git

2 - Go into the directory of the project and type the command:

python3 -m venv .venv

3 - activate the environment:

source .venv/bin/activate

4 - Install correct version of jaxlib

pip install jaxlib==0.4.17+cuda11.cudnn86 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

5 - Install jfsd and rest of dependencies

pip install ".[test]"

6 - Reboot the environment (needed fore pytest to work):

deactivate && source .venv/bin/activate

7 - Run the JFSD unit tests

pytest tests/test_class.py

How to Run Simulations

1 - Create a configuration .toml in files/ (use "files/example_configuration.toml" as reference)

2a - Run the main code, from the project directory, via:

jfsd -c files/<your_config_file.toml> -o <directory_name_for_output>

Replace <your_config_file.toml> with the name of the file created in step (1), and <directory_name_for_output> with the name of the folder where you want to store the trajetory data.

2b - To provide initial particle positions to jfsd, run it via:

jfsd -c files/<your_config_file.toml> -s <initial_particle_positions.npy> -o <directory_name_for_output>

Replace <initial_particle_positions.npy> with the name of your initial configuration, and modify <your_config_file.toml> to accept initial positions from file (instead of randomly creating it, see 'example_configuration.toml'). Note that this file must be a numpy array of shape (N_p, 3), with N_p the number of particles.

Analyzing Trajectory Data

During the simulation, the particles trajectories,velocities, and stresslets are saved in a numpy array of shape (N_s, N_p, N_c), with N_s the number of frames stored, N_p the number of particles and N_c the number of d.o.f. (3 for trajectories, 6 for velocities, 5 for stresslets). Lengths in the simulation are expressed in units of the particles radius a, and momenta in units of γa, with unity mass for each particle.

Common/Known Issues

"Unable to load cuPTI"

This issue can occur if the PATH or LD_LIBRARY_PATH environment variables are not set correctly to include your CUDA installation, especially the cuPTI (CUDA Profiling Tools Interface) library, which JAX requires.

To solve this (temporarily), add the following line to your environment:

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/extras/CUPTI/lib64

For a permanent solution, edit the .bashrc file in your home directory:

echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/extras/CUPTI/lib64' >> ~/.bashrc

and apply the changes:

source ~/.bashrc

Note that, if cuda is installed in a different location than '/usr/local/', you need to modify the lines above accordingly.

License

This project is licensed under the Apache License, Version 2.0. See the LICENSE file for the full terms and conditions. For a concise summary, refer to LICENSE_SHORT.