Julien2429/Derivatives

1import os2import numpy as np3import torch # type: ignore4import torch.nn as nn # type: ignore5import torch.nn.functional as F # type: ignore6 7os.environ.setdefault('KMP_DUPLICATE_LIB_OK', 'TRUE')8 9VOLTAGE_SPACE = np.linspace(-1.0, 0.0, 99)10VOLTAGE = VOLTAGE_SPACE[VOLTAGE_SPACE <= -0.40]11N_SIG = len(VOLTAGE)12N_INT_TS = 3 13N_INT_FEAT = 20 14 15VOLTAGE_INTERVALS = [16    ('int1', -1.00, -0.82),17    ('int2', -0.82, -0.62),18    ('int3', -0.62, -0.40),19]20 21FEATURE_SUFFIXES = [22    'valley',23    'kurtosis',24    'skewness', 25    'area',26    'valley_position', 27    'peak_width',28    'd1_max', 29    'd1_min', 30    'n_zero_crossings',31    'd2_max', 32    'd2_min',33    'mean', 34    'std', 35    'range', 36    'energy',37    'valley_to_mean',38    'asymmetry',39    'slope_start', 40    'slope_end', 41    'overall_slope',42]43 44INTERVAL_LABELS = [45    f"{v0:.2f}–{v1:.2f} V"46    for _, v0, v1 in VOLTAGE_INTERVALS47]48 49FEAT_NAMES = [50    f"int_({v0:.2f},{v1:.2f})_{feature}"51    for _, v0, v1 in VOLTAGE_INTERVALS52    for feature in FEATURE_SUFFIXES53]54 55INTERVAL_COLORS = ['#1f77b4', '#ff7f0e', '#2ca02c']56 57class MLPNet(nn.Module):58    def __init__(self, n_features, num_classes):59        super().__init__()60        self.fc1 = nn.Linear(n_features, 64)61        self.bn1 = nn.BatchNorm1d(64)62        self.drop1 = nn.Dropout(0.3)63        self.fc2 = nn.Linear(64, 32)64        self.bn2 = nn.BatchNorm1d(32)65        self.drop2 = nn.Dropout(0.3)66        self.fc3 = nn.Linear(32, 16)67        self.fc_out = nn.Linear(16, num_classes)68 69    def forward(self, x):70        x = self.drop1(F.relu(self.bn1(self.fc1(x))))71        x = self.drop2(F.relu(self.bn2(self.fc2(x))))72        return self.fc_out(F.relu(self.fc3(x)))73 74 75class LSTMWithAttn(nn.Module):76    def __init__(self, n_features, num_classes, hidden=64):77        super().__init__()78        self.lstm    = nn.LSTM(n_features, hidden, batch_first=True, bidirectional=True)79        self.norm    = nn.LayerNorm(hidden * 2)80        self.drop    = nn.Dropout(0.3)81        self.attn    = nn.Linear(hidden * 2, 1)82        self.fc1     = nn.Linear(hidden * 2, 32)83        self.drop_fc = nn.Dropout(0.2)84        self.fc_out  = nn.Linear(32, num_classes)85 86    def forward(self, x):87            x, _= self.lstm(x)88            x= self.drop(self.norm(x))89            weights = torch.softmax(self.attn(x), dim=1)90            pooled  = (weights * x).sum(dim=1)91            return self.fc_out(self.drop_fc(F.relu(self.fc1(pooled))))92 93 94class DualBranchLSTM(nn.Module):95    def __init__(self, n_sig_features, n_int_features, num_classes, hidden_a=64, hidden_b=32):96        super().__init__()97        self.lstm_a = nn.LSTM(n_sig_features, hidden_a, batch_first=True, bidirectional=True)98        self.norm_a = nn.LayerNorm(hidden_a * 2)99        self.drop_a = nn.Dropout(0.3)100        self.attn_a = nn.Linear(hidden_a * 2, 1)101 102        self.lstm_b = nn.LSTM(n_int_features, hidden_b, batch_first=True, bidirectional=True)103        self.norm_b = nn.LayerNorm(hidden_b * 2)104        self.drop_b = nn.Dropout(0.2)105        self.attn_b = nn.Linear(hidden_b * 2, 1)106 107        fused_dim = hidden_a * 2 + hidden_b * 2108        self.norm_fuse = nn.LayerNorm(fused_dim)109        self.drop_fuse = nn.Dropout(0.3)110        self.fc1 = nn.Linear(fused_dim, 32)111        self.fc_out = nn.Linear(32, num_classes)112 113    def forward(self, x_signal, x_intervals):114        a, _ = self.lstm_a(x_signal)115        a = self.drop_a(self.norm_a(a))116        weights = torch.softmax(self.attn_a(a), dim=1)117        a = (weights * a).sum(dim=1)118 119        b, _ = self.lstm_b(x_intervals)120        b = self.drop_b(self.norm_b(b))121        weights_b = torch.softmax(self.attn_b(b), dim=1)122        b = (weights_b * b).sum(dim=1)123 124        x = torch.cat([a, b], dim=1)125        x = self.drop_fuse(self.norm_fuse(x))126        return self.fc_out(F.relu(self.fc1(x)))127 128 129class MetaLearner(nn.Module):130    def __init__(self, n_base_models, num_classes):131        super().__init__()132        self.fc1 = nn.Linear(n_base_models * num_classes, 32)133        self.drop = nn.Dropout(0.3)134        self.fc2 = nn.Linear(32, num_classes)135 136    def forward(self, x):137        return self.fc2(self.drop(F.relu(self.fc1(x))))138 139 140class FlatWrapper(nn.Module):141    def __init__(self, model, T, F):142        super().__init__()143        self.model = model144        self.T = T145        self.F = F146 147    def forward(self, x):148        return self.model(x.reshape(-1, self.T, self.F))149 150 151class DualInputWrapper(nn.Module):152    def __init__(self, model, n_sig, n_int_ts, n_int_feat):153        super().__init__()154        self.model = model155        self.n_sig = n_sig156        self.n_int_ts = n_int_ts157        self.n_int_feat = n_int_feat158 159    def forward(self, x):160        sig = x[:, :self.n_sig].unsqueeze(-1)161        intervals = x[:, self.n_sig:].reshape(-1, self.n_int_ts, self.n_int_feat)162        return self.model(sig, intervals)