######################################################################
# pbcwrap                                                            #
#                                                                    #
# Wrap atoms of selection $sel around PBC unit cell boundaries.      #
#                                                                    #
# Author: Jan Saam                                                   #
#         Institute for Biochemistry                                 #
#         Charite Berlin                                             #
#         Germany                                                    #
#         saam@charite.de                                            #
#         ++49-30-450-528-449                                        #
#                                                                    #
# Feel free to send comments, bugs, etc.                             #
######################################################################

package provide pbcwrap 1.0


#########################################################################################
# pbcwrap                                                                               #
# ------------                                                                          #
# Wraps atoms of selection $sel around PBC unit cell boundaries.                        #
# Unit cell geometry is read from the xst-file $xst.                                    #
#                                                                                       #
# Options 'beg' and 'end' are used to denote the first and last frame to process,       #
# option 'firsttime' allows to specify the first timestep in the xst-file to take into  #
# account. With 'delta' you tell the program the time that lies between two frames in   #
# the dcd-file, the corresponding records in the xst-file are chosen accordingly.       #
# This can be useful if you preprocessed the trajectory with catdcd.                    #
# Note that the first timestep in $xst is omitted unless you demand it by using		#
# 'firsttime', because xst info starts at frame 0 while dcd files start at frame 1 	#
# (when both are generated by NAMD).                   					#
#											#
# If you specify 'outfile' the control output is written to this file instead of the 	#
# console. (The output helps checking if your dcd frames and xst times match.)		#
#                                                                                       #
# The wrapping of atoms is done by transforming the unit cell to a orthonormal cell     #
# which allows to easily select atoms out side the cell (x<1 or x>1, ...).              #
# After warping them along the coordinate axes into the cell, the system is             #
# transformed back.     								#
#                                                                                       #
# Usage:                                                                                #
# pbcwrap selection xstfile [beg $first] [end $last] [firsttime $firsttime]             #
#         [delta $delta] [outfile $outfile]						#
#                                                                                       #
# Example:                                                                              #
# set sel [atomselect top "segname OXY"]                                                #
# pbcwrap $sel mysimulation.xst beg 100 end 200 firsttime 50000 delta 1000              #
#											#
# (The first frame of the trajectory is identified with timestep 50000 fs of the 	#
# xst file. There are 1000 fs between two subsequent frames. Only frames 100-200 are	#
# processed.)										#
#########################################################################################

proc pbcwrap { sel xstfile args} {
   set beg       0
   set end       last
   set firsttime -1
   set delta	 -1
   set outfile   stdout
   set out 	 stdout	
   foreach {i j} $args {
      if {$i=="beg"}       then { set beg $j }
      if {$i=="end"}       then { set end $j }
      if {$i=="firsttime"} then { set firsttime $j }
      if {$i=="delta"}     then { set delta $j }
      if {$i=="outfile"}   then { set outfile $j }
   }
   set numframes [molinfo [$sel molid] get numframes]
   if {$end=="last"} then {set end [expr $numframes-1]}

   set xst [open $xstfile]
   if {$outfile!="stdout"} { set out [open $outfile a] }

   set dt 0
   set time 0
   set frame 0
   set numline 0
   while {![eof $xst]} {
      set line [gets $xst]

      if {[string first \# $line]==-1} {
	 #puts $line
	 incr numline
	 set oldtime $time;
	 set olddt $dt
 	 set time [lindex $line 0]
	 set dt [expr $time-$oldtime];

         # The first line is omitted if firsttime is not specified, 
	 # because xst info starts at frame 0 while dcd record starts at frame 1
	 if {$numline==1 && $firsttime==-1} { continue }
	 if {$numline==2 && $firsttime==-1} { 
	    set firsttime $time
            if  {$delta==-1} { set delta $dt }
	    puts $out "delta=$delta, firsttime=$firsttime"
	 }
         if {[expr fmod([expr $time],$delta)] != 0.0} { continue }
	 if {!$numline==1 && $dt!=$olddt} {puts "\nWARNING Stepsize changed! dt=$dt, olddt=$olddt\n"}
	 if {$time<$firsttime} { continue }
	 if {$frame<$beg}      { incr frame; continue }
	 if {$frame>$end}      { break }

	 # Get half length PBC vectors
	 set a1 [vecscale 0.5 [lrange $line 1 3]]
	 set a2 [vecscale 0.5 [lrange $line 4 6]]
	 set a3 [vecscale 0.5 [lrange $line 7 9]]
	 set ori [lrange $line 10 12]
	 puts $out "time=$time \t frame=$frame"

	 # Orthonormalize system:
	 # Find an orthonormal basis (in cartesian coords)
	 set obase [orthonormal_basis $a1 $a2 $a3]
	 set b_cartcoor  [basis_change $obase [list {1 0 0} {0 1 0} {0 0 1}] ]
	 set b2carti [trans_from_rotate $b_cartcoor]
	 set b2cart  [measure inverse $b2carti]

	 # Get coordinates of $a in terms of $obase
	 set m [basis_change [list $a1 $a2 $a3] $obase]
	 set rmat [measure inverse [trans_from_rotate $m]]
	 # actually: [transmult $b2cart $b2carti $rmat $b2cart]
	 set mat4 [transmult $rmat $b2cart]
	 

	 # Go to the right frame and transform into ONS
	 $sel frame $frame
	 $sel move $mat4

	 # Now it is easy to do the selection: 
	 # The unit cell is now a rectangular box of size 2x2x2 A.
	 # Let's wrap atoms around the periodic boundaries
         set outsel [atomselect top "([$sel text]) and (not same residue as (x>-1))" frame $frame]
	 $outsel moveby { 2  0  0}
         set outsel [atomselect top "([$sel text]) and (not same residue as (x<1))"  frame $frame]
	 $outsel moveby {-2  0  0}
         set outsel [atomselect top "([$sel text]) and (not same residue as (y>-1))" frame $frame]
	 $outsel moveby { 0  2  0}
         set outsel [atomselect top "([$sel text]) and (not same residue as (y<1))"  frame $frame]
	 $outsel moveby { 0 -2  0}
         set outsel [atomselect top "([$sel text]) and (not same residue as (z>-1))" frame $frame]
	 $outsel moveby { 0  0  2}
         set outsel [atomselect top "([$sel text]) and (not same residue as (z<1))"  frame $frame]
	 $outsel moveby { 0  0 -2}

	 # Transform back
	 $sel move [measure inverse $mat4]
	 incr frame
      }
   }
   close $xst
   if {$out!="stdout"} { close $out }
   return $time
}


########################################################################################
# check_pbcwrap                                                                        #
# --------------                                                                       #
# Check the results of pbc_wrap                                                        #
# It does a simple distance check for every frame                                      #
# If the wrapping was done right, atom distances below some value (default=1.9A)       #
# should not occur, as they would be avoided during the simulation by VDW interaction. #
# If you encounter a distance of lets say 0.6A, most likely something went wrong!      #
# (Note: In case you smoothed the trajectory beforehand, you'll see much smaller       #
# distances)                                                                           #
########################################################################################

proc check_pbc_wrap { sel {dist 1.9} } {
   set sum 0
   set molid [$sel molid]
   set numframes [molinfo $molid get numframes]
   set checksel [atomselect $molid "[$sel text] and within $dist of (not [$sel text])"]
   for {set i 0} {$i<$numframes} {incr i} {
      $checksel frame $i
      $checksel update
      set close [$checksel num]
      puts "frame $i: $close atom distances under $dist A" 
      incr sum $close
   }
   puts "Summary: $sum atom distances under $dist A in $i frames."
}


#############################################################################
# basis_change                                                              #
# ------------                                                              #         
# Returns $base in coordinates of $obase. $obase must be orthonormal        #
#############################################################################

proc basis_change { base obase } {
   set dim1 [llength $base]
   set dim2 [llength [lindex $base 0]]
   set cc ""
   foreach i $obase {
      set c ""
      foreach j $base {
	 lappend c [vecdot $j $i]
      }
      lappend cc $c
   }
   return $cc
}

#############################################################################
# orthonormal_basis                                                         #
# -----------------                                                         #         
# Find an orthononal basis R^3 with $ob1 || $b1                             #
#############################################################################

proc orthonormal_basis { b1 b2 b3 } {
   set ob1 $b1
   set e1  [vecnorm $ob1]
   set ob2 [vecsub $b2  [vecscale [vecdot $e1 $b2] $e1]]
   set e2  [vecnorm $ob2]
   set ob3 [vecsub $b3  [vecscale [vecdot $e1 $b3] $e1]]
   set ob3 [vecsub $ob3 [vecscale [vecdot $e2 $b3] $e2]]
   set e3  [vecnorm $ob3]
   return [list $e1 $e2 $e3]
}