diff --git a/python/damask/_colormap.py b/python/damask/_colormap.py index 8c897e6ab..40aab652c 100644 --- a/python/damask/_colormap.py +++ b/python/damask/_colormap.py @@ -2,6 +2,9 @@ import os import json import functools import colorsys +from pathlib import Path +from typing import Sequence, Union, TextIO + import numpy as np import matplotlib as mpl @@ -39,20 +42,20 @@ class Colormap(mpl.colors.ListedColormap): """ - def __add__(self,other): + def __add__(self, other: "Colormap") -> "Colormap": """Concatenate.""" return Colormap(np.vstack((self.colors,other.colors)), f'{self.name}+{other.name}') - def __iadd__(self,other): + def __iadd__(self, other: "Colormap") -> "Colormap": """Concatenate (in-place).""" return self.__add__(other) - def __invert__(self): + def __invert__(self) -> "Colormap": """Reverse.""" return self.reversed() - def __repr__(self): + def __repr__(self) -> str: """Show as matplotlib figure.""" fig = plt.figure(self.name,figsize=(5,.5)) ax1 = fig.add_axes([0, 0, 1, 1]) @@ -64,7 +67,11 @@ class Colormap(mpl.colors.ListedColormap): @staticmethod - def from_range(low,high,name='DAMASK colormap',N=256,model='rgb'): + def from_range(low: Sequence[float], + high: Sequence[float], + name: str = 'DAMASK colormap', + N: int = 256, + model: str = 'rgb') -> "Colormap": """ Create a perceptually uniform colormap between given (inclusive) bounds. @@ -145,7 +152,7 @@ class Colormap(mpl.colors.ListedColormap): @staticmethod - def from_predefined(name,N=256): + def from_predefined(name: str, N: int = 256) -> "Colormap": """ Select from a set of predefined colormaps. @@ -185,7 +192,10 @@ class Colormap(mpl.colors.ListedColormap): return Colormap.from_range(definition['low'],definition['high'],name,N) - def shade(self,field,bounds=None,gap=None): + def shade(self, + field: np.ndarray, + bounds: Sequence[float] = None, + gap: float = None) -> Image: """ Generate PIL image of 2D field using colormap. @@ -226,7 +236,7 @@ class Colormap(mpl.colors.ListedColormap): mode='RGBA') - def reversed(self,name=None): + def reversed(self, name: str = None) -> "Colormap": """ Reverse. @@ -251,7 +261,7 @@ class Colormap(mpl.colors.ListedColormap): return Colormap(np.array(rev.colors),rev.name[:-4] if rev.name.endswith('_r_r') else rev.name) - def _get_file_handle(self,fname,extension): + def _get_file_handle(self, fname: Union[TextIO, str, Path, None], suffix: str) -> TextIO: """ Provide file handle. @@ -259,8 +269,7 @@ class Colormap(mpl.colors.ListedColormap): ---------- fname : file, str, pathlib.Path, or None Filename or filehandle, will be name of the colormap+extension if None. - - extension: str + suffix: str Extension of the filename. Returns @@ -270,17 +279,14 @@ class Colormap(mpl.colors.ListedColormap): """ if fname is None: - fhandle = open(self.name.replace(' ','_')+'.'+extension,'w',newline='\n') + return open(self.name.replace(' ','_')+suffix, 'w', newline='\n') + elif isinstance(fname, (str, Path)): + return open(fname, 'w', newline='\n') else: - try: - fhandle = open(fname,'w',newline='\n') - except TypeError: - fhandle = fname - - return fhandle + return fname - def save_paraview(self,fname=None): + def save_paraview(self, fname: Union[TextIO, str, Path] = None): """ Save as JSON file for use in Paraview. @@ -303,10 +309,10 @@ class Colormap(mpl.colors.ListedColormap): 'RGBPoints':colors }] - json.dump(out,self._get_file_handle(fname,'json'),indent=4) + json.dump(out,self._get_file_handle(fname,'.json'),indent=4) - def save_ASCII(self,fname=None): + def save_ASCII(self, fname: Union[TextIO, str, Path] = None): """ Save as ASCII file. @@ -319,10 +325,10 @@ class Colormap(mpl.colors.ListedColormap): """ labels = {'RGBA':4} if self.colors.shape[1] == 4 else {'RGB': 3} t = Table(self.colors,labels,f'Creator: {util.execution_stamp("Colormap")}') - t.save(self._get_file_handle(fname,'txt')) + t.save(self._get_file_handle(fname,'.txt')) - def save_GOM(self,fname=None): + def save_GOM(self, fname: Union[TextIO, str, Path] = None): """ Save as ASCII file for use in GOM Aramis. @@ -340,10 +346,10 @@ class Colormap(mpl.colors.ListedColormap): + ' '.join([f' 0 {c[0]} {c[1]} {c[2]} 255 1' for c in reversed((self.colors*255).astype(int))]) \ + '\n' - self._get_file_handle(fname,'legend').write(GOM_str) + self._get_file_handle(fname,'.legend').write(GOM_str) - def save_gmsh(self,fname=None): + def save_gmsh(self, fname: Union[TextIO, str, Path] = None): """ Save as ASCII file for use in gmsh. @@ -358,11 +364,13 @@ class Colormap(mpl.colors.ListedColormap): gmsh_str = 'View.ColorTable = {\n' \ +'\n'.join([f'{c[0]},{c[1]},{c[2]},' for c in self.colors[:,:3]*255]) \ +'\n}\n' - self._get_file_handle(fname,'msh').write(gmsh_str) + self._get_file_handle(fname,'.msh').write(gmsh_str) @staticmethod - def _interpolate_msh(frac,low,high): + def _interpolate_msh(frac, + low: np.ndarray, + high: np.ndarray) -> np.ndarray: """ Interpolate in Msh color space. @@ -439,31 +447,31 @@ class Colormap(mpl.colors.ListedColormap): @staticmethod - def _hsv2rgb(hsv): + def _hsv2rgb(hsv: np.ndarray) -> np.ndarray: """H(ue) S(aturation) V(alue) to R(red) G(reen) B(lue).""" return np.array(colorsys.hsv_to_rgb(hsv[0]/360.,hsv[1],hsv[2])) @staticmethod - def _rgb2hsv(rgb): + def _rgb2hsv(rgb: np.ndarray) -> np.ndarray: """R(ed) G(reen) B(lue) to H(ue) S(aturation) V(alue).""" h,s,v = colorsys.rgb_to_hsv(rgb[0],rgb[1],rgb[2]) return np.array([h*360,s,v]) @staticmethod - def _hsl2rgb(hsl): + def _hsl2rgb(hsl: np.ndarray) -> np.ndarray: """H(ue) S(aturation) L(uminance) to R(red) G(reen) B(lue).""" return np.array(colorsys.hls_to_rgb(hsl[0]/360.,hsl[2],hsl[1])) @staticmethod - def _rgb2hsl(rgb): + def _rgb2hsl(rgb: np.ndarray) -> np.ndarray: """R(ed) G(reen) B(lue) to H(ue) S(aturation) L(uminance).""" h,l,s = colorsys.rgb_to_hls(rgb[0],rgb[1],rgb[2]) return np.array([h*360,s,l]) @staticmethod - def _xyz2rgb(xyz): + def _xyz2rgb(xyz: np.ndarray) -> np.ndarray: """ CIE Xyz to R(ed) G(reen) B(lue). @@ -483,7 +491,7 @@ class Colormap(mpl.colors.ListedColormap): return np.clip(rgb,0.,1.) @staticmethod - def _rgb2xyz(rgb): + def _rgb2xyz(rgb: np.ndarray) -> np.ndarray: """ R(ed) G(reen) B(lue) to CIE Xyz. @@ -501,7 +509,7 @@ class Colormap(mpl.colors.ListedColormap): @staticmethod - def _lab2xyz(lab,ref_white=None): + def _lab2xyz(lab: np.ndarray, ref_white: np.ndarray = None) -> np.ndarray: """ CIE Lab to CIE Xyz. @@ -520,7 +528,7 @@ class Colormap(mpl.colors.ListedColormap): ])*(ref_white if ref_white is not None else _ref_white) @staticmethod - def _xyz2lab(xyz,ref_white=None): + def _xyz2lab(xyz: np.ndarray, ref_white: np.ndarray = None) -> np.ndarray: """ CIE Xyz to CIE Lab. @@ -540,7 +548,7 @@ class Colormap(mpl.colors.ListedColormap): @staticmethod - def _lab2msh(lab): + def _lab2msh(lab: np.ndarray) -> np.ndarray: """ CIE Lab to Msh. @@ -558,7 +566,7 @@ class Colormap(mpl.colors.ListedColormap): ]) @staticmethod - def _msh2lab(msh): + def _msh2lab(msh: np.ndarray) -> np.ndarray: """ Msh to CIE Lab. @@ -575,29 +583,29 @@ class Colormap(mpl.colors.ListedColormap): ]) @staticmethod - def _lab2rgb(lab): + def _lab2rgb(lab: np.ndarray) -> np.ndarray: return Colormap._xyz2rgb(Colormap._lab2xyz(lab)) @staticmethod - def _rgb2lab(rgb): + def _rgb2lab(rgb: np.ndarray) -> np.ndarray: return Colormap._xyz2lab(Colormap._rgb2xyz(rgb)) @staticmethod - def _msh2rgb(msh): + def _msh2rgb(msh: np.ndarray) -> np.ndarray: return Colormap._lab2rgb(Colormap._msh2lab(msh)) @staticmethod - def _rgb2msh(rgb): + def _rgb2msh(rgb: np.ndarray) -> np.ndarray: return Colormap._lab2msh(Colormap._rgb2lab(rgb)) @staticmethod - def _hsv2msh(hsv): + def _hsv2msh(hsv: np.ndarray) -> np.ndarray: return Colormap._rgb2msh(Colormap._hsv2rgb(hsv)) @staticmethod - def _hsl2msh(hsl): + def _hsl2msh(hsl: np.ndarray) -> np.ndarray: return Colormap._rgb2msh(Colormap._hsl2rgb(hsl)) @staticmethod - def _xyz2msh(xyz): + def _xyz2msh(xyz: np.ndarray) -> np.ndarray: return Colormap._lab2msh(Colormap._xyz2lab(xyz)) diff --git a/python/tests/test_Colormap.py b/python/tests/test_Colormap.py index 81bf8d2f6..342165134 100644 --- a/python/tests/test_Colormap.py +++ b/python/tests/test_Colormap.py @@ -77,6 +77,12 @@ class TestColormap: # xyz2msh assert np.allclose(Colormap._xyz2msh(xyz),msh,atol=1.e-6,rtol=0) + @pytest.mark.parametrize('low,high',[((0,0,0),(1,1,1)), + ([0,0,0],[1,1,1])]) + def test_from_range_types(self,low,high): + a = Colormap.from_range(low,high) + b = Colormap.from_range(np.array(low),np.array(high)) + assert np.all(a.colors == b.colors) @pytest.mark.parametrize('format',['ASCII','paraview','GOM','gmsh']) @pytest.mark.parametrize('model',['rgb','hsv','hsl','xyz','lab','msh'])