# Script to estimate forces acting on harmonically restrained atoms in simulation. Supposing that one frame has coordinates of restraint targets, script compares these two frames of trajectory, calculates force acting on restrained atoms (according spring constants in the BETA field), puts the value into OCCUPANCY field and draws the color-coded force arrows.
# Andriy Anishkin (anishkin@icqmail.com) UMCP

#Selection to estimate
#set selection [atomselect top "beta > 0"]
#set selection [atomselect top backbone]
set selection [atomselect top "name CA"]

#Reference (initial) frame
set reference_frame 0

#Estimated frame
set estimated_frame [molinfo top get frame]
#set estimated_frame [expr {[molinfo top get numframes]-1}]

#Clear all the previously added graphics in VMD window
draw delete all

#Scaling coefficient for the arrow length
set scale_arrow 20

#Scaling coefficient for the arrow color
set scale_color 100


#Option to estimate energies (1) or forces (0)
#set estimate_energies 1
set estimate_energies 0

#Option to set VDW radii of restrained atoms proportional to standard deviation data stored in OCCUPANCY column
#set set_radii_to_std 1
set set_radii_to_std 0

#Scaling coefficient to (Option to set VDW radii of restrained atoms proportional to standard deviation data stored in OCCUPANCY column) - how many std to set
set scale_set_radii_to_std 1

#Option to save estimated vectors as x, y, z coordinates
#set save_vectors 1
set save_vectors 0

#Shift vectors are in coordinate columns as x, y, z (relative to the first frame, which is really coordinates)
#set vectors_instead_coordinates 1
set vectors_instead_coordinates 0

#Set the direction of the force applied - spring force <-1> or external force <1>
set show_spring_force -1
#set show_spring_force 1

#scale show_spring_force considering the scale_arrow constant
set scale_arrow [expr {$scale_arrow}]
if {$scale_arrow < 0} {
	set show_spring_force [expr {-1*$show_spring_force}]
	set scale_arrow [expr {-1*$scale_arrow}]
}


#----------- Go through the selection atoms and calculate forces
#set the reference selection
#puts "set the reference selection"
set reference_selection [atomselect top "index [$selection get index]"]
$reference_selection frame $reference_frame
$reference_selection update

#set the estimated selection
$selection frame $estimated_frame
$selection update
set estimated_selection $selection

if {$set_radii_to_std == 1} {
	#Set VDW radii of restrained atoms proportional to standard deviation data stored in OCCUPANCY column
	#Get radii from occupancies
	set radii [$selection get occupancy]
	#scale radii
	set radii [vecscale $scale_set_radii_to_std $radii]
	#Set new radii
	$selection set radius $radii
}


#cycle through atoms to calculate parameters
#puts "cycle through atoms to calculate parameters"
set new_occupancy {}
[atomselect top all] set occupancy 0
set max_arrow_length 0
set reconstructed_coordinates {}
foreach reference_coord [$reference_selection get {x y z}] estimated_coord [$estimated_selection get {x y z}] k [$reference_selection get beta] {
	#puts "cycle through atoms"
 	if {$vectors_instead_coordinates == 1} {
 		#There are Shift vectors are in coordinate columns as x, y, z (relative to the first frame, which is really coordinates)
 		#reconstruct real reference coordinates
 		#puts "reconstruct real reference coordinates"
 		set estimated_coord [vecadd $reference_coord $estimated_coord]
 		#remember reconstructed reference coordinates
 		lappend reconstructed_coordinates $estimated_coord
 	}

	# get the distance/direction
	set arrow_direction [vecsub $estimated_coord $reference_coord]
	#get the force/energy
	if {$estimate_energies == 1} {
		#calculate energy
		set arrow_length [expr {$k*[veclength2 $arrow_direction]/2}]
	} else {
		#calculate force
		set arrow_length [expr {$k*[veclength $arrow_direction]}]
	}
	lappend new_occupancy $arrow_length
	if {$max_arrow_length < $arrow_length} {set max_arrow_length $arrow_length}
}

if {$vectors_instead_coordinates == 1} {
	#record reconstructed reference coordinates as real reference x, y, z coordinates
	#puts "record reconstructed reference coordinates as real reference x, y, z coordinates"
	$estimated_selection set {x y z} $reconstructed_coordinates
}

# set the beta term
$reference_selection set occupancy $new_occupancy

#draw the color-coded arrows
set reconstructed_coordinates {}
foreach reference_coord [$reference_selection get {x y z}] estimated_coord [$estimated_selection get {x y z}] arrow_length [$reference_selection get occupancy] {
	# get the distance/direction
	set arrow_direction [vecscale $show_spring_force [vecsub $estimated_coord $reference_coord]]
	
	if {$save_vectors == 1} {
		#save estimated vectors as x, y, z coordinates
		#puts "save estimated vectors as x, y, z coordinates"
		lappend reconstructed_coordinates [vecscale $arrow_length [vecnorm $arrow_direction]]
	}
		
	#Draw the arrows
	if {$arrow_length > 0 & $save_vectors == 0} {
		#puts "arrow_length $arrow_length"
		draw color [expr {16 + 1007*$arrow_length/$max_arrow_length}]
		set arrow_length [expr {$arrow_length*$scale_arrow}]
		set arrow_end [vecadd $estimated_coord [vecscale $arrow_length [vecnorm $arrow_direction]]]
		set arrow_middle [vecadd $estimated_coord [vecscale 0.9 [vecsub $arrow_end $estimated_coord]]]
		set radius [expr {$arrow_length/10}]
		draw cone $arrow_middle $estimated_coord  radius [expr {$radius/1.5}]
		draw cone $arrow_middle $arrow_end radius $radius
	}
}

if {$save_vectors == 1} {
	#save estimated vectors as x, y, z coordinates
	#puts "save estimated vectors as x, y, z coordinates"
	$estimated_selection set {x y z} $reconstructed_coordinates 
}