Package: torch-gauge
Torch-GaugeA light-weightPyTorchextension for efficient gauge-equivariant learning.AboutTorch-Gaugeis a library to boost geometric learning on physical data structures
with Lie-group symmetry and beyond two-body interactions. The library is designed to be specifically
optimized for training and inference on GPUs, with mini-batch training utilities natively supported.UsageTorch-Gauge uses paddedVerlet listas the core for
manipulating relational data, which enables compact and highly customizable implementation of
geometric learning models.As an illustration, check out our implementation ofSchNet's
interaction module in just 20 lines of code:importtorchfromtorch.nnimportLinear,Parameterfromtorch_gauge.verlet_listimportVerletListfromtorch_gauge.nnimportSSPclassSchNetLayer(torch.nn.Module):def__init__(self,num_features):super().__init__()_nf=num_featuresself.gamma=10.0self.rbf_centers=Parameter(torch.linspace(0.1,30.1,300),requires_grad=False)self.cfconv=torch.nn.Sequential(Linear(300,_nf,bias=False),SSP(),Linear(_nf,_nf,bias=False),SSP())self.pre_conv=Linear(_nf,_nf)self.post_conv=torch.nn.Sequential(Linear(_nf,_nf),SSP(),Linear(_nf,_nf))defforward(self,vl:VerletList,l:int):xyz=vl.ndata["xyz"]pre_conv=self.pre_conv(vl.ndata[f"atomic_{l}"])d_ij=(vl.query_src(xyz)-xyz.unsqueeze(1)).norm(dim=2,keepdim=True)filters=self.cfconv(torch.exp(-self.gamma*(d_ij-self.rbf_centers.view(1,1,-1)).pow(2)))conv_out=(filters*vl.query_src(pre_conv)*vl.edge_mask.unsqueeze(2)).sum(1)vl.ndata[f"atomic_{l+1}"]=vl.ndata[f"atomic_{l}"]+self.post_conv(conv_out)returnvl.ndata[f"atomic_{l+1}"]And that's it. With theSphericalTensorsupport we can even quickly extend it to SE(3) equivariant
convolutions without pain, similar to the building block ofNequIP:importtorchfromtorch_gauge.nnimportSSP,IELinfromtorch.nnimportLinear,Parameterfromtorch_gauge.verlet_listimportVerletListfromtorch_gauge.o3.sphericalimportSphericalTensorfromtorch_gauge.o3.rshimportRSHxyzfromtorch_gauge.geometricimportpoly_envfromtorch_gauge.o3.clebsch_gordanimportLeviCivitaCouplerclassSE3Layer(torch.nn.Module):def__init__(self,num_features):super().__init__()_nf=num_featuresself.rbf_freqs=Parameter(torch.arange(16),requires_grad=False)self.rsh_mod=RSHxyz(max_l=1)self.coupler=LeviCivitaCoupler(torch.LongTensor([_nf,_nf]))self.filter_gen=torch.nn.Sequential(Linear(16,_nf,bias=False),SSP(),Linear(_nf,_nf,bias=False))self.pre_conv=IELin([_nf,_nf],[_nf,_nf])self.post_conv=torch.nn.ModuleList([IELin([2*_nf,3*_nf],[_nf,_nf]),SSP(),IELin([_nf,_nf],[_nf,_nf])])defforward(self,vl:VerletList,l:int):r_ij=vl.query_src(vl.ndata["xyz"])-vl.ndata["xyz"].unsqueeze(1)d_ij=r_ij.norm(dim=2,keepdim=True)r_ij=r_ij/d_ijfeat_in:SphericalTensor=vl.ndata[f"atomic_{l}"]pre_conv=vl.query_src(self.pre_conv(feat_in))filters_radial=self.filter_gen(torch.sin(d_ij/5*self.rbf_freqs.view(1,1,-1))/d_ij*poly_env(d_ij/5))filters=pre_conv.self_like(self.rsh_mod(r_ij).ten.unsqueeze(-1).mul(filters_radial).flatten(2,3))coupling_out=self.post_conv[0](self.coupler(pre_conv,filters,overlap_out=False))conv_out=feat_in.self_like(coupling_out.ten.mul(vl.edge_mask.unsqueeze(2)).sum(1))conv_out=conv_out.scalar_mul(self.post_conv[1](conv_out.invariant()))conv_out.ten=self.post_conv[2](conv_out).tenvl.ndata[f"atomic_{l+1}"]=feat_in+conv_outreturnvl.ndata[f"atomic_{l+1}"]Torch-Gauge's Verlet list also supports generating higher-order views (triplets, quartets, etc.) to
streamline implementing models with specialized non-2body message passing algorithms.Conventional molecular descriptors that were commonly used in Kernel methods can also be revisited
by Torch-Gauge, and examples will be provided in future versions.SetupsFrom pippip install torch-gaugeFrom SourceOnce PyTorch is installed, runningpip install versioneer && versioneer install
pip install -e .Project Structuretorch_gaugeHigh-level python functionstorch_gauge/o3O(3) group algebra functionalstorch_gauge/nn.pyTensorial neural network building blockstorch_gauge/verlet_list.pyVerlet neighbor-list operations for representing relational datatorch_gauge/geometric.pycontains geometric and algebra operations with autogradtorch_gauge/modelscontains exemplary implementations of GNN variants and descriptorstorch_gauge/testscontains pytestssrccontains C++/CUDA backend operations (under development)Testmake testQuestions?Please submit a question/issue oremail me.