def iter_vectors(self, with_indices=False):
last_dim = self.shape[-1]
flags = []
if with_indices:
flags.append('multi_index')
it = np.nditer(self, flags=flags)
curr_ary = []
# TODO make this yield a view (?)
for x in it:
curr_ary.append(x)
if len(curr_ary) == last_dim:
if with_indices:
yield LightVector(curr_ary), x.multi_index[:-1]
else:
yield LightVector(curr_ary)
curr_ary = []
def coordinate_with_cartesian_displacements(self, base_coord, disps):
disps = tuple(disps)
if all(d == 0 for d in disps):
return base_coord
ret_val = self.displaced_coordinates.get((base_coord, disps), None)
if ret_val is None:
new_atoms = []
for atom, disp in zip(base_coord.atoms, grouper(3, disps)):
new_atoms.append(atom.displaced(LightVector(disp)*self.delta))
ret_val = base_coord.copy_with_atoms(new_atoms)
# Cache it for later to avoid gratuitous creation of displaced Molecule instances
self.displaced_coordinates[(base_coord, disps)] = ret_val
return ret_val
def value_for_displacements(self, pairs):
# get the position vector
xyzvect = LightVector([a.position for a in self.atoms]).ravel()
# make pairs into a list in case it's a general iterator
pairs = list(pairs)
int_idxs = self.internal_indices_for_coordinates(*[pair[0] for pair in pairs])
for i, (__, ndeltas) in enumerate(pairs):
xyzvect[int_idxs[i]] += self.b_tensor_finite_difference_delta * ndeltas
return self.__class__.value_for_xyz(
Matrix(
list(grouper(3, xyzvect))
)
)
def transform_tensor(self, tensor, to_representation):
"""
"""
self.freeze()
to_representation.freeze()
shape = (len(to_representation), ) * len(tensor.shape)
ret_val = RepresentationDependentTensor(
shape=shape, representation=to_representation)
#--------------------------------------------------------------------------------#
if isinstance(to_representation, CartesianRepresentation):
if len(self) != len(to_representation):
raise ValueError(
"incompatible representation sizes ({} != {})".format(
len(self), len(to_representation)))
if tensor.representation is not self:
raise ValueError(
"representation {} can only transform a tensor whose representation attribute is the same "
" as itself (tensor.representation was {} ".format(
self, tensor.representation))
if self.molecule.is_linear():
raise NotImplementedError(
"linear molecule cartesian-to-cartesian transformation is not"
" yet implemented. Shouldn't be too hard...")
if sanity_checking_enabled:
#TODO use a recentered version of the 'from' representation if it is not centered
pass
# Make sure things are centered
#if not self.molecule.is_centered(cartesian_representation=self):
# raise ValueError("CartesianRepresentation objects transformed from and to"
# " must be centered at the center of mass ({} was not)".format(
# self
# ))
#elif not self.molecule.is_centered(cartesian_representation=to_representation):
# raise ValueError("CartesianRepresentation objects transformed from and to"
# " must be centered at the center of mass ({} was not)".format(
# to_representation
# ))
#----------------------------------------#
old = self.value
new = to_representation.value
unitconv = self.units.to(to_representation.units)
oldmat = Matrix(old).reshape((len(self) / 3, 3))
newmat = Matrix(new).reshape((len(self) / 3, 3))
#----------------------------------------#
# Check to see if the molecule is planar. If so, append the cross product of any
# two atoms' positions not colinear with the origin
if any(norm(oldmat[:dir].view(Vector)) < 1e-12 \
or norm(newmat[:dir].view(Vector)) < 1e-12 \
for dir in [X, Y, Z]):
first_atom = None
# TODO unit concious cutoff (e.g. this would fail if the units of the position matrix were meters)
nonorigin_cutoff = 1e-2
for i, v in enumerate(grouper(3, old)):
if norm(LightVector(v)) > nonorigin_cutoff:
first_atom = i
break
second_atom = None
for i in range(first_atom + 1, len(self) // 3):
#TODO make sure this works with Molecule.linear_cutoff
if norm(oldmat[i]) > nonorigin_cutoff and norm(newmat[i]) > nonorigin_cutoff \
and abs(angle_between_vectors(oldmat[first_atom], oldmat[i])) > 1e-5 \
and abs(angle_between_vectors(newmat[first_atom], newmat[i])) > 1e-5:
second_atom = i
break
oldmat = Matrix(
list(oldmat.rows) +
[cross(oldmat[first_atom], oldmat[second_atom])])
newmat = Matrix(
list(newmat.rows) +
[cross(newmat[first_atom], newmat[second_atom])])
#----------------------------------------#
rot_mat = newmat.T * np.linalg.pinv(oldmat.T)
# Divide by unit conversion because the representation dependent tensor
# is [some units] *per* representation units
rot_mat /= unitconv
trans_mat = Matrix(shape=(len(self), len(self)))
for idx in xrange(len(self) // 3):
off = idx * 3
trans_mat[off:off + 3, off:off + 3] = rot_mat
order = len(tensor.shape)
# This is basically impossible to read what I'm doing here, but work it out
# looking at the numpy.einsum documentation and you'll see that this is correct.
# Basically, we want to contract the row indices of the transformation matrix
# with each axis of the tensor to be transformed.
einsumargs = sum(
([trans_mat, [2 * i, 2 * i + 1]] for i in xrange(order)), [])
einsumargs += [
tensor, [i for i in xrange(2 * order) if i % 2 != 0]
]
einsumargs += [[i for i in xrange(2 * order) if i % 2 == 0]]
np.einsum(*einsumargs, out=ret_val)
return ret_val
#--------------------------------------------------------------------------------#
elif isinstance(to_representation, InternalRepresentation):
if len(tensor.shape) == 1:
A = to_representation.a_matrix
ret_val[...] = A.T * tensor.view(Vector)
else:
raise NotImplementedError("use transform_forcefield instead")
return ret_val
#--------------------------------------------------------------------------------#
elif isinstance(to_representation, NormalRepresentation):
B = to_representation.b_matrix
ret_val[...] = tensor.linearly_transformed(B)
return ret_val
else:
raise NotImplementedError(
"Transformation of arbitrary tensors from representation of type '{}' to "
" representations of type '{}' is not implemented.".format(
self.__class__.__name__,
to_representation.__class__.__name__))
def r(j, k):
return LightVector.l_sub(args[k - 1], args[j - 1]).magnitude()
def e(j, k):
ev = LightVector.l_sub(args[k - 1], args[j - 1])
ev.normalize()
return ev
def value_for_xyz(cls, xyz):
v1 = LightVector.l_sub(xyz[0], xyz[1])
v2 = LightVector.l_sub(xyz[2], xyz[1])
return angle_between_vectors(v1, v2)
def analytic_b_tensor_for_order(self, order):
# First check the cache
cache_key = (self.__class__, order) + tuple(a.pos for a in self.atoms)
cache_resp = SimpleInternalCoordinate._check_b_tensor_cache(*cache_key)
if cache_resp is not None:
return cache_resp
#--------------------------------------------------------------------------------#
B = partial(SimpleInternalCoordinate.b_tensor_element_reindexed, self)
#--------------------------------------------------------------------------------#
# First order is already done elsewhere...
if order == 1:
# We can't use the `b_vector` attribute since that vector is indexed in the parent
# molecule's indexing scheme
my_b = Vector(self.__class__.b_vector_for_positions(*[a.pos for a in self.atoms]))
# Return now, to avoid double-caching
return my_b
#--------------------------------------------------------------------------------#
# TODO Explicit (matrix/tensor based) implementations of 2nd, 3rd, and 4th order to speed things up substantially
else:
# We only have general formulas for terminal atoms as of now...
# So we can save a little bit of time by computing these terms and
# then doing only finite difference for everything else
# Torsions are composed of 4 atoms (3 CartesianCoordinates each),
# so the output will be a 12x12x...x12 (`order`-dimensional) tensor
my_b = np.ndarray(shape=(12,)*order)
if sanity_checking_enabled:
my_b = np.ones(shape=(12,)*order) * float('inf')
# some precomputed values
#========================================#
# Helper function needed for derivative:
def Bsin2phi(phi, *idx_alphas):
# some precomputed values
twophi = 2.0*phi.value*phi.units.to(Radians)
sin2phi = sin(twophi)
cos2phi = cos(twophi)
def h_K(k):
if k % 2 == 1:
if ((k-1) / 2) % 2 == 0:
return cos2phi
else:
return -cos2phi
else: # k is even
if k/2 % 2 == 0:
return sin2phi
else:
return -sin2phi
cumsum = 0.0
n = len(idx_alphas)
for k in xrange(1, n + 1):
inner_sum = 0.0
for s in I_lubar(n, k, idx_alphas):
inner_prod = 1.0
for i in range(k):
inner_prod *= B(phi, *s[i])
inner_sum += inner_prod
cumsum += 2.0**k * h_K(k) * inner_sum
return cumsum
#========================================#
# The aaaa...bbbbb...cccc... terms
phis = []
rbcs = []
for a_idx, a in zip([0, 3], self.terminal_atoms):
# Determine whihc terminal atom we're handling
if a_idx == 0:
b_idx, c_idx, d_idx = 1, 2, 3
else: # a_idx == 3
b_idx, c_idx, d_idx = 2, 1, 0
#----------------------------------------#
phi_abc = self.get_coord(BondAngle,
self.atoms[a_idx],
self.atoms[b_idx],
self.atoms[c_idx],
)
# units of Radians
ang_conv = phi_abc.units.to(Radians)
# TODO figure out what 'units' should be here
rbc = self.get_coord(BondLength, self.atoms[b_idx], self.atoms[c_idx])
# Keep them for later
phis.append(phi_abc)
rbcs.append(rbc)
#----------------------------------------#
# some precomputed values
sin2phi = sin(2.0 * phi_abc.value * ang_conv)
#----------------------------------------#
for num_a, num_b, num_c in unlabeled_balls_in_labeled_boxes(order-1, [order-1]*3):
num_a += 1
alph_iter = product(*map(lambda n: symmetric_product([X,Y,Z], n), (num_a, num_b, num_c)))
for alphas_a, alphas_b, alphas_c in alph_iter:
a_alphas = tuple(3*a_idx + alpha for alpha in alphas_a)
b_alphas = tuple(3*b_idx + alpha for alpha in alphas_b)
c_alphas = tuple(3*c_idx + alpha for alpha in alphas_c)
all_alphas = a_alphas + b_alphas + c_alphas
cum_sum = 0.0
for t1, t2 in I_ll(num_b + num_c, 2, b_alphas + c_alphas):
bphi = LightVector([
B(phi_abc, 3*a_idx + sigma, *(a_alphas[1:] + t1))
for sigma in [X, Y, Z]
])
brbc = LightVector([
B(rbc, 3*c_idx + sigma, *t2) for sigma in [X, Y, Z]]
)
cum_sum += cross(bphi, brbc)[alphas_a[0]]
cum_sum *= 2.0
cum_sum -= B(self, a_alphas[0]) * Bsin2phi(phi_abc, *all_alphas[1:])
for s1, s2 in I_llbar(num_a - 1 + num_b + num_c, 2, all_alphas[1:]):
cum_sum -= Bsin2phi(phi_abc, *s1) * B(self, *((a_alphas[0],) + s2))
cum_sum /= sin2phi
for perm in permutations(all_alphas):
my_b[perm] = cum_sum
#========================================#
# Note that the terminal-atom cross-derivatives are 0
if sanity_checking_enabled:
# Fill in the explicity zeros now, since we had infinity there to make sure
# uncomputed values weren't being used for something else.
for a_idx, d_idx in permutations([0, 3]):
for remaining_idxs in product([0, 1, 2, 3], repeat=order-2):
for alphas in product([X, Y, Z], repeat=order):
idx_alphas = tuple(3*atom_idx + alpha
for atom_idx, alpha in zip((a_idx, d_idx) + remaining_idxs, alphas))
for perm in permutations(idx_alphas):
my_b[perm] = 0.0
#========================================#
# Arbitrary order bbbbb.... terms
a_idx, b_idx, c_idx, d_idx = 0, 1, 2, 3
phi_abc = phis[0]
phi_bcd = self.get_coord(BondAngle,
self.atoms[b_idx],
self.atoms[c_idx],
self.atoms[d_idx],
)
ang_conv = phi_bcd.units.to(Radians)
# TODO figure out what 'units' should be here
r_ba = self.get_coord(BondLength,
self.atoms[b_idx],
self.atoms[a_idx]
)
r_cd = self.get_coord(BondLength,
self.atoms[c_idx],
self.atoms[d_idx]
)
#----------------------------------------#
def Bcscphi(phi, *b_alphas):
phi_val = phi.value * phi.units.to(Radians)
sinphi = sin(phi_val)
cscphi = 1.0 / sinphi
if len(b_alphas) == 0:
return cscphi
cotphi = cos(phi_val) / sinphi
#------------------------------------#
def dcsc_n(n):
def t(n_t, k):
if k == 0:
return 1
elif k <= n_t//2:
return (2*k + 1) * t(n_t-1, k) + (n_t - 2*k + 1) * t(n_t-1, k-1)
else:
return 0
#--------------------------------#
ret_val = 0.0
for kk in xrange(n//2 + 1):
ret_val += t(n, kk) * cotphi**(n - 2*kk) * cscphi**(2*kk + 1)
if n % 2 == 1:
return -ret_val
else:
return ret_val
#------------------------------------#
outer_sum = 0.0
for k in xrange(1, len(b_alphas) + 1):
inner_sum = 0.0
for idx_sets in labeled_balls_in_unlabeled_boxes(len(b_alphas), [len(b_alphas)]*k):
if any(len(st) == 0 for st in idx_sets):
continue
b_idx_sets = tuple(tuple(b_alphas[i] for i in idxset) for idxset in idx_sets)
product = 1.0
for b_idxs in b_idx_sets:
product *= B(phi, *b_idxs)
inner_sum += product
outer_sum += dcsc_n(k) * inner_sum
return outer_sum
#----------------------------------------#
for alphas in symmetric_product([X, Y, Z], order):
# here we go...
term_sum = first_term = second_term = third_term = 0.0
iter_alphas = alphas[1:]
b_alphas = tuple(3*b_idx + alpha for alpha in alphas)
#----------------------------------------#
for b_alphas1, b_alphas2, b_alphas3 in I_ll(order-1, 3, b_alphas[1:]):
#----------------------------------------#
# preconstruct the vectors we need for the cross products
b_a_phi_abc = LightVector([
B(phi_abc, 3*a_idx + sigma, *b_alphas2) for sigma in [X, Y, Z]
])
b_c_phi_abc = LightVector([
B(phi_abc, 3*c_idx + sigma, *b_alphas2) for sigma in [X, Y, Z]
])
b_a_r_ba = LightVector([
B(r_ba, 3*a_idx + sigma, *b_alphas3) for sigma in [X, Y, Z]
])
#----------------------------------------#
# save a little bit of time by only computing this part once per iteration
Bcscphi_abc = Bcscphi(phi_abc, *b_alphas1)
#----------------------------------------#
# now add the contribution from this set of indices
first_term += Bcscphi_abc * cross(b_a_phi_abc, b_a_r_ba)[alphas[0]]
second_term += Bcscphi_abc * cross(b_c_phi_abc, b_a_r_ba)[alphas[0]]
#----------------------------------------#
term_sum -= first_term + second_term
b_d_r_cd = LightVector([
B(r_cd, 3*d_idx + sigma) for sigma in [X, Y, Z]
])
for b_alphas1, b_alphas2 in I_ll(order-1, 2, b_alphas[1:]):
#----------------------------------------#
b_b_phi_bcd = LightVector([
B(phi_bcd, 3*b_idx + sigma, *b_alphas2) for sigma in [X, Y, Z]
])
#----------------------------------------#
third_term += Bcscphi(phi_bcd, *b_alphas1) * cross(b_b_phi_bcd, b_d_r_cd)[alphas[0]]
term_sum += third_term
#----------------------------------------#
# and spread it across permutations
for perm in permutations(tuple(3*b_idx + alpha for alpha in alphas)):
my_b[perm] = term_sum
#========================================#
# and fill in the bbbb...cccc... derivatives by translational invariance
# Range from one c index to all
for num_c in xrange(1, order+1):
num_b = order - num_c
alph_iter = product(*map(lambda n: symmetric_product([X,Y,Z], n), (num_b, num_c)))
for alphas_b, alphas_c in alph_iter:
b_alphas = tuple(3*b_idx + alph for alph in alphas_b)
c_alphas = tuple(3*c_idx + alph for alph in alphas_c)
alphas_all = alphas_b + alphas_c
currsum = 0.0
for repl_atoms in product([a_idx, b_idx, d_idx], repeat=num_c):
repl_alphas = tuple(3*repl_atom + alph
for repl_atom, alph in zip(repl_atoms, alphas_c))
currsum += my_b[repl_alphas + b_alphas]
if sanity_checking_enabled and math.isinf(my_b[b_alphas + repl_alphas]):
raise IndexError("indices not filled in: {}, needed for {}".format(
b_alphas + repl_alphas,
b_alphas + c_alphas
))
if num_c % 2 == 1:
currsum *= -1.0
#----------------------------------------#
# and spread this value over all permutations...
for perm in permutations(b_alphas + c_alphas):
my_b[perm] = currsum
#--------------------------------------------------------------------------------#
# Cache the value we got
SimpleInternalCoordinate._set_b_tensor_cache_entry(my_b, *cache_key)
#--------------------------------------------------------------------------------#
return my_b
def r(j, k):
return LightVector.l_sub(args[k-1], args[j-1]).magnitude()
def value_for_xyz(cls, xyz):
"""
"""
return LightVector.l_sub(xyz[1], xyz[0]).magnitude()
def e(j, k):
ev = LightVector.l_sub(xyz[k-1], xyz[j-1]); ev.normalize()
return ev
def value_for_xyz(cls, xyz):
"""
"""
return LightVector.l_sub(xyz[1], xyz[0]).magnitude()
def value_for_xyz(cls, xyz):
v1 = LightVector.l_sub(xyz[0], xyz[1])
v2 = LightVector.l_sub(xyz[2], xyz[1])
return angle_between_vectors(v1, v2)
def l_xyz(self):
"""
Returns the value of the representation as a light vector
"""
return LightVector([c.value for c in self])
def e(j, k):
ev = LightVector.l_sub(args[k-1], args[j-1]); ev.normalize()
return ev
def e(j, k):
ev = LightVector.l_sub(xyz[k - 1], xyz[j - 1])
ev.normalize()
return ev
def _l_xyz(self):
return LightVector([atom.pos for atom in self.atoms])
