# Script:  mkinput_lamellae_diblock.py 
# Purpose: Make input file of few layers lamellae of diblock copolymers in Kremer-Grest approach (JCP, 105, 1996); graft half-polymers on two facing walls and then mirror them to the opposite directions to complete one layer)
# Syntax: "python mkinput_lamellae_diblock.py"
# Author: Modified Mark Stevens's mkinput.py by Youngmi Seo (10/2014)

# chains are random walk.

#  LAMMPS types (all neutral)
#  1         soft polymer beads (A)
#  2         hard polymer beads (B)

import sys, string, math, itertools
from numpy import *
from random import *
#-------------------------------------------------------------------------

## 40 bead/monomes, 120 poly = 4800 beads
## ADJUSTED DENSITY
## INPUT PARAMETERs

#iseed = 9328             # random number seed
#seed(iseed)              # to generate same set of random numbers every time
nbeads = 40               # number of beads in monomer
hnbeads = nbeads/2        # half number of beads in monomer
nmonomersperpoly = 1      # total (including x) number of monomers in polymer
npoly = 480               # number of polymers
hnpoly = npoly/2          # half number of polymers
ntot = npoly*nbeads       # number of total beads
nlayers = 2               # number of lamellar layers
ntypes = 4                # number of monomer types
nbtypes = 2               # number of bond types
minsep = 1.0              # allowed separation in overlap check
dens = 0.85               # bead number density
cdens = 0.1               # coverage density on a planar surface
bond = 0.97               # bond length. depends on bond potential, but close to 1 is good enough

# Label schemes
# type names: conducting block(PEO), nonconducting block(PS)
#                       1                      2   

atype = ( '00', 'A' , 'B' )

# Define monomer sequence: diblock copolymers
sequence = [1]*20 + [2]*20


    
## END INPUT Parameters ------------------------------------------------------
# Files
INPUT_LAMMPS = open('input_ghost_lamellae.lammps', 'w')

# Simulation cell parameters for grafting polymers
gnpoly = npoly/nlayers   #for grafting polymers
gntot = gnpoly*hnbeads
gvol = gntot/dens
surf = gnpoly/2/cdens
ghx = surf**0.5
ghy = surf**0.5
ghz = gvol/surf
ghx2 = ghx/2.
ghy2 = ghy/2.
ghz2 = ghz/2.
gvol = ghx*ghy*ghz
gnbonds = gntot-gnpoly

# Simulation cell parameters
dim = ntot
vol = ntot/dens
hx = ghx
hy = ghy
hz = ghz*2*nlayers + (nlayers*2 - 1)
hx2 = hx/2.
hy2 = hy/2.
hz2 = hz/2.
vol = hx*hy*hz
nbonds = ntot-npoly

#Ghost Particle Constants
nlayers = 2                                      #number of lamellar layers
#xd, yd, and zd are the box dimensions generated from running npt without ghosts for 10mil timesteps
xd = 26.480464633825079
yd = 26.480464633825079
zd = 34.819861778846153
layersize = zd / (nlayers * 2.)                   #z dimension of one layer
ghostzlayers = 12                                #ghost particles in the z direction in a layer
ghostxlayers = int(ghostzlayers*(xd/layersize))  #ghost particles in the x direction in a layer
ghostylayers = int(ghostzlayers*(xd/layersize))  #ghost particles in the z direction in a layer
spacing = zd / (2. * nlayers) / ghostzlayers                  #amount of space between each ghost particle

#Determine Box Size, Redefine Constants
xd = spacing * ghostxlayers
yd = spacing * ghostylayers
zd = spacing * ghostzlayers * 2. * nlayers
#x, y, and z are half of the box dimensions (used as each dimension is centered at the origin)
x = xd/2.
y = yd/2.
z = zd/2.
layersize = zd / (nlayers * 2.)                   #z dimension of one layer
spacing = zd / (2. * nlayers) / ghostzlayers                  #amount of space between each ghost particle

#Determine nghost
nghost = 0                   #number of ghost particles (determined as they are generated)
for l in xrange(0, nlayers):
    for _x in xrange(0, ghostxlayers):
        for _y in xrange(0, ghostylayers):
            for _z in xrange(0, ghostzlayers):
                nghost += 1
nghost2 = nghost

# simulation cell parameters
hx = xd
hy = yd
hz = zd
hx2 = x
hy2 = y
hz2 = z
vol = hx*hy*hz
nbonds = ntot-npoly
ghx = hx
ghy = hy
ghz = (hz - (nlayers*2 -1))/(2. * nlayers)

 
print "Total number of beads:",ntot
print "Number of chains =", npoly
print "Number of Beads in monomer =", nbeads
print "Number of atoms types = ",ntypes
#print "seed = ", iseed
print " "
print "Geometry:"
print "dens = ", dens
print "vol = ", vol
print " "
print "metric: %10.4f %10.4f %10.4f\n\n" % (hx, hy, hz)

# Initial variables
xc=zeros(dim+nghost+nghost2+1,float32)
yc=zeros(dim+nghost+nghost2+1,float32)
zc=zeros(dim+nghost+nghost2+1,float32)
cx=zeros(dim+1)
cy=zeros(dim+1)
cz=zeros(dim+1)

rg2ave=0.0
rgave=0.0
rend2ave = 0.0
typeb=[0]*(dim+1)
molnum=[0]*(dim+1)
k=0

## BUILD POLYMERS
for nl in range(0,nlayers):
    for ix in range((hnpoly/nlayers)*nl*2,(hnpoly/nlayers)*(nl*2+1)):      
        lengthcurrentpoly = 0
        for iy in range(nmonomersperpoly):
            currentmonomer = ix*nmonomersperpoly + iy
            seqnum = 0
            seq = sequence
            for iz in seq:
                seqnum = seqnum + 1
                k = k + 1
                kk = k + 19 - ((k-1) % 20)*2
                lengthcurrentpoly = lengthcurrentpoly + 1
                if iy == 0 and seqnum == 1:
                    typeb[kk] = iz
                    molnum[kk] = ix + 1 
                    k1 = k
                    xc[kk] = random()*ghx
                    yc[kk] = random()*ghy
                    zc[kk] = 0.0000 + (ghz*nl*2 + nl*2)
                elif seqnum <= hnbeads:          # pick random direction; scale to be bond length
                    typeb[kk] = iz
                    molnum[kk] = ix + 1 
                    theta = random()*2*pi
                    dz = random()*2 - 1 
                    dx = sqrt(1-dz**2)*cos(theta)
                    dy = sqrt(1-dz**2)*sin(theta)                  
                    r = sqrt(dx*dx+dy*dy+dz*dz)
                    scale = bond/r
                    dx = scale*dx
                    dy = scale*dy
                    dz = scale*dz
                    xc[kk] = xc[kk+1] + dx
                    yc[kk] = yc[kk+1] + dy
                    zc[kk] = zc[kk+1] + dz        # don't need to shift, already done for the first bead
                    if zc[kk] >= (ghz + (ghz*nl*2 + nl*2)) or zc[kk] <= (ghz*nl*2 + nl*2):
                        zc[kk] = zc[kk+1] - dz
                else:
                    typeb[k] = iz
                    molnum[k] = ix + 1
                    xc[k] = xc[kk-20]
                    yc[k] = yc[kk-20]
                    zc[k] = zc[kk-20] - (zc[kk-20] - zc[kk-20 + ((k-21) % 20)] + 0.5)*2
                        
    for ix in range((hnpoly/nlayers)*(nl*2+1),(hnpoly/nlayers)*(nl*2+2)):
        for iy in range(nmonomersperpoly):
            currentmonomer = ix*nmonomersperpoly + iy
            seqnum = 0
            seq = sequence
            for iz in seq:
                seqnum = seqnum + 1
                k = k + 1
                kk = k + 19 - ((k-1) % 20)*2
                lengthcurrentpoly = lengthcurrentpoly + 1                
                if iy == 0 and seqnum == 1:
                    typeb[kk] = iz
                    molnum[kk] = ix + 1 
                    k1 = k
                    xc[kk] = random()*ghx
                    yc[kk] = random()*ghy
                    zc[kk] = ghz + (ghz*nl*2 + nl*2)
                elif seqnum <= hnbeads:        # pick random direction; scale to be bond length
                    typeb[kk] = iz
                    molnum[kk] = ix + 1 
                    theta = random()*2*pi
                    dz = random()*2 - 1
                    dx = sqrt(1-dz**2)*cos(theta)
                    dy = sqrt(1-dz**2)*sin(theta)
                    r = sqrt(dx*dx+dy*dy+dz*dz)
                    scale = bond/r
                    dx = scale*dx
                    dy = scale*dy
                    dz = scale*dz   
                    xc[kk] = xc[kk+1] + dx
                    yc[kk] = yc[kk+1] + dy
                    zc[kk] = zc[kk+1] + dz
                    if zc[kk] >= (ghz + (ghz*nl*2 + nl*2)) or zc[kk] <= (ghz*nl*2 + nl*2):
                        zc[kk] = zc[kk+1] - dz
                else:
                    typeb[k] = iz
                    molnum[k] = ix + 1
                    xc[k] = xc[kk-20]
                    yc[k] = yc[kk-20]
                    zc[k] = zc[kk-20] - (zc[kk-20] - zc[kk-20 + ((k-21) % 20)] - 0.5)*2

#PLACE GHOST PARTICLES
for l in xrange(0, nlayers):
    for _x in xrange(0, ghostxlayers):
        for _y in xrange(0, ghostylayers):
            for _z in xrange(0, ghostzlayers):
                k += 1
                xc[k] = (0.5 * spacing) - x + (_x * spacing)
                yc[k] = (0.5 * spacing) - y + (_y * spacing)
                zc[k] = (0.5 * spacing) - z + (_z * spacing) + (zd * (l * 1.0/ nlayers))
for m in xrange(0, nlayers):
    for _x in xrange(0, ghostxlayers):
        for _y in xrange(0, ghostylayers):
            for _z in xrange(0, ghostzlayers):
                k += 1
                xc[k] = (0.5 * spacing) - x + (_x * spacing)
                yc[k] = (0.5 * spacing) - y + (_y * spacing)
                zc[k] = (0.5 * spacing) - z + layersize + (_z * spacing) + (zd * (m * 1.0/ nlayers))


## IMAGE FLAGS
for k in xrange(1,ntot+1):
    if (xc[k] > hx):
        cx[k] = int(xc[k]/hx)
        xc[k] = xc[k] - cx[k]*hx - hx2
    elif (xc[k] < 0.0):
        cx[k] = -int((-xc[k]+hx)/hx)
        xc[k] = xc[k] - cx[k]*hx - hx2
    else:
        cx[k] = 0
        xc[k] = xc[k] - hx2
        
    if (yc[k] > hy):
        cy[k] = int(yc[k]/hy)
        yc[k] = yc[k] - cy[k]*hy - hy2
    elif (yc[k] < 0.0):
        cy[k] = -int((-yc[k]+hy)/hy)
        yc[k] = yc[k] - cy[k]*hy - hy2
    else:
        cy[k] = 0
        yc[k] = yc[k] - hy2
        
    if (zc[k] > hz):
        cz[k] = int(zc[k]/hz)
        zc[k] = zc[k] - cz[k]*hz - hz2
    elif (zc[k] < 0.0):
        cz[k] = -int((-zc[k]+hz)/hz)
        zc[k] = zc[k] - cz[k]*hz - hz2
    else:
        cz[k] = 0
        zc[k] = zc[k] - hz2

print "Polymers built."

## OUTPUT Headers ---------------------------------------------------------------

# Input.lammps header 
INPUT_LAMMPS.write("# Kremer-Grest type Lamellae YS 10/2014\n")
INPUT_LAMMPS.write("# Number of lamellar layers = %1.0f\n" % nlayers)
INPUT_LAMMPS.write("\n")
INPUT_LAMMPS.write("%10i    atoms\n" % (dim+nghost+nghost2))
INPUT_LAMMPS.write("%10i    bonds\n" % nbonds)
INPUT_LAMMPS.write("%10i    angles\n" % 0)
INPUT_LAMMPS.write("%10i    dihedrals\n" % 0)
INPUT_LAMMPS.write("\n")
INPUT_LAMMPS.write("%10i    atom types\n" % ntypes)
INPUT_LAMMPS.write("%10i    bond types\n" % nbtypes)
INPUT_LAMMPS.write("%10i    angle types\n" % 0)
INPUT_LAMMPS.write("%10i    dihedral types\n" % 0)
INPUT_LAMMPS.write("\n")
INPUT_LAMMPS.write(" %16.8f %16.8f   xlo xhi\n" % (-hx2,hx2))
INPUT_LAMMPS.write(" %16.8f %16.8f   ylo yhi\n" % (-hy2,hy2))
INPUT_LAMMPS.write(" %16.8f %16.8f   zlo zhi\n" % (-hz2,hz2))
INPUT_LAMMPS.write("\n")
INPUT_LAMMPS.write("Atoms\n")
INPUT_LAMMPS.write("\n")

## END OUTPUT Headers -----------------------------------------------------------

mass = 1.0
i = 0
imol = 0

# Positions 
for i in xrange(1,dim+1):
    itype = typeb[i]
    #aname = atype[itype]
    if itype != 3:
        imol = molnum[i]
    elif itype == 3:       #for the case there are penetrants or ions added
        imol = i-ntot+npoly 
    INPUT_LAMMPS.write("%6i %6i %2i %9.4f %9.4f %9.4f %6i %6i %6i\n" % (i, imol, typeb[i], xc[i], yc[i], zc[i], cx[i], cy[i], cz[i]))

for j in xrange(dim+1,dim+1+nghost):
    type3 = 3
    imol = npoly + j - ntot
    image_flag = 0
    INPUT_LAMMPS.write("%6i %6i %2i %9.4f %9.4f %9.4f %6i %6i %6i\n" % (j, imol, type3, xc[j], yc[j], zc[j], image_flag, image_flag, image_flag))

for n in xrange(dim+1+nghost,dim+1+nghost+nghost2):
    type4 = 4
    imol = imol + 1
    image_flag = 0
    INPUT_LAMMPS.write("%6i %6i %2i %9.4f %9.4f %9.4f %6i %6i %6i\n" % (n, imol, type4, xc[n], yc[n], zc[n], image_flag, image_flag, image_flag))


# Bonds
INPUT_LAMMPS.write("\n")
INPUT_LAMMPS.write("Bonds\n")
INPUT_LAMMPS.write("\n")
jbond1 = zeros(nbonds+1)
jbond2 = zeros(nbonds+1)
pbond1 = zeros(nbonds+1)
pbond2 = zeros(nbonds+1)
ibond=0
for i in xrange(1,ntot):       
        if molnum[i+1] == molnum[i]:           #if not at the end of the polymer
            ibond = ibond+1 #bond number
            j=i+1
            if typeb[i] != typeb[j]:
                INPUT_LAMMPS.write("%8i  2 %8i %8i\n" % (ibond,i,j))
            else: 
                INPUT_LAMMPS.write("%8i  1 %8i %8i\n" % (ibond,i,j))


# Masses
INPUT_LAMMPS.write("\n")
INPUT_LAMMPS.write("Masses\n")
INPUT_LAMMPS.write("\n")

for ii in xrange(1,ntypes+1):
    INPUT_LAMMPS.write("%3i  1.0\n" % ii)

# Close files
INPUT_LAMMPS.close()

print "LAMMPS output complete."



