"""
Misc functions used elsewhere
"""
from __future__ import annotations
import os
import sys
import inspect
import logging as log
from copy import deepcopy
from types import ModuleType
from contextlib import contextmanager
from typing import Any, Literal, Self, Generator, Iterable, Iterator, Generic, cast, overload
import torch
import numpy as np
from torch import nn, Tensor
from numpy import ndarray
from packaging.version import Version
from netloader.utils.types import ArrayT, ModuleT
[docs]
class Shapes(list[list[int] | list[list[int]]]):
"""
Shapes list that allows for indexing of multiple shapes.
"""
@overload # type: ignore[override]
def __getitem__(self, idx: int | tuple[int, int]) -> list[int]: ...
@overload
def __getitem__(self, idx: slice) -> Shapes: ...
def __getitem__(self, idx: int | tuple[int, int] | slice) -> list[int] | Shapes:
"""
Get shape from the shapes list.
Parameters
----------
idx : int | tuple[int, int] | slice
Index of the shape to get, if tuple then first index is the layer index and second is
the shape index if the layer has multiple shapes, if slice, returns a Shapes object with
the sliced shapes
Returns
-------
list[int] | Shapes
Shape at the provided index
"""
return self.get(idx, list_=False)
@overload # type: ignore[override]
def __setitem__(self, idx: slice, value: list[list[int] | list[list[int]]]) -> None: ...
@overload
def __setitem__(self, idx: tuple[int, int], value: list[int]) -> None: ...
@overload
def __setitem__(self, idx: int, value: list[int] | list[list[int]]) -> None: ...
def __setitem__(
self,
idx: int | tuple[int, int] | slice,
value: list[int] | list[list[int]] | list[list[int] | list[list[int]]]) -> None:
"""
Sets shape in the shapes list.
Parameters
----------
idx : int | tuple[int, int] | slice
Index of the shape to set, if tuple then first index is the layer index and second is
the shape index if the layer has multiple shapes, if slice, sets multiple shapes
value : list[int] | list[list[int]]
Shape to set at the provided index
"""
if isinstance(idx, tuple):
if not isinstance(self[idx[0]][0], list):
raise ValueError(f'Index {idx} does not contain multiple shapes, shape at index '
f'is {self[idx[0]]}')
if not isinstance(value[0], int):
raise ValueError(f'Index {idx} requires a single shape, got {value}')
cast(list[list[int]], self[idx[0]])[idx[1]] = cast(list[int], value)
elif isinstance(idx, int):
if isinstance(value[0], list) and isinstance(value[0][0], list):
raise ValueError(f'If index is an int, value must be a single shape or a list of '
f'shapes, got {value}')
super().__setitem__(idx, cast(list[int] | list[list[int]], value))
else:
if isinstance(value[0], int):
raise ValueError(f'If index is a slice, value must be a list of shapes, got '
f'{value}')
super().__setitem__(idx, cast(list[list[int] | list[list[int]]], value))
def __repr__(self) -> str:
return f'Shapes({super().__repr__()})'
[docs]
def check(self, idx: int) -> bool:
"""
Checks if the shape at the provided index is a single shape.
Parameters
----------
idx : int
Index of the shape to check
Returns
-------
bool
True if the shape at the provided index is a single shape
"""
if isinstance(self.get(idx, list_=True)[0], int):
return True
return False
@overload
def get(self, idx: slice, list_: bool) -> Shapes: ...
@overload
def get(self, idx: tuple[int, int], list_: bool) -> list[int]: ...
@overload
def get(self, idx: int, list_: Literal[False]) -> list[int]: ...
@overload
def get(self, idx: int, list_: Literal[True]) -> list[int] | list[list[int]]: ...
[docs]
def get(
self,
idx: int | tuple[int, int] | slice,
list_: bool = True) -> list[int] | list[list[int]] | Shapes:
"""
Get shape from the shapes list.
Parameters
----------
idx : int | tuple[int, int] | slice
Index of the shape to get, if tuple then first index is the layer index and second is
the shape index if the layer has multiple shapes, if slice, returns a Shapes object with
the sliced shapes
list_ : bool, Optional
If true, returns a list of shapes if the index contains multiple shapes, else returns a
single shape, default = True
Returns
-------
list[int] | list[list[int]] | Shapes
Shape at the provided index
"""
item: list[int] | list[list[int]] | list[list[int] | list[list[int]]] = super().__getitem__(
idx if isinstance(idx, (int, slice)) else idx[0],
)
if isinstance(idx, slice):
return Shapes(cast(list[list[int] | list[list[int]]], item))
if isinstance(idx, tuple) and isinstance(item[0], list):
return cast(list[int], item[idx[1]])
if isinstance(idx, tuple):
raise ValueError(f'Index {idx} does not contain multiple shapes, shape at index is '
f'{item}')
if isinstance(item[0], list) and not list_:
raise ValueError(f'Index {idx} does not contain a single shape, shape at index is '
f'{item}, if you want to return all shapes at that index, use '
f'Shapes.get(idx, list_=True) instead')
return cast(list[int], item)
[docs]
class TypedModuleList(nn.ModuleList, Generic[ModuleT]):
"""
A typed version of nn.ModuleList that allows for type checking of the modules in the list.
See nn.ModuleList for more information.
"""
# pylint: disable-next=abstract-method
def __init__(self, modules: Iterable[ModuleT] | None = None) -> None:
super().__init__(modules)
def __add__(self, other: Iterable[ModuleT]) -> TypedModuleList[ModuleT]: # type: ignore[override]
return cast(TypedModuleList[ModuleT], super().__add__(other))
@overload # type: ignore[override]
def __getitem__(self, idx: int) -> ModuleT: ...
@overload
def __getitem__(self, idx: slice) -> TypedModuleList[ModuleT]: ...
def __getitem__(self, index: int | slice) -> ModuleT | TypedModuleList[ModuleT]:
return cast(ModuleT | TypedModuleList[ModuleT], super().__getitem__(index))
def __iadd__(self, modules: Iterable[ModuleT]) -> Self: # type: ignore[override]
return super().__iadd__(modules)
def __iter__(self) -> Iterator[ModuleT]:
return cast(Iterator[ModuleT], super().__iter__())
def __setitem__(self, idx: int, module: ModuleT) -> None: # type: ignore[override]
super().__setitem__(idx, module)
[docs]
def append(self, module: ModuleT) -> Self: # type: ignore[override]
return super().append(module)
[docs]
def extend(self, modules: Iterable[ModuleT]) -> Self: # type: ignore[override]
return super().extend(modules)
[docs]
def insert(self, index: int, module: ModuleT) -> None: # type: ignore[override]
super().insert(index, module)
@overload # type: ignore[override]
def pop(self, key: int) -> ModuleT: ...
@overload
def pop(self, key: slice) -> TypedModuleList[ModuleT]: ...
def pop(self, key: int | slice) -> ModuleT | TypedModuleList[ModuleT]:
return cast(ModuleT | TypedModuleList[ModuleT], super().pop(key))
[docs]
def ascii_plot(
data: list[float],
*,
clear: bool = False,
x_num: int = 20,
height: int = 6,
text: str = '',
label: str = '',
data2: list[float] | None = None) -> None:
"""
Prints a simple ASCII graph of data per index.
Parameters
----------
data : list[float]
Data to plot
clear : bool, Optional
If true, clears the previous graph, WARNING: this is very buggy and may not work properly,
default = False
x_num : int, Optional
Maximum number of epochs to display in the graph, default = 20
height : int, Optional
Height of the graph in lines, default = 6
text : str, Optional
Optional text to display at the end of the graph
label : str, Optional
Label for the x-axis
data2 : list[float] | None, Optional
Secondary data to plot
"""
if len(data) < 2:
return
i: int
row: int
row2: int
spacing: int = int(np.ceil(np.log10(len(data)))) + 1
datum: float
datum2: float | None
line: str
symbol: str
symbol2: str
x_label: str = f'{label}:' if label else ''
grid: list[str]
symbols: list[str] = ['_', '-', '‾'] # low, mid, high
symbols2: list[str] = ['.', '·', '˙']
rows: ndarray
red_data2: ndarray
red_data: ndarray = np.array(data)[-x_num:]
if data2:
red_data2 = np.array(data2)[-x_num:]
range_ = (min(red_data.min(), red_data2.min()), max(red_data.max(), red_data2.max()))
red_data2 = (red_data2 - range_[0]) / (range_[1] - range_[0])
else:
red_data2 = np.array([None] * len(red_data))
range_ = (red_data.min(), red_data.max())
rows = np.linspace(0, 1, height + 1)
red_data = (red_data - range_[0]) / (range_[1] - range_[0])
grid = [''] * height
rows[-1] += 1
if clear:
sys.stdout.write(f'\033[{height + 2}F')
for i, (datum, datum2) in enumerate(zip(red_data, red_data2)):
row = int(np.digitize(datum, rows)) - 1
symbol = symbols[min(
int((datum - rows[row]) // (1 / len(symbols) / height)),
len(symbols) - 1,
)]
if datum2 is not None:
row2 = int(np.digitize(datum2, rows)) - 1
symbol2 = symbols2[min(
int((datum2 - rows[row2]) // (1 / len(symbols2) / height)),
len(symbols2) - 1,
)]
if row == row2:
symbol = '*'
else:
grid[row2] = f'{grid[row2]:<{i * spacing + len(x_label) + 1}}{symbol2 * spacing}'
grid[row] = f'{grid[row]:<{i * spacing + len(x_label) + 1}}{symbol * spacing}'
for line in grid[::-1]:
print(line)
print(
x_label,
''.join(f'{i + 1:^{spacing}}' for i in range(max(0, len(data) - x_num), len(data))),
f'\n{text}',
)
[docs]
def check_params(name: str, supported_params: list[str] | ndarray, in_params: ndarray) -> None:
"""
Checks if provided parameters are supported by the function.
Parameters
----------
name : str
Name of the function
supported_params : list[str] | ndarray
Parameters supported by the function
in_params : ndarray
Input parameters of shape (N), where N is the number of parameters
"""
bad_params: ndarray = in_params[~np.isin(in_params, supported_params)]
if len(bad_params):
log.getLogger(__name__).warning(f'Unknown parameters for {name}: {bad_params}')
[docs]
def compare_versions(cur_ver: str, min_ver: str) -> bool:
"""
Compares two version strings.
Parameters
----------
cur_ver : str
Current version
min_ver : str
Minimum version
Returns
-------
bool
True if current version is greater than or equal to minimum version
"""
if not bool(cur_ver):
return True
return not '<' in cur_ver and Version(cur_ver) >= Version(min_ver)
[docs]
def deep_merge(base: dict, new: dict) -> dict:
"""
Performs a deep merge of two dictionaries, equivalent to recursive base | new
Parameters
----------
base : dict
Base dictionary
new : dict
Dictionary to deep merge into base
Returns
-------
dict
Deep merged dictionary
"""
merged: dict = base.copy()
key: Any
value: Any
for key, value in new.items():
if key in merged and isinstance(merged[key], dict) and isinstance(value, dict):
merged[key] = deep_merge(merged[key], value)
else:
merged[key] = value
return merged
[docs]
def dict_list_append(
dict1: dict[str, Any],
dict2: dict[str, Any],
concat: bool = False,
balance: bool = False) -> dict[str, Any]:
"""
Merges two dictionaries.
If a key exists in both dictionaries and not concat, the data type must match.
If data type is list or ndarray, the values are extended.
If data type is any other type and not concat, the value from the second dictionary overrides
the first, else if concat, the values are converted to a list and appended.
If balance is True, all lists and ndarrays in each dictionary must have the same length, if a
key is missing from one dictionary, it is added with None values to match the length of the
other lists/ndarrays in that dictionary.
Parameters
----------
dict1 : dict[str, Any]
Primary dict to merge secondary dict into
dict2 : dict[str, Any]
Secondary dict to merge into primary dict
concat : bool, Optional
If true, concatenates non list/ndarray types instead of overriding them, default = False
balance : bool, Optional
If true, balances the lengths of lists/ndarrays in each dictionary so that each list/ndarray
in the returning dictionary has the same length, default = False
Returns
-------
dict[str, Any]
First dict with second dict merged into it
"""
dict1_len: int = 0
dict2_len: int = 0
key: str
dict2 = deepcopy(dict2)
if not concat:
for key in np.intersect1d(list(dict1.keys()), list(dict2.keys())).tolist():
if not isinstance(dict1[key], type(dict2[key])):
raise TypeError(f'Type mismatch for key "{key}": {type(dict1[key])} != '
f'{type(dict2[key])}')
if balance:
# If primary dict is not empty, find the length of a list/ndarray in the dictionary
if len(dict1) > 0:
for value in dict1.values():
if isinstance(value, (list, ndarray)):
dict1_len = len(value)
break
if any(len(value) != dict1_len for value in dict1.values()
if isinstance(value, (list, ndarray))):
raise ValueError('All lists/ndarrays in dict1 must have the same length')
# If secondary dict is not empty, find the length of a list/ndarray in the dictionary
if len(dict2) > 0:
for value in dict2.values():
if isinstance(value, (list, ndarray)):
dict2_len = len(value)
break
if any(len(value) != dict2_len for value in dict2.values()
if isinstance(value, (list, ndarray))):
raise ValueError('All lists/ndarrays in dict2 must have the same length')
# Merge two dictionaries
for key in np.unique(list(dict1.keys()) + list(dict2.keys())).tolist():
# If the secondary dict has a key not in the primary, pad with Nones
if key in dict1 or not balance:
pass
elif isinstance(dict2[key], ndarray):
dict1[key] = np.full(
[dict1_len] + (list(dict2[key][0].shape) if dict2[key].ndim > 1 else []),
None,
)
elif isinstance(dict2[key], list) or concat:
dict1[key] = [None] * dict1_len
if concat and key in dict1 and not isinstance(dict1[key], (list, ndarray)):
dict1[key] = [dict1[key]]
if concat and key in dict2 and not isinstance(dict2[key], (list, ndarray)):
dict2[key] = [dict2[key]]
# If the primary dict has a key not in the secondary dict, pad with Nones, else merge dicts
if key not in dict2 and not balance:
continue
if key not in dict1:
dict1[key] = dict2[key]
elif balance and key not in dict2 and isinstance(dict1[key], ndarray):
dict1[key] = np.concat(
(dict1[key], [np.full( # type: ignore[arg-type]
dict1[key][0].shape,
None,
) if dict1[key].ndim > 1 else None] * dict2_len),
axis=0,
)
elif balance and key not in dict2 and isinstance(dict1[key], list):
dict1[key].extend([None] * dict2_len)
elif isinstance(dict1[key], ndarray):
dict1[key] = np.concatenate((dict1[key], dict2[key]), axis=0)
elif isinstance(dict1[key], list):
dict1[key].extend(dict2[key])
elif key in dict2:
dict1[key] = dict2[key]
return dict1
[docs]
def dict_list_convert(data: dict[str, list[Any] | ndarray]) -> list[dict[str, Any]]:
"""
Converts a dictionary of lists to a list of dictionaries.
Parameters
----------
data : dict[str, list[Any] | ndarray]
Dictionary of lists to convert
Returns
-------
list[dict[str, Any]]
List of dictionaries
"""
return [{key: data[key][i:i + 1] for key in data} for i in range(len(list(data.values())[0]))]
[docs]
def get_device() -> tuple[dict[str, Any], torch.device]:
"""
Gets the device for PyTorch to use.
Returns
-------
tuple[dict[str, Any], torch.device]
Arguments for the PyTorch DataLoader to use when loading data into memory and PyTorch device
"""
device: torch.device = torch.device(
'cuda' if torch.cuda.is_available() else
'mps' if torch.backends.mps.is_available() else
'cpu',
)
kwargs: dict[str, Any] = {
'num_workers': 4,
'pin_memory': True,
'persistent_workers': True,
} if device == torch.device('cuda') else {}
return kwargs, device
[docs]
def label_change(
data: ArrayT,
in_label: ArrayT,
*,
one_hot: bool = False,
out_label: ArrayT | None = None) -> ArrayT:
"""
Converts an array or tensor of class values to an array or tensor of class indices.
Parameters
----------
data : ArrayT
Classes of shape (N) where N is the number of samples
in_label : ArrayT
Unique class values of size (C) where C is the number of classes
one_hot : bool, Optional
If the returned tensor should be 1D array of class indices or 2D one hot tensor if out_label
is None or is an int, default = False
out_label : ArrayT | None, Optional
Unique class values of size (C) to transform data into, if None, then values will be indexes
Returns
-------
ArrayT
ndarray or Tensor of class indices of shape (N) and type float, or if one_hot is True, one
hot tensor of shape (N,C) and type float
"""
data_one_hot: ArrayT
out_data: ArrayT
module: ModuleType
if isinstance(data, Tensor):
module = torch
elif isinstance(data, ndarray):
module = np
else:
raise TypeError(f'Data type {type(data)} not supported')
if out_label is None:
out_label = module.arange(len(in_label))
if isinstance(out_label, Tensor):
out_label = cast(ArrayT, out_label.to(data.device))
out_data = cast(ArrayT, out_label[module.searchsorted(in_label, data)])
if one_hot:
data_one_hot = module.zeros((len(data), len(in_label)))
data_one_hot[module.arange(len(data)), out_data] = 1
out_data = data_one_hot
return cast(ArrayT, out_data.to(data.device) if isinstance(out_data, Tensor) else out_data)
[docs]
def list_dict_convert(
data: list[dict[str, Any]],
balance: bool = False,
concat: bool = False) -> dict[str, list[Any]]:
"""
Converts a list of dictionaries to a dictionary of lists
Parameters
----------
data : list[dict[str, float | ndarray]]
List of dictionaries to convert
balance : bool, Optional
If true, balances the lengths of lists/ndarrays in each dictionary so that each list/ndarray
in the returning dictionary has the same length, default = False
concat : bool, Optional
If true, concatenates non list/ndarray types instead of overriding them, default = False
Returns
-------
dict[str, list[float | None] | list[ndarray]]
Dictionary of lists
"""
value: dict[str, float | ndarray]
new_data: dict[str, list[float] | list[ndarray]] = {}
for value in data:
dict_list_append(new_data, value, concat=concat, balance=balance)
return new_data
def progress_bar(i: int, total: int, *, text: str = '', **kwargs: Any) -> None:
"""
Terminal progress bar.
Parameters
----------
i : int
Current progress
total : int
Completion number
text : str, Optional
Optional text to place at the end of the progress bar
**kwargs
Optional keyword arguments to pass to print
"""
filled: int
length: int = 50
percent: float
bar_fill: str
i += 1
filled = int(i * length / total)
percent = i * 100 / total
bar_fill = '█' * filled + '-' * (length - filled)
print(f'\rProgress: |{bar_fill}| {int(percent)}%\t{text}\t', end='', **kwargs)
if i == total:
print()
[docs]
def safe_globals(
safe_package: str,
modules: ModuleType | list[ModuleType],
use_all: bool = False,
legacy_path_replace: tuple[str, str] | None = None) -> None:
"""
Adds all classes in the module from the specified package to the list of safe PyTorch classes
when loading saved architectures.
Parameters
----------
safe_package : str
Package that classes must belong to to be added to the safe globals
modules : ModuleType | list[ModuleType]
Module(s) to add all classes from to the list of safe PyTorch classes when loading saved
architectures
use_all : bool, Optional
If true, only adds classes declared in module.__all__, default = False
legacy_path_replace : tuple[str, str] | None, Optional
If not None, replaces the first string with the second string in the module path when
adding classes to the safe globals
"""
if isinstance(modules, ModuleType):
modules = [modules]
for module in modules:
if use_all and not hasattr(module, '__all__'):
raise AttributeError(f'Module {module.__name__} does not have an __all__ attribute')
torch.serialization.add_safe_globals([
member[1]
for member in inspect.getmembers(sys.modules[module.__name__], inspect.isclass)
if member[1].__module__.split('.', maxsplit=1)[0] == safe_package
and (not use_all or member[0] in module.__all__)
])
if legacy_path_replace:
torch.serialization.add_safe_globals([
(member[1], f'{member[1].__module__.replace(*legacy_path_replace)}.{member[0]}')
for member in inspect.getmembers(sys.modules[module.__name__], inspect.isclass)
if member[1].__module__.split('.', maxsplit=1)[0] == safe_package
and (not use_all or member[0] in module.__all__)
])
[docs]
def save_name(num: int | str, states_dir: str, name: str) -> str:
"""
Standardises the architecture save file naming.
Parameters
----------
num : int | str
File number or name
states_dir : str
Directory of network saves
name : str
Name of the network
Returns
-------
str
Path to the network save file
"""
return os.path.join(states_dir, f'{name}_{num}.pth')
[docs]
@contextmanager
def suppress_logger_warnings(logger_name: str) -> Generator[None, None, None]:
"""
Suppresses warnings from a specified logger.
Parameters
----------
logger_name : str
Name of the logger to suppress warnings from
"""
level: int
logger: log.Logger = log.getLogger(logger_name)
level = logger.level
logger.setLevel(log.ERROR)
try:
yield
finally:
logger.setLevel(level)
__all__ = [
'Shapes',
'TypedModuleList',
'ascii_plot',
'check_params',
'compare_versions',
'deep_merge',
'dict_list_append',
'dict_list_convert',
'get_device',
'label_change',
'list_dict_convert',
'progress_bar',
'safe_globals',
'save_name',
'suppress_logger_warnings',
]