text:
# -*- coding: utf-8 -*-

import time
import emcee
import numpy as np
from astropy.io import fits
import scipy.constants as sc
import warnings

warnings.filterwarnings("ignore")

class rotationmap:
    """
    Read in the velocity maps and initialize the class. To make the fitting
    quicker, we can clip the cube to a smaller region, or downsample the
    image to get a quicker initial fit.

    Args:
        path (str): Relative path to the rotation map you want to fit.
        uncertainty (optional[srt]): Relative path to the map of v0
            uncertainties. Must be a FITS file with the same shape as the
            data. If nothing is specified, will tried to find an uncerainty
            file following the ``bettermoments`` format. If this is not found,
            it will assume a 10% on all pixels.
        clip (optional[float]): If specified, clip the data down to a
            square with sides 2 x clip centered on the image center.
        downsample (optional[int]): Downsample the image by this factor for
            quicker fitting. For example, using ``downsample=4`` on an image
            which is 1000 x 1000 pixels will reduce it to 250 x 250 pixels.
            If you use ``downsample='beam'`` it will sample roughly
            spatially independent pixels using the beam major axis as the
            spacing.
        x0 (optional[float]): Recenter the image to this offset.
        y0 (optional[float]): Recenter the image to this offset.
        unit (optional[str]): Unit of the input data cube. By deafult
            assumed to be [m/s].
    """

    msun = 1.988e30
    fwhm = 2. * np.sqrt(2 * np.log(2))
    disk_coords_niter = 5
    priors = {}

    def __init__(self, path, uncertainty=None, clip=None, downsample=None,
                 x0=0.0, y0=0.0, unit='m/s'):
        # Read in the data and position axes.
        self.path = path
        self.data = np.squeeze(fits.getdata(path))
        self.header = fits.getheader(path)
        if uncertainty is not None:
            self.error = abs(np.squeeze(fits.getdata(uncertainty)))
        else:
            try:
                uncertainty = '_'.join(self.path.split('_')[:-1])
                uncertainty += '_d' + self.path.split('_')[-1]
                print("Assuming uncertainties in {}".format(uncertainty))
                self.error = abs(np.squeeze(fits.getdata(uncertainty)))
            except FileNotFoundError:
                print("No uncertainties found, assuming uncertainties of 10%.")
                print("Change this at any time with rotationmap.error.")
                self.error = 0.1 * abs((self.data - np.nanmedian(self.data)))
        self.error = np.where(np.isnan(self.error), 0.0, self.error)

        # Convert the data to [km/s].
        if unit.lower() == 'm/s':
            self.data *= 1e-3
            self.error *= 1e-3
        elif unit.lower() != 'km/s':
            raise ValueError("unit must me 'm/s' or 'km/s'.")

        # Read the position axes.
        self.xaxis = self._read_position_axis(a=1)
        self.yaxis = self._read_position_axis(a=2)
        self.dpix = abs(np.diff(self.xaxis)).mean()

        # Recenter the image if provided.
        if (x0 != 0.0) or (y0 != 0.0):
            self._shift_center(dx=x0, dy=y0)

        # Beam parameters.
        self._readbeam()

        # Clip and downsample the cube to speed things up.
        if downsample is not None:
            if downsample == 'beam':
                downsample = int(np.ceil(self.bmaj / self.dpix))
            self._downsample_cube(downsample)
            self.dpix = abs(np.diff(self.xaxis)).mean()
        if clip is not None:
            self._clip_cube(clip)
        self.nypix = self.yaxis.size
        self.nxpix = self.xaxis.size
        self.mask = np.isfinite(self.data)

        # Estimate the systemic velocity.
        self.vlsr = np.nanmedian(self.data)

        # Set priors.
        self._set_default_priors()

        # Set the defaults for 'shadowed' emission surfaces.
        self.shadowed = False
        self.shadowed_extend = 1.5
        self.shadowed_oversample = 2.0
        self.shadowed_method = 'nearest'

    def fit_map(self, p0, params, r_min=None, r_max=None, optimize=True,
                nwalkers=None, nburnin=300, nsteps=100, scatter=1e-3,
                plots=None, returns=None, pool=None, emcee_kwargs=None,
                niter=1, shadowed=False):
        """
        Fit a rotation profile to the data. Note that for a disk with
        a non-zero height, the sign of the inclination dictates the direction
        of the tilt: a positive tilt rotates the disk about around the x-axis
        such that the southern side of the disk is closer to the observer. The
        same definition is used for the warps.

        The function must be provided with a dictionary of parameters,
        ``params`` where key is the parameter and its value is either the fixed
        value this will take,

            ``params['PA'] = 45.0``

        would fix position angle to 45 degrees. Or, if an integeter, the index
        in the starting positions, ``p0``, such that,

            ``params['mstar'] = 0``

        would mean that the first value in ``p0`` is the guess for ``'mstar'``.

        For a list of the geometrical parameters, included flared emission
        surfaces or warps, see :func:`disk_coords`. In addition to these
        geometrical properties ``params`` also requires ``'mstar'``,
        ``'vlsr'``, ``'dist'``.

        To include a spatial convolution of the model rotation map you may also
        set ``params['beam']=True``, however this is only really necessary for
        low spatial resolution data, or for rotation maps made via the
        intensity-weighted average velocity.

        .. _Rosenfeld et al. (2013): https://ui.adsabs.harvard.edu/abs/2013ApJ...774...16R/abstract

        Args:
            p0 (list): List of the free parameters to fit.
            params (dictionary): Dictionary of the model parameters.
            r_min (optional[float]): Inner radius to fit in [arcsec].
            r_max (optional[float]): Outer radius to fit in [arcsec].
            optimize (optional[bool]): Use ``scipy.optimize`` to find the
                ``p0`` values which maximize the likelihood. Better results
                will likely be found. Note that for the masking the default
                ``p0`` and ``params`` values are used for the deprojection, or
                those foundvfrom the optimization. If this results in a poor
                initial mask, try with ``optimise=False`` or with a ``niter``
                value larger than 1.
            nwalkers (optional[int]): Number of walkers to use for the MCMC.
            nburnin (optional[int]): Number of steps to discard for burn-in.
            nsteps (optional[int]): Number of steps to use to sample the
                posterior distributions.
            scatter (optional[float]): Scatter used in distributing walker
                starting positions around the initial ``p0`` values.
            plots (optional[list]): List of the diagnostic plots to make. This
                can include ``'mask'``, ``'walkers'``, ``'corner'``,
                ``'bestfit'``, ``'residual'``, or ``'none'`` if no plots are to
                be plotted. By default, all are plotted.
            returns (optional[list]): List of items to return. Can contain
                ``'samples'``, ``'percentiles'``, ``'dict'`` or ``'none'``. By
                default only ``'percentiles'`` are returned.
            pool (optional): An object with a `map` method.
            emcee_kwargs (Optional[dict]): Dictionary to pass to the emcee
                ``EnsembleSampler``.
            niter (optional[int]): Number of iterations to perform using the
                median PDF values from the previous MCMC run as starting
                positions. This is probably only useful if you have no idea
                about the starting positions for the emission surface or if you
                want to remove walkers stuck in local minima.

        Returns:
            to_return (list): Depending on the returns list provided.
                ``'samples'`` will be all the samples of the posteriors (with
                the burn-in periods removed). ``'percentiles'`` will return the
                16th, 50th and 84th percentiles of each posterior functions.
                ``'dict'`` will return a dictionary of the median parameters
                which can be directly input to other functions.
        """

        # Check the dictionary. May need some more work.
        if r_min is not None:
            if 'r_min' in params.keys():
                print("Found `r_min` in `params`. Overwriting value.")
            params['r_min'] = r_min
        if r_max is not None:
            if 'r_max' in params.keys():
                print("Found `r_max` in `params`. Overwriting value.")
            params['r_max'] = r_max
        params = self.verify_params_dictionary(params)
        self.shadowed = shadowed

        # Generate the mask for fitting based on the params.
        p0 = np.squeeze(p0).astype(float)
        temp = rotationmap._populate_dictionary(p0, params)
        self.ivar = self._calc_ivar(temp)

        # Check what the parameters are.
        labels = rotationmap._get_labels(params)
        labels_raw = []
        for label in labels:
            label_raw = label.replace('$', '').replace('{', '')
            label_raw = label_raw.replace(r'\rm ', '').replace('}', '')
            labels_raw += [label_raw]
        if len(labels) != len(p0):
            raise ValueError("Mismatch in labels and p0. Check for integers.")
        print("Assuming:\n\tp0 = [%s]." % (', '.join(labels_raw)))

        # Run an initial optimization using scipy.minimize. Recalculate the
        # inverse variance mask.
        if optimize:
            p0 = self._optimize_p0(p0, params)

        # Make the mask for fitting.
        temp = rotationmap._populate_dictionary(p0, params)
        temp = self.verify_params_dictionary(temp)

        self.ivar = self._calc_ivar(temp)

        # Set up and run the MCMC with emcee.
        time.sleep(0.5)
        nwalkers = 2 * p0.size if nwalkers is None else nwalkers
        emcee_kwargs = {} if emcee_kwargs is None else emcee_kwargs
        emcee_kwargs['scatter'], emcee_kwargs['pool'] = scatter, pool
        for n in range(int(niter)):

            # Run the sampler.
            sampler = self._run_mcmc(p0=p0, params=params, nwalkers=nwalkers,
                                     nburnin=nburnin, nsteps=nsteps,
                                     **emcee_kwargs)
            if type(params['PA']) is int:
                sampler.chain[:, :, params['PA']] %= 360.0

            # Split off the samples.
            samples = sampler.chain[:, -int(nsteps):]
            samples = samples.reshape(-1, samples.shape[-1])
            p0 = np.median(samples, axis=0)
            medians = rotationmap._populate_dictionary(p0, params)
            medians = self.verify_params_dictionary(medians)

        # Diagnostic plots.
        if plots is None:
            plots = ['mask', 'walkers', 'corner', 'bestfit', 'residual']
        plots = np.atleast_1d(plots)
        if 'none' in plots:
            plots = []
        if 'mask' in plots:
            self.plot_data(ivar=self.ivar)
        if 'walkers' in plots:
            rotationmap._plot_walkers(sampler.chain.T, nburnin, labels)
        if 'corner' in plots:
            rotationmap._plot_corner(samples, labels)
        if 'bestfit' in plots:
            self._plot_bestfit(medians, ivar=self.ivar)
        if 'residual' in plots:
            self._plot_residual(medians, ivar=self.ivar)

        # Generate the output.
        if returns is None:
            returns = ['percentiles']
        returns = np.atleast_1d(returns)
        if 'none' in returns:
            return None
        to_return = []
        if 'samples' in returns:
            to_return += [samples]
        if 'lnprob' in returns:
            to_return += [sampler.lnprobability[nburnin:]]
        if 'percentiles' in returns:
            to_return += [np.percentile(samples, [16, 50, 84], axis=0)]
        if 'dict' in returns:
            to_return += [medians]
        self.shadowed = False
        return to_return if len(to_return) > 1 else to_return[0]

    def set_prior(self, param, args, type='flat'):
        """
        Set the prior for the given parameter. There are two types of priors
        currently usable, ``'flat'`` which requires ``args=[min, max]`` while
        for ``'gaussian'`` you need to specify ``args=[mu, sig]``.

        Args:
            param (str): Name of the parameter.
            args (list): Values to use depending on the type of prior.
            type (optional[str]): Type of prior to use.
        """
        type = type.lower()
        if type not in ['flat', 'gaussian']:
            raise ValueError("type must be 'flat' or 'gaussian'.")
        if type == 'flat':
            def prior(p):
                if not min(args) <= p <= max(args):
                    return -np.inf
                return np.log(1.0 / (args[1] - args[0]))
        else:
            def prior(p):
                lnp = np.exp(-0.5 * ((args[0] - p) / args[1])**2)
                return np.log(lnp / np.sqrt(2. * np.pi) / args[1])
        rotationmap.priors[param] = prior

    def disk_coords(self, x0=0.0, y0=0.0, inc=0.0, PA=0.0, z0=0.0, psi=0.0,
                    r_cavity=0.0, r_taper=None, q_taper=None, w_i=0.0, w_r=1.0,
                    w_t=0.0, frame='cylindrical', **_):
        r"""
        Get the disk coordinates given certain geometrical parameters and an
        emission surface. The emission surface is most simply described as a
        power law profile,

        .. math::

            z(r) = z_0 \times \left(\frac{r}{1^{\prime\prime}}\right)^{\psi}

        where ``z0`` and ``psi`` can be provided by the user. With the increase
        in spatial resolution afforded by interferometers such as ALMA there
        are a couple of modifications that can be used to provide a better
        match to the data.

        An inner cavity can be included with the ``r_cavity`` argument which
        makes the transformation:

        .. math::

            \tilde{r} = {\rm max}(0, r - r_{\rm cavity})

        Note that the inclusion of a cavity will mean that other parameters,
        such as ``z0``, would need to change as the radial axis has effectively
        been shifted.

        To account for the drop in emission surface in the outer disk where the
        gas surface density decreases there are two descriptions. The preferred
        way is to include an exponential taper to the power law profile,

        .. math::

            z_{\rm tapered}(r) = z(r) \times \exp\left( -\left[
            \frac{r}{r_{\rm taper}} \right]^{q_{\rm taper}} \right)

        where both ``r_taper`` and ``q_taper`` values must be set.

        We can also include a warp which is parameterized by,

        .. math::

            z_{\rm warp}(r,\, t) = r \times \tan \left(w_i \times \exp\left(-
            \frac{r^2}{2 w_r^2} \right) \times \sin(t - w_t)\right)

        where ``w_i`` is the inclination in [radians] describing the warp at
        the disk center. The width of the warp is given by ``w_r`` [arcsec] and
        ``w_t`` in [radians] is the angle of nodes (where the warp is zero),
        relative to the position angle of the disk, measured east of north.

        .. WARNING::

            The use of warps is largely untested. Use with caution!
            If you are using a warp, increase the number of iterations
            for the inference through ``self.disk_coords_niter`` (by default at
            5). For high inclinations, also set ``shadowed=True``.

        Args:
            x0 (optional[float]): Source right ascension offset [arcsec].
            y0 (optional[float]): Source declination offset [arcsec].
            inc (optional[float]): Source inclination [degrees].
            PA (optional[float]): Source position angle [degrees]. Measured
                between north and the red-shifted semi-major axis in an
                easterly direction.
            z0 (optional[float]): Aspect ratio at 1" for the emission surface.
                To get the far side of the disk, make this number negative.
            psi (optional[float]): Flaring angle for the emission surface.
            z1 (optional[float]): Correction term for ``z0``.
            phi (optional[float]): Flaring angle correction term for the
                emission surface.
            r_cavity (optional[float]): Outer radius of a cavity. Within this
                region the emission surface is taken to be zero.
            w_i (optional[float]): Warp inclination in [degrees] at the disk
                center.
            w_r (optional[float]): Scale radius of the warp in [arcsec].
            w_t (optional[float]): Angle of nodes of the warp in [degrees].
            frame (optional[str]): Frame of reference for the returned
                coordinates. Either ``'cartesian'`` or ``'cylindrical'``.

        Returns:
            array, array, array: Three coordinate arrays with ``(r, phi, z)``,
            in units of [arcsec], [radians], [arcsec], if
            ``frame='cylindrical'`` or ``(x, y, z)``, all in units of [arcsec]
            if ``frame='cartesian'``.
        """

        # Check the input variables.

        frame = frame.lower()
        if frame not in ['cylindrical', 'cartesian']:
            raise ValueError("frame must be 'cylindrical' or 'cartesian'.")

        # Apply the inclination concention.

        inc = inc if inc < 90.0 else inc - 180.0

        # Check that the necessary pairs are provided.

        msg = "Must specify either both or neither of `{}` and `{}`."
        if (r_taper is not None) != (q_taper is not None):
            raise ValueError(msg.format('r_taper', 'q_taper'))

        # Set the defaults.

        z0 = 0.0 if z0 is None else z0
        psi = 1.0 if psi is None else psi
        r_cavity = 0.0 if r_cavity is None else r_cavity
        r_taper = np.inf if r_taper is None else r_taper
        q_taper = 1.0 if q_taper is None else q_taper

        # Define the emission surface, z_func, and the warp function, w_func.

        def z_func(r_in):
            r = np.clip(r_in - r_cavity, a_min=0.0, a_max=None)
            z = z0 * np.power(r, psi) * np.exp(-np.power(r / r_taper, q_taper))
            return np.clip(z, a_min=0.0, a_max=None)

        def w_func(r_in, t):
            r = np.clip(r_in - r_cavity, a_min=0.0, a_max=None)
            warp = np.radians(w_i) * np.exp(-0.5 * (r / w_r)**2)
            return r * np.tan(warp * np.sin(t - np.radians(w_t)))

        # Calculate the pixel values.

        if self.shadowed:
            coords = self._get_shadowed_coords(x0, y0, inc, PA, z_func, w_func)
        else:
            coords = self._get_flared_coords(x0, y0, inc, PA, z_func, w_func)
        if frame == 'cylindrical':
            return coords
        r, t, z = coords
        return r * np.cos(t), r * np.sin(t), z

    def plot_data(self, levels=None, ivar=None, return_fig=False):
        """
        Plot the first moment map. By default will clip the velocity contours
        such that the velocities around the systemic velocity are highlighted.

        Args:
            levels (optional[list]): List of contour levels to use.
            ivar (optional[ndarray]): Inverse variances for each pixel. Will
                draw a solid contour around the regions with finite ``ivar``
                values and fill regions not considered.
            return_fig (optional[bool]): Return the figure.
        Returns:
            fig (Matplotlib figure): If ``return_fig`` is ``True``. Can access
                the axes through ``fig.axes`` for additional plotting.
        """
        import matplotlib.pyplot as plt
        fig, ax = plt.subplots()
        if levels is None:
            levels = np.nanpercentile(self.data, [2, 98]) - self.vlsr
            levels = max(abs(levels[0]), abs(levels[1]))
            levels = self.vlsr + np.linspace(-levels, levels, 30)
        im = ax.contourf(self.xaxis, self.yaxis, self.data, levels,
                         cmap=rotationmap.colormap(), extend='both', zorder=-9)
        cb = plt.colorbar(im, pad=0.03, format='%.2f')
        cb.minorticks_on()
        cb.set_label(r'${\rm v_{0} \quad (km\,s^{-1})}$',
                     rotation=270, labelpad=15)
        if ivar is not None:
            ax.contour(self.xaxis, self.yaxis, ivar,
                       [0.0], colors='k')
            ax.contourf(self.xaxis, self.yaxis, ivar,
                        [-1.0, 0.0], colors='k', alpha=0.5)
        self._gentrify_plot(ax)
        if return_fig:
            return fig

    # -- MCMC Functions -- #

    def _optimize_p0(self, theta, params, **kwargs):
        """Optimize the initial starting positions."""
        from scipy.optimize import minimize

        # TODO: implement bounds.
        # TODO: cycle through parameters one at a time for a better fit.

        # Negative log-likelihood function.
        def nlnL(theta):
            return -self._ln_probability(theta, params)

        method = kwargs.pop('method', 'TNC')
        options = kwargs.pop('options', {})
        options['maxiter'] = options.pop('maxiter', 10000)
        options['maxfun'] = options.pop('maxfun', 10000)
        options['ftol'] = options.pop('ftol', 1e-3)

        res = minimize(nlnL, x0=theta, method=method, options=options)
        if res.success:
            theta = res.x
            print("Optimized starting positions:")
        else:
            print("WARNING: scipy.optimize did not converge.")
            print("Starting positions:")
        print('\tp0 =', ['%.2e' % t for t in theta])
        time.sleep(.3)
        return theta

    def _run_mcmc(self, p0, params, nwalkers, nburnin, nsteps, **kwargs):
        """Run the MCMC sampling. Returns the sampler."""
        p0 = self._random_p0(p0, kwargs.pop('scatter', 1e-3), nwalkers)
        sampler = emcee.EnsembleSampler(nwalkers, p0.shape[1],
                                        self._ln_probability,
                                        args=[params, np.nan],
                                        **kwargs)
        if emcee.__version__ >= '3':
            progress = kwargs.pop('progress', True)
            sampler.run_mcmc(p0, nburnin + nsteps, progress=progress)
        else:
            sampler.run_mcmc(p0, nburnin + nsteps)
        return sampler

    @staticmethod
    def _random_p0(p0, scatter, nwalkers):
        """Get the starting positions."""
        p0 = np.squeeze(p0)
        dp0 = np.random.randn(nwalkers * len(p0)).reshape(nwalkers, len(p0))
        dp0 = np.where(p0 == 0.0, 1.0, p0)[None, :] * (1.0 + scatter * dp0)
        return np.where(p0[None, :] == 0.0, dp0 - 1.0, dp0)

    def _ln_likelihood(self, params):
        """Log-likelihood function. Simple chi-squared likelihood."""
        model = self._make_model(params) * 1e-3
        lnx2 = np.where(self.mask, np.power((self.data - model), 2), 0.0)
        lnx2 = -0.5 * np.sum(lnx2 * self.ivar)
        return lnx2 if np.isfinite(lnx2) else -np.inf

    def _ln_probability(self, theta, *params_in):
        """Log-probablility function."""
        model = rotationmap._populate_dictionary(theta, params_in[0])
        lnp = self._ln_prior(model)
        if np.isfinite(lnp):
            return lnp + self._ln_likelihood(model)
        return -np.inf

    def _set_default_priors(self):
        """Set the default priors."""

        # Basic Geometry.
        self.set_prior('x0', [-0.5, 0.5], 'flat')
        self.set_prior('y0', [-0.5, 0.5], 'flat')
        self.set_prior('inc', [-90.0, 90.0], 'flat')
        self.set_prior('PA', [-360.0, 360.0], 'flat')
        self.set_prior('mstar', [0.1, 5.0], 'flat')
        self.set_prior('vlsr', [np.nanmin(self.data) * 1e3,
                                np.nanmax(self.data) * 1e3], 'flat')

        # Emission Surface
        self.set_prior('z0', [0.0, 5.0], 'flat')
        self.set_prior('psi', [0.0, 5.0], 'flat')
        self.set_prior('r_cavity', [0.0, 1e30], 'flat')
        self.set_prior('r_taper', [0.0, 1e30], 'flat')
        self.set_prior('q_taper', [0.0, 5.0], 'flat')

        # Warp
        self.set_prior('w_i', [-90.0, 90.0], 'flat')
        self.set_prior('w_r', [0.0, self.xaxis.max()], 'flat')
        self.set_prior('w_t', [-180.0, 180.0], 'flat')

        # Velocity Profile
        self.set_prior('vp_100', [0.0, 1e4], 'flat')
        self.set_prior('vp_q', [-2.0, 0.0], 'flat')
        self.set_prior('vp_rtaper', [0.0, 1e30], 'flat')
        self.set_prior('vp_qtaper', [0.0, 5.0], 'flat')
        self.set_prior('vr_100', [-1e3, 1e3], 'flat')
        self.set_prior('vr_q', [-2.0, 2.0], 'flat')

    def _ln_prior(self, params):
        """Log-priors."""
        lnp = 0.0
        for key in params.keys():
            if key in rotationmap.priors.keys():
                lnp += rotationmap.priors[key](params[key])
                if not np.isfinite(lnp):
                    return lnp
        return lnp

    def _calc_ivar(self, params):
        """Calculate the inverse variance including radius mask."""

        # Check the error array is the same shape as the data.
        try:
            assert self.error.shape == self.data.shape
        except AttributeError:
            self.error = self.error * np.ones(self.data.shape)

        # Deprojected coordinates.
        r, t = self.disk_coords(**params)[:2]
        t = abs(t) if params['abs_PA'] else t

        # Radial mask.
        mask_r = np.logical_and(r >= params['r_min'], r <= params['r_max'])
        mask_r = ~mask_r if params['exclude_r'] else mask_r

        # Azimuthal mask.
        mask_t = np.logical_and(t >= params['PA_min'], t <= params['PA_max'])
        mask_t = ~mask_t if params['exclude_PA'] else mask_t

        # Finite value mask.
        mask_f = np.logical_and(np.isfinite(self.data), self.error > 0.0)

        # Velocity mask.
        v_min, v_max = params['v_min'], params['v_max']
        mask_v = np.logical_and(self.data >= v_min, self.data <= v_max)
        mask_v = ~mask_v if params['exclude_v'] else mask_v

        # Combine.
        mask = np.logical_and(np.logical_and(mask_v, mask_f),
                              np.logical_and(mask_r, mask_t))
        return np.where(mask, np.power(self.error, -2.0), 0.0)

    @staticmethod
    def _get_labels(params):
        """Return the labels of the parameters to fit."""
        idxs, labs = [], []
        for k in params.keys():
            if isinstance(params[k], int):
                if not isinstance(params[k], bool):
                    idxs.append(params[k])
                    try:
                        idx = k.index('_') + 1
                        label = k[:idx] + '{{' + k[idx:] + '}}'
                    except ValueError:
                        label = k
                    label = r'${{\rm {}}}$'.format(label)
                    labs.append(label)
        return np.array(labs)[np.argsort(idxs)]

    @staticmethod
    def _populate_dictionary(theta, dictionary_in):
        """Populate the dictionary of free parameters."""
        dictionary = dictionary_in.copy()
        for key in dictionary.keys():
            if isinstance(dictionary[key], int):
                if not isinstance(dictionary[key], bool):
                    dictionary[key] = theta[dictionary[key]]
        return dictionary

    def verify_params_dictionary(self, params):
        """
        Check that the minimum number of parameters are provided for the
        fitting. For non-essential parameters a default value is set.
        """

        # Basic geometrical properties.
        params['x0'] = params.pop('x0', 0.0)
        params['y0'] = params.pop('y0', 0.0)
        params['dist'] = params.pop('dist', 100.0)
        params['vlsr'] = params.pop('vlsr', self.vlsr)
        if params.get('inc', None) is None:
            raise KeyError("Must provide `'inc'`.")
        if params.get('PA', None) is None:
            raise KeyError("Must provide `'PA'`.")

        # Rotation profile.
        has_vp_100 = False if params.get('vp_100') is None else True
        has_mstar = False if params.get('mstar') is None else True
        if has_mstar and has_vp_100:
            raise KeyError("Only provide either `'mstar'` or `'vp_100'`.")
        if not has_mstar and not has_vp_100:
            raise KeyError("Must provide either `'mstar'` or `'vp_100'`.")
        if has_mstar:
            params['vfunc'] = self._proj_vkep
        else:
            params['vfunc'] = self._proj_vpow
        params['vp_q'] = params.pop('vp_q', -0.5)
        params['vp_rtaper'] = params.pop('vp_rtaper', 1e10)
        params['vp_qtaper'] = params.pop('vp_qtaper', 1.0)
        params['vr_100'] = params.pop('vr_100', 0.0)
        params['vr_q'] = params.pop('vr_q', 0.0)

        # Flared emission surface.
        params['z0'] = params.pop('z0', 0.0)
        params['psi'] = params.pop('psi', 1.0)
        params['r_taper'] = params.pop('r_taper', 1e10)
        params['q_taper'] = params.pop('q_taper', 1.0)
        params['r_cavity'] = params.pop('r_cavity', 0.0)

        # Warp parameters.
        params['w_i'] = params.pop('w_i', 0.0)
        params['w_r'] = params.pop('w_r', 1.0)
        params['w_t'] = params.pop('w_t', 0.0)
        params['shadowed'] = params.pop('shadowed', False)

        # Masking parameters.
        params['r_min'] = params.pop('r_min', 0.0)
        params['r_max'] = params.pop('r_max', 1e10)
        params['exclude_r'] = params.pop('exclude_r', False)
        params['PA_min'] = params.pop('PA_min', -np.pi)
        params['PA_max'] = params.pop('PA_max', np.pi)
        params['exclude_PA'] = params.pop('exclude_PA', False)
        params['abs_PA'] = params.pop('abs_PA', False)
        params['v_min'] = params.pop('v_min', np.nanmin(self.data))
        params['v_max'] = params.pop('v_max', np.nanmax(self.data))
        params['exclude_v'] = params.pop('exclude_v', False)

        if params['r_min'] >= params['r_max']:
            raise ValueError("`r_max` must be greater than `r_min`.")
        if params['PA_min'] >= params['PA_max']:
            raise ValueError("`PA_max` must be great than `PA_min`.")
        if params['v_min'] >= params['v_max']:
            raise ValueError("`v_max` must be greater than `v_min`.")

        # Beam convolution.
        params['beam'] = False

        return params

    def evaluate_models(self, samples, params, draws=0.5,
                        collapse_func=np.mean, coords_only=False):
        """
        Evaluate models based on the samples provided and the parameter
        dictionary. If ``draws`` is an integer, it represents the number of
        random draws from ``samples`` which are then averaged to return the
        model. Alternatively ``draws`` can be a float between 0 and 1,
        representing the percentile of the samples to use to evaluate the model
        at.

        Args:
            samples (ndarray): An array of samples returned from ``fit_map``.
            params (dict): The parameter dictionary passed to ``fit_map``.
            draws (Optional[int/float]): If an integer, describes the number of
                random draws averaged to form the returned model. If a float,
                represents the percentile used from the samples. Must be
                between 0 and 1 if a float.
            collapse_func (Optional[callable]): How to collapse the random
                number of samples. Must be a function which allows an ``axis``
                argument (as with most Numpy functions).
            coords_only (Optional[bool]): Return the deprojected coordinates
                rather than the v0 model. Default is False.

        Returns:
            model (ndarray): The sampled model, either the v0 model, or, if
                ``coords_only`` is True, the deprojected cylindrical
                coordinates, (r, t, z).
        """

        # Check the input.
        nparam = np.sum([isinstance(params[k], int) for k in params.keys()])
        if samples.shape[1] != nparam:
            warning = "Invalid number of free parameters in 'samples': {:d}."
            raise ValueError(warning.format(nparam))
        if not callable(collapse_func):
            raise ValueError("'collapse_func' must be callable.")
        verified_params = self.verify_params_dictionary(params.copy())

        # Avearge over a random draw of models.
        if isinstance(draws, int):
            models = []
            for idx in np.random.randint(0, samples.shape[0], draws):
                tmp = self._populate_dictionary(samples[idx], verified_params)
                if coords_only:
                    models += [self.disk_coords(**tmp)]
                else:
                    models += [self._make_model(tmp)]
            return collapse_func(models, axis=0)

        # Take a percentile of the samples.
        elif isinstance(draws, float):
            tmp = np.percentile(samples, draws, axis=0)
            tmp = self._populate_dictionary(tmp, verified_params)
            if coords_only:
                return self.disk_coords(**tmp)
            else:
                return self._make_model(tmp)

        else:
            raise ValueError("'draws' must be a float or integer.")

    # -- Deprojection Functions -- #

    @staticmethod
    def _rotate_coords(x, y, PA):
        """Rotate (x, y) by PA [deg]."""
        x_rot = y * np.cos(np.radians(PA)) + x * np.sin(np.radians(PA))
        y_rot = x * np.cos(np.radians(PA)) - y * np.sin(np.radians(PA))
        return x_rot, y_rot

    @staticmethod
    def _deproject_coords(x, y, inc):
        """Deproject (x, y) by inc [deg]."""
        return x, y / np.cos(np.radians(inc))

    def _get_cart_sky_coords(self, x0, y0):
        """Return cartesian sky coordinates in [arcsec, arcsec]."""
        return np.meshgrid(self.xaxis - x0, self.yaxis - y0)

    def _get_midplane_cart_coords(self, x0, y0, inc, PA):
        """Return cartesian coordaintes of midplane in [arcsec, arcsec]."""
        x_sky, y_sky = self._get_cart_sky_coords(x0, y0)
        x_rot, y_rot = self._rotate_coords(x_sky, y_sky, PA)
        return rotationmap._deproject_coords(x_rot, y_rot, inc)

    def _get_midplane_polar_coords(self, x0, y0, inc, PA):
        """Return the polar coordinates of midplane in [arcsec, radians]."""
        x_mid, y_mid = self._get_midplane_cart_coords(x0, y0, inc, PA)
        return np.hypot(y_mid, x_mid), np.arctan2(y_mid, x_mid)

    def _get_flared_coords(self, x0, y0, inc, PA, z_func, w_func):
        """Return cyclindrical coords of surface in [arcsec, rad, arcsec]."""
        x_mid, y_mid = self._get_midplane_cart_coords(x0, y0, inc, PA)
        r_tmp, t_tmp = np.hypot(x_mid, y_mid), np.arctan2(y_mid, x_mid)
        for _ in range(self.disk_coords_niter):
            z_tmp = z_func(r_tmp) + w_func(r_tmp, t_tmp)
            y_tmp = y_mid + z_tmp * np.tan(np.radians(inc))
            r_tmp = np.hypot(y_tmp, x_mid)
            t_tmp = np.arctan2(y_tmp, x_mid)
        return r_tmp, t_tmp, z_func(r_tmp)

    def _get_shadowed_coords(self, x0, y0, inc, PA, z_func, w_func):
        """Return cyclindrical coords of surface in [arcsec, rad, arcsec]."""

        # Make the disk-frame coordinates.
        extend = self.shadowed_extend
        oversample = self.shadowed_oversample
        diskframe_coords = self._get_diskframe_coords(extend, oversample)
        xdisk, ydisk, rdisk, tdisk = diskframe_coords
        zdisk = z_func(rdisk) + w_func(rdisk, tdisk)

        # Incline the disk.
        inc = np.radians(inc)
        x_dep = xdisk
        y_dep = ydisk * np.cos(inc) - zdisk * np.sin(inc)

        # Remove shadowed pixels.
        if inc < 0.0:
            y_dep = np.maximum.accumulate(y_dep, axis=0)
        else:
            y_dep = np.minimum.accumulate(y_dep[::-1], axis=0)[::-1]

        # Rotate and the disk.
        x_rot, y_rot = self._rotate_coords(x_dep, y_dep, PA)
        x_rot, y_rot = x_rot + x0, y_rot + y0

        # Grid the disk.
        from scipy.interpolate import griddata
        disk = (x_rot.flatten(), y_rot.flatten())
        grid = (self.xaxis[None, :], self.yaxis[:, None])
        r_obs = griddata(disk, rdisk.flatten(), grid,
                         method=self.shadowed_method)
        t_obs = griddata(disk, tdisk.flatten(), grid,
                         method=self.shadowed_method)
        return r_obs, t_obs, z_func(r_obs)

    def _get_diskframe_coords(self, extend=2.0, oversample=0.5):
        """Disk-frame coordinates based on the cube axes."""
        x_disk = np.linspace(extend * self.xaxis[0], extend * self.xaxis[-1],
                             int(self.nxpix * oversample))[::-1]
        y_disk = np.linspace(extend * self.yaxis[0], extend * self.yaxis[-1],
                             int(self.nypix * oversample))
        x_disk, y_disk = np.meshgrid(x_disk, y_disk)
        r_disk = np.hypot(x_disk, y_disk)
        t_disk = np.arctan2(y_disk, x_disk)
        return x_disk, y_disk, r_disk, t_disk

    # -- Functions to build projected velocity profiles. -- #

    def _vphi(self, params):
        """
        Return a rotation profile in [m/s]. If ``mstar`` is set, assume a
        Keplerian profile including corrections for non-zero emission heights.
        Otherwise, assume a powerlaw profile. See the documents in
        ``disk_coords`` for more details of the parameters.

        Args:
            params (dict): Dictionary of parameters describing the model.

        Returns:
            vproj (ndarray): Projected Keplerian rotation at each pixel (m/s).
        """
        rvals, tvals, zvals = self.disk_coords(**params)
        v_phi = params['vfunc'](rvals, tvals, zvals, params)
        return v_phi + params['vlsr']

    def _proj_vkep(self, rvals, tvals, zvals, params):
        """Projected Keplerian rotational velocity profile."""
        rvals *= sc.au * params['dist']
        zvals *= sc.au * params['dist']
        v_phi = sc.G * params['mstar'] * self.msun * np.power(rvals, 2)
        v_phi = np.sqrt(v_phi * np.power(np.hypot(rvals, zvals), -3))
        return self._proj_vphi(v_phi, tvals, params)

    def _proj_vpow(self, rvals, tvals, zvals, params):
        """Projected power-law rotational velocity profile."""
        v_phi = (rvals * params['dist'] / 100.)**params['vp_q']
        v_phi *= np.exp(-(rvals / params['vp_rtaper'])**params['vp_qtaper'])
        v_phi = self._proj_vphi(params['vp_100'] * v_phi, tvals, params)
        v_rad = (rvals * params['dist'] / 100.)**params['vr_q']
        v_rad = self._proj_vrad(params['vr_100'] * v_rad, tvals, params)
        return v_phi + v_rad

    def _proj_vphi(self, v_phi, tvals, params):
        """Project the rotational velocity."""
        return v_phi * np.cos(tvals) * abs(np.sin(np.radians(params['inc'])))

    def _proj_vrad(self, v_rad, tvals, params):
        """Project the radial velocity."""
        return v_rad * np.sin(tvals) * abs(np.sin(np.radians(params['inc'])))

    def _make_model(self, params):
        """Build the velocity model from the dictionary of parameters."""
        v_phi = self._vphi(params)
        if params['beam']:
            v_phi = rotationmap._convolve_image(v_phi, self._beamkernel())
        return v_phi

    # -- Functions to help determine the emission height. -- #

    def find_maxima(self, x0=0.0, y0=0.0, PA=0.0, vlsr=None, r_max=None,
                    r_min=None, smooth=False, through_center=True):
        """
        Find the node of velocity maxima along the major axis of the disk.

        Args:
            x0 (Optional[float]): Source center offset along x-axis in
                [arcsec].
            y0 (Optional[float]): Source center offset along y-axis in
                [arcsec].
            PA (Optioanl[float]): Source position angle in [deg].
            vlsr (Optional[float]): Systemic velocity in [m/s].
            r_max (Optional[float]): Maximum offset to consider in [arcsec].
            r_min (Optional[float]): Minimum offset to consider in [arcsec].
            smooth (Optional[bool/float]): Smooth the line of nodes. If
                ``True``, smoth with the beam kernel, otherwise ``smooth``
                describes the FWHM of the Gaussian convolution kernel in
                [arcsec].
            through_center (Optional[bool]): If ``True``, force the central
                pixel to go through ``(0, 0)``.

        Returns:
            array, array: Arrays of ``x_sky`` and ``y_sky`` values of the
            maxima.
        """

        # Default parameters.
        vlsr = np.nanmedian(self.data) if vlsr is None else vlsr
        r_max = 0.5 * self.xaxis.max() if r_max is None else r_max
        r_min = 0.0 if r_min is None else r_min

        # Shift and rotate the image.
        data = self._shift_center(dx=x0, dy=y0, save=False)
        data = self._rotate_image(PA=PA, data=data, save=False)

        # Find the maximum values. Apply some clipping to help.
        mask = np.maximum(0.3 * abs(self.xaxis), self.bmaj)
        mask = abs(self.yaxis)[:, None] > mask[None, :]
        resi = np.where(mask, 0.0, abs(data - vlsr))
        resi = np.take(self.yaxis, np.argmax(resi, axis=0))

        # Gentrification.
        if through_center:
            resi[abs(self.xaxis).argmin()] = 0.0
        if smooth:
            if isinstance(smooth, bool):
                kernel = np.hanning(self.bmaj / self.dpix)
                kernel /= kernel.sum()
            else:
                from astropy.convolution import Gaussian1DKernel
                kernel = Gaussian1DKernel(smooth / self.fwhm / self.dpix)
            resi = np.convolve(resi, kernel, mode='same')
        mask = np.logical_and(abs(self.xaxis) <= r_max,
                              abs(self.xaxis) >= r_min)
        x, y = self.xaxis[mask], resi[mask]

        # Rotate back to sky-plane and return.
        x, y = self._rotate_coords(x, -y, PA)
        return x + x0, y + y0

    def find_minima(self, x0=0.0, y0=0.0, PA=0.0, vlsr=None, r_max=None,
                    r_min=None, smooth=False, through_center=True):
        """
        Find the line of nodes where v0 = v_LSR along the minor axis.

        Args:
            x0 (Optional[float]): Source center offset along x-axis in
                [arcsec].
            y0 (Optional[float]): Source center offset along y-axis in
                [arcsec].
            PA (Optioanl[float]): Source position angle in [deg].
            vlsr (Optional[float]): Systemic velocity in [m/s].
            r_max (Optional[float]): Maximum offset to consider in [arcsec].
            r_min (Optional[float]): Minimum offset to consider in [arcsec].
            smooth (Optional[bool/float]): Smooth the line of nodes. If
                ``True``, smoth with the beam kernel, otherwise ``smooth``
                describes the FWHM of the Gaussian convolution kernel in
                [arcsec].
            through_center (Optional[bool]): If ``True``, force the central
                pixel to go through ``(0, 0)``.

        Returns:
            array, array: Arrays of ``x_sky`` and ``y_sky`` values of the
            velocity minima.
        """

        # Default parameters.
        vlsr = np.nanmedian(self.data) if vlsr is None else vlsr
        r_max = 0.5 * self.xaxis.max() if r_max is None else r_max
        r_min = 0.0 if r_min is None else r_min

        # Shift and rotate the image.
        data = self._shift_center(dx=x0, dy=y0, save=False)
        data = self._rotate_image(PA=PA, data=data, save=False)

        # Find the maximum values. Apply some clipping to help.
        mask = np.maximum(0.3 * abs(self.yaxis), self.bmaj)
        mask = abs(self.xaxis)[None, :] > mask[:, None]
        resi = np.where(mask, 1e10, abs(data - vlsr))
        resi = np.take(-self.yaxis, np.argmin(resi, axis=1))

        # Gentrification.
        if through_center:
            resi[abs(self.yaxis).argmin()] = 0.0
        if smooth:
            if isinstance(smooth, bool):
                kernel = np.hanning(self.bmaj / self.dpix)
                kernel /= kernel.sum()
            else:
                from astropy.convolution import Gaussian1DKernel
                kernel = Gaussian1DKernel(smooth / self.fwhm / self.dpix)
            resi = np.convolve(resi, kernel, mode='same')
        mask = np.logical_and(abs(self.yaxis) <= r_max,
                              abs(self.yaxis) >= r_min)
        x, y = resi[mask], self.yaxis[mask]

        # Rotate back to sky-plane and return.
        x, y = self._rotate_coords(x, -y, PA)
        return x + x0, y + y0

    def _fit_surface(self, inc, PA, x0=None, y0=None, r_min=None, r_max=None,
                     fit_z1=False, ycut=20.0, **kwargs):
        """
        Fit the emission surface with the parametric model based on the pixels
        of peak velocity.

        Args:
            inc (float): Disk inclination in [degrees].
            PA (float): Disk position angle in [degrees].
            x0 (Optional[float]): Disk center x-offset in [arcsec].
            y0 (Optional[float]): Disk center y-offset in [arcsec].
            r_min (Optional[float]): Minimum radius to fit in [arcsec].
            r_max (Optional[float]): Maximum radius to fit in [arcsec].
            fit_z1 (Optional[bool]): Include (z1, phi) in the fit.
            y_cut (Optional[float]): Only use the `ycut` fraction of the minor
                axis for fitting.
            kwargs (Optional[dict]): Additional kwargs for curve_fit.

        Returns:
            popt (list): Best-fit parameters of (z0, psi[, z1, phi]).
        """

        from scipy.optimize import curve_fit

        def z_func(x, z0, psi, z1=0.0, phi=1.0):
            return z0 * x**psi + z1 * x**phi

        # Get the coordinates to fit.
        r, z = self.get_peak_pix(PA=PA, inc=inc, x0=x0, y0=y0,
                                 frame='disk', ycut=ycut)

        # Mask the coordinates.
        r_min = r_min if r_min is not None else r[0]
        r_max = r_max if r_max is not None else r[-1]
        m1 = np.logical_and(np.isfinite(r), np.isfinite(z))
        m2 = (r >= r_min) & (r <= r_max)
        mask = m1 & m2
        r, z = r[mask], z[mask]

        # Initial guesses for the free params.
        p0 = [z[abs(r - 1.0).argmin()], 1.0]
        if fit_z1:
            p0 += [-0.05, 3.0]
        maxfev = kwargs.pop('maxfev', 100000)

        # First fit the inner half to get a better p0 value.
        mask = r <= 0.5 * r.max()
        p_in = curve_fit(z_func, r[mask], z[mask], p0=p0, maxfev=maxfev,
                         **kwargs)[0]

        # Remove the pixels which are 'negative' depending on the inclination.
        mask = z > 0.0 if np.sign(p_in[0]) else z < 0.0
        popt = curve_fit(z_func, r[mask], z[mask], p0=p_in, maxfev=maxfev,
                         **kwargs)[0]
        return popt

    # -- Axes Functions -- #

    def _clip_cube(self, radius):
        """Clip the cube to +/- radius arcseconds from the origin."""
        xa = abs(self.xaxis - radius).argmin()
        if self.xaxis[xa] < radius:
            xa -= 1
        xb = abs(self.xaxis + radius).argmin()
        if -self.xaxis[xb] < radius:
            xb += 1
        xb += 1
        ya = abs(self.yaxis + radius).argmin()
        if -self.yaxis[ya] < radius:
            ya -= 1
        yb = abs(self.yaxis - radius).argmin()
        if self.yaxis[yb] < radius:
            yb += 1
        yb += 1
        self.xaxis = self.xaxis[xa:xb]
        self.yaxis = self.yaxis[ya:yb]
        self.data = self.data[ya:yb, xa:xb]
        self.error = self.error[ya:yb, xa:xb]

    def _downsample_cube(self, N):
        """Downsample the cube to make faster calculations."""
        N0 = int(N / 2)
        self.xaxis = self.xaxis[N0::N]
        self.yaxis = self.yaxis[N0::N]
        self.data = self.data[N0::N, N0::N]
        self.error = self.error[N0::N, N0::N]

    def _read_position_axis(self, a=1):
        """Returns the position axis in [arcseconds]."""
        if a not in [1, 2]:
            raise ValueError("'a' must be in [0, 1].")
        a_len = self.header['naxis%d' % a]
        a_del = self.header['cdelt%d' % a]
        a_pix = self.header['crpix%d' % a]
        return 3600 * ((np.arange(a_len) - a_pix + 1.5) * a_del)

    def _shift_center(self, dx=0.0, dy=0.0, data=None, save=True):
        """
        Shift the center of the image.

        Args:
            dx (optional[float]): shift along x-axis [arcsec].
            dy (optional[float]): Shift along y-axis [arcsec].
            data (optional[ndarray]): Data to shift. If nothing is provided,
                will shift the attached ``rotationmap.data``.
            save (optional[bool]): If True, overwrite ``rotationmap.data``.
        """
        from scipy.ndimage import shift
        data = self.data.copy() if data is None else data
        to_shift = np.where(np.isfinite(data), data, 0.0)
        data = shift(to_shift, [-dy / self.dpix, dx / self.dpix])
        if save:
            self.data = data
        return data

    def _rotate_image(self, PA, data=None, save=True):
        """
        Rotate the image anticlockwise about the center.

        Args:
            PA (float): Rotation angle in [degrees].
            data (optional[ndarray]): Data to shift. If nothing is provided,
                will shift the attached ``rotationmap.data``.
            save (optional[bool]): If True, overwrite ``rotationmap.data``.
        """
        from scipy.ndimage import rotate
        data = self.data.copy() if data is None else data
        to_rotate = np.where(np.isfinite(data), data, 0.0)
        data = rotate(to_rotate, PA - 90.0, reshape=False)
        if save:
            self.data = data
        return data

    # -- Convolution functions. -- #

    def _readbeam(self):
        """Reads the beam properties from the header."""
        try:
            if self.header.get('CASAMBM', False):
                beam = fits.open(self.path)[1].data
                beam = np.median([b[:3] for b in beam.view()], axis=0)
                self.bmaj, self.bmin, self.bpa = beam
            else:
                self.bmaj = self.header['bmaj'] * 3600.
                self.bmin = self.header['bmin'] * 3600.
                self.bpa = self.header['bpa']
        except KeyError:
            self.bmaj = self.dpix
            self.bmin = self.dpix
            self.bpa = 0.0
            self.beamarea = self.dpix**2.0

    def _beamkernel(self):
        """Returns the 2D Gaussian kernel for convolution."""
        from astropy.convolution import Kernel
        bmaj = self.bmaj / self.dpix / self.fwhm
        bmin = self.bmin / self.dpix / self.fwhm
        bpa = np.radians(self.bpa)
        return Kernel(self._gaussian2D(bmin, bmaj, bpa + 90.).T)

    def _gaussian2D(self, dx, dy, PA=0.0):
        """2D Gaussian kernel in pixel coordinates."""
        xm = np.arange(-4*np.nanmax([dy, dx]), 4*np.nanmax([dy, dx])+1)
        x, y = np.meshgrid(xm, xm)
        x, y = self._rotate_coords(x, y, PA)
        k = np.power(x / dx, 2) + np.power(y / dy, 2)
        return np.exp(-0.5 * k) / 2. / np.pi / dx / dy

    @staticmethod
    def _convolve_image(image, kernel, fast=True):
        """Convolve the image with the provided kernel."""
        if fast:
            from astropy.convolution import convolve_fft
            return convolve_fft(image, kernel)
        from astropy.convolution import convolve
        return convolve(image, kernel)

    # -- Plotting functions. -- #

    @staticmethod
    def colormap():
        import matplotlib.pyplot as plt
        import matplotlib.colors as mcolors
        c2 = plt.cm.Reds(np.linspace(0, 1, 32))
        c1 = plt.cm.Blues_r(np.linspace(0, 1, 32))
        c1 = np.vstack([c1, [1, 1, 1, 1]])
        colors = np.vstack((c1, c2))
        return mcolors.LinearSegmentedColormap.from_list('eddymap', colors)

    @property
    def extent(self):
        return [self.xaxis[0], self.xaxis[-1], self.yaxis[0], self.yaxis[-1]]

    def _plot_bestfit(self, params, ivar=None, residual=False,
                      return_ax=False):
        """Plot the best-fit model."""
        import matplotlib.pyplot as plt
        ax = plt.subplots()[1]
        vkep = self._make_model(params) * 1e-3
        levels = np.nanpercentile(vkep, [2, 98])
        levels = np.linspace(levels[0], levels[1], 30)
        im = ax.contourf(self.xaxis, self.yaxis, vkep, levels,
                         cmap=rotationmap.colormap(), extend='both')
        if ivar is not None:
            ax.contour(self.xaxis, self.yaxis, ivar,
                       [0.0], colors='k')
            ax.contourf(self.xaxis, self.yaxis, ivar,
                        [-1.0, 0.0], colors='k', alpha=0.5)
        cb = plt.colorbar(im, pad=0.02, format='%.2f')
        cb.set_label(r'${\rm v_{mod} \quad (km\,s^{-1})}$',
                     rotation=270, labelpad=15)
        cb.minorticks_on()
        self._gentrify_plot(ax)
        if return_ax:
            return ax

    def _plot_residual(self, params, ivar=None, return_ax=False):
        """Plot the residual from the provided model."""
        import matplotlib.cm as cm
        import matplotlib.pyplot as plt
        ax = plt.subplots()[1]
        vres = self.data * 1e3 - self._make_model(params)
        levels = np.where(self.ivar != 0.0, vres, np.nan)
        levels = np.nanpercentile(levels, [2, 98])
        levels = max(abs(levels[0]), abs(levels[1]))
        levels = np.linspace(-levels, levels, 30)
        im = ax.contourf(self.xaxis, self.yaxis, vres, levels,
                         cmap=cm.RdBu_r, extend='both')
        if ivar is not None:
            ax.contour(self.xaxis, self.yaxis, ivar,
                       [0.0], colors='k')
            ax.contourf(self.xaxis, self.yaxis, ivar,
                        [-1.0, 0.0], colors='k', alpha=0.5)
        cb = plt.colorbar(im, pad=0.02, format='%d', ax=ax)
        cb.set_label(r'${\rm  v_{0} - v_{mod} \quad (m\,s^{-1})}$',
                     rotation=270, labelpad=15)
        cb.minorticks_on()
        self._gentrify_plot(ax)
        if return_ax:
            return ax

    def plot_surface(self, ax=None, x0=0.0, y0=0.0, inc=0.0, PA=0.0, z0=0.0,
                     psi=0.0, r_cavity=0.0, r_taper=None, q_taper=None,
                     w_i=0.0, w_r=1.0, w_t=0.0, r_min=0.0, r_max=None,
                     ntheta=24, nrad=15, check_mask=True, **kwargs):
        """
        Overplot the emission surface onto the provided axis.

        Args:
            ax (Optional[AxesSubplot]): Axis to plot onto.
            x0 (Optional[float]): Source right ascension offset [arcsec].
            y0 (Optional[float]): Source declination offset [arcsec].
            inc (Optional[float]): Source inclination [deg].
            PA (Optional[float]): Source position angle [deg]. Measured
                between north and the red-shifted semi-major axis in an
                easterly direction.
            z0 (Optional[float]): Aspect ratio at 1" for the emission surface.
                To get the far side of the disk, make this number negative.
            psi (Optional[float]): Flaring angle for the emission surface.
            z1 (Optional[float]): Aspect ratio correction term at 1" for the
                emission surface. Should be opposite sign to z0.
            phi (Optional[float]): Flaring angle correction term for the
                emission surface.
            w_i (Optional[float]): Warp inclination in [degrees] at the disk
                center.
            w_r (Optional[float]): Scale radius of the warp in [arcsec].
            w_t (Optional[float]): Angle of nodes of the warp in [degrees].
            r_min (Optional[float]): Inner radius to plot, default is 0.
            r_max (Optional[float]): Outer radius to plot.
            ntheta (Optional[int]): Number of theta contours to plot.
            nrad (Optional[int]): Number of radial contours to plot.
            check_mask (Optional[bool]): Mask regions which are like projection
                errors for highly flared surfaces.

        Returns:
            ax (AxesSubplot): Axis with the contours overplotted.
        """

        # Dummy axis to overplot.
        if ax is None:
            import matplotlib.pyplot as plt
            ax = plt.subplots()[1]

        # Front half of the disk.
        rf, tf, zf = self.disk_coords(x0=x0, y0=y0, inc=inc, PA=PA, z0=z0,
                                      psi=psi, r_cavity=r_cavity,
                                      r_taper=r_taper, q_taper=q_taper)
        mf = np.logical_and(zf >= 0.0, np.isfinite(self.data))
        rf = np.where(mf, rf, np.nan)
        tf = np.where(mf, tf, np.nan)

        # Rear half of the disk.
        rb, tb, zb = self.disk_coords(x0=x0, y0=y0, inc=inc, PA=PA, z0=-z0,
                                      psi=psi, r_cavity=r_cavity,
                                      r_taper=r_taper, q_taper=q_taper)
        mb = np.logical_and(zb <= 0.0, np.isfinite(self.data))
        rb = np.where(mb, rb, np.nan)
        tb = np.where(mb, tb, np.nan)

        # Flat disk for masking.
        rr, _, _ = self.disk_coords(x0=x0, y0=y0, inc=inc, PA=PA)

        # Make sure the bounds are OK.
        r_min = 0.0 if r_min is None else r_min
        r_max = rr.max() if r_max is None else r_max

        # Make sure the front side hides the rear.
        mf = np.logical_and(rf >= r_min, rf <= r_max)
        mb = np.logical_and(rb >= r_min, rb <= r_max)
        tf = np.where(mf, tf, np.nan)
        rb = np.where(~mf, rb, np.nan)
        tb = np.where(np.logical_and(np.isfinite(rb), mb), tb, np.nan)

        # For some geometries we want to make sure they're not doing funky
        # things in the outer disk when psi is large.
        if check_mask:
            mm = rr <= check_mask * r_max
            rf = np.where(mm, rf, np.nan)
            rb = np.where(mm, rb, np.nan)
            tf = np.where(mm, tf, np.nan)
            tb = np.where(mm, tb, np.nan)

        # Popluate the kwargs with defaults.
        lw = kwargs.pop('lw', kwargs.pop('linewidth', 1.0))
        zo = kwargs.pop('zorder', 10000)
        c = kwargs.pop('colors', kwargs.pop('c', 'k'))

        radii = np.linspace(0, r_max, int(nrad + 1))[1:]
        theta = np.linspace(-np.pi, np.pi, int(ntheta))
        theta += np.diff(theta)[0]

        # Do the plotting.
        ax.contour(self.xaxis, self.yaxis, rf, levels=radii, colors=c,
                   linewidths=lw, linestyles='-', zorder=zo, **kwargs)
        for tt in theta:
            tf_tmp = np.where(abs(tf - tt) <= 0.5, tf, np.nan)
            ax.contour(self.xaxis, self.yaxis, tf_tmp, levels=[tt], colors=c,
                       linewidths=lw, linestyles='-', zorder=zo, **kwargs)
        ax.contour(self.xaxis, self.yaxis, rb, levels=radii, colors=c,
                   linewidths=lw, linestyles='--', zorder=zo, **kwargs)
        for tt in theta:
            tb_tmp = np.where(abs(tb - tt) <= 0.5, tb, np.nan)
            ax.contour(self.xaxis, self.yaxis, tb_tmp, levels=theta, colors=c,
                       linewidths=lw, linestyles='--', zorder=zo)

        return ax

    def _plot_axes(self, ax, x0=0.0, y0=0.0, inc=0.0, PA=0.0, major=1.0,
                   plot_kwargs=None):
        """
        Plot the major and minor axes on the provided axis.

        Args:
            ax (Matplotlib axes): Axes instance to plot onto.
            x0 (Optional[float]): Relative x-location of the center [arcsec].
            y0 (Optional[float]): Relative y-location of the center [arcsec].
            inc (Optional[float]): Inclination of the disk in [degrees].
            PA (Optional[float]): Position angle of the disk in [degrees].
            major (Optional[float]): Size of the major axis line in [arcsec].
            plot_kwargs (Optional[dict]): Dictionary of parameters to pass to
                ``matplotlib.plot``.

        Returns:
            matplotlib axis: Matplotlib ax with axes drawn.
        """
        x = np.array([major, -major, 0.0, 0.0])
        y = np.array([0.0, 0.0, major * np.cos(np.radians(inc)),
                      -major * np.cos(np.radians(inc))])
        x, y = self._rotate_coords(x, y, PA=PA)
        x, y = x + x0, y + y0
        plot_kwargs = {} if plot_kwargs is None else plot_kwargs
        c = plot_kwargs.pop('c', plot_kwargs.pop('color', 'k'))
        ls = plot_kwargs.pop('ls', plot_kwargs.pop('linestyle', ':'))
        lw = plot_kwargs.pop('lw', plot_kwargs.pop('linewidth', 1.0))
        ax.plot(x[:2], y[:2], c=c, ls=ls, lw=lw)
        ax.plot(x[2:], y[2:], c=c, ls=ls, lw=lw)
        return ax

    def plot_maxima(self, x0=0.0, y0=0.0, inc=0.0, PA=0.0, vlsr=None,
                    r_max=1.0, r_min=None, smooth=False, through_center=True,
                    levels=None, plot_axes_kwargs=None, plot_kwargs=None,
                    return_fig=False):
        """
        Mark the position of the maximum velocity as a function of radius. This
        can help demonstrate if there is an appreciable bend in velocity
        contours indicative of a flared emission surface. This function pulls
        the results from fit_cube.get_peak_pix() for plotting.

        Args:
            x0 (Optional[float]): Source center x-offset in [arcsec].
            y0 (Optional[float]): Source center y-offset in [arcsec].
            inc (Optional[float]): Disk inclination in [degrees].
            PA (Optional[float]): Disk position angle in [degrees].
            vlsr (Optional[float]): Systemic velocity in [m/s].
            r_max (Optional[float]): Maximum offset to consider in [arcsec].
            r_min (Optional[float]): Minimum offset to consider in [arcsec].
                This is useful to skip large regions where beam convolution
                dominates the map.
            smooth (Optional[bool/float]): Smooth the line of nodes. If
                ``True``, smoth with the beam kernel, otherwise ``smooth``
                describes the FWHM of the Gaussian convolution kernel in
                [arcsec].
            through_center (Optional[bool]): If ``True``, force the central
                pixel to go through ``(0, 0)``.
            levels (Optional[array]): Levels to pass to ``plot_data``.
            plot_axes_kwargs (Optional[dict]): Dictionary of kwargs for the
                plot_axes function.
            plot_kwargs (Optional[dict]): Dictionary of kwargs for the
                plot of the maximum and minimum pixels.
            return_fig (Optional[bool]): If True, return the figure.

        Returns:
            (Matplotlib figure): If return_fig is ``True``. Can access the axis
                through ``fig.axes[0]``.
        """

        # Background figure.
        fig = self.plot_data(levels=levels, return_fig=True)
        ax = fig.axes[0]
        if plot_axes_kwargs is None:
            plot_axes_kwargs = dict()
        self._plot_axes(ax=ax, x0=x0, y0=y0, inc=inc, PA=PA,
                        major=plot_axes_kwargs.pop('major', r_max),
                        **plot_axes_kwargs)

        # Get the pixels for the maximum and minimum values.
        x_maj, y_maj = self.find_maxima(x0=x0, y0=y0, PA=PA, vlsr=vlsr,
                                        r_max=r_max, r_min=r_min,
                                        smooth=smooth,
                                        through_center=through_center)
        r_max = r_max * np.cos(np.radians(inc))
        x_min, y_min = self.find_minima(x0=x0, y0=y0, PA=PA, vlsr=vlsr,
                                        r_max=r_max, r_min=r_min,
                                        smooth=smooth,
                                        through_center=through_center)

        # Plot the lines.
        plot_kwargs = {} if plot_kwargs is None else plot_kwargs
        c = plot_kwargs.pop('c', plot_kwargs.pop('color', 'k'))
        lw = plot_kwargs.pop('lw', plot_kwargs.pop('linewidth', 1.0))

        ax.plot(x_maj, y_maj, c=c, lw=lw, **plot_kwargs)
        ax.plot(x_min, y_min, c=c, lw=lw, **plot_kwargs)

        if return_fig:
            return fig

    @staticmethod
    def _plot_walkers(samples, nburnin=None, labels=None, histogram=True):
        """Plot the walkers to check if they are burning in."""

        # Import matplotlib.
        import matplotlib.pyplot as plt
        from mpl_toolkits.axes_grid1.axes_divider import make_axes_locatable

        # Check the length of the label list.
        if labels is None:
            if samples.shape[0] != len(labels):
                raise ValueError("Not correct number of labels.")

        # Cycle through the plots.
        for s, sample in enumerate(samples):
            fig, ax = plt.subplots()
            for walker in sample.T:
                ax.plot(walker, alpha=0.1, color='k')
            ax.set_xlabel('Steps')
            if labels is not None:
                ax.set_ylabel(labels[s])
            if nburnin is not None:
                ax.axvline(nburnin, ls=':', color='r')

            # Include the histogram.
            if histogram:
                fig.set_size_inches(1.37 * fig.get_figwidth(),
                                    fig.get_figheight(), forward=True)
                ax_divider = make_axes_locatable(ax)
                bins = np.linspace(ax.get_ylim()[0], ax.get_ylim()[1], 50)
                hist, _ = np.histogram(sample[nburnin:].flatten(), bins=bins,
                                       density=True)
                bins = np.average([bins[1:], bins[:-1]], axis=0)
                ax1 = ax_divider.append_axes("right", size="35%", pad="2%")
                ax1.fill_betweenx(bins, hist, np.zeros(bins.size), step='mid',
                                  color='darkgray', lw=0.0)
                ax1.set_ylim(ax.get_ylim()[0], ax.get_ylim()[1])
                ax1.set_xlim(0, ax1.get_xlim()[1])
                ax1.set_yticklabels([])
                ax1.set_xticklabels([])
                ax1.tick_params(which='both', left=0, bottom=0, top=0, right=0)
                ax1.spines['right'].set_visible(False)
                ax1.spines['bottom'].set_visible(False)
                ax1.spines['top'].set_visible(False)

    @staticmethod
    def _plot_corner(samples, labels=None, quantiles=None):
        """Plot the corner plot to check for covariances."""
        import corner
        quantiles = [0.16, 0.5, 0.84] if quantiles is None else quantiles
        corner.corner(samples, labels=labels, title_fmt='.4f', bins=30,
                      quantiles=quantiles, show_titles=True)

    def plot_beam(self, ax, dx=0.125, dy=0.125, **kwargs):
        """Plot the sythensized beam on the provided axes."""
        from matplotlib.patches import Ellipse
        beam = Ellipse(ax.transLimits.inverted().transform((dx, dy)),
                       width=self.bmin, height=self.bmaj, angle=-self.bpa,
                       fill=kwargs.get('fill', False),
                       hatch=kwargs.get('hatch', '//////////'),
                       lw=kwargs.get('linewidth', kwargs.get('lw', 1)),
                       color=kwargs.get('color', kwargs.get('c', 'k')),
                       zorder=kwargs.get('zorder', 1000))
        ax.add_patch(beam)

    def _gentrify_plot(self, ax):
        """Gentrify the plot."""
        from matplotlib.ticker import MaxNLocator
        ax.set_aspect(1)
        ax.grid(ls=':', color='k', alpha=0.1, lw=0.5)
        ax.tick_params(which='both', right=True, top=True)
        ax.set_xlim(self.xaxis.max(), self.xaxis.min())
        ax.set_ylim(self.yaxis.min(), self.yaxis.max())
        ax.xaxis.set_major_locator(MaxNLocator(5, min_n_ticks=3))
        ax.yaxis.set_major_locator(MaxNLocator(5, min_n_ticks=3))
        ax.set_xlabel('Offset (arcsec)')
        ax.set_ylabel('Offset (arcsec)')
        if self.bmaj is not None:
            self.plot_beam(ax=ax)

max_stars_repo_name: elnjensen/eddy

max_issues_repo_name: elnjensen/eddy

max_forks_repo_name: elnjensen/eddy

reason: import numpy,import scipy,from scipy,from astropy


text:
# Author: Bichen Wu (bichen@berkeley.edu) 08/25/2016

"""Evaluation"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import cv2
from datetime import datetime
import os.path
import sys
import time

import numpy as np
from six.moves import xrange
import tensorflow as tf

from config import *
from dataset import pascal_voc, kitti, ball
from utils.util import bbox_transform, Timer
from nets import *

FLAGS = tf.app.flags.FLAGS

# tf.app.flags.DEFINE_string('dataset', 'KITTI', """Currently support PASCAL_VOC or KITTI dataset.""")
tf.app.flags.DEFINE_string('dataset', 'BALL',""" """)
tf.app.flags.DEFINE_string('data_path', 'D:\\Humanoid\\squeezeDet\\squeezeDet-master\\data\\val\\', """Root directory of data""")

tf.app.flags.DEFINE_string('image_set', 'val',
                           """Can be train, trainval, val, or test""")

tf.app.flags.DEFINE_string('eval_dir', 'logs\\squeezeDet\\val',
                            """Directory where to write event logs """)

tf.app.flags.DEFINE_string('checkpoint_path', '.\\logs\\squeezeDet\\train',#'/tmp/bichen/logs/squeezeDet/train',
                            """Path to the training checkpoint.""")
tf.app.flags.DEFINE_integer('eval_interval_secs', 60 * 1,
                             """How often to check if new cpt is saved.""")
tf.app.flags.DEFINE_boolean('run_once', False,
                             """Whether to run eval only once.""")
tf.app.flags.DEFINE_string('net', 'squeezeDet',
                           """Neural net architecture.""")
tf.app.flags.DEFINE_string('gpu', '0', """gpu id.""")

def eval_once(saver, ckpt_path, summary_writer, eval_summary_ops, eval_summary_phs, imdb, model):

  with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:

    # Restores from checkpoint
    saver.restore(sess, ckpt_path)
    # Assuming model_checkpoint_path looks something like:
    #   /ckpt_dir/model.ckpt-0,
    # extract global_step from it.
    global_step = ckpt_path.split('/')[-1].split('-')[-1]

    num_images = len(imdb.image_idx)

    all_boxes = [[[] for _ in xrange(num_images)]
                 for _ in xrange(imdb.num_classes)]

    _t = {'im_detect': Timer(), 'im_read': Timer(), 'misc': Timer()}

    num_detection = 0.0
    for i in xrange(num_images):
      _t['im_read'].tic()
      images, scales = imdb.read_image_batch(shuffle=False)
      _t['im_read'].toc()

      _t['im_detect'].tic()
      det_boxes, det_probs, det_class = sess.run(
          [model.det_boxes, model.det_probs, model.det_class],
          feed_dict={model.image_input:images})
      _t['im_detect'].toc()

      _t['misc'].tic()
      for j in range(len(det_boxes)): # batch
        # rescale
        det_boxes[j, :, 0::2] /= scales[j][0]
        det_boxes[j, :, 1::2] /= scales[j][1]

        det_bbox, score, det_class = model.filter_prediction(
            det_boxes[j], det_probs[j], det_class[j])

        num_detection += len(det_bbox)
        for c, b, s in zip(det_class, det_bbox, score):
          all_boxes[c][i].append(bbox_transform(b) + [s])
      _t['misc'].toc()

      print ('im_detect: {:d}/{:d} im_read: {:.3f}s '
             'detect: {:.3f}s misc: {:.3f}s'.format(
                i+1, num_images, _t['im_read'].average_time,
                _t['im_detect'].average_time, _t['misc'].average_time))

    print ('Evaluating detections...')
    aps, ap_names = imdb.evaluate_detections(
        FLAGS.eval_dir, global_step, all_boxes)

    print ('Evaluation summary:')
    print ('  Average number of detections per image: {}:'.format(
      num_detection/num_images))
    print ('  Timing:')
    print ('    im_read: {:.3f}s detect: {:.3f}s misc: {:.3f}s'.format(
      _t['im_read'].average_time, _t['im_detect'].average_time,
      _t['misc'].average_time))
    print ('  Average precisions:')

    feed_dict = {}
    for cls, ap in zip(ap_names, aps):
      feed_dict[eval_summary_phs['APs/'+cls]] = ap
      print ('    {}: {:.3f}'.format(cls,ap))

    print ('    Mean average precision: {:.3f}'.format(np.mean(aps)))
    feed_dict[eval_summary_phs['APs/mAP']] = np.mean(aps)
    feed_dict[eval_summary_phs['timing/im_detect']] = \
        _t['im_detect'].average_time
    feed_dict[eval_summary_phs['timing/im_read']] = \
        _t['im_read'].average_time
    feed_dict[eval_summary_phs['timing/post_proc']] = \
        _t['misc'].average_time
    feed_dict[eval_summary_phs['num_det_per_image']] = \
        num_detection/num_images

    print ('Analyzing detections...')
    stats, ims = imdb.do_detection_analysis_in_eval(
        FLAGS.eval_dir, global_step)

    eval_summary_str = sess.run(eval_summary_ops, feed_dict=feed_dict)
    for sum_str in eval_summary_str:
      summary_writer.add_summary(sum_str, global_step)

def evaluate():
  """Evaluate."""

  os.environ['CUDA_VISIBLE_DEVICES'] = FLAGS.gpu

  with tf.Graph().as_default() as g:

    if FLAGS.net == 'squeezeDet':
      mc = ball_config()
      mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
      mc.LOAD_PRETRAINED_MODEL = False
      model = SqueezeDet(mc)
    elif FLAGS.net == 'squeezeDet+':
      mc = ball_config()
      mc.BATCH_SIZE = 1 # TODO(bichen): allow batch size > 1
      mc.LOAD_PRETRAINED_MODEL = False
      model = SqueezeDetPlus(mc)

    imdb = ball(FLAGS.image_set, FLAGS.data_path, mc)

    # add summary ops and placeholders
    ap_names = []
    # for cls in imdb.classes:
      # ap_names.append(cls+'_easy')
      # ap_names.append(cls+'_medium')
      # ap_names.append(cls+'_hard')

    eval_summary_ops = []
    eval_summary_phs = {}
    # for ap_name in ap_names:
    ph = tf.placeholder(tf.float32)
    eval_summary_phs['APs'] = ph
    eval_summary_ops.append(tf.summary.scalar('APs', ph))

    ph = tf.placeholder(tf.float32)
    eval_summary_phs['APs/mAP'] = ph
    eval_summary_ops.append(tf.summary.scalar('APs/mAP', ph))

    ph = tf.placeholder(tf.float32)
    eval_summary_phs['timing/im_detect'] = ph
    eval_summary_ops.append(tf.summary.scalar('timing/im_detect', ph))

    ph = tf.placeholder(tf.float32)
    eval_summary_phs['timing/im_read'] = ph
    eval_summary_ops.append(tf.summary.scalar('timing/im_read', ph))

    ph = tf.placeholder(tf.float32)
    eval_summary_phs['timing/post_proc'] = ph
    eval_summary_ops.append(tf.summary.scalar('timing/post_proc', ph))

    ph = tf.placeholder(tf.float32)
    eval_summary_phs['num_det_per_image'] = ph
    eval_summary_ops.append(tf.summary.scalar('num_det_per_image', ph))

    saver = tf.train.Saver(model.model_params)

    summary_writer = tf.summary.FileWriter(FLAGS.eval_dir, g)

    ckpts = set()
    while True:
      if FLAGS.run_once:
        # When run_once is true, checkpoint_path should point to the exact
        # checkpoint file.
        eval_once(
            saver, FLAGS.checkpoint_path, summary_writer, eval_summary_ops,
            eval_summary_phs, imdb, model)
        return
      else:
        # When run_once is false, checkpoint_path should point to the directory
        # that stores checkpoint files.
        ckpt = tf.train.get_checkpoint_state(FLAGS.checkpoint_path)
        if ckpt and ckpt.model_checkpoint_path:
          if ckpt.model_checkpoint_path in ckpts:
            # Do not evaluate on the same checkpoint
            # print ('Wait {:d}s for new checkpoints to be saved ... '
            #           .format(FLAGS.eval_interval_secs))
            # time.sleep(FLAGS.eval_interval_secs)
            break
          else:
            ckpts.add(ckpt.model_checkpoint_path)
            print ('Evaluating {}...'.format(ckpt.model_checkpoint_path))
            eval_once(
                saver, ckpt.model_checkpoint_path, summary_writer,
                eval_summary_ops, eval_summary_phs, imdb, model)
        else:
          print('No checkpoint file found')
          if not FLAGS.run_once:
            print ('Wait {:d}s for new checkpoints to be saved ... '
                      .format(FLAGS.eval_interval_secs))
            time.sleep(FLAGS.eval_interval_secs)

def main(argv=None):  # pylint: disable=unused-argument
  if tf.gfile.Exists(FLAGS.eval_dir):
    tf.gfile.DeleteRecursively(FLAGS.eval_dir)
  tf.gfile.MakeDirs(FLAGS.eval_dir)
  evaluate()

if __name__ == '__main__':
  tf.app.run()

ext: py

max_stars_repo_name: AjJordy/squeezeDet-master

max_issues_repo_name: AjJordy/squeezeDet-master

max_forks_repo_name: AjJordy/squeezeDet-master

reason: import numpy


text:
# pylint: disable-all
from datetime import datetime

import pytest
import numpy as np

from replay.constants import LOG_SCHEMA
from replay.splitters import NewUsersSplitter
from tests.utils import spark

@pytest.fixture
def log(spark):
    return spark.createDataFrame(
        data=[
            ["user2", "item4", datetime(2019, 9, 14), 3.0],
            ["user2", "item1", datetime(2019, 9, 14), 3.0],
            ["user2", "item2", datetime(2019, 9, 15), 4.0],
            ["user1", "item4", datetime(2019, 9, 12), 1.0],
            ["user4", "item1", datetime(2019, 9, 12), 1.0],
            ["user4", "item2", datetime(2019, 9, 13), 2.0],
            ["user3", "item1", datetime(2019, 9, 16), 5.0],
            ["user3", "item4", datetime(2019, 9, 16), 5.0],
            ["user1", "item3", datetime(2019, 9, 17), 1.0],
        ],
        schema=LOG_SCHEMA,
    )

def test_users_are_cold(log):
    splitter = NewUsersSplitter(test_size=0.25, drop_cold_items=False)
    train, test = splitter.split(log)

    train_users = train.toPandas().user_id
    test_users = test.toPandas().user_id

    assert not np.isin(test_users, train_users).any()

@pytest.mark.parametrize("test_size", [-1.0, 2.0])
def test_bad_test_size(log, test_size):
    with pytest.raises(ValueError):
        NewUsersSplitter(test_size)

max_stars_repo_name: dev-rinchin/RePlay

max_issues_repo_name: dev-rinchin/RePlay

max_forks_repo_name: dev-rinchin/RePlay


text:
from warnings import warn
import numpy as np
from scipy.interpolate import RectBivariateSpline
from ..util import img_as_float
from ..filters import sobel

"""
简化后的代码，不可运行，只用于Snake算法逻辑参考
完整代码参见 skimage::segmentation::active_contour
源代码链接：https://github.com/scikit-image/scikit-image/blob/master/skimage/segmentation/active_contour_model.py
"""

def active_contour(image, snake, alpha=0.01, beta=0.1,
                   w_line=0, w_edge=1, gamma=0.01,
                   bc=None, max_px_move=1.0,
                   max_iterations=2500, convergence=0.1,
                   *,
                   boundary_condition='periodic',
                   coordinates=None):

    img = img_as_float(image)

    # 计算梯度幅值图像
    edge = [sobel(img)]

    # 叠加强度和边缘图像:
    img = w_line * img + w_edge * edge[0]

    x, y = snake[:, 0].astype(np.float), snake[:, 1].astype(np.float)
    n = len(x)
    convergence_order = 10
    xsave = np.empty((convergence_order, n))
    ysave = np.empty((convergence_order, n))

    # 计算迭代式中的五对角循环矩阵A
    a = np.roll(np.eye(n), -1, axis=0) + \
        np.roll(np.eye(n), -1, axis=1) - \
        2 * np.eye(n)  # second order derivative, central difference
    b = np.roll(np.eye(n), -2, axis=0) + \
        np.roll(np.eye(n), -2, axis=1) - \
        4 * np.roll(np.eye(n), -1, axis=0) - \
        4 * np.roll(np.eye(n), -1, axis=1) + \
        6 * np.eye(n)  # fourth order derivative, central difference
    A = -alpha * a + beta * b

    # 利用隐式格式来最小化能量函数只需要计算一次逆矩阵
    inv = np.linalg.inv(A + gamma * np.eye(n))

    # Explicit time stepping for image energy minimization:
    intp = RectBivariateSpline(np.arange(img.shape[1]),
                               np.arange(img.shape[0]),
                               img.T, kx=2, ky=2, s=0)
    for i in range(max_iterations):
        fx = intp(x, y, dx=1, grid=False)
        fy = intp(x, y, dy=1, grid=False)
        xn = inv @ (gamma * x + fx)
        yn = inv @ (gamma * y + fy)

        # Movements are capped to max_px_move per iteration:
        dx = max_px_move * np.tanh(xn - x)
        dy = max_px_move * np.tanh(yn - y)
        x += dx
        y += dy

        # Convergence criteria needs to compare to a number of previous
        # configurations since oscillations can occur.
        j = i % (convergence_order + 1)
        if j < convergence_order:
            xsave[j, :] = x
            ysave[j, :] = y
        else:
            dist = np.min(np.max(np.abs(xsave - x[None, :]) +
                                 np.abs(ysave - y[None, :]), 1))
            if dist < convergence:
                break

max_stars_repo_name: QCHighing/pyImageProcessing

max_issues_repo_name: QCHighing/pyImageProcessing

max_forks_repo_name: QCHighing/pyImageProcessing

reason: import numpy,from scipy


text:
"""
graph_diffusion.py
--------------------------

Graph diffusion distance, from

Hammond, D. K., Gur, Y., & Johnson, C. R. (2013, December). Graph diffusion
distance: A difference measure for weighted graphs based on the graph Laplacian
exponential kernel. In Global Conference on Signal and Information Processing,
2013 IEEE (pp 419-422). IEEE. https://doi.org/10.1109/GlobalSIP.2013.6736904

author: Brennan Klein
email: brennanjamesklein at gmail dot com
Submitted as part of the 2019 NetSI Collabathon.

"""

import numpy as np
import networkx as nx
from scipy.sparse.csgraph import laplacian
from netrd.distance import BaseDistance
from netrd.utilities import undirected

class GraphDiffusion(BaseDistance):
    """Find the maximally dissimilar diffusion kernels between two graphs."""

    @undirected
    def dist(self, G1, G2, thresh=1e-08, resolution=1000):
        r"""The graph diffusion distance between two graphs, :math:`G` and :math:`G'`,
        is a distance measure based on the notion of flow within each graph. As
        such, this measure uses the unnormalized Laplacian matrices of both
        graphs, :math:`\mathcal{L}` and :math:`\mathcal{L}'`, and uses them to
        construct time-varying Laplacian exponential diffusion kernels,
        :math:`e^{-t\mathcal{L}}` and :math:`e^{-t\mathcal{L}'}`, by
        effectively simulating a diffusion process for :math:`t` timesteps,
        creating a column vector of node-level activity at each timestep. The
        distance :math:`d_\texttt{GDD}(G, G')` is defined as the Frobenius norm
        between the two diffusion kernels at the timestep :math:`t^{*}` where
        the two kernels are maximally different. That is, we compute the
        Frobenius norms and their differences for each timestep, and return the
        maximum difference.

        .. math::
            D_{GDD}(G,G') = \sqrt{||e^{-t^{*}\mathcal{L}}-e^{-t^{*}\mathcal{L}'}||}

        The results dictionary also stores a 2-tuple of the underlying
        adjacency matrices in `adjacency_matrices`, the Laplacian matrices in
        `laplacian_matrices`, and the output of the optimization process
        (`peak_diffusion_time` and `peak_deviation`).

        Adapted from the authors' MATLAB code, available at: https://rb.gy/txbfrh

        Parameters
        ----------

        G1, G2 (nx.Graph)
            two networkx graphs to be compared.

        thresh (float)
            minimum value above which the eigenvalues will be considered.

        resolution (int)
            number of :math:`t` values to span through.

        Returns
        -------
        dist (float)
            the distance between `G1` and `G2`.

        References
        ----------

        .. [1] Hammond, D. K., Gur, Y., & Johnson, C. R. (2013, December).
               Graph diffusion distance: A difference measure for weighted graphs based on the
               graph Laplacian exponential kernel. In Global Conference on Signal and
               Information Processing, 2013 IEEE (pp 419-422). IEEE.
               https://doi.org/10.1109/GlobalSIP.2013.6736904

        """

        A1 = nx.to_numpy_array(G1)
        A2 = nx.to_numpy_array(G2)

        L1 = laplacian(A1)
        L2 = laplacian(A2)

        vals1, vecs1 = np.linalg.eig(L1)
        vals2, vecs2 = np.linalg.eig(L2)

        eigs = np.hstack((np.diag(vals1), np.diag(vals2)))
        eigs = eigs[np.where(eigs > thresh)]
        eigs = np.sort(eigs)

        if len(eigs) == 0:
            dist = 0
            self.results['dist'] = dist
            return dist

        t_upperbound = np.real(1.0 / eigs[0])
        ts = np.linspace(0, t_upperbound, resolution)

        # Find the Frobenius norms between all the diffusion kernels at
        # different times. Return the value and where this vector is minimized.
        E = -exponential_diffusion_diff(vecs1, vals1, vecs2, vals2, ts)
        f_val, t_star = (np.nanmin(E), np.argmin(E))

        dist = np.sqrt(-f_val)

        self.results['adjacency_matrices'] = A1, A2
        self.results['laplacian_matrices'] = L1, L2
        self.results['peak_diffusion_time'] = t_star
        self.results['peak_deviation'] = f_val

        self.results['dist'] = dist

        return dist

def exponential_diffusion_diff(vecs1, vals1, vecs2, vals2, ts):
    """
    Computes Frobenius norm of difference of Laplacian exponential diffusion
    kernels, at specified timepoints.

    Parameters
    ----------

    vecs1, vecs2 (np.array)
        eigenvectors of the Laplacians of `G1` and `G2`

    vals1, vals2 (np.array)
        eigenvalues of the Laplacians of `G1` and `G2`

    ts (np.array)
        times at which to compute the difference in Frobenius norms

    Returns
    -------

    diffs (np.array)
        same shape as :math:`t`, contains differences of Frobenius norms

    """

    diffs = np.zeros(len(ts))

    for kt, t in enumerate(ts):
        exp_diag_1 = np.diag(np.exp(-t * np.diag(vals1)))
        exp_diag_2 = np.diag(np.exp(-t * np.diag(vals2)))

        # multiply the eigenvectors element-wise by the appropriate diffusion value
        # before left-multiplying the eigenvectors again.
        norm1 = vecs1.dot(np.multiply(exp_diag_1, vecs1).T)
        norm2 = vecs2.dot(np.multiply(exp_diag_2, vecs2).T)
        diff = norm1 - norm2

        diffs[kt] = (diff ** 2).sum()

    return diffs

max_stars_repo_name: sdmccabe/wegd

max_issues_repo_name: sdmccabe/wegd

max_forks_repo_name: sdmccabe/wegd

reason: import numpy,from scipy,import networkx


text:
import numpy as np
from ..constants import molwt

def from_aviNOx(emissions, frac, E_ref=2.946, F_ref=0.0448, ref_isNO2=True):
    """Calculates contrail radiative forcing from emissions of aviation NOx.

    Inputs:
        emissions:  Raw emissions data.
        frac:       fraction of total NOx emissions due to aviation.
    Keywords:
        E_ref:      reference-year emissions of aviation NOx, Mt/yr. 2.946
                    is 2005 emissions of aviation NOx in RCP4.5 measured in
                    Mt NO2/yr.
        F_ref:      Forcing from linear persistent contrails + contrail
                    induced cirrus. The default of 0.0448 W/m2 for 2005 is
                    from Lee et al, 2009 (Atmos. Environ.,
                    doi:10.1016/j.atmosenv.2009.04.024).
        ref_isNO2:  True if E_ref is in units of NO2 rather than N.
    """

    em_NOx   = emissions[:,8]
    factor = 1

    if ref_isNO2:
        factor = molwt.NO2/molwt.N
    return em_NOx*frac * F_ref/E_ref * factor

max_stars_repo_name: OMS-NetZero/FAIR-pro

max_issues_repo_name: OMS-NetZero/FAIR-pro

max_forks_repo_name: OMS-NetZero/FAIR-pro


text:
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Author: Tran Phan

import os
import time
import traceback

import numpy as np
from beamngpy import BeamNGpy, Scenario, Vehicle
from beamngpy.sensors import Electrics, Camera

from src import definitions
from src.modules import config
from src.modules import logger
from src.modules.config import root
import pandas as pd

import imageio

class DataCollector:

    def __init__(self):
        self.collected_data = {"image_name": [], "steering": []}
        self.scenario = Scenario(config.get("data_collecting.scenario_level"),
                                 config.get("data_collecting.scenario_name"))
        self.vehicle = Vehicle('ego_vehicle', model=config.get("data_collecting.vehicle_model"), licence='PYTHON')
        self.bng = BeamNGpy(config.get("beamNG.host"), int(config.get("beamNG.port")), home=config.get("beamNG.home"))

    def launch_beam_ng(self, mode="manual_mode"):
        # Create an ETK800 with the licence plate 'PYTHON'
        electrics = Electrics()
        # attach to get steering angles
        self.vehicle.attach_sensor('electrics', electrics)
        # attach to get images from camera
        self.vehicle.attach_sensor('front_cam', self.create_camera_sensor())
        # Add it to our scenario at this position and rotation
        # self.scenario.add_vehicle(self.vehicle)
        self.scenario.add_vehicle(self.vehicle, pos=tuple(map(float, config.get("data_collecting.pos").split(","))),
                                  rot_quat=tuple(map(float, config.get("data_collecting.rot_quat").split(","))))
        # Place files defining our scenario for the simulator to read
        self.scenario.make(self.bng)
        # Launch BeamNG.research
        self.bng.open()
        # Load and start our scenario
        self.bng.load_scenario(self.scenario)
        self.bng.start_scenario()
        if mode == "ai_mode":
            # Make the vehicle's AI span the map
            self.vehicle.ai_drive_in_lane(True)
            self.vehicle.ai_set_mode('span')

    def save_image_manually(self, cam_name='front_cam'):
        img = self.bng.poll_sensors(self.vehicle)[cam_name]['colour']
        steering = self.bng.poll_sensors(self.vehicle)['electrics']['steering']
        self.save_data(img, steering)

    @staticmethod
    def create_camera_sensor(pos=(-0.3, 2, 1.0), direction=(0, 1, 0), fov=100, res=None):
        # Set up camera sensor
        resolution = res
        if res is None:
            resolution = (
                int(config.get("data_collecting.image_width")),
                int(config.get("data_collecting.image_height")))

        pos = pos
        direction = direction
        fov = fov
        front_camera = Camera(pos, direction, fov, resolution,
                              colour=True, depth=True, annotation=True)
        return front_camera

    def save_data(self, img, steering, i: str = "0"):
        file_name = str(int(time.time())) + i + ".jpg"
        try:
            image_path = definitions.ROOT_DIR + config.get('data_collecting.data_path') + file_name
            imageio.imwrite(image_path, np.asarray(img.convert('RGB')), "jpg")
            self.collected_data["image_name"].append(file_name)
            self.collected_data["steering"].append(steering)
        except Exception as ex:
            logger.info(f"Error while saving data -- {ex}")
            raise Exception

    def collect_data(self, number_of_images: int, mode="manual_mode"):
        self.launch_beam_ng(mode)
        logger.info("Start after 3 seconds...")
        time.sleep(5)
        logger.info(f"Start collecting {config.get('data_collecting.number_of_images')} training images")
        i = 0
        exit_normally = True
        try:
            while i < number_of_images:
                # image is training image and steering is label
                img = self.bng.poll_sensors(self.vehicle)['front_cam']['colour']
                steering = self.bng.poll_sensors(self.vehicle)['electrics']['steering']
                logger.info(f"Saving data {i + 1} ...")
                self.save_data(img, steering, str(i))
                logger.info("Saved data successfully")
                i = i + 1
                time.sleep(int(config.get("data_collecting.sleep")))

        except Exception as ex:
            exit_normally = False
            logger.info(f"Error while collecting data {ex}")
        finally:
            self.bng.close()
            return exit_normally

    def save_csv_data(self):
        logger.info("Start saving csv file......")
        csv_path = definitions.ROOT_DIR + config.get('data_collecting.csv_path')
        df = pd.DataFrame(self.collected_data, columns=['image_name', 'steering'])
        if not os.path.isfile(csv_path):
            df.to_csv(csv_path, index=False, header=True)
        else:  # else it exists so append without writing the header
            df.to_csv(csv_path, index=False, mode='a', header=False)

def run_collector(number_of_images=0, mode="manual_mode"):
    root.file_config("config.yml")
    if not number_of_images:
        number_of_images = int(config.get("data_collecting.number_of_images"))

    collector = DataCollector()
    exit_ok = collector.collect_data(number_of_images, mode)
    # Save csv data
    if exit_ok:
        collector.save_csv_data()
        logger.info("Data collected successfully!")
    else:
        logger.info("Saving data...")
        collector.save_csv_data()
        logger.info(traceback.format_exc())

if __name__ == "__main__":
    run_collector()

max_stars_repo_name: phantran/NeuronCoverageDNN

max_issues_repo_name: phantran/beamng-neuron-coverage

max_forks_repo_name: phantran/beamng-neuron-coverage


text:
import sys
import time
import numpy as np
import random
import matplotlib.pyplot as plt
import queue
from scipy.io.wavfile import read as wavread
from scipy.signal import blackmanharris
from pysoundcard import *
from math import log
from sys import float_info
from collections import deque

#takes a numpy vector.
#i am not sure what is in the vector. Hmm.

def get_frequency(in_data, RATE, chunk):
	# Take the fft and square each value

	windowed = in_data[:,0] * blackmanharris(len(in_data))
	data_after_fft = np.fft.rfft(windowed)

	# Find the peak and interpolate to get a more accurate peak
	i = np.argmax(abs(data_after_fft))

	# Convert to equivalent frequency
	thefreq= chunk * i / len(windowed)

	data_in_decibels = map (lambda x : - 30 if x<sys.float_info.min else 20* log(x) , data_after_fft)
	peak_intensity = max(data_in_decibels)

	return thefreq, peak_intensity

block_length = 2048

class DrawingBuffer:
	def __init__(self, bufferSize):
		#self.buffer = deque([-30]*bufferSize, maxlen=bufferSize)
		self.bufferSize = bufferSize
		self.current = 0
		plt.axis([0, bufferSize, 0, 200])
		plt.ion()
		plt.show()

	def newFrame(self, frequency, intensity): # for now do nothing with intensity
		self.current = (self.current + 1 ) % self.bufferSize
		plt.scatter(self.current, frequency)
		plt.draw()

class DrawingBuffer2:
	def __init__(self, bufferSize):
		self.buffer = deque([-30]*bufferSize, maxlen=bufferSize)
		self.bufferSize = bufferSize
		self.current = 0
		plt.axis([0, bufferSize, -30, 200])
		plt.ion()
		plt.show()

	def newFrame(self, frequency, intensity): # for now do nothing with intensity
		self.buffer.append(frequency)
		plt.rese
		plt.plot(list(self.buffer))
		plt.draw()

def make_callback_that_draws(drawing_object):
	def callback(in_data, frame_count, time_info, status):
		drawing_object.newFrame(*get_frequency(in_data, frame_count, block_length))
		return (in_data, continue_flag)
	return callback

def make_callback_that_puts_into_queue(queue):
	def callback(in_data, frame_count, time_info, status):
		frequency, intensity = get_frequency(in_data, frame_count, block_length)
		queue.put((frequency, intensity,time.time()))
		return (in_data, continue_flag)
	return callback

def callback(in_data, frame_count, time_info, status):
	print(get_frequency(in_data, frame_count, 2048))
	return (in_data, continue_flag)

queue_for_the_stream = queue.Queue()

s = Stream(sample_rate=44100, block_length=block_length,output_device=False, callback=make_callback_that_puts_into_queue(queue_for_the_stream))

s.start()

while True:
	t = time.time()
	while(time.time()-t < 5.0):
		print((lambda x, y,z : (x,y,time.time()-z) )(*queue_for_the_stream.get()))
	time.sleep(10.0)

time.sleep(1000)
s.stop()

max_stars_repo_name: meclav/whistle

max_issues_repo_name: meclav/whistle

max_forks_repo_name: meclav/whistle


text:
import argparse
import sys
from omegaconf import OmegaConf
import numpy as np
import torch
from src.models.model import Network
import matplotlib.pyplot as plt
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from src.data.make_dataset import MNISTDataset
import hydra
import logging
import wandb

@hydra.main(config_path="../configs", config_name="config.yaml")

def main(config):
    test_x = torch.randn([64,784])
    N_classes = 10

    model = Network(config["hyperparameters"]["input_size"],
                    config["hyperparameters"]["output_size"],
                    config["hyperparameters"]["hidden_layers"],
                    config["hyperparameters"]["dropout_probability"])

    out = model(test_x)

    assert out.shape[0] == test_x.shape[0]
    assert out.shape[1] == N_classes

if __name__ == '__main__':
    main()


text:
#import libraries
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt
import numpy as np

#set the tank/compartment volume (in Litres)
volume = '12124'

x = np.array([20820, 43910, 5254, 6272, 4343, 2380, 3372, 5678, 1575, 2978, 3675, 4155, 7948, 10510, 3186, 9464, 7949, 3785, 5678, 7949, 34000, 16670, 3747, 2300, 4500,11356,3785,7570,13249, 36000, 22976, 34000, 34000, 31797.46, 31000, 36000, 16655, 4921,14763, 18927.05, 31797.4]) #L
y = np.array([54, 88.17, 18, 20, 16, 13, 13,16.3, 6.6, 13, 7.9, 15, 23, 24.8, 12, 24, 17.8, 8.83, 11.9, 19.8, 83.94, 41.2, 12.4, 7.6, 8.7,26.85,7.62,15.4,30.94, 84.6, 55.57, 85.60, 85.00,75.5302, 63.7, 84.1,38.46, 9.94, 31.31, 56.11, 79.61]) #m2

f = interp1d(x,y,fill_value="extrapolate")

type(f)

f(volume)
print ('A tank/compartment with a volume of {} L has an estimated exposed surface area of: {} m2.'.format(volume,f(volume)))

max_stars_repo_name: jphandrigan/exposedsurface

max_issues_repo_name: jphandrigan/exposedsurface

max_forks_repo_name: jphandrigan/exposedsurface


text:
import fidimag
from fidimag.common import init_vector
import numpy as np
import unittest

def m_init(pos, t):
    x, y, z = pos
    return (np.sin(t), np.cos(t), 0.0)

def test_init_vector_function():
    import fidimag
    import numpy as np
    # Create simulation
    nx = ny = nz = 2
    mesh = fidimag.common.CuboidMesh(nx=nx, ny=ny, nz=nz,
                                     dx=1, dy=1, dz=1)

    test_vec = np.zeros(3 * nx * ny * nz)
    test_vec[:] = init_vector(m_init, mesh, 3, False, 0)
    assert test_vec[0] < 1e-15
    assert test_vec[1] - 1.0 < 1e-15
    test_vec[:] = init_vector(m_init, mesh, 3, False, np.pi/2.0)
    assert test_vec[0] - 1.0 < 1e-15
    assert test_vec[1] < 1e-15

max_stars_repo_name: computationalmodelling/fidimag

max_issues_repo_name: computationalmodelling/fidimag

max_forks_repo_name: computationalmodelling/fidimag