def test_against_standard_computations():
file_path = os.path.abspath(os.path.dirname(__file__))
data = torch.load(os.path.join(file_path, 'test_data_3d.pt'))
x = data['x']
scattering_ref = data['Sx']
J = data['J']
L = data['L']
integral_powers = data['integral_powers']
M = x.shape[1]
batch_size = x.shape[0]
N, O = M, M
sigma = 1
scattering = Scattering3D(M=M, N=N, O=O, J=J, L=L, sigma_0=sigma)
for device in devices:
if device == 'cpu':
x = x.cpu()
else:
x = x.cuda()
order_0 = compute_integrals(x, integral_powers)
order_1, order_2 = scattering(x, order_2=True,
method='integral', integral_powers=integral_powers)
order_0 = order_0.cpu().numpy().reshape((batch_size, -1))
start = 0
end = order_0.shape[1]
order_0_ref = scattering_ref[:,start:end].numpy()
order_1 = order_1.cpu().numpy().reshape((batch_size, -1))
start = end
end += order_1.shape[1]
order_1_ref = scattering_ref[:, start:end].numpy()
order_2 = order_2.cpu().numpy().reshape((batch_size, -1))
start = end
end += order_2.shape[1]
order_2_ref = scattering_ref[:, start:end].numpy()
order_0_diff_cpu = relative_difference(order_0_ref, order_0)
order_1_diff_cpu = relative_difference(order_1_ref, order_1)
order_2_diff_cpu = relative_difference(order_2_ref, order_2)
assert order_0_diff_cpu < 1e-6, "CPU : order 0 do not match, diff={}".format(order_0_diff_cpu)
assert order_1_diff_cpu < 1e-6, "CPU : order 1 do not match, diff={}".format(order_1_diff_cpu)
assert order_2_diff_cpu < 1e-6, "CPU : order 2 do not match, diff={}".format(order_2_diff_cpu)
def test_against_standard_computations():
file_path = os.path.abspath(os.path.dirname(__file__))
d = np.load(os.path.join(file_path, 'test_data_3d.pt'))
x = d.item()['x']
scattering_ref = d.item()['scat']
M, N, O, J, L, sigma = 64, 64, 64, 2, 2, 1.
integral_powers = [1., 2.]
scattering = Scattering3D(M=M, N=N, O=O, J=J, L=L, sigma_0=sigma)
x = torch.from_numpy(x)
for device in devices:
if device == 'cpu':
x = x.cpu()
else:
x = x.cuda()
order_0 = compute_integrals(x, integral_powers)
order_1, order_2 = scattering(x, order_2=True,
method='integral', integral_powers=integral_powers)
order_0 = order_0.cpu().numpy().reshape((1, -1))
start = 0
end = order_0.shape[1]
order_0_ref = scattering_ref[:,start:end]
order_1 = order_1.cpu().numpy().reshape((1, -1))
start = end
end += order_1.shape[1]
order_1_ref = scattering_ref[:, start:end]
order_2 = order_2.cpu().numpy().reshape((1, -1))
start = end
end += order_2.shape[1]
order_2_ref = scattering_ref[:, start:end]
order_0_diff_cpu = relative_difference(order_0_ref, order_0)
order_1_diff_cpu = relative_difference(order_1_ref, order_1)
order_2_diff_cpu = relative_difference(order_2_ref, order_2)
assert order_0_diff_cpu < 1e-6, "CPU : order 0 do not match, diff={}".format(order_0_diff_cpu)
assert order_1_diff_cpu < 1e-6, "CPU : order 1 do not match, diff={}".format(order_1_diff_cpu)
assert order_2_diff_cpu < 1e-6, "CPU : order 2 do not match, diff={}".format(order_2_diff_cpu)
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
def test_against_standard_computations():
file_path = os.path.abspath(os.path.dirname(__file__))
data = torch.load(os.path.join(file_path, 'test_data_3d.pt'))
x = data['x']
scattering_ref = data['Sx']
J = data['J']
L = data['L']
integral_powers = data['integral_powers']
M = x.shape[1]
batch_size = x.shape[0]
N, O = M, M
sigma = 1
scattering = HarmonicScattering3D(J=J, shape=(M, N, O), L=L, sigma_0=sigma)
for device in devices:
if device == 'cpu':
x = x.cpu()
scattering.cpu()
else:
x = x.cuda()
scattering.cuda()
order_0 = compute_integrals(x, integral_powers)
scattering.max_order = 2
scattering.method = 'integral'
scattering.integral_powers = integral_powers
orders_1_and_2 = scattering(x)
# WARNING: These are hard-coded values for the setting J = 2.
n_order_1 = 3
n_order_2 = 3
# Extract orders and make order axis the slowest in accordance with
# the stored reference scattering transform.
order_1 = orders_1_and_2[:,0:n_order_1,...]
order_2 = orders_1_and_2[:,n_order_1:n_order_1+n_order_2,...]
# Permute the axes since reference has (batch index, integral power, j,
# ell) while the computed transform has (batch index, j, ell, integral
# power).
order_1 = order_1.permute(0, 3, 1, 2)
order_2 = order_2.permute(0, 3, 1, 2)
order_1 = order_1.reshape((batch_size, -1))
order_2 = order_2.reshape((batch_size, -1))
orders_1_and_2 = torch.cat((order_1, order_2), 1)
order_0 = order_0.cpu().numpy().reshape((batch_size, -1))
start = 0
end = order_0.shape[1]
order_0_ref = scattering_ref[:,start:end].numpy()
orders_1_and_2 = orders_1_and_2.cpu().numpy().reshape((batch_size, -1))
start = end
end += orders_1_and_2.shape[1]
orders_1_and_2_ref = scattering_ref[:, start:end].numpy()
order_0_diff_cpu = relative_difference(order_0_ref, order_0)
orders_1_and_2_diff_cpu = relative_difference(
orders_1_and_2_ref, orders_1_and_2)
assert order_0_diff_cpu < 1e-6, "CPU : order 0 do not match, diff={}".format(order_0_diff_cpu)
assert orders_1_and_2_diff_cpu < 1e-6, "CPU : orders 1 and 2 do not match, diff={}".format(orders_1_and_2_diff_cpu)
