# closewater.tcl
# Justin Gullingsrud
# jgulling@mccammon.ucsd.edu
# 8 November 2004

# This script processes a trajectory to create a new file containing just
# a selection of atoms and the N closest water to that selection.  The N
# waters are recomputed for each timestep, and need not be the same waters
# in each timestep (in fact, they probably will not be); thus it is in general
# meaningless to analyze the dynamics of individual waters.  However, it
# may be useful for analyzing the distribution of waters around a relatively
# static protein or DNA chain.

# usage: closewater <molid> <selection text> <# waters> <filename prefix>

proc closewater { molid seltext nwat prefix } {

  set wat [atomselect $molid "name OH2"]
  if { [$wat num] < $nwat } {
    error "Molecule contains ony [wat $num] water molecules."
  }
  set nframes [molinfo $molid get numframes]
  for { set i 0 } { $i < $nframes } { incr i } {
    set numinner 0
    set inner [list]
    set cut 3
    while {1} {
      set sel [atomselect $molid "name OH2 and within $cut of ($seltext)" frame $i]
      set outer [$sel list]
      $sel delete
      set numouter [llength $outer]
      if { $numouter < $nwat } {
        set inner $outer
        set numinner $numouter
        incr cut
        continue
      }
      break
    }
    puts "Found $numouter waters at cutoff $cut"
    # take all waters from inner, and just enough from outer
    catch { unset ohash }
    foreach ind $outer { set ohash($ind) 1 }
    foreach ind $inner { unset ohash($ind) }
    set outer [lrange [array names ohash] 0 [expr $nwat - $numinner - 1]]
    set watind [concat $inner $outer]
    set sel [atomselect $molid "($seltext) or same residue as (index $watind)"]
    $sel frame $i
    $sel writepdb [format $prefix-%04d.pdb $i]
    $sel delete
  }
}