fdwt 0.1.0__py3-none-any.whl

dwt/DWT1D.py

@@ -0,0 +1,70 @@
+"""ॐ
+FDWT: Fast Multidimensional Discrete Wavelet Transform Layers (PyTorch).
+Copyright 2026 Kishore Kumar Tarafdar
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License."""
+
+import torch
+from dwt.layout import DWTNDlayout, IDWTNDlayout
+
+
+class DWT1D(DWTNDlayout):
+    """1D DWT analysis module.
+
+    clean=True:  (batch, N, C) -> (batch, N/2, C*2)   [L||H packed along channel axis]
+    clean=False: (batch, N, C) -> (batch, N, C)
+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        N = x.shape[1]
+        A = self._get_A(N, x.device)
+        out = torch.einsum('ij,bjc->bic', A, x)
+        if self.clean:
+            return self._extract_2subbands(out)
+        return out
+
+    def _extract_2subbands(self, x: torch.Tensor) -> torch.Tensor:
+        mid = x.shape[1] // 2
+        L = x[:, :mid, :]
+        H = x[:, mid:, :]
+        return torch.cat([L, H], dim=-1)
+
+
+class IDWT1D(IDWTNDlayout):
+    """1D IDWT synthesis module.
+
+    clean=True:  (batch, N/2, C*2) -> (batch, N, C)
+    clean=False: (batch, N, C)     -> (batch, N, C)
+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.clean:
+            x = self._join_2subbands(x)
+        N = x.shape[1]
+        S = self._get_S(N, x.device)
+        return torch.einsum('ij,bjc->bic', S, x)
+
+    def _join_2subbands(self, x: torch.Tensor) -> torch.Tensor:
+        L, H = torch.chunk(x, 2, dim=-1)
+        return torch.cat([L, H], dim=1)
+
+
+if __name__ == '__main__':
+    wave = 'haar'
+    dwt  = DWT1D(wave)
+    idwt = IDWT1D(wave)
+    x = torch.randn(2, 256, 1)
+    lh = dwt(x)
+    xhat = idwt(lh)
+    print('DWT output:', lh.shape)
+    print('Reconstruction error (max):', (x - xhat).abs().max().item())

dwt/DWT2D.py

@@ -0,0 +1,85 @@
+"""ॐ
+FDWT: Fast Multidimensional Discrete Wavelet Transform Layers (PyTorch).
+Copyright 2026 Kishore Kumar Tarafdar
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License."""
+
+import torch
+from dwt.layout import DWTNDlayout, IDWTNDlayout
+
+
+class DWT2D(DWTNDlayout):
+    """2D DWT analysis module.
+
+    clean=True:  (batch, H, W, C) -> (batch, H/2, W/2, C*4)   [LL|LH|HL|HH along channel]
+    clean=False: (batch, H, W, C) -> (batch, H, W, C)
+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        N = x.shape[1]
+        A = self._get_A(N, x.device)
+        # columns (swap H and W, apply A along H, swap back)
+        x = x.permute(0, 2, 1, 3)
+        x = torch.einsum('ij,bjkc->bikc', A, x)
+        x = x.permute(0, 2, 1, 3)
+        # rows
+        x = torch.einsum('ij,bjkc->bikc', A, x)
+        if self.clean:
+            return self._extract_4subbands(x)
+        return x
+
+    def _extract_4subbands(self, x: torch.Tensor) -> torch.Tensor:
+        mid = x.shape[1] // 2
+        LL = x[:, :mid, :mid, :]
+        LH = x[:, mid:, :mid, :]
+        HL = x[:, :mid, mid:, :]
+        HH = x[:, mid:, mid:, :]
+        return torch.cat([LL, LH, HL, HH], dim=-1)
+
+
+class IDWT2D(IDWTNDlayout):
+    """2D IDWT synthesis module.
+
+    clean=True:  (batch, H/2, W/2, C*4) -> (batch, H, W, C)
+    clean=False: (batch, H, W, C)        -> (batch, H, W, C)
+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.clean:
+            x = self._join_quadrants(x)
+        N = x.shape[1]
+        S = self._get_S(N, x.device)
+        # columns
+        x = x.permute(0, 2, 1, 3)
+        x = torch.einsum('ij,bjkc->bikc', S, x)
+        x = x.permute(0, 2, 1, 3)
+        # rows
+        x = torch.einsum('ij,bjkc->bikc', S, x)
+        return x
+
+    def _join_quadrants(self, x: torch.Tensor) -> torch.Tensor:
+        LL, LH, HL, HH = torch.chunk(x, 4, dim=-1)
+        top    = torch.cat([LL, HL], dim=2)
+        bottom = torch.cat([LH, HH], dim=2)
+        return torch.cat([top, bottom], dim=1)
+
+
+if __name__ == '__main__':
+    wave = 'haar'
+    dwt  = DWT2D(wave)
+    idwt = IDWT2D(wave)
+    x = torch.randn(2, 256, 256, 1)
+    lh = dwt(x)
+    xhat = idwt(lh)
+    print('DWT output:', lh.shape)
+    print('Reconstruction error (max):', (x - xhat).abs().max().item())

dwt/DWT3D.py

@@ -0,0 +1,101 @@
+"""ॐ
+FDWT: Fast Multidimensional Discrete Wavelet Transform Layers (PyTorch).
+Copyright 2026 Kishore Kumar Tarafdar
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License."""
+
+import torch
+from dwt.layout import DWTNDlayout, IDWTNDlayout
+
+
+class DWT3D(DWTNDlayout):
+    """3D DWT analysis module.
+
+    clean=True:  (batch, D, H, W, C) -> (batch, D/2, H/2, W/2, C*8)
+    clean=False: (batch, D, H, W, C) -> (batch, D, H, W, C)
+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        N = x.shape[1]
+        A = self._get_A(N, x.device)
+        # columns (axis 2): swap D and H
+        x = x.permute(0, 2, 1, 3, 4)
+        x = torch.einsum('ij,bjklc->biklc', A, x)
+        x = x.permute(0, 2, 1, 3, 4)
+        # rows (axis 1)
+        x = torch.einsum('ij,bjklc->biklc', A, x)
+        # depth (axis 3): bring W to front
+        x = x.permute(0, 3, 1, 2, 4)
+        x = torch.einsum('ij,bjklc->biklc', A, x)
+        x = x.permute(0, 2, 3, 1, 4)
+        if self.clean:
+            return self._extract_8subbands(x)
+        return x
+
+    def _extract_8subbands(self, x: torch.Tensor) -> torch.Tensor:
+        mid = x.shape[1] // 2
+        LLL = x[:, :mid, :mid, :mid, :]
+        LLH = x[:, :mid, :mid, mid:, :]
+        LHL = x[:, :mid, mid:, :mid, :]
+        LHH = x[:, :mid, mid:, mid:, :]
+        HLL = x[:, mid:, :mid, :mid, :]
+        HLH = x[:, mid:, :mid, mid:, :]
+        HHL = x[:, mid:, mid:, :mid, :]
+        HHH = x[:, mid:, mid:, mid:, :]
+        return torch.cat([LLL, LLH, LHL, LHH, HLL, HLH, HHL, HHH], dim=-1)
+
+
+class IDWT3D(IDWTNDlayout):
+    """3D IDWT synthesis module.
+
+    clean=True:  (batch, D/2, H/2, W/2, C*8) -> (batch, D, H, W, C)
+    clean=False: (batch, D, H, W, C)          -> (batch, D, H, W, C)
+    """
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.clean:
+            x = self._join_octants(x)
+        N = x.shape[1]
+        S = self._get_S(N, x.device)
+        # columns
+        x = x.permute(0, 2, 1, 3, 4)
+        x = torch.einsum('ij,bjklc->biklc', S, x)
+        x = x.permute(0, 2, 1, 3, 4)
+        # rows
+        x = torch.einsum('ij,bjklc->biklc', S, x)
+        # depth
+        x = x.permute(0, 3, 1, 2, 4)
+        x = torch.einsum('ij,bjklc->biklc', S, x)
+        x = x.permute(0, 2, 3, 1, 4)
+        return x
+
+    def _join_octants(self, x: torch.Tensor) -> torch.Tensor:
+        LLL, LLH, LHL, LHH, HLL, HLH, HHL, HHH = torch.chunk(x, 8, dim=-1)
+        front_top = torch.cat([LLL, LLH], dim=3)
+        front_bot = torch.cat([LHL, LHH], dim=3)
+        back_top  = torch.cat([HLL, HLH], dim=3)
+        back_bot  = torch.cat([HHL, HHH], dim=3)
+        front = torch.cat([front_top, front_bot], dim=2)
+        back  = torch.cat([back_top,  back_bot],  dim=2)
+        return torch.cat([front, back], dim=1)
+
+
+if __name__ == '__main__':
+    wave = 'bior1.5'
+    dwt  = DWT3D(wave)
+    idwt = IDWT3D(wave)
+    x = torch.randn(1, 32, 32, 32, 2)
+    lh = dwt(x)
+    xhat = idwt(lh)
+    print('DWT output:', lh.shape)
+    print('Reconstruction error (max):', (x - xhat).abs().max().item())

dwt/__init__.py

@@ -0,0 +1,26 @@
+"""ॐ
+FDWT: Fast Multidimensional Discrete Wavelet Transform Layers (PyTorch).
+Copyright 2026 Kishore Kumar Tarafdar
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License."""
+
+from dwt.filters import FetchAnalysisSynthesisFilters
+from dwt.dwt_op import make_dwt_operator_matrix_A
+from dwt.DWT1D import DWT1D, IDWT1D
+from dwt.DWT2D import DWT2D, IDWT2D
+from dwt.DWT3D import DWT3D, IDWT3D
+from dwt.multilevel.dwt1 import dwt, idwt
+from dwt.multilevel.dwt2 import dwt2, idwt2
+from dwt.multilevel.dwt3 import dwt3, idwt3
+
+__version__ = "0.1.0"