Bootstrap: docker
From: ghcr.io/nvidia/jax:jax-2024-07-24
# the above version is chose for compatibility with fastemriwaveforms-cuda12x and the local CUDA driver

%post
    python -m pip install --upgrade pip
    python -m pip install --no-cache-dir nvidia-cuda-runtime-cu12 astropy eryn fastemriwaveforms-cuda12x multiprocess optax matplotlib scipy jupyter interpax numba Cython lisaanalysistools tabulate
    python -c "import few; few.get_backend('cuda12x'); print('FEW installation successful')"

    # Set compilers explicitly and unset conda variables
    unset CC CXX CUDACXX
    export CC=/usr/bin/gcc
    export CXX=/usr/bin/g++
    export CUDACXX=/usr/local/cuda/bin/nvcc
    export NVCC_PREPEND_FLAGS='-ccbin /usr/bin/g++'
    
    # install lisa on gpu and StableEMRIFisher-package
    git clone https://github.com/cchapmanbird/EMRI-FoM.git emri_fom_temp
    cd emri_fom_temp/lisa-on-gpu/
    python setup.py install
    cd ../StableEMRIFisher-package/
    python -m pip install .
    cd ../
    # python -m unittest test_waveform_and_response.py
    # echo "LISA on GPU and StableEMRIFisher-package installation successful"