plotastrodata 1.8.11__tar.gz → 1.8.12__tar.gz

{plotastrodata-1.8.11/plotastrodata.egg-info → plotastrodata-1.8.12}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plotastrodata
-Version: 1.8.11
+Version: 1.8.12
 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.11 → plotastrodata-1.8.12}/plotastrodata/__init__.py

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

{plotastrodata-1.8.11 → plotastrodata-1.8.12}/plotastrodata/analysis_utils.py

@@ -8,7 +8,7 @@ from scipy.signal import convolve
 from plotastrodata import const_utils as cu
 from plotastrodata.coord_utils import coord2xy, rel2abs, xy2coord
 from plotastrodata.fits_utils import data2fits, FitsData, Jy2K
-from plotastrodata.fitting_utils import EmceeCorner
+from plotastrodata.fitting_utils import EmceeCorner, gaussfit1d
 from plotastrodata.matrix_utils import dot2d, Mfac, Mrot
 from plotastrodata.noise_utils import estimate_rms
 from plotastrodata.other_utils import (gaussian2d, isdeg,
@@ -442,7 +442,6 @@ class AstroData():
             xlist, ylist = coord2xy(coords, self.center) * 3600.
         nprof = len(xlist)
         v = self.v
-        dv = self.dv
         data, xf, yf = filled2d(self.data, self.x, self.y, ninterp)
         x, y = np.meshgrid(xf, yf)
         prof = np.empty((nprof, len(v)))
@@ -469,22 +468,12 @@ class AstroData():
                 prof *= np.abs(self.dx * self.dy) / Omega
         gfitres = {}
         if gaussfit:
-            xmin, xmax = np.min(v), np.max(v)
-            ymin, ymax = np.min(prof), np.max(prof)
-            bounds = [[ymin, xmin, np.abs(dv)], [ymax, xmax, xmax - xmin]]
-
-            def gauss(x, p, c, w):
-                return p * np.exp(-4. * np.log(2.) * ((x - c) / w)**2)
-
-            nprof = len(prof)
-            best, error = [None] * nprof, [None] * nprof
+            res = [None] * nprof
             for i in range(nprof):
-                popt, pcov = curve_fit(gauss, v, prof[i], bounds=bounds)
-                perr = np.sqrt(np.diag(pcov))
-                print('Gauss (peak, center, FWHM):', popt)
-                print('Gauss uncertainties:', perr)
-                best[i], error[i] = popt, perr
-            gfitres = {'best': best, 'error': error}
+                res[i] = gaussfit1d(xdata=v, ydata=prof[i],
+                                    sigma=None, show=True)
+            gfitres['best'] = [a['popt'][:3] for a in res]
+            gfitres['error'] = [a['perr'][:3] for a in res]
         return v, prof, gfitres
 
     def rotate(self, pa: float = 0, **kwargs):
@@ -642,34 +631,24 @@ class AstroFrame():
     def __post_init__(self):
         self.xdir = -1 if self.xflip else 1
         self.ydir = -1 if self.yflip else 1
-        if self.xmax is None:
-            self.xmax = self.rmax
-        if self.xmin is None:
-            self.xmin = -self.rmax
-        if self.ymax is None:
-            self.ymax = self.rmax
-        if self.ymin is None:
-            self.ymin = -self.rmax
-        if self.xdir == -1:
+        self.xmin = -self.rmax if self.xmin is None else self.xmin
+        self.xmax = self.rmax if self.xmax is None else self.xmax
+        self.ymin = -self.rmax if self.ymin is None else self.ymin
+        self.ymax = self.rmax if self.ymax is None else self.ymax
+        if self.xflip:
             self.xmin, self.xmax = self.xmax, self.xmin
-        if self.ydir == -1:
+        if self.yflip:
             self.ymin, self.ymax = self.ymax, self.ymin
         xlim = [self.xoff + self.xmin, self.xoff + self.xmax]
         ylim = [self.yoff + self.ymin, self.yoff + self.ymax]
         vlim = [self.vmin, self.vmax]
         if self.pv:
-            xlim = np.sort(xlim)
+            xlim = sorted(xlim)
             if not self.xflip:
-                xlim = xlim[::-1]
-        self.xlim = xlim
-        self.ylim = ylim
-        self.vlim = vlim
-        if self.pv:
-            self.Xlim = vlim if self.swapxy else xlim
-            self.Ylim = xlim if self.swapxy else vlim
-        else:
-            self.Xlim = ylim if self.swapxy else xlim
-            self.Ylim = xlim if self.swapxy else ylim
+                xlim.reverse()
+        self.xlim, self.ylim, self.vlim = xlim, ylim, vlim
+        _x, _y = (xlim, vlim) if self.pv else (xlim, ylim)
+        self.Xlim, self.Ylim = (_y, _x) if self.swapxy else (_x, _y)
         if self.quadrants is not None:
             self.Xlim = [0, self.rmax]
             self.Ylim = [0, min(self.vmax, -self.vmin)]

{plotastrodata-1.8.11 → plotastrodata-1.8.12}/plotastrodata/fits_utils.py

@@ -398,49 +398,39 @@ def fits2data(fitsimage: str, Tb: bool = False, log: bool = False,
     return fd.data, (fd.x, fd.y, fd.v), beam, bunit, rms
 
 
-def data2fits(d: np.ndarray | None = None, h: dict = {},
+def data2fits(d: np.ndarray, h: dict = {},
               templatefits: str | None = None,
               fitsimage: str = 'test') -> None:
     """Make a fits file from a N-D array.
 
     Args:
-        d (np.ndarray, optional): N-D array. Defaults to None.
-        h (dict, optional): Fits header. Defaults to {}.
-        templatefits (str, optional): Fits file to copy header. Defaults to None.
-        fitsimage (str, optional): Output name. Defaults to 'test'.
+        d (np.ndarray): N-D array.
+        h (dict, optional): Additional FITS header. Defaults to {}.
+        templatefits (str, optional): FITS file whose header is used as a temperate. Defaults to None.
+        fitsimage (str, optional): Output filename, with or without '.fits'. Defaults to 'test'.
     """
-    naxis = np.ndim(d)
-    w = wcs.WCS(naxis=naxis)
     _h = {} if templatefits is None else FitsData(templatefits).get_header()
     _h.update(h)
+    naxis = np.ndim(d)
+    w = wcs.WCS(naxis=naxis)
     if _h == {}:
         w.wcs.crpix = [0] * naxis
         w.wcs.crval = [0] * naxis
         w.wcs.cdelt = [1] * naxis
-    ctype = ['RA---SIN', 'DEC--SIN', 'FREQ']
-    if 'CTYPE1' in _h:
-        ctype[0] = _h['CTYPE1']
-    if 'CTYPE2' in _h:
-        ctype[1] = _h['CTYPE2']
-    if 'CTYPE3' in _h:
-        ctype[2] = _h['CTYPE3']
-    w.wcs.ctype = ctype[:naxis]
-    cunit = ['deg', 'deg', 'Hz']
-    if 'CUNIT1' in _h:
-        cunit[0] = _h['CUNIT1']
-    if 'CUNIT2' in _h:
-        cunit[1] = _h['CUNIT2']
-    if 'CUNIT3' in _h:
-        cunit[2] = _h['CUNIT3']
-    w.wcs.cunit = cunit[:naxis]
+    defaults = {'CTYPE': ['RA---SIN', 'DEC--SIN', 'FREQ'],
+                'CUNIT': ['deg', 'deg', 'Hz']}
+    for k, v in defaults.items():
+        for i in range(naxis):
+            _h.setdefault(f'{k}{i+1:d}', v[i])
+    w.wcs.ctype = [_h[f'CTYPE{i+1}'] for i in range(naxis)]
+    w.wcs.cunit = [_h[f'CUNIT{i+1}'] for i in range(naxis)]
+    _h.setdefault('BUNIT', 'Jy/beam')
+    if naxis >= 3:
+        _h.setdefault('SPECSYS', 'LSRK')
     header = w.to_header()
     hdu = fits.PrimaryHDU(d, header=header)
-    if 'BUNIT' not in _h:
-        _h['BUNIT'] = 'Jy/beam'
-    if naxis >= 3 and 'SPECSYS' not in _h:
-        _h['SPECSYS'] = 'LSRK'
-    for k in _h:
-        if not ('COMMENT' in k or 'HISTORY' in k) and _h[k] is not None:
-            hdu.header[k] = _h[k]
+    for k, v in _h.items():
+        if v is not None and 'COMMENT' not in k and 'HISTORY' not in k:
+            hdu.header[k] = v
     hdu = fits.HDUList([hdu])
-    hdu.writeto(fitsimage.replace('.fits', '') + '.fits', overwrite=True)
+    hdu.writeto(fitsimage.removesuffix('.fits') + '.fits', overwrite=True)

{plotastrodata-1.8.11 → plotastrodata-1.8.12}/plotastrodata/fitting_utils.py

@@ -393,3 +393,52 @@ class EmceeCorner():
         error = evidence * results.logzerr[-1]
         self.evidence = evidence
         return {'evidence': evidence, 'error': error}
+
+
+def gaussfit1d(xdata: np.ndarray, ydata: np.ndarray,
+               sigma: float | np.ndarray | None,
+               show: bool = False, **kwargs) -> dict:
+    """Gaussian fitting to a pair of 1D arrays.
+
+    Args:
+        xdata (np.ndarray): ydata is compared with Gauss(xdata).
+        ydata (np.ndarray): ydata is compared with Gauss(xdata).
+        sigma (float | np.ndarray | None): Noise level of ydata. If None is given, sigma is also a free parameter. Defaults to None.
+        show (bool, optional): True means to show the best-fit parameters and uncertainties. Defaults to False.
+
+    Returns:
+        dict: _description_
+    """
+    if sigma is not None and np.shape(sigma) == ():
+        sigma = [sigma] * len(xdata)
+    xmin, xmax = np.min(xdata), np.max(xdata)
+    ymin, ymax = np.min(ydata), np.max(ydata)
+    dx = np.abs(xdata[1] - xdata[0])
+    bounds = [[ymin, ymax], [xmin, xmax], [dx, xmax - xmin]]
+    if sigma is None:
+        bounds.append([np.log(ymax * 1e-6), np.log(ymax)])
+
+    def g(x, p):
+        a, c, w = p
+        return a * np.exp(-4. * np.log(2.) * ((x - c) / w)**2)
+
+    if sigma is None:
+        def logl(p):
+            sigmdl = np.exp(p[3])
+            chi2 = np.sum(((ydata - g(xdata, p[:3])) / sigmdl)**2)
+            return -0.5 * chi2 - p[3]
+    else:
+        def logl(p):
+            chi2 = np.sum(((ydata - g(xdata, p)) / sigma)**2)
+            return -0.5 * chi2
+
+    fitter = EmceeCorner(bounds=bounds, logl=logl)
+    fitter.fit(**kwargs)
+    popt = fitter.popt
+    plow = fitter.plow
+    phigh = fitter.phigh
+    perr = (phigh - plow) / 2
+    if show:
+        print('Gauss (peak, center, FWHM):', popt)
+        print('Gauss uncertainties:', perr)
+    return {'popt': popt, 'perr': perr}

{plotastrodata-1.8.11 → plotastrodata-1.8.12}/plotastrodata/other_utils.py

@@ -37,11 +37,11 @@ def isdeg(s: str) -> bool:
 
 def trim(data: np.ndarray | None = None, x: np.ndarray | None = None,
          y: np.ndarray | None = None, v: np.ndarray | None = None,
-         xlim: list[float, float] | None = None,
-         ylim: list[float, float] | None = None,
-         vlim: list[float, float] | None = None,
+         xlim: list[float] | None = None,
+         ylim: list[float] | None = None,
+         vlim: list[float] | None = None,
          pv: bool = False
-         ) -> tuple[np.ndarray, list[np.ndarray, np.ndarray, np.ndarray]]:
+         ) -> tuple[np.ndarray | None, list[np.ndarray | None]]:
     """Trim 2D or 3D data by given coordinates and their limits.
 
     Args:
@@ -56,47 +56,37 @@ def trim(data: np.ndarray | None = None, x: np.ndarray | None = None,
     Returns:
         tuple: Trimmed (data, [x,y,v]).
     """
-    xout, yout, vout, dataout = x, y, v, data
-    i0 = j0 = k0 = 0
-    i1 = j1 = k1 = 100000
-    if x is not None and xlim is not None:
-        if None not in xlim:
-            x0 = np.max([np.min(x), xlim[0]])
-            x1 = np.min([np.max(x), xlim[1]])
-            i0 = np.argmin(np.abs(x - x0))
-            i1 = np.argmin(np.abs(x - x1))
-            i0, i1 = sorted([i0, i1])
-            xout = x[i0:i1+1]
-    if y is not None and ylim is not None:
-        if None not in ylim:
-            y0 = np.max([np.min(y), ylim[0]])
-            y1 = np.min([np.max(y), ylim[1]])
-            j0 = np.argmin(np.abs(y - y0))
-            j1 = np.argmin(np.abs(y - y1))
-            j0, j1 = sorted([j0, j1])
-            yout = y[j0:j1+1]
-    if v is not None and vlim is not None:
-        if None not in vlim:
-            v0 = np.max([np.min(v), vlim[0]])
-            v1 = np.min([np.max(v), vlim[1]])
-            k0 = np.argmin(np.abs(v - v0))
-            k1 = np.argmin(np.abs(v - v1))
-            k0, k1 = sorted([k0, k1])
-            vout = v[k0:k1+1]
-    if data is not None:
-        d = np.squeeze(data)
-        if np.ndim(d) == 0:
-            print('data has only one pixel.')
-            d = data
-        if np.ndim(d) == 2:
-            if pv:
-                j0, j1 = k0, k1
-            dataout = d[j0:j1+1, i0:i1+1]
-        else:
-            d = np.moveaxis(d, [-3, -2, -1], [0, 1, 2])
-            d = d[k0:k1+1, j0:j1+1, i0:i1+1]
-            d = np.moveaxis(d, [0, 1, 2], [-3, -2, -1])
-            dataout = d
+    def get_bounds(arr, lim):
+        if arr is None or lim is None or None in lim:
+            return arr, 0, None
+        lo = np.argmin(np.abs(arr - max(np.min(arr), lim[0])))
+        hi = np.argmin(np.abs(arr - min(np.max(arr), lim[1])))
+        lo, hi = sorted((lo, hi))
+        return arr[lo:hi + 1], lo, hi + 1
+
+    xout, i0, i1 = get_bounds(x, xlim)
+    yout, j0, j1 = get_bounds(y, ylim)
+    vout, k0, k1 = get_bounds(v, vlim)
+
+    if data is None:
+        return None, [xout, yout, vout]
+
+    d = np.squeeze(data)
+
+    if d.ndim == 0:
+        print("data has only one pixel.")
+        return data, [xout, yout, vout]
+
+    if d.ndim == 2:
+        if pv:
+            j0, j1 = k0, k1
+        dataout = d[j0:j1, i0:i1]
+    else:
+        d = np.moveaxis(d, [-3, -2, -1], [0, 1, 2])
+        d = d[k0:k1, j0:j1, i0:i1]
+        d = np.moveaxis(d, [0, 1, 2], [-3, -2, -1])
+        dataout = d
+
     return dataout, [xout, yout, vout]
 
 

{plotastrodata-1.8.11 → plotastrodata-1.8.12}/plotastrodata/plot_utils.py

@@ -96,56 +96,110 @@ def logcbticks(vmin: float = 1e-3, vmax: float = 1e3
     return ticks[cond], ticklabels[cond]
 
 
-def do_stretch(x: list | np.ndarray,
-               stretch: str, stretchscale: float,
-               stretchpower: float) -> np.ndarray:
+@dataclass
+class Stretcher():
     """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): The output is asinh(data / stretchscale).
-        stretchpower (float): The output is data**stretchpower / stretchpower. 1 means the linear scale, while 0 means the logarithmic scale.
-
-    Returns:
-        np.ndarray: Output stretched array.
+        Args:
+            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): The output is asinh(data / stretchscale). Defaults to None.
+            stretchpower (float, optional): The output is data**stretchpower / stretchpower. 1 means the linear scale, while 0 means the logarithmic scale. Defaults to 0.5.
+            vmin (float, optional): The minimum value for Axes.pcolormesh() of matplotlib. Defaults to None.
+            vmax (float, optional): The maximum value for Axes.pcolormesh() of matplotlib. Defaults to None.
+            sigma (float, optional): Noise level. Defaults to 0.
     """
-    t = np.array(x)
-    match stretch:
-        case 'log':
-            t = np.log10(t)  # To be consistent with logcbticks().
-        case 'asinh':
-            t = np.arcsinh(t / stretchscale)
-        case 'power':
-            p = 1e-6 if stretchpower == 0 else stretchpower
-            t = t**p / p
-    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.
+    stretch: str = 'linear'
+    stretchscale: float | None = None
+    stretchpower: float = 0.5
+    vmin: float | None = None
+    vmax: float | None = None
+    sigma: float = 0
+
+    def __post_init__(self):
+        self.n = 1 if type(self.stretch) is str else len(self.stretch)
+        stretch = self.stretch
+        stsc = self.stretchscale
+        vmin = self.vmin
+        sigma = self.sigma
+        if self.n == 1:
+            if stsc is None:
+                self.stretchscale = sigma
+            if (stretch == 'log' or stretch == 'power') and vmin is None:
+                self.vmin = sigma
+        else:
+            getsigma = np.equal(stsc, None)
+            self.stretchscale = np.where(getsigma, sigma, stsc)
+            islog = np.equal(stretch, 'log')
+            ispower = np.equal(stretch, 'power')
+            novmin = np.equal(vmin, None)
+            getsigma = (islog + ispower) * novmin
+            self.vmin = np.where(getsigma, sigma, vmin)
 
-    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): The input is asinh(data / stretchscale).
-        stretchpower (float): The input is data**stretchpower / stretchpower. 1 means the linear scale, while 0 means the logarithmic scale.
+    def do(self, x: list | np.ndarray) -> np.ndarray:
+        """Get the stretched values.
 
-    Returns:
-        np.ndarray: Output array in the linear scale.
-    """
-    t = np.array(x)
-    match stretch:
-        case 'log':
-            t = 10**t  # To be consistent with logcbticks().
-        case 'asinh':
-            t = np.sinh(t) * stretchscale
-        case 'power':
-            p = 1e-6 if stretchpower == 0 else stretchpower
-            t = (t * p)**(1 / p)
-    return t
+        Args:
+            x (list | np.ndarray): Input array in the linear scale.
+
+        Returns:
+            np.ndarray: Output stretched array.
+        """
+        t = np.array(x)
+        match self.stretch:
+            case 'log':
+                t = np.log10(t)  # To be consistent with logcbticks().
+            case 'asinh':
+                t = np.arcsinh(t / self.stretchscale)
+            case 'power':
+                p = 1e-6 if self.stretchpower == 0 else self.stretchpower
+                t = t**p / p
+        return t
+
+    def undo(self, x: list | np.ndarray) -> np.ndarray:
+        """Get the linear values from the stretched values.
+
+        Args:
+            x (list | np.ndarray): Input stretched array.
+
+        Returns:
+            np.ndarray: Output array in the linear scale.
+        """
+        t = np.array(x)
+        match self.stretch:
+            case 'log':
+                t = 10**t  # To be consistent with logcbticks().
+            case 'asinh':
+                t = np.sinh(t) * self.stretchscale
+            case 'power':
+                p = 1e-6 if self.stretchpower == 0 else self.stretchpower
+                t = (t * p)**(1 / p)
+        return t
+
+    def set_minmax(self, data: np.ndarray
+                   ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Set vmin and vmax for color pcolormesh and RGB maps.
+
+        Args:
+            data (np.ndarray): 2D/3D data to plot.
+
+        Returns:
+            tuple[np.ndarray, np.ndarray, np.ndarray]: (Clipped stretched data, new vmin, new vmax).
+        """
+        single = self.n == 1
+        vminout = [self.vmin] if single else self.vmin
+        vmaxout = [self.vmax] if single else self.vmax
+        dataout = [data] if single else data
+        for i, (c, v0, v1) in enumerate(zip(dataout, vminout, vmaxout)):
+            dataout[i] = cout = self.do(c.clip(v0, v1))
+            vminout[i] = np.nanmin(cout)
+            vmaxout[i] = np.nanmax(cout)
+        if single:
+            dataout = dataout[0]
+            vminout = vminout[0]
+            vmaxout = vmaxout[0]
+        self.vmin = vminout
+        self.vmax = vmaxout
+        return dataout, vminout, vmaxout
 
 
 @dataclass
@@ -251,61 +305,6 @@ class PlotAxes2D():
                 ax.set_aspect(self.aspect)
 
 
-def set_minmax(data: np.ndarray, stretch: str, stretchscale: float,
-               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', 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): The input is asinh(data / stretchscale).
-        stretchpower (float): The input is data**stretchpower / stretchpower. 1 means the linear scale, while 0 means the logarithmic scale.
-        sigma (float): Noise level.
-        kw (dict): Probably like {'vmin':0, 'vmax':1}.
-
-    Returns:
-        np.ndarray: Data clipped with the vmin and vmax.
-    """
-    if type(stretch) is str:
-        data = [data]
-        sigma = [sigma]
-        stretch = [stretch]
-        stretchscale = [stretchscale]
-        stretchpower = [stretchpower]
-        for k in ['vmin', 'vmax']:
-            if k in kw:
-                kw[k] = [kw[k]]
-
-    n = len(data)
-    for k in ['vmin', 'vmax']:
-        if k not in kw:
-            kw[k] = [None] * n
-    isnan = np.equal(stretchscale, None)
-    stretchscale = np.where(isnan, sigma, stretchscale)
-    isnan = np.equal(kw['vmin'], None)
-    cmin = np.where(isnan, 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]
-        for k in ['vmin', 'vmax']:
-            kw[k] = kw[k][0]
-    return data
-
-
 def kwargs2AstroData(kw: dict) -> AstroData:
     """Get AstroData and remove its arguments from kwargs.
 
@@ -802,9 +801,7 @@ class PlotAstroData(AstroFrame):
     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):
+                      cblocation: str, st: Stretcher):
         if not show_cbar:
             return
 
@@ -823,27 +820,19 @@ class PlotAstroData(AstroFrame):
         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 None and stretch == 'log':
-            cbticks, cbticklabels = logcbticks(10**vmin, 10**vmax)
-        if cbticks is not None:
-            cbticks = do_stretch(cbticks, *stretch_args)
-        else:
-            cbticks = cb.get_ticks()
-        cond = (vmin <= cbticks) * (cbticks <= vmax)
+        if cbticks is None and st.stretch == 'log':
+            cbticks, cbticklabels = logcbticks(10**st.vmin, 10**st.vmax)
+        cbticks = cb.get_ticks() if cbticks is None else st.do(cbticks)
+        cond = (st.vmin <= cbticks) * (cbticks <= st.vmax)
         cbticks = cbticks[cond]
         cb.set_ticks(cbticks)
-        if cbticklabels is not None:
-            cbticklabels = np.array(cbticklabels)[cond]
+        if cbticklabels is None:
+            cbticklabels = [f'{t:{cbformat[1:]}}' for t in st.undo(cbticks)]
         else:
-            t = undo_stretch(cbticks, *stretch_args)
-            cbticklabels = [f'{d:{cbformat[1:]}}' for d in t]
+            cbticklabels = np.array(cbticklabels)[cond]
         cb.set_ticklabels(cbticklabels)
 
     def add_color(self,
-                  stretch: str = 'linear',
-                  stretchscale: float | None = None,
-                  stretchpower: float = 0.5,
                   show_cbar: bool = True,
                   cblabel: str | None = None,
                   cbformat: float = '%.1e',
@@ -851,12 +840,9 @@ class PlotAstroData(AstroFrame):
                   cbticklabels: list[str] | None = None,
                   cblocation: str = 'right',
                   **kwargs) -> None:
-        """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.
+        """Use Axes.pcolormesh of matplotlib. kwargs must include the arguments of AstroData to specify the data to be plotted. kwargs may include the arguments for Stretcher (stretch, stretchscale, and stretchpower) to specify the stretch parameters. 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', '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**stretchpower / stretchpower. 1 means the linear scale, while 0 means the logarithmic scale. Defaults to 0.5.
             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'.
@@ -865,7 +851,8 @@ class PlotAstroData(AstroFrame):
             cblocation (str, optional): 'left', 'top', 'left', 'right'. Only for 2D images. Defaults to 'right'.
         """
         self._kw = {'cmap': 'cubehelix', 'alpha': 1,
-                    'edgecolors': 'none', 'zorder': 1}
+                    'edgecolors': 'none', 'zorder': 1,
+                    'vmin': None, 'vmax': None}
         c, x, y, v, beam, sigma, bunit, _kw, beam_kwargs, singlepix \
             = self._map_init(kwargs)
         if singlepix:
@@ -874,11 +861,16 @@ class PlotAstroData(AstroFrame):
 
         if cblabel is None:
             cblabel = bunit
-        if stretchscale is None:
-            stretchscale = sigma
 
-        stretch_args = (stretch, stretchscale, stretchpower)
-        c = set_minmax(c, *stretch_args, sigma, _kw)
+        stretch_params = {}
+        for k in ['stretch', 'stretchscale', 'stretchpower']:
+            if k in _kw:
+                stretch_params[k] = _kw.pop(k)
+        st = Stretcher(vmin=_kw['vmin'], vmax=_kw['vmax'],
+                       sigma=sigma, **stretch_params)
+        c, cmin, cmax = st.set_minmax(c)
+        _kw['vmin'] = cmin
+        _kw['vmax'] = cmax
         c = self.vskipfill(c, v)
         if type(self.channelnumber) is int:
             c = [c[self.channelnumber]]
@@ -889,8 +881,7 @@ class PlotAstroData(AstroFrame):
                 p[ch] = pnow
         for ch in self.bottomleft:
             self._set_colorbar(p, ch, show_cbar, cblabel, cbformat,
-                               cbticks, cbticklabels, cblocation,
-                               *stretch_args, _kw['vmin'], _kw['vmax'])
+                               cbticks, cbticklabels, cblocation, st)
         self.add_beam(beam=beam, **beam_kwargs)
 
     def add_contour(self,
@@ -982,33 +973,34 @@ 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: list[float, float, float] = [1, 1, 1],
+                stretchpower: list[float, float, float] = [0.5, 0.5, 0.5],
                 **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.
+        """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 the arguments for Stretcher (stretch, stretchscale, and stretchpower; three-element array for each) to specify the stretch parameters. 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', '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**stretchpower / stretchpower. 1 means the linear scale, while 0 means the logarithmic scale. Defaults to 1.
+            stretchpower (float, optional): Color scale is data**stretchpower / stretchpower. 1 means the linear scale, while 0 means the logarithmic scale. Defaults to 0.5.
         """
         from PIL import Image
 
-        self._kw = {}
+        self._kw = {'vmin': [None] * 3, 'vmax': [None] * 3}
         c, x, y, v, beam, sigma, _, _kw, beam_kwargs, singlepix \
             = self._map_init(kwargs)
         if singlepix:
             print('No pixel size. Skip add_rgb.')
             return
 
-        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
 
-        for i, (cmin, cmax) in enumerate(zip(_kw['vmin'], _kw['vmax'])):
-            if cmax > cmin:
-                c[i] = (c[i] - cmin) / (cmax - cmin) * 255
+        st = Stretcher(stretch, stretchscale, stretchpower,
+                       _kw['vmin'], _kw['vmax'], sigma)
+        c, cmin, cmax = st.set_minmax(c)
+        for i in range(st.n):
+            if cmax[i] > cmin[i]:
+                c[i] = (c[i] - cmin[i]) / (cmax[i] - cmin[i]) * 255
             c[i] = self.vskipfill(c[i], v)
         size = np.shape(c[0][0])[::-1]
         c = np.moveaxis(c, 1, 0)[:, :, ::-self.ydir, ::-self.xdir]

{plotastrodata-1.8.11 → plotastrodata-1.8.12/plotastrodata.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plotastrodata
-Version: 1.8.11
+Version: 1.8.12
 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