local_frames

get_rotation_matrix_from_basis(basis: Tensor) Tensor[source]

Returns the rotation matrix from the basis.

Parameters:

basis (Tensor) – The basis of the local frame.

Returns:

The rotation matrix to the new basis from the standard basis. Assumed to be multiplied as ‘vector @ rot_mat.T’ or ‘torch.matmul(rot_mat, vector)’.

Return type:

Tensor

local_frames_from_positions(pos: Tensor, neighbor1_pos: Tensor, neighbor2_pos: Tensor, neighbor3_pos: Tensor | None = None) Tensor[source]

Returns the local frames of the atoms in the molecule. Warning: This function does not use the third neighbor position.

Parameters:
  • pos (Tensor) – The position of the atom. Shape (…, 3)

  • neighbor1_pos (Tensor) – The position of the first neighbor atom. Shape (…, 3)

  • neighbor2_pos (Tensor) – The position of the second neighbor atom. Shape (…, 3)

  • neighbor3_pos (Optional[Tensor]) – The position of the third neighbor atom. Shape (…, 3). If given, the basis system will be chosen such that the third axis is oriented towards the third neighbor.

Returns:

The local frames of the atoms in the molecule.

Return type:

torch.Tensor

local_frames_from_rel_positions(diff_vec_1: Tensor, diff_vec_2: Tensor, diff_vec_3: Tensor | None = None) Tensor[source]

Returns the local frames of the atoms in the molecule.

In case the vectors are parallel, the second vector is changed to a random vector.

Parameters:
  • diff_vec_1 (Tensor) – The relative position of the first neighbor atom. Shape (…, 3)

  • diff_vec_2 (Tensor) – The relative position of the second neighbor atom. Shape (…, 3)

  • diff_vec_3 (Optional[Tensor]) – The relative position of the third neighbor atom. Shape (…, 3). If given, the basis system will be chosen such that the third axis is oriented towards the third neighbor.

Returns:

The local frames of the atoms in the molecule.

Return type:

torch.Tensor

Notes

In case the vectors are parallel, the second vector is changed to a random vector.

pyscf_to_e3nn_local_frames_matrix(basis: Tensor, irreps: Irreps) Tensor[source]

Calculates transformation matrix which can be used to transform a coefficient vector in the pyscf convention to a coefficient vector in local frames in the e3nn convention.

A coefficient vector can be transformed by multiplying it with this matrix on the left.

Parameters:
  • basis – The basis vectors of the local frame. Shape (3, 3).

  • irreps – The irreps of the coefficient vector in the pyscf convention.

Returns:

The transformation matrix. Shape (irreps.dim, irreps.dim).

Return type:

torch.Tensor

pyscf_to_e3nn_matrix(irreps: Irreps, dtype: dtype = torch.float64) Tensor[source]

Calculates the transformation matrix which can be used to transform a coefficient vector in the pyscf convention to a coefficient vector in the e3nn convention.

For details on the conventions, compare https://docs.e3nn.org/en/stable/guide/change_of_basis.html and https://pyscf.org/user/gto.html#ordering-of-basis-functions. The difference lies in the orientation of the spherical harmonics for l>1.

A coefficient vector can be transformed by multiplying it with this matrix on the left.

Parameters:
  • irreps – The irreps of the coefficient vector in the pyscf convention.

  • dtype – The dtype of the matrix.

Returns:

The transformation matrix. Shape (irreps.dim, irreps.dim).

Return type:

torch.Tensor

transform_coeffs_to_local(coeffs: Tensor, irreps: Irreps, basis: Tensor) Tensor[source]

Transforms the coefficients into the local frame.

Parameters:
  • coeffs (Tensor) – Coefficients of the atom.

  • irreps (Irreps) – Irreps of the atom.

  • basis (Tensor) – Local basis of the atom.

Returns:

The transformed coefficients in the local frame.

Return type:

Tensor

transform_local_coeffs_to_global(coeffs: Tensor, irreps: Irreps, basis: Tensor) Tensor[source]

Transforms the coefficients (in the e3nn convention) back into the standard basis.

Parameters:
  • coeffs (Tensor) – Coefficients of the atom.

  • irreps (Irreps) – Irreps of the atom.

  • basis (Tensor) – Local basis of the atom.

Returns:

The transformed coefficients into the local frame.

Return type:

Tensor