Transformation under Euclidean symmetries#

In this section, we will show how the action of a symmetry element denoted by $g$ can be tailored to respect symmetries when knowing the irreducible representation type Irrep of the input data.

For example, we take the irreducible representation obtained through decomposing an arbitrary rank-2 tensor: “1x0e + 1x1o + 1x2o” from the previous section. Irrep objects are represented as a vector by concatenating irreducible componentes of increasing order $l$. Remember that a $l$-irrep is of dimension $\mathcal{R}^{2l+1}$, the irrep “1x0e + 1x1o + 1x2o” can be expressed as a vector of $\mathcal{R}^9$:

\[ Irrep = \begin{bmatrix} \vec{I^0} & \vec{I^1} & \vec{I^2} \end{bmatrix} =\begin{bmatrix} c_1 & v_1 & v_2 & v_3 & l_1 & l_2 & l_3 & l_4 & l_5 \end{bmatrix} \in \mathcal{R}^9 \]

Under this definition, any symmetry elements $g$ such as rotations can be represented as a block diagonal square matrix $D(g)$ that can be constructed based on the irreducible representation type. For the irrep “1x0e + 1x1o + 1x2o”, $D(g)$ will be a square matrix of $\mathcal{R}^{9\times 9}$:

\[\begin{split} D(g) = \begin{bmatrix} 1 & & \\ & D(g)^{l=1} & \\ & & D(g)^{l=2} \end{bmatrix} \text{ where } D(g)^{l} \in \mathcal{R}^{(2l+1)\times (2l+1)} \text{ is the transformation matrix for $l$-irrep} \end{split}\]

Notice that $D(g)^{l=0}$ for irrep “$n\times$ 0e” of n scalars will be $I_{n\times n}$ since scalar values do not change under rotation.

For the illustration below we will use the example of a particle sitting at location $(x, y, z)$ with four vector properties $v_0, v_1, v_2, v_3 \in \mathcal{R}^3$ (velocity, acceleration, spin etc.) This particle and all of its properties can be represented by an Irrep of “3x0e + 4x1o” expressed as a vector of $\mathcal{R}^{15}$. Its corresponding transformation matrix $D(g)$ given an arbitrary symmetry opertation $g$ will be of $\mathcal{R}^{15\times 15}$.

The illustration plotted here showcases how the block diagonal matrix $D(g)$ and the particle properties transforms under a rotation operation with $\alpha = \beta = \theta$ varing from $0$ to $2\pi$.

from numpy import *
import numpy as np

n_vec = 4
x = np.zeros(n_vec)
y = np.zeros(n_vec)
z = np.zeros(n_vec)

# phi = 2*np.pi*(np.random.random(n_vec)*0.05+0.125)
phi = 2*np.pi*(0.125)
# theta = np.arccos(-1 + 2*(np.random.random(n_vec)*0.05+sqrt(2)/2))
theta = np.arccos(-1 + 2*(sqrt(2)/2))
R = np.ones(n_vec)
u = np.abs(R * sin(theta) * cos(phi)) * np.array([-1,1,-1,1])
v = np.abs(R * sin(theta) * sin(phi)) * np.array([1,-1,-1,1])
w = np.abs(R * cos(theta)) * np.array([1,1,-1,-1])

import plotly.figure_factory as ff
import plotly.graph_objects as go
import plotly.express as px
from plotly.express.colors import sample_colorscale

c = px.colors.qualitative.Plotly
# c = sample_colorscale('Viridis', np.linspace(0,1,4))

vector_flattened = np.concatenate((u[None,:],v[None,:],w[None,:]), axis=0).flatten("F")
vector_in = np.concatenate(([0,0,0], vector_flattened))

from e3nn.o3 import Irreps

irreps = Irreps("3x0e + 4x1o")

import plotly.graph_objects as go
import numpy as np
from plotly.subplots import make_subplots
import torch
t = torch.tensor

# Create figure
fig = make_subplots(rows=1, cols=2,
    specs=[[{"type": "xy"}, {"type": "scene"}]],
    column_widths=[0.45, 0.55])

# Add traces, one for each slider step
rotations = np.linspace(0, 2*np.pi, 101)
for step in rotations:
    # a small rotation around the y axis
    D = irreps.D_from_angles(alpha=t(float(step)), beta=t(float(step)), gamma=t(float(step)), k=t(0))

    # Display Heatmap
    fig.add_trace(
        go.Heatmap(visible=False, z=D, colorscale="RdYlBu", zmin=-1, zmax=1, 
        colorbar=dict(x=0.40, 
            y=0, 
            xanchor='left', 
            yanchor='bottom',
            len=1, 
            thickness=20)),
        row=1, col=1)

    # Calculate vector after rotation
    v_out = vector_in @ np.array(D)

    # Plot Points and Vectors after rotation
    fig = fig.add_trace(go.Scatter3d(
        x=v_out[0:1], y=v_out[1:2], z=v_out[2:3],
        mode='markers',
        name='point (3x0e)',
        marker=dict(
            size=12,
            color=c[0]),
        visible=False),
        row=1,
        col=2
    )
    for idx in range(n_vec):
        x_vec = np.array([v_out[0], v_out[3*idx+3]])
        y_vec = np.array([v_out[1], v_out[3*idx+4]])
        z_vec = np.array([v_out[2], v_out[3*idx+5]])
        fig.add_trace(go.Scatter3d(
            visible=False,
            x=x_vec, y=y_vec, z=z_vec,
            mode='lines', name='vector {} (1x1o)'.format(idx),
            line=dict(color=c[idx],width=2)
        ), row=1, col=2)
        x_cone = [v_out[3*idx+3]*0.9]
        y_cone = [v_out[3*idx+4]*0.9]
        z_cone = [v_out[3*idx+5]*0.9]
        u_cone = [v_out[3*idx+3]*0.3]
        v_cone = [v_out[3*idx+4]*0.3]
        w_cone = [v_out[3*idx+5]*0.3]
        fig.add_trace(go.Cone(
            x=x_cone, y=y_cone, z=z_cone, u=u_cone, v=v_cone, w=w_cone,
            showscale=False, colorscale=[[0, c[idx]], [1, c[idx]]], sizemode="absolute", visible=False
        ), row=1, col=2)

# Make first set of traces visible
# for idx in range(410, 420):
for idx in range(10):
    fig.data[idx].visible = True

# Create and add slider
steps = []
for i in range(len(rotations)):
    step = dict(
        method="update",
        args=[{"visible": [False] * len(fig.data)},
              {"title": 'Rotation with angle &#945; = &#946; = &#947; = {:.2f}&#960;'.format(0.02*i)}],  # layout attribute
            # {"title":  '&#945; =&#946;=&#947;=&#960;'}]
    )
    for idx in range(10):
        step["args"][0]["visible"][10*i+idx] = True  # Toggle i'th trace to "visible"
    steps.append(step)

sliders = [dict(
    active=0,
    currentvalue={"prefix": "Rotation "},
    steps=steps
)]

fig.update_layout(
    sliders=sliders, 
    title_text='Rotation with angle &#945; = &#946; = &#947; = 0&#960;'
)
fig.update_yaxes(autorange="reversed")
fig.update_layout(
    scene = dict(
        xaxis = dict(nticks=4, range=[-1,1]),
        yaxis = dict(nticks=4, range=[-1,1]),
        zaxis = dict(nticks=4, range=[-1,1]),
    ),
    scene_aspectmode='cube',
    autosize=False,
    width=1000,
    height=500)

fig.show()
# fig.write_image('figs/transformation.pdf')

/Users/killiansheriff/opt/miniconda3/lib/python3.9/site-packages/e3nn/o3/_wigner.py:92: UserWarning:

An output with one or more elements was resized since it had shape [1, 3, 3], which does not match the required output shape [1, 1, 3, 3]. This behavior is deprecated, and in a future PyTorch release outputs will not be resized unless they have zero elements. You can explicitly reuse an out tensor t by resizing it, inplace, to zero elements with t.resize_(0). (Triggered internally at  /Users/runner/work/pytorch/pytorch/pytorch/aten/src/ATen/native/Resize.cpp:24.)

Demistiyfing E(3)-equivariant NN

Transformation under Euclidean symmetries

Transformation under Euclidean symmetries#