plotastrodata 1.8.6__tar.gz → 1.8.8__tar.gz

{plotastrodata-1.8.6/plotastrodata.egg-info → plotastrodata-1.8.8}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plotastrodata
-Version: 1.8.6
+Version: 1.8.8
 Summary: plotastrodata is a tool for astronomers to create figures from FITS files and perform fundamental data analyses with ease.
 Home-page: https://github.com/yusukeaso-astron/plotastrodata
 Download-URL: https://github.com/yusukeaso-astron/plotastrodata

{plotastrodata-1.8.6 → plotastrodata-1.8.8}/plotastrodata/__init__.py

@@ -1,4 +1,4 @@
 import warnings
 
 warnings.simplefilter('ignore', FutureWarning)
-__version__ = '1.8.6'
+__version__ = '1.8.8'

{plotastrodata-1.8.6 → plotastrodata-1.8.8}/plotastrodata/plot_utils.py

@@ -95,6 +95,86 @@ def logcbticks(vmin: float = -3.01, vmax: float = 3.01
     return ticks, ticklabels
 
 
+def pow10(x: np.ndarray, stretchpower: float) -> np.ndarray:
+    """A power-law function scaled by xmin. This function is used for the case of stretch='power' in PlotAstroData.add_color().
+
+    Args:
+        x (np.ndarray): Input in the linear scale.
+        stretchpower (float, optional): The output is (data**stretchpower - 1) / stretchpower / ln(10). 1 means the linear scale, while 0 means the logarithmic scale.
+
+    Returns:
+        np.ndarray: Output values.
+    """
+    p = 1e-6 if stretchpower == 0 else stretchpower
+    y = (x**p - 1) / p / np.log(10)
+    return y
+
+
+def ipow10(x: np.ndarray, stretchpower: float) -> np.ndarray:
+    """The inverse function of pow10. This function is used for the case of stretch='power' in PlotAstroData.add_color().
+
+    Args:
+        x (np.ndarray): Input values.
+        stretchpower (float, optional): The input is (data**stretchpower - 1) / stretchpower / ln(10). 1 means the linear scale, while 0 means the logarithmic scale.
+
+    Returns:
+        np.ndarray: Output values in the linear scale.
+    """
+    p = 1e-6 if stretchpower == 0 else stretchpower
+    y = (1 + p * np.log(10) * x)**(1 / p)
+    return y
+
+
+def do_stretch(x: list | np.ndarray,
+               stretch: str, stretchscale: float,
+               stretchpower: float) -> np.ndarray:
+    """Get the stretched values.
+
+    Args:
+        x (list | np.ndarray): Input array in the linear scale.
+        stretch (str): 'log', 'asinh', 'power', or 'linear'. Any other means 'linear'. 'log' means the mapped data are logarithmic. 'asinh' means the mapped data are arc sin hyperbolic. 'power' means the mapped data are power-law (see also stretchpower). Defaults to 'linear'.
+        stretchscale (float, optional): The output is asinh(data / stretchscale). Defaults to None.
+        stretchpower (float, optional): The output is ((data / stretchscale)**stretchpower - 1) / stretchpower / ln(10). 1 means the linear scale, while 0 means the logarithmic scale.
+
+    Returns:
+        np.ndarray: Output stretched array.
+    """
+    t = np.array(x)
+    match stretch:
+        case 'log':
+            t = np.log10(t)
+        case 'asinh':
+            t = np.arcsinh(t / stretchscale)
+        case 'power':
+            t = pow10(t / stretchscale, stretchpower)
+    return t
+
+
+def undo_stretch(x: list | np.ndarray,
+                 stretch: str, stretchscale: float,
+                 stretchpower: float) -> np.ndarray:
+    """Get the linear values from the stretched values.
+
+    Args:
+        x (list | np.ndarray): Input stretched array.
+        stretch (str): 'log', 'asinh', 'power', or 'linear'. Any other means 'linear'. 'log' means the mapped data are logarithmic. 'asinh' means the mapped data are arc sin hyperbolic. 'power' means the mapped data are power-law (see also stretchpower). Defaults to 'linear'.
+        stretchscale (float, optional): The input is asinh(data / stretchscale). Defaults to None.
+        stretchpower (float, optional): The input is ((data / stretchscale)**stretchpower - 1) / stretchpower / ln(10). 1 means the linear scale, while 0 means the logarithmic scale.
+
+    Returns:
+        np.ndarray: Output array in the linear scale.
+    """
+    t = np.array(x)
+    match stretch:
+        case 'log':
+            t = 10**t
+        case 'asinh':
+            t = np.sinh(t) * stretchscale
+        case 'power':
+            t = ipow10(t, stretchpower) * stretchscale
+    return t
+
+
 @dataclass
 class PlotAxes2D():
     """Use Axes.set_* to adjust x and y axes.
@@ -199,16 +279,16 @@ class PlotAxes2D():
 
 
 def set_minmax(data: np.ndarray, stretch: str, stretchscale: float,
-               stretchpower: float, rms: float, kw: dict
+               stretchpower: float, sigma: float, kw: dict
                ) -> np.ndarray:
     """Set vmin and vmax for color pcolormesh and RGB maps.
 
     Args:
         data (np.ndarray): Plotted data.
         stretch (str): 'log', 'asinh', 'power'. Any other means linear.
-        stretchscale (float): For the arcsinh strech.
+        stretchscale (float): For the arcsinh strech and the power stretch.
         stretchpower (float): For the power strech.
-        rms (float): RMS noise level.
+        sigma (float): Noise level.
         kw (dict): Probably like {'vmin':0, 'vmax':1}.
 
     Returns:
@@ -216,51 +296,38 @@ def set_minmax(data: np.ndarray, stretch: str, stretchscale: float,
     """
     if type(stretch) is str:
         data = [data]
-        rms = [rms]
+        sigma = [sigma]
         stretch = [stretch]
         stretchscale = [stretchscale]
-        if 'vmin' in kw:
-            kw['vmin'] = [kw['vmin']]
-        if 'vmax' in kw:
-            kw['vmax'] = [kw['vmax']]
-    z = (data, stretch, stretchscale, rms)
-    for i, (c, st, stsc, r) in enumerate(zip(*z)):
-        if stsc is None:
-            stsc = r
-        if st == 'log':
-            if np.any(c > 0):
-                c = np.log10(c.clip(np.nanmin(c[c > 0]), None))
-        elif st == 'asinh':
-            c = np.arcsinh(c / stsc)
-        elif st == 'power':
-            cmin = kw['min'][i] if 'vmin' in kw else r
-            c = c.clip(cmin, None)
-            p = 1 - stretchpower
-            c = ((c / cmin)**p - 1) / p / np.log(10)
-        data[i] = c
+        stretchpower = [stretchpower]
+        for k in ['vmin', 'vmax']:
+            if k in kw:
+                kw[k] = [kw[k]]
+
     n = len(data)
-    for m in ['vmin', 'vmax']:
-        if m in kw:
-            for i, (c, st, stsc, _) in enumerate(zip(*z)):
-                if st == 'log':
-                    kw[m][i] = np.log10(kw[m][i])
-                elif st == 'asinh':
-                    kw[m][i] = np.arcsinh(kw[m][i] / stsc)
-                elif st == 'power':
-                    p = 1 - stretchpower
-                    kw[m][i] = ((kw[m][i]/c.min())**p - 1) / p / np.log(10)
-        else:
-            kw[m] = [None] * n
-            for i, (c, st, _, r) in enumerate(zip(*z)):
-                if m == 'vmin':
-                    kw[m][i] = np.log10(r) if st == 'log' else np.nanmin(c)
-                else:
-                    kw[m][i] = np.nanmax(c)
+    for k in ['vmin', 'vmax']:
+        if k not in kw:
+            kw[k] = [None] * n
+    stretch = np.where(np.equal(stretch, None), sigma, stretch)
+    cmin = np.where(np.equal(kw['vmin'], None), sigma, kw['vmin'])
+
+    argslist = (stretch, stretchscale, stretchpower)
+    for i, stretch_args in enumerate(zip(*argslist)):
+        c = data[i]
+        if stretch_args[0] in ['log', 'power']:
+            c = c.clip(cmin[i], None)
+        c = do_stretch(c, *stretch_args)
+        data[i] = c
+        for k in ['vmin', 'vmax']:
+            if kw[k][i] is None:
+                kw[k][i] = np.nanmin(c) if k == 'vmin' else np.nanmax(c)
+            else:
+                kw[k][i] = do_stretch(kw[k][i], *stretch_args)
     data = [c.clip(a, b) for c, a, b in zip(data, kw['vmin'], kw['vmax'])]
     if n == 1:
         data = data[0]
-        kw['vmin'] = kw['vmin'][0]
-        kw['vmax'] = kw['vmax'][0]
+        for k in ['vmin', 'vmax']:
+            kw[k] = kw[k][0]
     return data
 
 
@@ -765,10 +832,49 @@ class PlotAstroData(AstroFrame):
             axnow.plot([x[0] - length/2., x[0] + length/2.], [y[0], y[0]],
                        '-', linewidth=linewidth, color=color)
 
+    def _set_colorbar(self, mappable, ch: int, show_cbar: bool,
+                      cblabel: str, cbformat: str,
+                      cbticks: list | None, cbticklabels: list | None,
+                      cblocation: str, stretch: str,
+                      stretchscale: float, stretchpower: float,
+                      vmin: float, vmax: float):
+        if not show_cbar:
+            return
+
+        if self.fig is None:
+            fig = plt.figure(ch // self.rowcol)
+        else:
+            fig = self.fig
+        if len(self.ax) == 1:
+            ax = self.ax[ch]
+            cb = fig.colorbar(mappable[ch], ax=ax, label=cblabel,
+                              format=cbformat, location=cblocation)
+        else:
+            cax = plt.axes([0.88, 0.105, 0.015, 0.77])
+            cb = fig.colorbar(mappable[ch], cax=cax, label=cblabel,
+                              format=cbformat)
+        cb.ax.tick_params(labelsize=14)
+        font = mpl.font_manager.FontProperties(size=16)
+        cb.ax.yaxis.label.set_font_properties(font)
+        stretch_args = (stretch, stretchscale, stretchpower)
+        if cbticks is not None and ch // self.rowcol == 0:
+            cbticks = do_stretch(cbticks, *stretch_args)
+        if cbticks is None and stretch == 'log':
+            cbticks, cbticklabels = logcbticks()
+        if cbticks is not None:
+            cond = (vmin < cbticks) * (cbticks < vmax)
+            cb.set_ticks(cbticks[cond])
+            if cbticklabels is not None:
+                tl = np.array(cbticklabels)[cond]
+            else:
+                t = undo_stretch(cb.get_ticks(), *stretch_args)
+                tl = [f'{d:{cbformat[1:]}}' for d in t]
+            cb.set_ticklabels(tl)
+
     def add_color(self,
                   stretch: str = 'linear',
                   stretchscale: float | None = None,
-                  stretchpower: float = 0,
+                  stretchpower: float = 1,
                   show_cbar: bool = True,
                   cblabel: str | None = None,
                   cbformat: float = '%.1e',
@@ -779,9 +885,9 @@ class PlotAstroData(AstroFrame):
         """Use Axes.pcolormesh of matplotlib. kwargs must include the arguments of AstroData to specify the data to be plotted. kwargs may include arguments for add_beam() and a dict of beam_kwargs to specify the beam patch in more detail. kwargs may include xskiip and yskip.
 
         Args:
-            stretch (str, optional): 'log' means the mapped data are logarithmic. 'asinh' means the mapped data are arc sin hyperbolic. 'power' means the mapped data are power-law (see also stretchpower). Defaults to 'linear'.
-            stretchscale (float, optional): color scale is asinh(data / stretchscale). Defaults to None.
-            stretchpower (float, optional): color scale is ((data / vmin)**(1 - stretchpower) - 1) / (1 - stretchpower) / ln(10). 0 means the linear scale. 1 means the logarithmic scale. Defaults to 0.
+            stretch (str, optional): 'log', 'asinh', 'power', or 'linear'. Any other means 'linear'. 'log' means the mapped data are logarithmic. 'asinh' means the mapped data are arc sin hyperbolic. 'power' means the mapped data are power-law (see also stretchpower). Defaults to 'linear'.
+            stretchscale (float, optional): Color scale is asinh(data / stretchscale). Defaults to None.
+            stretchpower (float, optional): Color scale is ((data / stretchscale)**stretchpower - 1) / stretchpower / ln(10). 1 means the linear scale, while 0 means the logarithmic scale. Defaults to 1.
             show_cbar (bool, optional): Show color bar. Defaults to True.
             cblabel (str, optional): Colorbar label. Defaults to None.
             cbformat (float, optional): Format for ticklabels of colorbar. Defaults to '%.1e'.
@@ -797,11 +903,11 @@ class PlotAstroData(AstroFrame):
             print('No pixel size. Skip add_color.')
             return
 
-        if stretchscale is None:
-            stretchscale = sigma
-        cmin_org = _kw['vmin'] if 'vmin' in _kw else sigma
-        c = set_minmax(c, stretch, stretchscale, stretchpower, sigma, _kw)
-        cmin, cmax = np.nanmin(c), np.nanmax(c)
+        if cblabel is None:
+            cblabel = bunit
+        
+        stretch_args = (stretch, stretchscale, stretchpower)
+        c = set_minmax(c, *stretch_args, sigma, _kw)
         c = self.vskipfill(c, v)
         if type(self.channelnumber) is int:
             c = [c[self.channelnumber]]
@@ -811,54 +917,9 @@ class PlotAstroData(AstroFrame):
             if ch in self.bottomleft:
                 p[ch] = pnow
         for ch in self.bottomleft:
-            if not show_cbar:
-                break
-            if cblabel is None:
-                cblabel = bunit
-            if self.fig is None:
-                fig = plt.figure(ch // self.rowcol)
-            else:
-                fig = self.fig
-            if len(self.ax) == 1:
-                ax = self.ax[ch]
-                cb = fig.colorbar(p[ch], ax=ax, label=cblabel,
-                                  format=cbformat, location=cblocation)
-            else:
-                cax = plt.axes([0.88, 0.105, 0.015, 0.77])
-                cb = fig.colorbar(p[ch], cax=cax, label=cblabel,
-                                  format=cbformat)
-            cb.ax.tick_params(labelsize=14)
-            font = mpl.font_manager.FontProperties(size=16)
-            cb.ax.yaxis.label.set_font_properties(font)
-            if cbticks is not None and ch // self.rowcol == 0:
-                cbticks = np.array(cbticks)
-                match stretch:
-                    case 'log':
-                        cbticks = np.log10(cbticks)
-                    case 'asinh':
-                        cbticks = np.arcsinh(cbticks / stretchscale)
-                    case 'power':
-                        p = 1 - stretchpower
-                        cbticks = (cbticks / cmin_org)**p - 1
-                        cbticks = cbticks / p / np.log(10)
-            if stretch == 'log' and cbticks is None:
-                cbticks, cbticklabels = logcbticks(cmin, cmax)
-            if cbticks is not None:
-                cb.set_ticks(cbticks)
-            if cbticklabels is not None:
-                cb.set_ticklabels(cbticklabels)
-            elif stretch in ['log', 'asinh', 'power']:
-                t = cb.get_ticks()
-                t = t[(cmin < t) * (t < cmax)]
-                cb.set_ticks(t)
-                if stretch == 'log':
-                    ticklin = 10**t
-                elif stretch == 'asinh':
-                    ticklin = np.sinh(t) * stretchscale
-                elif stretch == 'power':
-                    p = 1 - stretchpower
-                    ticklin = cmin_org * (1 + p * np.log(10) * t)**(1 / p)
-                cb.set_ticklabels([f'{d:{cbformat[1:]}}' for d in ticklin])
+            self._set_colorbar(p, ch, show_cbar, cblabel, cbformat,
+                               cbticks, cbticklabels, cblocation,
+                               *stretch_args, _kw['vmin'], _kw['vmax'])
         self.add_beam(beam=beam, **beam_kwargs)
 
     def add_contour(self,
@@ -950,14 +1011,14 @@ class PlotAstroData(AstroFrame):
     def add_rgb(self,
                 stretch: list[str, str, str] = ['linear'] * 3,
                 stretchscale: list[float | None, float | None, float | None] = [None] * 3,
-                stretchpower: float = 0,
+                stretchpower: list[float, float, float] = [1, 1, 1],
                 **kwargs) -> None:
         """Use PIL.Image and imshow of matplotlib. kwargs must include the arguments of AstroData to specify the data to be plotted. A three-element array ([red, green, blue]) is supposed for all arguments, except for xskip, yskip and show_beam, including vmax and vmin. kwargs may include arguments for add_beam() and a dict of beam_kwargs to specify the beam patch in more detail. kwargs may include xskiip and yskip.
 
         Args:
-            stretch (str, optional): 'log' means the mapped data are logarithmic. 'asinh' means the mapped data are arc sin hyperbolic. 'power' means the mapped data are power-law (see also stretchpower). Defaults to 'linear'.
-            stretchscale (float, optional): color scale is asinh(data / stretchscale). Defaults to None.
-            stretchpower (float, optional): color scale is ((data / vmin)**(1 - stretchpower) - 1) / (1 - stretchpower) / ln(10). 0 means the linear scale. 1 means the logarithmic scale. Defaults to 0.
+            stretch (str, optional): 'log', 'asinh', 'power', or 'linear'. Any other means 'linear'. 'log' means the mapped data are logarithmic. 'asinh' means the mapped data are arc sin hyperbolic. 'power' means the mapped data are power-law (see also stretchpower). Defaults to 'linear'.
+            stretchscale (float, optional): Color scale is asinh(data / stretchscale). Defaults to None.
+            stretchpower (float, optional): Color scale is ((data / stretchscale)**stretchpower - 1) / stretchpower / ln(10). 1 means the linear scale, while 0 means the logarithmic scale. Defaults to 1.
         """
         from PIL import Image
 
@@ -968,10 +1029,8 @@ class PlotAstroData(AstroFrame):
             print('No pixel size. Skip add_rgb.')
             return
 
-        for i in range(len(stretchscale)):
-            if stretchscale[i] is None:
-                stretchscale[i] = sigma[i]
-        c = set_minmax(c, stretch, stretchscale, stretchpower, sigma, _kw)
+        stretch_args = (stretch, stretchscale, stretchpower)
+        c = set_minmax(c, *stretch_args, sigma, _kw)
         if not (np.shape(c[0]) == np.shape(c[1]) == np.shape(c[2])):
             print('RGB shapes mismatch. Skip add_rgb.')
             return

{plotastrodata-1.8.6 → plotastrodata-1.8.8/plotastrodata.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plotastrodata
-Version: 1.8.6
+Version: 1.8.8
 Summary: plotastrodata is a tool for astronomers to create figures from FITS files and perform fundamental data analyses with ease.
 Home-page: https://github.com/yusukeaso-astron/plotastrodata
 Download-URL: https://github.com/yusukeaso-astron/plotastrodata