<commit_before>#!/usr/bin/env python
"""
generate_strain.py

Generate strained castep input files and 
.cijdat files for elastic constants calculation
for analysis with elastics.py

Copyright (c) 2010 <NAME> (<EMAIL>)
All rights reserved.
"""

import sys
import os
import optparse
import re
import math

def GetStrainPatterns(code):
	"""
	Given a code number for the crystal symmetry, 
	returns a list of strain patterns needed for
	the calculation of the elastic constants tensor.
	Each pattern is a 6 element list, the subscript
	reflects the strain in IRE notation

	Supported Strain Patterns
	-------------------------

	5 Cubic: e1+e4
	7 Hexagonal: e3 and e1+e4
	7 Trigonal-High (32, 3m, -3m): e1 and e3+e4
	6 Trigonal-Low (3, -3): e1 and e3+e4
	4 Tetragonal: e1+e4 and e3+e6
	3 Orthorhombic: e1+e4 and e2+e5 and e3+e6
	2 Monoclinic: e1+e4 and e3+e6 and e2 and e5
	1 Triclinic: e1 to e6 separately
	0 Unknown...
	"""

	if (code == 1):
		pattern = [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
		           [0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
		           [0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
		           [0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
		           [0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
		           [0.0, 0.0, 0.0, 0.0, 0.0, 1.0]]
	elif (code == 2):
		pattern = [[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
		           [0.0, 0.0, 1.0, 0.0, 0.0, 1.0],
		           [0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
		           [0.0, 0.0, 0.0, 0.0, 1.0, 0.0]]
	elif (code == 3):
		pattern = [[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
		           [0.0, 1.0, 0.0, 0.0, 1.0, 0.0],
		           [0.0, 0.0, 1.0, 0.0, 0.0, 1.0]]
	elif (code == 4):
		pattern = [[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
		           [0.0, 0.0, 1.0, 0.0, 0.0, 1.0]]
	elif (code == 5):
		pattern = [[1.0, 0.0, 0.0, 1.0, 0.0, 0.0]]
	elif (code == 6):
		pattern = [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
		           [0.0, 0.0, 1.0, 1.0, 0.0, 0.0]]
	elif (code == 7):
		# NB is this correct for hex and hig trig? - see missmatch above/
		# I suspect I have to rotate lattice for trig high?
		pattern = [[0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
		           [1.0, 0.0, 0.0, 1.0, 0.0, 0.0]]

	return pattern	

def main(input_options, libmode=False):

	# deal with options
	def get_options():
		if not libmode:
			p = optparse.OptionParser()
			p.add_option('--debug', '-d', action='store_true', \
			              help="Debug mode (output to stdout rather than file)")
			p.add_option('--num_steps', '-n', action='store', type='int', dest="numsteps")
			p.add_option('--strain_mag', '-s', action='store', type='float', dest="strain")
			p.add_option('--lattice_type', '-l', action='store', type='int', dest="lattice")
			options,arguments = p.parse_args(args=input_options)

		return options, arguments



	# regular expression to match the whole of the final cell from a .castep file
	dotcastep_latt_RE = re.compile("""\sBFGS:\sFinal\sConfiguration:\s*\n
	            =+\s*\n\s*\n\s+\-+\s*\n\s+Unit\sCell\s*\n\s+\-+\s*\n
		    \s+Real\sLattice\(A\)\s+Reciprocal\sLattice\(1/A\)\s*\n
	            \s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s*\n
	            \s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s*\n
	            \s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s*\n""",
	          re.VERBOSE)
	# Start of the 'final configuration'
	dotcastep_infinal_RE = re.compile("BFGS: Final Configuration:")
	# Once inside final configuation, this should only match a line with atoms
	dotcastep_atomline_RE = re.compile("x\s+(\w+)\s+\d+\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+x")
	def parse_dotcastep(seedname):
		dotCastep = open(seedname+".castep","r")
		# Find the lattice
		latticeblock = dotcastep_latt_RE.findall(dotCastep.read())[-1] # Get the last block - handle concat restarts
		lattice = []
		lattice.append([float(latticeblock[0]), float(latticeblock[1]), float(latticeblock[2])])
		lattice.append([float(latticeblock[6]), float(latticeblock[7]), float(latticeblock[8])])
		lattice.append([float(latticeblock[12]), float(latticeblock[13]), float(latticeblock[14])])
		# rewind and search for and final atomic positions (these will be absent if e.g. they are all on symmetry positions)
		dotCastep.seek(0)
		in_atoms = False
		atoms = []
		for line in dotCastep:
			if (in_atoms and dotcastep_atomline_RE.search(line)):
				atom_line = dotcastep_atomline_RE.search(line)
				atoms.append([atom_line.group(1), float(atom_line.group(2)), float(atom_line.group(3)), float(atom_line.group(4))])
			elif ((not in_atoms) and (dotcastep_infinal_RE.search(line))):
				in_atoms = True
			
                dotCastep.close()
		
		return (lattice, atoms)

	# Regular expressions to match a lattice block in a castep .cell file. Note that these
	# can be of the form %block lattice_abc or %block lattice_cart and are case insensitive
	dotcell_lattice_start_RE = re.compile("^\s*%BLOCK\s+LATTICE_(?:CART|ABC)",re.IGNORECASE)
	dotcell_lattice_end_RE = re.compile("^\s*%ENDBLOCK\s+LATTICE_(?:CART|ABC)",re.IGNORECASE)
	dotcell_atoms_start_RE = re.compile("^\s*%BLOCK\s+POSITIONS_(?:FRAC|ABS)", re.IGNORECASE)
	dotcell_atoms_end_RE = re.compile("^\s*%ENDBLOCK\s+POSITIONS_(?:FRAC|ABS)", re.IGNORECASE)
	def produce_dotcell(seedname, filename, newlattice, atoms):
		"""
		produce_dotcell: reads <seedname>.cell (CASTEP cell file
		and writes a new .cell file to <filename> replacing the 
		lattice block with a new crystalographic lattice <newlattice> 
		(which should be supplied as a list of three lists, each with 
		three elements). Also adds command to fix cell during optimization.
		"""
		in_lattice = False
		in_atoms = False
		have_atoms = (atoms != []) # If we have an empty list, no atoms were optimized so just leave them in the .cell file.
		inputfile = open(seedname+".cell", "r")
		outputfile = open(filename, "w")
		for line in inputfile:
			if (dotcell_lattice_end_RE.search(line) and in_lattice):
				in_lattice = False
			elif (dotcell_lattice_start_RE.search(line) and not in_lattice):
				outputfile.write("%block LATTICE_CART\n")
				outputfile.write(str(defcell[0][0]) + " " + str(defcell[0][1]) + " " + str(defcell[0][2]) + "\n")
				outputfile.write(str(defcell[1][0]) + " " + str(defcell[1][1]) + " " + str(defcell[1][2]) + "\n")
				outputfile.write(str(defcell[2][0]) + " " + str(defcell[2][1]) + " " + str(defcell[2][2]) + "\n")
				outputfile.write("%endblock LATTICE_CART\n")
				outputfile.write("FIX_ALL_CELL true\n")
				in_lattice = True
			elif (dotcell_atoms_end_RE.search(line) and in_atoms and have_atoms):
				in_atoms = False
			elif ((dotcell_atoms_start_RE.search(line)) and (not in_atoms) and have_atoms):
				outputfile.write("%block POSITIONS_FRAC\n")
				for atom in atoms:
					outputfile.write("  " + atom[0] + "  " + str(atom[1]) + "  " + str(atom[2]) + "  " + str(atom[3]) + "\n")
				outputfile.write("%endblock POSITIONS_FRAC\n")
				in_atoms = True
			elif(not (in_lattice or in_atoms)):
				outputfile.write(line)
		inputfile.close
		outputfile.close
			
	options, arguments = get_options()
	seedname = arguments[0]
	
	(cell,atoms) = parse_dotcastep(seedname)

	cijdat = open(seedname+".cijdat","w")
	print "\nWriting strain data to ", seedname+".cijdat\n"

	# Not sure why the lattice types are enumerated like this, but this is how .cijdat does it...
	latticeTypes = {0:"Unknown", 1:"Triclinic", 2:"Monoclinic", 3:"Orthorhombic", \
	                4:"Tetragonal", 5:"Cubic", 6:"Trigonal-low", 7:"Trigonal-high/Hexagonal"}
	maxstrain = options.strain
	numsteps = options.numsteps
	latticeCode = options.lattice<commit_msg>Put crystal on a useful cart ref frame<commit_after>#!/usr/bin/env python
"""
generate_strain.py

Generate strained castep input files and 
.cijdat files for elastic constants calculation
for analysis with elastics.py

Copyright (c) 2010 <NAME> (<EMAIL>)
All rights reserved.
"""

import sys
import os
import optparse
import re
import math

def GetStrainPatterns(code):
	"""
	Given a code number for the crystal symmetry, 
	returns a list of strain patterns needed for
	the calculation of the elastic constants tensor.
	Each pattern is a 6 element list, the subscript
	reflects the strain in IRE notation

	Supported Strain Patterns
	-------------------------

	5 Cubic: e1+e4
	7 Hexagonal: e3 and e1+e4
	7 Trigonal-High (32, 3m, -3m): e1 and e3+e4
	6 Trigonal-Low (3, -3): e1 and e3+e4
	4 Tetragonal: e1+e4 and e3+e6
	3 Orthorhombic: e1+e4 and e2+e5 and e3+e6
	2 Monoclinic: e1+e4 and e3+e6 and e2 and e5
	1 Triclinic: e1 to e6 separately
	0 Unknown...
	"""

	if (code == 1):
		pattern = [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
		           [0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
		           [0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
		           [0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
		           [0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
		           [0.0, 0.0, 0.0, 0.0, 0.0, 1.0]]
	elif (code == 2):
		pattern = [[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
		           [0.0, 0.0, 1.0, 0.0, 0.0, 1.0],
		           [0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
		           [0.0, 0.0, 0.0, 0.0, 1.0, 0.0]]
	elif (code == 3):
		pattern = [[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
		           [0.0, 1.0, 0.0, 0.0, 1.0, 0.0],
		           [0.0, 0.0, 1.0, 0.0, 0.0, 1.0]]
	elif (code == 4):
		pattern = [[1.0, 0.0, 0.0, 1.0, 0.0, 0.0],
		           [0.0, 0.0, 1.0, 0.0, 0.0, 1.0]]
	elif (code == 5):
		pattern = [[1.0, 0.0, 0.0, 1.0, 0.0, 0.0]]
	elif (code == 6):
		pattern = [[1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
		           [0.0, 0.0, 1.0, 1.0, 0.0, 0.0]]
	elif (code == 7):
		# NB is this correct for hex and hig trig? - see missmatch above/
		# I suspect I have to rotate lattice for trig high?
		pattern = [[0.0, 0.0, 1.0, 0.0, 0.0, 0.0],
		           [1.0, 0.0, 0.0, 1.0, 0.0, 0.0]]

	return pattern	

def main(input_options, libmode=False):

	# deal with options
	def get_options():
		if not libmode:
			p = optparse.OptionParser()
			p.add_option('--debug', '-d', action='store_true', \
			              help="Debug mode (output to stdout rather than file)")
			p.add_option('--num_steps', '-n', action='store', type='int', dest="numsteps")
			p.add_option('--strain_mag', '-s', action='store', type='float', dest="strain")
			p.add_option('--lattice_type', '-l', action='store', type='int', dest="lattice")
			options,arguments = p.parse_args(args=input_options)

		return options, arguments



	# regular expression to match the whole of the final cell from a .castep file
	dotcastep_latt_RE = re.compile("""\sBFGS:\sFinal\sConfiguration:\s*\n
	            =+\s*\n\s*\n\s+\-+\s*\n\s+Unit\sCell\s*\n\s+\-+\s*\n
		    \s+Real\sLattice\(A\)\s+Reciprocal\sLattice\(1/A\)\s*\n
	            \s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s*\n
	            \s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s*\n
	            \s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s*\n""",
	          re.VERBOSE)
	# Start of the 'final configuration'
	dotcastep_infinal_RE = re.compile("BFGS: Final Configuration:")
	# Once inside final configuation, this should only match a line with atoms
	dotcastep_atomline_RE = re.compile("x\s+(\w+)\s+\d+\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+([\+\-]?\d+.\d+)\s+x")
	def parse_dotcastep(seedname):
		dotCastep = open(seedname+".castep","r")
		# Find the lattice
		latticeblock = dotcastep_latt_RE.findall(dotCastep.read())[-1] # Get the last block - handle concat restarts
		lattice = []
		lattice.append([float(latticeblock[0]), float(latticeblock[1]), float(latticeblock[2])])
		lattice.append([float(latticeblock[6]), float(latticeblock[7]), float(latticeblock[8])])
		lattice.append([float(latticeblock[12]), float(latticeblock[13]), float(latticeblock[14])])
		# rewind and search for and final atomic positions (these will be absent if e.g. they are all on symmetry positions)
		dotCastep.seek(0)
		in_atoms = False
		atoms = []
		for line in dotCastep:
			if (in_atoms and dotcastep_atomline_RE.search(line)):
				atom_line = dotcastep_atomline_RE.search(line)
				atoms.append([atom_line.group(1), float(atom_line.group(2)), float(atom_line.group(3)), float(atom_line.group(4))])
			elif ((not in_atoms) and (dotcastep_infinal_RE.search(line))):
				in_atoms = True
			
                dotCastep.close()
		print lattice
		a, b, c, al, be, ga = cellCART2cellABC(lattice)
		print a, b, c, al, be, ga
		lattice = cellABC2cellCART(a, b, c, al, be, ga)
		print lattice
		
		return (lattice, atoms)

	def cellABC2cellCART (a, b, c, alp, bet, gam):
		import numpy as S
		# Get lattice vetors on cart frame from a, b, c and angles
		# For monoclinic, b // Y and c // Z
		if (alp == 90.0):
      			cosa = 0.0;
		else:
      			cosa = S.cos(S.radians(alp))
		if (bet == 90.0):
     			cosb = 0.0
		else:
			cosb = S.cos(S.radians(bet))
		if (gam == 90.0):
			sing = 1.0
			cosg = 0.0
		else:
			sing = S.sin(S.radians(gam))
			cosg = S.cos(S.radians(gam))
		rv21 = 0.0
		rv31 = 0.0
		rv32 = 0.0
		rv11 = a
		rv12 = b*cosg
		rv22 = b*sing
		rv13 = c*cosb
		rv23 = c*(cosa - cosg*cosb)/sing
		trm1 = rv23/c
		rv33 = c*S.sqrt(1.0 - cosb**2 - trm1**2)
		return [[rv11, rv12, rv13], [rv21, rv22, rv23], [rv31, rv32, rv33]]

	def cellCART2cellABC (lattice):
		import numpy as S
		# Does not care about orentation...
		a = S.sqrt(lattice[0][0]**2 + lattice[0][1]**2 + lattice[0][2]**2)
		b = S.sqrt(lattice[1][0]**2 + lattice[1][1]**2 + lattice[1][2]**2)
		c = S.sqrt(lattice[2][0]**2 + lattice[2][1]**2 + lattice[2][2]**2)
		gam = S.arccos(S.dot(lattice[0],lattice[1]) / (a*b))
		bet = S.arccos(S.dot(lattice[0],lattice[2]) / (a*c))
		alp = S.arccos(S.dot(lattice[1],lattice[2]) / (b*c))
		return a, b, c, S.degrees(alp), S.degrees(bet), S.degrees(gam)



	# Regular expressions to match a lattice block in a castep .cell file. Note that these
	# can be of the form %block lattice_abc or %block lattice_cart and are case insensitive
	dotcell_lattice_start_RE = re.compile("^\s*%BLOCK\s+LATTICE_(?:CART|ABC)",re.IGNORECASE)
	dotcell_lattice_end_RE = re.compile("^\s*%ENDBLOCK\s+LATTICE_(?:CART|ABC)",re.IGNORECASE)
	dotcell_atoms_start_RE = re.compile("^\s*%BLOCK\s+POSITIONS_(?:FRAC|ABS)", re.IGNORECASE)
	dotcell_atoms_end_RE = re.compile("^\s*%ENDBLOCK\s+POSITIONS_(?:FRAC|ABS)", re.IGNORECASE)
	def produce_dotcell(seedname, filename, newlattice, atoms):
		"""
		produce_dotcell: reads <seedname>.cell (CASTEP cell file
		and writes a new .cell file to <filename> replacing the 
		lattice block with a new crystalographic lattice <newlattice> 
		(which should be supplied as a list of three lists, each with 
		three elements). Also adds command to fix cell during optimization.
		"""
		in_lattice = False
		in_atoms = False
		have_atoms = (atoms != []) # If we have an empty list, no atoms were optimized so just leave them in the .cell file.
		inputfile = open(seedname+".cell", "r")
		outputfile = open(filename, "w")
		for line in inputfile:
			if (dotcell_lattice_end_RE.search(line) and in_lattice):
				in_lattice = False
			elif (dotcell_lattice_start_RE.search(line) and not in_lattice):
				outputfile.write("%block LATTICE_CART\n")
				outputfile.write(str(defcell[0][0]) + " " + str(defcell[0][1]) + " " + str(defcell[0][2]) + "\n")
				outputfile.write(str(defcell[1][0]) + " " + str(defcell[1][1]) + " " + str(defcell[1][2]) + "\n")
				outputfile.write(str(defcell[2][0]) + " " + str(defcell[2][1]) + " " + str(defcell[2][2]) + "\n")
				outputfile.write("%endblock LATTICE_CART\n")
				outputfile.write("FIX_ALL_CELL true\n")
				in_lattice = True
			elif (dotcell_atoms_end_RE.search(line) and in_atoms and have_atoms):
				in_atoms = False
			elif ((dotcell_atoms_start_RE.search(line)) and (not in_atoms) and have_atoms):
				outputfile.write("%block POSITIONS_FRAC\n")
				for atom in atoms:
					outputfile.write("  " + atom[0] + "  " + str(atom[1]) + "  " + str(atom[2]) + "  " + str(atom[3]) + "\n")
				outputfile.write("%endblock POSITIONS_FRAC\n")
				in_atoms = True
			elif(not (in_lattice or in_atoms)):
				outputfile.write(line)
		inputfile.close
		outputfile.close
		return()
			
	options, arguments = get_options()
	seedname = arguments[0]
	
	(cell,atoms) = parse_dotcastep(seedname)

	cijdat = open(seedname+".cijdat","w")
	print "\nWriting strain data to ", seedname+".cijdat\n"

	# Not sure why the lattice types are enumerated like this, but this is how .cijdat does it...
	latticeTypes = {0:"Unknown", 1:"Triclinic", 2:"Monoclinic", 3:"Orthorhombic", \
	                4:"Tetragonal", 5:"Cubic", 6:"Trigonal-low", 7:"Trigonal-high/Hexagonal"}
	maxstrain = options.strain
	numsteps = options.numsteps
	latticeCode = options.lattice
{'fraction_non_alphanumeric': 0.20447379380093997, 'fraction_numerical': 0.04735558331846035, 'mean_word_length': 3.1080156402737047, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1955320', 'n_tokens_mistral': 7804, 'n_tokens_neox': 6984, 'n_words': 1580}