"""Utilities module."""
import base64
import itertools
import logging
import os
import ssl
import zlib
from typing import Any, Optional, Tuple, Union
from typing import Dict as TypingDict
from typing import List as TypingList
from urllib.error import URLError
from urllib.request import Request, urlopen
import msgpack
import numpy as np
from traitlets import Float as _TraitFloat
from traitlets import Int as _TraitInt
from traitlets import TraitError
from traittypes import Array as _TraitArray
from ._protocol import get_protocol
# Set up module-level logger
logger = logging.getLogger(__name__)
if not logger.hasHandlers():
handler = logging.StreamHandler()
formatter = logging.Formatter(
"[%(asctime)s] %(levelname)s in %(module)s: %(message)s"
)
handler.setFormatter(formatter)
logger.addHandler(handler)
logger.setLevel(logging.INFO)
def check_unsigned_range(value: np.ndarray, target: np.dtype) -> None:
"""Raise TraitError if casting `value` into the unsigned dtype `target` would wrap.
The cast happens before any .valid() validator runs: -1 reaches a validator promising to
reject it as 65535, and 256 into a uint8 array becomes 0, an empty voxel.
"""
target = np.dtype(target)
if target.kind != "u" or value.dtype.kind not in "iuf" or not value.size:
return
if value.dtype.kind in "if" and (value < 0).any():
raise TraitError("Negative values cannot be stored in a %s array" % target.name)
# an unsigned array no wider than the target cannot overflow, so skip the scan
if value.dtype.kind == "u" and value.dtype.itemsize <= target.itemsize:
return
limit = 2 ** (8 * target.itemsize)
if (value >= limit).any():
raise TraitError(
"Values above %d cannot be stored in a %s array" % (limit - 1, target.name)
)
class Array(_TraitArray):
"""Array trait that converts silently in two cases: float64 narrowed to the float32 the
GPU takes, and a dtype differing only in byte order (legacy VTK files are big-endian),
which traittypes reports as 'float32 does not match float32' because it names dtypes
without their order. Any other mismatch still warns.
"""
def validate(self, obj, value):
if self.dtype is not None and isinstance(value, np.ndarray):
target = np.dtype(self.dtype)
check_unsigned_range(value, target)
if target == np.float32 and value.dtype == np.float64:
value = value.astype(np.float32)
elif value.dtype != target and value.dtype.name == target.name:
value = value.astype(target)
return super().validate(obj, value)
class Float(_TraitFloat):
"""Float trait that also takes numpy scalars: np.float64 subclasses Python float and
passes traitlets, np.float32 does not."""
def validate(self, obj, value):
if isinstance(value, (np.floating, np.integer)):
value = float(value)
return super().validate(obj, value)
class Int(_TraitInt):
"""An Int trait that also takes numpy integers, which are not Python ints."""
def validate(self, obj, value):
if isinstance(value, np.integer):
value = int(value)
return super().validate(obj, value)
# pylint: disable=unused-argument
# noinspection PyUnusedLocal
[docs]
def array_to_json(
ar: np.ndarray, compression_level: int = 0, force_contiguous: bool = True
) -> Union[str, TypingDict[str, Any]]:
"""
Return the serialization of a numpy array.
Parameters
----------
ar : ndarray
A numpy array.
compression_level : int, optional
Level of compression [-1, 9], by default 0.
force_contiguous : bool, optional
Make the array contiguous in memory, by default True.
Returns
-------
dict or str
Binary data of the array with its dtype and shape, or a base64 string if protocol is 'text'.
Raises
------
ValueError
If the dtype is unsupported.
"""
if ar.dtype.kind not in ["u", "i", "f"]: # ints and floats
logger.error(f"Unsupported dtype: {ar.dtype}")
raise ValueError(f"Unsupported dtype: {ar.dtype}")
if ar.dtype == np.float64: # WebGL does not support float64
logger.debug("Converting float64 array to float32 for WebGL compatibility.")
ar = ar.astype(np.float32)
elif ar.dtype == np.int64: # JS does not support int64
logger.debug("Converting int64 array to int32 for JS compatibility.")
ar = ar.astype(np.int32)
elif ar.dtype == np.uint64: # the JS deserializer has no uint64 typed array
logger.debug("Converting uint64 array to uint32 for JS compatibility.")
ar = ar.astype(np.uint32)
# JS typed arrays read a buffer little-endian and the deserializer keys off str(dtype),
# which for a big-endian array is '>f4' - a key typesToArray does not have
if ar.dtype.byteorder == ">":
ar = ar.astype(ar.dtype.newbyteorder("<"))
# make sure it's contiguous
if force_contiguous and not ar.flags["C_CONTIGUOUS"]:
ar = np.ascontiguousarray(ar)
# ravel, not flatten: contiguity is already ensured above, so this is a view, while flatten
# copies unconditionally. The uncompressed buffer therefore aliases the caller's array.
if compression_level > 0:
ret = {
"compressed_data": zlib.compress(ar.ravel(), compression_level),
"dtype": str(ar.dtype),
"shape": ar.shape,
}
else:
ret = {
"data": memoryview(ar.ravel()),
"dtype": str(ar.dtype),
"shape": ar.shape,
}
if get_protocol() == "text":
return "base64_" + base64.b64encode(
msgpack.packb(ret, use_bin_type=True)
).decode("ascii")
return ret
# noinspection PyUnusedLocal
[docs]
def json_to_array(
value: Optional[TypingDict[str, Any]], obj: Optional[Any] = None
) -> Optional[np.ndarray]:
"""
Return numpy array from serialization.
Parameters
----------
value : dict
Binary data of an array with its dtype and shape.
obj : dict, optional
Object, by default None.
Returns
-------
ndarray or None
Numpy array or None.
"""
if value:
data = (value["data"] if "data" in value
else bytearray(zlib.decompress(value["compressed_data"])))
ar = np.frombuffer(data, dtype=value["dtype"]).reshape(value["shape"])
if not ar.flags["WRITEABLE"]:
ar = ar.copy()
return ar
return None
[docs]
def to_json(
name: str, input: Any, obj: Optional[Any] = None, compression_level: int = 0
) -> Any:
"""
Return JSON object serialization.
Parameters
----------
name : str
Name of the property.
input : Any
Input data to serialize.
obj : Any, optional
Object containing property, by default None.
compression_level : int, optional
Compression level, by default 0.
Returns
-------
Any
Serialized JSON object.
"""
if hasattr(obj, "compression_level"):
compression_level = obj.compression_level
# click_callback/hover_callback: the browser only needs to know a handler is attached, and
# the function itself has no serialization. Matches what callback_serialization_wrap sends
# on the widget path, which this one has to agree with.
if callable(input):
return True
if isinstance(input, dict):
property = obj[name]
ret = {}
for key, value in input.items():
ret[str(key)] = to_json(key, value, property, compression_level)
return ret
if isinstance(input, np.ndarray) and input.dtype is np.dtype(object):
return to_json(name, input.tolist(), obj, compression_level)
if isinstance(input, list):
property = obj[name]
return [
to_json(idx, v, property, compression_level) for idx, v in enumerate(input)
]
if isinstance(input, bytes):
return array_to_json(np.frombuffer(input, dtype=np.uint8), compression_level)
if isinstance(input, np.ndarray):
return array_to_json(input, compression_level)
if isinstance(input, np.number):
return input.tolist()
return input
[docs]
def from_json(input: Any, obj: Optional[Any] = None) -> Any:
"""
Return JSON object deserialization.
Parameters
----------
input : Any
Input data to deserialize.
obj : Any, optional
Object, by default None.
Returns
-------
Any
Deserialized object.
"""
if isinstance(input, str) and input[0:7] == "base64_":
input = msgpack.unpackb(base64.b64decode(input[7:]))
if (
isinstance(input, dict)
and "dtype" in input
and ("data" in input or "compressed_data" in input)
and "shape" in input
):
return json_to_array(input, obj)
if isinstance(input, list):
return [from_json(i, obj) for i in input]
if isinstance(input, dict):
ret = {}
for key, value in input.items():
ret[key] = from_json(value, obj)
return ret
return input
def environment_to_json(value: Any, obj: Optional[Any] = None) -> Any:
"""Preset names travel as strings, user maps as a typed array."""
if value is None or isinstance(value, str):
return value
data = array_to_json(np.ascontiguousarray(np.asarray(value, dtype=np.float32)))
name = getattr(obj, "_environment_catalog_name", None)
if name is not None:
data["name"] = name
return data
def environment_from_json(value: Any, obj: Optional[Any] = None) -> Any:
if value is None or isinstance(value, str):
return value
return json_to_array(value)
[docs]
def array_serialization_wrap(name: str) -> TypingDict[str, Any]:
"""
Return a wrap of the serialization and deserialization functions for array objects.
Parameters
----------
name : str
Name of the property.
Returns
-------
dict
Dictionary with 'to_json' and 'from_json' functions.
"""
return {
"to_json": (lambda input, obj: to_json(name, input, obj)),
"from_json": from_json,
}
[docs]
def callback_serialization_wrap(name: str) -> TypingDict[str, Any]:
"""
Return a wrap of the serialization and deserialization functions for mouse actions.
Parameters
----------
name : str
Name of the property.
Returns
-------
dict
Dictionary with 'to_json' and 'from_json' functions.
"""
return {
"to_json": (lambda input, obj: obj[name] is not None),
"from_json": from_json,
}
# leading bytes of the image formats a browser decodes, by the name image/<name> takes
IMAGE_SIGNATURES = (
(b"\x89PNG\r\n\x1a\n", "png"),
(b"\xff\xd8\xff", "jpeg"),
(b"GIF87a", "gif"),
(b"GIF89a", "gif"),
(b"BM", "bmp"),
)
def image_format(data: Optional[bytes]) -> Optional[str]:
"""The image/ MIME subtype of encoded image bytes, None unless PNG, JPEG, GIF, WebP or BMP."""
if not data:
return None
head = bytes(data[:12])
if head[:4] == b"RIFF" and head[8:12] == b"WEBP":
return "webp"
for signature, name in IMAGE_SIGNATURES:
if head.startswith(signature):
return name
return None
def pack_colors(colors: Any) -> Tuple[Any, Optional[np.ndarray]]:
"""Packed 0xRRGGBB and alpha (or None) of (N, 3|4) colours - floats 0..1 or ints 0..255."""
if isinstance(colors, dict):
return colors, None
array = np.asarray(colors)
if array.ndim != 2 or array.shape[1] not in (3, 4):
return colors, None
if np.issubdtype(array.dtype, np.floating):
channels = np.clip(np.rint(array * 255.0), 0, 255).astype(np.uint32)
alpha = array[:, 3].astype(np.float32) if array.shape[1] == 4 else None
else:
channels = np.clip(array, 0, 255).astype(np.uint32)
alpha = (array[:, 3] / 255.0).astype(np.float32) if array.shape[1] == 4 else None
packed = (channels[:, 0] << 16) | (channels[:, 1] << 8) | channels[:, 2]
return packed.astype(np.uint32), alpha
[docs]
def download(url: str) -> str:
"""
Retrieve the file at url, save it locally and return its name.
Parameters
----------
url : str
URL.
Returns
-------
str
File path.
"""
basename = os.path.basename(url)
if os.path.exists(basename):
logger.info(f"File already exists locally: {basename}")
return basename
try:
# some hosts answer the default "Python-urllib" User-Agent with 403/406
request = Request(url, headers={"User-Agent": "K3D-jupyter", "Accept": "*/*"})
try:
with urlopen(request) as response, open(basename, "wb") as output:
output.write(response.read())
except URLError as error:
# an unverified retry is acceptable here: public data, never credentials
if not isinstance(getattr(error, "reason", None), ssl.SSLCertVerificationError):
raise
logger.warning(
f"Certificate verification failed for {url} - retrying without it"
)
unverified = ssl.create_default_context()
unverified.check_hostname = False
unverified.verify_mode = ssl.CERT_NONE
with urlopen(request, context=unverified) as response, open(basename, "wb") as output:
output.write(response.read())
logger.info(f"Downloaded file from {url} to {basename}")
except Exception as e:
logger.error(f"Failed to download {url}: {e}")
raise
return basename
def rgb_volume_channels(value: Any) -> int:
"""Number of colour channels a volume carries per voxel, 0 if it is a scalar field.
A trailing axis of 3 or 4 on an otherwise 3D array means the colour was measured - a
photographic cryosection, an RGB-encoded NIfTI - so there is nothing for a colormap to
map and the data goes to the GPU as colour.
"""
shape = getattr(value, "shape", ())
if len(shape) == 4 and shape[3] in (3, 4):
return int(shape[3])
return 0
[docs]
def minmax(arr: np.ndarray) -> TypingList[float]:
"""Return the minimum and maximum value of an array.
Parameters
----------
arr : array_like
Array of numbers.
Returns
-------
list
Array of two numbers.
"""
return [float(np.nanmin(arr)), float(np.nanmax(arr))]
[docs]
def check_attribute_color_range(
attribute: Union[np.ndarray, TypingDict[str, np.ndarray]],
color_range: Union[TypingList[float], Tuple[float, ...]] = (),
channels: bool = False,
) -> TypingList[float]:
"""Return color range versus provided attribute.
Parameters
----------
attribute : list or dict (for timeseries)
Array of numbers.
color_range : tuple, optional
Two numbers, by default ().
channels : bool, optional
Whether a list means one attribute per channel, by default False. Only volume_slice
takes several channels; everywhere else a list is one attribute written out by hand,
and reading it as channels returns a range per slice that no trait accepts.
Returns
-------
tuple
Color range.
"""
if color_range is None:
color_range = []
# The range is per channel: the browser reads color_range[2 * i] and [2 * i + 1] for
# channel i, so a two-element range over two channels would leave the second one reading
# undefined and its uniforms NaN.
if channels and isinstance(attribute, (list, tuple)):
if len(color_range) == 2 * len(attribute):
color_range = [float(v) for v in color_range]
for i in range(0, len(color_range), 2):
if color_range[i] == color_range[i + 1]:
color_range[i + 1] += 1.0
return color_range
# one range asked for over several channels means the same range on each of them,
# which is what a caller writing color_range=[0, 1] is asking for
if len(color_range) == 2:
low, high = float(color_range[0]), float(color_range[1])
if low == high:
high += 1.0
return [low, high] * len(attribute)
ranges = []
for channel in attribute:
low, high = minmax(channel)
if low == high:
high += 1.0
ranges += [low, high]
return ranges
if len(color_range) == 2:
low, high = float(color_range[0]), float(color_range[1])
if low == high:
high += 1.0
return [low, high]
# a nested list is one attribute written out by hand; the paths below index .size
if isinstance(attribute, (list, tuple)):
attribute = np.asarray(attribute)
if type(attribute) is dict:
t = [minmax(attribute[k]) for k in attribute]
color_range = [min([v[0] for v in t]), max([v[1] for v in t])]
elif attribute.size == 0:
return color_range
else:
color_range = minmax(attribute)
if color_range[0] == color_range[1]:
color_range[1] += 1.0
return color_range
[docs]
def map_colors(
attribute: np.ndarray,
color_map: Union[TypingList[TypingList[float]], np.ndarray],
color_range: Union[TypingList[float], Tuple[float, ...]] = (),
) -> np.ndarray:
"""Return color mapping according to an attribute and a colormap.
The attribute represents the data on which the colormap will be applied.
The color range allows to constraint the colormap between two values.
Parameters
----------
attribute : ndarray
Array of numbers.
color_map : array_like
Array of numbers.
color_range : tuple, optional
Two numbers, by default ().
Returns
-------
ndarray
Color mapping.
"""
a_min, a_max = check_attribute_color_range(attribute, color_range)
map_array = np.asarray(color_map)
map_array = map_array.reshape((map_array.size // 4, 4))
# normalizing attribute for range lookup
span = a_max - a_min
if span == 0:
attribute = np.full(np.shape(attribute), 0.5, dtype=np.float64)
else:
attribute = (np.asarray(attribute, dtype=np.float64) - a_min) / span
red, green, blue = [
np.array(
255 * np.interp(attribute, xp=map_array[:, 0], fp=map_array[:, i + 1]),
dtype=np.int32,
)
for i in range(3)
]
return (red << 16) + (green << 8) + blue
[docs]
def bounding_corners(
bounds: Union[TypingList[float], np.ndarray], z_bounds: Tuple[float, float] = (0, 1)
) -> np.ndarray:
"""Return corner point coordinates for bounds array.
`z_bounds` assigns Z points coordinates if bounds contains less than 5 items.
Parameters
----------
bounds : array_like
Array of numbers.
z_bounds : tuple, optional
Two numbers, by default (0, 1).
Returns
-------
ndarray
Corner points coordinates.
"""
# len(), not `or`: a numpy slice has no truth value and raises here
z = bounds[4:] if len(bounds[4:]) > 0 else z_bounds
return np.array(list(itertools.product(bounds[:2], bounds[2:4], z)))
[docs]
def min_bounding_dimension(bounds: Union[TypingList[float], np.ndarray]) -> float:
"""Return the minimal dimension along axis in a bounds array.
`bounds` must be of the form [min_x, max_x, min_y, max_y, min_z, max_z].
Parameters
----------
bounds : array_like
Array of numbers.
Returns
-------
number
Minimum value of the array.
"""
# the pairs are (min_x, max_x), (min_y, max_y), (min_z, max_z); zipping neighbours instead
# measures max_x against min_y and reports 0 for any two axes that meet
return min(abs(x1 - x0) for x0, x1 in zip(bounds[::2], bounds[1::2]))
[docs]
def shape_validation(*dimensions):
"""Create a validator callback ensuring array shape.
Returns
-------
function
Shape validator function.
Raises
------
TraitError
Expected an array of shape _ and got _.
"""
def validator(trait, value):
if np.shape(value) != dimensions:
raise TraitError(
"Expected an array of shape %s and got %s" % (dimensions, value.shape)
)
return value
return validator
[docs]
def sparse_voxels_validation():
"""Check sparse voxels for array shape and values.
Returns
-------
function
Sparse voxels validator function.
Raises
------
TraitError
Expected an array of shape (N, 4) and got _.
TraitError
Voxel coordinates and values must be non-negative.
"""
def validator(trait, value):
if len(value.shape) != 2 or value.shape[1] != 4:
raise TraitError(
"Expected an array of shape (N, 4) and got %s" % (value.shape,)
)
if (np.asarray(value) < 0).any():
raise TraitError("Voxel coordinates and values must be non-negative")
return value
return validator
[docs]
def quad(w: float, h: float) -> Tuple[np.ndarray, np.ndarray]:
"""Return the vertices and indices of a `w` * `h` quadrilateral.
Parameters
----------
w : number
Quadrilateral width.
h : number
Quadrilateral height.
Returns
-------
tuple
Array of vertices and indices.
"""
w /= 2
h /= 2
vertices = np.array([-w, -h, -0, w, -h, 0, w, h, 0, -w, h, 0], dtype=np.float32)
indices = np.array([0, 1, 2, 0, 2, 3], dtype=np.uint32)
return vertices, indices
[docs]
def get_bounding_box(model_matrix, boundary=[-0.5, 0.5, -0.5, 0.5, -0.5, 0.5]):
"""Return the boundaries of a model matrix.
Parameters
----------
model_matrix : ndarray
Matrix of numbers. Must have four columns.
boundary : list, optional
Array of numbers, by default [-0.5, 0.5, -0.5, 0.5, -0.5, 0.5].
Must be of the form [min_x, max_x, min_y, max_y, min_z, max_z].
Returns
-------
ndarray
Model matrix boundaries.
"""
# Homogeneous coordinate 1: these are points, not directions, so the matrix's
# translation column has to apply. All eight corners, because a rotation or a mirror
# maps the min corner past the max one and a box with min > max drops out of get_auto_grid.
corners = np.array([
[boundary[i], boundary[2 + j], boundary[4 + k], 1.0]
for i in (0, 1) for j in (0, 1) for k in (0, 1)
])
transformed = np.concatenate([
corners.dot(matrix.T)[:, 0:3] for matrix in _model_matrices(model_matrix)
])
return np.dstack([
np.nanmin(transformed, axis=0), np.nanmax(transformed, axis=0)
]).flatten()
def _model_matrices(model_matrix):
"""Every 4x4 matrix in `model_matrix`, whether it is one matrix or a time series of them."""
frames = model_matrix.values() if isinstance(model_matrix, dict) else [model_matrix]
return [np.asarray(matrix, dtype=np.float64).reshape(4, 4) for matrix in frames]
[docs]
def get_bounding_box_points(arr, model_matrix):
"""Return the minimum and maximum coordinates on x, y, z axes.
Parameters
----------
arr : ndarray
Array of vertices [x, y, z].
model_matrix : ndarray
Matrix of numbers. Must have four columns.
Returns
-------
ndarray
Array of numbers [min_x, max_x, min_y, max_y, min_z, max_z].
"""
d = _flatten_frames(arr)
if d.shape[0] < 3:
d = np.array([0, 0, 0])
# fmt: off
boundary = np.array([
np.nanmin(d[0::3]), np.nanmax(d[0::3]),
np.nanmin(d[1::3]), np.nanmax(d[1::3]),
np.nanmin(d[2::3]), np.nanmax(d[2::3]),
])
# fmt: on
return get_bounding_box(model_matrix, boundary)
def _flatten_frames(arr):
"""One flat array of coordinates, whether `arr` is an array or a time series of them."""
if isinstance(arr, dict):
frames = [np.asarray(v, dtype=np.float64).flatten() for v in arr.values()]
return np.concatenate(frames) if frames else np.array([])
return np.asarray(arr, dtype=np.float64).flatten()
[docs]
def get_bounding_box_point(position, model_matrix=None):
"""Return the boundaries of one or more 3D positions.
Parameters
----------
position : array_like
One position, n positions, or a time series of either.
model_matrix : ndarray, optional
Matrix of numbers the renderer applies to the position. Must have four columns.
Returns
-------
ndarray or None
Array of numbers [min_x, max_x, min_y, max_y, min_z, max_z], or None for a position
that is not three-dimensional (a 2D overlay has no place in the scene's box).
"""
d = _flatten_frames(position)
if d.shape[0] == 0 or d.shape[0] % 3 != 0:
return None
points = d.reshape(-1, 3)
boundary = np.dstack([np.nanmin(points, axis=0), np.nanmax(points, axis=0)]).flatten()
if model_matrix is None:
return boundary
return get_bounding_box(model_matrix, boundary)
def unify_color_map(cm):
cm[0::4] = (cm[0::4] - np.min(cm[0::4])) / (np.max(cm[0::4]) - np.min(cm[0::4]))
return cm
def contour(data, bounds, values, clustering_factor=0):
import pyvista
grid = pyvista.ImageData(
dimensions=data.shape[::-1],
spacing=(bounds[1::2] - bounds[::2]) / np.array(data.shape[::-1]),
origin=bounds[::2],
)
grid.point_data["values"] = data.flatten()
mesh = grid.contour(values, grid.point_data["values"], method="flying_edges")
mesh.compute_normals(inplace=True)
if clustering_factor > 1:
import pyacvd
clus = pyacvd.Clustering(mesh)
clus.cluster(mesh.n_points // clustering_factor)
return clus.create_mesh()
return mesh