Projekt_AI-Automatyczny_saper/venv/Lib/site-packages/pywt/_wavelet_packets.py
2021-06-01 17:38:31 +02:00

734 lines
24 KiB
Python

# Copyright (c) 2006-2012 Filip Wasilewski <http://en.ig.ma/>
# Copyright (c) 2012-2016 The PyWavelets Developers
# <https://github.com/PyWavelets/pywt>
# See COPYING for license details.
"""1D and 2D Wavelet packet transform module."""
from __future__ import division, print_function, absolute_import
__all__ = ["BaseNode", "Node", "WaveletPacket", "Node2D", "WaveletPacket2D"]
import numpy as np
from ._extensions._pywt import Wavelet, _check_dtype
from ._dwt import dwt, idwt, dwt_max_level
from ._multidim import dwt2, idwt2
def get_graycode_order(level, x='a', y='d'):
graycode_order = [x, y]
for i in range(level - 1):
graycode_order = [x + path for path in graycode_order] + \
[y + path for path in graycode_order[::-1]]
return graycode_order
class BaseNode(object):
"""
BaseNode for wavelet packet 1D and 2D tree nodes.
The BaseNode is a base class for `Node` and `Node2D`.
It should not be used directly unless creating a new transformation
type. It is included here to document the common interface of 1D
and 2D node and wavelet packet transform classes.
Parameters
----------
parent :
Parent node. If parent is None then the node is considered detached
(ie root).
data : 1D or 2D array
Data associated with the node. 1D or 2D numeric array, depending on the
transform type.
node_name :
A name identifying the coefficients type.
See `Node.node_name` and `Node2D.node_name`
for information on the accepted subnodes names.
"""
# PART_LEN and PARTS attributes that define path tokens for node[] lookup
# must be defined in subclasses.
PART_LEN = None
PARTS = None
def __init__(self, parent, data, node_name):
self.parent = parent
if parent is not None:
self.wavelet = parent.wavelet
self.mode = parent.mode
self.level = parent.level + 1
self._maxlevel = parent.maxlevel
self.path = parent.path + node_name
else:
self.wavelet = None
self.mode = None
self.path = ""
self.level = 0
# data - signal on level 0, coeffs on higher levels
self.data = data
# Need to retain original data size/shape so we can trim any excess
# boundary coefficients from the inverse transform.
if self.data is None:
self._data_shape = None
else:
self._data_shape = np.asarray(data).shape
self._init_subnodes()
def _init_subnodes(self):
for part in self.PARTS:
self._set_node(part, None)
def _create_subnode(self, part, data=None, overwrite=True):
raise NotImplementedError()
def _create_subnode_base(self, node_cls, part, data=None, overwrite=True):
self._validate_node_name(part)
if not overwrite and self._get_node(part) is not None:
return self._get_node(part)
node = node_cls(self, data, part)
self._set_node(part, node)
return node
def _get_node(self, part):
return getattr(self, part)
def _set_node(self, part, node):
setattr(self, part, node)
def _delete_node(self, part):
self._set_node(part, None)
def _validate_node_name(self, part):
if part not in self.PARTS:
raise ValueError("Subnode name must be in [%s], not '%s'." %
(', '.join("'%s'" % p for p in self.PARTS), part))
def _evaluate_maxlevel(self, evaluate_from='parent'):
"""
Try to find the value of maximum decomposition level if it is not
specified explicitly.
Parameters
----------
evaluate_from : {'parent', 'subnodes'}
"""
assert evaluate_from in ('parent', 'subnodes')
if self._maxlevel is not None:
return self._maxlevel
elif self.data is not None:
return self.level + dwt_max_level(
min(self.data.shape), self.wavelet)
if evaluate_from == 'parent':
if self.parent is not None:
return self.parent._evaluate_maxlevel(evaluate_from)
elif evaluate_from == 'subnodes':
for node_name in self.PARTS:
node = getattr(self, node_name, None)
if node is not None:
level = node._evaluate_maxlevel(evaluate_from)
if level is not None:
return level
return None
@property
def maxlevel(self):
if self._maxlevel is not None:
return self._maxlevel
# Try getting the maxlevel from parents first
self._maxlevel = self._evaluate_maxlevel(evaluate_from='parent')
# If not found, check whether it can be evaluated from subnodes
if self._maxlevel is None:
self._maxlevel = self._evaluate_maxlevel(evaluate_from='subnodes')
return self._maxlevel
@property
def node_name(self):
return self.path[-self.PART_LEN:]
def decompose(self):
"""
Decompose node data creating DWT coefficients subnodes.
Performs Discrete Wavelet Transform on the `~BaseNode.data` and
returns transform coefficients.
Note
----
Descends to subnodes and recursively
calls `~BaseNode.reconstruct` on them.
"""
if self.level < self.maxlevel:
return self._decompose()
else:
raise ValueError("Maximum decomposition level reached.")
def _decompose(self):
raise NotImplementedError()
def reconstruct(self, update=False):
"""
Reconstruct node from subnodes.
Parameters
----------
update : bool, optional
If True, then reconstructed data replaces the current
node data (default: False).
Returns:
- original node data if subnodes do not exist
- IDWT of subnodes otherwise.
"""
if not self.has_any_subnode:
return self.data
return self._reconstruct(update)
def _reconstruct(self):
raise NotImplementedError() # override this in subclasses
def get_subnode(self, part, decompose=True):
"""
Returns subnode or None (see `decomposition` flag description).
Parameters
----------
part :
Subnode name
decompose : bool, optional
If the param is True and corresponding subnode does not
exist, the subnode will be created using coefficients
from the DWT decomposition of the current node.
(default: True)
"""
self._validate_node_name(part)
subnode = self._get_node(part)
if subnode is None and decompose and not self.is_empty:
self.decompose()
subnode = self._get_node(part)
return subnode
def __getitem__(self, path):
"""
Find node represented by the given path.
Similar to `~BaseNode.get_subnode` method with `decompose=True`, but
can access nodes on any level in the decomposition tree.
Parameters
----------
path : str
String composed of node names. See `Node.node_name` and
`Node2D.node_name` for node naming convention.
Notes
-----
If node does not exist yet, it will be created by decomposition of its
parent node.
"""
if isinstance(path, str):
if (self.maxlevel is not None
and len(path) > self.maxlevel * self.PART_LEN):
raise IndexError("Path length is out of range.")
if path:
return self.get_subnode(path[0:self.PART_LEN], True)[
path[self.PART_LEN:]]
else:
return self
else:
raise TypeError("Invalid path parameter type - expected string but"
" got %s." % type(path))
def __setitem__(self, path, data):
"""
Set node or node's data in the decomposition tree. Nodes are
identified by string `path`.
Parameters
----------
path : str
String composed of node names.
data : array or BaseNode subclass.
"""
if isinstance(path, str):
if (
self.maxlevel is not None
and len(self.path) + len(path) > self.maxlevel * self.PART_LEN
):
raise IndexError("Path length out of range.")
if path:
subnode = self.get_subnode(path[0:self.PART_LEN], False)
if subnode is None:
self._create_subnode(path[0:self.PART_LEN], None)
subnode = self.get_subnode(path[0:self.PART_LEN], False)
subnode[path[self.PART_LEN:]] = data
else:
if isinstance(data, BaseNode):
self.data = np.asarray(data.data)
else:
self.data = np.asarray(data)
# convert data to nearest supported dtype
dtype = _check_dtype(data)
if self.data.dtype != dtype:
self.data = self.data.astype(dtype)
else:
raise TypeError("Invalid path parameter type - expected string but"
" got %s." % type(path))
def __delitem__(self, path):
"""
Remove node from the tree.
Parameters
----------
path : str
String composed of node names.
"""
node = self[path]
# don't clear node value and subnodes (node may still exist outside
# the tree)
# # node._init_subnodes()
# # node.data = None
parent = node.parent
node.parent = None # TODO
if parent and node.node_name:
parent._delete_node(node.node_name)
@property
def is_empty(self):
return self.data is None
@property
def has_any_subnode(self):
for part in self.PARTS:
if self._get_node(part) is not None: # and not .is_empty
return True
return False
def get_leaf_nodes(self, decompose=False):
"""
Returns leaf nodes.
Parameters
----------
decompose : bool, optional
(default: True)
"""
result = []
def collect(node):
if node.level == node.maxlevel and not node.is_empty:
result.append(node)
return False
if not decompose and not node.has_any_subnode:
result.append(node)
return False
return True
self.walk(collect, decompose=decompose)
return result
def walk(self, func, args=(), kwargs=None, decompose=True):
"""
Traverses the decomposition tree and calls
``func(node, *args, **kwargs)`` on every node. If `func` returns True,
descending to subnodes will continue.
Parameters
----------
func : callable
Callable accepting `BaseNode` as the first param and
optional positional and keyword arguments
args :
func params
kwargs :
func keyword params
decompose : bool, optional
If True (default), the method will also try to decompose the tree
up to the `maximum level <BaseNode.maxlevel>`.
"""
if kwargs is None:
kwargs = {}
if func(self, *args, **kwargs) and self.level < self.maxlevel:
for part in self.PARTS:
subnode = self.get_subnode(part, decompose)
if subnode is not None:
subnode.walk(func, args, kwargs, decompose)
def walk_depth(self, func, args=(), kwargs=None, decompose=True):
"""
Walk tree and call func on every node starting from the bottom-most
nodes.
Parameters
----------
func : callable
Callable accepting :class:`BaseNode` as the first param and
optional positional and keyword arguments
args :
func params
kwargs :
func keyword params
decompose : bool, optional
(default: False)
"""
if kwargs is None:
kwargs = {}
if self.level < self.maxlevel:
for part in self.PARTS:
subnode = self.get_subnode(part, decompose)
if subnode is not None:
subnode.walk_depth(func, args, kwargs, decompose)
func(self, *args, **kwargs)
def __str__(self):
return self.path + ": " + str(self.data)
class Node(BaseNode):
"""
WaveletPacket tree node.
Subnodes are called `a` and `d`, just like approximation
and detail coefficients in the Discrete Wavelet Transform.
"""
A = 'a'
D = 'd'
PARTS = A, D
PART_LEN = 1
def _create_subnode(self, part, data=None, overwrite=True):
return self._create_subnode_base(node_cls=Node, part=part, data=data,
overwrite=overwrite)
def _decompose(self):
"""
See also
--------
dwt : for 1D Discrete Wavelet Transform output coefficients.
"""
if self.is_empty:
data_a, data_d = None, None
if self._get_node(self.A) is None:
self._create_subnode(self.A, data_a)
if self._get_node(self.D) is None:
self._create_subnode(self.D, data_d)
else:
data_a, data_d = dwt(self.data, self.wavelet, self.mode)
self._create_subnode(self.A, data_a)
self._create_subnode(self.D, data_d)
return self._get_node(self.A), self._get_node(self.D)
def _reconstruct(self, update):
data_a, data_d = None, None
node_a, node_d = self._get_node(self.A), self._get_node(self.D)
if node_a is not None:
data_a = node_a.reconstruct() # TODO: (update) ???
if node_d is not None:
data_d = node_d.reconstruct() # TODO: (update) ???
if data_a is None and data_d is None:
raise ValueError("Node is a leaf node and cannot be reconstructed"
" from subnodes.")
else:
rec = idwt(data_a, data_d, self.wavelet, self.mode)
if self._data_shape is not None and (
rec.shape != self._data_shape):
rec = rec[tuple([slice(sz) for sz in self._data_shape])]
if update:
self.data = rec
return rec
class Node2D(BaseNode):
"""
WaveletPacket tree node.
Subnodes are called 'a' (LL), 'h' (HL), 'v' (LH) and 'd' (HH), like
approximation and detail coefficients in the 2D Discrete Wavelet Transform
"""
LL = 'a'
HL = 'h'
LH = 'v'
HH = 'd'
PARTS = LL, HL, LH, HH
PART_LEN = 1
def _create_subnode(self, part, data=None, overwrite=True):
return self._create_subnode_base(node_cls=Node2D, part=part, data=data,
overwrite=overwrite)
def _decompose(self):
"""
See also
--------
dwt2 : for 2D Discrete Wavelet Transform output coefficients.
"""
if self.is_empty:
data_ll, data_lh, data_hl, data_hh = None, None, None, None
else:
data_ll, (data_hl, data_lh, data_hh) =\
dwt2(self.data, self.wavelet, self.mode)
self._create_subnode(self.LL, data_ll)
self._create_subnode(self.LH, data_lh)
self._create_subnode(self.HL, data_hl)
self._create_subnode(self.HH, data_hh)
return (self._get_node(self.LL), self._get_node(self.HL),
self._get_node(self.LH), self._get_node(self.HH))
def _reconstruct(self, update):
data_ll, data_lh, data_hl, data_hh = None, None, None, None
node_ll, node_lh, node_hl, node_hh =\
self._get_node(self.LL), self._get_node(self.LH),\
self._get_node(self.HL), self._get_node(self.HH)
if node_ll is not None:
data_ll = node_ll.reconstruct()
if node_lh is not None:
data_lh = node_lh.reconstruct()
if node_hl is not None:
data_hl = node_hl.reconstruct()
if node_hh is not None:
data_hh = node_hh.reconstruct()
if (data_ll is None and data_lh is None
and data_hl is None and data_hh is None):
raise ValueError(
"Tree is missing data - all subnodes of `%s` node "
"are None. Cannot reconstruct node." % self.path
)
else:
coeffs = data_ll, (data_hl, data_lh, data_hh)
rec = idwt2(coeffs, self.wavelet, self.mode)
if self._data_shape is not None and (
rec.shape != self._data_shape):
rec = rec[tuple([slice(sz) for sz in self._data_shape])]
if update:
self.data = rec
return rec
def expand_2d_path(self, path):
expanded_paths = {
self.HH: 'hh',
self.HL: 'hl',
self.LH: 'lh',
self.LL: 'll'
}
return (''.join([expanded_paths[p][0] for p in path]),
''.join([expanded_paths[p][1] for p in path]))
class WaveletPacket(Node):
"""
Data structure representing Wavelet Packet decomposition of signal.
Parameters
----------
data : 1D ndarray
Original data (signal)
wavelet : Wavelet object or name string
Wavelet used in DWT decomposition and reconstruction
mode : str, optional
Signal extension mode for the `dwt` and `idwt` decomposition and
reconstruction functions.
maxlevel : int, optional
Maximum level of decomposition.
If None, it will be calculated based on the `wavelet` and `data`
length using `pywt.dwt_max_level`.
"""
def __init__(self, data, wavelet, mode='symmetric', maxlevel=None):
super(WaveletPacket, self).__init__(None, data, "")
if not isinstance(wavelet, Wavelet):
wavelet = Wavelet(wavelet)
self.wavelet = wavelet
self.mode = mode
if data is not None:
data = np.asarray(data)
assert data.ndim == 1
self.data_size = data.shape[0]
if maxlevel is None:
maxlevel = dwt_max_level(self.data_size, self.wavelet)
else:
self.data_size = None
self._maxlevel = maxlevel
def reconstruct(self, update=True):
"""
Reconstruct data value using coefficients from subnodes.
Parameters
----------
update : bool, optional
If True (default), then data values will be replaced by
reconstruction values, also in subnodes.
"""
if self.has_any_subnode:
data = super(WaveletPacket, self).reconstruct(update)
if update:
self.data = data
return data
return self.data # return original data
def get_level(self, level, order="natural", decompose=True):
"""
Returns all nodes on the specified level.
Parameters
----------
level : int
Specifies decomposition `level` from which the nodes will be
collected.
order : {'natural', 'freq'}, optional
- "natural" - left to right in tree (default)
- "freq" - band ordered
decompose : bool, optional
If set then the method will try to decompose the data up
to the specified `level` (default: True).
Notes
-----
If nodes at the given level are missing (i.e. the tree is partially
decomposed) and the `decompose` is set to False, only existing nodes
will be returned.
"""
assert order in ["natural", "freq"]
if level > self.maxlevel:
raise ValueError("The level cannot be greater than the maximum"
" decomposition level value (%d)" % self.maxlevel)
result = []
def collect(node):
if node.level == level:
result.append(node)
return False
return True
self.walk(collect, decompose=decompose)
if order == "natural":
return result
elif order == "freq":
result = dict((node.path, node) for node in result)
graycode_order = get_graycode_order(level)
return [result[path] for path in graycode_order if path in result]
else:
raise ValueError("Invalid order name - %s." % order)
class WaveletPacket2D(Node2D):
"""
Data structure representing 2D Wavelet Packet decomposition of signal.
Parameters
----------
data : 2D ndarray
Data associated with the node.
wavelet : Wavelet object or name string
Wavelet used in DWT decomposition and reconstruction
mode : str, optional
Signal extension mode for the `dwt` and `idwt` decomposition and
reconstruction functions.
maxlevel : int
Maximum level of decomposition.
If None, it will be calculated based on the `wavelet` and `data`
length using `pywt.dwt_max_level`.
"""
def __init__(self, data, wavelet, mode='smooth', maxlevel=None):
super(WaveletPacket2D, self).__init__(None, data, "")
if not isinstance(wavelet, Wavelet):
wavelet = Wavelet(wavelet)
self.wavelet = wavelet
self.mode = mode
if data is not None:
data = np.asarray(data)
assert data.ndim == 2
self.data_size = data.shape
if maxlevel is None:
maxlevel = dwt_max_level(min(self.data_size), self.wavelet)
else:
self.data_size = None
self._maxlevel = maxlevel
def reconstruct(self, update=True):
"""
Reconstruct data using coefficients from subnodes.
Parameters
----------
update : bool, optional
If True (default) then the coefficients of the current node
and its subnodes will be replaced with values from reconstruction.
"""
if self.has_any_subnode:
data = super(WaveletPacket2D, self).reconstruct(update)
if update:
self.data = data
return data
return self.data # return original data
def get_level(self, level, order="natural", decompose=True):
"""
Returns all nodes from specified level.
Parameters
----------
level : int
Decomposition `level` from which the nodes will be
collected.
order : {'natural', 'freq'}, optional
If `natural` (default) a flat list is returned.
If `freq`, a 2d structure with rows and cols
sorted by corresponding dimension frequency of 2d
coefficient array (adapted from 1d case).
decompose : bool, optional
If set then the method will try to decompose the data up
to the specified `level` (default: True).
"""
assert order in ["natural", "freq"]
if level > self.maxlevel:
raise ValueError("The level cannot be greater than the maximum"
" decomposition level value (%d)" % self.maxlevel)
result = []
def collect(node):
if node.level == level:
result.append(node)
return False
return True
self.walk(collect, decompose=decompose)
if order == "freq":
nodes = {}
for (row_path, col_path), node in [
(self.expand_2d_path(node.path), node) for node in result
]:
nodes.setdefault(row_path, {})[col_path] = node
graycode_order = get_graycode_order(level, x='l', y='h')
nodes = [nodes[path] for path in graycode_order if path in nodes]
result = []
for row in nodes:
result.append(
[row[path] for path in graycode_order if path in row]
)
return result