env_sparse module

Full Documentation for hippynn.custom_kernels.env_sparse module. Click here for a summary page.

Pure pytorch implementation of envsum operations

envsum(sensitivities, features, pair_first, pair_second)

featsum(env, sense, pair_first, pair_second)

make_sparse_sense(sensitivities, pair_first, pair_second, n_atom: int)

Construct sensitivities as a sparse matrix with shape (n_atoms * n_nu, n_atoms).

The n_atoms * n_nu is needed because of limitations in the hybrid-sparse matrix multiply routines available in pytorch. This function is implemented seprately because both envsum and sensesum use it. The to_csr call is done afterwards because envsum needs to be transposed.

Parameters:

• sensitivities

• pair_first

• pair_second

• n_atom

Returns:

sensesum(env, features, pair_first, pair_second)

Sparse sensesum implementation uses a sparsity matrix of zeros with shape (n_atoms x n_atoms), combined with the crucial function torch.sparse sampled_addmm.

Parameters:

• env

• features

• pair_first

• pair_second

Returns: