def grad(self, A, B, q): #Compute relative transformation pr = norm(q.x) if(pr >= A.radius + B.radius ): return array([0., 0., 0., 0.]) #Load shape parameters fa = A.pft fb = B.pft da = A.pdft db = B.pdft ea = A.energy eb = B.energy #Estimate cutoff threshold cutoff = self.__get_cutoff(A, B) #Compute coordinate coefficients m = 2.j * pi / (2. * self.SHAPE_R + 1) * pr phi = atan2(q.x[1], q.x[0]) #Set up initial sums s_0 = real(fa[0][0] * fb[0][0]) s_x = 0. s_y = 0. s_ta = 0. s_tb = 0. for r in range(1, self.R): #Compute theta terms dtheta = 2. * pi / len(fa[r]) theta = arange(len(fa[r])) * dtheta #Construct multiplier / v mult = exp((m * r) * cos(theta + phi)) * r * dtheta u = pds.shift(conjugate(fb[r]), q.theta) * mult v = fa[r] * u #Check for early out s_0 += sum(real(v)) if(s_0 + min(ea[r], eb[r]) <= cutoff): return array([0.,0.,0.,0.]) #Sum up gradient vectors v = real(1.j * v) s_x -= sum(v * sin(theta + phi) ) s_y -= sum(v * cos(theta + phi) ) s_t += sum(real(da[r] * u)) if(s_0 <= cutoff): return array([0., 0., 0., 0.]) return array([s_x, s_y, s_ta, s_tb, s_0])
def pft_rotate(pft, theta): R = len(pft) res = [] for k in range(R): res.append(shift(pft[k], theta)) return res