def rcut2vtk(self, fname, data, r, name): """ This routine transforms a radial cut of dimension (n_phi,n_theta) into a vts file. :param fname: file name :type fname: str :param data: input data :type data: numpy.ndarray :param r: radius of the selected cut :type r: float :param name: name of the physical field stored in the vts file :type name: str """ data = symmetrize(data, self.minc) phi = np.linspace(0., 2.*np.pi, data.shape[0]) X = np.zeros((data.shape[0], data.shape[1], 1), dtype=np.float32) Y = np.zeros_like(X) Z = np.zeros_like(X) ttheta, pphi = np.meshgrid(self.theta, phi) X[:, :, 0] = r*np.sin(ttheta)*np.cos(pphi) Y[:, :, 0] = r*np.sin(ttheta)*np.sin(pphi) Z[:, :, 0] = r*np.cos(ttheta) dat = np.zeros_like(X) point_data = {} dat[..., 0] = data point_data[name] = dat gridToVTK(fname, X, Y, Z, pointData=point_data) print('Store {}.vts'.format(fname))
def equat2vtk(self, fname, data, name): """ This routine transforms an equatorial cut of dimension (n_phi,n_r) into a vts file. :param fname: file name :type fname: str :param data: input data :type data: numpy.ndarray :param name: name of the physical field stored in the vts file :type name: str """ data = symmetrize(data, self.minc) phi = np.linspace(0., 2.*np.pi, data.shape[0]) X = np.zeros((1, data.shape[0], data.shape[1]), dtype=np.float32) Y = np.zeros_like(X) Z = np.zeros_like(X) rr, pphi = np.meshgrid(self.radius, phi) X[0, :, :] = rr*np.cos(pphi) Y[0, :, :] = rr*np.sin(pphi) dat = np.zeros_like(X) dat[0, ...] = data point_data = {} point_data[name] = dat gridToVTK(fname, X, Y, Z, pointData=point_data) print('Store {}.vts'.format(fname))
def scal3D2vtk(self, fname, data, name): """ This routine transforms a 3-D scalar field of dimension (n_phi,n_theta,n_r) into a vts file. :param fname: file name :type fname: str :param data: input data :type data: numpy.ndarray :param r: radius of the selected cut :type r: float :param name: name of the physical field stored in the vts file :type name: str """ data = symmetrize(data, self.minc) phi = np.linspace(0., 2.*np.pi, data.shape[0]) X = np.zeros(data.shape, dtype=np.float32) Y = np.zeros_like(X) Z = np.zeros_like(X) ttheta, pphi, rr = np.meshgrid(self.theta, phi, self.radius) X = rr*np.sin(ttheta)*np.cos(pphi) Y = rr*np.sin(ttheta)*np.sin(pphi) Z = rr*np.cos(ttheta)*np.ones_like(pphi) point_data = {} point_data[name] = data gridToVTK(fname, X, Y, Z, pointData=point_data) print('Store {}.vts'.format(fname))
def scal3D2vtk(self, fname, data, name): """ This routine transforms a 3-D scalar field of dimension (n_phi,n_theta,n_r) into a vts file. :param fname: file name :type fname: str :param data: input data :type data: numpy.ndarray :param r: radius of the selected cut :type r: float :param name: name of the physical field stored in the vts file :type name: str """ if self.rmin > 0 and self.rmax > 0: mask = np.where( abs(self.radius - self.rmin) == abs(self.radius - self.rmin).min(), 1, 0) idx2 = np.nonzero(mask)[0][0] mask = np.where( abs(self.radius - self.rmax) == abs(self.radius - self.rmax).min(), 1, 0) idx1 = np.nonzero(mask)[0][0] nr = idx2 - idx1 + 1 else: nr = data.shape[1] idx1 = 0 idx2 = data.shape[1] data = symmetrize(data, self.minc) phi = np.linspace(0., 2. * np.pi, data.shape[0]) X = np.zeros(data.shape, dtype=np.float32) Y = np.zeros_like(X) Z = np.zeros_like(X) ttheta, pphi, rr = np.meshgrid(self.theta, phi, self.radius[idx1:idx2 + 1]) X = rr * np.sin(ttheta) * np.cos(pphi) Y = rr * np.sin(ttheta) * np.sin(pphi) Z = rr * np.cos(ttheta) * np.ones_like(pphi) point_data = {} point_data[name] = data[..., idx1:idx2] gridToVTK(fname, X, Y, Z, pointData=point_data) print('Store {}.vts'.format(fname))
def equat2vtk(self, fname, data, name): """ This routine transforms an equatorial cut of dimension (n_phi,n_r) into a vts file. :param fname: file name :type fname: str :param data: input data :type data: numpy.ndarray :param name: name of the physical field stored in the vts file :type name: str """ if self.rmin > 0 and self.rmax > 0: mask = np.where( abs(self.radius - self.rmin) == abs(self.radius - self.rmin).min(), 1, 0) idx2 = np.nonzero(mask)[0][0] mask = np.where( abs(self.radius - self.rmax) == abs(self.radius - self.rmax).min(), 1, 0) idx1 = np.nonzero(mask)[0][0] nr = idx2 - idx1 + 1 else: nr = data.shape[1] idx1 = 0 idx2 = data.shape[1] data = symmetrize(data, self.minc) phi = np.linspace(0., 2. * np.pi, data.shape[0]) X = np.zeros((1, data.shape[0], nr), dtype=np.float32) Y = np.zeros_like(X) Z = np.zeros_like(X) rr, pphi = np.meshgrid(self.radius[idx1:idx2 + 1], phi) X[0, :, :] = rr * np.cos(pphi) Y[0, :, :] = rr * np.sin(pphi) dat = np.zeros_like(X) dat[0, ...] = data[:, idx1:idx2 + 1] point_data = {} point_data[name] = dat gridToVTK(fname, X, Y, Z, pointData=point_data) print('Store {}.vts'.format(fname))
def timeLongitude(self, ifield=0, removeMean=True, lat0=0., levels=12, cm='RdYlBu_r', deminc=True): """ Plot the time-longitude diagram (input latitude can be chosen) :param ifield: in case of a multiple-field movie file, you can change the default field displayed using the parameter ifield :type ifield: int :param lat0: value of the latitude :type lat0: float :param levels: number of contour levels :type levels: int :param cm: name of the colormap :type cm: str :param deminc: a logical to indicate if one wants do get rid of the possible azimuthal symmetry :type deminc: bool :param removeMean: remove the time-averaged part when set to True :type removeMean: bool """ if removeMean: datCut = self.data[ifield, ...] - self.data[ifield, ...].mean(axis=0) else: datCut = self.data[ifield, ...] th = np.linspace(np.pi / 2., -np.pi / 2., self.n_theta_max) lat0 *= np.pi / 180. mask = np.where(abs(th - lat0) == abs(th - lat0).min(), 1, 0) idx = np.nonzero(mask)[0][0] th = np.linspace(np.pi / 2., -np.pi / 2., self.n_theta_max) if deminc: phi = np.linspace(-np.pi, np.pi, self.minc * self.n_phi_tot + 1) else: phi = np.linspace(-np.pi / self.minc, np.pi / self.minc, self.n_phi_tot) if deminc: dat = np.zeros((self.nvar, self.minc * self.n_phi_tot + 1), np.float64) else: dat = np.zeros((self.nvar, self.n_phi_tot), np.float64) for k in range(self.nvar): if deminc: dat[k, :] = symmetrize(datCut[k, :, idx], self.minc) else: dat[k, :] = datCut[k, :, idx] fig = plt.figure() ax = fig.add_subplot(111) vmin = -max(abs(dat.max()), abs(dat.min())) vmax = -vmin cs = np.linspace(vmin, vmax, levels) ax.contourf(phi, self.time, dat, cs, cmap=plt.get_cmap(cm)) ax.set_xlabel('Longitude') ax.set_ylabel('Time') m_max = self.n_phi_tot / 3 w2 = np.fft.fft2(dat) w2 = abs(w2[1:self.nvar / 2 + 1, 0:m_max + 1]) dw = 2. * np.pi / (self.time[-1] - self.time[0]) omega = dw * np.arange(self.nvar) omega = omega[1:self.nvar / 2 + 1] ms = np.arange(m_max + 1) fig1 = plt.figure() ax1 = fig1.add_subplot(111) ax1.contourf(ms, omega, np.log10(w2), 65, cmap=plt.get_cmap('jet')) ax1.set_yscale('log') ax1.set_xlabel(r'Azimuthal wavenumber') ax1.set_ylabel(r'Frequency')
def plot(self, ifield=0, cut=0.5, levels=12, cmap='RdYlBu_r', png=False, step=1, normed=False, dpi=80, bgcolor=None, deminc=True, ic=False): """ Plotting function (it can also write the png files) :param ifield: in case of a multiple-field movie file, you can change the default field displayed using the parameter ifield :type ifield: int :param levels: number of contour levels :type levels: int :param cmap: name of the colormap :type cmap: 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 """ if png: plt.ioff() if not os.path.exists('movie'): os.mkdir('movie') else: plt.ion() if not normed: vmin = -max(abs(self.data[ifield, ...].max()), abs(self.data[ifield, ...].min())) vmin = cut * vmin vmax = -vmin # vmin, vmax = self.data.min(), self.data.max() cs = np.linspace(vmin, vmax, levels) if self.surftype == 'phi_constant': if self.n_theta_plot == self.n_theta_max: th = np.linspace(np.pi / 2., -np.pi / 2., self.n_theta_max) fig = plt.figure(figsize=(4, 8)) th0 = th else: th0 = np.linspace(np.pi / 2., -np.pi / 2., self.n_theta_max) th1 = np.linspace(np.pi / 2., 3. * np.pi / 2., self.n_theta_max) th = np.concatenate((th0, th1)) fig = plt.figure(figsize=(6.5, 6)) # Plotting trick using th0 th0 = np.linspace(np.pi / 2, np.pi / 2 + 2. * np.pi, 2 * self.n_theta_max) rr, tth = np.meshgrid(self.radius, th) xx = rr * np.cos(tth) yy = rr * np.sin(tth) xxout = rr.max() * np.cos(th0) yyout = rr.max() * np.sin(th0) xxin = rr.min() * np.cos(th0) yyin = rr.min() * np.sin(th0) if ic: rr, tth = np.meshgrid(self.radius_ic, th) xx_ic = rr * np.cos(tth) yy_ic = rr * np.sin(tth) elif self.surftype == 'r_constant': th = np.linspace(np.pi / 2., -np.pi / 2., self.n_theta_max) if deminc: phi = np.linspace(-np.pi, np.pi, self.n_phi_tot * self.minc + 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, self.n_phi_tot) 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) fig = plt.figure(figsize=(8, 4)) elif self.surftype == 'theta_constant': if deminc: phi = np.linspace(0., 2. * np.pi, self.n_phi_tot * self.minc + 1) else: phi = np.linspace(0., 2. * np.pi / self.minc, self.n_phi_tot) rr, pphi = np.meshgrid(self.radius, phi) xx = rr * np.cos(pphi) yy = rr * np.sin(pphi) xxout = rr.max() * np.cos(pphi) yyout = rr.max() * np.sin(pphi) xxin = rr.min() * np.cos(pphi) yyin = rr.min() * np.sin(pphi) if ic: rr, pphi = np.meshgrid(self.radius_ic, phi) xx_ic = rr * np.cos(pphi) yy_ic = rr * np.sin(pphi) fig = plt.figure(figsize=(6, 6)) elif self.surftype == '3d volume': self.data = self.data[ifield, ..., 0] th = np.linspace(np.pi / 2., -np.pi / 2., self.n_theta_max) phi = np.linspace(-np.pi, np.pi, self.n_phi_tot) ttheta, pphi = np.meshgrid(th, phi) xx, yy = hammer2cart(ttheta, pphi) xxout, yyout = hammer2cart(th, -np.pi) xxin, yyin = hammer2cart(th, np.pi) fig = plt.figure(figsize=(8, 4)) fig.subplots_adjust(top=0.99, right=0.99, bottom=0.01, left=0.01) ax = fig.add_subplot(111, frameon=False) for k in range(self.nvar): if k == 0: if normed: vmin = -max(abs(self.data[ifield, k, ...].max()), abs(self.data[ifield, k, ...].min())) vmin = cut * vmin vmax = -vmin cs = np.linspace(vmin, vmax, levels) if self.surftype in ['r_constant', 'theta_constant']: if deminc: dat = symmetrize(self.data[ifield, k, ...], self.minc) if ic: datic = symmetrize(self.data_ic[ifield, k, ...], self.minc) else: dat = self.data[ifield, k, ...] if ic: datic = self.data_ic[ifield, k, ...] im = ax.contourf(xx, yy, dat, cs, cmap=cmap, extend='both') if ic: ax.contourf(xx_ic, yy_ic, datic, cs, cmap=cmap, extend='both') else: im = ax.contourf(xx, yy, self.data[ifield, k, ...], cs, cmap=cmap, extend='both') if ic: im_ic = ax.contourf(xx_ic, yy_ic, self.data_ic[ifield, k, ...], cs, cmap=cmap, extend='both') ax.plot(xxout, yyout, 'k-', lw=1.5) ax.plot(xxin, yyin, 'k-', lw=1.5) ax.axis('off') man = plt.get_current_fig_manager() man.canvas.draw() if k != 0 and k % step == 0: if not png: print(k + self.var2 - self.nvar) plt.cla() if normed: vmin = -max(abs(self.data[ifield, k, ...].max()), abs(self.data[ifield, k, ...].min())) vmin = cut * vmin vmax = -vmin cs = np.linspace(vmin, vmax, levels) if self.surftype in ['r_constant', 'theta_constant']: if deminc: dat = symmetrize(self.data[ifield, k, ...], self.minc) if ic: datic = symmetrize(self.data_ic[ifield, k, ...], self.minc) else: dat = self.data[ifield, k, ...] if ic: datic = self.data_ic[ifield, k, ...] ax.contourf(xx, yy, dat, cs, cmap=cmap, extend='both') if ic: ax.contourf(xx_ic, yy_ic, datic, cs, cmap=cmap, extend='both') else: ax.contourf(xx, yy, self.data[ifield, k, ...], cs, cmap=cmap, extend='both') if ic: ax.contourf(xx_ic, yy_ic, self.data_ic[ifield, k, ...], cs, cmap=cmap, extend='both') ax.plot(xxout, yyout, 'k-', lw=1.5) ax.plot(xxin, yyin, 'k-', lw=1.5) 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)
def timeLongitude(self, removeMean=True, lat0=0.0, levels=12, cm="RdYlBu_r", deminc=True): """ Plot the time-longitude diagram (input latitude can be chosen) :param lat0: value of the latitude :type lat0: float :param levels: number of contour levels :type levels: int :param cm: name of the colormap :type cm: str :param deminc: a logical to indicate if one wants do get rid of the possible azimuthal symmetry :type deminc: bool :param removeMean: remove the time-averaged part when set to True :type removeMean: bool """ if removeMean: datCut = self.data - self.data.mean(axis=0) else: datCut = self.data th = np.linspace(np.pi / 2.0, -np.pi / 2.0, self.n_theta_max) lat0 *= np.pi / 180.0 mask = np.where(abs(th - lat0) == abs(th - lat0).min(), 1, 0) idx = np.nonzero(mask)[0][0] th = np.linspace(np.pi / 2.0, -np.pi / 2.0, self.n_theta_max) if deminc: phi = np.linspace(-np.pi, np.pi, self.minc * self.n_phi_tot + 1) else: phi = np.linspace(-np.pi / self.minc, np.pi / self.minc, self.n_phi_tot) if deminc: dat = np.zeros((self.nvar, self.minc * self.n_phi_tot + 1), "Float64") else: dat = np.zeros((self.nvar, self.n_phi_tot), "Float64") for k in range(self.nvar): if deminc: dat[k, :] = symmetrize(datCut[k, :, idx], self.minc) else: dat[k, :] = datCut[k, :, idx] fig = plt.figure() ax = fig.add_subplot(111) vmin = -max(abs(dat.max()), abs(dat.min())) vmax = -vmin cs = np.linspace(vmin, vmax, levels) ax.contourf(phi, self.time, dat, cs, cmap=plt.get_cmap(cm)) ax.set_xlabel("Longitude") ax.set_ylabel("Time") m_max = self.n_phi_tot / 3 w2 = np.fft.fft2(dat) w2 = abs(w2[1 : self.nvar / 2 + 1, 0 : m_max + 1]) dw = 2.0 * np.pi / (self.time[-1] - self.time[0]) omega = dw * np.arange(self.nvar) omega = omega[1 : self.nvar / 2 + 1] ms = np.arange(m_max + 1) fig1 = plt.figure() ax1 = fig1.add_subplot(111) ax1.contourf(ms, omega, np.log10(w2), 65, cmap=plt.get_cmap("jet")) ax1.set_yscale("log") # ax1.set_xlim(0,13) ax1.set_xlabel(r"Azimuthal wavenumber") ax1.set_ylabel(r"Frequency")
def plot( self, cut=0.5, levels=12, cmap="RdYlBu_r", png=False, step=1, normed=False, dpi=80, bgcolor=None, deminc=True ): """ Plotting function (it can also write the png files) :param levels: number of contour levels :type levels: int :param cmap: name of the colormap :type cmap: 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 """ if png: plt.ioff() if not os.path.exists("movie"): os.mkdir("movie") else: plt.ion() if not normed: vmin = -max(abs(self.data.max()), abs(self.data.min())) vmin = cut * vmin vmax = -vmin # vmin, vmax = self.data.min(), self.data.max() cs = np.linspace(vmin, vmax, levels) if self.surftype == "phi_constant": # if self.movtype in [1, 7]: # th = np.linspace(0., 2.*np.pi, 2*self.n_theta_max) # else: th = np.linspace(np.pi / 2.0, -np.pi / 2.0, self.n_theta_max) rr, tth = np.meshgrid(self.radius, th) xx = rr * np.cos(tth) yy = rr * np.sin(tth) xxout = rr.max() * np.cos(th) yyout = rr.max() * np.sin(th) xxin = rr.min() * np.cos(th) yyin = rr.min() * np.sin(th) fig = plt.figure(figsize=(4, 8)) elif self.surftype == "r_constant": th = np.linspace(np.pi / 2.0, -np.pi / 2.0, self.n_theta_max) if deminc: phi = np.linspace(-np.pi, np.pi, self.n_phi_tot * self.minc + 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, self.n_phi_tot) 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) fig = plt.figure(figsize=(8, 4)) elif self.surftype == "theta_constant": if deminc: phi = np.linspace(0.0, 2.0 * np.pi, self.n_phi_tot * self.minc + 1) else: phi = np.linspace(0.0, 2.0 * np.pi / self.minc, self.n_phi_tot) rr, pphi = np.meshgrid(self.radius, phi) xx = rr * np.cos(pphi) yy = rr * np.sin(pphi) xxout = rr.max() * np.cos(pphi) yyout = rr.max() * np.sin(pphi) xxin = rr.min() * np.cos(pphi) yyin = rr.min() * np.sin(pphi) fig = plt.figure(figsize=(6, 6)) elif self.surftype == "3d volume": self.data = self.data[..., 0] th = np.linspace(np.pi / 2.0, -np.pi / 2.0, self.n_theta_max) phi = np.linspace(-np.pi, np.pi, self.n_phi_tot) ttheta, pphi = np.meshgrid(th, phi) xx, yy = hammer2cart(ttheta, pphi) xxout, yyout = hammer2cart(th, -np.pi) xxin, yyin = hammer2cart(th, np.pi) fig = plt.figure(figsize=(8, 4)) fig.subplots_adjust(top=0.99, right=0.99, bottom=0.01, left=0.01) ax = fig.add_subplot(111, frameon=False) for k in range(self.nvar): if k == 0: if normed: vmin = -max(abs(self.data[k, ...].max()), abs(self.data[k, ...].min())) vmin = cut * vmin vmax = -vmin cs = np.linspace(vmin, vmax, levels) if self.surftype in ["r_constant", "theta_constant"]: if deminc: dat = symmetrize(self.data[k, ...], self.minc) else: dat = self.data[k, ...] im = ax.contourf(xx, yy, dat, cs, cmap=cmap, extend="both") else: im = ax.contourf(xx, yy, self.data[k, ...], cs, cmap=cmap, extend="both") ax.plot(xxout, yyout, "k-", lw=1.5) ax.plot(xxin, yyin, "k-", lw=1.5) ax.axis("off") man = plt.get_current_fig_manager() man.canvas.draw() if k != 0 and k % step == 0: if not png: print(k + self.var2 - self.nvar) plt.cla() if normed: vmin = -max(abs(self.data[k, ...].max()), abs(self.data[k, ...].min())) vmin = cut * vmin vmax = -vmin cs = np.linspace(vmin, vmax, levels) if self.surftype in ["r_constant", "theta_constant"]: if deminc: dat = symmetrize(self.data[k, ...], self.minc) else: dat = self.data[k, ...] im = ax.contourf(xx, yy, dat, cs, cmap=cmap, extend="both") else: im = ax.contourf(xx, yy, self.data[k, ...], cs, cmap=cmap, extend="both") ax.plot(xxout, yyout, "k-", lw=1.5) ax.plot(xxin, yyin, "k-", lw=1.5) 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 __init__(self, file=None, step=1, lastvar=None, nvar='all', nrout=48, ratio_out=2., potExtra=False, precision=np.float32): """ :param file: file name :type file: str :param nvar: the number of timesteps of the movie file we want to plot starting from the last line :type nvar: int :param lastvar: the number of the last timestep to be read :type lastvar: int :param step: the stepping between two timesteps :type step: int :param precision: precision of the input file, np.float32 for single precision, np.float64 for double precision :type precision: str :param potExtra: when set to True, potential extrapolation of the magnetic field outside the fluid domain is also computed :type potExtra: bool :param ratio_out: ratio of desired external radius to the CMB radius. This is is only used when potExtra=True :type ratio_out: float :param nrout: number of additional radial grid points to compute the potential extrapolation. This is only used when potExtra=True :type nrout: int """ if file is None: dat = glob.glob('*_mov.*') str1 = 'Which movie do you want ?\n' for k, movie in enumerate(dat): str1 += ' %i) %s\n' % (k+1, movie) index = int(input(str1)) try: filename = dat[index-1] except IndexError: print('Non valid index: %s has been chosen instead' % dat[0]) filename = dat[0] else: filename = file mot = re.compile(r'.*_mov\.(.*)') end = mot.findall(filename)[0] # DETERMINE THE NUMBER OF LINES BY READING THE LOG FILE logfile = open('log.%s' % end, 'r') mot = re.compile(r' ! WRITING MOVIE FRAME NO\s*(\d*).*') for line in logfile.readlines(): if mot.match(line): nlines = int(mot.findall(line)[0]) logfile.close() if lastvar is None: self.var2 = nlines else: self.var2 = lastvar if str(nvar) == 'all': self.nvar = nlines self.var2 = nlines else: self.nvar = nvar vecNames = np.r_[3] scalNames = np.r_[-1] # READ the movie file infile = npfile(filename, endian='B') # HEADER version = infile.fort_read('|S64') n_type, n_surface, const, n_fields = infile.fort_read(precision) n_fields = int(n_fields) n_surface = int(n_surface) if n_fields == 1: movtype = infile.fort_read(precision) self.movtype = int(movtype) else: movtype = infile.fort_read(precision) # RUN PARAMETERS runid = infile.fort_read('|S64') n_r_mov_tot, n_r_max, n_theta_max, n_phi_tot, minc, ra, \ ek, pr, prmag, radratio, tScale = infile.fort_read(precision) minc = int(minc) self.n_r_max = int(n_r_max) self.n_theta_max = int(n_theta_max) self.n_phi_tot = int(n_phi_tot) n_r_mov_tot = int(n_r_mov_tot) # GRID if potExtra: self.radius = np.zeros((n_r_mov_tot+1+nrout), np.float32) else: self.radius = np.zeros((n_r_mov_tot+2), np.float32) tmp = infile.fort_read(precision)/(1.-radratio) rcmb = tmp[0] self.radius[2:n_r_mov_tot+2] = tmp[::-1] self.theta = infile.fort_read(precision) self.phi = infile.fort_read(precision) shape = (n_r_mov_tot+2, self.n_theta_max, self.n_phi_tot) self.time = np.zeros(self.nvar, precision) if potExtra: scals = np.zeros((n_r_mov_tot+1+nrout, self.n_theta_max, self.n_phi_tot*minc+1, 1), np.float32) else: scals = np.zeros((n_r_mov_tot+2, self.n_theta_max, self.n_phi_tot*minc+1, 1), np.float32) vecr = np.zeros_like(scals) vect = np.zeros_like(scals) vecp = np.zeros_like(scals) # Potential extrapolation radii = self.radius scals = scals.T scals[0, ...] = radii startdir = os.getcwd() if not os.path.exists('vtsFiles'): os.mkdir('vtsFiles') os.chdir('vtsFiles') else: os.chdir('vtsFiles') for i in range(self.var2-self.nvar): n_frame, t_movieS, omega_ic, omega_ma, movieDipColat, \ movieDipLon, movieDipStrength, \ movieDipStrengthGeo = infile.fort_read(precision) tmp = infile.fort_read(precision, shape=shape) vecr[0:n_r_mov_tot+2, :, :, 0] = symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vect[0:n_r_mov_tot+2, :, :, 0] = symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vecp[0:n_r_mov_tot+2, :, :, 0] = symmetrize(tmp, minc, reversed=True) for k in range(self.nvar): n_frame, t_movieS, omega_ic, omega_ma, movieDipColat, \ movieDipLon, movieDipStrength, \ movieDipStrengthGeo = infile.fort_read(precision) self.time[k] = t_movieS if k % step == 0: # print(k+self.var2-self.nvar) tmp = infile.fort_read(precision, shape=shape) brCMB = np.zeros((self.n_theta_max, self.n_phi_tot), np.float32) brCMB = tmp[0, :, :] brCMB = brCMB.T if potExtra: vecr[nrout-1:nrout+n_r_mov_tot+2, :, :, 0] = \ symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vect[nrout-1:nrout+n_r_mov_tot+2, :, :, 0] = \ symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vecp[nrout-1:nrout+n_r_mov_tot+2, :, :, 0] = \ symmetrize(tmp, minc, reversed=True) else: vecr[0:n_r_mov_tot+2, :, :, 0] = \ symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vect[0:n_r_mov_tot+2, :, :, 0] = \ symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vecp[0:n_r_mov_tot+2, :, :, 0] = \ symmetrize(tmp, minc, reversed=True) filename = 'B3D_%05d' % k vecr = vecr[::-1, ...] vect = vect[::-1, ...] vecp = vecp[::-1, ...] br = vecr.T bt = vect.T bp = vecp.T if potExtra: pot = ExtraPot(rcmb, brCMB, minc, ratio_out=ratio_out, nrout=nrout, cutCMB=True, deminc=True) br[0, ..., n_r_mov_tot+2:] = pot.brout bt[0, ..., n_r_mov_tot+2:] = pot.btout bp[0, ..., n_r_mov_tot+2:] = pot.bpout radii[n_r_mov_tot+2:] = pot.rout else: radii = self.radius vts(filename, radii, br, bt, bp, scals, scalNames, vecNames, 1) print('write %s.vts' % filename) else: # Otherwise we read vecr = infile.fort_read(precision, shape=shape) vect = infile.fort_read(precision, shape=shape) vecp = infile.fort_read(precision, shape=shape) os.chdir(startdir)
def plot(self, cut=0.5, levels=12, cmap='RdYlBu_r', png=False, step=1, normed=False, dpi=80, bgcolor=None, deminc=True): """ Plotting function (it can also write the png files) :param levels: number of contour levels :type levels: int :param cmap: name of the colormap :type cmap: 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 """ if png: P.ioff() if not os.path.exists('movie'): os.mkdir('movie') else: P.ion() if not normed: vmin = -max(abs(self.data.max()), abs(self.data.min())) vmin = cut * vmin vmax = -vmin #vmin, vmax = self.data.min(), self.data.max() cs = N.linspace(vmin, vmax, levels) if self.surftype == 'phi_constant': #if self.movtype in [1, 7]: #th = N.linspace(0., 2.*N.pi, 2*self.n_theta_max) #else: th = N.linspace(N.pi / 2., -N.pi / 2., self.n_theta_max) rr, tth = N.meshgrid(self.radius, th) xx = rr * N.cos(tth) yy = rr * N.sin(tth) xxout = rr.max() * N.cos(th) yyout = rr.max() * N.sin(th) xxin = rr.min() * N.cos(th) yyin = rr.min() * N.sin(th) fig = P.figure(figsize=(4, 8)) elif self.surftype == 'r_constant': th = N.linspace(N.pi / 2., -N.pi / 2., self.n_theta_max) if deminc: phi = N.linspace(-N.pi, N.pi, self.n_phi_tot * self.minc + 1) xxout, yyout = hammer2cart(th, -N.pi) xxin, yyin = hammer2cart(th, N.pi) else: phi = N.linspace(-N.pi / self.minc, N.pi / self.minc, self.n_phi_tot) xxout, yyout = hammer2cart(th, -N.pi / self.minc) xxin, yyin = hammer2cart(th, N.pi / self.minc) ttheta, pphi = N.meshgrid(th, phi) xx, yy = hammer2cart(ttheta, pphi) fig = P.figure(figsize=(8, 4)) elif self.surftype == 'theta_constant': if deminc: phi = N.linspace(0., 2. * N.pi, self.n_phi_tot * self.minc + 1) else: phi = N.linspace(0., 2. * N.pi / self.minc, self.n_phi_tot) rr, pphi = N.meshgrid(self.radius, phi) xx = rr * N.cos(pphi) yy = rr * N.sin(pphi) xxout = rr.max() * N.cos(pphi) yyout = rr.max() * N.sin(pphi) xxin = rr.min() * N.cos(pphi) yyin = rr.min() * N.sin(pphi) fig = P.figure(figsize=(6, 6)) elif self.surftype == '3d volume': self.data = self.data[..., 0] th = N.linspace(N.pi / 2., -N.pi / 2., self.n_theta_max) phi = N.linspace(-N.pi, N.pi, self.n_phi_tot) ttheta, pphi = N.meshgrid(th, phi) xx, yy = hammer2cart(ttheta, pphi) xxout, yyout = hammer2cart(th, -N.pi) xxin, yyin = hammer2cart(th, N.pi) fig = P.figure(figsize=(8, 4)) fig.subplots_adjust(top=0.99, right=0.99, bottom=0.01, left=0.01) ax = fig.add_subplot(111, frameon=False) for k in range(self.nvar): if k == 0: if normed: vmin = -max(abs(self.data[k, ...].max()), abs(self.data[k, ...].min())) vmin = cut * vmin vmax = -vmin cs = N.linspace(vmin, vmax, levels) if self.surftype in ['r_constant', 'theta_constant']: if deminc: dat = symmetrize(self.data[k, ...], self.minc) else: dat = self.data[k, ...] im = ax.contourf(xx, yy, dat, cs, cmap=cmap, extend='both') else: im = ax.contourf(xx, yy, self.data[k, ...], cs, cmap=cmap, extend='both') ax.plot(xxout, yyout, 'k-', lw=1.5) ax.plot(xxin, yyin, 'k-', lw=1.5) ax.axis('off') man = P.get_current_fig_manager() man.canvas.draw() if k != 0 and k % step == 0: if not png: print(k + self.var2 - self.nvar) P.cla() if normed: vmin = -max(abs(self.data[k, ...].max()), abs(self.data[k, ...].min())) vmin = cut * vmin vmax = -vmin cs = N.linspace(vmin, vmax, levels) if self.surftype in ['r_constant', 'theta_constant']: if deminc: dat = symmetrize(self.data[k, ...], self.minc) else: dat = self.data[k, ...] im = ax.contourf(xx, yy, dat, cs, cmap=cmap, extend='both') else: im = ax.contourf(xx, yy, self.data[k, ...], cs, cmap=cmap, extend='both') ax.plot(xxout, yyout, 'k-', lw=1.5) ax.plot(xxin, yyin, 'k-', lw=1.5) 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 __init__(self, file=None, step=1, lastvar=None, nvar='all', nrout=48, ratio_out=2., potExtra=False, precision='Float32'): """ :param file: file name :type file: str :param nvar: the number of timesteps of the movie file we want to plot starting from the last line :type nvar: int :param lastvar: the number of the last timestep to be read :type lastvar: int :param step: the stepping between two timesteps :type step: int :param precision: precision of the input file, Float32 for single precision, Float64 for double precision :type precision: str :param potExtra: when set to True, potential extrapolation of the magnetic field outside the fluid domain is also computed :type potExtra: bool :param ratio_out: ratio of desired external radius to the CMB radius. This is is only used when potExtra=True :type ratio_out: float :param nrout: number of additional radial grid points to compute the potential extrapolation. This is only used when potExtra=True :type nrout: int """ if file == None: dat = glob.glob('*_mov.*') str1 = 'Which movie do you want ?\n' for k, movie in enumerate(dat): str1 += ' %i) %s\n' % (k+1, movie) index = int(input(str1)) try: filename = dat[index-1] except IndexError: print('Non valid index: %s has been chosen instead' % dat[0]) filename = dat[0] else: filename = file mot = re.compile(r'.*_mov\.(.*)') end = mot.findall(filename)[0] # DETERMINE THE NUMBER OF LINES BY READING THE LOG FILE logfile = open('log.%s' % end, 'r') mot = re.compile(r' ! WRITING MOVIE FRAME NO\s*(\d*).*') for line in logfile.readlines(): if mot.match(line): nlines = int(mot.findall(line)[0]) logfile.close() if lastvar is None: self.var2 = nlines else: self.var2 = lastvar if str(nvar) == 'all': self.nvar = nlines self.var2 = nlines else: self.nvar = nvar vecNames = np.r_[3] scalNames = np.r_[-1] # READ the movie file infile = npfile(filename, endian='B') # HEADER version = infile.fort_read('|S64') n_type, n_surface, const, n_fields = infile.fort_read(precision) n_fields = int(n_fields) n_surface = int(n_surface) if n_fields == 1: movtype = infile.fort_read(precision) self.movtype = int(movtype) else: movtype = infile.fort_read(precision) # RUN PARAMETERS runid = infile.fort_read('|S64') n_r_mov_tot, n_r_max, n_theta_max, n_phi_tot, minc, ra, \ ek, pr, prmag, radratio, tScale = infile.fort_read(precision) minc = int(minc) self.n_r_max = int(n_r_max) self.n_theta_max = int(n_theta_max) self.n_phi_tot = int(n_phi_tot) n_r_mov_tot = int(n_r_mov_tot) # GRID if potExtra: self.radius = np.zeros((n_r_mov_tot+1+nrout), 'f') else: self.radius = np.zeros((n_r_mov_tot+2), 'f') tmp = infile.fort_read(precision)/(1.-radratio) rcmb = tmp[0] self.radius[2:n_r_mov_tot+2] = tmp[::-1] self.theta = infile.fort_read(precision) self.phi = infile.fort_read(precision) shape = (n_r_mov_tot+2, self.n_theta_max, self.n_phi_tot) self.time = np.zeros(self.nvar, precision) if potExtra: scals = np.zeros((n_r_mov_tot+1+nrout, self.n_theta_max, self.n_phi_tot*minc+1,1),'f') else: scals = np.zeros((n_r_mov_tot+2, self.n_theta_max, self.n_phi_tot*minc+1,1),'f') vecr = np.zeros_like(scals) vect = np.zeros_like(scals) vecp = np.zeros_like(scals) # Potential extrapolation radii = self.radius scals = scals.T scals[0, ...] = radii startdir = os.getcwd() if not os.path.exists('vtsFiles'): os.mkdir('vtsFiles') os.chdir('vtsFiles') else: os.chdir('vtsFiles') for i in range(self.var2-self.nvar): n_frame, t_movieS, omega_ic, omega_ma, movieDipColat, \ movieDipLon, movieDipStrength, \ movieDipStrengthGeo = infile.fort_read(precision) tmp = infile.fort_read(precision, shape=shape) vecr[0:n_r_mov_tot+2,:,:,0] = symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vect[0:n_r_mov_tot+2,:,:,0] = symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vecp[0:n_r_mov_tot+2,:,:,0] = symmetrize(tmp, minc, reversed=True) for k in range(self.nvar): n_frame, t_movieS, omega_ic, omega_ma, movieDipColat, \ movieDipLon, movieDipStrength, \ movieDipStrengthGeo = infile.fort_read(precision) self.time[k] = t_movieS if k % step == 0: #print(k+self.var2-self.nvar) tmp = infile.fort_read(precision, shape=shape) brCMB = np.zeros((self.n_theta_max, self.n_phi_tot), 'f') brCMB = tmp[0, :, :] brCMB = brCMB.T if potExtra: vecr[nrout-1:nrout+n_r_mov_tot+2,:,:,0] = symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vect[nrout-1:nrout+n_r_mov_tot+2,:,:,0] = symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vecp[nrout-1:nrout+n_r_mov_tot+2,:,:,0] = symmetrize(tmp, minc, reversed=True) else: vecr[0:n_r_mov_tot+2,:,:,0] = symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vect[0:n_r_mov_tot+2,:,:,0] = symmetrize(tmp, minc, reversed=True) tmp = infile.fort_read(precision, shape=shape) vecp[0:n_r_mov_tot+2,:,:,0] = symmetrize(tmp, minc, reversed=True) filename = 'B3D_%05d' % k vecr = vecr[::-1, ...] vect = vect[::-1, ...] vecp = vecp[::-1, ...] br = vecr.T bt = vect.T bp = vecp.T if potExtra: pot = ExtraPot(rcmb, brCMB, minc, ratio_out=ratio_out, nrout=nrout, cutCMB=True, deminc=True) br[0, ..., n_r_mov_tot+2:] = pot.brout bt[0, ..., n_r_mov_tot+2:] = pot.btout bp[0, ..., n_r_mov_tot+2:] = pot.bpout radii[n_r_mov_tot+2:] = pot.rout else: radii = self.radius vts(filename, radii, br, bt, bp, scals, scalNames, vecNames, 1) print('write %s.vts' % filename) else: # Otherwise we read vecr = infile.fort_read(precision, shape=shape) vect = infile.fort_read(precision, shape=shape) vecp = infile.fort_read(precision, shape=shape) os.chdir(startdir)
print 'BtCMB: %.8e %.8e' % (btCMB.max(), btCMB.min()) print 'reconstruct BtCMB: %.8e %.8e' % (btsurf.max(), btsurf.min()) print 'BpCMB: %.8e %.8e' % (bpCMB.max(), bpCMB.min()) print 'reconstruct BpCMB: %.8e %.8e' % (bpsurf.max(), bpsurf.min()) fig = P.figure(figsize=(12, 4)) P.subplots_adjust(right=0.99, left=0.01, top=0.99, bottom=0.01, wspace=0.02, hspace=0.02) ax = fig.add_subplot(231, frameon=False) im = ax.contourf(x, y, symmetrize(brCMB, g.minc), 16, cmap=P.get_cmap('RdYlBu_r')) P.axis('off') ax = fig.add_subplot(232, frameon=False) im = ax.contourf(x, y, symmetrize(btCMB, g.minc), 16, cmap=P.get_cmap('RdYlBu_r')) P.axis('off') ax = fig.add_subplot(233, frameon=False) im = ax.contourf(x, y, symmetrize(bpCMB, g.minc), 16,
btsurf = btsurf.real bpsurf = N.fft.ifft(bpm, axis=0)*g.npI bpsurf = bpsurf.real print 'BrCMB: %.8e %.8e' % (brCMB.max(), brCMB.min()) print 'reconstruct BrCMB: %.8e %.8e' % (brsurf.max(), brsurf.min()) print 'BtCMB: %.8e %.8e' % (btCMB.max(), btCMB.min()) print 'reconstruct BtCMB: %.8e %.8e' % (btsurf.max(), btsurf.min()) print 'BpCMB: %.8e %.8e' % (bpCMB.max(), bpCMB.min()) print 'reconstruct BpCMB: %.8e %.8e' % (bpsurf.max(), bpsurf.min()) fig = P.figure(figsize=(12,4)) P.subplots_adjust(right=0.99, left=0.01, top=0.99, bottom=0.01, wspace=0.02, hspace=0.02) ax = fig.add_subplot(231, frameon=False) im = ax.contourf(x,y,symmetrize(brCMB, g.minc), 16, cmap=P.get_cmap('RdYlBu_r')) P.axis('off') ax = fig.add_subplot(232, frameon=False) im = ax.contourf(x,y,symmetrize(btCMB, g.minc), 16, cmap=P.get_cmap('RdYlBu_r')) P.axis('off') ax = fig.add_subplot(233, frameon=False) im = ax.contourf(x,y,symmetrize(bpCMB, g.minc), 16, cmap=P.get_cmap('RdYlBu_r')) P.axis('off') ax = fig.add_subplot(234, frameon=False) im = ax.contourf(x,y,symmetrize(brsurf, g.minc), 16, cmap=P.get_cmap('RdYlBu_r')) P.axis('off') ax = fig.add_subplot(235, frameon=False) im = ax.contourf(x,y,symmetrize(btsurf, g.minc), 16, cmap=P.get_cmap('RdYlBu_r')) P.axis('off') ax = fig.add_subplot(236, frameon=False) im = ax.contourf(x,y,symmetrize(bpsurf, g.minc), 16, cmap=P.get_cmap('RdYlBu_r'))