From 45e8ebf9955de53fd56c59da490d68765c32f38f Mon Sep 17 00:00:00 2001 From: RichieHakim Date: Sun, 9 Jun 2024 20:14:07 -0400 Subject: [PATCH] feat: Add complex_colormap function for generating RGB colormap from complex values --- bnpm/plotting_helpers.py | 45 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 45 insertions(+) diff --git a/bnpm/plotting_helpers.py b/bnpm/plotting_helpers.py index d63ad2a..99e6706 100644 --- a/bnpm/plotting_helpers.py +++ b/bnpm/plotting_helpers.py @@ -729,6 +729,51 @@ def __call__(self, x): return colors +def complex_colormap( + mags: np.ndarray, + angles: np.ndarray, + normalize_mags: bool = True, + color_sin: Tuple[int, int, int] = (255, 0, 0), + color_cos: Tuple[int, int, int] = (0, 0, 255), +) -> np.ndarray: + """ + Generates an RGB colormap for complex values based on magnitude and angle. + + The colors vary with the angle and the brightness varies with + the magnitude. + + Args: + mags (np.ndarray): + Array of magnitudes. + angles (np.ndarray): + Array of angles in radians. + normalize_mags (bool): + If True, applies min-max normalization to the magnitudes. + + Returns: + np.ndarray: + Array with RGB values. + """ + assert mags.shape == angles.shape, "The shapes of mags and angles must be the same." + + ## Normalize the magnitudes to the range [0, 1] + if normalize_mags: + mags_norm = (mags - np.min(mags)) / (np.max(mags) - np.min(mags)) + mags_norm = np.clip(mags_norm, 0, 1) + + ## Initialize the RGB array + rgb = np.zeros((mags.size, 3)) + + ## Apply the colors according to the angles + rgb += np.array(color_sin)[None, :] * np.sin(angles)[:, None] + rgb += np.array(color_cos)[None, :] * np.cos(angles)[:, None] + + ## Apply the brightness according to the normalized magnitudes + rgb *= mags_norm[:, None] + + return rgb + + class Figure_Saver: """ Class for saving figures