setup.sh

#!/bin/bash

set -e

# Inspired by https://github.com/YoshitakaMo/localcolabfold, with the 
# hopefully added bonus of producing a functioning installation on 
# our infrastructure...
#
# Hopefully provides tighter controls of installed versions and dependancies as well
VERSION="1.5.5"

# pick a snake, any snake...
#[[ "$(which conda 2>/dev/null)" ]] && CONDA='conda'
#[[ "$(which mamba 2>/dev/null)" ]] && CONDA='mamba'
CONDA='conda'

inDockerBuild=`uname -a|(grep -c buildkit || true)`

if [[ -z "$CONDA_PREFIX" ]]; then
	echo "Please activate your base environment prior to running this script..."
	exit 1
fi

if [[ -z "${CONDA}" ]]; then
	echo "Please install Miniconda3 or mamba prior to running this script..."
	exit 1
else
	source $CONDA_PREFIX/etc/profile.d/conda.sh

	# If installation is run on a non-gpu host, some cpu-centric packages will be installed instead
	# of cuda packages. This can be fixed with the 'CONDA_OVERRIDE_CUDA' variabe...
	CONDA_OVERRIDE_CUDA="12.6" $CONDA env create -f colabfold_batch.yaml

	source $(dirname $CONDA_EXE)/activate colabfold_batch_${VERSION}

	echo 'CUDNN_PATH=$(dirname $(python -c "import nvidia.cudnn;print(nvidia.cudnn.__file__)"))' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
	echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$CONDA_PREFIX/lib/:$CUDNN_PATH/lib' >> $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh

	py_ver=$(python --version|cut -f2 -d' '|sed 's/\.[0-9]*$//')
	cd $CONDA_PREFIX/lib/python${py_ver}/site-packages/colabfold

	# Use matplotlib Agg backend for headless operation
	sed -i -e "s#from matplotlib import pyplot as plt#import matplotlib\nmatplotlib.use('Agg')\nimport matplotlib.pyplot as plt#g" plot.py
	# Store alphafold weightings in $CONDA_PREFIX/share/colabfold rather than default user directory
	sed -i -e "s#appdirs.user_cache_dir(__package__ or \"colabfold\")#\"${CONDA_PREFIX}/share/colabfold\"#g" download.py
	# suppress warnings related to tensorflow
	sed -i -e "s#from io import StringIO#from io import StringIO\nfrom silence_tensorflow import silence_tensorflow\nsilence_tensorflow()#g" batch.py
	rm -rf __pycache__
	
	# Download alphafold weights 
	cd ${CONDA_PREFIX}/share
	mkdir -m 0777 -p colabfold
	python -m colabfold.download

	if [[ "$inDockerBuild" == '1' ]]; then
		mamba clean -a -y
		pip cache purge
	fi
fi