From c5f391d24bcf2a011e608388e6bfe8383627f387 Mon Sep 17 00:00:00 2001 From: paugier Date: Tue, 2 Jan 2024 15:45:49 +0100 Subject: [PATCH] Replace transonic.jit by boost --- .hgignore | 2 + Makefile | 6 +- doc/install.rst | 6 - fluidsim/__init__.py | 13 - fluidsim/solvers/ns2d/bouss/meson.build | 6 + fluidsim/solvers/ns2d/bouss/solver.py | 4 +- fluidsim/solvers/ns2d/strat/meson.build | 6 + fluidsim/solvers/ns2d/strat/solver.py | 4 +- fluidsim/solvers/plate2d/meson.build | 6 + fluidsim/solvers/plate2d/operators.py | 6 +- .../plate2d/output/correlations_freq.py | 6 +- fluidsim/solvers/plate2d/output/meson.build | 6 + fluidsim/solvers/sphere/sw1l/meson.build | 5 + fluidsim/solvers/sphere/sw1l/solver.py | 4 +- fluidsim/solvers/sw1l/meson.build | 9 + fluidsim/solvers/sw1l/operators.py | 6 +- fluidsim/solvers/sw1l/output/__init__.py | 263 +---------------- fluidsim/solvers/sw1l/output/base.py | 267 ++++++++++++++++++ fluidsim/solvers/sw1l/output/meson.build | 7 + fluidsim/solvers/sw1l/solver.py | 6 +- 20 files changed, 337 insertions(+), 301 deletions(-) create mode 100644 fluidsim/solvers/sw1l/output/base.py diff --git a/.hgignore b/.hgignore index a1a58bccb..62f0a8602 100644 --- a/.hgignore +++ b/.hgignore @@ -83,3 +83,5 @@ scratch fluidsim/_version.py doc/examples/htmlcov/* + +fluidsim/build_conf.txt diff --git a/Makefile b/Makefile index f56c11675..a3822c1b0 100644 --- a/Makefile +++ b/Makefile @@ -8,7 +8,7 @@ MPI_NUM_PROCS ?= 2 develop: develop_lib pip install meson-python ninja numpy "pythran>=0.9.7" pip install --force-reinstall transonic@hg+https://foss.heptapod.net/fluiddyn/transonic - pip install -e .[dev] --no-build-isolation + pip install -e ".[dev]" --no-build-isolation develop_lib: cd lib && pip install -e . @@ -41,10 +41,10 @@ shortlog: @hg log -M -r$(RELEASE): --template '- {desc|firstline} (:rev:`{node|short}`)\n' black: - black -l 82 fluidsim scripts bench doc lib --exclude "/(__pythran__|doc/_build|\.ipynb_checkpoints/*)/" + black -l 82 fluidsim scripts bench doc lib --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/" black_check: - black --check -l 82 fluidsim scripts bench doc lib --exclude "/(__pythran__|doc/_build|\.ipynb_checkpoints/*)/" + black --check -l 82 fluidsim scripts bench doc lib --exclude "/(__pythran__|__python__|__numba__|doc/_build|\.ipynb_checkpoints/*)/" tests: pytest -v lib diff --git a/doc/install.rst b/doc/install.rst index 0f8965856..16ea95069 100644 --- a/doc/install.rst +++ b/doc/install.rst @@ -100,12 +100,6 @@ Fluidsim is also sensitive to the environment variables: - ``FLUIDDYN_PATH_SCRATCH``: working directory (can be useful on some clusters). -- ``TRANSONIC_COMPILE_JIT``: set this variable to force JIT compilation using - ``transonic`` while running tests. This is not necessary, but could be useful - for troubleshooting if simulations freeze. For example:: - - TRANSONIC_COMPILE_JIT=1 fluidsim-test -m fluidsim.solvers.sw1l - Warning about re-installing fluidsim and fluidfft with new build options ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ diff --git a/fluidsim/__init__.py b/fluidsim/__init__.py index 0e7cf37eb..a07baae73 100644 --- a/fluidsim/__init__.py +++ b/fluidsim/__init__.py @@ -55,19 +55,6 @@ def _show(*args, **kwargs): plt.show = _show - if all( - env_var not in os.environ - for env_var in ("FLUID_COMPILE_CACHEDJIT", "TRANSONIC_COMPILE_JIT") - ): - mpi.printby0("Compilation of jit functions disabled.") - from transonic import set_compile_jit - - set_compile_jit(False) - elif "FLUID_COMPILE_CACHEDJIT" in os.environ: - mpi.printby0( - "WARNING: FLUID_COMPILE_CACHEDJIT is deprecated, use " - "TRANSONIC_COMPILE_JIT instead." - ) from ._version import __version__, get_local_version diff --git a/fluidsim/solvers/ns2d/bouss/meson.build b/fluidsim/solvers/ns2d/bouss/meson.build index db75efa12..385f4e08c 100644 --- a/fluidsim/solvers/ns2d/bouss/meson.build +++ b/fluidsim/solvers/ns2d/bouss/meson.build @@ -9,3 +9,9 @@ py.install_sources( python_sources, subdir: 'fluidsim/solvers/ns2d/bouss', ) + +run_command(['transonic', '--meson', '--backend', backend, 'solver.py'], check: true) + +foreach be : backends + subdir('__' + be + '__') +endforeach diff --git a/fluidsim/solvers/ns2d/bouss/solver.py b/fluidsim/solvers/ns2d/bouss/solver.py index 2407fd193..1933faf5d 100644 --- a/fluidsim/solvers/ns2d/bouss/solver.py +++ b/fluidsim/solvers/ns2d/bouss/solver.py @@ -8,7 +8,7 @@ """ import numpy as np -from transonic import jit, Array +from transonic import boost, Array from fluidsim.base.setofvariables import SetOfVariables from fluidsim.solvers.ns2d.solver import InfoSolverNS2D, Simul as SimulNS2D @@ -17,7 +17,7 @@ AF = Array[np.float64, "2d"] -@jit +@boost def tendencies_nonlin_ns2dbouss( ux: AF, uy: AF, px_rot: AF, py_rot: AF, px_b: AF, py_b: AF ): diff --git a/fluidsim/solvers/ns2d/strat/meson.build b/fluidsim/solvers/ns2d/strat/meson.build index 741a149e6..f04bcf2f7 100644 --- a/fluidsim/solvers/ns2d/strat/meson.build +++ b/fluidsim/solvers/ns2d/strat/meson.build @@ -14,3 +14,9 @@ py.install_sources( ) subdir('output') + +run_command(['transonic', '--meson', '--backend', backend, 'solver.py'], check: true) + +foreach be : backends + subdir('__' + be + '__') +endforeach diff --git a/fluidsim/solvers/ns2d/strat/solver.py b/fluidsim/solvers/ns2d/strat/solver.py index 8520d3958..59d47726f 100644 --- a/fluidsim/solvers/ns2d/strat/solver.py +++ b/fluidsim/solvers/ns2d/strat/solver.py @@ -9,7 +9,7 @@ import numpy as np -from transonic import jit, Array +from transonic import boost, Array from fluidsim.base.setofvariables import SetOfVariables @@ -18,7 +18,7 @@ AF = Array[np.float64, "2d"] -@jit +@boost def tendencies_nonlin_ns2dstrat( ux: AF, uy: AF, px_rot: AF, py_rot: AF, px_b: AF, py_b: AF, N: float ): diff --git a/fluidsim/solvers/plate2d/meson.build b/fluidsim/solvers/plate2d/meson.build index bff354224..22850b37c 100644 --- a/fluidsim/solvers/plate2d/meson.build +++ b/fluidsim/solvers/plate2d/meson.build @@ -16,3 +16,9 @@ py.install_sources( ) subdir('output') + +run_command(['transonic', '--meson', '--backend', backend, 'operators.py'], check: true) + +foreach be : backends + subdir('__' + be + '__') +endforeach diff --git a/fluidsim/solvers/plate2d/operators.py b/fluidsim/solvers/plate2d/operators.py index fb780218a..e8fa009e5 100644 --- a/fluidsim/solvers/plate2d/operators.py +++ b/fluidsim/solvers/plate2d/operators.py @@ -11,7 +11,7 @@ import numpy as np -from transonic import jit, Array +from transonic import boost, Array from ...operators.operators2d import OperatorsPseudoSpectral2D @@ -19,7 +19,7 @@ AF2 = Array[np.float64, "2d"] -@jit +@boost def monge_ampere_step0(a_fft: AC2, b_fft: AC2, KX2: AF2, KY2: AF2, KXKZ: AF2): pxx_a_fft = -a_fft * KX2 pyy_a_fft = -a_fft * KY2 @@ -30,7 +30,7 @@ def monge_ampere_step0(a_fft: AC2, b_fft: AC2, KX2: AF2, KY2: AF2, KXKZ: AF2): return pxx_a_fft, pyy_a_fft, pxy_a_fft, pxx_b_fft, pyy_b_fft, pxy_b_fft -@jit +@boost def monge_ampere_step1( pxx_a: AF2, pyy_a: AF2, pxy_a: AF2, pxx_b: AF2, pyy_b: AF2, pxy_b: AF2 ): diff --git a/fluidsim/solvers/plate2d/output/correlations_freq.py b/fluidsim/solvers/plate2d/output/correlations_freq.py index 450b6aadf..9a0ef9c5d 100644 --- a/fluidsim/solvers/plate2d/output/correlations_freq.py +++ b/fluidsim/solvers/plate2d/output/correlations_freq.py @@ -17,14 +17,14 @@ import h5py import matplotlib.pyplot as plt -from transonic import jit +from transonic import boost from fluiddyn.util import mpi from fluiddyn.calcul.easypyfft import FFTW1DReal2Complex from fluidsim.base.output.base import SpecificOutput -@jit +@boost def compute_correl4_seq( q_fftt: "complex128[][]", iomegas1: "int32[]", nb_omegas: int, nb_xs_seq: int ): @@ -94,7 +94,7 @@ def compute_correl4_seq( return corr4 -@jit +@boost def compute_correl2_seq( q_fftt: "complex128[][]", iomegas1: "int32[]", nb_omegas: int, nb_xs_seq: int ): diff --git a/fluidsim/solvers/plate2d/output/meson.build b/fluidsim/solvers/plate2d/output/meson.build index 4a4ec2784..e4e6e817b 100644 --- a/fluidsim/solvers/plate2d/output/meson.build +++ b/fluidsim/solvers/plate2d/output/meson.build @@ -10,3 +10,9 @@ py.install_sources( python_sources, subdir: 'fluidsim/solvers/plate2d/output' ) + +run_command(['transonic', '--meson', '--backend', backend, 'correlations_freq.py'], check: true) + +foreach be : backends + subdir('__' + be + '__') +endforeach diff --git a/fluidsim/solvers/sphere/sw1l/meson.build b/fluidsim/solvers/sphere/sw1l/meson.build index 502a407bf..91b97ae01 100644 --- a/fluidsim/solvers/sphere/sw1l/meson.build +++ b/fluidsim/solvers/sphere/sw1l/meson.build @@ -10,3 +10,8 @@ py.install_sources( subdir: 'fluidsim/solvers/sphere/sw1l' ) +run_command(['transonic', '--meson', '--backend', backend, 'solver.py'], check: true) + +foreach be : backends + subdir('__' + be + '__') +endforeach diff --git a/fluidsim/solvers/sphere/sw1l/solver.py b/fluidsim/solvers/sphere/sw1l/solver.py index f483b8040..d07b2a116 100644 --- a/fluidsim/solvers/sphere/sw1l/solver.py +++ b/fluidsim/solvers/sphere/sw1l/solver.py @@ -10,7 +10,7 @@ :private-members: """ -from transonic import jit +from transonic import boost from fluidsim.base.sphericalharmo.solver import ( InfoSolverSphericalHarmo, SimulSphericalHarmo, @@ -22,7 +22,7 @@ Af = "float64[:, :]" -@jit +@boost def compute_Frot(rot: Af, ux: Af, uy: Af, f_radial: Af): """Compute cross-product of absolute potential vorticity with velocity.""" rot_abs = rot + f_radial diff --git a/fluidsim/solvers/sw1l/meson.build b/fluidsim/solvers/sw1l/meson.build index 1488222a0..7caabdb7a 100644 --- a/fluidsim/solvers/sw1l/meson.build +++ b/fluidsim/solvers/sw1l/meson.build @@ -18,3 +18,12 @@ subdir('exactlin') subdir('modified') subdir('onlywaves') subdir('output') + +run_command( + ['transonic', '--meson', '--backend', backend, 'operators.py', 'solver.py'], + check: true +) + +foreach be : backends + subdir('__' + be + '__') +endforeach diff --git a/fluidsim/solvers/sw1l/operators.py b/fluidsim/solvers/sw1l/operators.py index ba2a2a190..841bf5c6f 100644 --- a/fluidsim/solvers/sw1l/operators.py +++ b/fluidsim/solvers/sw1l/operators.py @@ -12,7 +12,7 @@ import numpy as np from fluiddyn.util.compat import cached_property -from transonic import boost, jit, Array +from transonic import boost, Array from fluidsim.operators.operators2d import ( OperatorsPseudoSpectral2D, @@ -25,7 +25,7 @@ AF = Array[np.float64, "2d"] -@jit +@boost def _qapamfft_from_uxuyetafft( ux_fft: AC, uy_fft: AC, @@ -309,7 +309,7 @@ def apamfft_from_adfft(self, a_fft, d_fft): am_fft = 0.5 * (a_fft - Delta_a_fft) return ap_fft, am_fft - @jit + @boost def divfft_from_apamfft(self, ap_fft: AC, am_fft: AC): """Return div from the eigen modes ap and am.""" n0 = self.nK0_loc diff --git a/fluidsim/solvers/sw1l/output/__init__.py b/fluidsim/solvers/sw1l/output/__init__.py index ad301ad55..635f83cbe 100644 --- a/fluidsim/solvers/sw1l/output/__init__.py +++ b/fluidsim/solvers/sw1l/output/__init__.py @@ -1,10 +1,6 @@ """Output SW1L (:mod:`fluidsim.solvers.sw1l.output`) ==================================================== -.. autoclass:: OutputBaseSW1L - :members: - :private-members: - .. autosummary:: :toctree: @@ -16,261 +12,6 @@ """ -import numpy as np - -from transonic import jit - -from fluidsim.base.output import OutputBasePseudoSpectral - - -@jit -def linear_eigenmode_from_values_1k( - ux_fft: np.complex128, - uy_fft: np.complex128, - eta_fft: np.complex128, - kx: float, - ky: float, - f: "float or int", - c2: "float or int", -): - """Compute q, d, a (fft) for a single wavenumber.""" - div_fft = 1j * (kx * ux_fft + ky * uy_fft) - rot_fft = 1j * (kx * uy_fft - ky * ux_fft) - q_fft = rot_fft - f * eta_fft - k2 = kx**2 + ky**2 - ageo_fft = f * rot_fft / c2 + k2 * eta_fft - return q_fft, div_fft, ageo_fft - - -class OutputBaseSW1L(OutputBasePseudoSpectral): - @staticmethod - def _complete_info_solver(info_solver): - """Complete the ParamContainer *info_solver* with child classes to be - instantiated under *sim.output*. - - This is a static method! - """ - classes = info_solver.classes.Output._set_child("classes") - - package = "fluidsim.solvers.sw1l.output" - - classes._set_child( - "PrintStdOut", - attribs={ - "module_name": package + ".print_stdout", - "class_name": "PrintStdOutSW1L", - }, - ) - - classes._set_child( - "PhysFields", - attribs={ - "module_name": "fluidsim.base.output.phys_fields2d", - "class_name": "PhysFieldsBase2D", - }, - ) - - classes._set_child( - "Spectra", - attribs={ - "module_name": package + ".spectra", - "class_name": "SpectraSW1LNormalMode", - }, - ) - - classes._set_child( - "SpatialMeans", - attribs={ - "module_name": package + ".spatial_means", - "class_name": "SpatialMeansSW1L", - }, - ) - - attribs = { - "module_name": package + ".spect_energy_budget", - "class_name": "SpectralEnergyBudgetSW1L", - } - classes._set_child("SpectralEnergyBudget", attribs=attribs) - - attribs = { - "module_name": package + ".increments", - "class_name": "IncrementsSW1L", - } - classes._set_child("Increments", attribs=attribs) - - attribs = { - "module_name": "fluidsim.base.output.prob_dens_func", - "class_name": "ProbaDensityFunc", - } - classes._set_child("ProbaDensityFunc", attribs=attribs) - - attribs = { - "module_name": "fluidsim.base.output.time_signals_fft", - "class_name": "TimeSignalsK", - } - classes._set_child("TimeSignalsK", attribs=attribs) - - @staticmethod - def _complete_params_with_default(params, info_solver): - """This static method is used to complete the *params* container.""" - OutputBasePseudoSpectral._complete_params_with_default( - params, info_solver - ) - - params.output.phys_fields.field_to_plot = "rot" - - def linear_eigenmode_from_values_1k(self, ux_fft, uy_fft, eta_fft, kx, ky): - """Compute the linear eigenmodes for a single wavenumber.""" - return linear_eigenmode_from_values_1k( - ux_fft, uy_fft, eta_fft, kx, ky, self.sim.params.f, self.sim.params.c2 - ) - - def omega_from_wavenumber(self, k): - r"""Evaluates the dispersion relation and returns the linear frequency - - .. math:: \omega = \sqrt{f ^ 2 + (ck)^2} - - """ - return np.sqrt(self.sim.params.f**2 + self.sim.params.c2 * k**2) - - def compute_enstrophy_fft(self): - r"""Calculate enstrophy from vorticity in the spectral space.""" - rot_fft = self.sim.state.get_var("rot_fft") - return np.abs(rot_fft) ** 2 / 2.0 - - def compute_PV_fft(self): - r"""Compute Ertel and Charney (QG) potential vorticity. - - .. math:: \zeta_{er} = \frac{f + \zeta}{1 + \eta} - .. math:: \zeta_{ch} = \zeta - f \eta - - """ - get_var = self.sim.state.get_var - rot = get_var("rot") - eta = self.sim.state.state_phys.get_var("eta") - ErtelPV_fft = self.oper.fft2((self.sim.params.f + rot) / (1.0 + eta)) - rot_fft = get_var("rot_fft") - eta_fft = get_var("eta_fft") - CharneyPV_fft = rot_fft - self.sim.params.f * eta_fft - return ErtelPV_fft, CharneyPV_fft - - def compute_PE_fft(self): - """Compute Ertel and Charney (QG) potential enstrophy.""" - ErtelPV_fft, CharneyPV_fft = self.compute_PV_fft() - return (abs(ErtelPV_fft) ** 2 / 2.0, abs(CharneyPV_fft) ** 2 / 2.0) - - def compute_CharneyPE_fft(self): - """Compute Charney (QG) potential enstrophy.""" - rot_fft = self.sim.state.get_var("rot_fft") - eta_fft = self.sim.state.get_var("eta_fft") - CharneyPV_fft = rot_fft - self.sim.params.f * eta_fft - return abs(CharneyPV_fft) ** 2 / 2.0 - - def compute_energies(self): - """Compute kinetic, available potential and rotational kinetic energies.""" - energyK_fft, energyA_fft, energyKr_fft = self.compute_energies_fft() - return ( - self.sum_wavenumbers(energyK_fft), - self.sum_wavenumbers(energyA_fft), - self.sum_wavenumbers(energyKr_fft), - ) - - def compute_energiesKA(self): - """Compute K.E. and A.P.E.""" - energyK_fft, energyA_fft = self.compute_energiesKA_fft() - return ( - self.sum_wavenumbers(energyK_fft), - self.sum_wavenumbers(energyA_fft), - ) - - def compute_energy(self): - """Compute total energy by summing K.E. and A.P.E.""" - energyK_fft, energyA_fft = self.compute_energiesKA_fft() - return self.sum_wavenumbers(energyK_fft) + self.sum_wavenumbers( - energyA_fft - ) - - def compute_enstrophy(self): - """Compute total enstrophy.""" - enstrophy_fft = self.compute_enstrophy_fft() - return self.sum_wavenumbers(enstrophy_fft) - - def compute_lin_energies_fft(self): - r"""Compute quadratic energies decomposed into contributions from - potential vorticity (:math:`q`), divergence (:math:`\nabla.\mathbf u`), - and ageostrophic variable (:math:`a`). - - """ - get_var = self.sim.state.get_var - ux_fft = get_var("ux_fft") - uy_fft = get_var("uy_fft") - eta_fft = get_var("eta_fft") - - q_fft, div_fft, ageo_fft = self.oper.qdafft_from_uxuyetafft( - ux_fft, uy_fft, eta_fft - ) - - udx_fft, udy_fft = self.oper.vecfft_from_divfft(div_fft) - energy_dlin_fft = 0.5 * (np.abs(udx_fft) ** 2 + np.abs(udy_fft) ** 2) - - ugx_fft, ugy_fft, etag_fft = self.oper.uxuyetafft_from_qfft(q_fft) - energy_glin_fft = 0.5 * ( - np.abs(ugx_fft) ** 2 - + np.abs(ugy_fft) ** 2 - + self.sim.params.c2 * np.abs(etag_fft) ** 2 - ) - - uax_fft, uay_fft, etaa_fft = self.oper.uxuyetafft_from_afft(ageo_fft) - energy_alin_fft = 0.5 * ( - np.abs(uax_fft) ** 2 - + np.abs(uay_fft) ** 2 - + self.sim.params.c2 * np.abs(etaa_fft) ** 2 - ) - - return energy_glin_fft, energy_dlin_fft, energy_alin_fft - - -class OutputSW1L(OutputBaseSW1L): - def compute_energies_fft(self): - r"""Compute kinetic, available potential and rotational kinetic energies - in the spectral space. - - .. math:: E_K = (h \mathbf u).\mathbf u / 2 - .. math:: E_A = c^2 \eta^2) / 2 - .. math:: E_{K,r} = (h \mathbf u_r).\mathbf u_r / 2 - - """ - get_var = self.sim.state.get_var - eta_fft = get_var("eta_fft") - energyA_fft = self.sim.params.c2 * np.abs(eta_fft) ** 2 / 2 - Jx_fft = get_var("Jx_fft") - Jy_fft = get_var("Jy_fft") - ux_fft = get_var("ux_fft") - uy_fft = get_var("uy_fft") - energyK_fft = ( - np.real(Jx_fft.conj() * ux_fft + Jy_fft.conj() * uy_fft) / 2.0 - ) - - rot_fft = get_var("rot_fft") - uxr_fft, uyr_fft = self.oper.vecfft_from_rotfft(rot_fft) - rotJ_fft = self.oper.rotfft_from_vecfft(Jx_fft, Jy_fft) - Jxr_fft, Jyr_fft = self.oper.vecfft_from_rotfft(rotJ_fft) - energyKr_fft = ( - np.real(Jxr_fft.conj() * uxr_fft + Jyr_fft.conj() * uyr_fft) / 2.0 - ) - return energyK_fft, energyA_fft, energyKr_fft - - def compute_energiesKA_fft(self): - """Compute K.E. and A.P.E in the spectral space.""" - get_var = self.sim.state.get_var - eta_fft = get_var("eta_fft") - energyA_fft = self.sim.params.c2 * np.abs(eta_fft) ** 2 / 2 - Jx_fft = get_var("Jx_fft") - Jy_fft = get_var("Jy_fft") - ux_fft = get_var("ux_fft") - uy_fft = get_var("uy_fft") - energyK_fft = ( - np.real(Jx_fft.conj() * ux_fft + Jy_fft.conj() * uy_fft) / 2.0 - ) +from .base import OutputBaseSW1L, OutputSW1L - return energyK_fft, energyA_fft +__all__ = ["OutputBaseSW1L", "OutputSW1L"] diff --git a/fluidsim/solvers/sw1l/output/base.py b/fluidsim/solvers/sw1l/output/base.py new file mode 100644 index 000000000..8efa6d2a8 --- /dev/null +++ b/fluidsim/solvers/sw1l/output/base.py @@ -0,0 +1,267 @@ +"""Output SW1L (:mod:`fluidsim.solvers.sw1l.output.base`) +========================================================= + +.. autoclass:: OutputBaseSW1L + :members: + :private-members: + +""" + +import numpy as np + +from transonic import boost + +from fluidsim.base.output import OutputBasePseudoSpectral + + +@boost +def linear_eigenmode_from_values_1k( + ux_fft: np.complex128, + uy_fft: np.complex128, + eta_fft: np.complex128, + kx: float, + ky: float, + f: "float or int", + c2: "float or int", +): + """Compute q, d, a (fft) for a single wavenumber.""" + div_fft = 1j * (kx * ux_fft + ky * uy_fft) + rot_fft = 1j * (kx * uy_fft - ky * ux_fft) + q_fft = rot_fft - f * eta_fft + k2 = kx**2 + ky**2 + ageo_fft = f * rot_fft / c2 + k2 * eta_fft + return q_fft, div_fft, ageo_fft + + +class OutputBaseSW1L(OutputBasePseudoSpectral): + @staticmethod + def _complete_info_solver(info_solver): + """Complete the ParamContainer *info_solver* with child classes to be + instantiated under *sim.output*. + + This is a static method! + """ + classes = info_solver.classes.Output._set_child("classes") + + package = "fluidsim.solvers.sw1l.output" + + classes._set_child( + "PrintStdOut", + attribs={ + "module_name": package + ".print_stdout", + "class_name": "PrintStdOutSW1L", + }, + ) + + classes._set_child( + "PhysFields", + attribs={ + "module_name": "fluidsim.base.output.phys_fields2d", + "class_name": "PhysFieldsBase2D", + }, + ) + + classes._set_child( + "Spectra", + attribs={ + "module_name": package + ".spectra", + "class_name": "SpectraSW1LNormalMode", + }, + ) + + classes._set_child( + "SpatialMeans", + attribs={ + "module_name": package + ".spatial_means", + "class_name": "SpatialMeansSW1L", + }, + ) + + attribs = { + "module_name": package + ".spect_energy_budget", + "class_name": "SpectralEnergyBudgetSW1L", + } + classes._set_child("SpectralEnergyBudget", attribs=attribs) + + attribs = { + "module_name": package + ".increments", + "class_name": "IncrementsSW1L", + } + classes._set_child("Increments", attribs=attribs) + + attribs = { + "module_name": "fluidsim.base.output.prob_dens_func", + "class_name": "ProbaDensityFunc", + } + classes._set_child("ProbaDensityFunc", attribs=attribs) + + attribs = { + "module_name": "fluidsim.base.output.time_signals_fft", + "class_name": "TimeSignalsK", + } + classes._set_child("TimeSignalsK", attribs=attribs) + + @staticmethod + def _complete_params_with_default(params, info_solver): + """This static method is used to complete the *params* container.""" + OutputBasePseudoSpectral._complete_params_with_default( + params, info_solver + ) + + params.output.phys_fields.field_to_plot = "rot" + + def linear_eigenmode_from_values_1k(self, ux_fft, uy_fft, eta_fft, kx, ky): + """Compute the linear eigenmodes for a single wavenumber.""" + return linear_eigenmode_from_values_1k( + ux_fft, uy_fft, eta_fft, kx, ky, self.sim.params.f, self.sim.params.c2 + ) + + def omega_from_wavenumber(self, k): + r"""Evaluates the dispersion relation and returns the linear frequency + + .. math:: \omega = \sqrt{f ^ 2 + (ck)^2} + + """ + return np.sqrt(self.sim.params.f**2 + self.sim.params.c2 * k**2) + + def compute_enstrophy_fft(self): + r"""Calculate enstrophy from vorticity in the spectral space.""" + rot_fft = self.sim.state.get_var("rot_fft") + return np.abs(rot_fft) ** 2 / 2.0 + + def compute_PV_fft(self): + r"""Compute Ertel and Charney (QG) potential vorticity. + + .. math:: \zeta_{er} = \frac{f + \zeta}{1 + \eta} + .. math:: \zeta_{ch} = \zeta - f \eta + + """ + get_var = self.sim.state.get_var + rot = get_var("rot") + eta = self.sim.state.state_phys.get_var("eta") + ErtelPV_fft = self.oper.fft2((self.sim.params.f + rot) / (1.0 + eta)) + rot_fft = get_var("rot_fft") + eta_fft = get_var("eta_fft") + CharneyPV_fft = rot_fft - self.sim.params.f * eta_fft + return ErtelPV_fft, CharneyPV_fft + + def compute_PE_fft(self): + """Compute Ertel and Charney (QG) potential enstrophy.""" + ErtelPV_fft, CharneyPV_fft = self.compute_PV_fft() + return (abs(ErtelPV_fft) ** 2 / 2.0, abs(CharneyPV_fft) ** 2 / 2.0) + + def compute_CharneyPE_fft(self): + """Compute Charney (QG) potential enstrophy.""" + rot_fft = self.sim.state.get_var("rot_fft") + eta_fft = self.sim.state.get_var("eta_fft") + CharneyPV_fft = rot_fft - self.sim.params.f * eta_fft + return abs(CharneyPV_fft) ** 2 / 2.0 + + def compute_energies(self): + """Compute kinetic, available potential and rotational kinetic energies.""" + energyK_fft, energyA_fft, energyKr_fft = self.compute_energies_fft() + return ( + self.sum_wavenumbers(energyK_fft), + self.sum_wavenumbers(energyA_fft), + self.sum_wavenumbers(energyKr_fft), + ) + + def compute_energiesKA(self): + """Compute K.E. and A.P.E.""" + energyK_fft, energyA_fft = self.compute_energiesKA_fft() + return ( + self.sum_wavenumbers(energyK_fft), + self.sum_wavenumbers(energyA_fft), + ) + + def compute_energy(self): + """Compute total energy by summing K.E. and A.P.E.""" + energyK_fft, energyA_fft = self.compute_energiesKA_fft() + return self.sum_wavenumbers(energyK_fft) + self.sum_wavenumbers( + energyA_fft + ) + + def compute_enstrophy(self): + """Compute total enstrophy.""" + enstrophy_fft = self.compute_enstrophy_fft() + return self.sum_wavenumbers(enstrophy_fft) + + def compute_lin_energies_fft(self): + r"""Compute quadratic energies decomposed into contributions from + potential vorticity (:math:`q`), divergence (:math:`\nabla.\mathbf u`), + and ageostrophic variable (:math:`a`). + + """ + get_var = self.sim.state.get_var + ux_fft = get_var("ux_fft") + uy_fft = get_var("uy_fft") + eta_fft = get_var("eta_fft") + + q_fft, div_fft, ageo_fft = self.oper.qdafft_from_uxuyetafft( + ux_fft, uy_fft, eta_fft + ) + + udx_fft, udy_fft = self.oper.vecfft_from_divfft(div_fft) + energy_dlin_fft = 0.5 * (np.abs(udx_fft) ** 2 + np.abs(udy_fft) ** 2) + + ugx_fft, ugy_fft, etag_fft = self.oper.uxuyetafft_from_qfft(q_fft) + energy_glin_fft = 0.5 * ( + np.abs(ugx_fft) ** 2 + + np.abs(ugy_fft) ** 2 + + self.sim.params.c2 * np.abs(etag_fft) ** 2 + ) + + uax_fft, uay_fft, etaa_fft = self.oper.uxuyetafft_from_afft(ageo_fft) + energy_alin_fft = 0.5 * ( + np.abs(uax_fft) ** 2 + + np.abs(uay_fft) ** 2 + + self.sim.params.c2 * np.abs(etaa_fft) ** 2 + ) + + return energy_glin_fft, energy_dlin_fft, energy_alin_fft + + +class OutputSW1L(OutputBaseSW1L): + def compute_energies_fft(self): + r"""Compute kinetic, available potential and rotational kinetic energies + in the spectral space. + + .. math:: E_K = (h \mathbf u).\mathbf u / 2 + .. math:: E_A = c^2 \eta^2) / 2 + .. math:: E_{K,r} = (h \mathbf u_r).\mathbf u_r / 2 + + """ + get_var = self.sim.state.get_var + eta_fft = get_var("eta_fft") + energyA_fft = self.sim.params.c2 * np.abs(eta_fft) ** 2 / 2 + Jx_fft = get_var("Jx_fft") + Jy_fft = get_var("Jy_fft") + ux_fft = get_var("ux_fft") + uy_fft = get_var("uy_fft") + energyK_fft = ( + np.real(Jx_fft.conj() * ux_fft + Jy_fft.conj() * uy_fft) / 2.0 + ) + + rot_fft = get_var("rot_fft") + uxr_fft, uyr_fft = self.oper.vecfft_from_rotfft(rot_fft) + rotJ_fft = self.oper.rotfft_from_vecfft(Jx_fft, Jy_fft) + Jxr_fft, Jyr_fft = self.oper.vecfft_from_rotfft(rotJ_fft) + energyKr_fft = ( + np.real(Jxr_fft.conj() * uxr_fft + Jyr_fft.conj() * uyr_fft) / 2.0 + ) + return energyK_fft, energyA_fft, energyKr_fft + + def compute_energiesKA_fft(self): + """Compute K.E. and A.P.E in the spectral space.""" + get_var = self.sim.state.get_var + eta_fft = get_var("eta_fft") + energyA_fft = self.sim.params.c2 * np.abs(eta_fft) ** 2 / 2 + Jx_fft = get_var("Jx_fft") + Jy_fft = get_var("Jy_fft") + ux_fft = get_var("ux_fft") + uy_fft = get_var("uy_fft") + energyK_fft = ( + np.real(Jx_fft.conj() * ux_fft + Jy_fft.conj() * uy_fft) / 2.0 + ) + + return energyK_fft, energyA_fft diff --git a/fluidsim/solvers/sw1l/output/meson.build b/fluidsim/solvers/sw1l/output/meson.build index d8ba1395f..035920cf8 100644 --- a/fluidsim/solvers/sw1l/output/meson.build +++ b/fluidsim/solvers/sw1l/output/meson.build @@ -1,5 +1,6 @@ python_sources = [ '__init__.py', + 'base.py', 'increments.py', 'normal_mode.py', '_old_spatial_means.py', @@ -14,3 +15,9 @@ py.install_sources( python_sources, subdir: 'fluidsim/solvers/sw1l/output' ) + +run_command(['transonic', '--meson', '--backend', backend, 'base.py'], check: true) + +foreach be : backends + subdir('__' + be + '__') +endforeach diff --git a/fluidsim/solvers/sw1l/solver.py b/fluidsim/solvers/sw1l/solver.py index a138afc86..cac027e87 100644 --- a/fluidsim/solvers/sw1l/solver.py +++ b/fluidsim/solvers/sw1l/solver.py @@ -17,7 +17,7 @@ """ -from transonic import jit, Array +from transonic import boost, Array from fluiddyn.util import mpi from fluidsim.base.setofvariables import SetOfVariables @@ -29,7 +29,7 @@ A = Array[float, "2d"] -@jit +@boost def compute_Frot(rot: A, ux: A, uy: A, f: float): """Compute cross-product of absolute potential vorticity with velocity.""" if f != 0: @@ -43,7 +43,7 @@ def compute_Frot(rot: A, ux: A, uy: A, f: float): return F1x, F1y -@jit +@boost def compute_pressure(c2: float, eta: A, ux: A, uy: A): return c2 * eta + 0.5 * (ux**2 + uy**2)