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.