def test_solid_harmonic_scattering(): # Compare value to analytical formula in the case of a single Gaussian centers = torch.FloatTensor(1, 1, 3).fill_(0) weights = torch.FloatTensor(1, 1).fill_(1) sigma_gaussian = 3. sigma_0_wavelet = 3. M, N, O, J, L = 128, 128, 128, 1, 3 grid = torch.from_numpy( np.fft.ifftshift(np.mgrid[-M//2:-M//2+M, -N//2:-N//2+N, -O//2:-O//2+O].astype('float32'), axes=(1,2,3))) x = generate_weighted_sum_of_gaussians(grid, centers, weights, sigma_gaussian) scattering = HarmonicScattering3D(J=J, shape=(M, N, O), L=L, sigma_0=sigma_0_wavelet) scattering.max_order = 1 scattering.method = 'integral' scattering.integral_powers = [1] for device in devices: if device == 'cpu': x = x.cpu() scattering.cpu() else: x = x.cuda() scattering.cuda() s = scattering(x) for j in range(J+1): sigma_wavelet = sigma_0_wavelet*2**j k = sigma_wavelet / np.sqrt(sigma_wavelet**2 + sigma_gaussian**2) for l in range(1, L+1): err = torch.abs(s[0, j, l, 0] - k ** l).sum()/(1e-6+s[0, j, l, 0].abs().sum()) assert err<1e-4
def test_solid_harmonic_scattering(): # Compare value to analytical formula in the case of a single Gaussian centers = torch.FloatTensor(1, 1, 3).fill_(0) weights = torch.FloatTensor(1, 1).fill_(1) sigma_gaussian = 3. sigma_0_wavelet = 3. M, N, O, J, L = 128, 128, 128, 1, 3 grid = torch.from_numpy( np.fft.ifftshift(np.mgrid[-M//2:-M//2+M, -N//2:-N//2+N, -O//2:-O//2+O].astype('float32'), axes=(1,2,3))) x = generate_weighted_sum_of_gaussians(grid, centers, weights, sigma_gaussian) scattering = Scattering3D(M=M, N=N, O=O, J=J, L=L, sigma_0=sigma_0_wavelet) s = scattering(x, order_2=False, method='integral', integral_powers=[1]) for j in range(J+1): sigma_wavelet = sigma_0_wavelet*2**j k = sigma_wavelet / np.sqrt(sigma_wavelet**2 + sigma_gaussian**2) for l in range(1, L+1): err = torch.abs(s[0, 0, j, l] - k ** l).sum()/(1e-6+s[0, 0, j, l].abs().sum()) assert err<1e-4
def test_solid_harmonic_scattering(device, backend): if backend.name.endswith('_skcuda') and device == "cpu": pytest.skip("The skcuda backend does not support CPU tensors.") # Compare value to analytical formula in the case of a single Gaussian centers = np.zeros((1, 1, 3)) weights = np.ones((1, 1)) sigma_gaussian = 3. sigma_0_wavelet = 3. M, N, O, J, L = 128, 128, 128, 1, 3 grid = np.mgrid[-M // 2:-M // 2 + M, -N // 2:-N // 2 + N, -O // 2:-O // 2 + O] grid = grid.astype('float32') grid = np.fft.ifftshift(grid, axes=(1, 2, 3)) x = torch.from_numpy( generate_weighted_sum_of_gaussians(grid, centers, weights, sigma_gaussian)).to(device).float() scattering = HarmonicScattering3D(J=J, shape=(M, N, O), L=L, sigma_0=sigma_0_wavelet, max_order=1, method='integral', integral_powers=[1], frontend='torch', backend=backend).to(device) scattering.max_order = 1 scattering.method = 'integral' scattering.integral_powers = [1] s = scattering(x) for j in range(J + 1): sigma_wavelet = sigma_0_wavelet * 2**j k = sigma_wavelet / np.sqrt(sigma_wavelet**2 + sigma_gaussian**2) for l in range(1, L + 1): err = torch.abs(s[0, j, l, 0] - k**l).sum() / (1e-6 + s[0, j, l, 0].abs().sum()) assert err < 1e-4
def compute_qm7_solid_harmonic_scattering_coefficients(M=192, N=128, O=96, sigma=2., J=2, L=3, integral_powers=(0.5, 1., 2., 3.), batch_size=16): """ Computes the scattering coefficients of the molecules of the QM7 database. Channels used are full charges, valence charges and core charges. Linear regression of the qm7 energies with the given values gives MAE 2.75, RMSE 4.18 (kcal.mol-1). Parameters ---------- M, N, O: int dimensions of the numerical grid sigma : float width parameter of the Gaussian that represents a particle J: int maximal scale of the solid harmonic wavelets L: int maximal first order of the solid harmonic wavelets integral_powers: list of int powers for the integrals batch_size: int size of the batch for computations Returns ------- order_0: torch tensor array containing the order 0 scattering coefficients order_1: torch tensor array containing the order 1 scattering coefficients order_2: torch tensor array containing the order 2 scattering coefficients """ cuda = torch.cuda.is_available() grid = torch.from_numpy( np.fft.ifftshift(np.mgrid[-M // 2:-M // 2 + M, -N // 2:-N // 2 + N, -O // 2:-O // 2 + O].astype('float32'), axes=(1, 2, 3))) pos, full_charges, valence_charges = get_qm7_positions_and_charges(sigma) if cuda: grid = grid.cuda() pos = pos.cuda() full_charges = full_charges.cuda() valence_charges = valence_charges.cuda() n_molecules = pos.size(0) n_batches = np.ceil(n_molecules / batch_size).astype(int) scattering = Scattering3D(M=M, N=N, O=O, J=J, L=L, sigma_0=sigma) order_0, order_1, order_2 = [], [], [] print('Computing solid harmonic scattering coefficients of {} molecules ' 'of QM7 database on {}'.format(pos.size(0), 'GPU' if cuda else 'CPU')) print('sigma: {}, L: {}, J: {}, integral powers: {}'.format( sigma, L, J, integral_powers)) this_time = None last_time = None for i in range(n_batches): this_time = time.time() if last_time is not None: dt = this_time - last_time print("Iteration {} ETA: [{:02}:{:02}:{:02}]".format( i + 1, int(((n_batches - i - 1) * dt) // 3600), int((((n_batches - i - 1) * dt) // 60) % 60), int(((n_batches - i - 1) * dt) % 60)), end='\r') else: print("Iteration {} ETA: {}".format(i + 1, '-'), end='\r') last_time = this_time time.sleep(1) start, end = i * batch_size, min((i + 1) * batch_size, n_molecules) pos_batch = pos[start:end] full_batch = full_charges[start:end] val_batch = valence_charges[start:end] full_density_batch = generate_weighted_sum_of_gaussians(grid, pos_batch, full_batch, sigma, cuda=cuda) full_order_0 = compute_integrals(full_density_batch, integral_powers) full_order_1, full_order_2 = scattering( full_density_batch, order_2=True, method='integral', integral_powers=integral_powers) val_density_batch = generate_weighted_sum_of_gaussians(grid, pos_batch, val_batch, sigma, cuda=cuda) val_order_0 = compute_integrals(val_density_batch, integral_powers) val_order_1, val_order_2 = scattering(val_density_batch, order_2=True, method='integral', integral_powers=integral_powers) core_density_batch = full_density_batch - val_density_batch core_order_0 = compute_integrals(core_density_batch, integral_powers) core_order_1, core_order_2 = scattering( core_density_batch, order_2=True, method='integral', integral_powers=integral_powers) order_0.append( torch.stack([full_order_0, val_order_0, core_order_0], dim=-1)) order_1.append( torch.stack([full_order_1, val_order_1, core_order_1], dim=-1)) order_2.append( torch.stack([full_order_2, val_order_2, core_order_2], dim=-1)) order_0 = torch.cat(order_0, dim=0) order_1 = torch.cat(order_1, dim=0) order_2 = torch.cat(order_2, dim=0) return order_0, order_1, order_2
else: print("Iteration {} ETA: {}".format(i + 1, '-')) last_time = this_time time.sleep(1) # Extract the current batch. start = i * batch_size end = min(start + batch_size, n_molecules) pos_batch = pos[start:end] full_batch = full_charges[start:end] val_batch = valence_charges[start:end] # Calculate the density map for the nuclear charges and transfer # to PyTorch. full_density_batch = generate_weighted_sum_of_gaussians( grid, pos_batch, full_batch, sigma) full_density_batch = torch.from_numpy(full_density_batch) full_density_batch = full_density_batch.to(device).float() # Compute zeroth-order, first-order, and second-order scattering # coefficients of the nuclear charges. full_order_0 = compute_integrals(full_density_batch, integral_powers) full_scattering = scattering(full_density_batch) # Compute the map for valence charges. val_density_batch = generate_weighted_sum_of_gaussians( grid, pos_batch, val_batch, sigma) val_density_batch = torch.from_numpy(val_density_batch) val_density_batch = val_density_batch.to(device).float() # Compute scattering coefficients for the valence charges.