plotting.py

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

class Plotting:
    def __init__(self, tensor):
        self.tensor = tensor
    
    def plot(self):
        if self.tensor.dim() == 1:
            #line plot
            plt.plot(self.tensor.numpy())
            plt.title("1D Tensor Plot")
            plt.xlabel("Index")
            plt.ylabel("Value")

        elif self.tensor.dim() == 2:
            if self.tensor.shape[1] == 2:
            #scatter plot               #:, 0 means all rows from first column, :,1 means all rows from second column
                plt.scatter(self.tensor[:,0].numpy(), self.tensor[:,1].numpy()) #converts tensor to numpy array for plotting
                plt.title("2D Tensor Plot")
                plt.xlabel("X-axis")
                plt.ylabel("Y-axis")
            elif self.tensor.shape[1] == 3:
            #3D scatter plot
                fig = plt.figure()
                ax = fig.add_subplot(111, projection = '3d')
                ax.scatter(self.tensor[:,0].numpy(), self.tensor[:,1].numpy(), self.tensor[:,2].numpy())
                ax.set_title("3D Tensor Plot")
                ax.set_xlabel("X-axis")
                ax.set_ylabel("Y-axis")
                ax.set_zlabel("Z-axis")
        else:
            print("Plotting not supported.")
        plt.show()