[FIXED] Plotting a quadratic loss using norm function from Pytorch and showing it using Plotly

Issue

I have a quadratic loss z=(1/2)||Aw-b||^2 where A is 4x2 matrix, w=[x,y] is a 2d vector, and b is a 4d vector. If we plot z, there would be a surface in terms of x,y. I want to plot z using Plotly library. To do this, I want to use Pytorch and the function torch.norm for calculating the norm. Here is a worked example for plotting a 3d surface and I want to modify it as follows:

import plotly.graph_objects as go
import numpy as np

A = torch.tensor([[ 0.1542, -0.0682],
        [ 0.8631,  0.6762],
        [-1.4002,  1.1773],
        [ 0.4614,  0.2431]])

b = torch.tensor([-0.2332, -0.7453,  0.9061,  1.2118])

x = np.arange(-1,1,.01)
y = np.arange(-1,1,.01)
X,Y = np.meshgrid(x,y)

W = ??????

Z = 0.5*torch.norm(torch.matmul(A, W)-b)**2

fig = go.Figure(
    data=[go.Surface(z=Z, x=x, y=y, colorscale="Reds", opacity=0.5)])
fig.update_layout(
    title='My title', 
    autosize=False,
    width=500, 
    height=500,
    margin=dict(l=65, r=50, b=65, t=90), 
    scene_aspectmode='cube'
)
fig.show()

Question:

How should I modify W which includes x,y to plot the surface?

Solution

You could simply do:

Z = [[0.5 * torch.norm(torch.matmul(A, torch.tensor([float(xx), float(yy)]))-b)**2 for xx in x] for yy in y]

Update: You can improve the performance significantly by using torch's micro batch feature. For this you have to reshape your data to lists of matrices. That means you have to extend tensor A to a list that contains only one matrix and W to a list that contains all mesh points, each as matrix.

import plotly.graph_objects as go
import torch

A = torch.tensor([[[0.1542, -0.0682],
                   [0.8631,  0.6762],
                   [-1.4002, 1.1773],
                   [0.4614,  0.2431]]])

b = torch.tensor([-0.2332, -0.7453,  0.9061,  1.2118])

x = torch.arange(-1, 1, 0.01)
y = torch.arange(-1, 1, 0.01)
W = torch.reshape(torch.cartesian_prod(x, y), (len(x) * len(y), 2, 1))

V = torch.reshape(torch.matmul(A, W), (len(x), len(y), 4)) - b
Z = 0.5 * torch.norm(V, dim=2)**2

Answered By - Markus