import future.utils
import glob
import json as simplejson
import os
import os.path
import random
import shutil
import stat
from .viparr_plugins_util import * # noqa F403
simplejson.encoder.ESCAPE_DCT['/'] = '/'
[docs]def merge_rules(file1, file2):
# file1 = dest
# file2 = src
if not os.path.exists(file1):
shutil.copy2(file2, file1) # src, dest
else:
rules1 = parse_json(file1)
rules2 = parse_json(file2)
if rules1['vdw_comb_rule'] != rules2['vdw_comb_rule']:
raise AssertionError('vdw_comb_rule not compatible')
if rules1['es_scale'] != rules2['es_scale']:
raise AssertionError('es_scale not compatible')
if rules1['lj_scale'] != rules2['lj_scale']:
raise AssertionError('lj_scale not compatible')
# info
rules1['info'].extend(rules2.get('info', []))
# plugins
# assumes "atoms" is last in both lists
d = dict(rules1['plugins'])
for x in rules2['plugins']:
val = d.get(x[0])
if val == x[1]:
pass
elif val == None:
# this is a new key
rules1['plugins'].insert(0, x)
else:
raise AssertionError('plugins not compatible')
# write rules file
os.remove(file1)
print(simplejson.dumps(rules1, indent=3), file=open(file1, 'w'))
[docs]def merge_lists(file1, file2):
'''Load two files, concatenate the two lists, and store result in file1.
If file1 does not exist initially, then a simple copy is made.
file1 = generally the built-in force field
file2 = takes precedence in case of conflicts; generally the new force field
'''
if not os.path.exists(file1):
shutil.copy2(file2, file1) # src, dest
else:
list1 = parse_json(file1)
list2 = parse_json(file2)
# this method of determining number of keys may not work for new files
num = len([x for x in list1[0] if type(x) == type('string')])
d = dict() # dictionary of what is in list2
for item in list2:
d[tuple(item[:num])] = item
# traverse list1 and incrementally create new list
newlist = list()
for item in list1:
val = d.get(tuple(item[:num]))
if val == None:
# this is a new key in list1 that is not in list2
newlist.append(item)
elif val == item:
#print 'Ignoring identical parameters:', val
pass
else:
print(item, 'is being overridden by', val)
# now add items that are in list2
newlist.extend(list2)
# write list
os.remove(file1)
print(simplejson.dumps(newlist, indent=3), file=open(file1, 'w'))
[docs]def update_templates(templates, filename):
t = parse_json(filename)
for key in future.utils.viewkeys(t):
if key in templates:
print('Template <%s>: template in <%s> taking precedence' %
(key, filename))
templates.update(t)
# same as above, but fail on any identical template names
[docs]def update_templates_strict(templates, filename):
t = parse_json(filename)
for key in future.utils.viewkeys(t):
if key in templates:
raise AssertionError(
'Template <%s>: template in <%s> is a duplicate' %
(key, filename))
templates.update(t)
[docs]def merge(dir1, dir2):
'''Merge force fields in directories dir1 and dir2 and return name of
directory with merged force field.
dir1 = generally the built-in force field
dir2 = takes precedence in case of conflicts; generally the new force field
'''
if not os.path.isdir(dir1):
raise AssertionError('Not a directory: %s' % dir1)
if not os.path.isdir(dir2):
raise AssertionError('Not a directory: %s' % dir2)
# make destdir; assumes you don't have a ridiculous number of files in your /tmp
while True:
destdir = str(random.randrange(10000, 100000))
destdir = os.path.join('/tmp', 'viparr_' + destdir)
if not os.path.exists(destdir):
break
# copy all files from dir1 to destdir, but not the templates
shutil.copytree(dir1, destdir)
# following chmod is needed in case dir1 and thus destdir are not writable
# (but this may not work on Windows)
os.chmod(destdir, stat.S_IRWXU)
for filename in glob.glob(os.path.join(destdir, 'templates*')):
os.remove(filename)
# loop over templates in dir1
templates = {}
for filename in glob.glob(os.path.join(dir1, 'templates*')):
update_templates(templates, filename)
# loop over all files in dir2
# and decide what to do with each one
for filename in os.listdir(dir2):
print('*****Merging:', filename)
if filename == 'rules':
merge_rules(os.path.join(destdir, filename),
os.path.join(dir2, filename))
elif filename == 'cmap':
# if cmap already exists, we're going to override it
if os.path.exists(os.path.join(destdir, filename)):
os.remove(os.path.join(destdir, filename))
print('Warning: using cmap file in <%s> and not in <%s>' %
(dir2, dir1))
shutil.copy2(os.path.join(dir2, filename),
os.path.join(destdir, filename))
elif filename.startswith('templates'):
update_templates(templates, os.path.join(dir2, filename))
else:
merge_lists(os.path.join(destdir, filename),
os.path.join(dir2, filename))
# write the templates file
print(simplejson.dumps(templates, indent=3),
file=open(os.path.join(destdir, 'templates'), 'w'))
return destdir