gmx_ir.bsh

usage="\
>>>>>>>>>>>>>>>>    gmx_ir     <<<<<<<<<<<<<<<<
>>>>>>>>>>>>>>>>    Jicun LI   <<<<<<<<<<<<<<<<
>>>>>>>>>>     2023-03-07 21:27:44    <<<<<<<<<
>>   Usage: gmx_ir [-s File.tpr] [-f File.trr] [-n  File.ndx] [-g Group]
                   [-b timefrom] [-e timeto  ]
>> Default: gmx_ir  -s topol.tpr  -f traj.trr   -n index.ndx   -g System
>>
>>2023-03-07: fix bug for VSite (e.g. TIP4P)
>>            add time range option                                      \n\n"

# 默认设置
gmx="gmx"      # gmx命令
tpr=topol.tpr  # 拓扑文件
trr=traj.trr   # 轨迹文件
ndx=index.ndx  # 索引文件
grp=System     # 索引组名称

# 命令行选项
opt=($*); N=${#opt[@]}
for((i=0; i<N; i++)); do
	arg=${opt[$i]}; val=${opt[$((i+1))]}
	[[ $arg =~ -s ]] && { tpr=$val; }
	[[ $arg =~ -f ]] && { trr=$val; }
	[[ $arg =~ -n ]] && { ndx=$val; }
	[[ $arg =~ -g ]] && { grp=$val; }
	[[ $arg =~ -b ]] && { tb=$val; }
	[[ $arg =~ -e ]] && { te=$val; }
done

# 起止时间
[[ ! -z "$tb" ]] && { time="-b $tb";       ft="_$tb"; }
[[ ! -z "$te" ]] && { time="$time -e $te"; ft="$ft-$te"; }

[[ ! -f "$tpr" ]] && { echo -e "$usage !!! ERROR !!! topology File   $tpr  NOT Exist !\n"; exit; }
[[ ! -f "$trr" ]] && { echo -e "$usage !!! ERROR !!! trajectory File $trr  NOT Exist !\n"; exit; }
[[ ! -f "$ndx" ]] && { echo -e "$usage !!! ERROR !!! index File      $ndx  NOT Exist !\n"; exit; }

# 中间文件
val=${trr%.trr}
atm=$val~atm.ndx  # 单原子索引
tpx=$val~atm.tpr  # 单原子拓扑
chg=$val~chg.dat  # 电荷文件
vgro=$val~v.gro   # 含速度的gro文件
vqgro=$val~vq.gro # 含vq的gro文件
vqtrr=$val~vq.trr # 含vq的trr文件

# 获取电荷信息
$gmx dump -quiet -s $tpr | awk -v ndx=$ndx -v grp=$grp -v atm=$atm >$chg '

BEGIN {
	RS="[";
	while(getline < ndx) { gsub(/\s/,"", $1)
		if($1==grp)    for(i=3; i<=NF; i++) ndxGrp[$i+0]++
		if($1==grp"]") for(i=2; i<=NF; i++) ndxGrp[$i+0]++
	}
	RS="\r?\n"
}

/#molblock/  { Ntyp=$3+0 }

/moltype.+=/ { Imol=$3;
	gsub(/"/,"",$4); Name[Imol]=$4; #"
	getline; Nmol[Imol]=$3+0
}

/moltype.+\(/ {
	gsub(/[^0-9]/,"",$0); Imol=$0
	getline; getline;
	getline; gsub(/[^0-9]/,"",$0); Natm[Imol]=$0+0
	for(i=0; i<Natm[Imol]; i++) {
		getline;
		txt=$0; sub(/.+q=/, "", txt); sub(/,.+/, "", txt); Qatm[Imol, i]=txt
		txt=$0; sub(/.+ptype= */, "", txt); sub(/,.+/, "", txt); Patm[Imol, i]=txt
	}
	getline
	for(i=0; i<Natm[Imol]; i++) {
		getline; txt=$0
		sub(/.+=./, "", txt)
		sub(/..$/, "", txt)
		Tatm[Imol, i]=txt
	}
}

END {
	#print "#Mol  Name  #Atom #Num"
	Ntot=0
	for(i=0; i<Ntyp; i++) {
		#printf "%3d %6s %3d %5d\n",  i+1, Name[i], Natm[i], Nmol[i]
		for(n=0; n<Nmol[i]; n++) {
			for(j=0; j<Natm[i]; j++) {
				Ntot++
				if(Ntot in ndxGrp) {
					if(Patm[i, j]=="Atom") maxAtm=Ntot
					printf "%6d %15s %s %5d %s  %-s\n",
					Ntot, Qatm[i, j], n+1"."Name[i], j+1, Tatm[i, j], Patm[i, j]
				}
			}
		}
	}
	print "[ "grp"_LAST ]\n"maxAtm >atm
} '

# 获得带速度的gro文件
echo $grp | $gmx trjconv -s $tpr -f $trr -n $ndx -o $vgro -vel $time

# 使用vq替换速度
awk -v chg=$chg -v grp=$grp '
BEGIN {
	FMT="%s%15.8f%15.8f%15.8f%15.9f%15.9f%15.9f\r\n"
	i=0; while(getline<chg) { i++; q[i]=$2; }
}

{	print $0"\r"
	getline Natm; print "    1\r"
	qvx=0; qvy=0; qvz=0
	for(i=1; i<=Natm; i++) {
		getline
		qvx += q[i]*$(NF-2); qvy += q[i]*$(NF-1); qvz += q[i]*$NF
	}
	printf FMT, substr($0,1,20), $(NF-5), $(NF-4), $(NF-3), qvx, qvy, qvz

	getline; print $0"\r"
} ' $vgro >$vqgro

# 抽取单个原子的拓扑
echo "${grp}_LAST" | $gmx convert-tpr -s $tpr -n $atm -o $tpx

# 将gro转换为trr
echo 0 | $gmx trjconv -s $tpx -f $vqgro -o $vqtrr

# 计算光谱
echo 0 | $gmx velacc -f $vqtrr -s $tpx -o vac.xvg -os spectrum$ft.xvg -mol

# 清理中间文件
rm -f \#*
#rm -f $atm $tpx $chg $vgro $vqgro $vqtrr