### Install escnn from Git repository

Source: https://github.com/quva-lab/escnn/blob/master/README.md

This command installs escnn by cloning the latest code directly from the GitHub repository. This is useful for development or when the latest features are needed.

```bash
pip install git+https://github.com/QUVA-Lab/escnn
```

--------------------------------

### Install escnn using pip

Source: https://github.com/quva-lab/escnn/blob/master/README.md

This command installs the latest release of the escnn library from PyPI. Ensure you have Python 3.7+ and pip installed.

```bash
pip install escnn
```

--------------------------------

### Citation for E(N)-Equivariant Steerable CNNs

Source: https://github.com/quva-lab/escnn/blob/master/README.md

BibTeX entry for the paper 'A Program to Build E(N)-Equivariant Steerable CNNs', which is a foundational work for the escnn library.

```bibtex
@inproceedings{cesa2022a,
        title={A Program to Build {E(N)}-Equivariant Steerable {CNN}s },
        author={Gabriele Cesa and Leon Lang and Maurice Weiler},
        booktitle={International Conference on Learning Representations},
        year={2022},
        url={https://openreview.net/forum?id=WE4qe9xlnQw}
    }
```

--------------------------------

### Citation for General E(2)-Equivariant Steerable CNNs

Source: https://github.com/quva-lab/escnn/blob/master/README.md

BibTeX entry for the paper 'General E(2)-Equivariant Steerable CNNs', which is a key publication related to the escnn library and its development.

```bibtex
@inproceedings{e2cnn,
       title={{General E(2)-Equivariant Steerable CNNs}},
       author={Weiler, Maurice and Cesa, Gabriele},
       booktitle={Conference on Neural Information Processing Systems (NeurIPS)},
       year={2019},
       url={https://arxiv.org/abs/1911.08251}
   }
```

--------------------------------

### Initialize DataLoader for Training and Testing

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

Creates PyTorch DataLoader instances for both training and testing datasets. These DataLoaders efficiently manage batching and shuffling of data, which is crucial for the training loop.

```python
mnist_train = MnistRotDataset(mode='train', transform=train_transform)
train_loader = torch.utils.data.DataLoader(mnist_train, batch_size=64)

mnist_test = MnistRotDataset(mode='test', transform=test_transform)
test_loader = torch.utils.data.DataLoader(mnist_test, batch_size=64)
```

--------------------------------

### Running e2wrn.py for Equivariant Wide Resnet Models

Source: https://github.com/quva-lab/escnn/blob/master/README.md

Executes the e2wrn.py script to demonstrate complex 2D equivariant Wide Resnet models. This command navigates to the examples directory and runs the script, showcasing models that are equivariant under reflections.

```bash
cd examples
python e2wrn.py
```

--------------------------------

### Set up Loss Function and Optimizer

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

Initializes the loss function (CrossEntropyLoss) and the optimizer (Adam) for model training. These are standard components for training deep learning models, configured with specific learning rates and weight decay.

```python
loss_function = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=5e-5, weight_decay=1e-5)
```

--------------------------------

### Instantiate SO3MLPtensor Model on CPU

Source: https://github.com/quva-lab/escnn/blob/master/examples/mlp.ipynb

Demonstrates how to instantiate the SO3MLPtensor model and move it to the CPU. This is a common setup step before model training or evaluation. It checks for CUDA availability but defaults to CPU.

```python
device = 'cpu' # 'cuda' if torch.cuda.is_available() else 'cpu'
model = SO3MLPtensor().to(device)
```

--------------------------------

### Run Equivariant Wide Resnet Model (Python)

Source: https://github.com/quva-lab/escnn/blob/master/examples/README.md

Demonstrates how to execute an equivariant Wide Resnet model from the examples directory. This script can be run with optional arguments to enable specific equivariance properties, such as rotational symmetry.

```shell
cd examples
python e2wrn.py
```

```shell
python e2wrn.py --rot90
```

--------------------------------

### Running e2wrn.py with --rot90 flag

Source: https://github.com/quva-lab/escnn/blob/master/README.md

Executes the e2wrn.py script with the --rot90 flag to run a version of the Wide Resnet model that is simultaneously equivariant under reflections and rotations by multiples of 90 degrees.

```bash
python e2wrn.py --rot90
```

--------------------------------

### Import necessary libraries for E(n)-Equivariant CNNs

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

Imports the required modules from the `sys` and `escnn` libraries to build E(n)-Equivariant Steerable CNNs. This includes `gspaces` for defining the symmetry group and `nn` for neural network modules.

```python
import sys
sys.path.append('../')

import torch

from escnn import gspaces
from escnn import nn
```

--------------------------------

### Define Training Data Transformations

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

Sets up the data transformations for the training dataset. It includes padding, resizing, random rotation, and conversion to a tensor, essential for preparing data for an ESCNN model that handles rotation.

```python
train_transform = Compose([
    pad,
    resize1,
    RandomRotation(180., interpolation=InterpolationMode.BILINEAR, expand=False),
    resize2,
    totensor,
])
```

--------------------------------

### Download and Uncompress MNIST Rotation Dataset

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

This snippet demonstrates how to download the rotated MNIST dataset using wget and then uncompress it using the unzip command. It's typically used for preparing the dataset before training or testing models.

```bash
# download the dataset
!wget -nc http://www.iro.umontreal.ca/~lisa/icml2007data/mnist_rotation_new.zip
# uncompress the zip file
!unzip -n mnist_rotation_new.zip -d mnist_rotation_new
```

--------------------------------

### Retrieve Raw Test Data

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

Fetches the first image and its corresponding label from the raw test dataset iterator. This is useful for direct inspection or debugging of individual data samples.

```python
# retrieve the first image from the test set
x, y = next(iter(raw_mnist_test))
```

--------------------------------

### Run 3D CNN Equivariant to Rotations and Translations (Python)

Source: https://github.com/quva-lab/escnn/blob/master/examples/README.md

Shows how to run a 3D Convolutional Neural Network (CNN) example that is equivariant to rotations and translations in 3D space. This is useful for applications dealing with 3D volumetric data.

```shell
cd examples
python se3_3Dcnn.py
```

--------------------------------

### Model Training Loop

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

Implements the training loop for the ESCNN model over a specified number of epochs. It includes forward and backward passes, optimizer steps, and periodic evaluation of test accuracy.

```python
for epoch in range(31):
    model.train()
    for i, (x, t) in enumerate(train_loader):
        
        optimizer.zero_grad()

        x = x.to(device)
        t = t.to(device)

        y = model(x)

        loss = loss_function(y, t)

        loss.backward()

        optimizer.step()
    
    if epoch % 10 == 0:
        total = 0
        correct = 0
        with torch.no_grad():
            model.eval()
            for i, (x, t) in enumerate(test_loader):

                x = x.to(device)
                t = t.to(device)
                
                y = model(x)

                _, prediction = torch.max(y.data, 1)
                total += t.shape[0]
                correct += (prediction == t).sum().item()
        print(f"epoch {epoch} | test accuracy: {correct/total*100.}")
```

--------------------------------

### Equivariant Convolution on RGB Images with escnn

Source: https://github.com/quva-lab/escnn/blob/master/README.md

Performs an equivariant convolution on an RGB image using the escnn library. It defines a C8 rotation symmetry group, input/output feature types, and applies an R2Conv layer followed by a ReLU activation. The input is a random PyTorch tensor representing an RGB image batch, which is then wrapped into an escnn GeometricTensor.

```python3
from escnn import gspaces
from escnn import nn
import torch

r2_act = gspaces.rot2dOnR2(N=8)
feat_type_in  = nn.FieldType(r2_act,  3*[r2_act.trivial_repr])
feat_type_out = nn.FieldType(r2_act, 10*[r2_act.regular_repr])

conv = nn.R2Conv(feat_type_in, feat_type_out, kernel_size=5)
relu = nn.ReLU(feat_type_out)

x = torch.randn(16, 3, 32, 32)
x = feat_type_in(x)

y = relu(conv(x))
```

--------------------------------

### Build Equivariant 2D Convolutional Layer (escnn)

Source: https://context7.com/quva-lab/escnn/llms.txt

Provides an example of constructing an E(n)-equivariant 2D convolutional layer using escnn.nn.R2Conv. It shows how to define input and output feature types and demonstrates applying the convolution and verifying its equivariance property.

```python
from escnn import nn, gspaces
import torch

r2_act = gspaces.rot2dOnR2(N=8)
feat_type_in = nn.FieldType(r2_act, 3*[r2_act.trivial_repr])  # RGB input
feat_type_out = nn.FieldType(r2_act, 16*[r2_act.regular_repr])

# Create C8-equivariant convolution
conv = nn.R2Conv(feat_type_in, feat_type_out, kernel_size=7, padding=3, stride=1, bias=True)

# Wrap input tensor as geometric tensor
x = torch.randn(8, 3, 64, 64)  # batch of 8 RGB images
x_geo = feat_type_in(x)

# Apply convolution (output transforms correctly under rotations)
y = conv(x_geo)
print(f"Input: {x_geo.shape}, Output: {y.shape}")  # torch.Size([8, 3, 64, 64]) -> torch.Size([8, 128, 64, 64])

# Verify equivariance: conv(rotate(x)) == rotate(conv(x))
g = r2_act.fibergroup.element(1)  # 45-degree rotation
x_rotated = x_geo.transform(g)
y_from_rotated = conv(x_rotated)
y_rotated = y.transform(g)
assert torch.allclose(y_from_rotated.tensor, y_rotated.tensor, atol=1e-5)
```

--------------------------------

### Define Test Data Transformations

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

Defines the data transformations for the test dataset. This sequence includes padding and conversion to a tensor, used for evaluating the model's performance on unseen data without augmentation.

```python
test_transform = Compose([
    pad,
    totensor,
])
```

--------------------------------

### Test Model with Input Data

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

This function call evaluates the performance of a pre-trained 'model' using input data 'x'. It is typically used to assess the model's predictions or behavior on a given dataset.

```python
test_model(model, x)
```

--------------------------------

### Apply Equivariant Nonlinearities (escnn)

Source: https://context7.com/quva-lab/escnn/llms.txt

Shows how to apply equivariant non-linear activation functions within escnn models. This example defines feature types and prepares for applying nonlinearities, though the specific activation functions are not detailed here.

```python
from escnn import nn, gspaces
import torch

r2_act = gspaces.rot2dOnR2(N=4)
regular_field = nn.FieldType(r2_act, 8*[r2_act.regular_repr])
vector_field = nn.FieldType(r2_act, 4*[r2_act.irrep(1)])
```

--------------------------------

### Model Evaluation Setup in Python

Source: https://github.com/quva-lab/escnn/blob/master/examples/mlp.ipynb

Sets the model to evaluation mode and defines a batch size for testing. This is crucial for disabling training-specific layers like dropout and batch normalization, ensuring consistent results during inference.

```python
np.set_printoptions(linewidth=10000, precision=4, suppress=True)

model.eval()

B = 10
```

--------------------------------

### Initialize and Test ESCNN Model on Rotated MNIST

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

This section covers initializing the ESCNN model and defining a function to test its performance on rotated versions of input images. The `test_model` function evaluates the model's output logits for various rotation angles.

```python
import torch

# Assuming C8SteerableCNN is defined elsewhere
# from escnn.models import C8SteerableCNN 

# Placeholder for the actual model definition
class C8SteerableCNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        # Dummy layers for demonstration
        self.conv1 = torch.nn.Conv2d(1, 16, 3, 1)
        self.fc1 = torch.nn.Linear(16 * 27 * 27, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = torch.flatten(x, 1)
        x = self.fc1(x)
        return x

model = C8SteerableCNN().to(device)

def test_model(model: torch.nn.Module, x: Image):
    np.set_printoptions(linewidth=10000)
    
    # evaluate the `model` on 8 rotated versions of the input image `x`
    model.eval()
    
    x = resize1(pad(x)) # Using transform defined in MnistRotDataset snippet
    
    print()
    print('##########################################################################################')
    header = 'angle |  ' + '  '.join(["{:<6d}".format(d) for d in range(10)])
    print(header)
    with torch.no_grad():
        for r in range(8):
            # Assuming totensor and resize2 are defined in MnistRotDataset snippet
            x_transformed = totensor(resize2(x.rotate(r*45., Image.BILINEAR))).reshape(1, 1, 29, 29)
            x_transformed = x_transformed.to(device)

            y = model(x_transformed)
            y = y.to('cpu').numpy().squeeze()
            
            angle = r * 45
            print("{:5d} : {}".format(angle, y))
    print('##########################################################################################')
    print()

# build the test set    
raw_mnist_test = MnistRotDataset(mode='test')

# retrieve the first image from the test set
x, y = next(iter(raw_mnist_test))
```

--------------------------------

### Define Custom MNIST Rotated Dataset in PyTorch

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

Implements a custom PyTorch Dataset for the rotated MNIST dataset. It handles loading data from a .amat file, separating images and labels, and applying specified transformations.

```python
from torch.utils.data import Dataset
from torchvision.transforms import RandomRotation
from torchvision.transforms import Pad
from torchvision.transforms import Resize
from torchvision.transforms import ToTensor
from torchvision.transforms import Compose
from torchvision.transforms import InterpolationMode

import numpy as np

from PIL import Image

device = 'cuda' if torch.cuda.is_available() else 'cpu'

class MnistRotDataset(Dataset):
    
    def __init__(self, mode, transform=None):
        assert mode in ['train', 'test']
            
        if mode == "train":
            file = "mnist_rotation_new/mnist_all_rotation_normalized_float_train_valid.amat"
        else:
            file = "mnist_rotation_new/mnist_all_rotation_normalized_float_test.amat"
        
        self.transform = transform

        data = np.loadtxt(file, delimiter=' ')
            
        self.images = data[:, :-1].reshape(-1, 28, 28).astype(np.float32)
        self.labels = data[:, -1].astype(np.int64)
        self.num_samples = len(self.labels)
    
    def __getitem__(self, index):
        image, label = self.images[index], self.labels[index]
        image = Image.fromarray(image, mode='F')
        if self.transform is not None:
            image = self.transform(image)
        return image, label
    
    def __len__(self):
        return len(self.labels)

# images are padded to have shape 29x29.
# this allows to use odd-size filters with stride 2 when downsampling a feature map in the model
pad = Pad((0, 0, 1, 1), fill=0)

# to reduce interpolation artifacts (e.g. when testing the model on rotated images),
# we upsample an image by a factor of 3, rotate it and finally downsample it again
resize1 = Resize(87)
resize2 = Resize(29)

totensor = ToTensor()
```

--------------------------------

### Build 3D SE(3)-Equivariant ResNet Block with ESCNN

Source: https://context7.com/quva-lab/escnn/llms.txt

This Python code defines an SE(3)-equivariant ResNet block using ESCNN's R3Conv and IIDBatchNorm3d. It builds feature types using spherical harmonics, constructs the ResNet block with convolutional layers, batch normalization, and ReLU activations, and includes a skip connection. The example demonstrates building and testing the block.

```python
from escnn import nn, gspaces
import torch

r3_act = gspaces.rot3dOnR3()
so3 = r3_act.fibergroup

# Build feature types with spherical harmonics up to L=2
irreps_L2 = so3.bl_irreps(L=2)  # [L=0, L=1, L=2]
feat_in = nn.FieldType(r3_act, 8*[r3_act.irrep(0)] + 4*[r3_act.irrep(1)])
feat_hid = nn.FieldType(r3_act, 16*[r3_act.irrep(0)] + 8*[r3_act.irrep(1)] + 4*[r3_act.irrep(2)])

# Equivariant ResNet block
class SE3ResBlock(nn.EquivariantModule):
    def __init__(self, in_type, hidden_type, out_type):
        super().__init__()
        self.in_type = in_type
        self.out_type = out_type

        self.conv1 = nn.R3Conv(in_type, hidden_type, kernel_size=3, padding=1, bias=False)
        self.bn1 = nn.IIDBatchNorm3d(hidden_type, affine=True)
        self.relu1 = nn.ReLU(hidden_type)

        self.conv2 = nn.R3Conv(hidden_type, out_type, kernel_size=3, padding=1, bias=False)
        self.bn2 = nn.IIDBatchNorm3d(out_type, affine=True)

        if in_type != out_type:
            self.skip = nn.R3Conv(in_type, out_type, kernel_size=1, bias=False)
        else:
            self.skip = nn.IdentityModule(in_type)

        self.relu_out = nn.ReLU(out_type)

    def forward(self, x):
        identity = self.skip(x)
        out = self.relu1(self.bn1(self.conv1(x)))
        out = self.bn2(self.conv2(out))
        return self.relu_out(out + identity)

    def evaluate_output_shape(self, input_shape):
        return input_shape[:1] + (self.out_type.size,) + input_shape[2:]

# Build and test
resblock = SE3ResBlock(feat_in, feat_hid, feat_in)
x = feat_in(torch.randn(2, 44, 16, 16, 16))  # 8*1 + 4*3 + 16*1 + 8*3 = 44 channels
y = resblock(x)
print(f"SE3 ResBlock: {x.shape} -> {y.shape}")
```

--------------------------------

### Define C8SteerableCNN Model Architecture in Python

Source: https://github.com/quva-lab/escnn/blob/master/examples/model.ipynb

This Python code defines the C8SteerableCNN class, inheriting from torch.nn.Module. It sets up equivariant convolutional layers, batch normalization, ReLU activation, pooling, and a fully connected layer. The input is expected to be a PyTorch Tensor representing an image, which is then wrapped into a GeometricTensor for processing.

```python
class C8SteerableCNN(torch.nn.Module):
    
    def __init__(self, n_classes=10):
        
        super(C8SteerableCNN, self).__init__()
        
        # the model is equivariant under rotations by 45 degrees, modelled by C8
        self.r2_act = gspaces.rot2dOnR2(N=8)
        
        # the input image is a scalar field, corresponding to the trivial representation
        in_type = nn.FieldType(self.r2_act, [self.r2_act.trivial_repr])
        
        # we store the input type for wrapping the images into a geometric tensor during the forward pass
        self.input_type = in_type
        
        # convolution 1
        # first specify the output type of the convolutional layer
        # we choose 24 feature fields, each transforming under the regular representation of C8
        out_type = nn.FieldType(self.r2_act, 24*[self.r2_act.regular_repr])
        self.block1 = nn.SequentialModule(
            nn.MaskModule(in_type, 29, margin=1),
            nn.R2Conv(in_type, out_type, kernel_size=7, padding=1, bias=False),
            nn.InnerBatchNorm(out_type),
            nn.ReLU(out_type, inplace=True)
        )
        
        # convolution 2
        # the old output type is the input type to the next layer
        in_type = self.block1.out_type
        # the output type of the second convolution layer are 48 regular feature fields of C8
        out_type = nn.FieldType(self.r2_act, 48*[self.r2_act.regular_repr])
        self.block2 = nn.SequentialModule(
            nn.R2Conv(in_type, out_type, kernel_size=5, padding=2, bias=False),
            nn.InnerBatchNorm(out_type),
            nn.ReLU(out_type, inplace=True)
        )
        self.pool1 = nn.SequentialModule(
            nn.PointwiseAvgPoolAntialiased(out_type, sigma=0.66, stride=2)
        )
        
        # convolution 3
        # the old output type is the input type to the next layer
        in_type = self.block2.out_type
        # the output type of the third convolution layer are 48 regular feature fields of C8
        out_type = nn.FieldType(self.r2_act, 48*[self.r2_act.regular_repr])
        self.block3 = nn.SequentialModule(
            nn.R2Conv(in_type, out_type, kernel_size=5, padding=2, bias=False),
            nn.InnerBatchNorm(out_type),
            nn.ReLU(out_type, inplace=True)
        )
        
        # convolution 4
        # the old output type is the input type to the next layer
        in_type = self.block3.out_type
        # the output type of the fourth convolution layer are 96 regular feature fields of C8
        out_type = nn.FieldType(self.r2_act, 96*[self.r2_act.regular_repr])
        self.block4 = nn.SequentialModule(
            nn.R2Conv(in_type, out_type, kernel_size=5, padding=2, bias=False),
            nn.InnerBatchNorm(out_type),
            nn.ReLU(out_type, inplace=True)
        )
        self.pool2 = nn.SequentialModule(
            nn.PointwiseAvgPoolAntialiased(out_type, sigma=0.66, stride=2)
        )
        
        # convolution 5
        # the old output type is the input type to the next layer
        in_type = self.block4.out_type
        # the output type of the fifth convolution layer are 96 regular feature fields of C8
        out_type = nn.FieldType(self.r2_act, 96*[self.r2_act.regular_repr])
        self.block5 = nn.SequentialModule(
            nn.R2Conv(in_type, out_type, kernel_size=5, padding=2, bias=False),
            nn.InnerBatchNorm(out_type),
            nn.ReLU(out_type, inplace=True)
        )
        
        # convolution 6
        # the old output type is the input type to the next layer
        in_type = self.block5.out_type
        # the output type of the sixth convolution layer are 64 regular feature fields of C8
        out_type = nn.FieldType(self.r2_act, 64*[self.r2_act.regular_repr])
        self.block6 = nn.SequentialModule(
            nn.R2Conv(in_type, out_type, kernel_size=5, padding=1, bias=False),
            nn.InnerBatchNorm(out_type),
            nn.ReLU(out_type, inplace=True)
        )
        self.pool3 = nn.PointwiseAvgPoolAntialiased(out_type, sigma=0.66, stride=1, padding=0)
        
        self.gpool = nn.GroupPooling(out_type)
        
        # number of output channels
        c = self.gpool.out_type.size
        
        # Fully Connected
        self.fully_net = torch.nn.Sequential(
            torch.nn.Linear(c, 64),
            torch.nn.BatchNorm1d(64),
            torch.nn.ELU(inplace=True),
            torch.nn.Linear(64, n_classes),
        )
    
    def forward(self, input: torch.Tensor):
        # wrap the input tensor in a GeometricTensor
        # (associate it with the input type)
        x = nn.GeometricTensor(input, self.input_type)
        
        # apply each equivariant block
        
        # Each layer has an input and an output type
        # A layer takes a GeometricTensor in input.

```

--------------------------------

### Import escnn Libraries and Utilities

Source: https://github.com/quva-lab/escnn/blob/master/examples/introduction.ipynb

Imports core PyTorch functionality along with essential components from the escnn library for defining equivariant spaces and neural network layers. It also configures NumPy for precise and readable numerical output, which is useful for debugging and analysis.

```python
import torch

from escnn import gspaces
from escnn import nn

import numpy as np
np.set_printoptions(precision=3, linewidth=10000, suppress=True)
```

--------------------------------

### Model Initialization and Device Placement in Python

Source: https://github.com/quva-lab/escnn/blob/master/examples/mlp.ipynb

Initializes the SO3MLP model and places it on the appropriate device (CUDA if available, otherwise CPU). This is a standard practice for optimizing model performance by utilizing GPU acceleration when possible.

```python
device = 'cuda' if torch.cuda.is_available() else 'cpu'
model = SO3MLP().to(device)
```

--------------------------------

### Run an Input Through the Sequential Equivariant Model

Source: https://github.com/quva-lab/escnn/blob/master/examples/introduction.ipynb

This snippet shows how to prepare an input tensor (`x`) with the correct feature type and then pass it through the previously defined `SequentialModule` (`model`). The output `y` will inherit the output feature type of the last layer.

```python
x = torch.randn(1, 1, 17, 17)
x = feat_type_in(x)

y = model(x)
```

--------------------------------

### Build and Export Equivariant Model (Python)

Source: https://github.com/quva-lab/escnn/blob/master/docs/source/api/escnn.nn.rst

Demonstrates building a simple equivariant neural network using SequentialModule and then exporting it to a pure PyTorch module. It showcases the transformation of the network structure and verifies the equivalence of outputs between the equivariant and exported models.

```python
import escnn
import torch
from escnn.nn import SequentialModule, R2Conv, InnerBatchNorm, ReLU, PointwiseMaxPool, ELU, GroupPooling, FieldType, GeometricTensor

# Build a simple equivariant model using a SequentialModule
s = escnn.gspaces.rot2dOnR2(8)
c_in = FieldType(s, [s.trivial_repr]*3)
c_hid = FieldType(s, [s.regular_repr]*3)
c_out = FieldType(s, [s.regular_repr]*1)

net = SequentialModule(
    R2Conv(c_in, c_hid, 5, bias=False),
    InnerBatchNorm(c_hid),
    ReLU(c_hid, inplace=True),
    PointwiseMaxPool(c_hid, kernel_size=3, stride=2, padding=1),
    R2Conv(c_hid, c_out, 3, bias=False),
    InnerBatchNorm(c_out),
    ELU(c_out, inplace=True),
    GroupPooling(c_out)
)

# Export the model
net.eval()
net_exported = net.export()

print(net)
print(net_exported)

# Check that the two models are equivalent
x = torch.randn(10, c_in.size, 31, 31)
x = GeometricTensor(x, c_in)

y1 = net(x).tensor
y2 = net_exported(x.tensor)

assert torch.allclose(y1, y2)

```

--------------------------------

### SO2MLP Initialization and Layers

Source: https://github.com/quva-lab/escnn/blob/master/examples/mlp.ipynb

This snippet details the initialization of the SO2MLP class, including setting up the SO(2) group, defining the gspace, and configuring input types. It illustrates the construction of sequential blocks (block1, block2, block3) each comprising an equivariant linear layer, batch normalization, and a FourierELU activation function with specified channels and frequency band limits.

```python
class SO2MLP(nn.EquivariantModule):
    
    def __init__(self, n_classes=10):
        
        super(SO2MLP, self).__init__()
        
        # the model is equivariant to the group O(2)
        self.G = group.so2_group()
        
        # since we are building an MLP, there is no base-space
        self.gspace = gspaces.no_base_space(self.G)
        
        # the input contains the coordinates of a point in the 2D space
        self.in_type = self.gspace.type(self.G.standard_representation())
        
        # Layer 1
        # We will use the regular representation of SO(2) acting on signals over SO(2) itself, bandlimited to frequency 1
        # Most of the comments on the previous SO(3) network apply here as well
       
        activation1 = nn.FourierELU(
            self.gspace,
            channels=3, # specify the number of signals in the output features
            irreps=self.G.bl_regular_representation(L=1).irreps, # include all frequencies up to L=1
            inplace=True,
            # the following kwargs are used to build a discretization of the circle containing 6 equally distributed points
            type='regular', N=6,   
        )
        
        # map with an equivariant Linear layer to the input expected by the activation function, apply batchnorm and finally the activation
        self.block1 = nn.SequentialModule(
            nn.Linear(self.in_type, activation1.in_type),
            nn.IIDBatchNorm1d(activation1.in_type),
            activation1,
        )
        
        # Repeat a similar process for a few layers
        
        # 8 signals, bandlimited up to frequency 3
        activation2 = nn.FourierELU(
            self.gspace,
            channels=8, # specify the number of signals in the output features
            irreps=self.G.bl_regular_representation(L=3).irreps, # include all frequencies up to L=3
            inplace=True,
            # the following kwargs are used to build a discretization of the circle containing 16 equally distributed points
            type='regular', N=16,
        )
        self.block2 = nn.SequentialModule(
            nn.Linear(self.block1.out_type, activation2.in_type),
            nn.IIDBatchNorm1d(activation2.in_type),
            activation2,
        )
        
        # 8 signals, bandlimited up to frequency 3
        activation3 = nn.FourierELU(
            self.gspace,
            channels=8, # specify the number of signals in the output features
            irreps=self.G.bl_regular_representation(L=3).irreps, # include all frequencies up to L=3
            inplace=True,
            # the following kwargs are used to build a discretization of the circle containing 16 equally distributed points
            type='regular', N=16,
        )
        self.block3 = nn.SequentialModule(
            nn.Linear(self.block2.out_type, activation3.in_type),
            nn.IIDBatchNorm1d(activation3.in_type),
            activation3,
        )
        
        # 5 signals, bandlimited up to frequency 2
        activation4 = nn.FourierELU(
            self.gspace,
            channels=5, # specify the number of signals in the output features

```

--------------------------------

### SO2MLP Instantiation and Equivariance Testing

Source: https://github.com/quva-lab/escnn/blob/master/examples/mlp.ipynb

Demonstrates how to instantiate the SO2MLP model and test its equivariance property. It generates random input data, passes it through the model, and then applies random rotations to both the input and output to verify that the model's output transforms correctly according to group theory principles.

```python
import torch
import numpy as np
from e3nn import nn

# Assuming SO2MLP class is defined as above
# device = 'cuda' if torch.cuda.is_available() else 'cpu'
device = 'cpu'

model = SO2MLP().to(device)

np.set_printoptions(linewidth=10000, precision=4, suppress=True)

model.eval()

B = 10

# generates B random points in 2D and wrap them in a GeometricTensor of the right type
x = model.in_type(torch.randn(B, 2))


print('##########################################################################################')
with torch.no_grad():
    y = model(x.to(device)).to('cpu')
    print("Outputs' magnitudes")
    print(torch.linalg.norm(y.tensor, dim=1).numpy().reshape(-1))
    print('##########################################################################################')
    print("Errors' magnitudes")
    for r in range(8):
        # sample a random rotation
        g = model.G.sample()
        
        x_transformed = g @ x
        x_transformed = x_transformed.to(device)

        y_transformed = model(x_transformed).to('cpu')
        
        # verify that f(g@x) = g@f(x)=g@y
        print(torch.linalg.norm(y_transformed.tensor - (g@y).tensor, dim=1).numpy().reshape(-1))
        

print('##########################################################################################')
print()

```

--------------------------------

### Apply Equivariant Convolution and Verify Output Type

Source: https://github.com/quva-lab/escnn/blob/master/examples/introduction.ipynb

Demonstrates applying an equivariant convolution operation (`conv`) to an input geometric tensor (`x`) and verifying that the output tensor (`y`) matches the expected output feature type (`feat_type_out`). This showcases the basic usage of ESCNN's convolutional layers.

```python
y = conv(x)
assert y.type == feat_type_out
```

--------------------------------

### Build a Deeper Equivariant Model with SequentialModule

Source: https://github.com/quva-lab/escnn/blob/master/examples/introduction.ipynb

Demonstrates constructing a deeper, sequential equivariant neural network. It defines input, hidden, and output feature types and chains `R2Conv`, `InnerBatchNorm`, and `ReLU` layers using `nn.SequentialModule`. The entire model is confirmed to be equivariant if individual layers are compatible.

```python
feat_type_in = nn.FieldType(r2_act, [r2_act.trivial_repr])
feat_type_hid = nn.FieldType(r2_act, 8*[r2_act.regular_repr])
feat_type_out = nn.FieldType(r2_act, 2*[r2_act.regular_repr])

model = nn.SequentialModule(
    nn.R2Conv(feat_type_in, feat_type_hid, kernel_size=3),
    nn.InnerBatchNorm(feat_type_hid),
    nn.ReLU(feat_type_hid),
    nn.R2Conv(feat_type_hid, feat_type_hid, kernel_size=3),
    nn.InnerBatchNorm(feat_type_hid),
    nn.ReLU(feat_type_hid),
    nn.R2Conv(feat_type_hid, feat_type_out, kernel_size=3),
).eval()
```

--------------------------------

### Iterate and Assert Model Predictions with Transformations

Source: https://github.com/quva-lab/escnn/blob/master/examples/introduction.ipynb

This snippet iterates through testing elements, applies a transformation to the input, passes it through a model, and asserts that the model's output matches the transformed ground truth. It requires PyTorch and the ESCNN library for transformations and model execution.

```python
for g in r2_act.testing_elements:
    x_transformed = x.transform(g)
    y_from_x_transformed = model(x_transformed)
    
    y_transformed_from_x = y.transform(g)
    
    assert torch.allclose(y_from_x_transformed.tensor, y_transformed_from_x.tensor, atol=1e-5), g
```

--------------------------------

### Build Complete Equivariant Network (ESCNN)

Source: https://context7.com/quva-lab/escnn/llms.txt

Constructs a complete equivariant neural network using ESCNN's `nn.SequentialModule` and `nn.R2Conv`. The network processes input features and verifies C8-equivariance.

```python
from escnn import nn, gspaces
import torch

r2_act = gspaces.rot2dOnR2(N=8)

# Define feature types for each layer
feat_in = nn.FieldType(r2_act, [r2_act.trivial_repr])          # 1 channel input
feat_hid = nn.FieldType(r2_act, 32*[r2_act.regular_repr])      # 256 channels hidden
feat_out = nn.FieldType(r2_act, 16*[r2_act.regular_repr])      # 128 channels output

# Build sequential equivariant model
model = nn.SequentialModule(
    nn.R2Conv(feat_in, feat_hid, kernel_size=7, padding=3, bias=False),
    nn.InnerBatchNorm(feat_hid),
    nn.ReLU(feat_hid),

    nn.R2Conv(feat_hid, feat_hid, kernel_size=5, padding=2, bias=False),
    nn.InnerBatchNorm(feat_hid),
    nn.ReLU(feat_hid),

    nn.R2Conv(feat_hid, feat_out, kernel_size=3, padding=1, bias=False),
    nn.InnerBatchNorm(feat_out),
    nn.ReLU(feat_out),
).eval()

# Process input and verify equivariance
x = torch.randn(4, 1, 32, 32)
x_geo = feat_in(x)
y = model(x_geo)

for g in r2_act.testing_elements:
    x_transformed = x_geo.transform(g)
    y_from_transformed = model(x_transformed)
    y_transformed = y.transform(g)
    assert torch.allclose(y_from_transformed.tensor, y_transformed.tensor, atol=1e-5)

print(f"Model is C8-equivariant: {x_geo.shape} -> {y.shape}")
```

--------------------------------

### Prepare for Equivariance Testing

Source: https://github.com/quva-lab/escnn/blob/master/examples/mlp.ipynb

Sets up the numpy environment for numerical precision and evaluation mode for the model. It also defines a batch size 'B' which is typically used in subsequent testing steps to evaluate the model's equivariance properties.

```python
np.set_printoptions(linewidth=10000, precision=4, suppress=True)

model.eval()

B = 6
```

--------------------------------

### Build Equivariant 3D Convolutional Layer (escnn)

Source: https://context7.com/quva-lab/escnn/llms.txt

Demonstrates building an E(n)-equivariant 3D convolutional layer using escnn.nn.R3Conv. It defines input and output feature types suitable for 3D volumetric data and shows the application of the convolution.

```python
from escnn import nn, gspaces
import torch

r3_act = gspaces.rot3dOnR3()
feat_type_in = nn.FieldType(r3_act, [r3_act.trivial_repr])  # single scalar channel
feat_type_out = nn.FieldType(r3_act, 8*[r3_act.irrep(1)])   # 8 vector fields (L=1 spherical harmonics)

# Create SO(3)-equivariant 3D convolution
conv3d = nn.R3Conv(feat_type_in, feat_type_out, kernel_size=5, padding=2, bias=False)

# Process 3D volumetric data
x = torch.randn(4, 1, 32, 32, 32)  # batch of 4 volumes
x_geo = feat_type_in(x)
y = conv3d(x_geo)
print(f"3D Input: {x_geo.shape}, Output: {y.shape}")  # torch.Size([4, 1, 32, 32, 32]) -> torch.Size([4, 24, 32, 32, 32])
```

--------------------------------

### Build and Test Group-Equivariant Model in ESCNN

Source: https://github.com/quva-lab/escnn/blob/master/examples/introduction.ipynb

This snippet constructs a simple R2Conv model that maps scalar fields to vector fields. It defines input, hidden, and output field types, sets up convolutional layers with batch normalization and ReLU activation, and then prepares a random scalar input tensor for testing.

```python
import matplotlib.pyplot as plt
import numpy as np

feat_type_in = nn.FieldType(r2_act, [r2_act.trivial_repr])
feat_type_hid = nn.FieldType(r2_act, 8*[r2_act.regular_repr])

model = nn.SequentialModule(
    nn.R2Conv(feat_type_in, feat_type_hid, kernel_size=3),
    nn.InnerBatchNorm(feat_type_hid),
    nn.ReLU(feat_type_hid),
    nn.R2Conv(feat_type_hid, feat_type_hid, kernel_size=3),
    nn.InnerBatchNorm(feat_type_hid),
    nn.ReLU(feat_type_hid),
    nn.R2Conv(feat_type_hid, feat_type_out, kernel_size=3),
).eval()

S = 11
x = torch.randn(1, 1, S, S)
x = feat_type_in(x)

fig, axs = plt.subplots(1, r2_act.fibergroup.order(), sharex=True, sharey=True, figsize=(16, 4))

X, Y = np.meshgrid(range(S-6), range(S-7, -1, -1))
```

--------------------------------

### Create and Wrap Random Image Tensor

Source: https://github.com/quva-lab/escnn/blob/master/examples/introduction.ipynb

This code generates random tensor data representing grayscale images and wraps it into an nn.GeometricTensor. This tensor includes the associated feature type, enabling dynamic type checking within equivariant modules and PyTorch operations.

```python
import torch

x = torch.randn(4, 1, 32, 32)
x = feat_type_in(x)

assert isinstance(x.tensor, torch.Tensor)
```

--------------------------------

### Instantiate R2 Convolutional Layer

Source: https://github.com/quva-lab/escnn/blob/master/examples/introduction.ipynb

This code instantiates an R2Conv layer, which is an equivariant convolutional layer for the R2 group. It maps features from the input type (feat_type_in) to the output type (feat_type_out) using a specified kernel size.

```python
conv = nn.R2Conv(feat_type_in, feat_type_out, kernel_size=3)
```

--------------------------------

### Build Classification Model with Invariant Output (ESCNN)

Source: https://context7.com/quva-lab/escnn/llms.txt

Builds a complete classification model using ESCNN, featuring an equivariant backbone and an invariant output layer using `nn.GroupPooling` for classification.

```python
from escnn import nn, gspaces
import torch
import torch.nn.functional as F

r2_act = gspaces.rot2dOnR2(N=8)

feat_in = nn.FieldType(r2_act, 3*[r2_act.trivial_repr])
feat_h1 = nn.FieldType(r2_act, 16*[r2_act.regular_repr])
feat_h2 = nn.FieldType(r2_act, 32*[r2_act.regular_repr])
feat_h3 = nn.FieldType(r2_act, 64*[r2_act.regular_repr])

# Backbone: fully equivariant feature extraction
backbone = nn.SequentialModule(
    nn.R2Conv(feat_in, feat_h1, kernel_size=5, padding=2, bias=False),
    nn.InnerBatchNorm(feat_h1),
    nn.ReLU(feat_h1),
    nn.PointwiseMaxPool2D(feat_h1, kernel_size=2, stride=2),

    nn.R2Conv(feat_h1, feat_h2, kernel_size=5, padding=2, bias=False),
    nn.InnerBatchNorm(feat_h2),
    nn.ReLU(feat_h2),
    nn.PointwiseMaxPool2D(feat_h2, kernel_size=2, stride=2),

    nn.R2Conv(feat_h2, feat_h3, kernel_size=3, padding=1, bias=False),
    nn.InnerBatchNorm(feat_h3),
    nn.ReLU(feat_h3),
    nn.PointwiseAdaptiveAvgPool2D(feat_h3, output_size=4),
)

# Invariant mapping + classification head
class C8InvariantClassifier(torch.nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.backbone = backbone.eval()
        self.group_pool = nn.GroupPooling(feat_h3)
        self.fc = torch.nn.Linear(64 * 4 * 4, num_classes)

    def forward(self, x):
        x = feat_in(x)
        x = self.backbone(x)
        x = self.group_pool(x)  # make invariant
        x = x.tensor.reshape(x.shape[0], -1)
        return self.fc(x)

model = C8InvariantClassifier(num_classes=10)
x = torch.randn(4, 3, 32, 32)
logits = model(x)
print(f"Classification output: {x.shape} -> {logits.shape}")
```

--------------------------------

### Applying Norm-based Non-linearities with ESCNN

Source: https://github.com/quva-lab/escnn/blob/master/examples/introduction.ipynb

This Python code demonstrates how to apply norm-based non-linearities (NormNonLinearity) to model outputs. It includes transforming input data, applying the non-linearity, and asserting the equivariance by comparing transformed outputs.

```python
norm_relu = nn.NormNonLinearity(feat_type_out, 'n_relu')

y = norm_relu(model(x))

for g in r2_act.testing_elements:
    x_transformed = x.transform(g)
    y_from_x_transformed = norm_relu(model(x_transformed))
    
    y_transformed_from_x = y.transform(g)
    
    assert torch.allclose(y_from_x_transformed.tensor, y_transformed_from_x.tensor, atol=1e-5), g
```

--------------------------------

### Plotting Vector Fields with ESCNN

Source: https://github.com/quva-lab/escnn/blob/master/examples/introduction.ipynb

This snippet iterates through testing elements, transforms input data, applies a model, and plots the output vector field using matplotlib's quiver function. It demonstrates visualizing model outputs for different group elements.

```python
for i, g in enumerate(r2_act.testing_elements):
    # transform the input
    x_transformed = x.transform(g)
    
    y = model(x_transformed)
    y = y.tensor.detach().numpy().squeeze()
    
    # plot the output vector field
    axs[i].quiver(X, Y, y[0, ...], y[1, ...], units='xy')
    axs[i].set_title(g.to('int')*90)
    
plt.show()
```