# HG changeset patch # User calpe <miguel.calpe.linares@gmail.com> # Date 1543845807 -3600 # Mon Dec 03 15:03:27 2018 +0100 # Node ID 9a058df16f26c09f591fc6e6941ebf2b2cdcdd67 # Parent 89d7832e8837c98d93b862b59eb723bd585f658c Function plot module spectra_multidim.py diff --git a/fluidsim/solvers/ns2d/strat/output/spectra_multidim.py b/fluidsim/solvers/ns2d/strat/output/spectra_multidim.py --- a/fluidsim/solvers/ns2d/strat/output/spectra_multidim.py +++ b/fluidsim/solvers/ns2d/strat/output/spectra_multidim.py @@ -9,7 +9,10 @@ import h5py import numpy as np +import matplotlib.pyplot as plt +import matplotlib.patches as patches +from math import radians from fluidsim.base.output.spectra_multidim import SpectraMultiDim @@ -27,53 +30,151 @@ am_fft = self.sim.state.compute("am_fft") # Computes multidimensional spectra - spectrumkykx_E = self.oper.compute_spectrum_kykx(energy_fft) - spectrumkykx_EK = self.oper.compute_spectrum_kykx(energyK_fft) - spectrumkykx_EA = self.oper.compute_spectrum_kykx(energyA_fft) + spectrumkykx_E = self.oper.compute_spectrum_kykx(energy_fft, folded=False) + spectrumkykx_EK = self.oper.compute_spectrum_kykx(energyK_fft, folded=False) + spectrumkykx_EA = self.oper.compute_spectrum_kykx(energyA_fft, folded=False) # The function compute_spectrum_kykx does not supports complex variable... # Only works for the energy! - # spectrumkykx_ap_fft = self.oper.compute_spectrum_kykx(ap_fft) - # spectrumkykx_am_fft = self.oper.compute_spectrum_kykx(am_fft) + energy_ap_fft = abs(ap_fft)**2 + energy_am_fft = abs(am_fft)**2 - # Saves dictionary + spectrumkykx_ap_fft = self.oper.compute_spectrum_kykx(energy_ap_fft, folded=False) + spectrumkykx_am_fft = self.oper.compute_spectrum_kykx(energy_am_fft, folded=False) + dict_spectra = { "spectrumkykx_E": spectrumkykx_E, "spectrumkykx_EK": spectrumkykx_EK, - "spectrumkykx_EA": spectrumkykx_EA + "spectrumkykx_EA": spectrumkykx_EA, + "spectrumkykx_ap_fft": spectrumkykx_ap_fft, + "spectrumkykx_am_fft": spectrumkykx_am_fft } - # dict_spectra = { - # "spectrumkykx_E": spectrumkykx_E, - # "spectrumkykx_EK": spectrumkykx_EK, - # "spectrumkykx_EA": spectrumkykx_EA, - # "spectrumkykx_ap_fft": spectrumkykx_ap_fft, - # "spectrumkykx_am_fft": spectrumkykx_am_fft - # } - return dict_spectra def _online_plot_saving(self, dict_spectra): raise NotImplementedError("_online_plot_saving in not implemented.") - def plot(self, tmin=0, tmax=1000): - """Plots spectrumkykx averaged between tmin and tmax.""" + def plot(self, key=None, tmin=0, tmax=None): + """ + Plots spectrumkykx averaged between tmin and tmax. + + Parameters + ---------- + key : str + Key to plot the spectrum: E, EK, EA, ap_fft (default), am_fft + + """ + + oper = self.sim.params.oper - dict_results = self.load_mean(tmin, tmax) - kx = dict_results["kxE"] - ky = dict_results["kyE"] - spectrumkykx_E = dict_results["spectrumkykx_E"] + # Load data + with h5py.File(self.path_file, "r") as f: + times = f["times"].value + kx = f["kxE"].value + # kz = f["kyE"].value + if key == "E": + data = f["spectrumkykx_E"].value + elif key == "EK": + data = f["spectrumkykx_EK"].value + elif key == "EA": + data = f["spectrumkykx_EA"].value + elif key == "ap_fft" or not key: + data = f["spectrumkykx_ap_fft"].value + elif key == "am_fft": + data = f["spectrumkykx_am_fft"].value + else: + raise ValueError("Key unknown.") + + # Compute time average + if not tmax: + tmax = times[-1] + + itmin = np.argmin(abs(times - tmin)) + itmax = np.argmin(abs(times - tmax)) + + data_plot = np.mean(data[itmin : itmax, :, :], axis=0) + + # Create array kz with negative values + kz = 2 * np.pi * np.fft.fftfreq(oper.ny, oper.Ly / oper.ny) + kz[kz.shape[0]//2] *= -1 + + # Create mesh of wave-numbers + KX, KZ = np.meshgrid(kx, kz) + + ### Data + ikx = np.argmin(abs(kx - 200)) + ikz = np.argmin(abs(kz - 148)) + ikz_negative = np.argmin(abs(kz + 148)) + + + # Set figure parameters + fig, ax = plt.subplots() + ax.set_xlabel(r"$k_x$") + ax.set_ylabel(r"$k_z$") - fig, ax = self.output.figure_axe() - ax.set_xlabel("$k_x$") - ax.set_ylabel("$k_z$") + kz_modified = np.empty_like(kz) + kz_modified[0:kz_modified.shape[0]//2 - 1] = kz[kz_modified.shape[0]//2 + 1:] + kz_modified[kz_modified.shape[0]//2 - 1:] = kz[0:kz_modified.shape[0]//2 + 1] + + KX, KZ = np.meshgrid(kx, kz_modified) + + data_plot_modified = np.empty_like(data_plot) + data_plot_modified[0:kz_modified.shape[0]//2 - 1, :] = data_plot[kz_modified.shape[0]//2 + 1:, :] + data_plot_modified[kz_modified.shape[0]//2 - 1:, :] = data_plot[0:kz_modified.shape[0]//2 + 1, :] + + ax.pcolormesh(KX, KZ, data_plot_modified) + + + # Create a Rectangle patch + deltak = max(self.sim.oper.deltakx, self.sim.oper.deltaky) + + angle = radians(float(self.sim.params.forcing.tcrandom_anisotropic.angle.split("°")[0])) + + x_rect = np.sin(angle) * deltak * self.sim.params.forcing.nkmin_forcing + + z_rect = np.cos(angle) * deltak * self.sim.params.forcing.nkmin_forcing + + width = abs(x_rect - np.sin(angle) * deltak * self.sim.params.forcing.nkmax_forcing) + + height = abs(z_rect - np.cos(angle) * deltak * self.sim.params.forcing.nkmax_forcing) + + rect1 = patches.Rectangle((x_rect,z_rect),width,height,linewidth=1,edgecolor='r',facecolor='none') + + ax.add_patch(rect1) - KX, KY = np.meshgrid(kx, ky) - ax.pcolormesh( - KX, - KY, - spectrumkykx_E, - vmin=spectrumkykx_E.min(), - vmax=spectrumkykx_E.max(), + if self.sim.params.forcing.tcrandom_anisotropic.kz_negative_enable: + rect2 = patches.Rectangle( + (x_rect,-(z_rect + height)), width, height, linewidth=1, + edgecolor='r',facecolor='none') + + ax.add_patch(rect2) + + # Plot arc kmin and kmax forcing + ax.add_patch( + patches.Arc( + xy=(0, 0), + width=2 * self.sim.params.forcing.nkmin_forcing * deltak, + height=2 * self.sim.params.forcing.nkmin_forcing * deltak, + angle=0, + theta1=-90., + theta2=90., + linestyle="-.", + color="red" + ) ) + ax.add_patch( + patches.Arc( + xy=(0, 0), + width=2 * self.sim.params.forcing.nkmax_forcing * deltak, + height=2 * self.sim.params.forcing.nkmax_forcing * deltak, + angle=0, + theta1=-90, + theta2=90., + linestyle="-.", + color="red" + ) + ) + + ax.set_aspect("equal")