writelammps.tcl

package provide writelammps 0.11
namespace eval ::writelammps:: {
	variable DEBUG
	variable atom_type
}

################# Write Lammps Data #############################

# export internal structure data to a LAMMPS data file.
# this requires that a corresponding set of information
# is already present in VMD's memory.
# Arguments:
# mol = molecule id with matching coordinate data
# filename = name of data file
# style = atom style
# sel = selection function for the subset to be written out.
# flags = more flags. (currently not used)

proc ::writelammps::writelammpsdata {mol filename style sel parameters dipoles ellipsoids {flags none}} {
    if {[catch {open $filename w} fp]} {
        vmdcon -error "writelammpsdata: problem opening data file: $fp\n"
        return -1
    }
    
    # initialize system default settings
    array set lammps {
        atoms 0 atomtypes 0 bonds 0 bondtypes 0 angles 0 angletypes 0 
        dihedrals 0 dihedraltypes 0 impropers 0 impropertypes 0 xtrabond 0
        xlo 0 xhi 0 ylo 0 yhi 0 zlo 0 zhi 0 xy 0 xz 0 yz 0 style unknown
    }

    # gather available system information
    set lammps(atoms)         [$sel num]
    set lammps(ellipsoids)     [llength $ellipsoids]
    set lammps(bonds)         [bondinfo     numbonds     $sel]
    set lammps(angles)        [angleinfo    numangles    $sel]
    set lammps(dihedrals)     [dihedralinfo numdihedrals $sel]
    set lammps(impropers)     [improperinfo numimpropers $sel]
    set lammps(atomtypes)     [llength [lsort -ascii -unique [$sel get {type}]]]
    set lammps(bondtypes)     [bondinfo     numbondtypes     $sel]
    set lammps(angletypes)    [angleinfo    numangletypes    $sel]
    set lammps(dihedraltypes) [dihedralinfo numdihedraltypes $sel]
    set lammps(impropertypes) [improperinfo numimpropertypes $sel]
    set lammps(style) $style

    # initialize simulation cell dimensions from min/max search
    lassign [measure minmax $sel -withradii] min max
    lassign $min xlo ylo zlo
    lassign $max xhi yhi zhi
    lassign [molinfo $mol get {a b c alpha beta gamma}] boxx boxy boxz alpha beta gamma
    
    # override min/max settings where box info available. 
    # try to (mostly) preserve the center of the cell, by
    # deriving it from the preceeding min/max search.
	set lammps(xlo) [expr $xlo-1000]
	set lammps(ylo) [expr $ylo-1000]
	set lammps(zlo) [expr $zlo-1000]
	
	set lammps(xhi) [expr $xhi+1000]
	set lammps(yhi) [expr $yhi+1000]
	set lammps(zhi) [expr $zhi+1000]
    set small 0.0001
    if {$boxx > $small} {
        set lammps(xmid) [expr {($xlo + $xhi) * 0.5}]
        set lammps(xlo)  [expr {-0.5*$boxx + $lammps(xmid)}]
        set lammps(xhi)  [expr { 0.5*$boxx + $lammps(xmid)}]
    }
    if {$boxy > $small} {
        set lammps(ymid) [expr {($ylo + $yhi) * 0.5}]
        set lammps(ylo)  [expr {-0.5*$boxy + $lammps(ymid)}]
        set lammps(yhi)  [expr { 0.5*$boxy + $lammps(ymid)}]
    }
    if {$boxz > $small} {
        set lammps(zmid) [expr {($zlo + $zhi) * 0.5}]
        set lammps(zlo)  [expr {-0.5*$boxz + $lammps(zmid)}]
        set lammps(zhi)  [expr { 0.5*$boxz + $lammps(zmid)}]
    }
    # XXX: need check for non-orthogonal cells.

    # write out supported data file sections
    writelammpsheader $fp [array get lammps]

    # write out hints about type to number mappings
    # for coefficient settings
    writelammpscoeffhint $fp $sel atoms $parameters
    if {$lammps(bonds) > 0} {
        writelammpscoeffhint $fp $sel bonds $parameters
    }
    if {$lammps(angles) > 0} {
        writelammpscoeffhint $fp $sel angles $parameters
    }
    if {$lammps(dihedrals) > 0} {
        writelammpscoeffhint $fp $sel dihedrals $parameters
    }
    if {$lammps(impropers) > 0} {
        writelammpscoeffhint $fp $sel impropers $parameters
    }

    writelammpsmasses $fp $sel
    writelammpsatoms $fp $sel $style $dipoles
    writelammpsellipsoids $fp $ellipsoids
    set atomidmap  [$sel list]
    if {$lammps(bonds) > 0} {
        writelammpsbonds $fp $sel $atomidmap
    }
    if {$lammps(angles) > 0} {
        writelammpsangles $fp $sel $atomidmap
    }
    if {$lammps(dihedrals) > 0} {
        writelammpsdihedrals $fp $sel $atomidmap
    }
    if {$lammps(impropers) > 0} {
        writelammpsimpropers $fp $sel $atomidmap
    }
    
    close $fp
    return 0
}


# write lammps header section to open file
proc ::writelammps::writelammpsheader {fp flags} {
    array set lammps $flags
    # first header line is skipped.
    puts $fp "LAMMPS data file. atom_style $lammps(style) generated by VMD/TopoTools on [clock format [clock seconds]] LBM 2011"
	puts $fp ""  
    foreach key {atoms ellipsoids bonds angles dihedrals impropers} {
        puts $fp [format " %d %s" $lammps($key) $key]
    }

    foreach key {atomtypes bondtypes  angletypes dihedraltypes impropertypes} {
        puts $fp [format " %d %s" $lammps($key) [regsub types $key " &"]]
    }
#measure box size check if solvated + 50A
    puts $fp [format " %.6f %.6f  xlo xhi" $lammps(xlo) $lammps(xhi)]
    puts $fp [format " %.6f %.6f  ylo yhi" $lammps(ylo) $lammps(yhi)]
    puts $fp [format " %.6f %.6f  zlo zhi" $lammps(zlo) $lammps(zhi)]

    puts $fp ""
    return
}

# write masses section, but only if number of masses 
# matches the number of atom types and if no mass is < 0.01
proc ::writelammps::writelammpsmasses {fp sel} {
    # first run the checks and build list of masses

    variable atom_type
    set typemap  [lsort -unique -ascii [$sel get type]]
    set masslist {}
    set mol [$sel molid]
    set idx [$sel list]
    foreach type $typemap {
        set tsel [atomselect $mol "(type \"$type\")"]
        set mass [lsort -unique -real [$tsel get mass]]
        $tsel delete
        if {[llength $mass] != 1} return
        if {$mass < 0.01} return
        lappend masslist $mass
    }

    # we passed the test, write out what we learned.
    vmdcon -info "writing LAMMPS Masses section."
    puts $fp " Masses\n"
    set typeid 1
    foreach mass $masslist type $typemap {
        puts $fp [format " %d %.6f \# %s" $typeid $mass $type]
        set atom_array($type) $typeid
        incr typeid
    }
    puts $fp ""
    set atom_type [array get atom_array]
    return
}

# write atoms section
proc ::writelammps::writelammpsatoms {fp sel style dipoles} {
    vmdcon -info "writing LAMMPS Atoms section in style '$style'."
    puts $fp " Atoms\n"
    set typemap [lsort -unique -ascii [$sel get type]]
    set resmap  [lsort -unique -integer [$sel get residue]]
    set atomid 0
	set seg [lsort -unique [$sel get segname]]
	set i 0

    foreach adat [$sel get {type residue charge mass x y z resname occupancy beta segname}] d $dipoles {
        lassign $adat type residue charge mass x y z resname occupancy beta segname
        lassign $d dx dy dz
        set atomtype [lsearch -sorted -ascii $typemap $type]
        set resid    [lsearch -sorted -integer $resmap $residue]
	
        incr atomid
        incr atomtype
        incr resid
        if { $dx==0 && $dy==0 && $dz==0 } {
        	set e 0
        } else {
        	set e 1.0
        }
        switch $style {
            atomic -
            dpd       { 
                puts $fp [format "%d %d %.6f %.6f %.6f \# %s" \
                              $atomid        $atomtype  $x $y $z $type] 
            }
            bond  -
            angle -
            molecular { 
                puts $fp [format "%d %d %d %.6f %.6f %.6f \# %s %s" \
                              $atomid $resid $atomtype  $x $y $z $type $resname] 
            }
			multiscale {
                puts $fp [format "%d %d %.6f %.6f %.6f %d %d %d %.6f %.6f %.6f %.6f %.3f %.3f \# %s %s" \
                              $atomid $atomtype $x $y $z [expr int($occupancy)] [expr int($occupancy)] [expr int($beta)] $charge $dx $dy $dz $e $mass  $type $resname] 
            }
            charge    { 
                puts $fp [format "%d %d %.6f %.6f %.6f %.6f \# %s" \
                              $atomid $atomtype $charge $x $y $z $type] 
            }
            full      { 
                puts $fp [format "%d %d %d %.6f %.6f %.6f %.6f \# %s %s" \
                              $atomid $resid $atomtype $charge $x $y $z $type $resname] 
            }
            hybrid      { 
                puts $fp [format "%d %d %d %.6f %.6f %.6f %.6f \# %s %s" \
                              $atomid $atomtype $x $y $z $resid $charge $type $resname] 
            }
            default   {
            	puts "Unsupported style"
                # ignore this unsupported style
                # XXX: add a way to flag an error. actually the test for 
                #      supported lammps atom styles should be done on a
                #      much higher level, so that we don't do unneeded work.
            }
        }
    }
    puts $fp ""
    return
}
proc ::writelammps::writelammpsellipsoids {fp ellipsoids} {
	vmdcon -info "writing LAMMPS Ellipsoids section."
	puts $fp " Ellipsoids\n"
	set atomid 0
	foreach e $ellipsoids {
		lassign $e atomid sx sy sz q0 q1 q2 q3
		puts $fp [format "%d %.6f %.6f %.6f %.6f %.6f %.6f %.6f" \
        	                      $atomid $sx $sy $sz $q0 $q1 $q2 $q3]
    }
    puts $fp ""
    return
}

proc ::writelammps::bondinfo {infotype sel {flag none}} {

    set numbonds 0
    set bidxlist {}
    array set bondtypes {}

    set aidxlist [$sel list]
    set bondlist [$sel getbonds]
    set btyplist [$sel getbondtypes]
    set bordlist [$sel getbondorders]
    
    foreach a $aidxlist bl $bondlist tl $btyplist ol $bordlist {
        foreach b $bl t $tl o $ol {
            if {($a < $b) && ([lsearch -sorted -integer $aidxlist $b] != -1)} {
                incr numbonds
                switch $flag {
                    type   {lappend bidxlist [list $a $b $t]}
                    order  {lappend bidxlist [list $a $b $o]}
                    both   {lappend bidxlist [list $a $b $t $o]}
                    lammps {lappend bidxlist [list $numbonds $a $b $t]}
                    none   {lappend bidxlist [list $a $b]}
                }
            }
            set bondtypes($t) 1
        }
    }

    switch $infotype {
        numbonds      { return $numbonds }
        numbondtypes  { return [array size bondtypes] }
        bondtypenames { return [array names bondtypes] }
        getbondlist   { return $bidxlist }
        default       { return "bug? shoot the programmer!"}
    }
}

# write bond section
proc ::writelammps::writelammpsbonds {fp sel atomidmap} {
    set bonddata  [bondinfo getbondlist   $sel type]
    set bondtypes [bondinfo bondtypenames $sel type]
    vmdcon -info "writing LAMMPS Bonds section."
    puts $fp " Bonds\n"

    set bondid 0
    foreach bdat $bonddata {
        incr bondid
        lassign $bdat a b t
        set at1  [lsearch -integer -sorted $atomidmap $a]
        set at2  [lsearch -integer -sorted $atomidmap $b]
        set type [lsearch -ascii $bondtypes $t]

        # go from 0-based to 1-based indexing and write out
        incr at1; incr at2; incr type
        puts $fp [format "%d %d %d %d" $bondid $type $at1 $at2]
    }
    puts $fp ""
    return
}

proc ::writelammps::angleinfo {infotype sel {flag none}} {

    set numangles 0
    array set angletypes {}
    set atomindex [$sel list]
    set anglelist {}

    foreach angle [join [molinfo [$sel molid] get angles]] {
        lassign $angle t a b c 

        if {([lsearch -sorted -integer $atomindex $a] >= 0)          \
                && ([lsearch -sorted -integer $atomindex $b] >= 0)   \
                && ([lsearch -sorted -integer $atomindex $c] >= 0) } {
            set angletypes($t) 1
            incr numangles
            lappend anglelist $angle
        }
    }
    switch $infotype {

        numangles      { return $numangles }
        numangletypes  { return [array size angletypes] }
        angletypenames { return [array names angletypes] }
        getanglelist   { return $anglelist }
        default        { return "bug! shoot the programmer?"}
    }
}


# write angle section
proc ::writelammps::writelammpsangles {fp sel atomidmap} {
    set angledata  [angleinfo getanglelist   $sel]
    set angletypes [angleinfo angletypenames $sel]
    vmdcon -info "writing LAMMPS Angles section."
    puts $fp " Angles\n"

    set angleid 0
    foreach adat $angledata {
        incr angleid
        lassign $adat t a b c
        set at1  [lsearch -integer -sorted $atomidmap $a]
        set at2  [lsearch -integer -sorted $atomidmap $b]
        set at3  [lsearch -integer -sorted $atomidmap $c]
        set type [lsearch -ascii $angletypes $t]
   
        # go from 0-based to 1-based indexing and write out
        incr at1; incr at2; incr at3; incr type
        puts $fp [format "%d %d %d %d %d" $angleid $type $at1 $at2 $at3]
    }
    puts $fp ""
    return
}

proc ::writelammps::dihedralinfo {infotype sel {flag none}} {

    set numdihedrals 0
    array set dihedraltypes {}
    set atomindex [$sel list]
    set dihedrallist {}

    foreach dihedral [join [molinfo [$sel molid] get dihedrals]] {
        lassign $dihedral t a b c d

        if {([lsearch -sorted -integer $atomindex $a] >= 0)          \
                && ([lsearch -sorted -integer $atomindex $b] >= 0)   \
                && ([lsearch -sorted -integer $atomindex $c] >= 0)   \
                && ([lsearch -sorted -integer $atomindex $d] >= 0) } {
            set dihedraltypes($t) 1
            incr numdihedrals
            lappend dihedrallist $dihedral
        }
    }
    switch $infotype {

        numdihedrals      { return $numdihedrals }
        numdihedraltypes  { return [array size dihedraltypes] }
        dihedraltypenames { return [array names dihedraltypes] }
        getdihedrallist   { return $dihedrallist }
        default        { return "bug! shoot the programmer?"}
    }
}


# write dihedral section
proc ::writelammps::writelammpsdihedrals {fp sel atomidmap} {
    set dihedraldata  [dihedralinfo getdihedrallist   $sel]
    set dihedraltypes [dihedralinfo dihedraltypenames $sel]
    vmdcon -info "writing LAMMPS Dihedrals section."
    puts $fp " Dihedrals\n"

    set dihedralid 0
    foreach adat $dihedraldata {
        incr dihedralid
        lassign $adat t a b c d
        set at1  [lsearch -integer -sorted $atomidmap $a]
        set at2  [lsearch -integer -sorted $atomidmap $b]
        set at3  [lsearch -integer -sorted $atomidmap $c]
        set at4  [lsearch -integer -sorted $atomidmap $d]
        set type [lsearch -ascii $dihedraltypes $t]
   
        # go from 0-based to 1-based indexing and write out
        incr at1; incr at2; incr at3; incr at4; incr type
        puts $fp [format "%d %d %d %d %d %d" $dihedralid $type $at1 $at2 $at3 $at4]
    }
    puts $fp ""
    return
}

proc ::writelammps::improperinfo {infotype sel {flag none}} {

    set numimpropers 0
    array set impropertypes {}
    set atomindex [$sel list]
    set improperlist {}

    foreach improper [join [molinfo [$sel molid] get impropers]] {
        lassign $improper t a b c d

        if {([lsearch -sorted -integer $atomindex $a] >= 0)          \
                && ([lsearch -sorted -integer $atomindex $b] >= 0)   \
                && ([lsearch -sorted -integer $atomindex $c] >= 0)   \
                && ([lsearch -sorted -integer $atomindex $d] >= 0) } {
            set impropertypes($t) 1
            incr numimpropers
            lappend improperlist $improper
        }
    }
    switch $infotype {

        numimpropers      { return $numimpropers }
        numimpropertypes  { return [array size impropertypes] }
        impropertypenames { return [array names impropertypes] }
        getimproperlist   { return $improperlist }
        default        { return "bug! shoot the programmer?"}
    }
}



# write improper section
proc ::writelammps::writelammpsimpropers {fp sel atomidmap} {
    set improperdata  [improperinfo getimproperlist   $sel]
    set impropertypes [improperinfo impropertypenames $sel]
    vmdcon -info "writing LAMMPS Impropers section."
    puts $fp " Impropers\n"

    set improperid 0
    foreach adat $improperdata {
        incr improperid
        lassign $adat t a b c d
        set at1  [lsearch -integer -sorted $atomidmap $a]
        set at2  [lsearch -integer -sorted $atomidmap $b]
        set at3  [lsearch -integer -sorted $atomidmap $c]
        set at4  [lsearch -integer -sorted $atomidmap $d]
        set type [lsearch -ascii $impropertypes $t]
   
        # go from 0-based to 1-based indexing and write out
        incr at1; incr at2; incr at3; incr at4; incr type
        puts $fp [format "%d %d %d %d %d %d" $improperid $type $at1 $at2 $at3 $at4]
    }
    puts $fp ""
    return
}

# returns 0 if lammps atom style is supported by topotools and 1 if not.
proc ::writelammps::checklammpsstyle {style} {
    switch $style {

        atomic -
        dpd   -
        hybrid -
        bond  -
        angle -
        molecular -
				multiscale -
        charge -
        full {
            return 0
        }

        default {
            return 1
        }
    }
}

# write hints about type coefficient mappings
proc ::writelammps::writelammpscoeffhint {fp sel type parameters} {
	lassign $parameters pairparameters dipoleparameters ellipsoidparameters bondparameters angleparameters improperparameters dihedralparameters
    switch $type {
        atoms {
        	array set pair_array $pairparameters
            puts $fp "Pair Coeffs\n"
            set aid 1
            set atlist [lsort -ascii -unique [$sel get {type}]]
            foreach at $atlist {
                puts $fp "$aid [lreplace $pair_array($at) 0 0 ]"
                incr aid
            }
        }
        bonds {
        	array set bond_array $bondparameters
            puts $fp "Bond Coeffs\n"
            set bid 1
            foreach bt [bondinfo bondtypenames $sel type] {
            	if {$bt!=""} {
	                puts $fp "$bid $bond_array($bt)"
    	            incr bid
    	        } else {
    	        	puts "no coefficient were found. Please provide paramter file"
    	        }
            }
        }
        angles {
        	array set angle_array $angleparameters
            puts $fp "Angle Coeffs\n"
            set aid 1
            foreach at [angleinfo angletypenames $sel] {           	
            	lassign $angle_array($at) a k1 k2 k3 u b
            	if {$k2==""} {
            		set k2 0.0
            	}
            	if {$k3 == ""} {
            		set k3 0.0
            	}
            	if {$u=="" || $b==""} {
            		set u 0.0
            		set b 0.0
            	}
                puts $fp [format "%d %.3f %.3f %.3f %.3f %.3f %.3f" $aid $a $k1 $k2 $k3 $u $b]
                incr aid
            }
        }
        dihedrals {
        	array set dihedral_array $dihedralparameters
            puts $fp "Dihedral Coeffs\n"
            set did 1
            foreach dt [dihedralinfo dihedraltypenames $sel] {
	            	lassign $dihedral_array($dt) k1 k2 k3 a1 a2 a3
	            	puts $fp [format "%d %.3f %.3f %.3f %d %d %d %d" $did $k1 $k2 $k3 [expr int($a1)] [expr int($a2)] [expr int($a3)] 0 ]
                	incr did
            }
        }
        impropers {
        	array set improper_array $improperparameters
            puts $fp "Improper Coeffs\n"
            set iid 1
            foreach it [improperinfo impropertypenames $sel] {
            	lassign $improper_array($it) k1 k2 k3 a1 a2 a3
            	if {$a2==""} {
            		set a2 0
            	}
            	if { $a3=="" } {
            		set a3 0
            	}
	            puts $fp [format "%d %.3f %.3f %.3f %d %d %d" $iid $k1 $k2 $k3 [expr int($a1)] [expr int($a2)] [expr int($a3)] ]
                incr iid
            }
        }
        default {
            vmdcon -warn "writelammpscoeffhint: don't know how to write hints for '$type'"
            return 1
        }
    }
    puts $fp ""
    return
}


proc ::writelammps::writelammpsinput { input sc_dipole standard} {
    if {[catch {open "inputlammps.in" w} fp]} {
        vmdcon -error "writelammpsdata: problem opening data file: $fp\n"
        return -1
    }
  	variable atom_type
  	array set atom_array $atom_type
  	set sc [list]
  	foreach d $sc_dipole {
		lappend sc $atom_array($d)
	}
	set sc [lsort -unique $sc]
	set std [list]
	foreach s $standard {
		lappend std $atom_array($s)
	}
	set std [lsort -unique $std]
	 puts $fp "\r"
     puts $fp "units\treal\r"
     puts $fp "neigh_modify\tdelay 2 every 1\r"
     puts $fp "\r"
     puts $fp "atom_style\thybrid multiscale ellipsoid\r"
     puts $fp "\r"
     puts $fp "bond_style\tmultiscale\r"
     puts $fp "angle_style\tmultiscale\r"
     puts $fp "dihedral_style\tmultiscale\r"
     puts $fp "improper_style\tmultiscale\r"
     puts $fp "pair_style\tmultiscale  triplet.dat 15.0 30.0\r"
     puts $fp "dielectric     1.0"
     puts $fp "boundary f f f"
     puts $fp "read_data\t$input"
     puts $fp "\r"
     puts $fp "# ----  timestep ---- \r"
     puts $fp "timestep 1.0\r"
     puts $fp "\r"
     puts $fp "# ---- integration parameters ---- \r"
     puts $fp "group aa molecule 0 1  \r"
     puts $fp "group cg molecule 3  \r"
     puts $fp "group hyb molecule 1 \r"
     puts $fp "group standard type $std \r"
     puts $fp "group sc_dipole type $sc \r"
     puts $fp "\r"   
     puts $fp "#  ---- print commands ---- \r"
     puts $fp "thermo 100\r"
     puts $fp "thermo_style multi\r"
     puts $fp "\r"
     puts $fp "# \r"
     puts $fp "#---- minimization ----\r"
     puts $fp "minimize 1.0e-4 1.0e-6 10000 10000\r"    
     puts $fp "\r"
     puts $fp "dump  1 all dcd 50 hybrid.dcd\r"
     puts $fp "#  ---- Integration parameters  ---- \r"
     puts $fp "fix 1 normal nve\r"
     puts $fp "fix 2 sc_dipole nve/asphere update dipole\r"
     puts $fp "fix 3 standard langevin 300.0 300.0 10.0 98459\r"
     puts $fp "#  ---- Respa parameters  ---- \r"
     puts $fp "run_style respa_multiscale 3 5 2 inner 1 middle 2 outer 3\r"
     puts $fp "# ----  run ---- \r"
     puts $fp "run 50000\r"
     close $fp
}