AniSOAP descriptors vary very weakly with configuration of molecules

[elindgren]

Hi! I'm trying to use AniSOAP descriptors for representing trimer configurations of molecules, see example of a configuration in Fig 1. I've followed the tutorial with benzene molecules for how to create a suitable descriptor for my system. One difference is that I don't want a global descriptor, but rather a molecule-centered descriptor. Thus, I skip the averaging when computing the power spectrum, power_spectrum_tensor = calculator.power_spectrum(frames, mean_over_samples=False). I then compute these descriptors for a large set of trimers, and standardize them as in the tutorial using a StandardFlexibleScaler. However, when I then try to test the descriptor on a test trimer where I simply plot the change in the descriptor as a function of displacement of one of the molecules, I barely get any change. Specifically, I use the following similarity metric $\Gamma$ to compute the change:

$$ \Gamma = \exp{ ( - (d_{ref} - d_i)^2) }, $$

so a Gaussian similarity where $d_{ref}$ is the descriptor for the unperturbed trimer and $d_i$ is the descriptor for the perturbed system where one of the molecules have been translated. This similarity metric is essentially 1; the MSE in the exponent only changes by on the order of 1e-6. Note that both descriptors have been standardized using the scaler; if I don't standardize, the change is ~1e-12.

I suspect I'm doing something wrong, and I suspect it might be in how I compute the descriptor, but I cannot find what it is. Any help is gladly appreciated! Please find attached my code below.

Fig 1: Trimer of molecules

Fig 2: Similarity as a function of distance

def get_orientation_and_anisotropy(structure):
    # Get gyration tensor, in the lab frame (x, y, and z directions)
    G = get_gyration_tensor(structure)

    # Then extract the principal directions, in the lab frame
    eigenvalues, eigenvectors = np.linalg.eigh(G)

    # Finally, sort the eigenvectors according to the eigenvalues.
    # This interprets the eigenvectors as the rotation matrix
    # that rotates a molecule oriented with the long axis along
    # the x-axis, the short axis along the y-axis, and the
    # thickness along the z-axis to the current molecular orientation.

    idx = np.argsort(eigenvalues)[::-1]
    sorted_eigenvalues = eigenvalues[idx]
    eigenvectors = eigenvectors[:, idx]
    # Check the determinant: if negative == improper reordering of axis, flip direction of last eigenvector
    det = np.linalg.det(eigenvectors)
    if det < 0:
        eigenvectors[:, -1] *= -1

    # However, Anisoap defines rotations to transform from the body frame to the labe frame
    # Hence, we need to take the inverse = transpose of our rotation matrix
    orientation = Rotation.from_matrix(eigenvectors.T)

    # We interpret the (sorted) eigenvalues as the MVGs in the AniSOAP formalism,
    # and thus skip tuning them. However, since the molecules are point clouds,
    # if they are flat then the anisotropy in the thickness direction will be
    # essentially zero. Set the minimum value of the anisotropy to 0.5 Å, which
    # was used for the thickness of the benzene monomer in the AniSOAP preprint.
    sorted_eigenvalues = np.maximum(sorted_eigenvalues, 0.5)

    return orientation, sorted_eigenvalues

def get_anisoap(structure, cutoff: float, width=1.5, lmax=5, nmax=3):
    """
    Computes the AniSOAP descriptor for a cluster of molecules.
    """
    # Identify molecules
    molecules = get_molecules_without_centering(structure)
    # get COM, orientation, and MVG for each molecule
    coms = []
    quaternions = []
    ellipsoid_diameters = []
    atoms_in_each_molecule = []
    for i, molecule in enumerate(molecules):
        # First get orientation
        coms.append(molecule.get_center_of_mass())
        # Center molecule to ensure that no bonds are crossing the PBC boundaries
        centered_molecule = get_centered_molecule(molecule)
        
        R, mvg = get_orientation_and_anisotropy(centered_molecule)
        Q = R.as_quat(canonical=True, scalar_first=True)
        quaternions.append(Q)
        ellipsoid_diameters.append(mvg)  # sigma_x, sigma_y, sigma_z

        original_atom_indices = molecule.get_array("original_index")
        for original_atom_index in original_atom_indices:
            atoms_in_each_molecule.extend([i, original_atom_index])
    coms = np.array(coms)
    quaternions = np.array(quaternions)
    ellipsoid_diameters = np.array(ellipsoid_diameters)
    frame = Atoms(
        symbols="".join(["X" for _ in molecules]),
        positions=coms,
        cell=structure.cell,
        pbc=structure.pbc,
    )

    # Add the information to the Atoms object
    frame.arrays["c_q"] = quaternions
    frame.arrays["c_diameter[1]"] = ellipsoid_diameters[:, 0]
    frame.arrays["c_diameter[2]"] = ellipsoid_diameters[:, 1]
    frame.arrays["c_diameter[3]"] = ellipsoid_diameters[:, 2]

    frames = [frame]  # only one frame at a time
    # Then use AniSOAP to compute the descriptor

    AniSOAP_HYPERS = {
        "max_angular": lmax,
        "max_radial": nmax,
        "radial_basis_name": "gto",
        "subtract_center_contribution": True,
        "rotation_type": "quaternion",
        "rotation_key": "c_q",
        "cutoff_radius": float(cutoff),
        "radial_gaussian_width": float(width),
        "basis_rcond": 1e-8,
        "basis_tol": 1e-3,
    }
    calculator = EllipsoidalDensityProjection(**AniSOAP_HYPERS)
    power_spectrum_tensor = calculator.power_spectrum(frames, mean_over_samples=False)
    # Extract the anisoap vector with shape (n_molecules, n_dim)
    # The anisoap tensor has a single block
    block = power_spectrum_tensor[0]
    descriptor = block.values[
        :, 0, :
    ]  # Has shape N_molecules, spherical_harmonices_m, N_dim

    molecular_structure = frame
    return molecular_structure, descriptor, np.array(atoms_in_each_molecule).T

[elindgren]

I'm in particular skeptical against how I extract the descriptors from the TensorMap. Afaict the tensor only has one angular channel, but that seems weird to me?

power_spectrum_tensor = calculator.power_spectrum(frames, mean_over_samples=False)
  # Extract the anisoap vector with shape (n_molecules, n_dim)
  # The anisoap tensor has a single block
  block = power_spectrum_tensor[0]
  descriptor = block.values[
      :, 0, :
  ]  # Has shape N_molecules, spherical_harmonices_m, N_dim

[elindgren]

I'm not sure if the MVGs should be the eigenvalues or the square root of the eigenvalues of the gyration tensor, like below:

def get_orientation_and_anisotropy(structure):
    ...
    sigmas = np.sqrt(sorted_eigenvalues)
    mvgs = np.maximum(sigmas, 0.5)

    return orientation, mvgs

I've tried both approaches and both give insensitive descriptors.

[elindgren]

I tried rotating the same molecule that was displaced above around the central molecule in the trimer (called pitch in Fig 1), as well as rotating it around the axis connecting the center molecule and the molecule (Fig 2). In neither case AniSOAP gives a large difference in the computed descriptor. I suspect something is seriously wrong with my setup, but I don't know what it is.

Fig 1: Panning molecule around the center molecule

Fig 2: Rotating one of the molecules around the axis connecting it to the central molecule

[arthur-lin1027]

Hi @elindgren,

Thanks for bringing this issue up and sorry it took so long for me to respond! Can you please attach your datafile so that I can take a look at your entire pipeline?

I'm currently deep in the weeds of the internals, and I myself found a bug in our internal calculations. I will aim to do the following:

1. Run your code to verify your bug and to ensure nothing is wrong with your inputs.
2. Assuming your inputs are correct and the bug does indeed exist, see if my bugfix I make here alleviates your issues.

Thank you!

[elindgren]

Hi @arthur-lin1027! Please find attached the trimers used to generate the displacement plot above. Each atom has a label molecule corresponding to which of the molecules in each frame it corresponds to. You should then be able to directly replace the line

molecules = get_molecules_without_centering(structure)

with

molecules = [trimer[indices] for indices in [trimer.get_array('molecule') == i for i in range(3)]]

trimers-linear-displacement.xyz.zip

[arthur-lin1027]

Hi Eric,

I ran through your Trimers code and I was able to obtain varying descriptors. I have a few suggestions:

1. Once you generate your coarse-grained ellipsoidal frames, visualize them (I use the Ovito software) to ensure to ensure that the ellipsoids match up with your all atom representation.
2. Make sure you are using a large enough cutoff radius for neighbors to be able to see each other - a cutoff of 16.0 Angstroms (which is quite large), is necessary, especially for the 50th frame where the perturbed trimer's center of mass is quite far from the other two.

I used the following notebook (along with some code that is in the anisoap repo but not yet in main (it's in the branch CGRepHelper) to generate the coarse-grained beads and to run your similarity analysis. I also attach the coarse-grained beads xyz file for your convenience, although you'll almost certainly have to run the following code before putting it into the anisoap calculator:

cg_frames = ase.read("cg_frames.xyz", ":")
# Annoying, but ase automatically strips square brackets
for frame in cg_frames:
    frame.arrays["c_diameter[1]"] = frame.arrays["c_diameter1"]
    frame.arrays["c_diameter[2]"] = frame.arrays["c_diameter2"]
    frame.arrays["c_diameter[3]"] = frame.arrays["c_diameter3"]

Let me know how this works for you; I see different descriptors for each trimer, i.e. anisoap seems to be sensitive to distance perturbations.

analysis.ipynb

cg_frames.xyz.zip

[arthur-lin1027]

I'm not sure if the MVGs should be the eigenvalues or the square root of the eigenvalues of the gyration tensor, like below:

def get_orientation_and_anisotropy(structure):
    ...
    sigmas = np.sqrt(sorted_eigenvalues)
    mvgs = np.maximum(sigmas, 0.5)

    return orientation, mvgs

I've tried both approaches and both give insensitive descriptors.

From our experience, the specific dimensions don't matter that much, as long as they "make sense". See Figure 5 in the anisoap paper here: https://pubs.aip.org/aip/jcp/article/161/7/074112/3308992/Expanding-density-correlation-machine-learning.
In this figure, we vary the dimensions of benzene crystals, but we see that in general, as long as we are in the oblate ellipsoid regime (sigma1=sigma2 < sigma3), then we able to capture the atomistic details fairly well (what the Global Feature Reconstruction Errror i.e. GFRE describes)

[elindgren]

Thank you for the testing @arthur-lin1027! I've managed to identify several issues with my setup, including:

1. I use the gyration tensor, not the intertia tensor.
2. I reordered the eigenvectors of the gyration tensor according to the eigvenalues.
3. My semiaxes were incorrectly calculated.

I corrected these issues by following your implementation in CGRep.py. I also had an issue with how I computed my similarity meause; I averaged over the length of the vector, not summed (as np.linalg.norm would do).

All of these changes combine lead to a descriptor that varies more noticeably, but still weakly with the displacement. The main reason for my descriptor varying so weakly is my StandardScaler. I fit it on a large dataset of trimers (10k structures) with very varying configurations. As a result, the descriptors for all of these structures appear more similar than when fitting a scaler only on these displaced structures.

Btw, these two lines can be replaced with quat = R.from_matrix(evecs.T).as_quat(scalar_first=True) instead of np.roll.

AniSOAP/anisoap/asecg/CGRep.py

Lines 120 to 121 in e313ea1

	quat = R.from_matrix(evecs.T).as_quat() # IT's the ROWS! Not the columns, that are the vectors that we care about.
	quat = np.roll(quat, 1)

[arthur-lin1027]

Hi @elindgren,

I think everything you said above makes sense. Just to address some of your points above:

1. I'm not sure how different your semiaxes and orientation would be if you used switch between gyration tensor and moment of inertia tensor, because your molecular geometry is dominated by the planar arrangement of the carbons. I'm glad you figured it out though.
2. I think you have to be careful with how you standardize your features - I typically use StandardFlexibleScaler from the scikit-matter package, with column_wise=False (the default kwarg). On the other hand, StandardScaler from sklearn is equivalent to StandardFlexibleScaler with column_wise=True. Using StandardFlexibleScaler(column_wise=False) scales the entire matrix by the sum of the variance of each column, which preserves the scale between each feature (each column of the matrix). I would be curious to see how much your features differ with your gaussian similarity if you use StandardFlexibleScaler(column_wise=False). This requires installing skmatter
3. I think comparing the vector norm between two vectors makes more sense than averaging over the length of a vector as it represents more of a euclidean distance between two vectors.
4. Thanks for pointing out the new scaler_first=True keyword in the quaternion code - this was actually added fairly recently to scipy and I don't think I had that functionality earlier.

[elindgren]

@arthur-lin1027

1. Yeah this didn't make much of a difference, but I just wanted to get my approach in line with the "recommended" AniSOAP setup :). In addition to the switch from gyration -> inertia tensor, I also had made some errors in how I compute the semiaxes which I corrected.
2. Sorry I was unclear, I've been using the StandardFlexibleScaler from skmatter for all of my examples above. I didn't use the StandardScaler (or StandardFlexibleScaler(column_wise=True)) since there seems to be some columns that have zero variance in the AniSOAP descriptors.
3. Agreed 👍
4. 👌