LTC/utils.py at master · MBrzosko/LTC · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
import numpy as np
import matplotlib.pyplot as plt
import torch
from torch.utils.data import Dataset


class TimeSeriesDataset(Dataset):
    def __init__(self, data):
        self.data = torch.tensor(data, dtype=torch.float32)

    def __len__(self):
        return len(self.data) - 1

    def __getitem__(self, idx):
        x = self.data[idx]
        y = self.data[idx + 1]
        return x, y


class SequenceDataSet(Dataset):
    def __init__(self, data, seq_length: int = 1):
        self.data = torch.tensor(data, dtype=torch.float32)
        self.seq_length = seq_length

    def __len__(self):
        return len(self.data) - self.seq_length - 1

    def __getitem__(self, idx):
        x = self.data[idx : idx + self.seq_length, :]
        y = self.data[idx + self.seq_length + 1, (0, 1)]
        return x, y


def get_data(path: str) -> np.ndarray:
    data = np.loadtxt(path, delimiter=",", skiprows=1)

    return data


def create_sequences(data, seq_length):
    sequences = []
    targets = []
    for i in range(len(data) - seq_length):
        seq = data[i : i + seq_length]
        target = data[i + seq_length]
        sequences.append(seq)
        targets.append(target)
    return torch.tensor(np.array(sequences), dtype=torch.float32), torch.tensor(
        np.array(targets), dtype=torch.float32
    )


def create_sliced_dat(data, seq_length: int = 1):
    start_idx = 0
    x_sliced = []
    y = []

    while (start_idx + seq_length) < len(data):
        x_sliced.append(data[start_idx : start_idx + seq_length, :])
        y.append(data[start_idx + seq_length + 1, (0, 1)])
        start_idx += seq_length

    return x_sliced, y


def create_plot(x: np.ndarray, y: np.ndarray, x_label: str, y_label: str) -> None:
    plt.plot(x, y)
    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.grid(linestyle="dashed")
    plt.show()


def plot_data(time, data, plot_trajectories: bool = True) -> None:
    fig, axs = plt.subplots(3, 1, figsize=(15, 8))

    axs[0].plot(time, data[:, 1])
    axs[0].plot(time[0], data[0, 1], "ro")
    axs[0].set_ylabel("q1")
    axs[0].grid(linestyle="dashed")

    axs[1].plot(time, data[:, 2])
    axs[1].plot(time[0], data[0, 2], "ro")
    axs[1].set_ylabel("q2")
    axs[1].grid(linestyle="dashed")

    axs[2].plot(time, data[:, 3])
    axs[2].plot(time[0], data[0, 3], "ro")
    axs[2].set_xlabel("time")
    axs[2].set_ylabel("Forcing")
    axs[2].grid(linestyle="dashed")

    plt.tight_layout()
    plt.show()

    if plot_trajectories:
        create_plot(x=data[:, 1], y=data[:, 2], x_label="q1", y_label="q2")


def signal_sq_smp(dat, trans):
    return np.max(dat[trans:] ** 2)


def signal_sq_mean(dat, trans):
    return np.mean(dat[trans:] ** 2)


def signal_characteristic(dat, trans):
    return np.array([[signal_sq_smp(dat, trans)], [signal_sq_mean(dat, trans)]])