def I_ul(nballs, nboxes, labels):
if nballs == 0:
yield (tuple(),) * nboxes
else:
for occs in unlabeled_balls_in_labeled_boxes(nballs, [nballs]*nboxes):
stop_points = (0,) + tuple(cummulative_sum(occs))
yield tuple(labels[stop_points[i]:stop_points[i+1]] for i in xrange(nboxes))
def I_ul(nballs, nboxes, labels):
if nballs == 0:
yield (tuple(), ) * nboxes
else:
for occs in unlabeled_balls_in_labeled_boxes(nballs,
[nballs] * nboxes):
stop_points = (0, ) + tuple(cummulative_sum(occs))
yield tuple(labels[stop_points[i]:stop_points[i + 1]]
for i in xrange(nboxes))
def all_partitions_iter(sequence, length, allow_zero=True):
"""
Iterate over all possible partitions of a sequence
:Examples:
"""
n = len(sequence)
for partitions in unlabeled_balls_in_labeled_boxes(n, [n]*length):
if not allow_zero and 0 in partitions:
continue
to_yield = tuple(partitioned(sequence, partitions))[:length]
yield to_yield
def all_partitions_iter(sequence, length, allow_zero=True):
"""
Iterate over all possible partitions of a sequence
:Examples:
"""
n = len(sequence)
for partitions in unlabeled_balls_in_labeled_boxes(n, [n] * length):
if not allow_zero and 0 in partitions:
continue
to_yield = tuple(partitioned(sequence, partitions))[:length]
yield to_yield
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
