Beispiel #1
0
    def movieRad(self,
                 cut=0.5,
                 levels=12,
                 cm='RdYlBu_r',
                 png=False,
                 step=1,
                 normed=False,
                 dpi=80,
                 bgcolor=None,
                 deminc=True,
                 removeMean=False,
                 precision='Float64',
                 contour=False,
                 mer=False):
        """
        Plotting function (it can also write the png files)

        .. warning:: the python bindings of `SHTns <https://bitbucket.org/bputigny/shtns-magic>`_ are mandatory to use this plotting function!

        :param levels: number of contour levels
        :type levels: int
        :param cm: name of the colormap
        :type cm: str
        :param cut: adjust the contour extrema to max(abs(data))*cut
        :type cut: float
        :param png: save the movie as a series of png files when
                    set to True
        :type png: bool
        :param dpi: dot per inch when saving PNGs
        :type dpi: int
        :param bgcolor: background color of the figure
        :type bgcolor: str
        :param normed: the colormap is rescaled every timestep when set to True,
                       otherwise it is calculated from the global extrema
        :type normed: bool
        :param step: the stepping between two timesteps
        :type step: int
        :param deminc: a logical to indicate if one wants do get rid of the
                       possible azimuthal symmetry
        :type deminc: bool
        :param precision: single or double precision
        :type precision: char
        :param contour: also display the solid contour levels when set to True
        :type contour: bool
        :param mer: display meridians and circles when set to True
        :type mer: bool
        :param removeMean: remove the time-averaged part when set to True
        :type removeMean: bool
        """

        if removeMean:
            dataCut = self.wlm - self.wlm.mean(axis=0)
        else:
            dataCut = self.wlm

        nlat = max(int(self.l_max_r * (3. / 2. / 2.) * 2.), 192)
        nphi = 2 * nlat / self.minc

        # Define spectral transform setup
        sh = SpectralTransforms(l_max=self.l_max_r,
                                minc=self.minc,
                                lm_max=self.lm_max_r,
                                n_theta_max=nlat)
        """
        # The python bindings of shtns are mandatory to use this function !!!
        import shtns

        # Define shtns setup
        sh = shtns.sht(int(self.l_max_r), int(self.m_max_r/self.minc),
                       mres=int(self.minc),
                       norm=shtns.sht_orthonormal | shtns.SHT_NO_CS_PHASE)
        """

        # Transform data on grid space
        data = np.zeros((self.nstep, nphi, nlat), precision)
        print('Spectral -> Spatial transform')
        for k in progressbar(range(self.nstep)):
            data[k, ...] = sh.spec_spat(dataCut[k, :] * self.ell *
                                        (self.ell + 1) / self.radius**2)
        print('Done')

        if png:
            plt.ioff()
            if not os.path.exists('movie'):
                os.mkdir('movie')
        else:
            plt.ion()

        if not normed:
            vmin = -max(abs(data.max()), abs(data.min()))
            vmin = cut * vmin
            vmax = -vmin
            cs = np.linspace(vmin, vmax, levels)

        th = np.linspace(np.pi / 2., -np.pi / 2., nlat)
        if deminc:
            phi = np.linspace(-np.pi, np.pi, self.minc * nphi + 1)
            xxout, yyout = hammer2cart(th, -np.pi)
            xxin, yyin = hammer2cart(th, np.pi)
        else:
            phi = np.linspace(-np.pi / self.minc, np.pi / self.minc, nphi)
            xxout, yyout = hammer2cart(th, -np.pi / self.minc)
            xxin, yyin = hammer2cart(th, np.pi / self.minc)
        ttheta, pphi = np.meshgrid(th, phi)
        xx, yy = hammer2cart(ttheta, pphi)
        if deminc:
            fig = plt.figure(figsize=(8, 4))
        else:
            fig = plt.figure(figsize=(8 / self.minc, 4))
        fig.subplots_adjust(top=0.99, right=0.99, bottom=0.01, left=0.01)
        ax = fig.add_subplot(111, frameon=False)

        if mer:
            theta = np.linspace(np.pi / 2, -np.pi / 2, nlat)
            meridians = np.r_[-120, -60, 0, 60, 120]
            circles = np.r_[60, 30, 0, -30, -60]

        for k in range(self.nstep):
            if k == 0:
                if normed:
                    vmin = -max(abs(data[k, ...].max()), abs(data[k,
                                                                  ...].min()))
                    vmin = cut * vmin
                    vmax = -vmin
                    cs = np.linspace(vmin, vmax, levels)
                if deminc:
                    dat = symmetrize(data[k, ...], self.minc)
                else:
                    dat = data[k, ...]
                im = ax.contourf(xx,
                                 yy,
                                 dat,
                                 cs,
                                 cmap=plt.get_cmap(cm),
                                 extend='both')
                if contour:
                    ax.contour(xx,
                               yy,
                               dat,
                               cs,
                               linestyles=['-', '-'],
                               colors=['k', 'k'],
                               linewidths=[0.7, 0.7])
                ax.plot(xxout, yyout, 'k-', lw=1.5)
                ax.plot(xxin, yyin, 'k-', lw=1.5)

                if mer:
                    for lat0 in circles:
                        x0, y0 = hammer2cart(lat0 * np.pi / 180., phi)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)
                    for lon0 in meridians:
                        x0, y0 = hammer2cart(theta, lon0 * np.pi / 180.)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)

                ax.axis('off')
                man = plt.get_current_fig_manager()
                man.canvas.draw()

                if png:
                    filename = 'movie/img%05d.png' % k
                    print('write %s' % filename)
                    #st = 'echo %i' % ivar + ' > movie/imgmax'
                    if bgcolor is not None:
                        fig.savefig(filename, facecolor=bgcolor, dpi=dpi)
                    else:
                        fig.savefig(filename, dpi=dpi)

            elif k != 0 and k % step == 0:
                if not png:
                    print(k)
                plt.cla()
                if normed:
                    vmin = -max(abs(data[k, ...].max()), abs(data[k,
                                                                  ...].min()))
                    vmin = cut * vmin
                    vmax = -vmin
                    cs = np.linspace(vmin, vmax, levels)
                if deminc:
                    dat = symmetrize(data[k, ...], self.minc)
                else:
                    dat = data[k, ...]
                im = ax.contourf(xx,
                                 yy,
                                 dat,
                                 cs,
                                 cmap=plt.get_cmap(cm),
                                 extend='both')
                if contour:
                    ax.contour(xx,
                               yy,
                               dat,
                               cs,
                               colors=['k'],
                               linestyles=['-', '-'],
                               linewidths=[0.7, 0.7])
                ax.plot(xxout, yyout, 'k-', lw=1.5)
                ax.plot(xxin, yyin, 'k-', lw=1.5)

                if mer:
                    for lat0 in circles:
                        x0, y0 = hammer2cart(lat0 * np.pi / 180., phi)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)
                    for lon0 in meridians:
                        x0, y0 = hammer2cart(theta, lon0 * np.pi / 180.)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)

                ax.axis('off')
                man.canvas.draw()
                if png:
                    filename = 'movie/img%05d.png' % k
                    print('write %s' % filename)
                    #st = 'echo %i' % ivar + ' > movie/imgmax'
                    if bgcolor is not None:
                        fig.savefig(filename, facecolor=bgcolor, dpi=dpi)
                    else:
                        fig.savefig(filename, dpi=dpi)
    def movieCmb(self,
                 cut=0.5,
                 levels=12,
                 cm='RdYlBu_r',
                 png=False,
                 step=1,
                 normed=False,
                 dpi=80,
                 bgcolor=None,
                 deminc=True,
                 removeMean=False,
                 precision=np.float64,
                 contour=False,
                 mer=False):
        """
        Plotting function (it can also write the png files)

        :param levels: number of contour levels
        :type levels: int
        :param cm: name of the colormap
        :type cm: str
        :param cut: adjust the contour extrema to max(abs(data))*cut
        :type cut: float
        :param png: save the movie as a series of png files when
                    set to True
        :type png: bool
        :param dpi: dot per inch when saving PNGs
        :type dpi: int
        :param bgcolor: background color of the figure
        :type bgcolor: str
        :param normed: the colormap is rescaled every timestep when set to True,
                       otherwise it is calculated from the global extrema
        :type normed: bool
        :param step: the stepping between two timesteps
        :type step: int
        :param deminc: a logical to indicate if one wants do get rid of the
                       possible azimuthal symmetry
        :type deminc: bool
        :param precision: single or double precision
        :type precision: char
        :param contour: also display the solid contour levels when set to True
        :type contour: bool
        :param mer: display meridians and circles when set to True
        :type mer: bool
        :param removeMean: remove the time-averaged part when set to True
        :type removeMean: bool
        """

        if removeMean:
            blmCut = self.blm - self.blm.mean(axis=0)
        else:
            blmCut = self.blm

        nlat = int(max(int(self.l_max_cmb * (3. / 2. / 2.) * 2.), 192))
        if np.mod(nlat, 2) == 1:
            nlat += 1
        nphi = int(2 * nlat / self.minc)

        # Define spectral transform setup
        sh = SpectralTransforms(l_max=self.l_max_cmb,
                                minc=self.minc,
                                lm_max=self.lm_max_cmb,
                                n_theta_max=nlat)

        # Transform data on grid space
        BrCMB = np.zeros((self.nstep, nphi, nlat), precision)
        print('Spectral -> Spatial transform')
        for k in progressbar(range(self.nstep)):
            BrCMB[k, ...] = sh.spec_spat(blmCut[k, :] * self.ell *
                                         (sh.ell + 1) / self.rcmb**2)
        print('Done')

        if png:
            plt.ioff()
            if not os.path.exists('movie'):
                os.mkdir('movie')
        else:
            plt.ion()

        if not normed:
            vmin = -max(abs(BrCMB.max()), abs(BrCMB.min()))
            vmin = cut * vmin
            vmax = -vmin
            cs = np.linspace(vmin, vmax, levels)

        th = np.linspace(np.pi / 2., -np.pi / 2., nlat)
        if deminc:
            phi = np.linspace(-np.pi, np.pi, self.minc * nphi + 1)
            xxout, yyout = hammer2cart(th, -np.pi)
            xxin, yyin = hammer2cart(th, np.pi)
        else:
            phi = np.linspace(-np.pi / self.minc, np.pi / self.minc, nphi)
            xxout, yyout = hammer2cart(th, -np.pi / self.minc)
            xxin, yyin = hammer2cart(th, np.pi / self.minc)
        ttheta, pphi = np.meshgrid(th, phi)
        xx, yy = hammer2cart(ttheta, pphi)
        if deminc:
            fig = plt.figure(figsize=(8, 4))
        else:
            fig = plt.figure(figsize=(8 / self.minc, 4))
        fig.subplots_adjust(top=0.99, right=0.99, bottom=0.01, left=0.01)
        ax = fig.add_subplot(111, frameon=False)

        if mer:
            theta = np.linspace(np.pi / 2, -np.pi / 2, nlat)
            meridians = np.r_[-120, -60, 0, 60, 120]
            circles = np.r_[60, 30, 0, -30, -60]

        for k in range(self.nstep):
            if k == 0:
                if normed:
                    vmin = -max(abs(BrCMB[k, ...].max()),
                                abs(BrCMB[k, ...].min()))
                    vmin = cut * vmin
                    vmax = -vmin
                    cs = np.linspace(vmin, vmax, levels)
                if deminc:
                    dat = symmetrize(BrCMB[k, ...], self.minc)
                else:
                    dat = BrCMB[k, ...]
                im = ax.contourf(xx,
                                 yy,
                                 dat,
                                 cs,
                                 cmap=plt.get_cmap(cm),
                                 extend='both')
                if contour:
                    ax.contour(xx,
                               yy,
                               dat,
                               cs,
                               linestyles=['-', '-'],
                               colors=['k', 'k'],
                               linewidths=[0.7, 0.7])
                ax.plot(xxout, yyout, 'k-', lw=1.5)
                ax.plot(xxin, yyin, 'k-', lw=1.5)
                #ax.text(0.12, 0.9, 't={:.6f}'.format(self.time[0]), fontsize=16,
                #horizontalalignment='right',
                #verticalalignment='center', transform = ax.transAxes)

                if mer:
                    for lat0 in circles:
                        x0, y0 = hammer2cart(lat0 * np.pi / 180., phi)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)
                    for lon0 in meridians:
                        x0, y0 = hammer2cart(theta, lon0 * np.pi / 180.)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)

                ax.axis('off')
                man = plt.get_current_fig_manager()
                man.canvas.draw()
            if k != 0 and k % step == 0:
                if not png:
                    print(k)
                plt.cla()
                if normed:
                    vmin = -max(abs(BrCMB[k, ...].max()),
                                abs(BrCMB[k, ...].min()))
                    vmin = cut * vmin
                    vmax = -vmin
                    cs = np.linspace(vmin, vmax, levels)
                if deminc:
                    dat = symmetrize(BrCMB[k, ...], self.minc)
                else:
                    dat = BrCMB[k, ...]
                im = ax.contourf(xx,
                                 yy,
                                 dat,
                                 cs,
                                 cmap=plt.get_cmap(cm),
                                 extend='both')
                if contour:
                    ax.contour(xx,
                               yy,
                               dat,
                               cs,
                               colors=['k'],
                               linestyles=['-', '-'],
                               linewidths=[0.7, 0.7])
                ax.plot(xxout, yyout, 'k-', lw=1.5)
                ax.plot(xxin, yyin, 'k-', lw=1.5)
                #ax.text(0.12, 0.9, 't={:.6f}'.format(self.time[k]), fontsize=16,
                #horizontalalignment='right',
                #verticalalignment='center', transform = ax.transAxes)

                if mer:
                    for lat0 in circles:
                        x0, y0 = hammer2cart(lat0 * np.pi / 180., phi)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)
                    for lon0 in meridians:
                        x0, y0 = hammer2cart(theta, lon0 * np.pi / 180.)
                        ax.plot(x0, y0, 'k:', linewidth=0.7)

                ax.axis('off')
                man.canvas.draw()
            if png:
                filename = 'movie/img{:05d}.png'.format(k)
                print('write {}'.format(filename))
                #st = 'echo {}'.format(ivar) + ' > movie/imgmax'
                if bgcolor is not None:
                    fig.savefig(filename, facecolor=bgcolor, dpi=dpi)
                else:
                    fig.savefig(filename, dpi=dpi)