diff --git a/modules/chempy/cpv.py b/modules/chempy/cpv.py index 33e8fae54..175a9bf8e 100644 --- a/modules/chempy/cpv.py +++ b/modules/chempy/cpv.py @@ -21,43 +21,46 @@ # # TODO: documentation! +from typing import Sequence import math import random -import copy RSMALL4 = 0.0001 #------------------------------------------------------------------------------ -def get_null(): +def get_null() -> list[float]: return [0.0,0.0,0.0] #------------------------------------------------------------------------------ -def get_identity(): +def get_identity() -> list[list[float]]: return [[1.0,0.0,0.0],[0.0,1.0,0.0],[0.0,0.0,1.0]] #------------------------------------------------------------------------------ -def distance_sq(v1, v2): +def distance_sq(v1: Sequence[float], v2: Sequence[float]) -> float: d0 = v2[0] - v1[0] d1 = v2[1] - v1[1] d2 = v2[2] - v1[2] return (d0*d0) + (d1*d1) + (d2*d2) #------------------------------------------------------------------------------ -def distance(v1, v2): +def distance(v1: Sequence[float], v2: Sequence[float]) -> float: d0 = v2[0] - v1[0] d1 = v2[1] - v1[1] d2 = v2[2] - v1[2] return math.sqrt((d0*d0) + (d1*d1) + (d2*d2)) #------------------------------------------------------------------------------ -def length(v): +def length(v: Sequence[float]) -> float: return math.sqrt(v[0]*v[0] + v[1]*v[1] + v[2]*v[2]) #------------------------------------------------------------------------------ -def random_displacement(v,radius): - r_vect = lambda r=random.random:[r()-0.5,r()-0.5,r()-0.5] +def _random_vector() -> list[float]: + r = random.random + return [r()-0.5,r()-0.5,r()-0.5] + +def random_displacement(v: Sequence[float], radius: float) -> list[float]: while 1: - vect = r_vect() + vect = _random_vector() v_len = length(vect) if (v_len<=0.5): break; @@ -68,92 +71,85 @@ def random_displacement(v,radius): return v #------------------------------------------------------------------------------ -def random_sphere(v,radius): - r_vect = lambda r=random.random:[r()-0.5,r()-0.5,r()-0.5] +def random_sphere(v: Sequence[float], radius: float) -> list[float]: while 1: - vect = r_vect() + vect = _random_vector() v_len = length(vect) if (v_len<=0.5) and (v_len!=0.0): break; return add(v,scale([vect[0], vect[1], vect[2]],2*radius/v_len)) #------------------------------------------------------------------------------ -def random_vector(): - r_vect = lambda r=random.random:[r()-0.5,r()-0.5,r()-0.5] +def random_vector() -> list[float]: while 1: - vect = r_vect() + vect = _random_vector() if length(vect)<=0.5: break; return scale([vect[0], vect[1], vect[2]],2.0) #------------------------------------------------------------------------------ -def add(v1,v2): +def add(v1: Sequence[float], v2: Sequence[float]) -> list[float]: return [v1[0]+v2[0],v1[1]+v2[1],v1[2]+v2[2]] #------------------------------------------------------------------------------ -def average(v1,v2): +def average(v1: Sequence[float], v2: Sequence[float]) -> list[float]: return [(v1[0]+v2[0])/2.0,(v1[1]+v2[1])/2.0,(v1[2]+v2[2])/2.0] #------------------------------------------------------------------------------ -def scale(v,factor): +def scale(v: Sequence[float], factor: float) -> list[float]: return [v[0]*factor,v[1]*factor,v[2]*factor] #------------------------------------------------------------------------------ -def negate(v): +def negate(v: Sequence[float]) -> list[float]: return [-v[0],-v[1],-v[2]] #------------------------------------------------------------------------------ -def sub(v1,v2): +def reverse(v: Sequence[float]) -> list[float]: + return [ -v[0], -v[1], -v[2] ] + +#------------------------------------------------------------------------------ +def sub(v1: Sequence[float], v2: Sequence[float]) -> list[float]: return [v1[0]-v2[0],v1[1]-v2[1],v1[2]-v2[2]] #------------------------------------------------------------------------------ -def dot_product(v1,v2): +def dot_product(v1: Sequence[float], v2: Sequence[float]) -> float: return v1[0]*v2[0]+v1[1]*v2[1]+v1[2]*v2[2] #------------------------------------------------------------------------------ -def cross_product(v1,v2): +def cross_product(v1: Sequence[float], v2: Sequence[float]) -> list[float]: return [(v1[1]*v2[2]) - (v1[2]*v2[1]), (v1[2]*v2[0]) - (v1[0]*v2[2]), (v1[0]*v2[1]) - (v1[1]*v2[0])] #------------------------------------------------------------------------------ -def transform(m,v): +def transform(m: Sequence[Sequence[float]], v: Sequence[float]) -> list[float]: return [m[0][0]*v[0] + m[0][1]*v[1] + m[0][2]*v[2], m[1][0]*v[0] + m[1][1]*v[1] + m[1][2]*v[2], m[2][0]*v[0] + m[2][1]*v[1] + m[2][2]*v[2]] #------------------------------------------------------------------------------ -def inverse_transform(m,v): +def inverse_transform(m: Sequence[Sequence[float]], v: Sequence[float]) -> list[float]: return [m[0][0]*v[0] + m[1][0]*v[1] + m[2][0]*v[2], m[0][1]*v[0] + m[1][1]*v[1] + m[2][1]*v[2], m[0][2]*v[0] + m[1][2]*v[1] + m[2][2]*v[2]] #------------------------------------------------------------------------------ -def multiply(m1,m2): # HAVEN'T YET VERIFIED THAT THIS CONFORMS TO STANDARD DEFT - return [[m1[0][0]*m2[0][0] + m1[0][1]*m2[1][0] + m1[0][2]*m2[2][0], - m1[1][0]*m2[0][0] + m1[1][1]*m2[1][0] + m1[1][2]*m2[2][0], - m1[2][0]*m2[0][0] + m1[2][1]*m2[1][0] + m1[2][2]*m2[2][0]], - [m1[0][0]*m2[0][1] + m1[0][1]*m2[1][1] + m1[0][2]*m2[2][1], - m1[1][0]*m2[0][1] + m1[1][1]*m2[1][1] + m1[1][2]*m2[2][1], - m1[2][0]*m2[0][1] + m1[2][1]*m2[1][1] + m1[2][2]*m2[2][1]], - [m1[0][0]*m2[0][2] + m1[0][1]*m2[1][2] + m1[0][2]*m2[2][2], - m1[1][0]*m2[0][2] + m1[1][1]*m2[1][2] + m1[1][2]*m2[2][2], - m1[2][0]*m2[0][2] + m1[2][1]*m2[1][2] + m1[2][2]*m2[2][2]]] +def multiply(m1: Sequence[Sequence[float]], m2: Sequence[Sequence[float]]) -> list[list[float]]: + # HAVEN'T YET VERIFIED THAT THIS CONFORMS TO STANDARD DEFT + # upd: no, it's not(fixed) -#------------------------------------------------------------------------------ -def transpose(m1): - return [[m1[0][0], - m1[1][0], - m1[2][0]], - [m1[0][1], - m1[1][1], - m1[2][1]], - [m1[0][2], - m1[1][2], - m1[2][2]]] + return [[m1[0][0]*m2[0][0] + m1[0][1]*m2[1][0] + m1[0][2]*m2[2][0], + m1[0][0]*m2[0][1] + m1[0][1]*m2[1][1] + m1[0][2]*m2[2][1], + m1[0][0]*m2[0][2] + m1[0][1]*m2[1][2] + m1[0][2]*m2[2][2]], + [m1[1][0]*m2[0][0] + m1[1][1]*m2[1][0] + m1[1][2]*m2[2][0], + m1[1][0]*m2[0][1] + m1[1][1]*m2[1][1] + m1[1][2]*m2[2][1], + m1[1][0]*m2[0][2] + m1[1][1]*m2[1][2] + m1[1][2]*m2[2][2]], + [m1[2][0]*m2[0][0] + m1[2][1]*m2[1][0] + m1[2][2]*m2[2][0], + m1[2][0]*m2[0][1] + m1[2][1]*m2[1][1] + m1[2][2]*m2[2][1], + m1[2][0]*m2[0][2] + m1[2][1]*m2[1][2] + m1[2][2]*m2[2][2]]] #------------------------------------------------------------------------------ -def get_system2(x,y): +def get_system2(x: Sequence[float],y: Sequence[float]) -> list[list[float]]: z = cross_product(x,y) z = normalize(z) y = cross_product(z,x); @@ -162,24 +158,24 @@ def get_system2(x,y): return [x,y,z] #------------------------------------------------------------------------------ -def scale_system(s,factor): +def scale_system(s: Sequence[Sequence[float]], factor: float) -> list[list[float]]: r = [] for a in s: r.append([a[0]*factor,a[1]*factor,a[2]*factor]) return r #------------------------------------------------------------------------------ -def transpose(m): +def transpose(m: Sequence[Sequence[float]]) -> list[list[float]]: return [[m[0][0], m[1][0], m[2][0]], [m[0][1], m[1][1], m[2][1]], [m[0][2], m[1][2], m[2][2]]] #------------------------------------------------------------------------------ -def transform_about_point(m,v,p): +def transform_about_point(m: Sequence[Sequence[float]], v: Sequence[float], p: Sequence[float]) -> list[float]: return add(transform(m,sub(v,p)),p) #------------------------------------------------------------------------------ -def get_angle(v1,v2): # v1,v2 must be unit vectors +def get_angle(v1: Sequence[float], v2: Sequence[float]) -> float: # v1,v2 must be unit vectors denom = (math.sqrt(((v1[0]*v1[0]) + (v1[1]*v1[1]) + (v1[2]*v1[2]))) * math.sqrt(((v2[0]*v2[0]) + (v2[1]*v2[1]) + (v2[2]*v2[2])))) if denom>1e-10: @@ -190,7 +186,7 @@ def get_angle(v1,v2): # v1,v2 must be unit vectors return result #------------------------------------------------------------------------------ -def get_angle_formed_by(p1,p2,p3): # angle formed by three positions in space +def get_angle_formed_by(p1: Sequence[float], p2: Sequence[float], p3: Sequence[float]) -> float: # angle formed by three positions in space # based on code submitted by Paul Sherwood r1 = distance(p1,p2) @@ -207,17 +203,17 @@ def get_angle_formed_by(p1,p2,p3): # angle formed by three positions in space return theta; #------------------------------------------------------------------------------ -def project(v,n): +def project(v: Sequence[float], n: Sequence[float]) -> list[float]: dot = v[0]*n[0] + v[1]*n[1] + v[2]*n[2] return [ dot * n[0], dot * n[1], dot * n[2] ] #------------------------------------------------------------------------------ -def remove_component(v, n): +def remove_component(v: Sequence[float], n: Sequence[float]) -> list[float]: dot = v[0]*n[0] + v[1]*n[1] + v[2]*n[2] return [v[0] - dot * n[0], v[1] - dot * n[1], v[2] - dot * n[2]] #------------------------------------------------------------------------------ -def normalize(v): +def normalize(v: Sequence[float]) -> list[float]: vlen = math.sqrt((v[0]*v[0]) + (v[1]*v[1]) + (v[2]*v[2])) if vlen>RSMALL4: return [v[0]/vlen,v[1]/vlen,v[2]/vlen] @@ -225,11 +221,7 @@ def normalize(v): return get_null() #------------------------------------------------------------------------------ -def reverse(v): - return [ -v[0], -v[1], -v[2] ] - -#------------------------------------------------------------------------------ -def normalize_failsafe(v): +def normalize_failsafe(v: Sequence[float]) -> list[float]: vlen = math.sqrt((v[0]*v[0]) + (v[1]*v[1]) + (v[2]*v[2])) if vlen>RSMALL4: return [v[0]/vlen,v[1]/vlen,v[2]/vlen] @@ -237,7 +229,7 @@ def normalize_failsafe(v): return [1.0,0.0,0.0] #------------------------------------------------------------------------------ -def rotation_matrix(angle,axis): +def rotation_matrix(angle: float, axis: Sequence[float]) -> list[list[float]]: x=axis[0] y=axis[1] @@ -271,7 +263,7 @@ def rotation_matrix(angle,axis): [ (one_c * zx) - ys, (one_c * yz) + xs, (one_c * zz) + c ]] #------------------------------------------------------------------------------ -def transform_array(rot_mtx,vec_array): +def transform_array(rot_mtx: Sequence[Sequence[float]], vec_array: Sequence[Sequence[float]]) -> list[list[float]]: '''transform_array( matrix, vector_array ) -> vector_array @@ -280,7 +272,7 @@ def transform_array(rot_mtx,vec_array): return [transform(rot_mtx, x) for x in vec_array] #------------------------------------------------------------------------------ -def translate_array(trans_vec,vec_array): +def translate_array(trans_vec: Sequence[float], vec_array: Sequence[Sequence[float]]) -> list[list[float]]: '''translate_array(trans_vec,vec_array) -> vec_array @@ -291,9 +283,9 @@ def translate_array(trans_vec,vec_array): return [add(trans_vec, x) for x in vec_array] #------------------------------------------------------------------------------ -def fit_apply(fit_result,vec_array): +def fit_apply(fit_result: tuple[Sequence[float], Sequence[float], Sequence[Sequence[float]], float], vec_array: Sequence[Sequence[float]]) -> list[list[float]]: '''fit_apply(fir_result,vec_array) -> vec_array - + Applies a fit result to an array of vectors ''' @@ -301,7 +293,7 @@ def fit_apply(fit_result,vec_array): return [add(t1, transform(m, add(mt2, x))) for x in vec_array] #------------------------------------------------------------------------------ -def fit(target_array, source_array): +def fit(target_array: Sequence[Sequence[float]], source_array: Sequence[Sequence[float]]) -> tuple[list[float], list[float], list[list[float]], float]: '''fit(target_array, source_array) -> (t1, t2, rot_mtx, rmsd) [fit_result] @@ -370,7 +362,6 @@ def fit(target_array, source_array): iters = 0 while (iters < maxiter): iters = iters + 1 - ix = (iters-1)%ndim iy = iters%ndim iz = (iters+1)%ndim sig = corr_mtx[iz][iy] - corr_mtx[iy][iz] @@ -403,5 +394,4 @@ def fit(target_array, source_array): return(t1, t2, rot_mtx, rmsd) # Too many iterations; something wrong. - print ("Error: Too many iterations in RMS fit.") - raise ValueError + raise ValueError("Error: Too many iterations in RMS fit.") diff --git a/modules/chempy/models.py b/modules/chempy/models.py index 6aa0d5929..ebffc56ae 100644 --- a/modules/chempy/models.py +++ b/modules/chempy/models.py @@ -1,705 +1,612 @@ -#A* ------------------------------------------------------------------- -#B* This file contains source code for the PyMOL computer program -#C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific. -#D* ------------------------------------------------------------------- -#E* It is unlawful to modify or remove this copyright notice. -#F* ------------------------------------------------------------------- -#G* Please see the accompanying LICENSE file for further information. -#H* ------------------------------------------------------------------- -#I* Additional authors of this source file include: -#-* -#-* -#-* -#Z* ------------------------------------------------------------------- - -import chempy +# A* ------------------------------------------------------------------- +# B* This file contains source code for the PyMOL computer program +# C* copyright 1998-2000 by Warren Lyford Delano of DeLano Scientific. +# D* ------------------------------------------------------------------- +# E* It is unlawful to modify or remove this copyright notice. +# F* ------------------------------------------------------------------- +# G* Please see the accompanying LICENSE file for further information. +# H* ------------------------------------------------------------------- +# I* Additional authors of this source file include: +# -* +# -* +# -* +# Z* ------------------------------------------------------------------- + +import abc import copy -from .cpv import * +from typing import Generic, Optional, TypeVar, Union -class Base: +import numpy as np - @property - def nAtom(self): - return len(self.atom) +from chempy import Atom, Bond, Molecule, cpv, feedback - @property - def nBond(self): - return len(self.bond) +T = TypeVar("T", Bond, list[Bond]) -#------------------------------------------------------------------------------ - def update_index(self): - if chempy.feedback['verbose']: - print(" "+str(self.__class__)+": updating indexes...") - c = 0 - self.index = {} - idx = self.index - for a in self.atom: - idx[id(a)] = c - c = c + 1 -#------------------------------------------------------------------------------ - def get_residues(self): - list = [] - if self.nAtom: - last = self.atom[0] - c = 0 - start = 0 - for a in self.atom: - if not a.in_same_residue(last): - list.append((start,c)) - start = c - last = a - c = c + 1 - if (c-start>1): - list.append((start,c)) - return list - -#------------------------------------------------------------------------------ - - def get_coord_list(self): - lst = [] - for a in self.atom: - lst.append(a.coord) - return lst - -#------------------------------------------------------------------------------ - def get_mass(self): - sm = 0.0 - for a in self.atom: - sm = sm + a.get_mass() - return sm - -#------------------------------------------------------------------------------ - def get_nuclear_charges(self): - '''Return the sum of nuclear charges of all atoms in a molecule.''' - sm = 0 - for a in self.atom: - sm = sm + a.get_number() - return sm +class Base(abc.ABC, Generic[T]): -#------------------------------------------------------------------------------ - def list(self): - for a in self.atom: - print(a.symbol, a.name, a.coord) - for a in self.bond: - print(a.index) - -#------------------------------------------------------------------------------ - def get_implicit_mass(self): - # mass calculation for implicit models - - valence = [0]*len(self.atom) - - for a in self.bond: - v = 1.5 if a.order == 4 else a.order - valence[a.index[0]] += v - valence[a.index[1]] += v - - h_count = sum(a.get_free_valence(v) - for (a, v) in zip(self.atom, valence)) - - hydrogen = chempy.Atom() - hydrogen.symbol='H' - return self.get_mass()+hydrogen.get_mass()*h_count - -#------------------------------------------------------------------------------ - def assign_names(self,preserve=1): - dct = {} - cnt = {} - if preserve: - for a in self.atom: - if a.has('name'): - dct[a['name']] = 1 - else: - for a in self.atom: - if hasattr(a,'name'): - del a.name - for a in self.atom: - if not a.has('name'): - if a.symbol not in cnt: - c = 1 - else: - c = cnt[a.symbol] - while 1: - nam = a.symbol+str(c) - c = c + 1 - if nam not in dct: - break - cnt[a.symbol]=c - a.name=nam - dct[nam]=1 -#------------------------------------------------------------------------------ -#------------------------------------------------------------------------------ - -class Indexed(Base): - -#------------------------------------------------------------------------------ def __init__(self): - self.reset() + self.index: Optional[dict[int, int]] = None + self.molecule = Molecule() + self.atom: list[Atom] = [] + self.bond: list[T] = [] -#------------------------------------------------------------------------------ - def reset(self): - self.index = None - self.molecule = chempy.Molecule() - self.atom = [] - self.bond = [] - -#------------------------------------------------------------------------------ - def get_min_max(self): - if len(self.atom): - mn = copy.deepcopy(self.atom[0].coord) - mx = copy.deepcopy(self.atom[0].coord) - for a in self.atom: - ac = a.coord - if mn[0]>ac[0]: mn[0]=ac[0] - if mn[1]>ac[1]: mn[1]=ac[1] - if mn[2]>ac[2]: mn[2]=ac[2] - if mx[0] list[list[Bond]]: + raise NotImplementedError() -#-------------------------------------------------------------------------------- - def delete_atom(self,index): - if chempy.feedback['atoms']: - print(" "+str(self.__class__)+": deleting atom %d." % index) + def _handle_new_atom(self) -> None: + """In case of Connected class we need to add an empty list to slef.bond""" + pass + + def add_atom(self, atom: Atom) -> int: + self.send_feedback(key="atoms", message=f": adding atom atom {atom.name}") + self.atom.append(atom) - nAtom=self.nAtom + self._handle_new_atom() -# update index if it exists + index = len(self.atom) - 1 if self.index: - idx = self.index - for k in list(idx.keys()): - if idx[k] > index: - idx[k] = idx[k] - 1 - del idx[id(self.atom[index])] + self.index[id(atom)] = index + return index + + def delete_atom(self, index: int) -> None: + self.send_feedback(key="atoms", message=f": deleting atom {index}.") -# delete atom + # update index if it exists + if self.index is not None: + for key, value in self.index.items(): + if value > index: + self.index[key] = value - 1 + del self.index[id(self.atom[index])] + + # delete atom del self.atom[index] -# delete bonds associated with this atom + # delete bonds associated with this atom + + for bond_group in self._bond_groups: + templist = [i for i, bond in enumerate(bond_group) if index in bond.index] + for i, j in enumerate(templist): + del bond_group[j - i] + + # re-index bond table + for bond_group in self._bond_groups: + for bond in bond_group: + if bond.index[0] > index: + bond.index[0] -= 1 + if bond.index[1] > index: + bond.index[1] -= 1 + + def insert_atom(self, index: int, atom: Atom) -> None: + self.send_feedback( + key="atoms", message=f": inserting atom {atom.name} before {index}." + ) + self.atom.insert(index, atom) + + # re-index bond table + for bond_pair in self._bond_groups: + for bond in bond_pair: + if bond.index[0] >= index: + bond.index[0] = bond.index[0] + 1 + if bond.index[1] >= index: + bond.index[1] = bond.index[1] + 1 + + # update index if it exists + if self.index is not None: + + for key, value in self.index.items(): + if value >= index: + self.index[key] = value + 1 - nBond = len(self.bond) - templist = [] - for i in range(nBond): - if index in self.bond[i].index: - templist.append(i) - for i in range(len(templist)): - j = templist[i] - i - del self.bond[j] + self.index[id(atom)] = index -# re-index bond table - for b in self.bond: - if b.index[0] > index: - b.index[0] = b.index[0] - 1 - if b.index[1] > index: - b.index[1] = b.index[1] - 1 + def reset(self): + self.index = None + self.molecule = Molecule() + self.atom.clear() + self.bond.clear() -#-------------------------------------------------------------------------------- - def delete_list(self,list): # delete a list of indexed atoms + def update_index(self) -> None: + self.send_feedback(key="verbose", message="updating indexes...") + self.index = {id(atom): index for index, atom in enumerate(self.atom)} - if chempy.feedback['atoms']: - print(" "+str(self.__class__)+": deleting atoms %s." % str(list)) + def get_residues(self) -> list[tuple[int, int]]: + result = [] - nAtom=self.nAtom + if not self.atom: + return result - lrev = copy.deepcopy(list) - lrev.sort() - lrev.reverse() + last = self.atom[0] + counter = 0 + start = 0 - # generate cross-reference tables + for atom in self.atom: + if not atom.in_same_residue(last): + result.append((start, counter)) + start = counter + last = atom + counter += 1 + if counter - start > 1: + result.append((start, counter)) - o2n = {} # old to new - if len(lrev): - nxt = lrev.pop() - else: - nxt = None - shft = 0 - for i in range(nAtom): - if i == nxt: - o2n[i]=-1 - if len(lrev): - nxt = lrev.pop() - else: - nxt = None - shft = shft - 1 + return result + + def get_coord_list(self) -> list[tuple[float, float, float]]: + return [atom.coord for atom in self.atom] + + def get_mass(self) -> float: + return sum(atom.get_mass() for atom in self.atom) + + def get_nuclear_charges(self): + """Return the sum of nuclear charges of all atoms in a molecule.""" + return sum(atom.get_number() for atom in self.atom) + + def assign_names(self, preserve: bool = True) -> None: + names = set() + counter = {} + + for atom in self.atom: + if preserve: + if atom.has("name"): + names.add(atom.name) else: - ixs = i + shft - o2n[i] = ixs - - if shft: - - # delete atoms - - new_atom = [] - for i in range(nAtom): - if o2n[i]>=0: - new_atom.append(self.atom[i]) - self.atom = new_atom - - # delete bonds - - new_bond = [] - for b in self.bond: - b0 = b.index[0] - b1 = b.index[1] - if (o2n[b0]>=0) and (o2n[b1]>=0): - b.index[0] = o2n[b0] - b.index[1] = o2n[b1] - new_bond.append(b) - self.bond = new_bond - - # update index if it exists - if self.index: - self.index = {} - idx = self.index - i = 0 - for a in self.atom: - idx[id(a)] = i - i = i + 1 - -#------------------------------------------------------------------------------ - def insert_atom(self,index,atom): - if chempy.feedback['atoms']: - print(" "+str(self.__class__)+': inserting atom "%s" before %d.' % ( - atom.name,index)) - - nAtom=self.nAtom - self.atom.insert(index,atom) - -# re-index bond table - for b in self.bond: - if b.index[0] >= index: - b.index[0] = b.index[0] + 1 - if b.index[1] >= index: - b.index[1] = b.index[1] + 1 + if hasattr(atom, "name"): + del atom.name -# update index if it exists - if self.index: - idx = self.index - for k in list(idx.keys()): - if idx[k] >= index: - idx[k] = idx[k] + 1 - idx[id(atom)] = index - -#------------------------------------------------------------------------------ - def index_atom(self,atom): - c = 0 - id_at = id(atom) - for a in self.atom: - if id(a)==id_at: - return c - c = c + 1 - return -1 + for atom in self.atom: + + if atom.has("name"): + continue + + c = counter[atom.symbol] if atom.symbol in counter else 1 + name = f"{atom.symbol}{c}" + + while name in names: + name = f"{atom.symbol}{c}" + c += 1 + + counter[atom.symbol] = c + atom.name = name + + names.add(name) + + def send_feedback(self, key: str, message: str) -> None: + if feedback[key]: + print(f" {self.__class__}): {message}") + + +class Indexed(Base[Bond]): + + @property + def _bond_groups(self) -> list[list[Bond]]: + return [self.bond] + + def get_min_max(self) -> Union[list[list[float]], list[tuple[float, float, float]]]: + + if not self.atom: + return [[0.0, 0.0, 0.0], [0.0, 0.0, 0.0]] + + coords = np.array([atom.coord for atom in self.atom]) + + min_coord = np.min(coords, axis=0).tolist() + max_coord = np.max(coords, axis=0).tolist() + + return [min_coord, max_coord] + + def merge(self, other: "Indexed") -> None: + """steals atom objects from 'other' and resets 'other'""" + self.send_feedback(key="actions", message="merging models...") + + self.atom.extend(other.atom) + + for bond in other.bond: + bond.index[0] += len(self.atom) + bond.index[1] += len(self.atom) + self.bond.append(bond) + + other.reset() -#------------------------------------------------------------------------------ - def add_atom(self,atom): - if chempy.feedback['atoms']: - print(" "+str(self.__class__)+': adding atom "%s".' % atom.name) - self.atom.append(atom) - index = self.nAtom - 1 if self.index: - self.index[id(atom)] = index - return index + self.update_index() + + def delete_list(self, indexes: list[int]) -> None: + """delete a list of indexed atoms.""" + + if not indexes: + return + + self.send_feedback(key="atoms", message=f": deleting atoms {indexes}") + + indexes_copy = copy.deepcopy(indexes) + indexes_copy.sort() + indexes_copy.reverse() -#------------------------------------------------------------------------------ - def add_bond(self,bond): - if chempy.feedback['bonds']: - print(" "+str(self.__class__)+": adding bond (%d,%d)." % \ - (bond.index[0],bond.index[1])) + # generate cross-reference tables + shft = 0 + old_to_new = {} + next = indexes_copy.pop() + + for i in range(len(self.atom)): + if i == next: + old_to_new[i] = -1 + next = indexes_copy.pop() if indexes_copy else None + shft = shft - 1 + else: + old_to_new[i] = i + shft + + if shft == 0: + return + + # delete atoms + new_atoms = [] + for i in range(len(self.atom)): + if old_to_new[i] >= 0: + new_atoms.append(self.atom[i]) + self.atom = new_atoms + + # delete bonds + new_bonds: list[Bond] = [] + for bond in self.bond: + b0 = bond.index[0] + b1 = bond.index[1] + if old_to_new[b0] >= 0 and old_to_new[b1] >= 0: + bond.index[0] = old_to_new[b0] + bond.index[1] = old_to_new[b1] + new_bonds.append(bond) + self.bond = new_bonds + + # update index if it exists + if self.index is not None: + self.index = {self.index[id(atom)]: i for i, atom in enumerate(self.atom)} + + def index_atom(self, atom: Atom): + return self.atom.index(atom) if atom in self.atom else -1 + + def add_bond(self, bond: Bond) -> None: + self.send_feedback( + key="bonds", message=f": adding bond ({bond.index[0]}, {bond.index[1]})" + ) self.bond.append(bond) -#------------------------------------------------------------------------------ - def remove_bond(self,index): - if chempy.feedback['bonds']: - print(" "+str(self.__class__)+": removing bond %d." % index) - nBond=len(self.Bond) + def remove_bond(self, index: int) -> None: + self.send_feedback(key="bonds", message=f": removing bond {index}") del self.bond[index] + def convert_to_connected(self) -> "Connected": + self.send_feedback(key="verbose", message=": converting to connected model...") -#------------------------------------------------------------------------------ - def convert_to_connected(self): - if chempy.feedback['verbose']: - print(" "+str(self.__class__)+": converting to connected model...") model = Connected() model.molecule = self.molecule model.atom = self.atom model.bond = [] model.index = None - for a in model.atom: + + for _ in model.atom: model.bond.append([]) - for b in self.bond: - model.bond[b.index[0]].append(b) # note two refs to same object - model.bond[b.index[1]].append(b) # note two refs to same object + + for bond in self.bond: + model.bond[bond.index[0]].append(bond) # note two refs to same object + model.bond[bond.index[1]].append(bond) # note two refs to same object + self.reset() + return model -#------------------------------------------------------------------------------ - def from_molobj(self,molobj): + + def from_molobj(self, molobj: Molecule) -> None: self.reset() - mol = self.molecule + if len(molobj.title): - mol.title = molobj.title + self.molecule.title = molobj.title + if len(molobj.comments): - mol.comments = molobj.comments - mol.chiral = molobj.chiral - mol.dim_code = molobj.dimcode + self.molecule.comments = molobj.comments + + self.molecule.chiral = molobj.chiral + self.molecule.dim_code = molobj.dim_code + for a in molobj.atom: - at = chempy.Atom() + at = Atom() at.symbol = a.symbol at.name = a.name - if a.resn != chempy.Atom.defaults['resn']: + + if a.resn != Atom.defaults["resn"]: at.resn = a.resn - if a.resn_code != chempy.Atom.defaults['resn_code']: + + if a.resn_code != Atom.defaults["resn_code"]: at.resn_code = a.resn_code + at.resi = a.resi at.resi_number = a.resi_number at.b = a.b at.q = a.q at.alt = a.alt at.hetatm = a.hetatm - if a.segi != chempy.Atom.defaults['segi']: + + if a.segi != Atom.defaults["segi"]: at.segi = a.segi - if a.chain != chempy.Atom.defaults['chain']: + + if a.chain != Atom.defaults["chain"]: at.chain = a.chain + at.color_code = a.color_code at.coord = a.coord at.formal_charge = a.formal_charge at.partial_charge = a.partial_charge + if a.numeric_type != -99: at.numeric_type = a.numeric_type - if a.text_type != 'UNKNOWN': + + if a.text_type != "UNKNOWN": at.text_type = a.text_type + at.stereo = a.stereo - if hasattr(a,'flags'): + + if hasattr(a, "flags"): at.flags = a.flags - if hasattr(a,'vdw'): + if hasattr(a, "vdw"): at.vdw = a.vdw + self.atom.append(at) + for b in molobj.bond: - bnd = chempy.Bond() - bnd.index = [b.atom[0],b.atom[1]] + bnd = Bond() + bnd.index = [b.atom[0], b.atom[1]] bnd.order = b.order bnd.stereo = b.stereo self.bond.append(bnd) -#------------------------------------------------------------------------------ + def sort(self): - if chempy.feedback['verbose']: - print(" "+__name__+": sorting...") - if not self.index: + self.send_feedback(key="verbose", message=f": sorting...") + + if self.index is None: self.update_index() + old_index = self.index self.atom.sort() self.update_index() + xref = {} new_index = self.index + for a in list(new_index.keys()): xref[old_index[a]] = new_index[a] + for b in self.bond: b.index[0] = xref[b.index[0]] b.index[1] = xref[b.index[1]] - del old_index - del xref - -#------------------------------------------------------------------------------ - def get_internal_tuples(self): - # generates raw atom sets needed to construct an internal coordinate - # description of the molecule - model = self - # get a connected version too - cmodel = copy.deepcopy(model).convert_to_connected() - center = [0.0,0.0,0.0] - nAtom = model.nAtom - to_go = nAtom + + def get_internal_tuples(self) -> list: + """ + generates raw atom sets needed to construct an internal coordinate + description of the molecule + get a connected version too + """ + model_copy = copy.deepcopy(self).convert_to_connected() + center = [0.0, 0.0, 0.0] + to_go = len(self.atom) done = {} - if to_go<3: - z_set = [(0),(1,0)] - else: - # get center of molecule - for a in model.atom: - center = add(center,a.coord) - center = scale(center,1.0/nAtom) - # find most central multivalent atom - min_a = -1 - c = 0 - for a in model.atom: - if len(cmodel.bond[c])>1: # must have at least two neighbors - d = distance(a.coord,center) - if min_a < 0: - min_d = d - min_a = c - elif d1: - nxt = nbr + + if to_go < 3: + return [(0), (1, 0)] + + # get center of molecule + for atom in self.atom: + center = cpv.add(center, atom.coord) + center = cpv.scale(center, 1.0 / len(self.atom)) + + # find most central multivalent atom + min_a = -1 + for counter, atom in enumerate(self.atom): + if len(model_copy.bond[counter]) > 1: # must have at least two neighbors + d = cpv.distance(atom.coord, center) + if min_a < 0: + min_d = d + min_a = counter + elif d < min_d: + # ^ not exists + min_d = d + min_a = counter + + # make this our first atom + fst = min_a + z_set = [(fst,)] + done[fst] = 1 + to_go = to_go - 1 + + # for the second atom, try to choose different multivalent neighbor + nxt = -1 + for bond in model_copy.bond[fst]: + neighbor = bond.index[0] + if neighbor == fst: + neighbor = bond.index[1] + if len(model_copy.bond[neighbor]) > 1: + nxt = neighbor + break + + # safety, choose any neighbor + if nxt < 0: + neighbor = bond.index[0] + if neighbor == fst: + neighbor = bond.index[1] + nxt = neighbor + z_set.append((nxt, fst)) + done[nxt] = 1 + to_go = to_go - 1 + + # for the third atom, choose a different multivalent neighbor + trd = -1 + for bond in model_copy.bond[fst]: + neighbor = bond.index[0] + if neighbor == fst: + neighbor = bond.index[1] + if len(model_copy.bond[neighbor]) > 1: + if neighbor not in done: + trd = neighbor break - # safety, choose any neighbor - if nxt<0: - nbr = b.index[0] - if nbr == fst: - nbr = b.index[1] - nxt = nbr - z_set.append((nxt,fst)) - done[nxt] = 1 - to_go = to_go - 1 - # for the third atom, choose a different multivalent neighbor - trd = -1 - for b in cmodel.bond[fst]: - nbr = b.index[0] - if nbr == fst: - nbr = b.index[1] - if len(cmodel.bond[nbr])>1: - if nbr not in done: - trd = nbr - break - # safety, choose any unchosen neighbor - if trd<0: - for b in cmodel.bond[fst]: - nbr = b.index[0] - if nbr == fst: - nbr = b.index[1] - if nbr not in done: - trd = nbr - break - z_set.append((trd,fst,nxt)) - done[trd] = 1 - result = 1 - to_go = to_go - 1 - if to_go: - # now find all torsions in the molecule - tors = {} - for b in model.bond: # use bond as center of torsion - a1 = b.index[0] - a2 = b.index[1] - for c in cmodel.bond[a1]: - a0 = c.index[0] - if a0 not in (a1,a2): # outside atom - for d in cmodel.bond[a2]: - a3 = d.index[0] - if a3 not in (a0,a1,a2): # outside atom - if a0 < a3: - to = (a0,a1,a2,a3) - else: - to = (a3,a2,a1,a0) - tors[to] = 1 - a3 = d.index[1] - if a3 not in (a0,a1,a2): # outside atom - if a0 < a3: - to = (a0,a1,a2,a3) - else: - to = (a3,a2,a1,a0) - tors[to] = 1 - a0 = c.index[1] - if a0 not in (a1,a2): # outside atom - for d in cmodel.bond[a2]: - a3 = d.index[0] - if a3 not in (a0,a1,a2): # outside atom - if a0 < a3: - to = (a0,a1,a2,a3) - else: - to = (a3,a2,a1,a0) - tors[to] = 1 - a3 = d.index[1] - if a3 not in (a0,a1,a2): # outside atom - if a0 < a3: - to = (a0,a1,a2,a3) - else: - to = (a3,a2,a1,a0) - tors[to] = 1 - if len(tors): - # choose remaining atoms based on existing atoms using torsion - while to_go: - for tor in list(tors.keys()): - a0 = tor[0] - a1 = tor[1] - a2 = tor[2] - a3 = tor[3] - dh0 = a0 in done - dh1 = a1 in done - dh2 = a2 in done - dh3 = a3 in done - if ( (not dh0) and dh1 and dh2 and dh3 ): - z_set.append((a0,a1,a2,a3)) - done[a0] = 1 - to_go = to_go - 1 - elif ( dh0 and dh1 and dh2 and (not dh3) ): - z_set.append((a3,a2,a1,a0)) - done[a3] = 1 - to_go = to_go - 1 - else: # for molecules with no torsions (dichloromethane, etc.) - # we have to generate torsions which include one virtual - # bond - for b in cmodel.bond[fst]: - nbr = b.index[0] - if nbr in [fst,nxt,trd]: - nbr = b.index[1] - if nbr not in done: - z_set.append((nbr,trd,fst,nxt)) - to_go = to_go - 1 - done[nbr] = 1 - return z_set -#------------------------------------------------------------------------------ -#------------------------------------------------------------------------------ -#------------------------------------------------------------------------------ -class Connected(Base): -#------------------------------------------------------------------------------ - def __init__(self): - self.reset() + # safety, choose any unchosen neighbor + if trd < 0: + for bond in model_copy.bond[fst]: + neighbor = bond.index[0] + if neighbor == fst: + neighbor = bond.index[1] + if neighbor not in done: + trd = neighbor + break -#------------------------------------------------------------------------------ - def reset(self): - self.index = None - self.molecule = chempy.Molecule() - self.atom = [] - self.bond = [] - -#------------------------------------------------------------------------------ - def convert_to_indexed(self): - if chempy.feedback['verbose']: - print(" "+str(self.__class__)+": converting to indexed model...") - indexed = Indexed() - indexed.atom = self.atom - indexed.molecule = self.molecule - c = 0 - for a in self.bond: - for b in a: - if b.index[0] == c: - indexed.bond.append(b) - c = c + 1 - self.reset() - return indexed + z_set.append((trd, fst, nxt)) + done[trd] = 1 + to_go = to_go - 1 + if to_go: + # now find all torsions in the molecule + tors = {} + for bond in self.bond: # use bond as center of torsion + a1 = bond.index[0] + a2 = bond.index[1] + for counter in model_copy.bond[a1]: + a0 = counter.index[0] + if a0 not in (a1, a2): # outside atom + for d in model_copy.bond[a2]: + a3 = d.index[0] + if a3 not in (a0, a1, a2): # outside atom + if a0 < a3: + to = (a0, a1, a2, a3) + else: + to = (a3, a2, a1, a0) + tors[to] = 1 + a3 = d.index[1] + if a3 not in (a0, a1, a2): # outside atom + if a0 < a3: + to = (a0, a1, a2, a3) + else: + to = (a3, a2, a1, a0) + tors[to] = 1 + a0 = counter.index[1] + if a0 not in (a1, a2): # outside atom + for d in model_copy.bond[a2]: + a3 = d.index[0] + if a3 not in (a0, a1, a2): # outside atom + if a0 < a3: + to = (a0, a1, a2, a3) + else: + to = (a3, a2, a1, a0) + tors[to] = 1 + a3 = d.index[1] + if a3 not in (a0, a1, a2): # outside atom + if a0 < a3: + to = (a0, a1, a2, a3) + else: + to = (a3, a2, a1, a0) + tors[to] = 1 + if len(tors): + # choose remaining atoms based on existing atoms using torsion + while to_go: + for tor in list(tors.keys()): + a0 = tor[0] + a1 = tor[1] + a2 = tor[2] + a3 = tor[3] + dh0 = a0 in done + dh1 = a1 in done + dh2 = a2 in done + dh3 = a3 in done + if (not dh0) and dh1 and dh2 and dh3: + z_set.append((a0, a1, a2, a3)) + done[a0] = 1 + to_go = to_go - 1 + elif dh0 and dh1 and dh2 and (not dh3): + z_set.append((a3, a2, a1, a0)) + done[a3] = 1 + to_go = to_go - 1 + else: # for molecules with no torsions (dichloromethane, etc.) + # we have to generate torsions which include one virtual + # bond + for bond in model_copy.bond[fst]: + neighbor = bond.index[0] + if neighbor in [fst, nxt, trd]: + neighbor = bond.index[1] + if neighbor not in done: + z_set.append((neighbor, trd, fst, nxt)) + to_go = to_go - 1 + done[neighbor] = 1 + return z_set -#------------------------------------------------------------------------------ - def insert_atom(self,index,atom): - if chempy.feedback['atoms']: - print(" "+str(self.__class__)+': inserting atom "%s" before %d.' % ( - atom.name,index)) + def list(self): + for atom in self.atom: + print(atom.symbol, atom.name, atom.coord) + for bond in self.bond: + print(bond.index) - nAtom=self.nAtom - self.atom.insert(index,atom) + def get_implicit_mass(self): + """mass calculation for implicit models""" -# re-index bond table - for a in self.bonds: - for b in a: - if b.index[0] >= index: - b.index[0] = b.index[0] + 1 - if b.index[1] >= index: - b.index[1] = b.index[1] + 1 + valence = [0] * len(self.atom) -# update index if it exists - if self.index: - idx = self.index - for k in list(idx.keys()): - if idx[k] >= index: - idx[k] = idx[k] + 1 - idx[id(atom)] = index + for bond in self.bond: + v = 1.5 if bond.order == 4 else bond.order + valence[bond.index[0]] += v + valence[bond.index[1]] += v -#------------------------------------------------------------------------------ - def delete_atom(self,index): - if chempy.feedback['atoms']: - print(" "+str(self.__class__)+": deleting atom %d." % index) + h_count = sum(atom.get_free_valence(v) for (atom, v) in zip(self.atom, valence)) - nAtom=self.nAtom + hydrogen = Atom() + hydrogen.symbol = "H" + return self.get_mass() + hydrogen.get_mass() * h_count -# update index if it exists - if self.index: - idx = self.index - for k in list(idx.keys()): - if idx[k] > index: - idx[k] = idx[k] - 1 - del idx[id(self.atom[index])] -# delete atom - del self.atom[index] +class Connected(Base[list[Bond]]): -# delete bonds associated with this atom + @property + def _bond_groups(self) -> list[list[Bond]]: + return self.bond - nBond = len(self.bond) + def _handle_new_atom(self) -> None: + self.bond.append([]) - for a in self.bond: - i = 0 - templist = [] - for b in a: - if index in b.index: - templist.append(i) - i = i + 1 - for i in range(len(templist)): - j = templist[i] - i - del a[j] + def convert_to_indexed(self) -> Indexed: + self.send_feedback(key="actions", message=": converting to indexed model...") -# re-index bond table - for b in self.bond: - if b.index[0] > index: - b.index[0] = b.index[0] - 1 - if b.index[1] > index: - b.index[1] = b.index[1] - 1 - -#------------------------------------------------------------------------------ - def add_atom(self,atom): - if chempy.feedback['atoms']: - print(" "+str(self.__class__)+': adding atom "%s".' % atom.name) - self.atom.append(atom) - self.bond.append([]) - index = self.nAtom - 1 - if self.index: - self.index[id(atom)] = index - return index + indexed = Indexed() + indexed.atom = self.atom + indexed.molecule = self.molecule + for counter, bond_pair in enumerate(self.bond): + for bond in bond_pair: + if bond.index[0] == counter: + indexed.bond.append(bond) + # TODO: remove next line, otherwise indexed model will be empty + # self.reset() + return indexed -#------------------------------------------------------------------------------ def sort(self): - if chempy.feedback['verbose']: - print(" "+__name__+": sorting...") - if not self.index: + self.send_feedback(key="verbose", message=f": sorting...") + + if self.index is None: self.update_index() + old_index = self.index self.atom.sort() self.update_index() + xref = {} new_index = self.index + for a in new_index.keys(): xref[old_index[a]] = new_index[a] - new_bond = [None] * self.nAtom + + new_bond = [None] * len(self.atom) c = 0 tmp_list = [] - for a in self.bond: - for b in a: - if c==b.index[0]: - tmp_list.append(b) - new_bond[xref[c]] = a + for bond_pair in self.bond: + for bond in bond_pair: + if c == bond.index[0]: + tmp_list.append(bond) + new_bond[xref[c]] = bond_pair c = c + 1 + for b in tmp_list: b.index[0] = xref[b.index[0]] b.index[1] = xref[b.index[1]] - del self.bond + self.bond = new_bond - del old_index - del xref diff --git a/testing/tests/chempy/models/test_connected.py b/testing/tests/chempy/models/test_connected.py new file mode 100644 index 000000000..34da2ee7f --- /dev/null +++ b/testing/tests/chempy/models/test_connected.py @@ -0,0 +1,96 @@ +from chempy import Atom, Bond +import pytest + +from chempy.models import Indexed, Connected + + +@pytest.fixture +def connected_model() -> Connected: + connected_model = Connected() + + for i in range(3): + atom = Atom() + atom.symbol = "H" if i == 0 else "He" + atom.chain = "AB" if i == 0 else "BA" + + atom.coord = [float(i), 0.0, 0.0] + connected_model.add_atom(atom) + + for i in range(len(connected_model.atom)): + bond = Bond() + bond.index = [i, i] + bond.order = i + # note two refs to same object + connected_model.bond[bond.index[0]].append(bond) + connected_model.bond[bond.index[1]].append(bond) + return connected_model + + +def test_base_model_methods(connected_model): + assert len(connected_model.atom) == 3 + assert len(connected_model.bond) == 3 + assert connected_model.get_residues() == [(0, 1), (1, 3)] + assert connected_model.get_coord_list() == [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [2.0, 0.0, 0.0], + ] + assert connected_model.get_mass() == 9.013144 + assert connected_model.get_nuclear_charges() == 5 + + assert connected_model.index is None + connected_model.update_index() + assert sorted(list(connected_model.index.values())) == [0, 1, 2] + + +def test_add_atom(connected_model): + new_atom = Atom() + new_atom.symbol = "Ne" + + assert len(connected_model.atom) == 3 + connected_model.update_index() + previos_last_atom = connected_model.atom[0] + + connected_model.add_atom(new_atom) + assert connected_model.atom[-1] is new_atom + assert connected_model.atom[-2] is not previos_last_atom + + assert len(connected_model.atom) == 4 + + +def test_delete_atom(connected_model): + + assert len(connected_model.atom) == 3 + assert len(connected_model.bond) == 3 + assert connected_model.atom[0].symbol == "H" + + connected_model.delete_atom(0) + + assert len(connected_model.atom) == 2 + assert len(connected_model.bond) == 3 + assert connected_model.bond[0] == [] + assert connected_model.atom[0].symbol == "He" + assert connected_model.bond[1][0].order == 1 + + +def test_insert_atom(connected_model): + new_atom = Atom() + new_atom.symbol = "Ne" + + assert len(connected_model.atom) == 3 + previos_second_atom = connected_model.atom[2] + + connected_model.insert_atom(2, new_atom) + + assert connected_model.atom[2] is new_atom + assert previos_second_atom in connected_model.atom + assert connected_model.atom.index(previos_second_atom) == 3 + assert len(connected_model.atom) == 4 + + +def test_convert_to_indexed(connected_model): + indexed_model = connected_model.convert_to_indexed() + assert isinstance(indexed_model, Indexed) + assert len(indexed_model.atom) == 3 + # maybe a bug ? + assert len(indexed_model.bond) == 6 diff --git a/testing/tests/chempy/models/test_indexed.py b/testing/tests/chempy/models/test_indexed.py new file mode 100644 index 000000000..e222a56be --- /dev/null +++ b/testing/tests/chempy/models/test_indexed.py @@ -0,0 +1,188 @@ +from chempy import Atom, Bond +import pytest + +from chempy.models import Indexed, Connected + + +@pytest.fixture +def indexed_model() -> Indexed: + indexed_model = Indexed() + + for i in range(3): + atom = Atom() + atom.symbol = "H" if i == 0 else "He" + atom.chain = "AB" if i == 0 else "BA" + + atom.coord = [float(i), 0.0, 0.0] + indexed_model.add_atom(atom) + + bond = Bond() + bond.index = [i, i] + bond.order = i + indexed_model.add_bond(bond) + + return indexed_model + + +def test_base_model_methods(indexed_model): + assert indexed_model.get_residues() == [(0, 1), (1, 3)] + assert indexed_model.get_coord_list() == [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [2.0, 0.0, 0.0], + ] + assert indexed_model.get_mass() == 9.013144 + assert indexed_model.get_nuclear_charges() == 5 + + assert indexed_model.index is None + indexed_model.update_index() + assert sorted(list(indexed_model.index.values())) == [0, 1, 2] + + +def test_get_implicit_mass(indexed_model): + assert indexed_model.get_implicit_mass() == 10.021084 + + +def test_list(indexed_model, capsys): + indexed_model.list() + captured = capsys.readouterr() + assert captured.out == ( + "H [0.0, 0.0, 0.0]\n" + "He [1.0, 0.0, 0.0]\n" + "He [2.0, 0.0, 0.0]\n" + "[0, 0]\n" + "[1, 1]\n" + "[2, 2]\n" + ) + + +def test_get_min_max(indexed_model): + assert indexed_model.get_min_max() == [[0.0, 0.0, 0.0], [2.0, 0.0, 0.0]] + + +def test_indexed_merge(indexed_model): + other_indexed_model = Indexed() + + atom = Atom() + atom.symbol = "Ne" + atom.chain = "F" + + atom.coord = [0.0, 0.0, 7.0] + other_indexed_model.add_atom(atom) + + bond = Bond() + bond.index = [1, 1] + bond.order = 1 + other_indexed_model.add_bond(bond) + + indexed_model.update_index() + assert sorted(list(indexed_model.index.values())) == [0, 1, 2] + + indexed_model.merge(other_indexed_model) + + assert sorted(list(indexed_model.index.values())) == [0, 1, 2, 3] + assert len(indexed_model.atom) == 4 + assert len(indexed_model.bond) == 4 + + assert other_indexed_model.index is None + assert other_indexed_model.bond == [] + assert other_indexed_model.bond == [] + + +def test_add_atom(indexed_model): + new_atom = Atom() + new_atom.symbol = "Ne" + + assert len(indexed_model.atom) == 3 + indexed_model.update_index() + previos_last_atom = indexed_model.atom[-1] + + indexed_model.add_atom(new_atom) + assert indexed_model.atom[-1] is new_atom + assert indexed_model.atom[-2] is previos_last_atom + + assert len(indexed_model.atom) == 4 + + +def test_delete_atom(indexed_model): + + assert len(indexed_model.atom) == 3 + assert len(indexed_model.bond) == 3 + assert indexed_model.atom[0].symbol == "H" + + indexed_model.delete_atom(0) + + assert len(indexed_model.atom) == 2 + assert len(indexed_model.bond) == 2 + assert indexed_model.atom[0].symbol == "He" + assert indexed_model.bond[0].order == 1 + + +def test_delete_list(indexed_model): + + assert len(indexed_model.atom) == 3 + assert len(indexed_model.bond) == 3 + assert indexed_model.atom[0].symbol == "H" + + indexed_model.delete_list([1, 2]) + + assert len(indexed_model.atom) == 1 + assert len(indexed_model.bond) == 1 + assert indexed_model.atom[0].symbol == "H" + assert indexed_model.bond[0].order == 0 + + +def test_insert_atom(indexed_model): + new_atom = Atom() + new_atom.symbol = "Ne" + + assert len(indexed_model.atom) == 3 + previos_second_atom = indexed_model.atom[2] + + indexed_model.insert_atom(2, new_atom) + + assert indexed_model.atom[2] is new_atom + assert previos_second_atom in indexed_model.atom + assert indexed_model.atom.index(previos_second_atom) == 3 + assert indexed_model.index_atom(previos_second_atom) == 3 + assert len(indexed_model.atom) == 4 + + +def test_index_atom(indexed_model): + new_atom = Atom() + new_atom.symbol = "Ne" + + indexed_model.add_atom(new_atom) + + assert indexed_model.index_atom(new_atom) == 3 + + +def test_add_bond(indexed_model): + assert len(indexed_model.atom) == 3 + assert len(indexed_model.bond) == 3 + assert indexed_model.bond[-1].order == 2 + + new_bond = Bond() + new_bond.index = [42, 42] + new_bond.order = 42 + indexed_model.add_bond(new_bond) + + assert len(indexed_model.atom) == 3 + assert len(indexed_model.bond) == 4 + assert indexed_model.bond[-1].order == 42 + + +def test_convert_to_connected(indexed_model): + connected_model = indexed_model.convert_to_connected() + assert isinstance(connected_model, Connected) + assert len(connected_model.bond) == 3 + assert connected_model.bond[0] is connected_model.bond[0] + # maybe a bug ? + + assert indexed_model.atom == [] + assert indexed_model.bond == [] + assert indexed_model.index == None + + +def test_get_internal_tuples(indexed_model): + assert indexed_model.get_internal_tuples() == [(1,), (1, 1), (-1, 1, 1)] diff --git a/testing/tests/chempy/test_cpv.py b/testing/tests/chempy/test_cpv.py new file mode 100644 index 000000000..37c0b7211 --- /dev/null +++ b/testing/tests/chempy/test_cpv.py @@ -0,0 +1,510 @@ +from unittest.mock import patch +import random + +import pytest +import numpy as np + +from chempy import cpv + + +@pytest.fixture +def vector_a() -> list[float]: + return [3.0, 4.0, 0.0] + + +@pytest.fixture +def vector_b() -> list[float]: + return [6.0, 8.0, 0.0] + + +@pytest.fixture +def matrix_a() -> list[list[float]]: + return [[4, 8, 15], [16, 23, 42], [1984, 1, 10]] + + +@pytest.fixture +def matrix_b() -> list[list[float]]: + return [[5, 1, 3], [1, 1, 1], [1, 2, 1]] + + +def test_get_null(): + assert cpv.get_null() == np.zeros(shape=3, dtype=np.float64).tolist() + assert cpv.get_null() == np.array([0, 0, 0], dtype=np.float64).tolist() + + +def test_get_identity(): + assert cpv.get_identity() == np.identity(3, dtype=np.float64).tolist() + assert ( + cpv.get_identity() + == np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.float64).tolist() + ) + + +def test_distance_sq(vector_a, vector_b): + + assert ( + cpv.distance_sq(vector_a, vector_b) + == np.linalg.norm(np.array(vector_a) - np.array(vector_b)) ** 2 + ) + assert cpv.distance_sq([3, 3, 3], [2, 2, 2]) == 3.0 + + +def test_distance(vector_a, vector_b): + assert cpv.distance(vector_a, vector_b) == np.linalg.norm( + np.array(vector_a) - np.array(vector_b) + ) + assert cpv.distance(vector_a, vector_b) == 5.0 + + +def test_length(vector_a): + assert cpv.length(vector_a) == np.linalg.norm(vector_a) + assert cpv.length(vector_a) == 5.0 + + +def test_random_displacement(vector_a): + radius = 2 + + with ( + patch.object(cpv, "_random_vector", return_value=[0.1, 0.1, 0.1]), + patch("numpy.random.random", return_value=np.array([0.1, 0.1, 0.1])), + patch("random.random", return_value=1), + ): + rand_np_vector = np.random.random(3) + v_len = random.random() * radius / np.linalg.norm(rand_np_vector) + assert ( + cpv.random_displacement(vector_a, radius) + == (np.array(vector_a) + rand_np_vector * v_len).tolist() + ) + assert ( + cpv.random_displacement(vector_a, radius) + == np.array( + [4.1547005383792515, 5.1547005383792515, 1.1547005383792512] + ).tolist() + ) + + +def test_random_sphere(vector_a): + radius = 2 + + with ( + patch.object(cpv, "_random_vector", return_value=[0.1, 0.1, 0.1]), + patch("numpy.random.random", return_value=np.array([0.1, 0.1, 0.1])), + ): + rand_np_vector = np.random.random(3) + assert ( + cpv.random_sphere(vector_a, radius) + == ( + np.array(vector_a) + + rand_np_vector * (2 * radius / np.linalg.norm(rand_np_vector)) + ).tolist() + ) + assert ( + cpv.random_sphere(vector_a, radius) + == np.array( + [ + 5.309401076758503, + 6.309401076758503, + 2.3094010767585025, + ] + ).tolist() + ) + + +def test_random_vector(): + with ( + patch.object(cpv, "_random_vector", return_value=[0.1, 0.1, 0.1]), + patch("numpy.random.random", return_value=np.array([0.1, 0.1, 0.1])), + ): + assert cpv.random_vector() == (np.random.random(3) * 2).tolist() + assert cpv.random_vector() == np.array([0.2, 0.2, 0.2]).tolist() + + +def test_add(vector_a, vector_b): + assert ( + cpv.add(vector_a, vector_b) + == (np.array(vector_a) + np.array(vector_b)).tolist() + ) + assert cpv.add(vector_a, vector_b) == np.array([9.0, 12.0, 0.0]).tolist() + + +def test_sub(vector_a, vector_b): + assert ( + cpv.sub(vector_a, vector_b) + == (np.array(vector_a) - np.array(vector_b)).tolist() + ) + assert cpv.sub(vector_a, vector_b) == (np.array([-3.0, -4.0, 0.0]).tolist()) + + +def test_average(vector_a, vector_b): + assert ( + cpv.average(vector_a, vector_b) + == np.average([np.array(vector_a), np.array(vector_b)], axis=0).tolist() + ) + assert cpv.average(vector_a, vector_b) == np.array([4.5, 6.0, 0.0]).tolist() + + +def test_scale(vector_a): + assert cpv.scale(vector_a, 2) == (np.array(vector_a) * 2).tolist() + assert cpv.scale(vector_a, 2) == np.array([6.0, 8.0, 0.0]).tolist() + + +def test_negate(vector_a): + assert cpv.negate(vector_a) == (np.array(vector_a) * -1).tolist() + assert cpv.negate(vector_a) == np.array([-3.0, -4.0, -0.0]).tolist() + + +def test_dot_product(vector_a, vector_b): + assert cpv.dot_product(vector_a, vector_b) == np.dot( + np.array(vector_a), np.array(vector_b) + ) + assert cpv.dot_product(vector_a, vector_b) == 50.0 + + +def test_cross_product(vector_a, vector_b): + assert ( + cpv.cross_product(vector_a, vector_b) + == np.cross(np.array(vector_a), np.array(vector_b)).tolist() + ) + assert cpv.cross_product(vector_a, vector_b) == np.array([0.0, 0.0, 0.0]).tolist() + + +def test_transform(matrix_a, vector_a): + assert ( + cpv.transform(matrix_a, vector_a) + == np.dot(np.array(matrix_a), np.array(vector_a)).tolist() + ) + assert cpv.transform(matrix_a, vector_a) == np.array([44.0, 140.0, 5956.0]).tolist() + + +def test_inverse_transform(matrix_a, vector_a): + assert ( + cpv.inverse_transform(matrix_a, vector_a) + == np.dot(np.array(matrix_a).T, np.array(vector_a)).tolist() + ) + assert ( + cpv.inverse_transform(matrix_a, vector_a) + == np.array([76.0, 116.0, 213.0]).tolist() + ) + + +def test_multiply(matrix_a, matrix_b): + assert ( + cpv.multiply(matrix_a, matrix_b) + == np.dot(np.array(matrix_a), np.array(matrix_b)).tolist() + ) + assert ( + cpv.multiply(matrix_a, matrix_b) + == np.array([[43, 42, 35], [145, 123, 113], [9931, 2005, 5963]]).tolist() + ) + + +def test_transpose(matrix_a): + assert cpv.transpose(matrix_a) == np.array(matrix_a).T.tolist() + assert ( + cpv.transpose(matrix_a) + == np.array([[4, 16, 1984], [8, 23, 1], [15, 42, 10]]).tolist() + ) + + +def test_get_system2(vector_a, vector_b): + def normalize(v: np.ndarray) -> np.ndarray: + return ( + np.array(v) / vector_len + if (vector_len := np.linalg.norm(np.array(v))) > cpv.RSMALL4 + else np.zeros(shape=3, dtype=np.float64) + ) + + z = normalize(np.cross(np.array(vector_a), np.array(vector_b))) + y = normalize(np.cross(z, np.array(vector_a))) + x = normalize(vector_a) + assert cpv.get_system2(vector_a, vector_b) == np.array([x, y, z]).tolist() + assert ( + cpv.get_system2(vector_a, vector_b) + == np.array([[0.6, 0.8, 0.0], [0.0, 0.0, 0.0], [0.0, 0.0, 0.0]]).tolist() + ) + + +# ------------------------------------------------------------------------------ +def test_scale_system(matrix_a): + assert cpv.scale_system(matrix_a, 1984) == (np.array(matrix_a) * 1984).tolist() + assert ( + cpv.scale_system(matrix_a, 1985) + == np.array( + [[7940, 15880, 29775], [31760, 45655, 83370], [3938240, 1985, 19850]] + ).tolist() + ) + + +def test_transform_about_point(matrix_a, vector_a, vector_b): + assert ( + cpv.transform_about_point(matrix_a, vector_a, vector_b) + == ( + np.dot(np.array(matrix_a), np.array(vector_a) - np.array(vector_b)) + + np.array(vector_b) + ).tolist() + ) + assert ( + cpv.transform_about_point(matrix_a, vector_a, vector_b) + == np.array([-38.0, -132.0, -5956.0]).tolist() + ) + + +def test_get_angle(vector_a, vector_b): + np_result = ( + np.dot(np.array(vector_a), np.array(vector_b)) / denom + if ( + denom := np.linalg.norm(np.array(vector_a)) + * np.linalg.norm(np.array(vector_b)) + ) + > 1e-10 + else 0.0 + ) + + assert cpv.get_angle(vector_a, vector_b) == np.arccos(np_result) + assert cpv.get_angle(vector_a, vector_b) == 0.0 + + +def test_get_angle_formed_by(vector_a, vector_b): + r1 = np.linalg.norm(np.array(vector_a) - np.array(vector_b)) + r2 = np.linalg.norm(np.array(vector_b) - np.array(vector_b)) + r3 = np.linalg.norm(np.array(vector_a) - np.array(vector_b)) + + theta = ( + np.pi + if (r1 + r2 - r3) < 1.0e-10 + else np.arccos((r1**2 + r2**2 - r3**2) / (2.0 * r1 * r2)) + ) + + assert cpv.get_angle_formed_by(vector_a, vector_b, vector_b) == theta + assert cpv.get_angle_formed_by(vector_a, vector_b, vector_b) == np.pi + + +def test_project(vector_a, vector_b): + np_version_dot = np.dot(np.array(vector_a), np.array(vector_b)) + + assert ( + cpv.project(vector_a, vector_b) + == (np.array(vector_b) * np_version_dot).tolist() + ) + assert cpv.project(vector_a, vector_b) == np.array([300.0, 400.0, 0.0]).tolist() + + +def test_remove_component(vector_a, vector_b): + np_version_dot = np.dot(np.array(vector_a), np.array(vector_b)) + + assert ( + cpv.remove_component(vector_a, vector_b) + == (np.array(vector_a) - np.array(vector_b) * np_version_dot).tolist() + ) + assert ( + cpv.remove_component(vector_a, vector_b) + == np.array([-297.0, -396.0, 0.0]).tolist() + ) + + +def test_normalize(vector_a): + np_result = ( + np.array(vector_a) / vector_len + if (vector_len := np.linalg.norm(np.array(vector_a))) > cpv.RSMALL4 + else np.zeros(shape=3, dtype=np.float64) + ) + + assert cpv.normalize(vector_a) == np_result.tolist() + assert cpv.normalize(vector_a) == np.array([0.6, 0.8, 0.0]).tolist() + + +def test_normalize_failsafe(vector_a): + np_result = ( + np.array(vector_a) / vector_len + if (vector_len := np.linalg.norm(np.array(vector_a))) > cpv.RSMALL4 + else np.array([1.0, 0.0, 0.0]) + ) + + assert cpv.normalize_failsafe(vector_a) == np_result.tolist() + assert cpv.normalize_failsafe(vector_a) == np.array([0.6, 0.8, 0.0]).tolist() + + +def test_rotation_matrix(): + angle = 30 + axis = [4, 8, 15] + axis = np.array(axis) + s = np.sin(angle) + c = np.cos(angle) + + mag = np.linalg.norm(axis) + + if abs(mag) < cpv.RSMALL4: + return np.eye(3) + + axis = axis / mag + + # Rodrigues' rotation formula + R = ( + c * np.eye(3) + + (1 - c) * np.outer(axis, axis) + + s + * np.array( + [[0, -axis[2], axis[1]], [axis[2], 0, -axis[0]], [-axis[1], axis[0], 0]] + ) + ) + + assert cpv.rotation_matrix(angle, [4, 8, 15]) == R.tolist() + assert ( + cpv.rotation_matrix(angle, [4, 8, 15]) + == np.array( + [ + [0.1986185869426616, 0.9373521674176533, -0.28621944580745806], + [-0.759883619197343, 0.3317199981078943, 0.5590516327950812], + [0.6189729737205398, 0.1064554230310823, 0.7781643147246124], + ], + ).tolist() + ) + + +def test_transform_array(matrix_a, matrix_b): + + assert ( + cpv.transform_array(matrix_a, matrix_b) + == np.apply_along_axis( + lambda vertex: np.dot(np.array(matrix_a), np.array(vertex)), + 1, + np.array(matrix_b), + ).tolist() + ) + assert ( + cpv.transform_array(matrix_a, matrix_b) + == np.array( + [ + [73, 229, 9951], + [27, 81, 1995], + [35, 104, 1996], + ] + ).tolist() + ) + + +def test_translate_array(vector_a, matrix_a): + assert ( + cpv.translate_array(vector_a, matrix_a) + == np.apply_along_axis( + lambda vertex: np.array(vector_a) + np.array(vertex), 1, np.array(matrix_a) + ).tolist() + ) + assert ( + cpv.translate_array(vector_a, matrix_a) + == np.array( + [[7.0, 12.0, 15.0], [19.0, 27.0, 42.0], [1987.0, 5.0, 10.0]] + ).tolist() + ) + + +def test_fit(matrix_a, matrix_b): + + def fit( + target_array: np.ndarray, source_array: np.ndarray + ) -> tuple[np.ndarray, np.ndarray, np.ndarray, float]: + """fit(target_array, source_array) -> (t1, t2, rot_mtx, rmsd) [fit_result] + + Calculates the translation vectors and rotation matrix required + to superimpose source_array onto target_array. Original arrays are + not modified. NOTE: Currently assumes 3-dimensional coordinates + """ + # Check dimensions of input arrays + if target_array.shape != source_array.shape: + raise ValueError("Error: arrays must be of same length for RMS fitting.") + if target_array.shape[1] != 3 or source_array.shape[1] != 3: + raise ValueError("Error: arrays must be dimension 3 for RMS fitting.") + + ndim = 3 + maxiter = 200 + tol = 0.001 + + # Calculate translation vectors (center-of-mass) + t1 = np.mean(target_array, axis=0) + t2 = np.mean(source_array, axis=0) + + # Calculate correlation matrix + tvec1 = target_array - t1 + tvec2 = source_array - t2 + corr_mtx = np.dot(tvec2.T, tvec1) + + # Initial rotation matrix (identity matrix) + rot_mtx = np.eye(ndim) + + # Main iteration scheme (hardwired for 3X3 matrix) + for iters in range(maxiter): + iy = iters % ndim + iz = (iters + 1) % ndim + sig = corr_mtx[iz, iy] - corr_mtx[iy, iz] + gam = corr_mtx[iy, iy] + corr_mtx[iz, iz] + + sg = np.sqrt(sig**2 + gam**2) + + if sg != 0.0 and abs(sig) > tol * abs(gam): + sg = 1.0 / sg + for i in range(ndim): + + bb = gam * corr_mtx[iy, i] + sig * corr_mtx[iz, i] + cc = gam * corr_mtx[iz, i] - sig * corr_mtx[iy, i] + corr_mtx[iy, i] = bb * sg + corr_mtx[iz, i] = cc * sg + + bb = gam * rot_mtx[iy, i] + sig * rot_mtx[iz, i] + cc = gam * rot_mtx[iz, i] - sig * rot_mtx[iy, i] + rot_mtx[iy, i] = bb * sg + rot_mtx[iz, i] = cc * sg + + continue + + # Calculate RMS deviation + vt = np.dot(source_array - t2, rot_mtx.T) + t1 + rmsd = np.sqrt(np.mean(np.sum((target_array - vt) ** 2, axis=1))) + + return t1, t2, rot_mtx, rmsd + + t1_np, t2_np, rot_mtx_np, rmsd_np = fit(np.array(matrix_a), np.array(matrix_b)) + result = cpv.fit(matrix_a, matrix_b) + + if result is None: + raise + t1, t2, rot_mtx, rmsd = result + assert t1 == t1_np.tolist() + assert t2 == t2_np.tolist() + for i in range(len(rot_mtx)): + assert rot_mtx[i] == pytest.approx(rot_mtx_np[i].tolist(), 0.001) + assert rmsd == pytest.approx(rmsd_np) + + +def test_fit_apply(matrix_a, matrix_b): + result = cpv.fit(matrix_a, matrix_b) + + if result is None: + raise + t1, mt2, m = result[:3] + + origin_result = cpv.fit_apply(result, matrix_b) + np_result = np.array(t1) + np.dot(np.array(m), (np.array(mt2) + np.array(matrix_b))) + + for i in range(len(origin_result)): + assert origin_result[i] == pytest.approx(np_result[i].tolist(), 1) + + np.testing.assert_array_almost_equal( + np.array(origin_result), + np.array( + [ + [661.0827157827758, 7.129705092201009, 17.79008759059663], + [665.1572144529815, 8.019304893172958, 19.404730826919444], + [665.5694270238597, 7.575498193306849, 18.609044701890262], + ] + ).tolist() + ) + np.testing.assert_array_almost_equal( + np_result, + np.array( + [ + [662.0405274477379, 8.365106579876256, 18.542585300247588], + [665.1798712452187, 7.2261452638093076, 19.82182248600216], + [662.2963554892253, 8.460848531966128, 18.41893138496441], + ] + ), + )