def tendencies_nonlin(self, state_spect=None, old=None): """Compute the "nonlinear" tendencies.""" oper = self.oper if state_spect is None: state_spect = self.state.state_spect if old is None: tendencies_fft = SetOfVariables( like=self.state.state_spect, info="tendencies_nonlin" ) else: tendencies_fft = old w_fft = state_spect.get_var("w_fft") z_fft = state_spect.get_var("z_fft") F_fft = tendencies_fft.get_var("w_fft") tendencies_fft.set_var("z_fft", w_fft) mamp_zz = oper.monge_ampere_from_fft(z_fft, z_fft) chi_fft = -oper.invlaplacian_fft(oper.fft2(mamp_zz), order=4) mamp_zchi = oper.monge_ampere_from_fft(z_fft, chi_fft) Nw_fft = oper.fft2(mamp_zchi) lap2z_fft = oper.laplacian_fft(z_fft, order=4) F_fft[:] = -lap2z_fft + Nw_fft oper.dealiasing(tendencies_fft) # ratio = self.test_tendencies_nonlin( # tendencies_fft, w_fft, z_fft, chi_fft) # print('ratio:', ratio) if self.params.forcing.enable: tendencies_fft += self.forcing.get_forcing() return tendencies_fft
class SpecificForcingPseudoSpectralCoarse(SpecificForcing): """Specific forcing for pseudo-spectra solvers""" tag = "pseudo_spectral" _key_forced_default = "rot_fft" @staticmethod def _check_forcing_shape(shape_forcing, shape): """Check if shape of the forcing array exceeds the shape of the global array. Parameters ---------- shape_forcing: array-like A single-element array containing index of largest forcing wavenumber or a tuple indincating shape of the forcing array. shape: array-like A tuple indicating the shape of an array or Operators instance. """ if any(np.greater(shape_forcing, shape)): raise NotImplementedError( "The resolution is too small for the required forcing: " "any(np.greater({}, {}))".format(shape_forcing, shape) ) def __init__(self, sim): super().__init__(sim) params = sim.params self.forcing_fft = SetOfVariables( like=sim.state.state_spect, info="forcing_fft", value=0.0 ) if params.forcing.nkmax_forcing < params.forcing.nkmin_forcing: raise ValueError( "params.forcing.nkmax_forcing < \n" "params.forcing.nkmin_forcing" ) self.kmax_forcing = self.oper.deltak * params.forcing.nkmax_forcing self.kmin_forcing = self.oper.deltak * params.forcing.nkmin_forcing self.forcing_rate = params.forcing.forcing_rate if params.forcing.key_forced is not None: self.key_forced = params.forcing.key_forced else: self.key_forced = self._key_forced_default try: n = 2 * fftw_grid_size(params.forcing.nkmax_forcing) except ImportError: warn("To use smaller forcing arrays: pip install pulp") i = 0 while 2 * params.forcing.nkmax_forcing > 2 ** i: i += 1 n = 2 ** i self._check_forcing_shape([n], sim.oper.shapeX_seq) try: angle = self.params.forcing[self.tag].angle except AttributeError: pass else: if isinstance(angle, str): if angle.endswith("°"): angle = radians(float(angle[:-1])) else: raise ValueError( "Angle should be a string with \n" + "the degree symbol or a float in radians" ) self.angle = angle self.kxmax_forcing = np.sin(angle) * self.kmax_forcing self.kymax_forcing = np.cos(angle) * self.kmax_forcing if mpi.rank == 0: params_coarse = deepcopy(params) params_coarse.oper.nx = n # The "+ 1" aims to give some gap between the kxmax and # the boundary of the oper_coarse. try: params_coarse.oper.nx = 2 * fftw_grid_size( int(self.kxmax_forcing / self.oper.deltakx) + 1 ) except AttributeError: pass try: params_coarse.oper.ny = n try: params_coarse.oper.ny = 2 * fftw_grid_size( int(self.kymax_forcing / self.oper.deltaky) + 1 ) except AttributeError: pass except AttributeError: pass try: params_coarse.oper.nz = n except AttributeError: pass params_coarse.oper.type_fft = "sequential" params_coarse.oper.coef_dealiasing = 1.0 self.oper_coarse = sim.oper.__class__(params=params_coarse) self.shapeK_loc_coarse = self.oper_coarse.shapeK_loc self.COND_NO_F = self._compute_cond_no_forcing() self.nb_forced_modes = ( self.COND_NO_F.size - np.array(self.COND_NO_F, dtype=np.int32).sum() ) if not self.nb_forced_modes: raise ValueError("0 modes forced.") try: hasattr(self, "plot_forcing_region") except NotImplementedError: pass else: mpi.printby0( "To plot the forcing modes, you can use:\n" "sim.forcing.forcing_maker.plot_forcing_region()" ) self.ind_forcing = ( np.logical_not(self.COND_NO_F).flatten().nonzero()[0] ) self.fstate_coarse = sim.state.__class__(sim, oper=self.oper_coarse) else: self.shapeK_loc_coarse = None if mpi.nb_proc > 1: self.shapeK_loc_coarse = mpi.comm.bcast( self.shapeK_loc_coarse, root=0 ) def _compute_cond_no_forcing(self): if hasattr(self.oper_coarse, "K"): K = self.oper_coarse.K else: K = np.sqrt(self.oper_coarse.K2) COND_NO_F = np.logical_or(K > self.kmax_forcing, K < self.kmin_forcing) COND_NO_F[self.oper_coarse.shapeK_loc[0] // 2] = True COND_NO_F[:, self.oper_coarse.shapeK_loc[1] - 1] = True return COND_NO_F def put_forcingc_in_forcing(self): """Copy data from self.fstate_coarse.state_spect into forcing_fft.""" if mpi.rank == 0: state_spect = self.fstate_coarse.state_spect oper_coarse = self.oper_coarse else: state_spect = None oper_coarse = None self.oper.put_coarse_array_in_array_fft( state_spect, self.forcing_fft, oper_coarse, self.shapeK_loc_coarse ) def verify_injection_rate(self): """Verify injection rate.""" f_fft = self.forcing_fft.get_var(self.key_forced) var_fft = self.sim.state.state_spect.get_var(self.key_forced) deltat = self.sim.time_stepping.deltat P_forcing1 = np.real(f_fft.conj() * var_fft) P_forcing2 = abs(f_fft) ** 2 P_forcing2 = deltat / 2 * self.oper.sum_wavenumbers(P_forcing2) P_forcing1 = self.oper.sum_wavenumbers(P_forcing1) if mpi.rank == 0: print( "P_f = {:9.4e} ; P_1 = {:9.4e}; P_2 = {:9.4e}".format( P_forcing1 + P_forcing2, P_forcing1, P_forcing2 ) ) def verify_injection_rate_coarse(self, var_fft=None): """Verify injection rate.""" if var_fft is None: var_fft = self.sim.state.state_spect.get_var(self.key_forced) var_fft = self.oper.coarse_seq_from_fft_loc( var_fft, self.shapeK_loc_coarse ) if mpi.rank == 0: f_fft = self.fstate_coarse.get_var(self.key_forced) deltat = self.sim.time_stepping.deltat P_forcing1 = np.real(f_fft.conj() * var_fft) P_forcing2 = abs(f_fft) ** 2 P_forcing2 = deltat / 2 * self.oper_coarse.sum_wavenumbers(P_forcing2) P_forcing1 = self.oper_coarse.sum_wavenumbers(P_forcing1) print( "P_f = {:9.4e} ; P_1 = {:9.4e}; P_2 = {:9.4e} (coarse)".format( P_forcing1 + P_forcing2, P_forcing1, P_forcing2 ) ) def compute(self): """compute a forcing normalize with a 2nd degree eq.""" if mpi.rank == 0: obj = self.compute_forcingc_fft_each_time() if isinstance(obj, dict): kwargs = obj else: kwargs = {self.key_forced: obj} self.fstate_coarse.init_statespect_from(**kwargs) self.oper_coarse.dealiasing_setofvar(self.fstate_coarse.state_spect) self.put_forcingc_in_forcing() def compute_forcingc_fft_each_time(self): raise NotImplementedError
def tendencies_nonlin(self, state_spect=None, old=None): r"""Compute the nonlinear tendencies of the solver ns2d.strat. Parameters ---------- state_spect : :class:`fluidsim.base.setofvariables.SetOfVariables` optional Array containing the state, i.e. the vorticity, in Fourier space. If `state_spect`, the variables vorticity and the velocity are computed from it, otherwise, they are taken from the global state of the simulation, `self.state`. These two possibilities are used during the Runge-Kutta time-stepping. Returns ------- tendencies_fft : :class:`fluidsim.base.setofvariables.SetOfVariables` An array containing the tendencies for the vorticity and the buoyancy perturbation. Notes ----- .. |p| mathmacro:: \partial The 2D Navier-Stokes equation can be written .. math:: \p_t \hat\zeta = \hat N(\zeta) - \sigma(k) \hat \zeta, and .. math:: \p_t \hat b = \hat N(b) - \sigma(k) \hat b This function compute the nonlinear terms ("tendencies") :math:`N(\zeta) = - \mathbf{u}\cdot \mathbf{\nabla} \zeta + \mathbf{\nabla}\wedge b\mathbf{e_z} = - \mathbf{u}\cdot \mathbf{\nabla} \zeta + \p_x b` and :math:`N(b) = - \mathbf{u}\cdot \mathbf{\nabla} b + N^2u_y`. """ oper = self.oper fft_as_arg = oper.fft_as_arg ifft_as_arg = oper.ifft_as_arg if old is None: tendencies_fft = SetOfVariables(like=self.state.state_spect) else: tendencies_fft = old f_rot_fft = tendencies_fft.get_var("rot_fft") f_b_fft = tendencies_fft.get_var("b_fft") if state_spect is None: rot_fft = self.state.state_spect.get_var("rot_fft") b_fft = self.state.state_spect.get_var("b_fft") ux = self.state.state_phys.get_var("ux") uy = self.state.state_phys.get_var("uy") else: rot_fft = state_spect.get_var("rot_fft") b_fft = state_spect.get_var("b_fft") ux_fft, uy_fft = oper.vecfft_from_rotfft(rot_fft) ux = self.state.field_tmp0 uy = self.state.field_tmp1 ifft_as_arg(ux_fft, ux) ifft_as_arg(uy_fft, uy) px_rot_fft, py_rot_fft = oper.gradfft_from_fft(rot_fft) px_b_fft, py_b_fft = oper.gradfft_from_fft(b_fft) px_rot = self.state.field_tmp2 py_rot = self.state.field_tmp3 px_b = self.state.field_tmp4 py_b = self.state.field_tmp5 ifft_as_arg(px_rot_fft, px_rot) ifft_as_arg(py_rot_fft, py_rot) ifft_as_arg(px_b_fft, px_b) ifft_as_arg(py_b_fft, py_b) f_rot, f_b = tendencies_nonlin_ns2dstrat(ux, uy, px_rot, py_rot, px_b, py_b, self.params.N) fft_as_arg(f_b, f_b_fft) fft_as_arg(f_rot, f_rot_fft) f_rot_fft += px_b_fft oper.dealiasing(tendencies_fft) # # CHECK NON-LINEAR TRANSFER rot # T_rot = np.real(f_rot_fft.conj()*rot_fft) # print(('sum(T_rot) = {0:9.4e} ; sum(abs(T_rot)) = {1:9.4e}').format( # self.oper.sum_wavenumbers(T_rot), # self.oper.sum_wavenumbers(abs(T_rot)))) # # CHECK NON-LINEAR TRANSFER b # T_b = np.real(f_b_fft.conj()*b_fft) # print(('sum(T_b) = {0:9.4e} ; sum(abs(T_b)) = {1:9.4e}').format( # self.oper.sum_wavenumbers(T_b), # self.oper.sum_wavenumbers(abs(T_b)))) if self.params.forcing.enable: tendencies_fft += self.forcing.get_forcing() # CHECK ENERGY CONSERVATION # self.check_energy_conservation(rot_fft, b_fft, f_rot_fft, f_b_fft) return tendencies_fft
def tendencies_nonlin(self, state_spect=None, old=None): oper = self.oper ifft_as_arg = oper.ifft_as_arg ifft_as_arg_destroy = oper.ifft_as_arg_destroy fft_as_arg = oper.fft_as_arg if state_spect is None: spect_get_var = self.state.state_spect.get_var else: spect_get_var = state_spect.get_var vx_fft = spect_get_var("vx_fft") vy_fft = spect_get_var("vy_fft") vz_fft = spect_get_var("vz_fft") b_fft = spect_get_var("b_fft") omegax_fft, omegay_fft, omegaz_fft = oper.rotfft_from_vecfft( vx_fft, vy_fft, vz_fft ) if self.params.f is not None: self._modif_omegafft_with_f(omegax_fft, omegay_fft, omegaz_fft) omegax = self.state.fields_tmp[3] omegay = self.state.fields_tmp[4] omegaz = self.state.fields_tmp[5] ifft_as_arg_destroy(omegax_fft, omegax) ifft_as_arg_destroy(omegay_fft, omegay) ifft_as_arg_destroy(omegaz_fft, omegaz) if state_spect is None: vx = self.state.state_phys.get_var("vx") vy = self.state.state_phys.get_var("vy") vz = self.state.state_phys.get_var("vz") else: vx = self.state.fields_tmp[0] vy = self.state.fields_tmp[1] vz = self.state.fields_tmp[2] ifft_as_arg(vx_fft, vx) ifft_as_arg(vy_fft, vy) ifft_as_arg(vz_fft, vz) fx, fy, fz = vector_product(vx, vy, vz, omegax, omegay, omegaz) if old is None: tendencies_fft = SetOfVariables( like=self.state.state_spect, info="tendencies_nonlin" ) else: tendencies_fft = old fx_fft = tendencies_fft.get_var("vx_fft") fy_fft = tendencies_fft.get_var("vy_fft") fz_fft = tendencies_fft.get_var("vz_fft") fft_as_arg(fx, fx_fft) fft_as_arg(fy, fy_fft) fft_as_arg(fz, fz_fft) fz_fft += b_fft oper.project_perpk3d(fx_fft, fy_fft, fz_fft) if state_spect is None: b = self.state.state_phys.get_var("b") else: b = self.state.fields_tmp[3] ifft_as_arg(b_fft, b) fb_fft = ( -oper.div_vb_fft_from_vb(vx, vy, vz, b) - self.params.N ** 2 * vz_fft ) tendencies_fft.set_var("b_fft", fb_fft) if self.is_forcing_enabled: tendencies_fft += self.forcing.get_forcing() return tendencies_fft
def tendencies_nonlin(self, state_spect=None, old=None): r"""Compute the nonlinear tendencies. Parameters ---------- state_spect : :class:`fluidsim.base.setofvariables.SetOfVariables` optional Array containing the state, i.e. the vorticity, the divergence and displacement in the spectral space. If ``state_spect`` is provided, the variables in the physical space are computed from it, otherwise, they are taken from the global state of the simulation, ``self.state.state_phys``. These two possibilities are used during the time-stepping. old : :class:`fluidsim.base.setofvariables.SetOfVariables` optional Array containing the previous ``tendencies_sh``. This array is reused to save memory and improve performance. Returns ------- tendencies_sh : :class:`fluidsim.base.setofvariables.SetOfVariables` An array containing the tendencies for the vorticity. Notes ----- .. |p| mathmacro:: \partial The 1-layer shallow water equations are solved in the vector-invariant form (see section 2.2.6 Vallis 2nd edition). .. math:: \p_t \hat\zeta = \hat N_\zeta - \sigma(k) \hat \zeta, \p_t \hat\delta = \hat N_\delta - \sigma(k) \hat \delta, \p_t \hat\eta = \hat N_\eta - \sigma(k) \hat \eta, This function computes the nonlinear term ("tendencies"). The algorithm is as follows, - Compute :math:`N_u` and :math:`N_v`, the tendencies for the velocities. - Take divergence and curl of the above to obtain :math:`N_\zeta, N_\delta`. - Subtract laplacian of total energy K.E. + hydrostatic pressure from :math:`N_\delta`. - Compute :math:`N_\eta = -\nabla.((\eta + 1)\mathbf{u})` """ # the spherical harmonics operator oper = self.oper if state_spect is None: state_spect = self.state.state_spect state_phys = self.state.state_phys else: state_phys = self.state.return_statephys_from_statespect( state_spect) # get or compute rot_sh, ux and uy ux = state_phys.get_var("ux") uy = state_phys.get_var("uy") eta = state_phys.get_var("eta") rot = state_phys.get_var("rot") if old is None: tendencies_sh = SetOfVariables(like=self.state.state_spect) else: tendencies_sh = old Frot_sh = tendencies_sh.get_var("rot_sh") Fdiv_sh = tendencies_sh.get_var("div_sh") Feta_sh = tendencies_sh.get_var("eta_sh") c2 = self.params.c2 # Absolute vorticity x Velocity F*x, Fuy = compute_Frot(rot, ux, uy, oper.f_radial) oper.divrotsh_from_vec(F*x, Fuy, Fdiv_sh, Frot_sh) # Subtract laplacian of K.E. + hydrostatic pressure term from # divergence tendency Fdiv_sh += oper.laplacian_sh(oper.sht(0.5 * (ux**2 + uy**2) + c2 * eta), negative=True) # Calculate Feta_sh = \nabla.(hu) = \nabla.((1 + \eta)u) oper.divrotsh_from_vec(-(eta + 1) * ux, -(eta + 1) * uy, Feta_sh) # oper.dealiasing(tendencies_sh) # def check_conservation(Fvar_sh, var_sh, var_str): # import numpy as np # T = np.real(Fvar_sh.conj() * var_sh + Fvar_sh * var_sh.conj()) / 2.0 # print( # ("sum(T_{2}) = {0:9.4e} ; sum(abs(T_{2})) = {1:9.4e}").format( # self.oper.sum_wavenumbers(T), # self.oper.sum_wavenumbers(abs(T)), # var_str # ), # end =" | " # ) # rot_sh = state_spect.get_var("rot_sh") # div_sh = state_spect.get_var("div_sh") # check_conservation(Frot_sh, rot_sh, "rot") # check_conservation(Fdiv_sh, div_sh, "div") # print() if self.params.forcing.enable: tendencies_sh += self.forcing.get_forcing() return tendencies_sh
class StatePseudoSpectral(StateBase): """Contains the state variables and handles the access to fields. This is the general class for the pseudo-spectral solvers. Parameters ---------- sim : child of :class:`fluidsim.base.solvers.base.SimulBase` oper : Optional[operators] """ @staticmethod def _complete_info_solver(info_solver): """Complete the ParamContainer info_solver. This is a static method! """ StateBase._complete_info_solver(info_solver) info_solver.classes.State.keys_state_phys = ["ux", "uy"] info_solver.classes.State.keys_phys_needed = ["ux", "uy"] info_solver.classes.State._set_attribs( {"keys_state_spect": ["ux_fft", "uy_fft"]}) def __init__(self, sim, oper=None): super().__init__(sim, oper) self.keys_state_spect = sim.info.solver.classes.State.keys_state_spect self.state_spect = SetOfVariables( keys=self.keys_state_spect, shape_variable=self.oper.shapeK_loc, dtype=np.complex128, info="state_spect", ) def has_vars(self, *keys): """Checks if all of the keys are present in the union of ``keys_state_phys``, ``keys_computable``, and ``keys_state_spect``. Parameters ---------- keys: str, str ... Strings indicating state variable names. Returns ------- bool Examples -------- >>> sim.state.has_vars('ux', 'uy', 'ux_fft') >>> sim.state.has_vars('rot') """ keys_state = set(self.keys_state_phys + self.keys_computable + self.keys_state_spect) keys = set(keys) return keys.issubset(keys_state) def get_var(self, key): """Get a variable (from the storage arrays or computed). This is one of the main method of the state classes. It tries to return the array corresponding to a physical variable. If it is stored in the main storage arrays (`state_phys` and `state_spec`) of the state class, it is directly returned. Otherwise, we try to compute the quantity with the method :func:`compute`. It should not be necessary to redefine this method in child class. """ if key in self.keys_state_spect: return self.state_spect.get_var(key) elif key in self.keys_state_phys: return self.state_phys.get_var(key) else: it = self.sim.time_stepping.it if key in self.vars_computed and it == self.it_computed[key]: return self.vars_computed[key] else: value = self.compute(key) self.vars_computed[key] = value self.it_computed[key] = it return value def __setitem__(self, key, value): """General setter function to set the value of a variable It should not be necessary to redefine this method in child class. """ if key in self.keys_state_spect: self.state_spect.set_var(key, value) elif key in self.keys_state_phys: self.state_phys.set_var(key, value) else: raise ValueError('key "' + key + '" is not known') def statespect_from_statephys(self): """Compute the spectral variables from the physical variables. When you implement a new solver, check that this method does the job! """ for ik in range(self.state_spect.nvar): self.oper.fft_as_arg(self.state_phys[ik], self.state_spect[ik]) def statephys_from_statespect(self): """Compute the physical variables from the spectral variables. When you implement a new solver, check that this method does the job! """ for ik in range(self.state_spect.nvar): self.oper.ifft_as_arg(self.state_spect[ik], self.state_phys[ik]) def return_statephys_from_statespect(self, state_spect=None): """Return the physical variables computed from the spectral variables.""" ifft = self.oper.ifft if state_spect is None: state_spect = self.state_spect state_phys = SetOfVariables(like=self.state_phys) for ik in range(self.state_spect.nvar): state_phys[ik] = ifft(state_spect[ik]) return state_phys def can_this_key_be_obtained(self, key): return (key in self.keys_state_phys or key in self.keys_computable or key in self.keys_state_spect) def init_statespect_from(self, **kwargs): """Initialize `state_spect` from arrays. Parameters ---------- **kwargs : {key: array, ...} keys and arrays used for the initialization. The other keys are set to zero. Examples -------- .. code-block:: python kwargs = {'a_fft': Fa_fft} init_statespect_from(**kwargs) ux_fft, uy_fft, eta_fft = oper.uxuyetafft_from_qfft(q_fft) init_statespect_from(ux_fft=ux_fft, uy_fft=uy_fft, eta_fft=eta_fft) """ self.state_spect[:] = 0.0 for key, value in list(kwargs.items()): if key not in self.keys_state_spect: raise ValueError( f'Do not know how to initialize with key "{key}".') self.state_spect.set_var(key, value)
class StateBase: """Contains the state variables and handles the access to fields. Parameters ---------- sim : child of :class:`fluidsim.base.solvers.base.SimulBase` oper : Optional[operators] """ @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_phys": ["X"], "keys_computable": [], "keys_phys_needed": ["X"], }) def __init__(self, sim, oper=None): self.sim = sim self.params = sim.params if oper is None: self.oper = sim.oper else: self.oper = oper # creation of the SetOfVariables state_spect and state_phys self.keys_state_phys = sim.info.solver.classes.State.keys_state_phys try: self.keys_computable = sim.info.solver.classes.State.keys_computable except AttributeError: self.keys_computable = [] self.state_phys = SetOfVariables( keys=self.keys_state_phys, shape_variable=self.oper.shapeX_loc, dtype=np.float64, info="state_phys", ) self.vars_computed = {} self.it_computed = {} self.is_initialized = False def compute(self, key): """Compute a not stored variable from the stored variables""" raise ValueError('No method to compute key "' + key + '"') def clear_computed(self): """Clear the stored computed variables.""" self.vars_computed.clear() def has_vars(self, *keys): """Checks if all of the keys are present in the union of ``keys_state_phys`` and ``keys_computable``. Parameters ---------- keys: str, str ... Strings indicating state variable names. Returns ------- bool Examples -------- >>> sim.state.has_vars('ux', 'uy') >>> sim.state.has_vars('ux') >>> sim.state.has_vars('ux', 'vx', strict=False) .. todo:: ``strict=True`` can be a Python 3 compatible keywords-only argument with the function like:: def has_vars(self, *keys, strict=True): ... if strict: return keys.issubset(keys_state) else: return len(keys.intersection(keys_state)) > 0 When ``True``, checks if all keys form a subset of state keys. When ``False``, checks if the intersection of the keys and the state keys has atleast one member. """ keys_state = set(self.keys_state_phys + self.keys_computable) keys = set(keys) return keys.issubset(keys_state) def get_var(self, key): """Get a physical variable (from the storage array or computed). This is one of the main method of the state classes. It tries to return the array corresponding to a physical variable. If it is stored in the main storage array of the state class, it is directly returned. Otherwise, we try to compute the quantity with the method :func:`compute`. It should not be necessary to redefine this method in child class. """ if key in self.keys_state_phys: return self.state_phys.get_var(key) else: it = self.sim.time_stepping.it if key in self.vars_computed and it == self.it_computed[key]: return self.vars_computed[key] else: value = self.compute(key) self.vars_computed[key] = value self.it_computed[key] = it return value def __call__(self, key): raise DeprecationWarning("Do not call a state object. " "Instead, use get_var method.") def __setitem__(self, key, value): """General setter function to set the value of a variable It should not be necessary to redefine this method in child class. """ if key in self.keys_state_phys: self.state_phys.set_var(key, value) else: raise ValueError('key "' + key + '" is not known') def can_this_key_be_obtained(self, key): """To check whether a variable can be obtained. .. deprecated:: 0.2.0 Use ``has_vars`` method instead. """ raise DeprecationWarning("Do not call can_this_key_be_obtained. " "Instead, use has_vars method.") def init_statephys_from(self, **kwargs): """Initialize `state_phys` from arrays. Parameters ---------- **kwargs : {key: array, ...} keys and arrays used for the initialization. The other keys are set to zero. Examples -------- .. code-block:: python kwargs = {'a': Fa} init_statespect_from(**kwargs) init_statespect_from(ux=ux, uy=uy, eta=eta) """ self.state_phys[:] = 0.0 for key, value in list(kwargs.items()): if key not in self.keys_state_phys: raise ValueError( f'Do not know how to initialize with key "{key}".') self.state_phys.set_var(key, value)
def tendencies_nonlin(self, state_spect=None, old=None): r"""Compute the nonlinear tendencies. Parameters ---------- state_spect : :class:`fluidsim.base.setofvariables.SetOfVariables` optional Array containing the state, i.e. the vorticity, in Fourier space. If `state_spect`, the variables vorticity and the velocity are computed from it, otherwise, they are taken from the global state of the simulation, `self.state`. These two possibilities are used during the Runge-Kutta time-stepping. Returns ------- tendencies_fft : :class:`fluidsim.base.setofvariables.SetOfVariables` An array containing the tendencies for the vorticity. Notes ----- .. |p| mathmacro:: \partial The 2D Navier-Stokes equation can be written as: .. math:: \p_t \hat\zeta = \hat N(\zeta) - \sigma(k) \hat \zeta, This function compute the nonlinear term ("tendencies") :math:`N(\zeta) = - \mathbf{u}\cdot \mathbf{\nabla} \zeta`. """ # the operator and the fast Fourier transform oper = self.oper ifft_as_arg = oper.ifft_as_arg # get or compute rot_fft, ux and uy if state_spect is None: rot_fft = self.state.state_spect.get_var("rot_fft") ux = self.state.state_phys.get_var("ux") uy = self.state.state_phys.get_var("uy") else: rot_fft = state_spect.get_var("rot_fft") ux_fft, uy_fft = oper.vecfft_from_rotfft(rot_fft) ux = self.state.field_tmp0 uy = self.state.field_tmp1 ifft_as_arg(ux_fft, ux) ifft_as_arg(uy_fft, uy) # "px" like $\partial_x$ px_rot_fft, py_rot_fft = oper.gradfft_from_fft(rot_fft) px_rot = self.state.field_tmp2 py_rot = self.state.field_tmp3 ifft_as_arg(px_rot_fft, px_rot) ifft_as_arg(py_rot_fft, py_rot) Frot = compute_Frot(ux, uy, px_rot, py_rot, self.params.beta) if old is None: tendencies_fft = SetOfVariables(like=self.state.state_spect) else: tendencies_fft = old Frot_fft = tendencies_fft.get_var("rot_fft") oper.fft_as_arg(Frot, Frot_fft) oper.dealiasing(Frot_fft) # import numpy as np # T_rot = np.real(Frot_fft.conj()*rot_fft + Frot_fft*rot_fft.conj())/2. # print(('sum(T_rot) = {0:9.4e} ; sum(abs(T_rot)) = {1:9.4e}' # ).format(self.oper.sum_wavenumbers(T_rot), # self.oper.sum_wavenumbers(abs(T_rot)))) if self.params.forcing.enable: tendencies_fft += self.forcing.get_forcing() return tendencies_fft
def tendencies_nonlin(self, state_spect=None, old=None): r"""Compute the nonlinear tendencies. Parameters ---------- state_spect : :class:`fluidsim.base.setofvariables.SetOfVariables` optional Array containing the state, i.e. the horizontal velocities and surface displacement scalar in Fourier space. When `state_spect` is provided, the variables vorticity and the velocities and surface displacement are computed from it, otherwise, they are taken from the global state of the simulation, `self.state`. These two possibilities are used during the Runge-Kutta time-stepping. Returns ------- tendencies_fft : :class:`fluidsim.base.setofvariables.SetOfVariables` An array containing the tendencies for the velocities :math:`\mathbf u` and the surface displacement scalar :math:`\eta`. Notes ----- .. |p| mathmacro:: \partial The one-layer shallow water equations can be written as: .. math:: \partial_t \mathbf u = - \nabla (|\mathbf u|^2/2 + c^2 \eta) - (\zeta + f) \times \mathbf u .. math:: \partial_t \eta = - \nabla. ((\eta + H) \mathbf u) This function computes all terms on the RHS of the equations above, including the nonlinear term. The linear dissipation term (not shown above) is computed implicitly, as demonstrated in :class:`fluidsim.base.time_stepping.pseudo_spect.TimeSteppingPseudoSpectral`. """ oper = self.oper if state_spect is None: state_phys = self.state.state_phys else: state_phys = self.state.return_statephys_from_statespect( state_spect) ux = state_phys.get_var("ux") uy = state_phys.get_var("uy") eta = state_phys.get_var("eta") rot = state_phys.get_var("rot") if old is None: tendencies_fft = SetOfVariables(like=self.state.state_spect, info="tendencies_nonlin") else: tendencies_fft = old F1x, F1y = compute_Frot(rot, ux, uy, self.params.f) gradx_fft, grady_fft = oper.gradfft_from_fft( oper.fft2(compute_pressure(self.params.c2, eta, ux, uy))) oper.dealiasing(gradx_fft, grady_fft) Fx_fft = tendencies_fft.get_var("ux_fft") Fy_fft = tendencies_fft.get_var("uy_fft") Feta_fft = tendencies_fft.get_var("eta_fft") Fx_fft[:] = oper.fft2(F1x) - gradx_fft Fy_fft[:] = oper.fft2(F1y) - grady_fft Feta_fft[:] = -oper.divfft_from_vecfft(oper.fft2( (eta + 1) * ux), oper.fft2((eta + 1) * uy)) oper.dealiasing(tendencies_fft) if self.params.forcing.enable: tendencies_fft += self.forcing.get_forcing() return tendencies_fft