Skip to content

GCNN Improvements

Vicentini Filippo requested to merge github/fork/chrisrothUT/GCNN-Improvements into master

Created by: chrisrothUT

This PR builds upon the GCNN interface with the help of the symmetry interface of #724. This code has been written in collaboration with @attila-i-szabo.

Changes:

  1. GCNNs can now be implemented three different ways, with fast Fourier transforms over the translation group, with a discrete group Fourier transform, and by constructing the full kernel matrix.

The new versions are much faster. If we define F as the feature dimension, P as the number of point group elements, T as the number of translation group elements and d_i as the dimensions of the irreps of the space group (where \sum_i d_i^2 = PT) then we have the following scaling

Full Kernel: (FPT)^2 Irrep Fourier Transform: F (PT)^2 + F^2 \sum_i d^3 Fast Fourier Transfrom FP T log(T) + (FP)^2 T

So computing the full kernel is always the slowest, and choosing between the other two depends on how big the translation group is.

  1. GCNNs can be given a parity eigenvalue of +/-1 for two level systems (since parity always commutes with the space group)

  2. The GCNNs now take an optional character table that can be generated using #724

  3. GCNN/DenseSymm/DenseEquivariant now use constructor functions

I understand that we recently got rid of constructor functions, but with all of the added modalities it felt like a necessary evil. Additionally, I think it's generally good practice to avoid doing unnecessary computation in the setup of Flax Modules.

  1. GCNNs can be initialized with a Graph, PermutationGroup, or the correct arrays.

  2. There's an argument imag_part which lets you return just the imaginary part of the model (i.e. you want to pre-optimize the phases)

Example: GCNN on the triangular lattice projected onto the k=[0,0], E_2 irrep of the d6 point group

graph = nk.graph.Triangular([3,3])
characters = graph.space_group_builder().space_group_irreps([0,0])[5]
machine = nk.models.GCNN(symmetry_info=graph,mode="irrep",features=features,layers=4,parity=1,characters=characters,dtype=complex)

I'll be adding tests and finishing this up throughout the week

Merge request reports