Source code for fluidsim.solvers.sw1l.modified.state

"""State class for the sw1l.modified solver
(:mod:`fluidsim.solvers.sw1l.modified.state`)
===================================================


Provides:

.. autoclass:: StateSW1LModified
   :members:
   :private-members:

"""

from fluidsim.base.setofvariables import SetOfVariables

from fluidsim.solvers.sw1l.state import StateSW1L

from fluiddyn.util import mpi


[docs] class StateSW1LModified(StateSW1L): """Contains the variables corresponding to the state and handles the access to other fields for the solver MSW1L. """
[docs] @staticmethod def _complete_info_solver(info_solver): """Complete the ParamContainer info_solver. This is a static method! """ info_solver.classes.State._set_attribs( { "keys_state_spect": ["ux_fft", "uy_fft", "eta_fft"], "keys_state_phys": ["ux", "uy", "eta"], "keys_computable": [], "keys_phys_needed": ["ux", "uy", "eta"], "keys_linear_eigenmodes": ["q_fft", "a_fft", "d_fft"], } )
[docs] def compute(self, key, SAVE_IN_DICT=True, RAISE_ERROR=True): it = self.sim.time_stepping.it if key in self.vars_computed and it == self.it_computed[key]: return self.vars_computed[key] if key == "ux_fft": result = self.oper.fft2(self.state_phys.get_var("ux")) elif key == "uy_fft": result = self.oper.fft2(self.state_phys.get_var("ux")) elif key == "rot_fft": ux_fft = self.compute("ux_fft") uy_fft = self.compute("uy_fft") result = self.oper.rotfft_from_vecfft(ux_fft, uy_fft) elif key == "div_fft": ux_fft = self.compute("ux_fft") uy_fft = self.compute("uy_fft") result = self.oper.divfft_from_vecfft(ux_fft, uy_fft) elif key == "rot": rot_fft = self.compute("rot_fft") result = self.oper.ifft2(rot_fft) elif key == "div": div_fft = self.compute("div_fft") result = self.oper.ifft2(div_fft) elif key == "q": rot = self.compute("rot") eta = self.sim.state.state_phys.get_var("eta") result = rot - self.params.f * eta elif key == "q_fft": ux_fft = self.state_spect.get_var("ux_fft") uy_fft = self.state_spect.get_var("uy_fft") eta_fft = self.state_spect.get_var("eta_fft") rot_fft = self.oper.rotfft_from_vecfft(ux_fft, uy_fft) result = rot_fft - self.params.f * eta_fft elif key == "a_fft": ux_fft = self.state_spect.get_var("ux_fft") uy_fft = self.state_spect.get_var("uy_fft") eta_fft = self.state_spect.get_var("eta_fft") result = self.oper.afft_from_uxuyetafft(ux_fft, uy_fft, eta_fft) else: to_print = 'Do not know how to compute "' + key + '".' if RAISE_ERROR: raise ValueError(to_print) else: if mpi.rank == 0: print(to_print + "\nreturn an array of zeros.") result = self.oper.create_arrayX(value=0.0) if SAVE_IN_DICT: self.vars_computed[key] = result self.it_computed[key] = it return result
[docs] def init_from_uxuyfft(self, ux_fft, uy_fft): oper = self.oper ifft2 = oper.ifft2 oper.projection_perp(ux_fft, uy_fft) oper.dealiasing(ux_fft, uy_fft) ux = ifft2(ux_fft) uy = ifft2(uy_fft) rot_fft = oper.rotfft_from_vecfft(ux_fft, uy_fft) rot = ifft2(rot_fft) eta_fft = self._etafft_no_div(ux, uy, rot) eta = ifft2(eta_fft) state_spect = self.state_spect state_spect.set_var("ux_fft", ux_fft) state_spect.set_var("uy_fft", uy_fft) state_spect.set_var("eta_fft", eta_fft) state_phys = self.state_phys state_phys.set_var("ux", ux) state_phys.set_var("uy", uy) state_phys.set_var("eta", eta)
[docs] def statephys_from_statespect(self): """Compute the state in physical space.""" ifft2 = self.oper.ifft2 ux_fft = self.state_spect.get_var("ux_fft") uy_fft = self.state_spect.get_var("uy_fft") eta_fft = self.state_spect.get_var("eta_fft") self.state_phys.set_var("ux", ifft2(ux_fft)) self.state_phys.set_var("uy", ifft2(uy_fft)) self.state_phys.set_var("eta", ifft2(eta_fft))
[docs] def return_statephys_from_statespect(self, state_spect=None): """Return the state in physical space.""" ifft2 = self.oper.ifft2 if state_spect is None: state_spect = self.state_spect ux_fft = state_spect.get_var("ux_fft") uy_fft = state_spect.get_var("uy_fft") eta_fft = state_spect.get_var("eta_fft") state_phys = SetOfVariables(like=self.state_phys) state_phys.set_var("ux", ifft2(ux_fft)) state_phys.set_var("uy", ifft2(uy_fft)) state_phys.set_var("eta", ifft2(eta_fft)) return state_phys