tensor_product_animated.py

# pylint: disable=not-callable, no-member, invalid-name, missing-docstring, line-too-long
import argparse
import math
import os
import shutil
import subprocess

import plotly.graph_objs as go
import torch
from tqdm.auto import tqdm
from e3nn import o3, rs
from e3nn.tensor import IrrepTensor, SphericalTensor


def get_cmap(x):
    if x == 'bwr':
        return [[0, 'rgb(0,50,255)'], [0.5, 'rgb(200,200,200)'], [1, 'rgb(255,50,0)']]
    if x == 'plasma':
        return [[0, '#9F1A9B'], [0.25, '#0D1286'], [0.5, '#000000'], [0.75, '#F58C45'], [1, '#F0F524']]


def surf(args, x, center):
    x = IrrepTensor(x, len(x) // 2)
    x = SphericalTensor.from_irrep_tensor(x)

    return go.Surface(
        **x.plotly_surface(args.res, center=center, normalization='norm'),
        showscale=False,
        cmin=-0.33,
        cmax=0.33,
        colorscale=get_cmap(args.cmap),
    )


def main(args):
    if args.out is None:
        p = {-1: 'o', 1: 'e'}
        args.out = f"{args.l1}{p[args.p1]}{args.l2}{p[args.p2]}"

    if os.path.exists(args.out):
        shutil.rmtree(args.out)
    os.makedirs(args.out)

    Rs1 = [(1, args.l1, args.p1)]
    Rs2 = [(1, args.l2, args.p2)]
    tp = rs.TensorProduct(Rs1, Rs2, o3.selection_rule, normalization='norm')
    Rs_out = list(rs.split_by_mul(tp.Rs_out))

    x1 = rs.randn(Rs1)
    x1 = 0.27 * x1 / x1.norm()
    x2 = rs.randn(Rs2)
    x2 = 0.27 * x2 / x2.norm()

    for i, t in enumerate(tqdm(torch.linspace(0, 1, args.steps + 1)[:-1])):

        if args.animation == "rotation":
            a = 4 * math.pi * t
            b = 0
            c = 0
            p = 0 if t < 0.5 else 1
            gx1 = rs.rep(Rs1, a, b, c, p) @ x1
            gx2 = rs.rep(Rs2, a, b, c, p) @ x2
        if args.animation == "random":
            gx1 = rs.randn(Rs1)
            gx2 = rs.randn(Rs2)

        xleft = -0.5 * (len(Rs_out) + 2 - 1)
        data = [
            surf(args, gx1, torch.tensor([xleft, 0, 0.0])),
            surf(args, gx2, torch.tensor([xleft + 1, 0, 0.0])),
        ]

        outs = rs.cut(4. * tp(gx1, gx2), *Rs_out)
        data += [
            surf(args, out, torch.tensor([xleft + 2 + i, 0, 0.0]))
            for i, out in enumerate(outs)
        ]

        axis = dict(
            showbackground=False,
            showticklabels=False,
            showgrid=False,
            zeroline=False,
            title='',
            nticks=3,
        )

        layout = dict(
            width=(2 + len(Rs_out)) * args.height,
            height=args.height,
            scene=dict(
                xaxis=dict(
                    **axis,
                    range=[xleft - 0.5, -xleft + 0.5]
                ),
                yaxis=dict(
                    **axis,
                    range=[-0.5, 0.5]
                ),
                zaxis=dict(
                    **axis,
                    range=[-0.5, 0.5]
                ),
                aspectmode='manual',
                aspectratio=dict(x=2 * (2 + len(Rs_out)), y=2, z=2),
                camera=dict(
                    up=dict(x=0, y=0, z=1),
                    center=dict(x=0, y=0, z=0),
                    eye=dict(x=0, y=-10, z=0),
                    projection=dict(type='orthographic'),
                ),
            ),
            paper_bgcolor=args.color_bg,
            plot_bgcolor=args.color_bg,
            margin=dict(l=0, r=0, t=0, b=0)
        )

        fig = go.Figure(data=data, layout=layout)
        fig.write_image('{}/{:03d}.png'.format(args.out, i))

    subprocess.check_output(["convert", "-delay", "5", "-loop", "0", "-dispose", "2", "{}/*.png".format(args.out), "{}.gif".format(args.out)])


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument("--height", type=int, default=200)
    parser.add_argument("--res", type=int, default=100)
    parser.add_argument("--steps", type=int, default=90)
    parser.add_argument("--color_bg", type=str, default="rgba(0,0,0,0)")
    parser.add_argument("--color_text", type=str, default="rgb(255,255,255)")
    parser.add_argument("--cmap", type=str, default="plasma")
    parser.add_argument("--animation", type=str, default="rotation")
    parser.add_argument("--out", type=str)
    parser.add_argument("--l1", type=int, default=1)
    parser.add_argument("--l2", type=int, default=1)
    parser.add_argument("--p1", type=int, default=-1)
    parser.add_argument("--p2", type=int, default=-1)

    main(parser.parse_args())