| 294 | |
| 295 | # ## |
| 296 | class MWT_CZ1d(nn.Module): |
| 297 | def __init__(self, |
| 298 | k=3, alpha=64, |
| 299 | L=0, c=1, |
| 300 | base='legendre', |
| 301 | initializer=None, |
| 302 | **kwargs): |
| 303 | super(MWT_CZ1d, self).__init__() |
| 304 | |
| 305 | self.k = k |
| 306 | self.L = L |
| 307 | H0, H1, G0, G1, PHI0, PHI1 = get_filter(base, k) |
| 308 | H0r = H0 @ PHI0 |
| 309 | G0r = G0 @ PHI0 |
| 310 | H1r = H1 @ PHI1 |
| 311 | G1r = G1 @ PHI1 |
| 312 | |
| 313 | H0r[np.abs(H0r) < 1e-8] = 0 |
| 314 | H1r[np.abs(H1r) < 1e-8] = 0 |
| 315 | G0r[np.abs(G0r) < 1e-8] = 0 |
| 316 | G1r[np.abs(G1r) < 1e-8] = 0 |
| 317 | self.max_item = 3 |
| 318 | |
| 319 | self.A = sparseKernelFT1d(k, alpha, c) |
| 320 | self.B = sparseKernelFT1d(k, alpha, c) |
| 321 | self.C = sparseKernelFT1d(k, alpha, c) |
| 322 | |
| 323 | self.T0 = nn.Linear(k, k) |
| 324 | |
| 325 | self.register_buffer('ec_s', torch.Tensor( |
| 326 | np.concatenate((H0.T, H1.T), axis=0))) |
| 327 | self.register_buffer('ec_d', torch.Tensor( |
| 328 | np.concatenate((G0.T, G1.T), axis=0))) |
| 329 | |
| 330 | self.register_buffer('rc_e', torch.Tensor( |
| 331 | np.concatenate((H0r, G0r), axis=0))) |
| 332 | self.register_buffer('rc_o', torch.Tensor( |
| 333 | np.concatenate((H1r, G1r), axis=0))) |
| 334 | |
| 335 | def forward(self, x): |
| 336 | B, N, c, k = x.shape # (B, N, k) |
| 337 | ns = math.floor(np.log2(N)) |
| 338 | nl = pow(2, math.ceil(np.log2(N))) |
| 339 | extra_x = x[:, 0:nl - N, :, :] |
| 340 | x = torch.cat([x, extra_x], 1) |
| 341 | Ud = torch.jit.annotate(List[Tensor], []) |
| 342 | Us = torch.jit.annotate(List[Tensor], []) |
| 343 | # decompose |
| 344 | for i in range(ns - self.L): |
| 345 | # print('x shape',x.shape) |
| 346 | d, x = self.wavelet_transform(x) |
| 347 | Ud += [self.A(d) + self.B(x)] |
| 348 | Us += [self.C(d)] |
| 349 | x = self.T0(x) # coarsest scale transform |
| 350 | |
| 351 | # reconstruct |
| 352 | for i in range(ns - 1 - self.L, -1, -1): |
| 353 | x = x + Us[i] |