torchwavelets 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (98) hide show
  1. torchwavelets-0.1.0/LICENSE +29 -0
  2. torchwavelets-0.1.0/MANIFEST.in +7 -0
  3. torchwavelets-0.1.0/PKG-INFO +303 -0
  4. torchwavelets-0.1.0/README.md +260 -0
  5. torchwavelets-0.1.0/docs/api.md +86 -0
  6. torchwavelets-0.1.0/docs/examples.md +188 -0
  7. torchwavelets-0.1.0/docs/figures/cifar10_scattering.png +0 -0
  8. torchwavelets-0.1.0/docs/figures/decomposition.png +0 -0
  9. torchwavelets-0.1.0/docs/figures/denoising.png +0 -0
  10. torchwavelets-0.1.0/docs/figures/dtcwt.png +0 -0
  11. torchwavelets-0.1.0/docs/figures/ecg_denoising.png +0 -0
  12. torchwavelets-0.1.0/docs/figures/lightning_training.png +0 -0
  13. torchwavelets-0.1.0/docs/figures/modes.png +0 -0
  14. torchwavelets-0.1.0/docs/figures/precision.png +0 -0
  15. torchwavelets-0.1.0/docs/figures/scattering.png +0 -0
  16. torchwavelets-0.1.0/docs/figures/transform_gallery.png +0 -0
  17. torchwavelets-0.1.0/docs/index.md +3 -0
  18. torchwavelets-0.1.0/examples/README.md +37 -0
  19. torchwavelets-0.1.0/examples/_data.py +115 -0
  20. torchwavelets-0.1.0/examples/application_cifar10.py +137 -0
  21. torchwavelets-0.1.0/examples/application_denoising.py +128 -0
  22. torchwavelets-0.1.0/examples/application_ecg.py +116 -0
  23. torchwavelets-0.1.0/examples/application_lightning.py +196 -0
  24. torchwavelets-0.1.0/examples/plot_decomposition.py +76 -0
  25. torchwavelets-0.1.0/examples/plot_dtcwt.py +62 -0
  26. torchwavelets-0.1.0/examples/plot_modes.py +51 -0
  27. torchwavelets-0.1.0/examples/plot_precision.py +62 -0
  28. torchwavelets-0.1.0/examples/plot_scattering.py +70 -0
  29. torchwavelets-0.1.0/examples/plot_transform_gallery.py +184 -0
  30. torchwavelets-0.1.0/pyproject.toml +74 -0
  31. torchwavelets-0.1.0/setup.cfg +4 -0
  32. torchwavelets-0.1.0/src/torchwavelets/__init__.py +62 -0
  33. torchwavelets-0.1.0/src/torchwavelets/_version.py +1 -0
  34. torchwavelets-0.1.0/src/torchwavelets/compat.py +136 -0
  35. torchwavelets-0.1.0/src/torchwavelets/core/__init__.py +13 -0
  36. torchwavelets-0.1.0/src/torchwavelets/core/filters.py +211 -0
  37. torchwavelets-0.1.0/src/torchwavelets/core/modes.py +109 -0
  38. torchwavelets-0.1.0/src/torchwavelets/dtcwt/__init__.py +24 -0
  39. torchwavelets-0.1.0/src/torchwavelets/dtcwt/_autocast.py +24 -0
  40. torchwavelets-0.1.0/src/torchwavelets/dtcwt/coeffs.py +76 -0
  41. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/__init__.py +5 -0
  42. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/antonini.npz +0 -0
  43. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/farras.npz +0 -0
  44. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/legall.npz +0 -0
  45. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/near_sym_a.npz +0 -0
  46. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/near_sym_a2.npz +0 -0
  47. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/near_sym_b.npz +0 -0
  48. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/near_sym_b_bp.npz +0 -0
  49. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/qshift_06.npz +0 -0
  50. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/qshift_32.npz +0 -0
  51. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/qshift_a.npz +0 -0
  52. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/qshift_b.npz +0 -0
  53. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/qshift_b_bp.npz +0 -0
  54. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/qshift_c.npz +0 -0
  55. torchwavelets-0.1.0/src/torchwavelets/dtcwt/data/qshift_d.npz +0 -0
  56. torchwavelets-0.1.0/src/torchwavelets/dtcwt/lowlevel.py +296 -0
  57. torchwavelets-0.1.0/src/torchwavelets/dtcwt/transform1d.py +253 -0
  58. torchwavelets-0.1.0/src/torchwavelets/dtcwt/transform2d.py +298 -0
  59. torchwavelets-0.1.0/src/torchwavelets/dtcwt/transform3d.py +420 -0
  60. torchwavelets-0.1.0/src/torchwavelets/dtcwt/transform_funcs.py +93 -0
  61. torchwavelets-0.1.0/src/torchwavelets/dwt/__init__.py +6 -0
  62. torchwavelets-0.1.0/src/torchwavelets/dwt/_validation.py +59 -0
  63. torchwavelets-0.1.0/src/torchwavelets/dwt/lowlevel.py +153 -0
  64. torchwavelets-0.1.0/src/torchwavelets/dwt/transform1d.py +175 -0
  65. torchwavelets-0.1.0/src/torchwavelets/dwt/transform2d.py +195 -0
  66. torchwavelets-0.1.0/src/torchwavelets/packets.py +439 -0
  67. torchwavelets-0.1.0/src/torchwavelets/py.typed +0 -0
  68. torchwavelets-0.1.0/src/torchwavelets/scatter/__init__.py +6 -0
  69. torchwavelets-0.1.0/src/torchwavelets/scatter/lowlevel.py +221 -0
  70. torchwavelets-0.1.0/src/torchwavelets/scatter/network.py +204 -0
  71. torchwavelets-0.1.0/src/torchwavelets/swt/__init__.py +14 -0
  72. torchwavelets-0.1.0/src/torchwavelets/swt/lowlevel.py +53 -0
  73. torchwavelets-0.1.0/src/torchwavelets/swt/transform.py +386 -0
  74. torchwavelets-0.1.0/src/torchwavelets.egg-info/PKG-INFO +303 -0
  75. torchwavelets-0.1.0/src/torchwavelets.egg-info/SOURCES.txt +96 -0
  76. torchwavelets-0.1.0/src/torchwavelets.egg-info/dependency_links.txt +1 -0
  77. torchwavelets-0.1.0/src/torchwavelets.egg-info/requires.txt +26 -0
  78. torchwavelets-0.1.0/src/torchwavelets.egg-info/top_level.txt +1 -0
  79. torchwavelets-0.1.0/tests/test_compat.py +60 -0
  80. torchwavelets-0.1.0/tests/test_compile.py +169 -0
  81. torchwavelets-0.1.0/tests/test_devices.py +153 -0
  82. torchwavelets-0.1.0/tests/test_dtcwt.py +120 -0
  83. torchwavelets-0.1.0/tests/test_dtcwt1d.py +121 -0
  84. torchwavelets-0.1.0/tests/test_dtcwt3d.py +107 -0
  85. torchwavelets-0.1.0/tests/test_dtcwt_compat.py +72 -0
  86. torchwavelets-0.1.0/tests/test_dtypes.py +90 -0
  87. torchwavelets-0.1.0/tests/test_edge_cases.py +157 -0
  88. torchwavelets-0.1.0/tests/test_gradients.py +70 -0
  89. torchwavelets-0.1.0/tests/test_lightning.py +144 -0
  90. torchwavelets-0.1.0/tests/test_packets.py +177 -0
  91. torchwavelets-0.1.0/tests/test_parity_1d.py +114 -0
  92. torchwavelets-0.1.0/tests/test_parity_2d.py +98 -0
  93. torchwavelets-0.1.0/tests/test_review_fixes.py +287 -0
  94. torchwavelets-0.1.0/tests/test_roundtrip.py +81 -0
  95. torchwavelets-0.1.0/tests/test_scatter.py +171 -0
  96. torchwavelets-0.1.0/tests/test_scatter_network.py +170 -0
  97. torchwavelets-0.1.0/tests/test_standard_usage.py +160 -0
  98. torchwavelets-0.1.0/tests/test_swt.py +168 -0
@@ -0,0 +1,29 @@
1
+ BSD 3-Clause License
2
+
3
+ Copyright (c) 2026, Rachid El Montassir
4
+ All rights reserved.
5
+
6
+ Redistribution and use in source and binary forms, with or without
7
+ modification, are permitted provided that the following conditions are met:
8
+
9
+ 1. Redistributions of source code must retain the above copyright notice, this
10
+ list of conditions and the following disclaimer.
11
+
12
+ 2. Redistributions in binary form must reproduce the above copyright notice,
13
+ this list of conditions and the following disclaimer in the documentation
14
+ and/or other materials provided with the distribution.
15
+
16
+ 3. Neither the name of the copyright holder nor the names of its
17
+ contributors may be used to endorse or promote products derived from
18
+ this software without specific prior written permission.
19
+
20
+ THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
21
+ AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
22
+ IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
23
+ DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
24
+ FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
25
+ DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
26
+ SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
27
+ CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
28
+ OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
29
+ OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
@@ -0,0 +1,7 @@
1
+ # Ensure the bundled DTCWT filter data ships in sdists and wheels.
2
+ include LICENSE
3
+ include README.md
4
+ recursive-include src/torchwavelets/dtcwt/data *.npz
5
+ recursive-include docs *.md *.png
6
+ graft examples
7
+ global-exclude *.pyc
@@ -0,0 +1,303 @@
1
+ Metadata-Version: 2.4
2
+ Name: torchwavelets
3
+ Version: 0.1.0
4
+ Summary: Torch-native successor to pytorch_wavelets: differentiable DWT (1D/2D), dual-tree complex wavelet transform (1D/2D/3D), DTCWT scattering, SWT and wavelet packets.
5
+ Author: Rachid El Montassir
6
+ License-Expression: BSD-3-Clause
7
+ Project-URL: Homepage, https://github.com/relmonta/torchwavelets
8
+ Keywords: wavelet,dwt,dtcwt,pytorch,deep-learning,signal-processing
9
+ Classifier: Development Status :: 3 - Alpha
10
+ Classifier: Intended Audience :: Science/Research
11
+ Classifier: Programming Language :: Python :: 3
12
+ Classifier: Programming Language :: Python :: 3.10
13
+ Classifier: Programming Language :: Python :: 3.11
14
+ Classifier: Programming Language :: Python :: 3.12
15
+ Classifier: Programming Language :: Python :: 3.13
16
+ Classifier: Topic :: Scientific/Engineering :: Mathematics
17
+ Classifier: Topic :: Scientific/Engineering :: Image Processing
18
+ Requires-Python: >=3.10
19
+ Description-Content-Type: text/markdown
20
+ License-File: LICENSE
21
+ Requires-Dist: numpy>=1.23
22
+ Requires-Dist: torch>=2.6
23
+ Requires-Dist: PyWavelets>=1.4
24
+ Provides-Extra: examples
25
+ Requires-Dist: matplotlib>=3.6; extra == "examples"
26
+ Requires-Dist: scikit-learn>=1.2; extra == "examples"
27
+ Requires-Dist: Pillow>=9.0; extra == "examples"
28
+ Provides-Extra: docs
29
+ Requires-Dist: mkdocs>=1.6; extra == "docs"
30
+ Requires-Dist: mkdocs-material>=9.5; extra == "docs"
31
+ Requires-Dist: mkdocstrings[python]>=0.26; extra == "docs"
32
+ Requires-Dist: pymdown-extensions>=10.8; extra == "docs"
33
+ Provides-Extra: lightning
34
+ Requires-Dist: lightning>=2.0; extra == "lightning"
35
+ Provides-Extra: dev
36
+ Requires-Dist: pre-commit>=3.6; extra == "dev"
37
+ Requires-Dist: mypy>=1.10; extra == "dev"
38
+ Provides-Extra: test
39
+ Requires-Dist: pytest>=7.0; extra == "test"
40
+ Requires-Dist: pytest-cov>=4.0; extra == "test"
41
+ Requires-Dist: ruff>=0.4; extra == "test"
42
+ Dynamic: license-file
43
+
44
+ # torchwavelets
45
+
46
+ [![CI](https://github.com/relmonta/torchwavelets/actions/workflows/ci.yml/badge.svg)](https://github.com/relmonta/torchwavelets/actions/workflows/ci.yml)
47
+ [![Docs](https://github.com/relmonta/torchwavelets/actions/workflows/docs.yml/badge.svg)](https://relmonta.github.io/torchwavelets/)
48
+ [![PyPI version](https://img.shields.io/pypi/v/torchwavelets.svg)](https://pypi.org/project/torchwavelets/)
49
+ [![Python versions](https://img.shields.io/pypi/pyversions/torchwavelets.svg)](https://pypi.org/project/torchwavelets/)
50
+ [![License](https://img.shields.io/pypi/l/torchwavelets.svg)](https://github.com/relmonta/torchwavelets/blob/main/LICENSE)
51
+ [![Code style: ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
52
+
53
+ Differentiable discrete wavelet transforms for PyTorch.
54
+
55
+ `torchwavelets` provides 1D and 2D discrete wavelet transforms (DWT), their
56
+ inverses, the 2D dual-tree complex wavelet transform (DTCWT) with native
57
+ complex tensors, DTCWT scattering networks, the stationary wavelet transform
58
+ (SWT) and wavelet packets — built to respect PyTorch's precision, device and
59
+ autocast semantics.
60
+
61
+ ## Why another wavelet library?
62
+
63
+ The reference implementation, [`fbcotter/pytorch_wavelets`](https://github.com/fbcotter/pytorch_wavelets),
64
+ is unmaintained (last commit 2023) and has well-known problems:
65
+
66
+ - **AMP / mixed precision is broken** — the custom `autograd.Function`s run
67
+ their backward in float32 against float16 gradients (`expected scalar type
68
+ Half but found Float`, issue #63).
69
+ - **Filters are baked in at construction time** via `torch.get_default_dtype()`;
70
+ feeding a float16 tensor to a float32-built transform fails.
71
+ - **Boundary handling** does not match PyWavelets exactly and fails for small
72
+ signals; the stationary wavelet inverse is dead code.
73
+ - No 1D/3D DTCWT, no wavelet packets, no native complex tensors, stale CI.
74
+
75
+ `torchwavelets` fixes these at the design level:
76
+
77
+ ## Design principles
78
+
79
+ 1. **Output dtype == input dtype.** fp16/bf16/fp32/fp64 in → the same out.
80
+ Filters are resolved lazily per `(dtype, device)` and cached; nothing is
81
+ mutated during forward, so `torch.compile`, DDP and mixed-precision
82
+ training all behave. (One documented exception: the DTCWT family computes
83
+ half and bfloat16 inputs in float32 — complex16 kernels do not exist — and
84
+ runs in an autocast-disabled scope, returning float32.)
85
+ 2. **Autocast-native.** The transforms are built from plain `conv2d` /
86
+ `conv_transpose2d` ops (no custom `autograd.Function`), so `torch.autocast`
87
+ and `torch.amp.GradScaler` work without any special-casing. The DTCWT
88
+ family computes in an autocast-disabled scope (its complex bandpasses have
89
+ no bf16 kernels) and returns float32 for half/bf16 inputs.
90
+ 3. **PyWavelets-exact boundary handling.** All modes (`zero`, `symmetric`,
91
+ `reflect`, `periodization`, `periodic`, `constant`, `replicate`) match
92
+ `pywt.dwt` / `pywt.idwt` bit-for-bit, including odd-length and tiny signals
93
+ (index-based padding, no size limits).
94
+ 4. **Pure torch ops.** Index-based padding keeps everything differentiable,
95
+ higher-order differentiable, and `torch.compile`-friendly.
96
+
97
+ ## Installation
98
+
99
+ ```bash
100
+ pip install -e .[test,examples]
101
+ ```
102
+
103
+ ## Usage
104
+
105
+ ```python
106
+ import torch
107
+ import torchwavelets as tw
108
+
109
+ x = torch.randn(4, 3, 64, 64) # (N, C, H, W)
110
+
111
+ # Functional API
112
+ yl, yh = tw.dwt2d(x, wave="db4", mode="symmetric", level=3)
113
+ xr = tw.idwt2d((yl, yh), wave="db4", mode="symmetric")
114
+
115
+ # Module API
116
+ dwt = tw.DWT2D(wave="db4", J=3, mode="symmetric")
117
+ yl, yh = dwt(x)
118
+ xr = tw.IDWT2D(wave="db4", mode="symmetric")((yl, yh))
119
+
120
+ # 1D
121
+ y = torch.randn(4, 3, 1024)
122
+ yl, yh = tw.dwt1d(y, wave="db6", level=3)
123
+
124
+ # Mixed precision just works
125
+ x16 = x.half()
126
+ yl16, _ = dwt(x16) # -> float16 output
127
+ with torch.autocast("cuda", dtype=torch.bfloat16):
128
+ yl_bf, _ = dwt(x.float()) # -> bfloat16 output (on CUDA)
129
+ ```
130
+
131
+ ## Compatibility with `pytorch_wavelets`
132
+
133
+ ```python
134
+ from torchwavelets.compat import DWTForward, DWTInverse, DWT1DForward, DWT1DInverse
135
+
136
+ xfm = DWTForward(J=3, wave="db3", mode="zero")
137
+ ```
138
+
139
+ The compat layer keeps the old API but fixes the old semantics: a float16
140
+ input works regardless of the dtype used when the transform was constructed.
141
+
142
+ ## Status
143
+
144
+ - [x] DWT 1D / 2D, analysis + synthesis, all boundary modes
145
+ - [x] float16 / bfloat16 / float32 / float64
146
+ - [x] PyWavelets numerical parity tests, gradient checks, AMP tests
147
+ - [x] `torch.compile`-compatible module API (`fullgraph=True`)
148
+ - [x] Device-portable: CPU, CUDA, multi-GPU, hybrid and Grace (GH200) nodes
149
+ - [x] DTCWT 1D / 2D / 3D (native complex tensors, `near_sym_*` / `qshift_*` banks)
150
+ - [x] DTCWT scattering layers (`ScatLayer`, `ScatLayerj2`)
151
+ - [x] General arbitrary-order scattering network (`DTCWTScattering`)
152
+ - [x] Stationary (undecimated) wavelet transform 1D / 2D
153
+ - [x] Wavelet packets 1D / 2D (`WaveletPacket`, `WaveletPacket2D`)
154
+ - [x] Scattering network (deeper, arbitrary-order)
155
+
156
+ ## Examples and applications
157
+
158
+ Self-contained, reproducible demos with real data (photograph, CIFAR-10,
159
+ MIT-BIH ECG) — run `python examples/<script>.py` from the repository root
160
+ (data is cached in `~/data/torchwavelets`, figures land in `docs/figures/`).
161
+ See [`examples/README.md`](examples/README.md) for the quick reference and
162
+ [`docs/examples.md`](docs/examples.md) for the full illustrated walkthrough
163
+ with results and figures.
164
+
165
+ <p align="center">
166
+ <a href="https://github.com/relmonta/torchwavelets/blob/main/docs/figures/transform_gallery.png">
167
+ <img src="https://raw.githubusercontent.com/relmonta/torchwavelets/main/docs/figures/transform_gallery.png" alt="transform gallery" width="98%">
168
+ </a>
169
+ </p>
170
+
171
+ **Application 1 — image denoising (DWT vs DTCWT).** A photograph corrupted
172
+ with Gaussian noise (sigma = 25/255, PSNR 20.3 dB) is restored by per-level
173
+ 3-sigma soft thresholding of the wavelet coefficients. The redundant, nearly
174
+ shift-invariant DTCWT gains **+3.1 dB PSNR** over the critically sampled DWT
175
+ (29.3 dB vs 26.2 dB).
176
+
177
+ <p align="center">
178
+ <a href="https://github.com/relmonta/torchwavelets/blob/main/docs/figures/denoising.png">
179
+ <img src="https://raw.githubusercontent.com/relmonta/torchwavelets/main/docs/figures/denoising.png" alt="image denoising" width="98%">
180
+ </a>
181
+ </p>
182
+
183
+ **Application 2 — ECG denoising (SWT vs DWT).** Real 360 Hz ECG (MIT-BIH
184
+ Arrhythmia record 100) plus 50 uV noise (SNR 10.6 dB), restored by the same
185
+ per-level soft thresholding. The undecimated SWT gains **+1.9 dB SNR** over
186
+ the DWT (14.9 dB vs 13.0 dB).
187
+
188
+ <p align="center">
189
+ <a href="https://github.com/relmonta/torchwavelets/blob/main/docs/figures/ecg_denoising.png">
190
+ <img src="https://raw.githubusercontent.com/relmonta/torchwavelets/main/docs/figures/ecg_denoising.png" alt="ecg denoising" width="98%">
191
+ </a>
192
+ </p>
193
+
194
+ **Application 3 — CIFAR-10 classification with scattering features.** A fixed
195
+ (untrained) DTCWT scattering front-end (`DTCWTScattering(J=2, max_order=2)`)
196
+ extracts 588 shift-invariant dimensions per 32x32 image (147 invariants,
197
+ 2x2-average pooled); a logistic regression reaches **56.5 %** test accuracy,
198
+ versus **31.6 %** on 3072 raw pixels — with 5x fewer dimensions and no
199
+ learned front-end.
200
+
201
+ <p align="center">
202
+ <a href="https://github.com/relmonta/torchwavelets/blob/main/docs/figures/cifar10_scattering.png">
203
+ <img src="https://raw.githubusercontent.com/relmonta/torchwavelets/main/docs/figures/cifar10_scattering.png" alt="cifar10 scattering" width="98%">
204
+ </a>
205
+ </p>
206
+
207
+ ## `torch.compile`
208
+
209
+ The module API (`DWT2D`, `IDWT2D`, `DWT1D`, `IDWT1D`) resolves its filters
210
+ into buffers at construction time and runs only torch ops in `forward`, so it
211
+ is fully traceable:
212
+
213
+ ```python
214
+ import torch
215
+ import torchwavelets as tw
216
+
217
+ dwt = torch.compile(tw.DWT2D(wave="db4", J=3, mode="symmetric"), fullgraph=True)
218
+ yl, yh = dwt(x)
219
+ ```
220
+
221
+ The functional API (`tw.dwt2d(x, wave=...)`) is eager-only for the filter
222
+ resolution step; use the module API under `torch.compile`.
223
+
224
+ ## Ecosystem notes
225
+
226
+ * **Serialization**: all transforms are plain `nn.Module`s — `state_dict`
227
+ save/load and `torch.save`/`torch.load` work as usual.
228
+ * **DataLoader / datasets**: the module API takes batched `(N, C, ...)`
229
+ tensors, so it slots directly into `DataLoader` pipelines.
230
+ * **`torch.func`**: `jacrev`/`jacfwd` and double-backward work through the
231
+ transforms (all plain torch ops).
232
+ * **`torch.distributed` (DDP)**: modules hold no state beyond filter buffers
233
+ and never mutate inputs, so DDP training works (smoke-tested with gloo).
234
+ * **`torch.jit.script` is not supported** (the transforms use dynamic loops
235
+ over levels); use `torch.compile` instead, which is fully supported.
236
+ * **`torch.vmap`**: the API is already batched (`N, C, ...` required), so
237
+ vmap is unnecessary; the batch dimension plays that role.
238
+ * **`torch.inference_mode`**: supported everywhere.
239
+
240
+ ## Device portability
241
+
242
+ Tests parametrize over every available device (`CPU` plus all CUDA devices),
243
+ so the suite automatically covers hybrid CPU/GPU nodes and NVIDIA
244
+ Grace-Hopper (ARM64) machines; GPU results are reported in the CHANGELOG.
245
+
246
+ ## Credits and acknowledgements
247
+
248
+ `torchwavelets` is an independent implementation, but it stands on the
249
+ shoulders of the wavelet community. The parts of the code that are ports or
250
+ adaptations are credited here; everything else was written from the
251
+ mathematics and the public conventions.
252
+
253
+ * **Nick Kingsbury (Cambridge University)** — the dual-tree complex wavelet
254
+ transform itself. The filter banks (`near_sym_a/b`, `qshift_a/b`, ...,
255
+ shipped in `torchwavelets/dtcwt/data/`), the quarter-sample-shift design,
256
+ the orientation convention and the `symmetric` (half-sample) boundary
257
+ handling originate from his DTCWT toolbox, which is **free for research
258
+ use only**. The 3-D DTCWT follows the same work-grid construction.
259
+ * **[`fbcotter/pytorch_wavelets`](https://github.com/fbcotter/pytorch_wavelets)** —
260
+ the unmaintained reference implementation this package replaces. Its
261
+ low-level DTCWT operations (`colfilter`, `rowfilter`, `coldfilt`,
262
+ `rowdfilt`, `colifilt`, `rowifilt`, `q2c`, `c2q`, orientation assembly),
263
+ the two-scale scattering layers (`ScatLayer`, `ScatLayerj2`) and the
264
+ `pytorch_wavelets`-compatible API (`torchwavelets.compat`) are ports of
265
+ that codebase, rewritten on plain torch ops to fix its AMP, dtype and
266
+ boundary problems.
267
+ * **[PyWavelets](https://pywavelets.readthedocs.io/)** — the parity oracle
268
+ for the DWT, SWT and wavelet packets. The boundary modes, coefficient
269
+ shapes, `swt` conventions, packet paths/sequency ordering and the
270
+ "drop extra sample" reconstruction rule all mirror PyWavelets exactly;
271
+ the test suite is bit-exact against it.
272
+ * **[`dtcwt`](https://github.com/rgommers/dtcwt)** (numpy reference
273
+ implementation) — the DTCWT parity oracle for 1-D/2-D/3-D; the test suite
274
+ is bit-exact against it to ~1e-15.
275
+ * **Stéphane Mallat** — the scattering-transform theory. The scattering
276
+ layers and the arbitrary-order `DTCWTScattering` network implement the
277
+ Mallat recipe (wavelet modulus + smoothing) with a DTCWT front-end.
278
+ * **PyTorch ecosystem** — `torch.compile`, autocast and `nn.Module`
279
+ conventions; the tests follow the standard PyTorch-ecosystem usage
280
+ patterns (serialization, DataLoader, `torch.func`, DDP).
281
+
282
+ If you use this package in research, please cite it — see
283
+ [`CITATION.cff`](CITATION.cff).
284
+
285
+ ### Primary references
286
+
287
+ 1. N. G. Kingsbury, *The dual-tree complex wavelet transform: a new efficient
288
+ tool for image restoration and enhancement*, EUSIPCO 1998.
289
+ 2. N. G. Kingsbury, *Complex wavelets for shift invariant analysis and
290
+ filtering of signals*, Applied and Computational Harmonic Analysis, 2001.
291
+ 3. I. W. Selesnick, R. G. Baraniuk, N. G. Kingsbury, *The dual-tree complex
292
+ wavelet transform*, IEEE Signal Processing Magazine, 2005.
293
+ 4. S. Mallat, *Group invariant scattering*, Communications on Pure and
294
+ Applied Mathematics, 2012.
295
+ 5. I. Daubechies, *Ten Lectures on Wavelets*, SIAM, 1992.
296
+
297
+ ## License
298
+
299
+ BSD-3-Clause (see `LICENSE`). See
300
+ [Credits and acknowledgements](#credits-and-acknowledgements) for the
301
+ upstream sources: the DTCWT filter banks originate from Nick Kingsbury's
302
+ MATLAB toolbox, which is free for **research use only** — relevant if you
303
+ plan commercial use of the DTCWT parts.
@@ -0,0 +1,260 @@
1
+ # torchwavelets
2
+
3
+ [![CI](https://github.com/relmonta/torchwavelets/actions/workflows/ci.yml/badge.svg)](https://github.com/relmonta/torchwavelets/actions/workflows/ci.yml)
4
+ [![Docs](https://github.com/relmonta/torchwavelets/actions/workflows/docs.yml/badge.svg)](https://relmonta.github.io/torchwavelets/)
5
+ [![PyPI version](https://img.shields.io/pypi/v/torchwavelets.svg)](https://pypi.org/project/torchwavelets/)
6
+ [![Python versions](https://img.shields.io/pypi/pyversions/torchwavelets.svg)](https://pypi.org/project/torchwavelets/)
7
+ [![License](https://img.shields.io/pypi/l/torchwavelets.svg)](https://github.com/relmonta/torchwavelets/blob/main/LICENSE)
8
+ [![Code style: ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)
9
+
10
+ Differentiable discrete wavelet transforms for PyTorch.
11
+
12
+ `torchwavelets` provides 1D and 2D discrete wavelet transforms (DWT), their
13
+ inverses, the 2D dual-tree complex wavelet transform (DTCWT) with native
14
+ complex tensors, DTCWT scattering networks, the stationary wavelet transform
15
+ (SWT) and wavelet packets — built to respect PyTorch's precision, device and
16
+ autocast semantics.
17
+
18
+ ## Why another wavelet library?
19
+
20
+ The reference implementation, [`fbcotter/pytorch_wavelets`](https://github.com/fbcotter/pytorch_wavelets),
21
+ is unmaintained (last commit 2023) and has well-known problems:
22
+
23
+ - **AMP / mixed precision is broken** — the custom `autograd.Function`s run
24
+ their backward in float32 against float16 gradients (`expected scalar type
25
+ Half but found Float`, issue #63).
26
+ - **Filters are baked in at construction time** via `torch.get_default_dtype()`;
27
+ feeding a float16 tensor to a float32-built transform fails.
28
+ - **Boundary handling** does not match PyWavelets exactly and fails for small
29
+ signals; the stationary wavelet inverse is dead code.
30
+ - No 1D/3D DTCWT, no wavelet packets, no native complex tensors, stale CI.
31
+
32
+ `torchwavelets` fixes these at the design level:
33
+
34
+ ## Design principles
35
+
36
+ 1. **Output dtype == input dtype.** fp16/bf16/fp32/fp64 in → the same out.
37
+ Filters are resolved lazily per `(dtype, device)` and cached; nothing is
38
+ mutated during forward, so `torch.compile`, DDP and mixed-precision
39
+ training all behave. (One documented exception: the DTCWT family computes
40
+ half and bfloat16 inputs in float32 — complex16 kernels do not exist — and
41
+ runs in an autocast-disabled scope, returning float32.)
42
+ 2. **Autocast-native.** The transforms are built from plain `conv2d` /
43
+ `conv_transpose2d` ops (no custom `autograd.Function`), so `torch.autocast`
44
+ and `torch.amp.GradScaler` work without any special-casing. The DTCWT
45
+ family computes in an autocast-disabled scope (its complex bandpasses have
46
+ no bf16 kernels) and returns float32 for half/bf16 inputs.
47
+ 3. **PyWavelets-exact boundary handling.** All modes (`zero`, `symmetric`,
48
+ `reflect`, `periodization`, `periodic`, `constant`, `replicate`) match
49
+ `pywt.dwt` / `pywt.idwt` bit-for-bit, including odd-length and tiny signals
50
+ (index-based padding, no size limits).
51
+ 4. **Pure torch ops.** Index-based padding keeps everything differentiable,
52
+ higher-order differentiable, and `torch.compile`-friendly.
53
+
54
+ ## Installation
55
+
56
+ ```bash
57
+ pip install -e .[test,examples]
58
+ ```
59
+
60
+ ## Usage
61
+
62
+ ```python
63
+ import torch
64
+ import torchwavelets as tw
65
+
66
+ x = torch.randn(4, 3, 64, 64) # (N, C, H, W)
67
+
68
+ # Functional API
69
+ yl, yh = tw.dwt2d(x, wave="db4", mode="symmetric", level=3)
70
+ xr = tw.idwt2d((yl, yh), wave="db4", mode="symmetric")
71
+
72
+ # Module API
73
+ dwt = tw.DWT2D(wave="db4", J=3, mode="symmetric")
74
+ yl, yh = dwt(x)
75
+ xr = tw.IDWT2D(wave="db4", mode="symmetric")((yl, yh))
76
+
77
+ # 1D
78
+ y = torch.randn(4, 3, 1024)
79
+ yl, yh = tw.dwt1d(y, wave="db6", level=3)
80
+
81
+ # Mixed precision just works
82
+ x16 = x.half()
83
+ yl16, _ = dwt(x16) # -> float16 output
84
+ with torch.autocast("cuda", dtype=torch.bfloat16):
85
+ yl_bf, _ = dwt(x.float()) # -> bfloat16 output (on CUDA)
86
+ ```
87
+
88
+ ## Compatibility with `pytorch_wavelets`
89
+
90
+ ```python
91
+ from torchwavelets.compat import DWTForward, DWTInverse, DWT1DForward, DWT1DInverse
92
+
93
+ xfm = DWTForward(J=3, wave="db3", mode="zero")
94
+ ```
95
+
96
+ The compat layer keeps the old API but fixes the old semantics: a float16
97
+ input works regardless of the dtype used when the transform was constructed.
98
+
99
+ ## Status
100
+
101
+ - [x] DWT 1D / 2D, analysis + synthesis, all boundary modes
102
+ - [x] float16 / bfloat16 / float32 / float64
103
+ - [x] PyWavelets numerical parity tests, gradient checks, AMP tests
104
+ - [x] `torch.compile`-compatible module API (`fullgraph=True`)
105
+ - [x] Device-portable: CPU, CUDA, multi-GPU, hybrid and Grace (GH200) nodes
106
+ - [x] DTCWT 1D / 2D / 3D (native complex tensors, `near_sym_*` / `qshift_*` banks)
107
+ - [x] DTCWT scattering layers (`ScatLayer`, `ScatLayerj2`)
108
+ - [x] General arbitrary-order scattering network (`DTCWTScattering`)
109
+ - [x] Stationary (undecimated) wavelet transform 1D / 2D
110
+ - [x] Wavelet packets 1D / 2D (`WaveletPacket`, `WaveletPacket2D`)
111
+ - [x] Scattering network (deeper, arbitrary-order)
112
+
113
+ ## Examples and applications
114
+
115
+ Self-contained, reproducible demos with real data (photograph, CIFAR-10,
116
+ MIT-BIH ECG) — run `python examples/<script>.py` from the repository root
117
+ (data is cached in `~/data/torchwavelets`, figures land in `docs/figures/`).
118
+ See [`examples/README.md`](examples/README.md) for the quick reference and
119
+ [`docs/examples.md`](docs/examples.md) for the full illustrated walkthrough
120
+ with results and figures.
121
+
122
+ <p align="center">
123
+ <a href="https://github.com/relmonta/torchwavelets/blob/main/docs/figures/transform_gallery.png">
124
+ <img src="https://raw.githubusercontent.com/relmonta/torchwavelets/main/docs/figures/transform_gallery.png" alt="transform gallery" width="98%">
125
+ </a>
126
+ </p>
127
+
128
+ **Application 1 — image denoising (DWT vs DTCWT).** A photograph corrupted
129
+ with Gaussian noise (sigma = 25/255, PSNR 20.3 dB) is restored by per-level
130
+ 3-sigma soft thresholding of the wavelet coefficients. The redundant, nearly
131
+ shift-invariant DTCWT gains **+3.1 dB PSNR** over the critically sampled DWT
132
+ (29.3 dB vs 26.2 dB).
133
+
134
+ <p align="center">
135
+ <a href="https://github.com/relmonta/torchwavelets/blob/main/docs/figures/denoising.png">
136
+ <img src="https://raw.githubusercontent.com/relmonta/torchwavelets/main/docs/figures/denoising.png" alt="image denoising" width="98%">
137
+ </a>
138
+ </p>
139
+
140
+ **Application 2 — ECG denoising (SWT vs DWT).** Real 360 Hz ECG (MIT-BIH
141
+ Arrhythmia record 100) plus 50 uV noise (SNR 10.6 dB), restored by the same
142
+ per-level soft thresholding. The undecimated SWT gains **+1.9 dB SNR** over
143
+ the DWT (14.9 dB vs 13.0 dB).
144
+
145
+ <p align="center">
146
+ <a href="https://github.com/relmonta/torchwavelets/blob/main/docs/figures/ecg_denoising.png">
147
+ <img src="https://raw.githubusercontent.com/relmonta/torchwavelets/main/docs/figures/ecg_denoising.png" alt="ecg denoising" width="98%">
148
+ </a>
149
+ </p>
150
+
151
+ **Application 3 — CIFAR-10 classification with scattering features.** A fixed
152
+ (untrained) DTCWT scattering front-end (`DTCWTScattering(J=2, max_order=2)`)
153
+ extracts 588 shift-invariant dimensions per 32x32 image (147 invariants,
154
+ 2x2-average pooled); a logistic regression reaches **56.5 %** test accuracy,
155
+ versus **31.6 %** on 3072 raw pixels — with 5x fewer dimensions and no
156
+ learned front-end.
157
+
158
+ <p align="center">
159
+ <a href="https://github.com/relmonta/torchwavelets/blob/main/docs/figures/cifar10_scattering.png">
160
+ <img src="https://raw.githubusercontent.com/relmonta/torchwavelets/main/docs/figures/cifar10_scattering.png" alt="cifar10 scattering" width="98%">
161
+ </a>
162
+ </p>
163
+
164
+ ## `torch.compile`
165
+
166
+ The module API (`DWT2D`, `IDWT2D`, `DWT1D`, `IDWT1D`) resolves its filters
167
+ into buffers at construction time and runs only torch ops in `forward`, so it
168
+ is fully traceable:
169
+
170
+ ```python
171
+ import torch
172
+ import torchwavelets as tw
173
+
174
+ dwt = torch.compile(tw.DWT2D(wave="db4", J=3, mode="symmetric"), fullgraph=True)
175
+ yl, yh = dwt(x)
176
+ ```
177
+
178
+ The functional API (`tw.dwt2d(x, wave=...)`) is eager-only for the filter
179
+ resolution step; use the module API under `torch.compile`.
180
+
181
+ ## Ecosystem notes
182
+
183
+ * **Serialization**: all transforms are plain `nn.Module`s — `state_dict`
184
+ save/load and `torch.save`/`torch.load` work as usual.
185
+ * **DataLoader / datasets**: the module API takes batched `(N, C, ...)`
186
+ tensors, so it slots directly into `DataLoader` pipelines.
187
+ * **`torch.func`**: `jacrev`/`jacfwd` and double-backward work through the
188
+ transforms (all plain torch ops).
189
+ * **`torch.distributed` (DDP)**: modules hold no state beyond filter buffers
190
+ and never mutate inputs, so DDP training works (smoke-tested with gloo).
191
+ * **`torch.jit.script` is not supported** (the transforms use dynamic loops
192
+ over levels); use `torch.compile` instead, which is fully supported.
193
+ * **`torch.vmap`**: the API is already batched (`N, C, ...` required), so
194
+ vmap is unnecessary; the batch dimension plays that role.
195
+ * **`torch.inference_mode`**: supported everywhere.
196
+
197
+ ## Device portability
198
+
199
+ Tests parametrize over every available device (`CPU` plus all CUDA devices),
200
+ so the suite automatically covers hybrid CPU/GPU nodes and NVIDIA
201
+ Grace-Hopper (ARM64) machines; GPU results are reported in the CHANGELOG.
202
+
203
+ ## Credits and acknowledgements
204
+
205
+ `torchwavelets` is an independent implementation, but it stands on the
206
+ shoulders of the wavelet community. The parts of the code that are ports or
207
+ adaptations are credited here; everything else was written from the
208
+ mathematics and the public conventions.
209
+
210
+ * **Nick Kingsbury (Cambridge University)** — the dual-tree complex wavelet
211
+ transform itself. The filter banks (`near_sym_a/b`, `qshift_a/b`, ...,
212
+ shipped in `torchwavelets/dtcwt/data/`), the quarter-sample-shift design,
213
+ the orientation convention and the `symmetric` (half-sample) boundary
214
+ handling originate from his DTCWT toolbox, which is **free for research
215
+ use only**. The 3-D DTCWT follows the same work-grid construction.
216
+ * **[`fbcotter/pytorch_wavelets`](https://github.com/fbcotter/pytorch_wavelets)** —
217
+ the unmaintained reference implementation this package replaces. Its
218
+ low-level DTCWT operations (`colfilter`, `rowfilter`, `coldfilt`,
219
+ `rowdfilt`, `colifilt`, `rowifilt`, `q2c`, `c2q`, orientation assembly),
220
+ the two-scale scattering layers (`ScatLayer`, `ScatLayerj2`) and the
221
+ `pytorch_wavelets`-compatible API (`torchwavelets.compat`) are ports of
222
+ that codebase, rewritten on plain torch ops to fix its AMP, dtype and
223
+ boundary problems.
224
+ * **[PyWavelets](https://pywavelets.readthedocs.io/)** — the parity oracle
225
+ for the DWT, SWT and wavelet packets. The boundary modes, coefficient
226
+ shapes, `swt` conventions, packet paths/sequency ordering and the
227
+ "drop extra sample" reconstruction rule all mirror PyWavelets exactly;
228
+ the test suite is bit-exact against it.
229
+ * **[`dtcwt`](https://github.com/rgommers/dtcwt)** (numpy reference
230
+ implementation) — the DTCWT parity oracle for 1-D/2-D/3-D; the test suite
231
+ is bit-exact against it to ~1e-15.
232
+ * **Stéphane Mallat** — the scattering-transform theory. The scattering
233
+ layers and the arbitrary-order `DTCWTScattering` network implement the
234
+ Mallat recipe (wavelet modulus + smoothing) with a DTCWT front-end.
235
+ * **PyTorch ecosystem** — `torch.compile`, autocast and `nn.Module`
236
+ conventions; the tests follow the standard PyTorch-ecosystem usage
237
+ patterns (serialization, DataLoader, `torch.func`, DDP).
238
+
239
+ If you use this package in research, please cite it — see
240
+ [`CITATION.cff`](CITATION.cff).
241
+
242
+ ### Primary references
243
+
244
+ 1. N. G. Kingsbury, *The dual-tree complex wavelet transform: a new efficient
245
+ tool for image restoration and enhancement*, EUSIPCO 1998.
246
+ 2. N. G. Kingsbury, *Complex wavelets for shift invariant analysis and
247
+ filtering of signals*, Applied and Computational Harmonic Analysis, 2001.
248
+ 3. I. W. Selesnick, R. G. Baraniuk, N. G. Kingsbury, *The dual-tree complex
249
+ wavelet transform*, IEEE Signal Processing Magazine, 2005.
250
+ 4. S. Mallat, *Group invariant scattering*, Communications on Pure and
251
+ Applied Mathematics, 2012.
252
+ 5. I. Daubechies, *Ten Lectures on Wavelets*, SIAM, 1992.
253
+
254
+ ## License
255
+
256
+ BSD-3-Clause (see `LICENSE`). See
257
+ [Credits and acknowledgements](#credits-and-acknowledgements) for the
258
+ upstream sources: the DTCWT filter banks originate from Nick Kingsbury's
259
+ MATLAB toolbox, which is free for **research use only** — relevant if you
260
+ plan commercial use of the DTCWT parts.