# -*- coding: utf-8 -*-
import sys
import numpy
import logging
from copy import deepcopy
try:
from lib import libnanostruct
except:
print "Didn't find libnanostruct!"
N_LINKS = 10
EPS = 1e-6
[docs]class Structure(object):
r"""
* atom_names - Array of names of atoms, i.e. 'Carbon', 'Hydrogen' and etc.
* atom_types - Numbers, that are correspondind to atoms names
* n_atoms - Number of atoms in structure
* n_links - Number of maximum possible links for atom. The second dimension of links matrix
* coords - Array of coordinates of each atom in structure
* velocity - Array of atoms velocities in form :math:`\left[V_x, V_y, V_z\right]`
* tension - Array of tensions for all atoms
* charges - Array of electric charge on all atoms
* links - Array of links of each atom with all others (one-sided links)
* num_links - Number of links for each atom
* full_links - Array of links for all atoms (two-sided bonds)
* num_full_links - Array of numbers of two-sided links for each atom
"""
def __init__(self, struct_data, auto_calculate = True):
li_par_1 = {'coords', 'links', 'num_links'}
ndarray_t = "<type 'numpy.ndarray'>"
self.n_atoms = struct_data['n_atoms']
for param in li_par_1:
if type(struct_data[param]) == ndarray_t:
setattr(self, param, struct_data[param])
else:
setattr(self, param, numpy.array(struct_data[param]))
if "n_links" in struct_data:
self.n_links = struct_data["n_links"]
if self.n_links != N_LINKS:
logging.warning("There is an original n_links in file instead of default (10). It may cause some problems!")
else:
self.n_links = N_LINKS
li_par_2 = {"atom_types", "atom_names", "charges", "limits", "lim", "group_names", "group_ids", "mols", "tension"}
for param in li_par_2:
if param in struct_data:
setattr(self, param, struct_data[param])
if 'el_field' in struct_data:
self.el_field = struct_data["el_field"]
else:
self.el_field = numpy.zeros(3, numpy.float)
if 'velocity' in struct_data:
self.velocity = struct_data["velocity"]
else:
self.velocity = numpy.zeros_like(self.coords)
if "full_links" and "num_full_links" in struct_data:
self.full_links, self.num_full_links = struct_data["full_links"], struct_data["num_full_links"]
else:
self.full_links, self.num_full_links = self.get_full_links(self.links)
if auto_calculate:
self.update_links_params()
def update_links_params(self):
self.neighbours, self.num_neighbours, self.n_neighbours = self.get_neighbours(self.full_links)
self.tors_seq, self.tors_seq_i = self.get_tors_sequences(self.full_links)
[docs] def get_full_links(self, links):
"""
It makes 2D array of reciprocal links between atoms.
"""
n = self.n_atoms
m = self.n_links
links_new = numpy.ones((n, m), numpy.int32) * -1
last = numpy.zeros(n, numpy.int32)
if sys.platform[0] == 'l':
self.full_links = links_new
self.num_full_links = last
libnanostruct.fill_full_links(self.__dict__)
else:
for i in xrange(n):
for j in xrange(len(links[i])):
ln = links[i, j]
if ln != -1:
links_new[i, last[i]] = ln
last[i] += 1
links_new[ln, last[ln]] = i
last[ln] += 1
return links_new, last
[docs] def get_neighbours(self, full_links):
"""
It returns array of neighbours of all atoms, array of numbers of all atoms, full number of atoms
"""
nb = []
ln = 0
for i in xrange(len(full_links)):
nb.append([i])
for j in full_links[i]:
if j != -1:
nb[i].append(j)
for k in full_links[j]:
if k != -1 and k != i:
nb[i].append(k)
ln = max(ln, len(nb[i]))
ln = ln if ln > 0 else 1
neighbours = numpy.ones((len(full_links), ln), numpy.int32)
neighbours *= -1
num_neighbours = numpy.zeros(len(full_links), numpy.int32)
for i, I in enumerate(nb):
for j, val in enumerate(I):
neighbours[i, j] = val
num_neighbours[i] += 1
return neighbours, num_neighbours, ln
[docs] def update_neighbours(self):
"""
It updates neighbours matrix, matrix of number of neighbours
"""
from lib import libnanostruct
max_num_neighbours = self.num_full_links.max() ** 2 + 1
self.neighbours = numpy.ones((self.n_atoms, max_num_neighbours), numpy.int32) * -1
self.num_neighbours = numpy.zeros(self.n_atoms, numpy.int32)
libnanostruct.fill_neighbours(self.__dict__)
[docs] def get_tors_sequences(self, full_links):
"""
It finds sequences of atoms, that form dihedral angles
"""
seq = []
used = [False for i in xrange(len(full_links))]
for i in xrange(len(full_links)):
used[i] = True
self.tors_dfs(full_links, i, used, [i], seq)
seq_i = []
for i in xrange(len(full_links)):
seq_i.append([])
for s in seq:
for i in s:
seq_i[i].append(s)
return seq, seq_i
[docs] def tors_dfs(self, full_links, pos, used, cur_way, result):
"""
Depth-first search to find dihedral angles
"""
for i in full_links[pos]:
if i != -1:
if i not in cur_way:
if len(cur_way) == 3:
if not used[i]:
result.append(cur_way + [int(i)])
else:
self.tors_dfs(full_links, i, used, cur_way + [int(i)], result)
def __iadd__(self, mol):
n = len(self.coords)
m = len(mol.coords)
self.n_atoms = n + m
self.coords = numpy.concatenate((self.coords, mol.coords))
self.links = numpy.concatenate((self.links, mol.links))
def shift(x):
if x > -1:
x += n
return x
vshift = numpy.vectorize(shift)
links = self.links[n:]
nx, ny = links.shape[0], links.shape[1]
links = links.reshape(nx * ny)
links = vshift(links)
links = links.reshape(nx, ny)
self.links[n:] = links
self.num_links = numpy.concatenate((self.num_links, mol.num_links))
li_par_1 = set(['velocity', 'tension'])
for param in li_par_1:
if hasattr(self, param):
if hasattr(mol, param):
spam = numpy.concatenate((getattr(self, param), getattr(mol, param)))
setattr(self, param, spam)
else:
spam = numpy.concatenate((getattr(self, param), numpy.zeros((mol.n_atoms, getattr(self, param).shape[1]))))
setattr(self, param, spam)
linear = ['charges']
for param in linear:
if hasattr(self, param):
if hasattr(mol, param):
spam = numpy.concatenate((getattr(self, param), getattr(mol, param)))
setattr(self, param, spam)
else:
spam = numpy.concatenate((getattr(self, param), numpy.zeros(mol.n_atoms)))
setattr(self, param, spam)
li_par_2 = {'atom_types', 'atom_names'}
for param in li_par_2:
if hasattr(self, param) and hasattr(mol, param):
setattr(self, param, numpy.concatenate((getattr(self, param), getattr(mol, param))))
if hasattr(self, 'full_links') and hasattr(mol, 'full_links'):
self.full_links = numpy.concatenate((self.full_links, mol.full_links))
for i in xrange(n, m + n):
for j in xrange(self.n_links):
if self.full_links[i, j] != -1:
self.full_links[i, j] += n
else:
break
if hasattr(self, 'num_full_links'):
if hasattr(mol, 'num_full_links'):
self.num_full_links = numpy.concatenate((self.num_full_links, mol.num_full_links))
else:
self.num_full_links = numpy.concatenate((self.num_full_links, numpy.zeros(mol.n_atoms, numpy.float)))
return self
def __add__(self, mol):
from copy import deepcopy
c = deepcopy(self)
c += mol
return c
def __getslice__(self, a, b):
struct_data = {}
struct_data['n_atoms'] = n = b - a + 1
struct_data['coords'] = self.coords[a:b+1]
struct_data['atom_names'] = self.atom_names[a:b+1]
li_par = set(['velocity','tension','charges','atom_types'])
for param in li_par:
if hasattr(self, param):
struct_data[param] = getattr(self, param)[a:b+1]
new_links = numpy.zeros_like(self.links[a:b+1])-1
new_num_links = numpy.zeros(n, numpy.int32)
for i in xrange(a,b+1):
for j in xrange(self.num_links[i]):
k = self.links[i][j]
if k >= a and k <= b:
new_links[i - a][new_num_links[i - a]] = k - a
new_num_links[i - a] += 1
struct_data['links'] = new_links
struct_data['num_links'] = new_num_links
struct_data['full_links'], struct_data['num_full_links'] = self.get_full_links(self.links)
return Structure(struct_data)
[docs] def getsubstr(self, selected):
"""
It forms new structure from selected atoms. For example:
Some_structure.getsubstr([(0,20),(30,45)]) - it returns new object of
class Structure, that is made from 0-20 and 30-45 atoms of initial structure.
"""
struct_data = {}
struct_data['n_atoms'] = int(sum([b - a + 1 for a, b in selected]))
struct_data['n_links'] = self.n_links
li_par_1 = {'coords', 'atom_names'}
for param in li_par_1:
struct_data[param] = numpy.concatenate([getattr(self, param)[a : b + 1] for a,b in selected])
li_par_2 = {'velocity', 'tension', 'charges', 'atom_types'}
for param in li_par_2:
if hasattr(self, param):
struct_data[param] = numpy.concatenate([getattr(self, param)[a : b + 1] for a,b in selected])
new_links = numpy.zeros((struct_data['n_atoms'], struct_data['n_links']), numpy.int32) - 1
new_num_links = numpy.zeros(struct_data['n_atoms'], numpy.int32)
ind = 0
for a,b in selected:
for i in xrange(a, b + 1):
for j in xrange(self.num_links[i]):
l = self.links[i,j]
shift = 0
for c,d in selected:
if l >= c and l <= d:
new_links[ind, new_num_links[ind]] = l - c + shift
new_num_links[ind] += 1
break
shift += d - c + 1
ind += 1
struct_data['links'] = new_links
struct_data['num_links'] = new_num_links
return Structure(struct_data)
def __delitem__(self, selected_del):
#TODO do not use numpy.array in selected list
selected = []
if selected_del[0][0] != 0:
selected.append(0)
for a,b in selected_del:
if a != 0:
selected.append(a - 1)
if b != len(self.coords) - 1:
selected.append(b + 1)
if selected_del[-1][1] != len(self.coords)-1:
selected.append(len(self.coords)-1)
selected = numpy.array(selected).reshape(len(selected)/2, 2)
selected_l = []
for x in selected:
selected_l.append(tuple(x))
a = self.getsubstr(selected_l)
for param in a.__dict__:
setattr(self, param, getattr(a, param))
[docs] def compare_topology_with(self, b):
"""
It only compares number of atoms, links and identity of atom types in
two structures.
"""
a = self
li_par = {'coords', 'links', 'num_links','atoms_types'}
for x in li_par:
a_p = getattr(a, x)
b_p = getattr(b, x)
l_a = len(a_p)
l_b = len(b_p)
if l_a != l_b or numpy.linalg.norm(a_p-b_p) > EPS:
return False
return True
[docs] def compare_with(self, b):
"""
If the structures are the same, it returns True.
"""
a = self
li_par = {'n_atoms', 'n_links', 'atom_types'}
for x in li_par:
a_g = getattr(a,x)
b_g = getattr(b,x)
try:
if a_g != b_g:
return False
except:
if len(a_g) != len(b_g) or numpy.linalg.norm(a_g - b_g) > EPS:
return False
return True
[docs] def copy(self):
"""
Returns new structure which is equal to the initial
"""
A = deepcopy(self)
return A