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'))
