Aluode/PerceptionLabPortable
0
1# Copyright (c) 2006-2012 Filip Wasilewski <http://en.ig.ma/>
2# Copyright (c) 2012-2016 The PyWavelets Developers
3# <https://github.com/PyWavelets/pywt>
4# See COPYING for license details.
5
6"""1D and 2D Wavelet packet transform module."""
7
8
9__all__ = ["BaseNode", "Node", "WaveletPacket", "Node2D", "WaveletPacket2D",
10 "NodeND", "WaveletPacketND"]
11
12from collections import OrderedDict
13from itertools import product
14
15import numpy as np
16
17from ._dwt import dwt, dwt_max_level, idwt
18from ._extensions._pywt import Wavelet, _check_dtype
19from ._multidim import dwt2, dwtn, idwt2, idwtn
20
21
22def get_graycode_order(level, x='a', y='d'):
23 graycode_order = [x, y]
24 for i in range(level - 1):
25 graycode_order = [x + path for path in graycode_order] + \
26 [y + path for path in graycode_order[::-1]]
27 return graycode_order
28
29
30class BaseNode:
31 """
32 BaseNode for wavelet packet 1D and 2D tree nodes.
33
34 The BaseNode is a base class for `Node` and `Node2D`.
35 It should not be used directly unless creating a new transformation
36 type. It is included here to document the common interface of 1D
37 and 2D node and wavelet packet transform classes.
38
39 Parameters
40 ----------
41 parent :
42 Parent node. If parent is None then the node is considered detached
43 (ie root).
44 data : 1D or 2D array
45 Data associated with the node. 1D or 2D numeric array, depending on the
46 transform type.
47 node_name :
48 A name identifying the coefficients type.
49 See `Node.node_name` and `Node2D.node_name`
50 for information on the accepted subnodes names.
51 """
52
53 # PART_LEN and PARTS attributes that define path tokens for node[] lookup
54 # must be defined in subclasses.
55 PART_LEN = None
56 PARTS = None
57
58 def __init__(self, parent, data, node_name):
59 self.parent = parent
60 if parent is not None:
61 self.wavelet = parent.wavelet
62 self.mode = parent.mode
63 self.level = parent.level + 1
64 self._maxlevel = parent.maxlevel
65 self.path = parent.path + node_name
66 self.axes = parent.axes
67 else:
68 self.wavelet = None
69 self.mode = None
70 self.axes = None
71 self.path = ""
72 self.level = 0
73
74 # data - signal on level 0, coeffs on higher levels
75 self.data = data
76 # Need to retain original data size/shape so we can trim any excess
77 # boundary coefficients from the inverse transform.
78 if self.data is None:
79 self._data_shape = None
80 else:
81 self._data_shape = np.asarray(data).shape
82
83 self._init_subnodes()
84
85 def _init_subnodes(self):
86 for part in self.PARTS:
87 self._set_node(part, None)
88
89 def _create_subnode(self, part, data=None, overwrite=True):
90 raise NotImplementedError()
91
92 def _create_subnode_base(self, node_cls, part, data=None, overwrite=True,
93 **kwargs):
94 self._validate_node_name(part)
95 if not overwrite and self._get_node(part) is not None:
96 return self._get_node(part)
97 node = node_cls(self, data, part, **kwargs)
98 self._set_node(part, node)
99 return node
100
101 def _get_node(self, part):
102 return getattr(self, part)
103
104 def _set_node(self, part, node):
105 setattr(self, part, node)
106
107 def _delete_node(self, part):
108 self._set_node(part, None)
109
110 def _validate_node_name(self, part):
111 if part not in self.PARTS:
112 raise ValueError("Subnode name must be in [{}], not '{}'.".format(', '.join("'%s'" % p for p in self.PARTS), part))
113
114 @property
115 def path_tuple(self):
116 """The path to the current node in tuple form.
117
118 The length of the tuple is equal to the number of decomposition levels.
119 """
120 path = self.path
121 nlev = len(path)//self.PART_LEN
122 return tuple([path[(n-1)*self.PART_LEN:n*self.PART_LEN]
123 for n in range(1, nlev+1)])
124
125 def _evaluate_maxlevel(self, evaluate_from='parent'):
126 """
127 Try to find the value of maximum decomposition level if it is not
128 specified explicitly.
129
130 Parameters
131 ----------
132 evaluate_from : {'parent', 'subnodes'}
133 """
134 assert evaluate_from in ('parent', 'subnodes')
135
136 if self._maxlevel is not None:
137 return self._maxlevel
138 elif self.data is not None:
139 return self.level + dwt_max_level(
140 min(self.data.shape), self.wavelet)
141
142 if evaluate_from == 'parent':
143 if self.parent is not None:
144 return self.parent._evaluate_maxlevel(evaluate_from)
145 elif evaluate_from == 'subnodes':
146 for node_name in self.PARTS:
147 node = getattr(self, node_name, None)
148 if node is not None:
149 level = node._evaluate_maxlevel(evaluate_from)
150 if level is not None:
151 return level
152 return None
153
154 @property
155 def maxlevel(self):
156 if self._maxlevel is not None:
157 return self._maxlevel
158
159 # Try getting the maxlevel from parents first
160 self._maxlevel = self._evaluate_maxlevel(evaluate_from='parent')
161
162 # If not found, check whether it can be evaluated from subnodes
163 if self._maxlevel is None:
164 self._maxlevel = self._evaluate_maxlevel(evaluate_from='subnodes')
165 return self._maxlevel
166
167 @property
168 def node_name(self):
169 return self.path[-self.PART_LEN:]
170
171 def decompose(self):
172 """
173 Decompose node data creating DWT coefficients subnodes.
174
175 Performs Discrete Wavelet Transform on the `~BaseNode.data` and
176 returns transform coefficients.
177
178 Note
179 ----
180 Descends to subnodes and recursively
181 calls `~BaseNode.reconstruct` on them.
182
183 """
184 if self.level < self.maxlevel:
185 return self._decompose()
186 else:
187 raise ValueError("Maximum decomposition level reached.")
188
189 def _decompose(self):
190 raise NotImplementedError()
191
192 def reconstruct(self, update=False):
193 """
194 Reconstruct node from subnodes.
195
196 Parameters
197 ----------
198 update : bool, optional
199 If True, then reconstructed data replaces the current
200 node data (default: False).
201
202 Returns:
203 - original node data if subnodes do not exist
204 - IDWT of subnodes otherwise.
205 """
206 if not self.has_any_subnode:
207 return self.data
208 return self._reconstruct(update)
209
210 def _reconstruct(self):
211 raise NotImplementedError() # override this in subclasses
212
213 def get_subnode(self, part, decompose=True):
214 """
215 Returns subnode or None (see `decomposition` flag description).
216
217 Parameters
218 ----------
219 part :
220 Subnode name
221 decompose : bool, optional
222 If the param is True and corresponding subnode does not
223 exist, the subnode will be created using coefficients
224 from the DWT decomposition of the current node.
225 (default: True)
226 """
227 self._validate_node_name(part)
228 subnode = self._get_node(part)
229 if subnode is None and decompose and not self.is_empty:
230 self.decompose()
231 subnode = self._get_node(part)
232 return subnode
233
234 def __getitem__(self, path):
235 """
236 Find node represented by the given path.
237
238 Similar to `~BaseNode.get_subnode` method with `decompose=True`, but
239 can access nodes on any level in the decomposition tree.
240
241 Parameters
242 ----------
243 path : str
244 String composed of node names. See `Node.node_name` and
245 `Node2D.node_name` for node naming convention.
246
247 Notes
248 -----
249 If node does not exist yet, it will be created by decomposition of its
250 parent node.
251 """
252 errmsg = ("Invalid path parameter type - expected string or "
253 "tuple of strings but got %s." % type(path))
254 if isinstance(path, tuple):
255 # concatenate tuple of strings into a single string
256 try:
257 path = ''.join(path)
258 except TypeError:
259 raise TypeError(errmsg)
260 if isinstance(path, str):
261 if (self.maxlevel is not None and
262 len(path) > self.maxlevel * self.PART_LEN):
263 raise IndexError("Path length is out of range.")
264 if path:
265 return self.get_subnode(path[0:self.PART_LEN], True)[
266 path[self.PART_LEN:]]
267 else:
268 return self
269 else:
270 raise TypeError(errmsg)
271
272 def __setitem__(self, path, data):
273 """
274 Set node or node's data in the decomposition tree. Nodes are
275 identified by string `path`.
276
277 Parameters
278 ----------
279 path : str
280 String composed of node names.
281 data : array or BaseNode subclass.
282 """
283
284 if isinstance(path, str):
285 if (
286 self.maxlevel is not None and
287 len(self.path) + len(path) > self.maxlevel * self.PART_LEN
288 ):
289 raise IndexError("Path length out of range.")
290 if path:
291 subnode = self.get_subnode(path[0:self.PART_LEN], False)
292 if subnode is None:
293 self._create_subnode(path[0:self.PART_LEN], None)
294 subnode = self.get_subnode(path[0:self.PART_LEN], False)
295 subnode[path[self.PART_LEN:]] = data
296 else:
297 if isinstance(data, BaseNode):
298 self.data = np.asarray(data.data)
299 else:
300 self.data = np.asarray(data)
301 # convert data to nearest supported dtype
302 dtype = _check_dtype(data)
303 if self.data.dtype != dtype:
304 self.data = self.data.astype(dtype)
305 else:
306 raise TypeError("Invalid path parameter type - expected string but"
307 " got %s." % type(path))
308
309 def __delitem__(self, path):
310 """
311 Remove node from the tree.
312
313 Parameters
314 ----------
315 path : str
316 String composed of node names.
317 """
318 node = self[path]
319 # don't clear node value and subnodes (node may still exist outside
320 # the tree)
321 # # node._init_subnodes()
322 # # node.data = None
323 parent = node.parent
324 node.parent = None # TODO
325 if parent and node.node_name:
326 parent._delete_node(node.node_name)
327
328 @property
329 def is_empty(self):
330 return self.data is None
331
332 @property
333 def has_any_subnode(self):
334 return any(self._get_node(part) is not None for part in self.PARTS)
335
336 def get_leaf_nodes(self, decompose=False):
337 """
338 Returns leaf nodes.
339
340 Parameters
341 ----------
342 decompose : bool, optional
343 (default: True)
344 """
345 result = []
346
347 def collect(node):
348 if node.level == node.maxlevel and not node.is_empty:
349 result.append(node)
350 return False
351 if not decompose and not node.has_any_subnode:
352 result.append(node)
353 return False
354 return True
355 self.walk(collect, decompose=decompose)
356 return result
357
358 def walk(self, func, args=(), kwargs=None, decompose=True):
359 """
360 Traverses the decomposition tree and calls
361 ``func(node, *args, **kwargs)`` on every node. If `func` returns True,
362 descending to subnodes will continue.
363
364 Parameters
365 ----------
366 func : callable
367 Callable accepting `BaseNode` as the first param and
368 optional positional and keyword arguments
369 args :
370 func params
371 kwargs :
372 func keyword params
373 decompose : bool, optional
374 If True (default), the method will also try to decompose the tree
375 up to the `maximum level <BaseNode.maxlevel>`.
376 """
377 if kwargs is None:
378 kwargs = {}
379 if func(self, *args, **kwargs) and self.level < self.maxlevel:
380 for part in self.PARTS:
381 subnode = self.get_subnode(part, decompose)
382 if subnode is not None:
383 subnode.walk(func, args, kwargs, decompose)
384
385 def walk_depth(self, func, args=(), kwargs=None, decompose=True):
386 """
387 Walk tree and call func on every node starting from the bottom-most
388 nodes.
389
390 Parameters
391 ----------
392 func : callable
393 Callable accepting :class:`BaseNode` as the first param and
394 optional positional and keyword arguments
395 args :
396 func params
397 kwargs :
398 func keyword params
399 decompose : bool, optional
400 (default: False)
401 """
402 if kwargs is None:
403 kwargs = {}
404 if self.level < self.maxlevel:
405 for part in self.PARTS:
406 subnode = self.get_subnode(part, decompose)
407 if subnode is not None:
408 subnode.walk_depth(func, args, kwargs, decompose)
409 func(self, *args, **kwargs)
410
411 def __str__(self):
412 return self.path + ": " + str(self.data)
413
414
415class Node(BaseNode):
416 """
417 WaveletPacket tree node.
418
419 Subnodes are called `a` and `d`, just like approximation
420 and detail coefficients in the Discrete Wavelet Transform.
421 """
422
423 A = 'a'
424 D = 'd'
425 PARTS = A, D
426 PART_LEN = 1
427
428 def _create_subnode(self, part, data=None, overwrite=True):
429 return self._create_subnode_base(node_cls=Node, part=part, data=data,
430 overwrite=overwrite)
431
432 def _decompose(self):
433 """
434
435 See also
436 --------
437 dwt : for 1D Discrete Wavelet Transform output coefficients.
438 """
439 if self.is_empty:
440 data_a, data_d = None, None
441 if self._get_node(self.A) is None:
442 self._create_subnode(self.A, data_a)
443 if self._get_node(self.D) is None:
444 self._create_subnode(self.D, data_d)
445 else:
446 data_a, data_d = dwt(self.data, self.wavelet, self.mode,
447 axis=self.axes)
448 self._create_subnode(self.A, data_a)
449 self._create_subnode(self.D, data_d)
450 return self._get_node(self.A), self._get_node(self.D)
451
452 def _reconstruct(self, update):
453 data_a, data_d = None, None
454 node_a, node_d = self._get_node(self.A), self._get_node(self.D)
455
456 if node_a is not None:
457 data_a = node_a.reconstruct() # TODO: (update) ???
458 if node_d is not None:
459 data_d = node_d.reconstruct() # TODO: (update) ???
460
461 if data_a is None and data_d is None:
462 raise ValueError("Node is a leaf node and cannot be reconstructed"
463 " from subnodes.")
464 else:
465 rec = idwt(data_a, data_d, self.wavelet, self.mode, axis=self.axes)
466 if self._data_shape is not None and (
467 rec.shape != self._data_shape):
468 rec = rec[tuple([slice(sz) for sz in self._data_shape])]
469 if update:
470 self.data = rec
471 return rec
472
473
474class Node2D(BaseNode):
475 """
476 WaveletPacket tree node.
477
478 Subnodes are called 'a' (LL), 'h' (HL), 'v' (LH) and 'd' (HH), like
479 approximation and detail coefficients in the 2D Discrete Wavelet Transform
480 """
481
482 LL = 'a'
483 HL = 'h'
484 LH = 'v'
485 HH = 'd'
486
487 PARTS = LL, HL, LH, HH
488 PART_LEN = 1
489
490 def _create_subnode(self, part, data=None, overwrite=True):
491 return self._create_subnode_base(node_cls=Node2D, part=part, data=data,
492 overwrite=overwrite)
493
494 def _decompose(self):
495 """
496 See also
497 --------
498 dwt2 : for 2D Discrete Wavelet Transform output coefficients.
499 """
500 if self.is_empty:
501 data_ll, data_lh, data_hl, data_hh = None, None, None, None
502 else:
503 data_ll, (data_hl, data_lh, data_hh) =\
504 dwt2(self.data, self.wavelet, self.mode, axes=self.axes)
505 self._create_subnode(self.LL, data_ll)
506 self._create_subnode(self.LH, data_lh)
507 self._create_subnode(self.HL, data_hl)
508 self._create_subnode(self.HH, data_hh)
509 return (self._get_node(self.LL), self._get_node(self.HL),
510 self._get_node(self.LH), self._get_node(self.HH))
511
512 def _reconstruct(self, update):
513 data_ll, data_lh, data_hl, data_hh = None, None, None, None
514
515 node_ll, node_lh, node_hl, node_hh =\
516 self._get_node(self.LL), self._get_node(self.LH),\
517 self._get_node(self.HL), self._get_node(self.HH)
518
519 if node_ll is not None:
520 data_ll = node_ll.reconstruct()
521 if node_lh is not None:
522 data_lh = node_lh.reconstruct()
523 if node_hl is not None:
524 data_hl = node_hl.reconstruct()
525 if node_hh is not None:
526 data_hh = node_hh.reconstruct()
527
528 if (data_ll is None and data_lh is None and
529 data_hl is None and data_hh is None):
530 raise ValueError(
531 "Tree is missing data - all subnodes of `%s` node "
532 "are None. Cannot reconstruct node." % self.path
533 )
534 else:
535 coeffs = data_ll, (data_hl, data_lh, data_hh)
536 rec = idwt2(coeffs, self.wavelet, self.mode, axes=self.axes)
537 if self._data_shape is not None and (
538 rec.shape != self._data_shape):
539 rec = rec[tuple([slice(sz) for sz in self._data_shape])]
540 if update:
541 self.data = rec
542 return rec
543
544 def expand_2d_path(self, path):
545 expanded_paths = {
546 self.HH: 'hh',
547 self.HL: 'hl',
548 self.LH: 'lh',
549 self.LL: 'll'
550 }
551 return (''.join([expanded_paths[p][0] for p in path]),
552 ''.join([expanded_paths[p][1] for p in path]))
553
554
555class NodeND(BaseNode):
556 """
557 WaveletPacket tree node.
558
559 Unlike Node and Node2D self.PARTS is a dictionary.
560 For 1D: self.PARTS has keys 'a' and 'd'
561 For 2D: self.PARTS has keys 'aa', 'ad', 'da', 'dd'
562 For 3D: self.PARTS has keys 'aaa', 'aad', 'ada', 'daa', ..., 'ddd'
563
564 Parameters
565 ----------
566 parent :
567 Parent node. If parent is None then the node is considered detached
568 (ie root).
569 data : 1D or 2D array
570 Data associated with the node. 1D or 2D numeric array, depending on the
571 transform type.
572 node_name : string
573 A name identifying the coefficients type.
574 See `Node.node_name` and `Node2D.node_name`
575 for information on the accepted subnodes names.
576 ndim : int
577 The number of data dimensions.
578 ndim_transform : int
579 The number of dimensions that are to be transformed.
580
581 """
582 def __init__(self, parent, data, node_name, ndim, ndim_transform):
583 super().__init__(parent=parent, data=data,
584 node_name=node_name)
585 self.PART_LEN = ndim_transform
586 self.PARTS = OrderedDict()
587 for key in product(*(('ad', )*self.PART_LEN)):
588 self.PARTS[''.join(key)] = None
589 self.ndim = ndim
590 self.ndim_transform = ndim_transform
591
592 def _init_subnodes(self):
593 # need this empty so BaseNode's _init_subnodes isn't called during
594 # __init__. We use a dictionary for PARTS instead for the nd case.
595 pass
596
597 def _get_node(self, part):
598 return self.PARTS[part]
599
600 def _set_node(self, part, node):
601 if part not in self.PARTS:
602 raise ValueError("invalid part")
603 self.PARTS[part] = node
604
605 def _delete_node(self, part):
606 self._set_node(part, None)
607
608 def _validate_node_name(self, part):
609 if part not in self.PARTS:
610 raise ValueError(
611 "Subnode name must be in [{}], not '{}'.".format(', '.join("'%s'" % p for p in list(self.PARTS.keys())), part))
612
613 def _create_subnode(self, part, data=None, overwrite=True):
614 return self._create_subnode_base(node_cls=NodeND, part=part, data=data,
615 overwrite=overwrite, ndim=self.ndim,
616 ndim_transform=self.ndim_transform)
617
618 def _evaluate_maxlevel(self, evaluate_from='parent'):
619 """
620 Try to find the value of maximum decomposition level if it is not
621 specified explicitly.
622
623 Parameters
624 ----------
625 evaluate_from : {'parent', 'subnodes'}
626 """
627 assert evaluate_from in ('parent', 'subnodes')
628
629 if self._maxlevel is not None:
630 return self._maxlevel
631 elif self.data is not None:
632 return self.level + dwt_max_level(
633 min(self.data.shape), self.wavelet)
634
635 if evaluate_from == 'parent':
636 if self.parent is not None:
637 return self.parent._evaluate_maxlevel(evaluate_from)
638 elif evaluate_from == 'subnodes':
639 for node_name, node in self.PARTS.items():
640 if node is not None:
641 level = node._evaluate_maxlevel(evaluate_from)
642 if level is not None:
643 return level
644 return None
645
646 def _decompose(self):
647 """
648 See also
649 --------
650 dwt2 : for 2D Discrete Wavelet Transform output coefficients.
651 """
652 if self.is_empty:
653 coefs = {key: None for key in self.PARTS}
654 else:
655 coefs = dwtn(self.data, self.wavelet, self.mode, axes=self.axes)
656
657 for key, data in coefs.items():
658 self._create_subnode(key, data)
659 return (self._get_node(key) for key in self.PARTS)
660
661 def _reconstruct(self, update):
662 coeffs = {key: None for key in self.PARTS}
663
664 nnodes = 0
665 for key in self.PARTS:
666 node = self._get_node(key)
667 if node is not None:
668 nnodes += 1
669 coeffs[key] = node.reconstruct()
670
671 if nnodes == 0:
672 raise ValueError(
673 "Tree is missing data - all subnodes of `%s` node "
674 "are None. Cannot reconstruct node." % self.path
675 )
676 else:
677 rec = idwtn(coeffs, self.wavelet, self.mode, axes=self.axes)
678 if update:
679 self.data = rec
680 return rec
681
682
683class WaveletPacket(Node):
684 """
685 Data structure representing Wavelet Packet decomposition of signal.
686
687 Parameters
688 ----------
689 data : 1D ndarray
690 Original data (signal)
691 wavelet : Wavelet object or name string
692 Wavelet used in DWT decomposition and reconstruction
693 mode : str, optional
694 Signal extension mode for the `dwt` and `idwt` decomposition and
695 reconstruction functions.
696 maxlevel : int, optional
697 Maximum level of decomposition.
698 If None, it will be calculated based on the `wavelet` and `data`
699 length using `pywt.dwt_max_level`.
700 axis : int, optional
701 The axis to transform.
702 """
703 def __init__(self, data, wavelet, mode='symmetric', maxlevel=None,
704 axis=-1):
705 super().__init__(None, data, "")
706
707 if not isinstance(wavelet, Wavelet):
708 wavelet = Wavelet(wavelet)
709 self.wavelet = wavelet
710 self.mode = mode
711 self.axes = axis # self.axes is just an integer for 1D transforms
712
713 if data is not None:
714 data = np.asarray(data)
715 if self.axes < 0:
716 self.axes = self.axes + data.ndim
717 if not 0 <= self.axes < data.ndim:
718 raise ValueError("Axis greater than data dimensions")
719 self.data_size = data.shape
720 if maxlevel is None:
721 maxlevel = dwt_max_level(data.shape[self.axes], self.wavelet)
722 else:
723 self.data_size = None
724
725 self._maxlevel = maxlevel
726
727 def __reduce__(self):
728 return (WaveletPacket,
729 (self.data, self.wavelet, self.mode, self.maxlevel))
730
731 def reconstruct(self, update=True):
732 """
733 Reconstruct data value using coefficients from subnodes.
734
735 Parameters
736 ----------
737 update : bool, optional
738 If True (default), then data values will be replaced by
739 reconstruction values, also in subnodes.
740 """
741 if self.has_any_subnode:
742 data = super().reconstruct(update)
743 if self.data_size is not None and (data.shape != self.data_size):
744 data = data[[slice(sz) for sz in self.data_size]]
745 if update:
746 self.data = data
747 return data
748 return self.data # return original data
749
750 def get_level(self, level, order="natural", decompose=True):
751 """
752 Returns all nodes on the specified level.
753
754 Parameters
755 ----------
756 level : int
757 Specifies decomposition `level` from which the nodes will be
758 collected.
759 order : {'natural', 'freq'}, optional
760 - "natural" - left to right in tree (default)
761 - "freq" - band ordered
762 decompose : bool, optional
763 If set then the method will try to decompose the data up
764 to the specified `level` (default: True).
765
766 Notes
767 -----
768 If nodes at the given level are missing (i.e. the tree is partially
769 decomposed) and `decompose` is set to False, only existing nodes
770 will be returned.
771
772 Frequency order (``order="freq"``) is also known as sequency order
773 and "natural" order is sometimes referred to as Paley order. A detailed
774 discussion of these orderings is also given in [1]_, [2]_.
775
776 References
777 ----------
778 ..[1] M.V. Wickerhauser. Adapted Wavelet Analysis from Theory to
779 Software. Wellesley. Massachusetts: A K Peters. 1994.
780 ..[2] D.B. Percival and A.T. Walden. Wavelet Methods for Time Series
781 Analysis. Cambridge University Press. 2000.
782 DOI:10.1017/CBO9780511841040
783 """
784 if order not in ["natural", "freq"]:
785 raise ValueError(f"Invalid order: {order}")
786 if level > self.maxlevel:
787 raise ValueError("The level cannot be greater than the maximum"
788 " decomposition level value (%d)" % self.maxlevel)
789
790 result = []
791
792 def collect(node):
793 if node.level == level:
794 result.append(node)
795 return False
796 return True
797
798 self.walk(collect, decompose=decompose)
799 if order == "natural":
800 return result
801 elif order == "freq":
802 result = {node.path: node for node in result}
803 graycode_order = get_graycode_order(level)
804 return [result[path] for path in graycode_order if path in result]
805 else:
806 raise ValueError(f"Invalid order name - {order}.")
807
808
809class WaveletPacket2D(Node2D):
810 """
811 Data structure representing 2D Wavelet Packet decomposition of signal.
812
813 Parameters
814 ----------
815 data : 2D ndarray
816 Data associated with the node.
817 wavelet : Wavelet object or name string
818 Wavelet used in DWT decomposition and reconstruction
819 mode : str, optional
820 Signal extension mode for the `dwt` and `idwt` decomposition and
821 reconstruction functions.
822 maxlevel : int
823 Maximum level of decomposition.
824 If None, it will be calculated based on the `wavelet` and `data`
825 length using `pywt.dwt_max_level`.
826 axes : 2-tuple of ints, optional
827 The axes that will be transformed.
828 """
829 def __init__(self, data, wavelet, mode='smooth', maxlevel=None,
830 axes=(-2, -1)):
831 super().__init__(None, data, "")
832
833 if not isinstance(wavelet, Wavelet):
834 wavelet = Wavelet(wavelet)
835 self.wavelet = wavelet
836 self.mode = mode
837 self.axes = tuple(axes)
838 if len(np.unique(self.axes)) != 2:
839 raise ValueError("Expected two unique axes.")
840 if data is not None:
841 data = np.asarray(data)
842 if data.ndim < 2:
843 raise ValueError(
844 "WaveletPacket2D requires data with 2 or more dimensions.")
845 self.data_size = data.shape
846 transform_size = [data.shape[ax] for ax in self.axes]
847 if maxlevel is None:
848 maxlevel = dwt_max_level(min(transform_size), self.wavelet)
849 else:
850 self.data_size = None
851 self._maxlevel = maxlevel
852
853 def __reduce__(self):
854 return (WaveletPacket2D,
855 (self.data, self.wavelet, self.mode, self.maxlevel))
856
857 def reconstruct(self, update=True):
858 """
859 Reconstruct data using coefficients from subnodes.
860
861 Parameters
862 ----------
863 update : bool, optional
864 If True (default) then the coefficients of the current node
865 and its subnodes will be replaced with values from reconstruction.
866 """
867 if self.has_any_subnode:
868 data = super().reconstruct(update)
869 if self.data_size is not None and (data.shape != self.data_size):
870 data = data[[slice(sz) for sz in self.data_size]]
871 if update:
872 self.data = data
873 return data
874 return self.data # return original data
875
876 def get_level(self, level, order="natural", decompose=True):
877 """
878 Returns all nodes from specified level.
879
880 Parameters
881 ----------
882 level : int
883 Decomposition `level` from which the nodes will be
884 collected.
885 order : {'natural', 'freq'}, optional
886 If `natural` (default) a flat list is returned.
887 If `freq`, a 2d structure with rows and cols
888 sorted by corresponding dimension frequency of 2d
889 coefficient array (adapted from 1d case).
890 decompose : bool, optional
891 If set then the method will try to decompose the data up
892 to the specified `level` (default: True).
893
894 Notes
895 -----
896 Frequency order (``order="freq"``) is also known as as sequency order
897 and "natural" order is sometimes referred to as Paley order. A detailed
898 discussion of these orderings is also given in [1]_, [2]_.
899
900 References
901 ----------
902 ..[1] M.V. Wickerhauser. Adapted Wavelet Analysis from Theory to
903 Software. Wellesley. Massachusetts: A K Peters. 1994.
904 ..[2] D.B. Percival and A.T. Walden. Wavelet Methods for Time Series
905 Analysis. Cambridge University Press. 2000.
906 DOI:10.1017/CBO9780511841040
907 """
908 if order not in ["natural", "freq"]:
909 raise ValueError(f"Invalid order: {order}")
910 if level > self.maxlevel:
911 raise ValueError("The level cannot be greater than the maximum"
912 " decomposition level value (%d)" % self.maxlevel)
913
914 result = []
915
916 def collect(node):
917 if node.level == level:
918 result.append(node)
919 return False
920 return True
921
922 self.walk(collect, decompose=decompose)
923
924 if order == "freq":
925 nodes = {}
926 for (row_path, col_path), node in [
927 (self.expand_2d_path(node.path), node) for node in result
928 ]:
929 nodes.setdefault(row_path, {})[col_path] = node
930 graycode_order = get_graycode_order(level, x='l', y='h')
931 nodes = [nodes[path] for path in graycode_order if path in nodes]
932 result = []
933 for row in nodes:
934 result.append(
935 [row[path] for path in graycode_order if path in row]
936 )
937 return result
938
939
940class WaveletPacketND(NodeND):
941 """
942 Data structure representing ND Wavelet Packet decomposition of signal.
943
944 Parameters
945 ----------
946 data : ND ndarray
947 Data associated with the node.
948 wavelet : Wavelet object or name string
949 Wavelet used in DWT decomposition and reconstruction
950 mode : str, optional
951 Signal extension mode for the `dwt` and `idwt` decomposition and
952 reconstruction functions.
953 maxlevel : int, optional
954 Maximum level of decomposition.
955 If None, it will be calculated based on the `wavelet` and `data`
956 length using `pywt.dwt_max_level`.
957 axes : tuple of int, optional
958 The axes to transform. The default value of `None` corresponds to all
959 axes.
960 """
961 def __init__(self, data, wavelet, mode='smooth', maxlevel=None,
962 axes=None):
963 if (data is None) and (axes is None):
964 # ndim is required to create a NodeND object
965 raise ValueError("If data is None, axes must be specified")
966
967 # axes determines the number of transform dimensions
968 if axes is None:
969 axes = range(data.ndim)
970 elif np.isscalar(axes):
971 axes = (axes, )
972 axes = tuple(axes)
973 if len(np.unique(axes)) != len(axes):
974 raise ValueError("Expected a set of unique axes.")
975 ndim_transform = len(axes)
976
977 if data is not None:
978 data = np.asarray(data)
979 if data.ndim == 0:
980 raise ValueError("data must be at least 1D")
981 ndim = data.ndim
982 else:
983 ndim = len(axes)
984
985 super().__init__(None, data, "", ndim,
986 ndim_transform)
987 if not isinstance(wavelet, Wavelet):
988 wavelet = Wavelet(wavelet)
989 self.wavelet = wavelet
990 self.mode = mode
991 self.axes = axes
992 self.ndim_transform = ndim_transform
993 if data is not None:
994 if data.ndim < len(axes):
995 raise ValueError("The number of axes exceeds the number of "
996 "data dimensions.")
997 self.data_size = data.shape
998 transform_size = [data.shape[ax] for ax in self.axes]
999 if maxlevel is None:
1000 maxlevel = dwt_max_level(min(transform_size), self.wavelet)
1001 else:
1002 self.data_size = None
1003 self._maxlevel = maxlevel
1004
1005 def reconstruct(self, update=True):
1006 """
1007 Reconstruct data using coefficients from subnodes.
1008
1009 Parameters
1010 ----------
1011 update : bool, optional
1012 If True (default) then the coefficients of the current node
1013 and its subnodes will be replaced with values from reconstruction.
1014 """
1015 if self.has_any_subnode:
1016 data = super().reconstruct(update)
1017 if self.data_size is not None and (data.shape != self.data_size):
1018 data = data[[slice(sz) for sz in self.data_size]]
1019 if update:
1020 self.data = data
1021 return data
1022 return self.data # return original data
1023
1024 def get_level(self, level, decompose=True):
1025 """
1026 Returns all nodes from specified level.
1027
1028 Parameters
1029 ----------
1030 level : int
1031 Decomposition `level` from which the nodes will be
1032 collected.
1033 decompose : bool, optional
1034 If set then the method will try to decompose the data up
1035 to the specified `level` (default: True).
1036 """
1037 if level > self.maxlevel:
1038 raise ValueError("The level cannot be greater than the maximum"
1039 " decomposition level value (%d)" % self.maxlevel)
1040
1041 result = []
1042
1043 def collect(node):
1044 if node.level == level:
1045 result.append(node)
1046 return False
1047 return True
1048
1049 self.walk(collect, decompose=decompose)
1050
1051 return result
1052