ftdmp-relax-with-openmm

#!/bin/bash

function print_help_and_exit
{
cat >&2 << 'EOF'

'ftdmp-relax-with-openmm' script relaxes a molecular structure using OpenMM.

Options:
    --input                   string  *  input file path
    --output                  string  *  output file path, setting to '_same_as_input' will overwrite input file
    --focus                   string     focus mode, default is 'whole_structure', others are: 'interface_side_chains', 'whole_interface', 'not_interface'
    --focus-base              string     selection of atoms that can be in a partial structure focus, default is '[]'
    --focus-first-sel         string     first selection of atoms to define interface not by chains, default is ''
    --focus-second-sel        string     second selection of atoms to define interface not by chains, default is ''
    --conda-path              string     conda installation path, default is ''
    --conda-env               string     conda environment name, default is ''
    --forcefield              string     forcefield combo name, default is 'amber99sb', others are: 'amber14-all', 'amber14-all-no-water', 'charmm36'
    --main-forcefield         string     main forcefield name, default is defined by the combo name, some others are: 'amber99sb', 'amber14-all', 'charmm36'
    --water-forcefield        string     water forcefiled name, default is defined by the combo name, some others are: '', 'amber99_obc', 'amber14/tip3pfb', 'charmm36/water'
    --max-iterations          number     max number of iterations, default is 100
    --score-at-end            string     mode for scoring interface at the end, default is '', others are: 'fast_iface', 'full_iface', 'full'
    --scoring-params          string     additional parameters for scoring, default is ''
    --multiple-tries          number     number of tries to generate and score interfaces, default is ''
    --cache-dir               string     cache directory path to store results of past calls
    --force-cuda                         flag to force the platform to be CUDA
    --trim-output                        flag to restrict output to atoms of proteins and nucleic acids
    --no-preparation                     flag to not run any preparation of input structure before simulations
    --limit-preparation                  flag to only add solvent if needed in the preparation stage
    --full-preparation                   flag to turn off all preparation disabling flags
    --no-simulation                      flag to not run any simulations
    --help | -h                          flag to display help message and exit

Standard output:
    space-separated table of scores for both input and output
    
Examples:

    ftdmp-relax-with-openmm --input model.pdb --output relaxed_model.pdb
    
    ftdmp-relax-with-openmm --conda-path ~/anaconda3 --conda-env alphafold2 \
      --forcefield amber14-all -i model.pdb -o relaxed_model.pdb --score-at-end fast_iface --trim-output

EOF
exit 1
}

readonly ZEROARG=$0
ALLARGS=("$@")

if [ -z "$1" ]
then
	print_help_and_exit
fi

if [ -z "$FTDMPDIR" ]
then
	export FTDMPDIR="$( cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
	export PATH="${FTDMPDIR}/core/voronota/expansion_js:${FTDMPDIR}:${PATH}"
fi

INFILE=""
OUTFILE=""
FOCUS_MODE="whole_structure"
FOCUS_BASE="[]"
FOCUS_FIRST_SEL=""
FOCUS_SECOND_SEL=""
CONDA_PATH=""
CONDA_ENV=""
FORCEFIELD="amber99sb"
MAIN_FORCEFIELD="default"
WATER_FORCEFIELD="default"
MAX_ITERATIONS="100"
SCORE_AT_END=""
SCORING_PARAMS=""
MULTIPLE_TRIES=""
CACHE_DIR=""
FORCE_CUDA="false"
TRIM_OUTPUT="false"
NO_PREPARATION="false"
LIMIT_PREPARATION="false"
NO_SIMULATION="false"
JUST_OUTPUT_FROM_DIR=""
ARCHIVE_WHEN_JUST_OUTPUT_FROM_DIR=""
HELP_MODE="false"

while [[ $# > 0 ]]
do
	OPTION="$1"
	OPTARG="$2"
	shift
	case $OPTION in
	--input)
		INFILE="$OPTARG"
		shift
		;;
	--output)
		OUTFILE="$OPTARG"
		shift
		;;
	--focus)
		FOCUS_MODE="$OPTARG"
		shift
		;;
	--focus-base)
		FOCUS_BASE="$OPTARG"
		shift
		;;
	--focus-first-sel)
		FOCUS_FIRST_SEL="$OPTARG"
		shift
		;;
	--focus-second-sel)
		FOCUS_SECOND_SEL="$OPTARG"
		shift
		;;
	--conda-path)
		CONDA_PATH="$OPTARG"
		shift
		;;
	--conda-env)
		CONDA_ENV="$OPTARG"
		shift
		;;
	--forcefield)
		FORCEFIELD="$OPTARG"
		shift
		;;
	--main-forcefield)
		MAIN_FORCEFIELD="$OPTARG"
		shift
		;;
	--water-forcefield)
		WATER_FORCEFIELD="$OPTARG"
		shift
		;;
	--max-iterations)
		MAX_ITERATIONS="$OPTARG"
		shift
		;;
	--score-at-end)
		SCORE_AT_END="$OPTARG"
		shift
		;;
	--scoring-params)
		SCORING_PARAMS="$OPTARG"
		shift
		;;
	--multiple-tries)
		MULTIPLE_TRIES="$OPTARG"
		shift
		;;
	--cache-dir)
		CACHE_DIR="$OPTARG"
		shift
		;;
	--force-cuda)
		FORCE_CUDA="true"
		;;
	--trim-output)
		TRIM_OUTPUT="true"
		;;
	--no-preparation)
		NO_PREPARATION="true"
		LIMIT_PREPARATION="false"
		;;
	--limit-preparation)
		NO_PREPARATION="false"
		LIMIT_PREPARATION="true"
		;;
	--full-preparation)
		NO_PREPARATION="false"
		LIMIT_PREPARATION="false"
		;;
	--no-simulation)
		NO_SIMULATION="true"
		;;
	--jofd)
		JUST_OUTPUT_FROM_DIR="$OPTARG"
		shift
		;;
	--awjofd)
		ARCHIVE_WHEN_JUST_OUTPUT_FROM_DIR="$OPTARG"
		shift
		;;
	-h|--help)
		HELP_MODE="true"
		;;
	*)
		echo >&2 "Error: invalid command line option '$OPTION'"
		exit 1
		;;
	esac
done

if [ "$HELP_MODE" == "true" ]
then
	print_help_and_exit
fi

if [ -z "$INFILE" ]
then
	echo >&2 "Error: input file path not provided"
	exit 1
fi

if [ -z "$OUTFILE" ]
then
	echo >&2 "Error: output file path not provided"
	exit 1
fi

if [ ! -s "$INFILE" ]
then
	echo >&2 "Error: input file '$INFILE' does not exist"
	exit 1
fi

if [ -n "$JUST_OUTPUT_FROM_DIR" ]
then
	if [ -s "$JUST_OUTPUT_FROM_DIR/scoring_results.txt" ]
	then
		cat "$JUST_OUTPUT_FROM_DIR/scoring_results.txt"
	fi
	
	if [ -s "$JUST_OUTPUT_FROM_DIR/openmm_errors.txt" ]
	then
		cat "$JUST_OUTPUT_FROM_DIR/openmm_errors.txt" >&2
	fi
	
	if [ -s "$JUST_OUTPUT_FROM_DIR/relaxed.pdb" ]
	then
		if [ -n "$OUTFILE" ] && [ "$OUTFILE" != "/dev/null" ]
		then
			if [ "$OUTFILE" == "_same_as_input" ]
			then
				cat "$JUST_OUTPUT_FROM_DIR/relaxed.pdb" > "$INFILE"
			else
				mkdir -p "$(dirname "$OUTFILE")"
				cat "$JUST_OUTPUT_FROM_DIR/relaxed.pdb" > "$OUTFILE"
			fi
		fi
	else
		echo >&2 "Error: no relaxed structure produced"
	fi
	
	if [ -n "$ARCHIVE_WHEN_JUST_OUTPUT_FROM_DIR" ] && [ ! -s "$ARCHIVE_WHEN_JUST_OUTPUT_FROM_DIR" ]
	then
		mkdir -p "$(dirname "$ARCHIVE_WHEN_JUST_OUTPUT_FROM_DIR")"
		tar -czf "$ARCHIVE_WHEN_JUST_OUTPUT_FROM_DIR" \
		  -C "$JUST_OUTPUT_FROM_DIR" \
		  $(find "$JUST_OUTPUT_FROM_DIR" -type f | xargs -L 1 basename | egrep '^openmm_errors.txt$|^relaxed.pdb$|^scoring_results.txt$')
	fi
	
	if [ -s "$JUST_OUTPUT_FROM_DIR/openmm_errors.txt" ] || [ ! -s "$JUST_OUTPUT_FROM_DIR/relaxed.pdb" ]
	then
		exit 1
	fi
	
	exit 0
fi

if [ "$MAIN_FORCEFIELD" == "default" ] || [ "$WATER_FORCEFIELD" == "default" ]
then
	if [ "$FORCEFIELD" != "amber99sb" ] && [ "$FORCEFIELD" != "amber14-all" ] && [ "$FORCEFIELD" != "amber14-all-no-water" ] && [ "$FORCEFIELD" != "charmm36" ]
	then
		echo >&2 "Error: invalid forcefield combo mode '$FORCEFIELD'"
		exit 1
	fi
	
	if [ "$MAIN_FORCEFIELD" == "default" ]
	then
		if [ "$FORCEFIELD" == "amber99sb" ]
		then
			MAIN_FORCEFIELD="amber99sb"
		fi
		
		if [ "$FORCEFIELD" == "amber14-all" ]
		then
			MAIN_FORCEFIELD="amber14-all"
		fi
		
		if [ "$FORCEFIELD" == "amber14-all-no-water" ]
		then
			MAIN_FORCEFIELD="amber14-all"
		fi
		
		if [ "$FORCEFIELD" == "charmm36" ]
		then
			MAIN_FORCEFIELD="charmm36"
		fi
	fi
	
	if [ "$WATER_FORCEFIELD" == "default" ]
	then
		if [ "$FORCEFIELD" == "amber99sb" ]
		then
			WATER_FORCEFIELD="amber99_obc"
		fi
		
		if [ "$FORCEFIELD" == "amber14-all" ]
		then
			WATER_FORCEFIELD="amber14/tip3pfb"
		fi
		
		if [ "$FORCEFIELD" == "amber14-all-no-water" ]
		then
			WATER_FORCEFIELD=""
		fi
		
		if [ "$FORCEFIELD" == "charmm36" ]
		then
			WATER_FORCEFIELD="charmm36/water"
		fi
	fi
fi

if [ "$FOCUS_MODE" != "interface_side_chains" ] && [ "$FOCUS_MODE" != "whole_interface" ] && [ "$FOCUS_MODE" != "not_interface" ] && [ "$FOCUS_MODE" != "whole_structure" ]
then
	echo >&2 "Error: invalid focus mode '$FOCUS_MODE'"
	exit 1
fi

if [ "$FOCUS_MODE" != "whole_structure" ] && [ -n "$WATER_FORCEFIELD" ] && [[ $WATER_FORCEFIELD != *"_obc"* ]]
then
	echo >&2 "Error: focus mode '$FOCUS_MODE' is incompatible with explicit solvent simulations"
	exit 1
fi

if [ "$FOCUS_MODE" != "whole_structure" ] || [ "$TRIM_OUTPUT" == "true" ]
then
	command -v voronota-js &> /dev/null || { echo >&2 "Error: 'voronota-js' executable not in binaries path"; exit 1; }
fi

if [ "$SCORE_AT_END" != "" ] && [ "$SCORE_AT_END" != "fast_iface" ] && [ "$SCORE_AT_END" != "full_iface" ] && [ "$SCORE_AT_END" != "full" ]
then
	echo >&2 "Error: invalid scoring mode '$SCORE_AT_END'"
	exit 1
fi

if [ "$SCORE_AT_END" == "fast_iface" ]
then
	command -v voronota-js-fast-iface-voromqa &> /dev/null || { echo >&2 "Error: 'voronota-js-fast-iface-voromqa' executable not in binaries path"; exit 1; }
fi

if [ "$SCORE_AT_END" == "full_iface" ] || [ "$SCORE_AT_END" == "full" ]
then
	command -v voronota-js-voromqa &> /dev/null || { echo >&2 "Error: 'voronota-js-voromqa' executable not in binaries path"; exit 1; }
fi

if [ -n "$MULTIPLE_TRIES" ] && [ "$SCORE_AT_END" != "fast_iface" ]
then
	echo >&2 "Error: multiple tries are allowed only with 'fast_iface' scoring mode"
	exit 1
fi

if [ -z "$CONDA_DEFAULT_ENV" ]
then
	if [ -z "$CONDA_PATH" ]
	then
		echo >&2 "Error: not in conda environment, and the conda path is not provided"
		exit 1
	fi
	if [ ! -s "${CONDA_PATH}/bin/activate" ]
	then
		echo >&2 "Error: no conda activation script '${CONDA_PATH}/bin/activate'"
		exit 1
	fi
	source "${CONDA_PATH}/bin/activate"
fi

if [ -n "$CONDA_ENV" ]
then
	if [ "$CONDA_DEFAULT_ENV" != "$CONDA_ENV" ]
	then
		conda activate "$CONDA_ENV"
	fi
	
	if [ "$CONDA_DEFAULT_ENV" != "$CONDA_ENV" ]
	then
		echo >&2 "Error: no '$CONDA_ENV' environment"
		exit 1
	fi
fi

readonly TMPLDIR=$(mktemp -d)
trap "rm -r $TMPLDIR" EXIT

CACHE_ARCHIVE=""

if [ -n "$CACHE_DIR" ]
then
	{
		cat "$INFILE"
		echo "$FOCUS_MODE"
		echo "$FOCUS_BASE"
		echo "$FOCUS_FIRST_SEL"
		echo "$FOCUS_SECOND_SEL"
		echo "$MAIN_FORCEFIELD"
		echo "$WATER_FORCEFIELD"
		echo "$MAX_ITERATIONS"
		echo "$SCORE_AT_END"
		echo "$SCORING_PARAMS"
		echo "$MULTIPLE_TRIES"
		echo "$TRIM_OUTPUT"
		echo "$NO_PREPARATION"
		echo "$LIMIT_PREPARATION"
		echo "$NO_SIMULATION"
	} \
	| md5sum | awk '{print $1}' \
	> "${TMPLDIR}/hashsum.txt"
	
	HASHSUM="ftdmp-relax-with-openmm-$(cat ${TMPLDIR}/hashsum.txt)"
	CACHE_ARCHIVE="${CACHE_DIR}/${HASHSUM}.tar.gz"
	
	if [ -s "$CACHE_ARCHIVE" ]
	then
		mkdir -p "${TMPLDIR}/saved_results"
		
		cp "$CACHE_ARCHIVE" "${TMPLDIR}/saved_results/archive.tar.gz"
		
		cd "${TMPLDIR}/saved_results"
		tar -xf "./archive.tar.gz"
		cd - &> /dev/null
		
		"$ZEROARG" "${ALLARGS[@]}" --jofd "$TMPLDIR/saved_results"
		
		exit 0
	fi
fi

cat "$INFILE" > "$TMPLDIR/input.pdb"

cd "$TMPLDIR"

true > openmm_errors.txt

if [ "$NO_PREPARATION" == "true" ]
then
	cp "input.pdb" "prepared.pdb"
else
	{
cat << 'EOF'
from openmm.app import *
from openmm import *
from openmm.unit import *
from pdbfixer import PDBFixer
import random
fixer = PDBFixer(filename='input.pdb')
EOF

	if [ "$LIMIT_PREPARATION" != "true" ]
	then
cat << 'EOF'
fixer.findMissingResidues()
fixer.findNonstandardResidues()
fixer.replaceNonstandardResidues()
fixer.removeHeterogens(keepWater=False)
fixer.findMissingAtoms()
fixer.addMissingAtoms(seed=999)
EOF
	fi

	if [ -n "$WATER_FORCEFIELD" ]
	then
cat << 'EOF'
forcefield = ForceField('MAIN_FORCEFIELD.xml', 'WATER_FORCEFIELD.xml')
EOF
	else
cat << 'EOF'
forcefield = ForceField('MAIN_FORCEFIELD.xml')
EOF
	fi

cat << 'EOF'
modeller = Modeller(fixer.topology, fixer.positions)
random.seed(999);
EOF

	if [ "$LIMIT_PREPARATION" != "true" ]
	then
cat << 'EOF'
modeller.addHydrogens(forcefield)
EOF
	else
cat << 'EOF'
modeller.deleteWater()
EOF
	fi

	if [ -n "$WATER_FORCEFIELD" ] && [[ $WATER_FORCEFIELD != *"_obc"* ]]
	then
cat << 'EOF'
modeller.addSolvent(forcefield, padding=0.2*nanometers)
EOF
	fi

cat << 'EOF'
PDBFile.writeFile(modeller.topology, modeller.positions, open('prepared.pdb', 'w'), keepIds=True)
EOF
	} \
	| sed "s|MAIN_FORCEFIELD|${MAIN_FORCEFIELD}|" \
	| sed "s|WATER_FORCEFIELD|${WATER_FORCEFIELD}|" \
	> prepare.py
	
	python prepare.py &>> openmm_errors.txt
fi

if [ ! -s "prepared.pdb" ]
then
	echo "Error: failed to prepare input with OpenMM" &>> openmm_errors.txt
	cd - &> /dev/null
	"$ZEROARG" "${ALLARGS[@]}" --jofd "$TMPLDIR" --awjofd "$CACHE_ARCHIVE"
	exit 1
fi

if [ "$NO_SIMULATION" == "true" ]
then
	cp "prepared.pdb" "relaxed.pdb"
else
	if [ "$FOCUS_MODE" != "whole_structure" ]
	then
		{
cat << EOF
var params={}
params.focus_base='$FOCUS_BASE';
params.focus_first_sel='$FOCUS_FIRST_SEL';
params.focus_second_sel='$FOCUS_SECOND_SEL';
EOF

cat << 'EOF'
voronota_import("-file", "prepared.pdb", "-include-hydrogens", "-include-heteroatoms");
voronota_assert_partial_success("Failed to import file");

voronota_select_atoms_close_to_interchain_interface("-name", "actii", "-use", params.focus_base, "-atoms-first", params.focus_first_sel, "-atoms-second", params.focus_second_sel);
voronota_assert_full_success("Failed to select interface atoms");

voronota_set_tag_of_atoms("-use", "[]", "-tag", "whole_structure");
voronota_assert_full_success("Failed to tag atoms");

voronota_set_tag_of_atoms("-use", "[actii]", "-tag", "whole_interface");
voronota_assert_full_success("Failed to tag atoms");

voronota_set_tag_of_atoms("-use", "(not [actii])", "-tag", "not_interface");
voronota_assert_full_success("Failed to tag atoms");

voronota_set_tag_of_atoms("-use", "([actii] and (not [-aname C,N,O,CA]))", "-tag", "interface_side_chains");
voronota_assert_full_success("Failed to tag atoms");

voronota_export_atoms("-file", "annotated.pa");
voronota_assert_full_success("Failed to export annotated atoms");
EOF
		} \
		> annotate.vs
	
		voronota-js --no-setup-defaults < annotate.vs
		
		if [ ! -f "annotated.pa" ]
		then
			echo >&2 "Error: failed to annotate input"
			exit 1
		fi
		
		cat annotated.pa | awk '{print (NR-1) " " $0}' | grep -v "$FOCUS_MODE" | awk '{print "system.setParticleMass(" $1 ", 0.0)"}' > freezing.py
	fi
	
	{
cat << 'EOF'
from openmm.app import *
from openmm import *
from openmm.unit import *
pdb = PDBFile('prepared.pdb')
EOF

	if [ -n "$WATER_FORCEFIELD" ]
	then
cat << 'EOF'
forcefield = ForceField('MAIN_FORCEFIELD.xml', 'WATER_FORCEFIELD.xml')
EOF
	else
cat << 'EOF'
forcefield = ForceField('MAIN_FORCEFIELD.xml')
EOF
	fi

cat << 'EOF'
system = forcefield.createSystem(pdb.topology)
EOF

	if [ -s "freezing.py" ]
	then
		cat freezing.py
	fi

cat << 'EOF'
integrator = VerletIntegrator(0.001*picoseconds)
EOF

	if [ "$FORCE_CUDA" == "true" ]
	then
cat << 'EOF'
platform = Platform.getPlatformByName('CUDA')
simulation = Simulation(pdb.topology, system, integrator, platform)
EOF
	else
cat << 'EOF'
simulation = Simulation(pdb.topology, system, integrator)
EOF
	fi

cat << 'EOF'
simulation.context.setPositions(pdb.positions)
simulation.minimizeEnergy(maxIterations=MAX_ITERATIONS)
positions = simulation.context.getState(getPositions=True).getPositions()
PDBFile.writeFile(simulation.topology, positions, open('relaxed.pdb', 'w'), keepIds=True)
EOF
	} \
	| sed "s|MAIN_FORCEFIELD|${MAIN_FORCEFIELD}|" \
	| sed "s|WATER_FORCEFIELD|${WATER_FORCEFIELD}|" \
	| sed "s|MAX_ITERATIONS|${MAX_ITERATIONS}|" \
	> relax.py
	
	if [ -n "$MULTIPLE_TRIES" ]
	then
		cp "prepared.pdb" "original_prepared.pdb"
		seq 1 "$MULTIPLE_TRIES" \
		| while read -r TRYNUM
		do
			python relax.py &>> openmm_errors.txt
			if [ ! -s "relaxed.pdb" ]
			then
				exit 1
			fi
			cp "relaxed.pdb" "prepared.pdb"
			mv "relaxed.pdb" "relaxed${TRYNUM}.pdb"
		done
		mv "original_prepared.pdb" "prepared.pdb"
	else
		python relax.py &>> openmm_errors.txt
	fi
fi

if [ ! -s "relaxed.pdb" ] && [ ! -s "relaxed1.pdb" ]
then
	echo "Error: failed to relax with OpenMM" &>> openmm_errors.txt
	cd - &> /dev/null
	"$ZEROARG" "${ALLARGS[@]}" --jofd "$TMPLDIR" --awjofd "$CACHE_ARCHIVE"
	exit 1
fi

{
if [ "$NO_PREPARATION" == "true" ]
then
	ls input.pdb relaxed*.pdb
else
	ls input.pdb prepared.pdb relaxed*.pdb
fi
} > list_to_score.txt

{
	if [ "$SCORE_AT_END" == "fast_iface" ]
	then
		if [ -n "$MULTIPLE_TRIES" ]
		then
			cat list_to_score.txt | voronota-js-fast-iface-voromqa -i _list $SCORING_PARAMS | tee scores.txt | column -t
			cp "$(cat scores.txt | awk '{print $1}' | egrep '^relaxed' | head -1)" "relaxed.pdb"
		else
			cat list_to_score.txt | voronota-js-fast-iface-voromqa -i _list $SCORING_PARAMS | column -t
		fi
	fi
	
	if [ "$SCORE_AT_END" == "full_iface" ]
	then
		cat list_to_score.txt | voronota-js-voromqa -i _list --inter-chain --processors 2 $SCORING_PARAMS | column -t
	fi
	
	if [ "$SCORE_AT_END" == "full" ]
	then
		cat list_to_score.txt | voronota-js-voromqa -i _list --processors 2 $SCORING_PARAMS | column -t
	fi
} > scoring_results.txt

if [ "$TRIM_OUTPUT" == "true" ]
then
	{
cat << 'EOF'
voronota_import("-file", "relaxed.pdb", "-include-hydrogens");
voronota_assert_partial_success("Failed to import file");

voronota_export_atoms("-file", "trimmed_relaxed.pdb", "-as-pdb");
voronota_assert_full_success("Failed to export atoms");
EOF
	} \
	| voronota-js --no-setup-defaults
	
	if [ ! -s "trimmed_relaxed.pdb" ]
	then
		echo >&2 "Error: failed to trim output"
		exit 1
	fi

	mv "trimmed_relaxed.pdb" "relaxed.pdb"
fi

cd - &> /dev/null

"$ZEROARG" "${ALLARGS[@]}" --jofd "$TMPLDIR" --awjofd "$CACHE_ARCHIVE"

exit 0