Use ArrayLike for numpy >= 1.20

This commit is contained in:
Martin Diehl 2021-12-18 17:21:46 +00:00 committed by Philip Eisenlohr
parent 430b947c7e
commit 0468bfd3e1
1 changed files with 16 additions and 12 deletions

View File

@ -6,6 +6,10 @@ from pathlib import Path
from typing import Sequence, Union, TextIO
import numpy as np
try:
from numpy.typing import ArrayLike
except ImportError:
ArrayLike = Union[np.ndarray,Sequence[float]] # type: ignore
import scipy.interpolate as interp
import matplotlib as mpl
if os.name == 'posix' and 'DISPLAY' not in os.environ:
@ -78,8 +82,8 @@ class Colormap(mpl.colors.ListedColormap):
@staticmethod
def from_range(low: Sequence[float],
high: Sequence[float],
def from_range(low: ArrayLike,
high: ArrayLike,
name: str = 'DAMASK colormap',
N: int = 256,
model: str = 'rgb') -> 'Colormap':
@ -129,7 +133,7 @@ class Colormap(mpl.colors.ListedColormap):
if model.lower() not in toMsh:
raise ValueError(f'Invalid color model: {model}.')
low_high = np.vstack((low,high))
low_high = np.vstack((low,high)).astype(float)
out_of_bounds = np.bool_(False)
if model.lower() == 'rgb':
@ -142,7 +146,7 @@ class Colormap(mpl.colors.ListedColormap):
out_of_bounds = np.any(low_high[:,0]<0)
if out_of_bounds:
raise ValueError(f'{model.upper()} colors {low} | {high} are out of bounds.')
raise ValueError(f'{model.upper()} colors {low_high[0]} | {low_high[1]} are out of bounds.')
low_,high_ = map(toMsh[model.lower()],low_high)
msh = map(functools.partial(Colormap._interpolate_msh,low=low_,high=high_),np.linspace(0,1,N))
@ -225,7 +229,7 @@ class Colormap(mpl.colors.ListedColormap):
def shade(self,
field: np.ndarray,
bounds: Sequence[float] = None,
bounds: ArrayLike = None,
gap: float = None) -> Image:
"""
Generate PIL image of 2D field using colormap.
@ -235,7 +239,7 @@ class Colormap(mpl.colors.ListedColormap):
field : numpy.array, shape (:,:)
Data to be shaded.
bounds : sequence of float, len (2), optional
Value range (low,high) spanned by colormap.
Value range (left,right) spanned by colormap.
gap : field.dtype, optional
Transparent value. NaN will always be rendered transparent.
@ -248,17 +252,17 @@ class Colormap(mpl.colors.ListedColormap):
mask = np.logical_not(np.isnan(field) if gap is None else \
np.logical_or (np.isnan(field), field == gap)) # mask NaN (and gap if present)
lo,hi = (field[mask].min(),field[mask].max()) if bounds is None else \
(min(bounds[:2]),max(bounds[:2]))
l,r = (field[mask].min(),field[mask].max()) if bounds is None else \
np.array(bounds,float)[:2]
delta,avg = hi-lo,0.5*(hi+lo)
delta,avg = r-l,0.5*abs(r+l)
if delta * 1e8 <= avg: # delta is similar to numerical noise
hi,lo = hi+0.5*avg,lo-0.5*avg # extend range to have actual data centered within
if abs(delta) * 1e8 <= avg: # delta is similar to numerical noise
l,r = l-0.5*avg*np.sign(delta),r+0.5*avg*np.sign(delta), # extend range to have actual data centered within
return Image.fromarray(
(np.dstack((
self.colors[(np.round(np.clip((field-lo)/(hi-lo),0.0,1.0)*(self.N-1))).astype(np.uint16),:3],
self.colors[(np.round(np.clip((field-l)/delta,0.0,1.0)*(self.N-1))).astype(np.uint16),:3],
mask.astype(float)
)
)*255