Bootstrap: docker
From: nvidia/cuda:12.8.0-devel-ubuntu24.04

%environment
    export PATH="/opt/ctffind/bin:/opt/relion/bin:/opt/conda/bin:${PATH}"
    export TORCH_HOME="/opt/torch/"

%post
    export DEBIAN_FRONTEND=noninteractive 
    export PATH="/opt/relion/bin:/opt/conda/bin:${PATH}" 
    export TORCH_HOME="/opt/torch/"
    apt-get update && apt install -y curl wget cmake git build-essential mpi-default-bin mpi-default-dev gfortran libfftw3-dev libtiff-dev libpng-dev ghostscript libxft-dev libwxgtk3.2-dev

    # Install CTFFIND
    export CTFFIND_URL="https://grigoriefflab.umassmed.edu/sites/default/files/ctffind-4.1.14.tar.gz"
    wget ${CTFFIND_URL} && tar -xzf ctffind-4.1.14.tar.gz && rm ctffind-4.1.14.tar.gz && cd ctffind-4.1.14
    # Patch ctffind source code
    CTFFIND_FILE="src/programs/ctffind/ctffind.cpp"
    sed -i -e 's/\bbool ComputeRotationalAverageOfPowerSpectrum/void ComputeRotationalAverageOfPowerSpectrum/g' -e 's/\bbool RescaleSpectrumAndRotationalAverage/void RescaleSpectrumAndRotationalAverage/g' "$CTFFIND_FILE"
    MATRIX_FILE="src/core/matrix.cpp"
    sed -i -e 's|#define _AL_SINCOS(x, s, c)  __asm__ ("fsincos" : "=t" (c), "=u" (s) : "0" (x))|#define _AL_SINCOS(x, s, c)  do { (s) = sinf(x); (c) = cosf(x); } while(0)|' -e 's|#define FLOATSINCOS(x, s, c)  _AL_SINCOS((x) * AL_PI / 128.0, s ,c)|#define FLOATSINCOS(x, s, c)  _AL_SINCOS((x) * AL_PI / 128.0f, s, c)|' "$MATRIX_FILE"
    # Build ctffind
    mkdir build && cd build
    ../configure --prefix="/opt/ctffind" --enable-static --disable-debug --enable-optimisations --enable-openmp FC=gfortran F77=gfortran CXXFLAGS="-g -O2 -Wno-maybe-uninitialized -Wno-error=maybe-uninitialized" CFLAGS="-g -O2 -Wno-maybe-uninitialized -Wno-error=maybe-uninitialized"
    make -j 72 && make install
    cd /

    # Install Miniforge Conda
    curl --location --remote-name "https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-$(uname)-$(uname -m).sh"
    bash Miniforge3-$(uname)-$(uname -m).sh -b -p /opt/conda && rm Miniforge3-$(uname)-$(uname -m).sh
    
    # Install Relion
    git clone -b ver5.0 https://github.com/3dem/relion.git && cd relion
    # patch conda environment file
    sed -i 's/^\(      \)- napari\[all\]/\1# - napari[all]/' environment_blackwell.yml
    sed -i 's/^\(      \)- tsnecuda/\1# - tsnecuda/' environment_blackwell.yml
    sed -i 's/^\(      \)- git+/\1# - git+/' environment_blackwell.yml
    sed -i 's/^\(      \)- PyQt5/\1# - PyQt5/' environment_blackwell.yml
    sed -i 's/torchvision==0.22.1+cu128/torchvision/' environment_blackwell.yml
    sed -i '/^  - setuptools=/a\  - pyqt=5.15.9' environment_blackwell.yml
    # install pip dependencies
    conda env create -f environment_blackwell.yml
    . /opt/conda/bin/activate && conda activate relion-5.0
    pip install git+https://github.com/3dem/relion-classranker git+https://github.com/3dem/relion-blush git+https://github.com/3dem/topaz
    conda deactivate
    # Build relion
    mkdir $TORCH_HOME && mkdir build && cd build
    cmake .. -DCMAKE_INSTALL_PREFIX=/opt/relion -DCUDA_ARCH=90 -DGUI=OFF -DFETCH_WEIGHTS=OFF && make -j 72