def propagate_plane_to_curved_flat_arbitrary(k, rs_support, Eri, z, xo, yo, qs_center=(0, 0), kz_mode='local_xy'): """Propagate from uniformly sampled plane to arbitrarily sampled curved surface. Args: k (scalar): Wavenumber. rs_support (scalar or pair of): Input aperture along x and y. num_pointss (int or tuple of): Number of input samples along x and y. Referred to as K and L below. z (M*N array): Propagation distances. xo (Mx1 array): Output x values. yo (M array): Output y values. qs_center (tuple or pair of): Center of transverse wavenumber support. kz_mode (str): 'paraxial' or 'local_xy'. Returns: Ero (MxN array): Output field. gradxyEo (tuple of MxN arrays): Partial derivatives of output field w.r.t. x and y at constant z (not along the surface). """ factors = math.calc_plane_to_curved_flat_arbitrary_factors( k, rs_support, Eri.shape, z, xo, yo, qs_center, kz_mode) invTx, gradxinvTx, invTy, gradyinvTy, Px, Py, Tx, Ty = factors # xo=i, yo=j, kx=k, ky=l, xi=m, yi=n. Ero = opt_einsum.contract('ik, jl, ijk, ijl, km, ln, mn -> ij', invTx, invTy, Px, Py, Tx, Ty, Eri) gradxEo = opt_einsum.contract('ik, jl, ijk, ijl, km, ln, mn -> ij', gradxinvTx, invTy, Px, Py, Tx, Ty, Eri) gradyEo = opt_einsum.contract('ik, jl, ijk, ijl, km, ln, mn -> ij', invTx, gradyinvTy, Px, Py, Tx, Ty, Eri) return Ero, (gradxEo, gradyEo)
def mode_coeff_calculation_multiprocessing_wrapper(args): #Integrates F(x1,x2,...,xn) * sin((m_1+1)*pi*x_1/L_1) * ... * sin((m_n+1)*pi*x_n/L_n)) dx_1 ... dx_n. Written as a separate function to permit multiprocessing pool map F = args[0] domain_volume = args[1] mplus1_pi_over_L = args[2] integrator = args[3] two_to_the_power_ndims = args[4] coefficients = [] for i in range(int(mplus1_pi_over_L.shape[0])): integrand = lambda *vars: F(*vars) * tf.reduce_prod(tf.sin(oe.contract('i...,i->i...',tf.stack(vars),mplus1_pi_over_L[i],backend = 'tensorflow')),axis=0) coefficients.append(-integrator(integrand) * tf.cast(two_to_the_power_ndims,tf.keras.backend.floatx()) / (tf.cast(domain_volume,tf.keras.backend.floatx()) * tf.reduce_sum(tf.square(mplus1_pi_over_L[i])))) return tf.constant(np.array(coefficients), dtype = tf.keras.backend.floatx())
def forward(self, x, y, z): r""" Args: x (torch.Tensor): ``[batch_size, seq_len, n_in]``. y (torch.Tensor): ``[batch_size, seq_len, n_in]``. z (torch.Tensor): ``[batch_size, seq_len, n_in]``. Returns: ~torch.Tensor: A scoring tensor of shape ``[batch_size, n_out, seq_len, seq_len, seq_len]``. If ``n_out=1``, the dimension for ``n_out`` will be squeezed automatically. """ if self.bias_x: x = torch.cat((x, torch.ones_like(x[..., :1])), -1) if self.bias_y: y = torch.cat((y, torch.ones_like(y[..., :1])), -1) w = contract('bzk,oikj->bozij', z, self.weight) # [batch_size, n_out, seq_len, seq_len, seq_len] s = contract('bxi,bozij,byj->bozxy', x, w, y) / self.n_in**self.scale # remove dim 1 if n_out == 1 s = s.squeeze(1) return s
def fold_img_with_L_inv(img, mu, L_inv, scale, threshold, normalize=True): device = img.get_device() bn, nc, h, w = img.shape bn, nk, _ = mu.shape # Stop Gradient Flow mu_stop = mu.detach() # Get Scaled Heatmap heat_scal = get_heat_map(mu_stop, scale * L_inv, False, h) heat_scal = contract('bkij -> bij', heat_scal) heat_scal = torch.clamp(heat_scal, min=0., max=1.) heat_scal = torch.where(heat_scal > threshold, heat_scal, torch.zeros_like(heat_scal)) # Normalize norm = torch.sum(heat_scal.reshape(bn, -1), dim=1).unsqueeze(1).unsqueeze(1) if normalize: heat_scal = heat_scal / norm # Return Folded Image around Part Means folded_img = contract('bcij, bij -> bcij', img, heat_scal) return folded_img
def pseudo_energy(self, l_ijab): ''' Compute the CCSD pseudoenergy :param l_ijab: Current iteration T2 :type l_ijab: numpy array :return: pseudoenergy :rtype: double ''' o = slice(0, self.no_occ) v = slice(self.no_occ, self.no_mo) # E = 1/2 <ab|ij> l_ijab E_pseudo = 0.5 * contract('abij,ijab->', self.MO[v, v, o, o], l_ijab) return E_pseudo
def test_cupy(string): # pragma: no cover views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) opt = expr(*views, backend='cupy') assert np.allclose(ein, opt) # test non-conversion mode cupy_views = backends.convert_arrays_to_cupy(views) cupy_opt = expr(*cupy_views, backend='cupy') assert isinstance(cupy_opt, cupy.ndarray) assert np.allclose(ein, cupy.asnumpy(cupy_opt))
def hess_WuYang(self, v): Vks_a = psi4.core.Matrix.from_array(np.einsum("ijk,k->ij", self.three_overlap, v[:self.nbf]) + self.initial_guess) Vks_b = psi4.core.Matrix.from_array(np.einsum("ijk,k->ij", self.three_overlap, v[self.nbf:]) + self.initial_guess) self.frags[0].scf( maxiter=0, #Solved non-self consistently hamiltonian=["kinetic", "external", "xxxtra"], xfock_nm=[Vks_a, Vks_b], get_ingredients=True) epsilon_occ_a = self.frags[0].eigs_a.np[:self.frags[0].nalpha, None] epsilon_occ_b = self.frags[0].eigs_b.np[:self.frags[0].nbeta, None] epsilon_unocc_a = self.frags[0].eigs_a.np[self.frags[0].nalpha:] epsilon_unocc_b = self.frags[0].eigs_b.np[self.frags[0].nbeta:] epsilon_a = epsilon_occ_a - epsilon_unocc_a epsilon_b = epsilon_occ_b - epsilon_unocc_b hess = np.zeros((self.nbf*2, self.nbf*2)) # Alpha electrons hess[0:self.nbf, 0:self.nbf] = - 1.0 * contract('ai,bj,ci,dj,ij,abm,cdn -> mn', self.frags[0].Ca.np[:, :self.frags[0].nalpha], self.frags[0].Ca.np[:, self.frags[0].nalpha:], self.frags[0].Ca.np[:, :self.frags[0].nalpha], self.frags[0].Ca.np[:, self.frags[0].nalpha:], np.reciprocal(epsilon_a), self.three_overlap, self.three_overlap) # Beta electrons hess[self.nbf:, self.nbf:] = - 1.0 * contract('ai,bj,ci,dj,ij,abm,cdn -> mn', self.frags[0].Cb.np[:, :self.frags[0].nbeta], self.frags[0].Cb.np[:, self.frags[0].nbeta:], self.frags[0].Cb.np[:, :self.frags[0].nbeta], self.frags[0].Cb.np[:, self.frags[0].nbeta:], np.reciprocal(epsilon_b),self.three_overlap, self.three_overlap) hess = (hess + hess.T) return hess
def test_sequential_logmatmulexp(batch_shape, state_dim, num_steps): logits = torch.randn(batch_shape + (num_steps, state_dim, state_dim)) actual = _sequential_logmatmulexp(logits) assert actual.shape == batch_shape + (state_dim, state_dim) # Check against einsum. operands = list(logits.unbind(-3)) symbol = (opt_einsum.get_symbol(i) for i in range(1000)) batch_symbols = ''.join(next(symbol) for _ in batch_shape) state_symbols = [next(symbol) for _ in range(num_steps + 1)] equation = (','.join(batch_symbols + state_symbols[t] + state_symbols[t + 1] for t in range(num_steps)) + '->' + batch_symbols + state_symbols[0] + state_symbols[-1]) expected = opt_einsum.contract(equation, *operands, backend='pyro.ops.einsum.torch_log') assert_close(actual, expected)
def solve(self, b): """Summary: Under a given b vector, construct a fock matrix and diagonalize it F = F0 + V(b) FC = SCE resulting mo_coeff(C), mo_energy(E), and density matrix are stored as instance attributes """ if not np.allclose(b, self.internal_b, rtol=1e-12, atol=1e-12): t = time.time() self.internal_b = b.copy() Fa = self.F0[0]+contract('t,ijt->ij', b[:self.npot], self.Sijt) Fb = self.F0[1]+contract('t,ijt->ij', b[self.npot:], self.Sijt) self.fock = ((Fa, Fb)) e_a, c_a = scf.hf.eig(Fa, self.S) e_b, c_b = scf.hf.eig(Fb, self.S) self.mo_coeff = np.array((c_a, c_b)) self.mo_energy = np.array((e_a, e_b)) self.mo_occ = self.get_occ(self.mo_energy, self.mo_coeff) self.dm = self.make_rdm1(self.mo_coeff, self.mo_occ) self.t_eig += time.time()-t t = time.time() if kf_imported: if self.Smnt is None: #AO basis for WY is same with AO basis of target density self.grad_a = kspies_fort.einsum_ij_ijt_2t((self.dm[0]-self.dm_tar[0]), self.Sijt, self.nbas, self.npot) self.grad_b = kspies_fort.einsum_ij_ijt_2t((self.dm[1]-self.dm_tar[1]), self.Sijt, self.nbas, self.npot) else: self.grad_a = kspies_fort.einsum_ij_ijt_2t(self.dm[0], self.Sijt, self.nbas, self.npot) self.grad_a -= kspies_fort.einsum_ij_ijt_2t(self.dm_tar[0], self.Smnt, len(self.Smnt[:, 0, 0]), self.npot) self.grad_b = kspies_fort.einsum_ij_ijt_2t(self.dm[1], self.Sijt, self.nbas, self.npot) self.grad_b -= kspies_fort.einsum_ij_ijt_2t(self.dm_tar[1], self.Smnt, len(self.Smnt[:, 0, 0]), self.npot) else: if self.Smnt is None: self.grad_a = contract('ij,ijt->t', (self.dm[0]-self.dm_tar[0]), self.Sijt) self.grad_b = contract('ij,ijt->t', (self.dm[1]-self.dm_tar[1]), self.Sijt) else: self.grad_a = contract('ij,ijt->t', self.dm[0], self.Sijt) self.grad_a -= contract('ij,ijt->t', self.dm_tar[0], self.Smnt) self.grad_b = contract('ij,ijt->t', self.dm[1], self.Sijt) self.grad_b -= contract('ij,ijt->t', self.dm_tar[1], self.Smnt) self.t_gd += time.time()-t
def make_input_tps_param(tps_param, move_point=None, scal_point=None): ''' ''' coord = tps_param.coord vector = tps_param.vector offset = tps_param.offset offset_2 = tps_param.offset_2 rot_mat = tps_param.rot_mat t_scal = tps_param.t_scal scaled_coord = contract('bk,bck->bck', t_scal, coord + vector - offset) + offset t_vector = contract('blk,bck->bcl', rot_mat, scaled_coord - offset_2) + offset_2 - coord if move_point is not None and scal_point is not None: coord = contract('bk,bck->bck', scal_point, coord + move_point) t_vector = contract('bk,bck->bck', scal_point, t_vector) else: assert (move_point is None and scal_point is None) return coord, t_vector
def test_custom_random_optimizer(): class NaiveRandomOptimizer(oe.path_random.RandomOptimizer): @staticmethod def random_path(r, n, inputs, output, size_dict): """Picks a completely random contraction order.""" np.random.seed(r) ssa_path = [] remaining = set(range(n)) while len(remaining) > 1: i, j = np.random.choice(list(remaining), size=2, replace=False) remaining.add(n + len(ssa_path)) remaining.remove(i) remaining.remove(j) ssa_path.append((i, j)) cost, size = oe.path_random.ssa_path_compute_cost( ssa_path, inputs, output, size_dict) return ssa_path, cost, size def setup(self, inputs, output, size_dict): self.was_used = True n = len(inputs) trial_fn = self.random_path trial_args = (n, inputs, output, size_dict) return trial_fn, trial_args eq, shapes = oe.helpers.rand_equation(5, 3, seed=42, d_max=3) views = list(map(np.ones, shapes)) exp = oe.contract(eq, *views, optimize=False) optimizer = NaiveRandomOptimizer(max_repeats=16) out = oe.contract(eq, *views, optimize=optimizer) assert exp == out assert optimizer.was_used assert len(optimizer.costs) == 16
def compute_ccsd_dipole(self, t1=None,t2=None,l1=None,l2=None): if t1 is None: t1 = self.t1 if t2 is None: t2 = self.t2 if l1 is None: l1 = self.l1 if l2 is None: l2 = self.l2 D = self.compute_ccsd_density(t1,t2,l1,l2) # Compute CCSD correlated dipole dipoles_elec = [] for n in range(3): d = contract('ui,uv,vj->ij',self.npC,np.asarray(self.ints[n]),self.npC) mu = ndot('ij,ji->',d,D) dipoles_elec.append(mu) return dipoles_elec
def calculate_lambda_tensor(self, alpha, scattering_inverse): # TODO: replace with same caching strategy as rest of code n_k_points = self.n_k_points third_bandwidth = self.phonons.third_bandwidth if self.phonons.is_classic: stat_label = 'c' else: stat_label = 'q' str_to_add = str(n_k_points) + '_' + str(alpha) + '_' + str( int(self.phonons.temperature)) + '_' + stat_label lamdb_filename = 'lamdb_' + '_' + str_to_add psi_filename = 'psi_' + str_to_add psi_inv_filename = 'psi_inv_' + str_to_add if third_bandwidth is not None: lamdb_filename = lamdb_filename + '_' + str(third_bandwidth) psi_filename = psi_filename + '_' + str(third_bandwidth) psi_inv_filename = psi_inv_filename + '_' + str(third_bandwidth) lamdb_filename = lamdb_filename + '.npy' psi_filename = psi_filename + '.npy' psi_inv_filename = psi_inv_filename + '.npy' try: self._lambd = np.load(lamdb_filename, allow_pickle=True) self._psi = np.load(psi_filename, allow_pickle=True) self._psi_inv = np.load(psi_inv_filename, allow_pickle=True) except FileNotFoundError as err: logging.info(err) n_phonons = self.n_phonons physical_mode = self.phonons.physical_mode.reshape(n_phonons) velocity = self.phonons.velocity.real.reshape( (n_phonons, 3))[physical_mode, :] heat_capacity = self.phonons.heat_capacity.flatten()[physical_mode] sqr_heat_capacity = heat_capacity**0.5 v_new = velocity[:, alpha] lambd_tensor = contract('m,m,mn,n->mn', sqr_heat_capacity, v_new, scattering_inverse, 1 / sqr_heat_capacity) # evals and evect equations # lambd_tensor = psi.dot(np.diag(lambd)).dot(psi_inv) # lambd_tensor.dot(psi) = psi.dot(np.diag(lambd)) self._lambd, self._psi = np.linalg.eig(lambd_tensor) self._psi_inv = np.linalg.inv(self._psi) np.save(lamdb_filename, self._lambd) np.save(psi_filename, self._psi) np.save(psi_inv_filename, self._psi_inv)
def vp_wy_r(self, dd_a, dd_b): """ Performs the Wu-Yang Method on the grid """ dvp = np.zeros_like(self.molecule.Da) dd = dd_a + dd_b #Bring grid information points_func = self.molecule.Vpot.properties()[0] #Calculate denominator for block in range(self.molecule.Vpot.nblocks()): grid_block = self.molecule.Vpot.get_block(block) points_func.compute_points(grid_block) npoints = grid_block.npoints() lpos = np.array(grid_block.functions_local_to_global()) w = np.array(grid_block.w()) phi = np.array(points_func.basis_values()["PHI"])[:npoints, :lpos.shape[0]] x_a = np.zeros((npoints, npoints)) x_b = np.zeros((npoints, npoints)) for frag in self.frags: orb_a = frag.orbitals["alpha_r"] orb_b = frag.orbitals["beta_r"] for i_occ in range(0,frag.nalpha): for i_vir in range(frag.nalpha, frag.nbf): den = frag.eigs_a.np[i_occ] - frag.eigs_a.np[i_vir] num = np.zeros((npoints, npoints)) for r1 in range(npoints): for r2 in range(npoints): num[r1, r2] = orb_a[str(i_occ)][block][r1] * orb_a[str(i_vir)][block][r1] * orb_a[str(i_vir)][block][r2] * orb_a[str(i_occ)][block][r2] x_a += num / den #Assume x_b = x_a dvp_block = np.zeros((npoints)) for r1 in range(npoints): dvp_block += (1 / (x_a[r1, :] + x_a[r1, :])) * dd[block] * w vtmp = contract('pb,p,p,pa->ab', phi, dvp_block, w, phi) dvp[(lpos[:, None], lpos)] += 0.5 * (vtmp + vtmp.T) return dvp, dvp
def test_torch(string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) opt = expr(*views, backend='torch') assert np.allclose(ein, opt) # test non-conversion mode torch_views = [backends.to_torch(view) for view in views] torch_opt = expr(*torch_views) assert isinstance(torch_opt, torch.Tensor) assert np.allclose(ein, torch_opt.cpu().numpy())
def ELBO_prior(self, W, q_A, q_gamma, q_alpha, q_mu, q_pi, phi): ''' taking the expectation of log p(A) + sum_k [ log p(mu_k) + log p(Sigma_k) ] + sum_i [log p(pi_i) + log p (gamma_i )] + sum_i sum_j sum_k z_ijk [ log pi_ik + log p(alpha_ij) + log p(beta_ij)] ***NOTE THAT THIS TERM IS OFF BY CONSTANTS*** not the true elbo ''' #first_term = -0.5* ((q_A**2).sum() + (q_mu**2).sum() + (q_gamma**2).sum()) second_term1 = torch.log(q_pi + self.epsilon) second_term = contract('ijk,ik->', phi, second_term1) #third_term = -0.5* contract('ijk,ij->',phi, q_alpha.exp()**2) return second_term #+ third_term + first_term
def reintegrate_ao(self, function): f_nm = np.zeros((self.nbf, self.nbf)) points_func = self.frags[0].Vpot.properties()[0] for block in range(self.nblocks): grid_block = self.frags[0].Vpot.get_block(block) points_func.compute_points(grid_block) npoints = grid_block.npoints() lpos = np.array(grid_block.functions_local_to_global()) w = np.array(grid_block.w()) phi = np.array(points_func.basis_values()["PHI"])[:npoints, :lpos.shape[0]] print(phi.shape) print(function[block]) vtmp = contract('pb,p,p,pa->ab', phi, function[block], w, phi) f_nm[(lpos[:, None], lpos)] += 0.5 * (vtmp + vtmp.T) return f_nm
def test_tensorflow(string): views = helpers.build_views(string) ein = contract(string, *views, optimize=False, use_blas=False) opt = np.empty_like(ein) shps = [v.shape for v in views] expr = contract_expression(string, *shps, optimize=True) sess = tf.Session() with sess.as_default(): expr(*views, backend='tensorflow', out=opt) assert np.allclose(ein, opt) # test non-conversion mode tensorflow_views = backends.convert_arrays_to_tensorflow(views) expr(*tensorflow_views, backend='tensorflow')
def get_heat_map(mu, L_inv, device): h, w, nk = 64, 64, L_inv.shape[1] y_t = torch.linspace(-1., 1., h).reshape(h, 1).repeat(1, w).unsqueeze(-1) x_t = torch.linspace(-1., 1., w).reshape(1, w).repeat(h, 1).unsqueeze(-1) y_t_flat = y_t.reshape(1, 1, 1, -1) x_t_flat = x_t.reshape(1, 1, 1, -1) mesh = torch.cat((y_t_flat, x_t_flat), dim=-2).to(device) dist = mesh - mu.unsqueeze(-1) proj_precision = contract('bnik, bnkf -> bnif', L_inv, dist) ** 2 # tf.matmul(precision, dist)**2 proj_precision = torch.sum(proj_precision, -2) # sum x and y axis heat = 1 / (1 + proj_precision) heat = heat.reshape(-1, nk, h, w) # bn number parts width height return heat
def test_naive_ubersum(equation, plates): inputs, outputs, operands, sizes = make_example(equation) actual = naive_ubersum(equation, *operands, plates=plates) assert isinstance(actual, tuple) assert len(actual) == len(outputs) for output, actual_part in zip(outputs, actual): expected_shape = tuple(sizes[dim] for dim in output) assert actual_part.shape == expected_shape if not plates: equation_part = ','.join(inputs) + '->' + output expected_part = opt_einsum.contract(equation_part, *operands, backend='pyro.ops.einsum.torch_log') assert_equal(expected_part, actual_part, msg=u"For output '{}':\nExpected:\n{}\nActual:\n{}".format( output, expected_part.detach().cpu(), actual_part.detach().cpu()))
def contract(names, *tensors): args = [] ids = {} seen_names = [] for t in tensors: group = [] for name in t._schema._names: if name not in ids: ids[name] = len(ids) seen_names.append(name) group.append(ids[name]) args.append(t._tensor) args.append(group) names = names.split() keep = [n for n in seen_names if n not in names] args.append([ids[n] for n in keep]) return NamedTensor(oe.contract(*args, backend="torch"), keep)
def s_term_bernoulli(y_batch, gamma_batch, eta): g0 = gamma_batch.sum(dim=1, keepdim=True) g = gamma_batch / g0 probs = contract("mk,k->m", g, eta, backend="torch").sigmoid() # to prevent overflows in log probs_cpy = probs if probs.min() <= 0: c = probs.min().detach() probs = probs - c + self.epsilon s_term1 = (y_batch * probs.log()).sum() probs = probs_cpy if probs.max() >= 1: c = probs.max().detach() probs = probs - (c - 1) - self.epsilon s_term2 = ((1 - y_batch) * (1 - probs).log()).sum() s_term = s_term1 + s_term2 return s_term
def build_Dov(self, t1, t2, l1, l2): # complete if self.ccwfn.model == 'CCD': Dov = np.zeros_like(t1) else: Dov = 2.0 * t1.copy() Dov += 2.0 * contract('me,imae->ia', l1, t2) Dov -= contract('me,miae->ia', l1, self.ccwfn.build_tau(t1, t2)) tmp = contract('mnef,inef->mi', l2, t2) Dov -= contract('mi,ma->ia', tmp, t1) tmp = contract('mnef,mnaf->ea', l2, t2) Dov -= contract('ea,ie->ia', tmp, t1) return Dov
def test_contract_expressions(string, optimize, use_blas, out_spec): views = helpers.build_views(string) shapes = [view.shape for view in views] expected = contract(string, *views, optimize=False, use_blas=False) expr = contract_expression(string, *shapes, optimize=optimize, use_blas=use_blas) if out_spec and ("->" in string) and (string[-2:] != "->"): out, = helpers.build_views(string.split('->')[1]) expr(*views, out=out) else: out = expr(*views) assert np.allclose(out, expected) # check representations assert string in expr.__repr__() assert string in expr.__str__()
def build_Wabef(t1, t2): """Builds [Stanton:1991:4334] Eqn. 7""" # Rate limiting step written using tensordot, ~10x faster # The commented out lines are consistent with the paper Wabef = MO[v, v, v, v].copy() Pab = contract('baef->abef', np.tensordot(t1, MO[v, o, v, v], axes=(0, 1))) # Pab = np.einsum('mb,amef->abef', t1, MO[v, o, v, v]) Wabef -= Pab Wabef += Pab.swapaxes(0, 1) tmp_tau = build_tau(t1, t2) Wabef += 0.25 * np.tensordot(tmp_tau, MO[v, v, o, o], axes=((0, 1), (2, 3))) # Wabef += 0.25 * np.einsum('mnab,mnef->abef', tmp_tau, MO[o, o, v, v]) return Wabef
def test_eps_single_pixel_output() -> None: input = torch.randn((2, 3, 2, 2, 2), dtype=torch.float64) core = torch.rand((*(2 for _ in range(8)), 4), dtype=torch.float64) eps_result = rearrange(eps_one_by_one(core, input), "b () () o -> b o") assert eps_result.shape == (3, 4) oe_result = oe.contract( "01234567θ,b0,b1,b2,b3,b4,b5,b6,b7->bθ", core, input[0, :, 0, 0], input[1, :, 0, 0], input[0, :, 0, 1], input[1, :, 0, 1], input[0, :, 1, 0], input[1, :, 1, 0], input[0, :, 1, 1], input[1, :, 1, 1], ) assert torch.allclose(eps_result, oe_result)
def modified_llr_codeword(LLR_Info_bits): # Generator matrix of shape (m+1, 2^m) for RM(m, 1) code is needed here. So load it before hand # LLR_Info_bits is of shape (batch*num_sparse, m + 1) required_LLR_info = contract( 'ij , jk ->ikj', LLR_Info_bits, Generator_Matrix_cuda) # (batch*num_sparse, 2^m, m+1) sign_matrix = (-1)**((required_LLR_info < 0).sum(2)).float() # (batch*num_sparse, 2^m+1) min_abs_LLR_info, _ = torch.min(torch.where( required_LLR_info == 0., torch.max(required_LLR_info.abs()) + 1, required_LLR_info.abs()), dim=2) return sign_matrix * min_abs_LLR_info
def test_slicer(): eq, shapes = oe.helpers.rand_equation(30, reg=5, seed=42, d_max=3) arrays = [np.random.uniform(size=s) for s in shapes] path, info = oe.contract_path(eq, *shapes, shapes=True) expected = oe.contract(eq, *arrays, optimize=path) sf = ctg.SliceFinder(info, target_size=1_000_000, target_overhead=None) inds, ccost = sf.search() assert info.largest_intermediate > 1_000_000 assert ccost.size <= 1_000_000 assert ccost.total_flops > info.opt_cost assert len(inds) > 1 sc = sf.SlicedContractor(arrays) assert sc.total_flops == ccost.total_flops assert sc.contract_all() == pytest.approx(expected)
def T2eq_rhs_CC2(self, t1, t2, F): v = self.vir o = self.occ TEI = self.TEI fae = F[v, v].copy() fmi = F[o, o].copy() wabef_2 = self.Wabef_2(t1, t2, F) wmnij_2 = self.Wmnij_2(t1, t2, F) #All terms in the T2 Equation term1 = TEI[o, o, v, v].copy() term2tmp = fae #- 0.5 *contract('me,mb->be', fme, t1) term2a = contract('be,ijae->ijab', term2tmp, t2) term2 = term2a - term2a.swapaxes(2, 3) #swap ab term3temp = fmi #+ 0.5 *contract('me,je->mj', fme, t1) term3a = -contract('mj,imab->ijab', term3temp, t2) term3 = term3a - term3a.swapaxes(0, 1) #swap ij tau = contract('ma,nb->mnab', t1, t1) - contract('na,mb->mnab', t1, t1) term44 = 0.5 * contract('mnij,mnab->ijab', wmnij_2, tau) term55 = 0.5 * contract('abef,ijef->ijab', wabef_2, tau) term6tmp = -contract('mbej,ie,ma->ijab', TEI[o, v, v, o], t1, t1) term6 = term6tmp - term6tmp.swapaxes(2, 3) - term6tmp.swapaxes( 0, 1) + term6tmp.swapaxes(0, 1).swapaxes(2, 3) term7tmp = contract('abej,ie->ijab', TEI[v, v, v, o], t1) term7 = term7tmp - term7tmp.swapaxes(0, 1) #swap ij term8tmp = -contract('mbij,ma->ijab', TEI[o, v, o, o], t1) term8 = term8tmp - term8tmp.swapaxes(2, 3) #swap ab total = term1 + term2 + term3 + term44 + term55 + term6 + term7 + term8 return total
def LRWibjm(self, t1, t2, F): v = self.vir o = self.occ TEI = self.TEI Fme = self.Fme(t1, t2, F) Wmnij = self.LSWmnij(t1, t2, F) term1 = -0.5*TEI[o, v, o, o].copy() term2 = 0.5*contract('ie,jmbe->ibjm', Fme, t2) term3 = contract('injm,nb->ibjm', Wmnij, t1) term4a = -TEI[o, v, v, o].copy() - contract('inef,nmfb->ibem', TEI[o, o, v, v], t2) term4 = contract('ibem,je->ibjm', term4a, t1) tau = 0.25*t2 + 0.5*contract('ia,jb->ijab', t1, t1) #-contract('ib,ja->ijab', t1, t1) term5 = -contract('ibef,jmef->ibjm', TEI[o, v, v, v], tau) term6 = contract('inem,jneb->ibjm', TEI[o, o, v, o], t2) total = term1 + (term2 + term3 + term4 + term5 + term6) return total