Source code for kvazar.core.structure

# -*- 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_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