diff --git a/.coveragerc b/.coveragerc index 9670e542..625961df 100644 --- a/.coveragerc +++ b/.coveragerc @@ -1,9 +1,9 @@ [run] source = src omit = - src/data_assimilation/* - src/inputs/* - src/utils/debug_helpers.py + src/pybella/data_assimilation/* + src/pybella/inputs/* + src/pybella/utils/debug_helpers.py */tests/* */test_* setup.py diff --git a/README.md b/README.md index 6a02af2b..5475dda5 100644 --- a/README.md +++ b/README.md @@ -8,8 +8,8 @@

- -GitHub Actions: docs + +GitHub Actions: docs open issues diff --git a/run_scripts/test_suite.py b/run_scripts/test_suite.py index 4cd596a5..8598f146 100644 --- a/run_scripts/test_suite.py +++ b/run_scripts/test_suite.py @@ -43,7 +43,6 @@ diag = td.compare_sol("gen_target") diag.update_targets() - ud["output_type"] = "test" # Do diagnostics ud["diag"] = True diff --git a/src/pybella/__main__.py b/src/pybella/__main__.py index b27d2848..253323b7 100644 --- a/src/pybella/__main__.py +++ b/src/pybella/__main__.py @@ -4,6 +4,8 @@ import numpy as np +from .flow_solver.utils import prepare + # dependencies of the atmospheric flow solver from .flow_solver.discretisation import time_update as dis_time_update @@ -16,7 +18,7 @@ from .data_assimilation import prepare as da_prepare, analysis as da_analysis # package imports -from .utils import prepare, io, sim_params as params +from .utils import io, sim_params as params ########################################################## @@ -67,23 +69,21 @@ def main(): debug_writer = io.create_debug_writer(params.debug, writer, mem) logging.info("For ensemble member = %i..." % cnt) - mem = dis_time_update.do( - sst, - mem, - tout, - blend, - step_writer, - debug_writer - ) + mem = dis_time_update.do(mem, sst.ud, tout, blend, step_writer, debug_writer) if sst.ud.diag: if sst.ud.diag_updt_targets: - strip_target.do(step_writer.OUTPUT_FILENAME + step_writer.BASE_NAME + step_writer.SUFFIX + '.h5', sst.ud, time_increment=sst.diag_comparison.time_increment) + strip_target.do( + step_writer.OUTPUT_FILENAME + + step_writer.BASE_NAME + + step_writer.SUFFIX + + ".h5", + sst.ud, + time_increment=sst.diag_comparison.time_increment, + ) sst.diag_comparison.update_targets() else: - sst.diag_comparison.test_do( - mem, sst.ud - ) + sst.diag_comparison.test_do(mem, sst.ud) futures.append(mem) diff --git a/src/pybella/data_assimilation/analysis.py b/src/pybella/data_assimilation/analysis.py index 69d2e7fd..0fcf4aa4 100644 --- a/src/pybella/data_assimilation/analysis.py +++ b/src/pybella/data_assimilation/analysis.py @@ -120,7 +120,7 @@ def do_for_window(tout, outer_step, results, sst, writer): # Update ensemble with analysis ###################################################### for mem in results: - elem, node, Sol, _, mpv, th, _ = mem + elem, node, Sol, _, mpv, th, _, _ = mem bdry.set_explicit_boundary_data(Sol, elem, sst.ud, th, mpv) p2_nodes = mpv.p2_nodes bdry.set_ghostnodes_p2(p2_nodes, node, sst.ud) diff --git a/src/pybella/data_assimilation/letkf.py b/src/pybella/data_assimilation/letkf.py index 07ef0fa4..1b7ba719 100644 --- a/src/pybella/data_assimilation/letkf.py +++ b/src/pybella/data_assimilation/letkf.py @@ -11,7 +11,7 @@ import matplotlib.pyplot as plt -from ..flow_solver.utils import options as opts +from ..utils import options as opts from . import utils debug_cnt = 0 diff --git a/src/pybella/data_assimilation/utils.py b/src/pybella/data_assimilation/utils.py index f81dc133..b1a593d2 100644 --- a/src/pybella/data_assimilation/utils.py +++ b/src/pybella/data_assimilation/utils.py @@ -5,7 +5,9 @@ import matplotlib.pyplot as plt -from ..flow_solver.utils import options as opts, boundary as bdry +from ..utils import options as opts + +from ..flow_solver.utils import boundary as bdry class ensemble(object): diff --git a/src/pybella/flow_solver/discretisation/grid.py b/src/pybella/flow_solver/discretisation/grid.py index b95bc7a1..0e4f82ab 100644 --- a/src/pybella/flow_solver/discretisation/grid.py +++ b/src/pybella/flow_solver/discretisation/grid.py @@ -1,6 +1,6 @@ import numpy as np -from ..utils import options as opts +from ...utils import options as opts def grid_init(ud): diff --git a/src/pybella/flow_solver/discretisation/time_update.py b/src/pybella/flow_solver/discretisation/time_update.py index e0f17a85..fb62afc1 100644 --- a/src/pybella/flow_solver/discretisation/time_update.py +++ b/src/pybella/flow_solver/discretisation/time_update.py @@ -3,9 +3,10 @@ import numpy as np +from ...utils import options as opts + # dependencies of the flow solver subpackage -from . import grid as dis_grid -from ..utils import boundary as bdry, options as opts +from ..utils import boundary as bdry from ..physics.gas_dynamics import ( numerical_flux as gd_flux, eos as gd_eos, @@ -18,9 +19,10 @@ from ...interfaces.dynamics_blending import schemes from ...interfaces.time_stepper import prestep + def do( - sst, mem, + ud, tout, bld=None, writer=None, @@ -29,91 +31,45 @@ def do( """ For more details, refer to the write-up :ref:`time-stepping`. - Does a time-step for the atmospheric solver. - - Parameters - ---------- - Sol : :class:`management.variable.Vars` - Solution data container. - flux : :class:`management.variable.States` - Data container for the fluxes. - mpv : :class:`physics.low_mach.mpv.MPV` - Variables relating to the elliptic solver. - t : float - Current time - tout : float - Next output time - ud : :class:`inputs.user_data.UserDataInit` - Data container for the initial conditions - elem : :class:`discretization.kgrid.ElemSpaceDiscr` - Cells grid. - node : :class:`discretization.kgrid.NodeSpaceDiscr` - Nodes grid. - step : int - Current step. - th : :class:`physics.gas_dynamics.thermodynamic.init` - Thermodynamic variables of the system - bld : :class:`data_assimilation.blending.Blend()` - Blending class used to initalise interface blending methods. - writer : :class:`management.io.io`, optional - `default == None`. If given, output after each time-step will be written in the hdf5 format. - debug : boolean, optional - `default == False`. If `True`, then writer will output `Sol`: - 1. before flux calculation - 2. before advection routine - 3. after advection routine - 4. after explicit solver - 5. after implicit solver - - during both the half-step for the prediction of advective flux and the full-step. - - Returns - ------- - list - A list of `[Sol,flux,mpv,[window_step,step]]` data containers at time `tout`. - """ - ud = sst.ud - elem, node, Sol, flux, mpv, th, time = mem + Does a time-step for the flow solver. - window_step = time.window_step - step = time.step - t = time.t + """ swe_to_lake = False - while (t < tout) and (step < ud.stepmax): - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) + while (mem.time.t < tout) and (mem.time.step < ud.stepmax): + bdry.set_explicit_boundary_data(mem.sol, mem.elem, ud, mem.th, mem.mpv) - label = "%.3d" % step + label = "%.3d" % mem.time.step - if step == 0 and writer != None: + if mem.time.step == 0 and writer != None: writer.write_all(mem, str(label) + "_ic") - dt, cfl, cfl_ac = gd_cfl.dynamic_timestep(Sol, t, tout, elem, ud, th, step) + dt, cfl, cfl_ac = gd_cfl.dynamic_timestep(mem.sol, mem.time.t, tout, mem.elem, ud, mem.th, mem.time.step) - dt = prestep.apply_modifcations(dt, ud, step) + dt = prestep.apply_modifcations(dt, ud, mem.time.step) ###################################################### # Blending : Do blending before timestep ###################################################### - swe_to_lake, Sol, mpv, t = schemes.prepare_blending( - sst, + swe_to_lake, mem.sol, mem.mpv, mem.time.t = schemes.prepare_blending( mem, + ud, bld, label, writer, - step, - window_step, - t, + mem.time.step, + mem.time.window_step, + mem.time.t, dt, swe_to_lake, debug_writer, ) - ud.is_nonhydrostatic = gd_eos.is_nonhydrostatic(ud, window_step) - ud.nonhydrostasy = gd_eos.nonhydrostasy(ud, t, window_step) + ud.is_nonhydrostatic = gd_eos.is_nonhydrostatic(ud, mem.time.window_step) + ud.nonhydrostasy = gd_eos.nonhydrostasy(ud, mem.time.t, mem.time.window_step) if ud.continuous_blending or ud.initial_blending: - logging.info(f"step = {step}, window_step = {window_step}") + logging.info(f"step = {mem.time.step}, window_step = {mem.time.window_step}") logging.info( f""" @@ -122,53 +78,45 @@ def do( compressibility = {ud.compressibility:.3f}, nonhydrostasy = {ud.nonhydrostasy:.3f} ------- """ - ) + ) - Sol0 = copy.deepcopy(Sol) + Sol0 = copy.deepcopy(mem.sol) debug_writer.write(f"{label}_before_flux") - gd_flux.recompute_advective_fluxes(flux, Sol) + gd_flux.recompute_advective_fluxes(mem) - debug_writer.populate_flux_components(f"{label}_before_advect", flux, elem) + debug_writer.populate_flux_components(f"{label}_before_advect", mem.flux, mem.elem) debug_writer.write(f"{label}_before_advect") if ud.do_advection: gd_explicit.advect_rk( - Sol, - flux, - 0.5 * dt, - elem, - step % 2, + mem, ud, - th, - mpv, - node, - str(label) + "_half", - writer, + 0.5 * dt, ) debug_writer.write(f"{label}_after_advect") - debug_writer.populate(f"{label}_after_full_step", "p2_nodes", mpv.p2_nodes) + debug_writer.populate(f"{label}_after_full_step", "p2_nodes", mem.mpv.p2_nodes) - mpv.p2_nodes0[...] = mpv.p2_nodes + mem.mpv.p2_nodes0[...] = mem.mpv.p2_nodes - lm_sp.euler_backward_non_advective_expl_part(Sol, mpv, elem, 0.5 * dt, ud, th) + lm_sp.euler_backward_non_advective_expl_part(mem, ud, 0.5 * dt) debug_writer.write(f"{label}_after_ebnaexp") Sol0_increment = Sol0 if ud.is_compressible == 0 else None lm_sp.euler_backward_non_advective_impl_part( - Sol, - mpv, - elem, - node, + mem.sol, + mem.mpv, + mem.elem, + mem.node, ud, - th, - t, + mem.th, + mem.time.t, 0.5 * dt, - 1.0, + mem, Sol0=Sol0_increment, label=f"{label}_after_ebnaimp", writer=writer, @@ -176,47 +124,43 @@ def do( if ud.bdry_type[1] == opts.BdryType.RAYLEIGH: # top rayleight damping - bdry.rayleigh_damping(Sol, mpv, ud, elem, node) + bdry.rayleigh_damping(mem.sol, mem.mpv, ud, mem.elem, mem.node) bdry.apply_rayleigh_forcing( - Sol, - mpv, + mem.sol, + mem.mpv, ud, - elem, - node, - t, - step, + mem.elem, + mem.node, + mem.time.t, + mem.time.step, dt, - th, + mem.th, bdry, ) debug_writer.write(f"{label}_after_ebnaimp") - gd_flux.recompute_advective_fluxes(flux, Sol) + gd_flux.recompute_advective_fluxes(mem) - debug_writer.populate_flux_components(f"{label}_after_half_step", flux, elem) + debug_writer.populate_flux_components(f"{label}_after_half_step", mem.flux, mem.elem) debug_writer.write(f"{label}_after_half_step") - Sol_half_new = copy.deepcopy(Sol) - mpv_half_new = copy.deepcopy(mpv) - mpv.p2_nodes_half = np.copy(mpv.p2_nodes) + Sol_half_new = copy.deepcopy(mem.sol) + mpv_half_new = copy.deepcopy(mem.mpv) + mem.mpv.p2_nodes_half = np.copy(mem.mpv.p2_nodes) if ud.is_nonhydrostatic == 0 or ( ud.is_compressible == 1 and ud.is_nonhydrostatic == 1 ): - mpv.p2_nodes[...] = mpv.p2_nodes0 + mem.mpv.p2_nodes[...] = mem.mpv.p2_nodes0 - Sol = copy.deepcopy(Sol0) + mem.sol = copy.deepcopy(Sol0) lm_sp.euler_forward_non_advective( - Sol, - mpv, - elem, - node, - 0.5 * dt, + mem, ud, - th, + 0.5 * dt, writer=writer, label=str(label) + "_after_efna", ) @@ -225,54 +169,49 @@ def do( if ud.do_advection: gd_explicit.advect( - Sol, - flux, - dt, - elem, - step % 2, + mem, ud, - th, - mpv, - node, + dt, + mem.time.step % 2, str(label) + "_full", writer, ) debug_writer.write(f"{label}_after_full_advect") - lm_sp.euler_backward_non_advective_expl_part(Sol, mpv, elem, 0.5 * dt, ud, th) + lm_sp.euler_backward_non_advective_expl_part(mem, ud, 0.5 * dt) debug_writer.write(f"{label}_after_full_ebnaexp") lm_sp.euler_backward_non_advective_impl_part( - Sol, - mpv, - elem, - node, + mem.sol, + mem.mpv, + mem.elem, + mem.node, ud, - th, - t, + mem.th, + mem.time.t, 0.5 * dt, - 2.0, + mem, writer=writer, label=str(label) + "_after_full_step", ) if ud.bdry_type[1] == opts.BdryType.RAYLEIGH: # top rayleight damping - bdry.rayleigh_damping(Sol, mpv, ud, elem, node) + bdry.rayleigh_damping(mem.sol, mem.mpv, ud, mem.elem, mem.node) # bottom rayleigh forcing bdry.apply_rayleigh_forcing( - Sol, - mpv, + mem.sol, + mem.mpv, ud, - elem, - node, - t, - step, + mem.elem, + mem.node, + mem.time.t, + mem.time.step, dt, - th, + mem.th, bdry, half=False, Sol_half_new=Sol_half_new, @@ -282,31 +221,27 @@ def do( ###################################################### # Blending : Do blending after timestep ###################################################### - Sol, mpv = schemes.blending_after_timestep( - Sol, - flux, - mpv, + mem.sol, mem.mpv = schemes.blending_after_timestep( + mem.sol, + mem.flux, + mem.mpv, bld, - elem, - node, - th, + mem.elem, + mem.node, + mem.th, ud, label, writer, - step, - window_step, - t, + mem.time.step, + mem.time.window_step, + mem.time.t, dt, swe_to_lake, debug_writer, ) - mem.sol = Sol - mem.flux = flux - mem.mpv = mpv - if writer != None: - writer.time = t + writer.time = mem.time.t writer.write_all(mem, str(label) + "_after_full_step") logging.info( @@ -314,19 +249,14 @@ def do( ) logging.info( "step %i done, t = %.12f, dt = %.12f, CFL = %.8f, CFL_ac = %.8f" - % (step, t, dt, cfl, cfl_ac) + % (mem.time.step, mem.time.t, dt, cfl, cfl_ac) ) logging.info( "###############################################################################################" ) - t += dt - step += 1 - window_step += 1 - - mem.time.t = t - mem.time.step = step - mem.time.window_step = window_step + mem.time.t += dt + mem.time.step += 1 + mem.time.window_step += 1 return mem - # return [Sol, flux, mpv, [window_step, step]] diff --git a/src/pybella/flow_solver/physics/gas_dynamics/cfl.py b/src/pybella/flow_solver/physics/gas_dynamics/cfl.py index 1badfcf4..e1eb709f 100644 --- a/src/pybella/flow_solver/physics/gas_dynamics/cfl.py +++ b/src/pybella/flow_solver/physics/gas_dynamics/cfl.py @@ -1,78 +1,111 @@ import numpy as np -machine_epsilon = np.finfo(float).eps - - def dynamic_timestep(Sol, time, time_output, elem, ud, th, step): - global machine_epsilon - - gamm = th.gamm - - Minv = 1.0 / np.sqrt(ud.Msq) - CFL = ud.CFL - - u_max = machine_epsilon - v_max = machine_epsilon - w_max = machine_epsilon - - upc_max = machine_epsilon - vpc_max = machine_epsilon - wpc_max = machine_epsilon - - p = Sol.rhoY**gamm - c = np.sqrt(gamm * p / Sol.rho) * Minv + """ + Calculate dynamic timestep for CFD simulation. + + Documentation to be homogenised. + + Args: + Sol: Solution object containing flow variables + time: Current simulation time + time_output: Target output time + elem: Element object with grid spacing + ud: User data object with CFL and timestep parameters + th: Thermodynamic properties + step: Current step number + + Returns: + If acoustic_timestep == 1: dt (float) + Else: tuple of (dt, cfl, cfl_ac) + """ + machine_epsilon = np.finfo(float).eps + + # Calculate thermodynamic properties + p = Sol.rhoY ** th.gamm + c = np.sqrt(th.gamm * p / Sol.rho) / np.sqrt(ud.Msq) + + # Calculate velocity components u = np.abs(Sol.rhou / Sol.rho) v = np.abs(Sol.rhov / Sol.rho) w = np.abs(Sol.rhow / Sol.rho) - - u_max = max(u.max(), u_max) - v_max = max(v.max(), v_max) - w_max = max(w.max(), w_max) - - upc_max = max((u + c).max(), upc_max) - vpc_max = max((v + c).max(), vpc_max) - wpc_max = max((w + c).max(), wpc_max) - + + # Find maximum velocities (with minimum threshold) + u_max = max(u.max(), machine_epsilon) + v_max = max(v.max(), machine_epsilon) + w_max = max(w.max(), machine_epsilon) + + # Calculate acoustic velocities + upc_max = max((u + c).max(), machine_epsilon) + vpc_max = max((v + c).max(), machine_epsilon) + wpc_max = max((w + c).max(), machine_epsilon) + if ud.acoustic_timestep == 1: - dtx = CFL * elem.dx / upc_max - dty = CFL * elem.dy / vpc_max - dtz = CFL * elem.dz / wpc_max - - dt_cfl = min(min(dtx, dty), dtz) - dt = min(dt_cfl, ud.dtfixed0 + min(step, 1.0) * (ud.dtfixed - ud.dtfixed0)) - - # if (2.0*dt > time_output - time): - # dt = 0.5 * (time_output - time) + machine_epsilon - if dt > (time_output - time): - dt = (time_output - time) + machine_epsilon - - return dt + return _calculate_acoustic_timestep( + ud.CFL, elem, upc_max, vpc_max, wpc_max, + time, time_output, ud, step, machine_epsilon + ) else: - dtx = CFL * elem.dx / u_max - dty = CFL * elem.dy / v_max - dtz = CFL * elem.dz / w_max - - # print("dtx=%.8f, dty=%.8f, dtz=%.8f" %(dtx,dty,dtz)) - # print("u_max=%.8f, v_max=%.8f, w_max=%.8f" %(u_max, v_max, w_max)) - - dt_cfl = min(min(dtx, dty), dtz) - # dt = min(dt_cfl, ud.dtfixed0 + min(step, 1.) * (ud.dtfixed - ud.dtfixed0)) - if step >= 0: - dt = min(dt_cfl, ud.dtfixed0 + min(step, 1.0) * (ud.dtfixed - ud.dtfixed0)) - dt *= min(float(step + 1), 1.0) - # else: - # dt = min(dt_cfl, ud.dtfixed0 + 1. * (ud.dtfixed - ud.dtfixed0)) - - # f = open("log0.txt", "a") - # f.write(str(dt_cfl) + "\n") - # f.close() - - if dt > (time_output - time): - dt = time_output - time # + machine_epsilon - - cfl = CFL * dt / dt_cfl - cfl_ac = max( - dt * upc_max / elem.dx, dt * vpc_max / elem.dy, dt * wpc_max / elem.dz + return _calculate_convective_timestep( + ud.CFL, elem, u_max, v_max, w_max, upc_max, vpc_max, wpc_max, + time, time_output, ud, step, machine_epsilon ) - return dt, cfl, cfl_ac + +def _calculate_acoustic_timestep(CFL, elem, upc_max, vpc_max, wpc_max, + time, time_output, ud, step, machine_epsilon): + """Calculate timestep based on acoustic CFL condition.""" + # Calculate directional timesteps + dt_directions = [ + CFL * elem.dx / upc_max, + CFL * elem.dy / vpc_max, + CFL * elem.dz / wpc_max + ] + dt_cfl = min(dt_directions) + + # Apply ramping factor + ramp_factor = ud.dtfixed0 + min(step, 1.0) * (ud.dtfixed - ud.dtfixed0) + dt = min(dt_cfl, ramp_factor) + + # Ensure we don't overshoot the output time + remaining_time = time_output - time + if dt > remaining_time: + dt = remaining_time + machine_epsilon + + return dt + + +def _calculate_convective_timestep(CFL, elem, u_max, v_max, w_max, + upc_max, vpc_max, wpc_max, + time, time_output, ud, step, machine_epsilon): + """Calculate timestep based on convective CFL condition.""" + # Calculate directional timesteps + dt_directions = [ + CFL * elem.dx / u_max, + CFL * elem.dy / v_max, + CFL * elem.dz / w_max + ] + dt_cfl = min(dt_directions) + + # Apply ramping and step scaling for positive steps + if step >= 0: + ramp_factor = ud.dtfixed0 + min(step, 1.0) * (ud.dtfixed - ud.dtfixed0) + dt = min(dt_cfl, ramp_factor) + dt *= min(float(step + 1), 1.0) + else: + dt = dt_cfl + + # Ensure we don't overshoot the output time + remaining_time = time_output - time + if dt > remaining_time: + dt = remaining_time + + # Calculate CFL numbers for output + cfl = CFL * dt / dt_cfl + cfl_ac = max( + dt * upc_max / elem.dx, + dt * vpc_max / elem.dy, + dt * wpc_max / elem.dz + ) + + return dt, cfl, cfl_ac \ No newline at end of file diff --git a/src/pybella/flow_solver/physics/gas_dynamics/explicit.py b/src/pybella/flow_solver/physics/gas_dynamics/explicit.py index 2d205416..2e64d80f 100644 --- a/src/pybella/flow_solver/physics/gas_dynamics/explicit.py +++ b/src/pybella/flow_solver/physics/gas_dynamics/explicit.py @@ -1,278 +1,143 @@ +from ....utils.slices import get_neighbor_indices from ...utils import boundary as bdry from . import recovery as gd_recovery from . import numerical_flux as gd_flux -def advect(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer=None): +def advect(mem, ud, dt, odd, label, writer=None): """ - Function that runs the advection routine with Strang-splitting. This function updates the `Sol` solution container with the advected solution in-place. - - Parameters - ---------- - Sol : :py:class:`management.variable.Vars` - Solution data container. - flux : :py:class:`management.variable.States` - Fluxes data container - dt : float - Time-step size. - elem : :py:class:`discretization.kgrid.ElemSpaceDiscr` - Container for cell-grid properties. - odd : int - Is current step odd or even? - ud : :py:class:`inputs.user_data.UserDataInit` - Class container for initial conditions - th : :py:class:`physics.gas_dynamics.thermodynamic.init` - Class container for thermodynamic quantities. - mpv : :py:class:`physics.low_mach.mpv.MPV` - Container for Exner pressure. - node : :py:class:`discretization.kgrid.NodeSpaceDiscr` - Container for node-grid properties. - label : string - Tag label for the output array - writer : :py:class:`management.io.io`, optional - Writer class for I/O operations, by default None + Concise implementation of Strang-splitting advection. + This function updates the `Sol` solution container with the advected solution in-place. """ - # double strang sweep time_step = 0.5 * dt - ndim = elem.ndim - stage = 0 - - # Sol.rhoX -= Sol.rho * mpv.HydroState.S0 - - if odd: - for split in range(ndim): - lmbda = time_step / elem.dxyz[split] - Sol.flip_forward() - if elem.iisc[split] > 1: - explicit_step_and_flux( - Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv - ) - else: - for i_split in range(ndim): - split = elem.ndim - 1 - i_split - lmbda = time_step / elem.dxyz[split] - if elem.iisc[split] > 1: - explicit_step_and_flux( - Sol, - flux[split], - lmbda, - elem, - split, - stage, - ud, - th, - mpv, - [writer, node, label], - ) - Sol.flip_backward() - - stage = 1 - if odd: - for i_split in range(ndim): - split = elem.ndim - 1 - i_split - lmbda = time_step / elem.dxyz[split] - if elem.iisc[split] > 1: - explicit_step_and_flux( - Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv - ) - Sol.flip_backward() - else: - for split in range(ndim): - lmbda = time_step / elem.dxyz[split] - Sol.flip_forward() - if elem.iisc[split] > 1: - explicit_step_and_flux( - Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv - ) - - # Sol.rhoX += Sol.rho * mpv.HydroState.S0 - - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - - -def explicit_step_and_flux( - Sol, flux, lmbda, elem, split_step, stage, ud, th, mpv, writer=None, tag=None -): + diagnostics = [writer, mem.node, label] if writer is not None else None + + # Define sweep configurations: (reverse_order, use_diagnostics) + sweeps = [(not odd, not odd), (odd, odd)] + + for reverse_order, use_diagnostics in sweeps: + _perform_dimensional_sweep( + mem, ud, time_step, + reverse=reverse_order, + diagnostics=diagnostics if use_diagnostics else None + ) + + bdry.set_explicit_boundary_data(mem.sol, mem.elem, ud, mem.th, mem.mpv) + + +def explicit_step_and_flux(mem, ud, lmbda, split_step, tag=None): """ - For each advection substep, solve the advection problem. For more details, see :ref:`advection_routine`. This function updates the solution `Sol` container in-place if a Strang-splitting is used, or returns the `flux` data container if a Runge-Kutta method is used. - - Parameters - ---------- - Sol : :py:class:`management.variable.Vars` - Solution data container. - flux : :py:class:`management.variable.States` - Fluxes data container - lmbda : float - :math:`\\frac{dt}{dx}`, where :math:`dx` is the grid-size in the direction of the substep. - elem : :py:class:`discretization.kgrid.ElemSpaceDiscr` - Container for cell-grid properties. - split_step : int - Tracks the substep in the Strang-splitting. - stage : int - Tracks whether the substep order goes in x-y-z or z-y-x. - ud : :py:class:`inputs.user_data.UserDataInit` - Class container for initial conditions - th : :py:class:`physics.gas_dynamics.thermodynamic.init` - Class container for thermodynamic quantities. - mpv : :py:class:`physics.low_mach.mpv.MPV` - Container for Exner pressure. - writer : :py:class:`management.io.io`, optional - Writer class for I/O operations, by default None - tag : str, optional - If `rk` then the advection routine is solved by means of a first-order Runge-Kutta method, by default None - - Returns - ------- - :py:class:`management.variable.States` - `flux` data container. + For each advection substep, solve the advection problem. For more details, see :ref:`advection_routine`. + This function updates the solution `Sol` container in-place if a Strang-splitting is used, + or returns the `flux` data container if a Runge-Kutta method is used. """ - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv, step=split_step) - - Lefts, Rights = gd_recovery.do(Sol, flux, lmbda, ud, th, elem, split_step, tag) - - # Lefts, Rights, u, Diffs, Ampls, Slopes = gd_recovery.do(Sol, flux, lmbda, ud, th, elem, split_step, tag) - - # if writer is not None: - # writer[0].write_all(Sol,mpv,elem,writer[1],th,str(writer[2])+'_split_%i' %split_step) - - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'u',u) - - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Leftsu',Lefts.u) - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Rightsu',Rights.u) - - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Leftsv',Lefts.v) - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Rightsv',Rights.v) - - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Leftsw',Lefts.w) - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Rightsw',Rights.w) - - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'LeftsrhoY',Lefts.rhoY) - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'RightsrhoY',Rights.rhoY) - - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Diffsu',Diffs.u) - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Diffsv',Diffs.v) - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Diffsw',Diffs.w) - - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Amplsu',Ampls.u) - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Amplsv',Ampls.v) - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Amplsw',Ampls.w) - - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Slopesu',Slopes.u) - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Slopesv',Slopes.v) - # if writer is not None: writer[0].populate(str(writer[2])+'_split_%i' %split_step,'Slopesw',Slopes.w) - - # skipped check_flux_bcs for now; first debug other functions - # check_flux_bcs(Lefts, Rights, elem, split_step, ud) - - flux = gd_flux.hll_solver(flux, Lefts, Rights, Sol, lmbda, ud, th) - - ndim = elem.ndim - left_idx, right_idx = [slice(None)] * ndim, [slice(None)] * ndim - right_idx[-1] = slice(1, None) - left_idx[-1] = slice(0, -1) - left_idx, right_idx = tuple(left_idx), tuple(right_idx) + flux = _compute_flux_and_recovery(mem, ud, lmbda, split_step, tag) + + # Cache neighbor indices (consider moving this to initialization if called frequently) + left_idx, right_idx = get_neighbor_indices(mem.elem.ndim) if tag != "rk": - Sol.rho += lmbda * (flux.rho[left_idx] - flux.rho[right_idx]) - Sol.rhou += lmbda * (flux.rhou[left_idx] - flux.rhou[right_idx]) - Sol.rhov += lmbda * (flux.rhov[left_idx] - flux.rhov[right_idx]) - Sol.rhow += lmbda * (flux.rhow[left_idx] - flux.rhow[right_idx]) - Sol.rhoX += lmbda * (flux.rhoX[left_idx] - flux.rhoX[right_idx]) - Sol.rhoY += lmbda * (flux.rhoY[left_idx] - flux.rhoY[right_idx]) - - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv, step=split_step) + _update_solution_variables(mem.sol, flux, lmbda, left_idx, right_idx) + + bdry.set_explicit_boundary_data(mem.sol, mem.elem, ud, mem.th, mem.mpv, step=split_step) if tag == "rk": return flux -def advect_rk(Sol, flux, dt, elem, odd, ud, th, mpv, node, label, writer=None): +def advect_rk(mem, ud, dt): """ - Function that runs the advection routine with a first-order Runge-Kutta update. This function updates the `Sol` solution container with the advected solution in-place. - - Parameters - ---------- - Sol : :py:class:`management.variable.Vars` - Solution data container. - flux : :py:class:`management.variable.States` - Fluxes data container - dt : float - Time-step size. - elem : :py:class:`discretization.kgrid.ElemSpaceDiscr` - Container for cell-grid properties. - odd : int - Is current step odd or even? - ud : :py:class:`inputs.user_data.UserDataInit` - Class container for initial conditions - th : :py:class:`physics.gas_dynamics.thermodynamic.init` - Class container for thermodynamic quantities. - mpv : :py:class:`physics.low_mach.mpv.MPV` - Container for Exner pressure. - node : :py:class:`discretization.kgrid.NodeSpaceDiscr` - Container for node-grid properties. - label : string - Tag label for the output array - writer : :py:class:`management.io.io`, optional - Writer class for I/O operations, by default None + Function that runs the advection routine with a first-order Runge-Kutta update. + This function updates the `Sol` solution container with the advected solution in-place. Attention --------- This function is not usually called unless commented out in the :py:meth:`management.data.time_update` routine. - """ - # Do 1-stages Runge-Kutta. time_step = dt - ndim = elem.ndim + ndim = mem.elem.ndim - stage = 0 - # Get RK update + # Compute fluxes for all dimensions for split in range(ndim): - lmbda = time_step / elem.dxyz[split] - Sol.flip_forward() - if elem.iisc[split] > 1: - flux[split] = explicit_step_and_flux( - Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv, tag="rk" - ) + lmbda = time_step / mem.elem.dxyz[split] + mem.sol.flip_forward() + if mem.elem.iisc[split] > 1: + mem.flux[split] = explicit_step_and_flux(mem, ud, lmbda, split, tag="rk") - ndim = elem.ndim - left_idx, right_idx = [slice(None)] * ndim, [slice(None)] * ndim - right_idx[-1] = slice(1, None) - left_idx[-1] = slice(0, -1) - left_idx, right_idx = tuple(left_idx), tuple(right_idx) + # Cache neighbor indices once + left_idx, right_idx = get_neighbor_indices(mem.elem.ndim) + # Apply flux updates for all dimensions for dim in range(ndim): - lmbda = time_step / elem.dxyz[dim] - Sol.flip_forward() - Sol.rho += lmbda * (flux[dim].rho[left_idx] - flux[dim].rho[right_idx]) - Sol.rhou += lmbda * (flux[dim].rhou[left_idx] - flux[dim].rhou[right_idx]) - Sol.rhov += lmbda * (flux[dim].rhov[left_idx] - flux[dim].rhov[right_idx]) - Sol.rhow += lmbda * (flux[dim].rhow[left_idx] - flux[dim].rhow[right_idx]) - Sol.rhoX += lmbda * (flux[dim].rhoX[left_idx] - flux[dim].rhoX[right_idx]) - Sol.rhoY += lmbda * (flux[dim].rhoY[left_idx] - flux[dim].rhoY[right_idx]) - - if dim == 1: # vertical axis - updt = lmbda * (flux[dim].rhoX[left_idx] - flux[dim].rhoX[right_idx]) - setattr(Sol, "pwchi", updt) - - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) + _apply_dimensional_flux_update(mem, dim, time_step, left_idx, right_idx) - # stage = 1 - # for split in range(ndim): - # lmbda = dt / elem.dxyz[split] - # Sol.flip_forward() - # if elem.iisc[split] > 1: - # flux[split] = explicit_step_and_flux(Sol, flux[split], lmbda, elem, split, stage, ud, th, mpv, tag='rk') + bdry.set_explicit_boundary_data(mem.sol, mem.elem, ud, mem.th, mem.mpv) - # Sol = deepcopy(Sol0) +def _update_solution_variables(sol, flux, lmbda, left_idx, right_idx, variables=None): + """ + Helper function to update solution variables with flux differences. + + """ + if variables is None: + variables = ['rho', 'rhou', 'rhov', 'rhow', 'rhoX', 'rhoY'] + + for var in variables: + flux_diff = getattr(flux, var)[left_idx] - getattr(flux, var)[right_idx] + current_val = getattr(sol, var) + setattr(sol, var, current_val + lmbda * flux_diff) - # for dim in range(ndim): - # lmbda = dt / elem.dxyz[dim] - # Sol.flip_forward() - # Sol.rho += lmbda * (flux[dim].rho[left_idx] - flux[dim].rho[right_idx]) - # Sol.rhou += lmbda * (flux[dim].rhou[left_idx] - flux[dim].rhou[right_idx]) - # Sol.rhov += lmbda * (flux[dim].rhov[left_idx] - flux[dim].rhov[right_idx]) - # Sol.rhow += lmbda * (flux[dim].rhow[left_idx] - flux[dim].rhow[right_idx]) - # Sol.rhoX += lmbda * (flux[dim].rhoX[left_idx] - flux[dim].rhoX[right_idx]) - # Sol.rhoY += lmbda * (flux[dim].rhoY[left_idx] - flux[dim].rhoY[right_idx]) - # set_explicit_boundary_data(Sol, elem, ud, th, mpv) +def _compute_flux_and_recovery(mem, ud, lmbda, split_step, tag=None): + """ + Helper function to compute flux using gradient recovery and HLL solver. + + Returns: + flux: Computed flux container + """ + flux = mem.flux[split_step] + + bdry.set_explicit_boundary_data(mem.sol, mem.elem, ud, mem.th, mem.mpv, step=split_step) + + Lefts, Rights = gd_recovery.do(mem, ud, lmbda, split_step, tag) + + flux = gd_flux.hll_solver(mem, flux, Lefts, Rights) + + return flux + +def _apply_dimensional_flux_update(mem, dim, time_step, left_idx, right_idx): + """ + Apply flux update for a specific dimension. + """ + lmbda = time_step / mem.elem.dxyz[dim] + mem.sol.flip_forward() + + _update_solution_variables(mem.sol, mem.flux[dim], lmbda, left_idx, right_idx) + + # Handle special case for vertical axis + if dim == 1: + updt = lmbda * (mem.flux[dim].rhoX[left_idx] - mem.flux[dim].rhoX[right_idx]) + setattr(mem.sol, "pwchi", updt) + +def _perform_dimensional_sweep(mem, ud, time_step, reverse=False, diagnostics=None): + """ + Perform a dimensional sweep in either forward or reverse order. + + """ + elem, Sol = mem.elem, mem.sol + ndim = elem.ndim + + # Determine dimension order + dim_range = range(ndim-1, -1, -1) if reverse else range(ndim) + + for split in dim_range: + lmbda = time_step / elem.dxyz[split] + + # Handle solution flipping based on sweep direction + if reverse: + if elem.iisc[split] > 1: + explicit_step_and_flux(mem, ud, lmbda, split, diagnostics) + Sol.flip_backward() + else: + Sol.flip_forward() + if elem.iisc[split] > 1: + explicit_step_and_flux(mem, ud, lmbda, split, diagnostics) diff --git a/src/pybella/flow_solver/physics/gas_dynamics/numerical_flux.py b/src/pybella/flow_solver/physics/gas_dynamics/numerical_flux.py index 3d5098a6..80b31e04 100644 --- a/src/pybella/flow_solver/physics/gas_dynamics/numerical_flux.py +++ b/src/pybella/flow_solver/physics/gas_dynamics/numerical_flux.py @@ -1,134 +1,106 @@ # -*- coding: utf-8 -*- import numpy as np -import scipy as sp +from numba import njit +from ....utils.operators import get_flux_convolution_kernels, apply_directional_convolution +from ....utils.slices import get_inner_slice, get_interface_indices, get_last_dim_inner_slice -def recompute_advective_fluxes(flux, Sol, *args, **kwargs): - """ - Recompute the advective fluxes at the cell interfaces, i.e. the faces. This function updates the `flux` container in-place. - +def recompute_advective_fluxes(mem, **kwargs): + """Recompute the advective fluxes at the cell interfaces. + Parameters ---------- - flux : :py:class:`management.variable.States` - Data container for the fluxes at the cell interfaces. - Sol : :py:class:`management.variable.States` - Data container for the Solution. - - Attention - --------- - This function is a mess and requires cleaning up. - + mem : object + Memory object containing sol and flux attributes + **kwargs + Optional pre-computed velocity components ('u', 'v', 'w') """ - ndim = Sol.rho.ndim - inner_idx = tuple([slice(1, -1)] * ndim) - - if ndim == 2: - kernel_u = np.array([[0.5, 1.0, 0.5], [0.5, 1.0, 0.5]]) - kernel_v = kernel_u.T - elif ndim == 3: - kernel_u = np.array( - [[[1, 2, 1], [2, 4, 2], [1, 2, 1]], [[1, 2, 1], [2, 4, 2], [1, 2, 1]]] - ) - kernel_v = np.swapaxes(kernel_u, 1, 0) - kernel_w = np.swapaxes(kernel_u, 2, 0) - - rhoYw = Sol.rhoY * Sol.rhow / Sol.rho - flux[2].rhoY[inner_idx] = ( - sp.signal.fftconvolve(rhoYw, kernel_w, mode="valid") / kernel_w.sum() + ndim = mem.sol.rho.ndim + inner_idx = get_inner_slice(ndim) + kernels = get_flux_convolution_kernels(ndim) + + # Define the component order and corresponding flux indices + components = ['u', 'v'] if ndim == 2 else ['u', 'v', 'w'] + rho_components = ['rhou', 'rhov'] if ndim == 2 else ['rhou', 'rhov', 'rhow'] + + for i, (comp, rho_comp) in enumerate(zip(components, rho_components)): + # Use provided velocity or compute from momentum + if comp in kwargs: + rhoY_vel = kwargs[comp] + else: + momentum = getattr(mem.sol, rho_comp) + rhoY_vel = mem.sol.rhoY * momentum / mem.sol.rho + + # Apply directional convolution + mem.flux[i].rhoY[inner_idx] = apply_directional_convolution( + rhoY_vel, kernels[comp], comp, ndim ) - else: - assert 0, "Unsupported dimension in recompute_advective_flux" - rhoYu = kwargs.get("u", Sol.rhoY * Sol.rhou / Sol.rho) - - flux[0].rhoY[inner_idx] = np.moveaxis( - sp.signal.fftconvolve(rhoYu, kernel_u, mode="valid") / kernel_u.sum(), 0, -1 +@njit(cache=True) +def _compute_flux_component(flux_values, rhoY_values, left_weight, right_weight, + left_val, right_val, remove_cols_idx): + """Numba-optimised flux component computation.""" + flux_values[remove_cols_idx] = rhoY_values[remove_cols_idx] * ( + left_weight * left_val + right_weight * right_val ) - rhoYv = kwargs.get("v", Sol.rhoY * Sol.rhov / Sol.rho) - if ndim == 2: - flux[1].rhoY[inner_idx] = ( - sp.signal.fftconvolve(rhoYv, kernel_v, mode="valid") / kernel_v.sum() - ) - elif ndim == 3: - flux[1].rhoY[inner_idx] = np.moveaxis( - sp.signal.fftconvolve(rhoYv, kernel_v, mode="valid") / kernel_v.sum(), -1, 0 - ) - # flux[1].rhoY[...,-1] = 0. - - -def hll_solver(flux, Lefts, Rights, Sol, lmbda, ud, th): +def hll_solver(mem, flux, Lefts, Rights): """ HLL solver for the Riemann problem. Chooses the advected quantities from `Lefts` or `Rights` based on the direction given by `flux`. - Parameters - ---------- - flux : :py:class:`management.variable.States` - Data container for fluxes. - Lefts : :py:class:`management.variable.States` - Container for the quantities on the left of the cell interfaces. - Rights : :py:class:`management.variable.States` - Container for the quantities on the right of the cell interfaces. - Sol : :py:class:`management.variable.Vars` - Solution data container. - lmbda : float - :math:`\\frac{dt}{dx}`, where :math:`dx` is the grid-size in the direction of the substep. - ud : :py:class:`inputs.user_data.UserDataInit` - Class container for the initial condition. - th : :py:class:`physics.gas_dynamics.thermodynamic.init` - Class container for the thermodynamical constants. - Returns ------- :py:class:`management.variable.States` `flux` data container with the solution of the Riemann problem. + """ - # flux: index 1 to end = Left[inner_idx]: index 0 to -1 = Right[inner_idx]: index 1 to end - - ndim = Sol.rho.ndim - left_idx, right_idx, remove_cols_idx = ( - [slice(None)] * ndim, - [slice(None)] * ndim, - [slice(None)] * ndim, - ) + ndim = mem.sol.rho.ndim + left_idx, right_idx, _ = get_interface_indices(ndim) + remove_cols_idx = get_last_dim_inner_slice(ndim) - remove_cols_idx[-1] = slice(1, -1) - left_idx[-1] = slice(0, -1) - right_idx[-1] = slice(1, None) - - left_idx, right_idx, remove_cols_idx = ( - tuple(left_idx), - tuple(right_idx), - tuple(remove_cols_idx), - ) - - Lefts.primitives(th) - Rights.primitives(th) + # Compute primitive variables + Lefts.primitives(mem.th) + Rights.primitives(mem.th) + # Compute upwind weights upwind = 0.5 * (1.0 + np.sign(flux.rhoY)) upl = upwind[right_idx] upr = 1.0 - upwind[left_idx] + + # Pre-compute common weights + left_weight = upl[left_idx] / Lefts.Y[left_idx] + right_weight = upr[right_idx] / Rights.Y[right_idx] + + # Define flux components to compute + flux_components = [ + ('rhou', 'u'), + ('rho', None, 1.0), # state_value=1.0 + ('rhov', 'v'), + ('rhow', 'w'), + ('rhoX', 'X') + ] + + # Compute all flux components + for component in flux_components: + flux_attr = component[0] + state_attr = component[1] if len(component) > 1 else None + state_value = component[2] if len(component) > 2 else 1.0 + + # Get state values + if state_attr is not None: + left_val = getattr(Lefts, state_attr)[left_idx] + right_val = getattr(Rights, state_attr)[right_idx] + else: + left_val = right_val = state_value + + _compute_flux_component( + getattr(flux, flux_attr), + flux.rhoY, + left_weight, + right_weight, + left_val, + right_val, + remove_cols_idx + ) - flux.rhou[remove_cols_idx] = flux.rhoY[remove_cols_idx] * ( - upl[left_idx] / Lefts.Y[left_idx] * Lefts.u[left_idx] - + upr[right_idx] / Rights.Y[right_idx] * Rights.u[right_idx] - ) - flux.rho[remove_cols_idx] = flux.rhoY[remove_cols_idx] * ( - upl[left_idx] / Lefts.Y[left_idx] * 1.0 - + upr[right_idx] / Rights.Y[right_idx] * 1.0 - ) - - flux.rhov[remove_cols_idx] = flux.rhoY[remove_cols_idx] * ( - upl[left_idx] / Lefts.Y[left_idx] * Lefts.v[left_idx] - + upr[right_idx] / Rights.Y[right_idx] * Rights.v[right_idx] - ) - flux.rhow[remove_cols_idx] = flux.rhoY[remove_cols_idx] * ( - upl[left_idx] / Lefts.Y[left_idx] * Lefts.w[left_idx] - + upr[right_idx] / Rights.Y[right_idx] * Rights.w[right_idx] - ) - flux.rhoX[remove_cols_idx] = flux.rhoY[remove_cols_idx] * ( - upl[left_idx] / Lefts.Y[left_idx] * Lefts.X[left_idx] - + upr[right_idx] / Rights.Y[right_idx] * Rights.X[right_idx] - ) - - return flux + return flux \ No newline at end of file diff --git a/src/pybella/flow_solver/physics/gas_dynamics/recovery.py b/src/pybella/flow_solver/physics/gas_dynamics/recovery.py index 3bfb9030..483ad62e 100644 --- a/src/pybella/flow_solver/physics/gas_dynamics/recovery.py +++ b/src/pybella/flow_solver/physics/gas_dynamics/recovery.py @@ -1,222 +1,109 @@ import numpy as np +from numba import njit -from ...utils import options as opts, variable as var +from ....utils import options as opts +from ....utils.slices import get_neighbor_indices, get_interface_indices +from ...utils import variable as var -def do(Sol, flux, lmbda, ud, th, elem, split_step, tag): +def do(mem, ud, lmbda, split_step, tag=None): """ Reconstruct the limited slopes at the cell interfaces. - Parameters - ---------- - Sol : :py:class:`management.variable.Vars` - Solution data container on cell centers. - flux : :py:class:`management.variable.States` - Flux data container on cell interfaces. - lmbda : float - :math:`\\frac{dt}{dx}`, where :math:`dx` is the grid-size in the direction of the substep. - ud : :py:class:`inputs.user_data.UserDataInit` - Class container for the initial condition. - th : :py:class:`physics.gas_dynamics.thermodynamic.init` - Class container for the thermodynamical constants. - elem : :py:class:`discretization.kgrid.ElemSpaceDiscr` - Class container for the cell-grid. - split_step : int - Tracks the substep in the Strang-splitting. - tag : `None` or `rk` - Default is `None` which uses a second-order Strang-splitting. `rk` toggles a first-order Runge-Kutta update for the advection scheme. - - Returns - ------- - :py:class:`management.variable.States`, :py:class:`management.variable.States` - Lefts, Rights are containers for the advected quantities at to the left and the right of the cell interfaces. - """ - gamm = th.gamm + flux = mem.flux[split_step] + gamm = mem.th.gamm order_two = 1 # always 1 - Sol.primitives(th) + mem.sol.primitives(mem.th) if tag == "rk": lmbda = 0.0 - ndim = elem.ndim - lefts_idx, rights_idx, inner_idx = ( - [ - slice( - None, - ) - ] - * ndim, - [ - slice( - None, - ) - ] - * ndim, - [slice(1, -1)] * ndim, - ) - lefts_idx[-1] = slice(0, -1) - rights_idx[-1] = slice(1, None) - lefts_idx, rights_idx, inner_idx = ( - tuple(lefts_idx), - tuple(rights_idx), - tuple(inner_idx), - ) + # TBD: Consider moving these to the cache + lefts_idx, rights_idx, face_inner_idx = get_interface_indices(mem.elem.ndim) # inner_idx here are where the interface fluxes are calculated with non-zero values. - face_inner_idx = inner_idx - u = np.zeros_like(Sol.rhoY) - u[inner_idx] = ( + # TBD: Move to the cache + u = np.zeros_like(mem.sol.rhoY) + u[face_inner_idx] = ( 0.5 * (flux.rhoY[face_inner_idx][lefts_idx] + flux.rhoY[face_inner_idx][rights_idx]) - / Sol.rhoY[inner_idx] + / mem.sol.rhoY[face_inner_idx] ) - shape = Sol.u.shape - Diffs = var.States(shape, ud) - Ampls = var.Characters(shape) - Lefts = var.States(shape, ud) - Rights = var.States(shape, ud) - - Diffs.u[..., :-1] = Sol.u[rights_idx] - Sol.u[lefts_idx] - Diffs.v[..., :-1] = Sol.v[rights_idx] - Sol.v[lefts_idx] - Diffs.w[..., :-1] = Sol.w[rights_idx] - Sol.w[lefts_idx] - Diffs.X[..., :-1] = Sol.X[rights_idx] - Sol.X[lefts_idx] - Diffs.Y[..., :-1] = 1.0 / Sol.Y[rights_idx] - 1.0 / Sol.Y[lefts_idx] - - Slopes = slopes(Sol, Diffs, ud, elem) - - Ampls.u[...] = 0.5 * Slopes.u * (1.0 - lmbda * u) - Ampls.v[...] = 0.5 * Slopes.v * (1.0 - lmbda * u) - Ampls.w[...] = 0.5 * Slopes.w * (1.0 - lmbda * u) - Ampls.X[...] = 0.5 * Slopes.X * (1.0 - lmbda * u) - Ampls.Y[...] = 0.5 * Slopes.Y * (1.0 - lmbda * u) - - Lefts.u[...] = Sol.u + order_two * Ampls.u - Lefts.v[...] = Sol.v + order_two * Ampls.v - Lefts.w[...] = Sol.w + order_two * Ampls.w - Lefts.X[...] = Sol.X + order_two * Ampls.X - Lefts.Y[...] = 1.0 / (1.0 / Sol.Y + order_two * Ampls.Y) - - Ampls.u[...] = -0.5 * Slopes.u * (1.0 + lmbda * u) - Ampls.v[...] = -0.5 * Slopes.v * (1.0 + lmbda * u) - Ampls.w[...] = -0.5 * Slopes.w * (1.0 + lmbda * u) - Ampls.X[...] = -0.5 * Slopes.X * (1.0 + lmbda * u) - Ampls.Y[...] = -0.5 * Slopes.Y * (1.0 + lmbda * u) - - Rights.u[...] = Sol.u + order_two * Ampls.u - Rights.v[...] = Sol.v + order_two * Ampls.v - Rights.w[...] = Sol.w + order_two * Ampls.w - Rights.X[...] = Sol.X + order_two * Ampls.X - Rights.Y[...] = 1.0 / (1.0 / Sol.Y + order_two * Ampls.Y) - - vel = [Sol.u, Sol.v, Sol.w] - - # Lefts.rhoY[lefts_idx] = Rights.rhoY[rights_idx] = 0.5 * (Sol.rhoY[lefts_idx] + Sol.rhoY[rights_idx]) \ - # - order_two * 0.5 * lmbda * (Sol.u[rights_idx] * Sol.rhoY[rights_idx] - Sol.u[lefts_idx] * Sol.rhoY[lefts_idx]) - Lefts.rhoY[lefts_idx] = Rights.rhoY[rights_idx] = 0.5 * ( - Sol.rhoY[lefts_idx] + Sol.rhoY[rights_idx] - ) - order_two * 0.5 * lmbda * ( - vel[split_step][rights_idx] * Sol.rhoY[rights_idx] - - vel[split_step][lefts_idx] * Sol.rhoY[lefts_idx] + shape = mem.sol.u.shape + + # Get cached objects + cache = mem.cache.get_recovery_objects(shape, ud) + Diffs, Ampls, Lefts, Rights, Slopes = cache['Diffs'], cache['Ampls'], cache['Lefts'], cache['Rights'], cache['Slopes'] + + # Compute differences + _compute_differences(mem.sol, rights_idx, lefts_idx, Diffs) + + # Compute slopes + slopes_obj = _slopes(Diffs, Slopes, ud, mem.elem) + + # Compute left-side amplitudes and values + _compute_amplitudes(slopes_obj, lmbda, u, Ampls, sign_factor=1.0, lambda_factor=-1.0) + _compute_reconstructed_values(mem.sol, Ampls, order_two, Lefts) + + # Compute right-side amplitudes and values + _compute_amplitudes(slopes_obj, lmbda, u, Ampls, sign_factor=-1.0, lambda_factor=1.0) + _compute_reconstructed_values(mem.sol, Ampls, order_two,Rights) + + # Return velocity components + vel = [mem.sol.u, mem.sol.v, mem.sol.w] + + # Compute rhoY reconstruction + reconstructed_rhoy = _compute_rhoy_reconstruction( + mem, Lefts, Rights, lefts_idx, rights_idx, vel, split_step, order_two, lmbda + ) + + # Compute pressure reconstruction + _compute_pressure_reconstruction( + Lefts, Rights, lefts_idx, rights_idx, reconstructed_rhoy, gamm ) - Lefts.p0[lefts_idx] = Rights.p0[rights_idx] = Lefts.rhoY[lefts_idx] ** gamm - - get_conservatives(Rights, ud, th) - get_conservatives(Lefts, ud, th) + _get_conservatives(Rights) + _get_conservatives(Lefts) - # return Lefts, Rights, u, Diffs, Ampls, Slopes return Lefts, Rights - -def slopes(Sol, Diffs, ud, elem): - """ - Reconstruct the piecewise linear slopes in the cells from the piecewise constants in the cell and its neighbours. - - Parameters - ---------- - Sol : :py:class:`management.variable.Vars` - Solution data container - Diffs : :py:class:`management.variable.States` - Data container for the difference in the quantities between adjacent cells, (right - left). - ud : :py:class:`inputs.user_data.UserDataInit` - Data container for the initial conditions - elem : :py:class:`discretization.kgrid.ElemSpaceDiscr` - Data container for the cell grid. - - Returns - ------- - :py:class:`management.variable.Characters` - Reconstructed piecewise linear slopes in the cell. - """ - limiter_type_velocity = ud.limiter_type_velocity - limiter_type_scalar = ud.limiter_type_scalars - - ndim = elem.ndim - lefts_idx, rights_idx = [ - slice( - None, - ) - ] * ndim, [ - slice( - None, - ) - ] * ndim - lefts_idx[-1] = slice(0, -1) - rights_idx[-1] = slice(1, None) - lefts_idx, rights_idx = tuple(lefts_idx), tuple(rights_idx) - - # amplitudes of the state differences: - # first lefts_idx removes the zero at the end - # since differences always result in len-1 - # and the second indexing selects lefts and rights - aul = Diffs.u[lefts_idx][lefts_idx] - avl = Diffs.v[lefts_idx][lefts_idx] - awl = Diffs.w[lefts_idx][lefts_idx] - aXl = Diffs.X[lefts_idx][lefts_idx] - aYl = Diffs.Y[lefts_idx][lefts_idx] - - aur = Diffs.u[lefts_idx][rights_idx] - avr = Diffs.v[lefts_idx][rights_idx] - awr = Diffs.w[lefts_idx][rights_idx] - aXr = Diffs.X[lefts_idx][rights_idx] - aYr = Diffs.Y[lefts_idx][rights_idx] - - Slopes = var.Characters(Diffs.u.shape) - - Slopes.u[..., 1:-1] = limiters(limiter_type_velocity, aul, aur) - Slopes.v[..., 1:-1] = limiters(limiter_type_velocity, avl, avr) - Slopes.w[..., 1:-1] = limiters(limiter_type_velocity, awl, awr) - Slopes.X[..., 1:-1] = limiters(limiter_type_scalar, aXl, aXr) - Slopes.Y[..., 1:-1] = limiters(limiter_type_scalar, aYl, aYr) - +def _slopes(Diffs, Slopes, ud, elem): + """Reconstruct piecewise linear slopes in cells.""" + # Configuration + variable_config = { + 'u': ud.limiter_type_velocity, + 'v': ud.limiter_type_velocity, + 'w': ud.limiter_type_velocity, + 'X': ud.limiter_type_scalars, + 'Y': ud.limiter_type_scalars + } + + # TBD: Consider moving indices to cache + lefts_idx, rights_idx = get_neighbor_indices(elem.ndim) + + # Process each variable + for var_name, limiter_type in variable_config.items(): + diff_data = getattr(Diffs, var_name) + + # Extract amplitudes + al = diff_data[lefts_idx][lefts_idx] + ar = diff_data[lefts_idx][rights_idx] + + # Calculate and assign slopes + slope_data = _limiters(limiter_type, al, ar) + getattr(Slopes, var_name)[..., 1:-1] = slope_data + return Slopes - -def limiters(limiter_type, al, ar): +@njit +def _limiters(limiter_type, al, ar): """ Applies the limiter type specified in the initial conditions to recovery the slope. - Parameters - ---------- - limiter_type : :py:class:`management.enumerator.LimiterType` - LimiterType list - al : :py:class:`management.variable.States` - Left indices of the `Diffs` array for the respective quantities. - ar : :py:class:`management.variable.States` - Right indices of the `Diffs` array for the respective quantities. - - Returns - ------- - :py:class:`management.variable.States` - The reconstructed slope in the cell - - Attention - --------- - For now, only the limiter type `NONE` is supported. This takes $\\frac{(al + ar)}{2}$. """ # write switch for limiter types # for now, just use LimiterType == None @@ -224,7 +111,7 @@ def limiters(limiter_type, al, ar): return 0.5 * (al + ar) -def get_conservatives(U, ud, th): +def _get_conservatives(U): """ Get advected (conservative) quantities at the left and right of the cell interfaces. @@ -232,10 +119,6 @@ def get_conservatives(U, ud, th): ---------- U : :py:class:`management.variable.States` `Lefts` and `Rights` corresponding to the values at the cell interfaces. - ud : :py:class:`inputs.user_data.UserDataInit` - Data container for the initial conditions - th : :py:class:`physics.gas_dynamics.thermodynamic.init` - Class container for the thermodynamical constants. """ U.rho = U.rhoY / U.Y U.rhou = U.u * U.rho @@ -244,5 +127,61 @@ def get_conservatives(U, ud, th): U.rhoY = U.Y * U.rho U.rhoX = U.X * U.rho - sgn = np.sign(U.rhoY) - p = sgn * np.abs(U.rhoY) ** th.gamminv + +def _compute_differences(sol, rights_idx, lefts_idx, diffs): + """Compute differences between right and left indices for all fields.""" + fields = ['u', 'v', 'w', 'X'] + for field in fields: + getattr(diffs, field)[..., :-1] = getattr(sol, field)[rights_idx] - getattr(sol, field)[lefts_idx] + + # Y field has special handling (reciprocal differences) + diffs.Y[..., :-1] = 1.0 / sol.Y[rights_idx] - 1.0 / sol.Y[lefts_idx] + + +def _compute_amplitudes(slopes, lmbda, u, ampls, sign_factor=1.0, lambda_factor=1.0): + """Compute amplitudes for all fields with given sign and lambda factors.""" + fields = ['u', 'v', 'w', 'X', 'Y'] + factor = sign_factor * 0.5 * (1.0 + lambda_factor * lmbda * u) + + for field in fields: + getattr(ampls, field)[...] = factor * getattr(slopes, field) + + +def _compute_reconstructed_values(sol, ampls, order_two, result): + """Compute reconstructed values for all fields.""" + fields = ['u', 'v', 'w', 'X'] + for field in fields: + getattr(result, field)[...] = getattr(sol, field) + order_two * getattr(ampls, field) + + # Y field has special handling (reciprocal computation) + result.Y[...] = 1.0 / (1.0 / sol.Y + order_two * ampls.Y) + +def _compute_rhoy_reconstruction(mem, lefts, rights, lefts_idx, rights_idx, + vel, split_step, order_two, lmbda): + """Compute rhoY reconstruction for both left and right sides.""" + # Use existing arrays for in-place operations + # First compute on lefts.rhoY, then copy to rights.rhoY + + # Start with average in lefts array + lefts.rhoY[lefts_idx] = 0.5 * (mem.sol.rhoY[lefts_idx] + mem.sol.rhoY[rights_idx]) + + # Subtract velocity correction in-place + lefts.rhoY[lefts_idx] -= order_two * 0.5 * lmbda * ( + vel[split_step][rights_idx] * mem.sol.rhoY[rights_idx] - + vel[split_step][lefts_idx] * mem.sol.rhoY[lefts_idx] + ) + + # Copy result to rights + rights.rhoY[rights_idx] = lefts.rhoY[lefts_idx] + + return lefts.rhoY[lefts_idx] + + +def _compute_pressure_reconstruction(lefts, rights, lefts_idx, rights_idx, + reconstructed_rhoy, gamm): + """Compute pressure reconstruction using power law.""" + reconstructed_p0 = reconstructed_rhoy ** gamm + lefts.p0[lefts_idx] = reconstructed_p0 + rights.p0[rights_idx] = reconstructed_p0 + + return reconstructed_p0 \ No newline at end of file diff --git a/src/pybella/flow_solver/physics/hydrostatics.py b/src/pybella/flow_solver/physics/hydrostatics.py index 64c720ba..6aa8e215 100644 --- a/src/pybella/flow_solver/physics/hydrostatics.py +++ b/src/pybella/flow_solver/physics/hydrostatics.py @@ -79,6 +79,7 @@ def column(HydroState, HydroState_n, Y, Y_n, elem, node, th, ud): HydroState_n.p0[xc_idx, igy + 1 :] = rhoY_hydro_n[:, igy:] ** th.gamm HydroState_n.p20[xc_idx, igy + 1 :] = pi_hydro_n[:, igy:] / ud.Msq + def integrated_state(mpv, elem, node, th, ud): """ Compute hydrostatic background state for atmospheric model. @@ -90,55 +91,52 @@ def integrated_state(mpv, elem, node, th, ud): gm1 = th.gm1 Gamma_inv = 1.0 / Gamma gm1_inv = 1.0 / gm1 - + # Grid parameters icy = elem.icy igy = elem.igy - + # Reference state at y=0 rhoY0 = 1.0 g = ud.gravity_strength[1] p0 = rhoY0**gamm pi0 = rhoY0**gm1 - + if g != 0.0: ########################### # Update cell hydrostates ########################### - + # Define midpoint quadrature along vertical (y-axis) - dys = np.hstack(( - np.ones(igy-1) * -elem.dy, - [-elem.dy/2], - [elem.dy/2], - np.ones(icy-3) * elem.dy - )) - + dys = np.hstack( + ( + np.ones(igy - 1) * -elem.dy, + [-elem.dy / 2], + [elem.dy / 2], + np.ones(icy - 3) * elem.dy, + ) + ) + # Cell centers and midpoints for integration y_ps = elem.y - y_ms = np.hstack(( - elem.y[1:igy], - node.y[igy], - node.y[igy], - elem.y[igy:-1] - )) - + y_ms = np.hstack((elem.y[1:igy], node.y[igy], node.y[igy], elem.y[igy:-1])) + # Get inverse stratification at each point S_ps = 1.0 / ud.stratification(y_ps) S_ms = 1.0 / ud.stratification(y_ms) - + # Trapezoidal integration over inverse stratification S_integral_p = 0.5 * dys * (S_ms + S_ps) - + # Cumulative integration (split at boundary igy) S_integral_p[:igy] = np.cumsum(S_integral_p[:igy][::-1])[::-1] S_integral_p[igy:] = np.cumsum(S_integral_p[igy:]) - + # Calculate hydrostatic fields pi_hydro = pi0 - Gamma * g * S_integral_p p_hydro = pi_hydro**Gamma_inv rhoY_hydro = pi_hydro**gm1_inv - + # Update cell solutions mpv.HydroState.rhoY0[:] = rhoY_hydro mpv.HydroState.rho0[:] = rhoY_hydro * S_ps @@ -147,11 +145,11 @@ def integrated_state(mpv, elem, node, th, ud): mpv.HydroState.S0[:] = S_ps mpv.HydroState.S10[:] = 0.0 mpv.HydroState.Y0[:] = 1.0 / S_ps - + ############################ # Update node hydrostates ############################ - + # Bottom reference node (y=0) mpv.HydroState_n.Y0[igy] = ud.stratification(0.0) mpv.HydroState_n.rhoY0[igy] = rhoY0 @@ -159,43 +157,49 @@ def integrated_state(mpv, elem, node, th, ud): mpv.HydroState_n.S0[igy] = 1.0 / mpv.HydroState_n.Y0[igy] mpv.HydroState_n.p0[igy] = p0 mpv.HydroState_n.p20[igy] = pi0 / ud.Msq - + # Ghost cells below bottom (negative heights) Sn_integral_p = np.zeros(igy) yn_p = node.y[:igy] - node.dy - yn_m = node.y[1:igy+1] - node.dy - - Sn_integral_p[:] = -node.dy * 1.0 / ud.stratification(0.5*(yn_p + yn_m)) + yn_m = node.y[1 : igy + 1] - node.dy + + Sn_integral_p[:] = -node.dy * 1.0 / ud.stratification(0.5 * (yn_p + yn_m)) Sn_integral_p = np.cumsum(Sn_integral_p[:igy][::-1])[::-1] - + # Bulk domain above reference level - yn_p = node.y[igy+1:] + yn_p = node.y[igy + 1 :] yn_m = np.zeros_like(yn_p) yn_m[1:] = yn_p[:-1] - + Sn_p = 1.0 / ud.stratification(0.5 * (yn_p + yn_m)) Sn_integral_p = np.hstack((Sn_integral_p, np.cumsum(elem.dy * Sn_p))) - + # Calculate nodal hydrostatic fields pi_hydro_n = pi0 - Gamma * g * Sn_integral_p rhoY_hydro_n = pi_hydro_n**gm1_inv - + # Update node solutions - below reference mpv.HydroState_n.rhoY0[:igy] = rhoY_hydro_n[:igy] - mpv.HydroState_n.Y0[:igy+1] = ud.stratification(0.5 * (y_ps[:igy+1] + y_ps[:igy+1] - elem.dy)) + mpv.HydroState_n.Y0[: igy + 1] = ud.stratification( + 0.5 * (y_ps[: igy + 1] + y_ps[: igy + 1] - elem.dy) + ) mpv.HydroState_n.rho0[:igy] = rhoY_hydro_n[:igy] / mpv.HydroState_n.Y0[:igy] mpv.HydroState_n.S0[:igy] = 1.0 / mpv.HydroState_n.Y0[:igy] - mpv.HydroState_n.p0[:igy] = rhoY_hydro_n[:igy]**th.gamm + mpv.HydroState_n.p0[:igy] = rhoY_hydro_n[:igy] ** th.gamm mpv.HydroState_n.p20[:igy] = pi_hydro_n[:igy] / ud.Msq - + # Update node solutions - above reference - mpv.HydroState_n.rhoY0[igy+1:] = rhoY_hydro_n[igy:] - mpv.HydroState_n.Y0[igy+1:] = ud.stratification(0.5 * (y_ps[igy:] + y_ps[igy:] + elem.dy)) - mpv.HydroState_n.rho0[igy+1:] = rhoY_hydro_n[igy:] / mpv.HydroState_n.Y0[igy+1:] - mpv.HydroState_n.S0[igy+1:] = 1.0 / mpv.HydroState_n.Y0[igy+1:] - mpv.HydroState_n.p0[igy+1:] = rhoY_hydro_n[igy:]**th.gamm - mpv.HydroState_n.p20[igy+1:] = pi_hydro_n[igy:] / ud.Msq - + mpv.HydroState_n.rhoY0[igy + 1 :] = rhoY_hydro_n[igy:] + mpv.HydroState_n.Y0[igy + 1 :] = ud.stratification( + 0.5 * (y_ps[igy:] + y_ps[igy:] + elem.dy) + ) + mpv.HydroState_n.rho0[igy + 1 :] = ( + rhoY_hydro_n[igy:] / mpv.HydroState_n.Y0[igy + 1 :] + ) + mpv.HydroState_n.S0[igy + 1 :] = 1.0 / mpv.HydroState_n.Y0[igy + 1 :] + mpv.HydroState_n.p0[igy + 1 :] = rhoY_hydro_n[igy:] ** th.gamm + mpv.HydroState_n.p20[igy + 1 :] = pi_hydro_n[igy:] / ud.Msq + else: # No gravity case - uniform atmosphere mpv.HydroState.p20[:] = 1.0 @@ -205,7 +209,7 @@ def integrated_state(mpv, elem, node, th, ud): mpv.HydroState.Y0[:] = 1.0 mpv.HydroState.S0[:] = 1.0 mpv.HydroState.S10[:] = 0.0 - + mpv.HydroState_n.p20[:] = 1.0 mpv.HydroState_n.p0[:] = 1.0 mpv.HydroState_n.rho0[:] = 1.0 @@ -224,7 +228,7 @@ def analytical_state(mpv, elem, node, th, ud): pi_nm = np.exp(-(node.y - 0.5 * dy) / Hex) pi_n = np.exp(-(node.y) / Hex) - Y_n = - Gamma * g * dy / (pi_np - pi_nm) + Y_n = -Gamma * g * dy / (pi_np - pi_nm) P_n = pi_n**th.gm1inv p_n = pi_n**th.Gammainv rho_n = P_n / Y_n @@ -238,9 +242,9 @@ def analytical_state(mpv, elem, node, th, ud): pi_cp = np.exp(-(elem.y + 0.5 * dy) / Hex) pi_cm = np.exp(-(elem.y - 0.5 * dy) / Hex) - pi_c = np.exp(-(elem.y) / Hex) + pi_c = np.exp(-(elem.y) / Hex) - Y_c = - Gamma * g * dy / (pi_cp - pi_cm) + Y_c = -Gamma * g * dy / (pi_cp - pi_cm) P_c = pi_c**th.gm1inv p_c = pi_c**th.Gammainv rho_c = P_c / Y_c @@ -253,7 +257,6 @@ def analytical_state(mpv, elem, node, th, ud): mpv.HydroState.S0[...] = 1.0 / Y_c - def initial_pressure(Sol, mpv, elem, node, ud, th): Gammainv = th.Gammainv igy = node.igy diff --git a/src/pybella/flow_solver/physics/low_mach/laplacian.py b/src/pybella/flow_solver/physics/low_mach/laplacian.py index cdf982e9..dcc1821b 100644 --- a/src/pybella/flow_solver/physics/low_mach/laplacian.py +++ b/src/pybella/flow_solver/physics/low_mach/laplacian.py @@ -2,7 +2,7 @@ import scipy as sp import numba as nb -from ...utils import options as opts +from ....utils import options as opts def stencil_9pt(elem, node, mpv, Sol, ud, diag_inv, dt, coriolis_params): diff --git a/src/pybella/flow_solver/physics/low_mach/second_projection.py b/src/pybella/flow_solver/physics/low_mach/second_projection.py index 3f208c94..fc9e4523 100644 --- a/src/pybella/flow_solver/physics/low_mach/second_projection.py +++ b/src/pybella/flow_solver/physics/low_mach/second_projection.py @@ -1,10 +1,13 @@ import itertools as it -import logging import numpy as np import scipy as sp +from numba import njit -from ...utils import options as opts, boundary as bdry +from ....utils import options as opts +from ....utils import operators + +from ...utils import boundary as bdry from . import laplacian as lm_lp @@ -22,97 +25,96 @@ def __call__(self, rk=None): self.rk = rk -def euler_forward_non_advective( - Sol, mpv, elem, node, dt, ud, th, writer=None, label=None, debug=False -): +def euler_forward_non_advective(mem, ud, dt, writer=None, label=None, debug=False): + # Unpack frequently used variables + th, sol, mpv, node, elem = mem.th, mem.sol, mem.mpv, mem.node, mem.elem + ndim = elem.ndim + nonhydro = ud.nonhydrostasy - g = ud.gravity_strength[1] - Msq = ud.Msq + g, Msq = ud.gravity_strength[1], ud.Msq Ginv = th.Gammainv - corr_h1 = ud.coriolis_strength[0] - corr_v = ud.coriolis_strength[1] - corr_h2 = ud.coriolis_strength[2] - u0 = ud.u_wind_speed - v0 = ud.v_wind_speed - w0 = ud.w_wind_speed - - p2n = np.copy(mpv.p2_nodes) + corr_h1, corr_v, corr_h2 = ud.coriolis_strength + u0, v0, w0 = ud.u_wind_speed, ud.v_wind_speed, ud.w_wind_speed + + # Reusable derived quantities + rho, rhoY, rhoX = sol.rho, sol.rhoY, sol.rhoX + rhou, rhov, rhow = sol.rhou, sol.rhov, sol.rhow + + # Pressure and derivatives + p2n = mpv.p2_nodes dp2n = np.zeros_like(p2n) - ndim = elem.ndim S0c = mpv.HydroState.get_S0c(elem) dSdy = mpv.HydroState_n.get_dSdy(elem, node) - mpv.rhs[...] = divergence_nodes(mpv.rhs, elem, node, Sol, ud) + # Compute divergence + mpv.rhs[...] = divergence_nodes(mpv.rhs, elem, node, sol, ud) if not hasattr(ud, "ATMOSPHERIC_EXTENSION"): - scale_wall_node_values(mpv.rhs, node, ud, 2.0) - div = mpv.rhs + bdry.scale_wall_node_values(mpv.rhs, node, ud, 2.0) if debug: - writer.populate(str(label), "rhs", div) - - rhoY = Sol.rhoY ** (th.gamm - 2.0) - dpidP_kernel = np.ones([2] * ndim) - dpidP = ( - (th.gm1 / ud.Msq) - * sp.signal.fftconvolve(rhoY, dpidP_kernel, mode="valid") - / dpidP_kernel.sum() + writer.populate(str(label), "rhs", mpv.rhs) + + # Compute compressibility kernel + kernel = operators.get_averaging_kernel(ndim, width=2) + dpidP = (th.gm1 / Msq) * operators.apply_convolution_kernel( + rhoY ** (th.gamm - 2.0), + kernel=kernel, + normalize=True, + use_numba=True ) - rhoYovG = Ginv * Sol.rhoY - dbuoy = Sol.rhoY * (Sol.rhoX / Sol.rho) - dpdx, dpdy, dpdz = grad_nodes(p2n, elem.ndim, node.dxyz) + rhoYovG = Ginv * rhoY + dbuoy = rhoY * (rhoX / rho) - drhou = Sol.rhou - u0 * Sol.rho - drhov = Sol.rhov - v0 * Sol.rho - drhow = Sol.rhow - w0 * Sol.rho - v = Sol.rhov / Sol.rho + # Pressure gradients + dpdx, dpdy, dpdz = operators.compute_gradient_nodes(p2n, ndim, node.dxyz) - Sol.rhou = Sol.rhou - dt * (rhoYovG * dpdx - corr_h2 * drhov + corr_v * drhow) + # Wind perturbations + drhou = rhou - u0 * rho + drhov = rhov - v0 * rho + drhow = rhow - w0 * rho + v = rhov / rho - Sol.rhov = Sol.rhov - dt * ( + # Momentum update (u, v, w) + rhou -= dt * (rhoYovG * dpdx - corr_h2 * drhov + corr_v * drhow) + rhov -= dt * ( rhoYovG * dpdy + (g / Msq) * dbuoy * nonhydro - corr_h1 * drhow + corr_h2 * drhou ) * (1 - ud.is_ArakawaKonor) - Sol.rhow = Sol.rhow - dt * (rhoYovG * dpdz - corr_v * drhou + corr_h1 * drhov) * ( - ndim == 3 - ) - - Sol.rhoX = (Sol.rho * (Sol.rho / Sol.rhoY - S0c)) - dt * (v * dSdy) * Sol.rho + if ndim == 3: + rhow -= dt * (rhoYovG * dpdz - corr_v * drhou + corr_h1 * drhov) - dp2n[node.i1] -= dt * dpidP * div # [node.i1] + # Scalar update (rhoX) + sol.rhoX[...] = (rho * (rho / rhoY - S0c)) - dt * (v * dSdy) * rho - weight = ud.compressibility - mpv.p2_nodes += weight * dp2n + # Compressibility correction to p2 + dp2n[node.i1] -= dt * dpidP * mpv.rhs + mpv.p2_nodes += ud.compressibility * dp2n + # Boundary conditions bdry.set_ghostnodes_p2(mpv.p2_nodes, node, ud) - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) + bdry.set_explicit_boundary_data(sol, elem, ud, th, mpv) -def euler_backward_non_advective_expl_part(Sol, mpv, elem, dt, ud, th): +def euler_backward_non_advective_expl_part(mem, ud, dt): nonhydro = ud.nonhydrostasy g = ud.gravity_strength[1] Msq = ud.Msq - dbuoy = Sol.rhoY * (Sol.rhoX / Sol.rho) - Sol.rhov = (nonhydro * Sol.rhov) - dt * (g / Msq) * dbuoy - # Sol.rhov[np.where(Sol.rhov < 1e-15)] = 0.0 + dbuoy = mem.sol.rhoY * (mem.sol.rhoX / mem.sol.rho) + mem.sol.rhov = (nonhydro * mem.sol.rhov) - dt * (g / Msq) * dbuoy - Sol.mod_bg_wind(ud, -1.0) + mem.sol.mod_bg_wind(ud, -1.0) - multiply_inverse_coriolis(Sol, Sol, mpv, ud, elem, elem, dt) + multiply_inverse_coriolis(mem.sol, mem, ud, dt) - Sol.mod_bg_wind(ud, +1.0) - - bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - - -total_iter = 0 -total_calls = 0 + mem.sol.mod_bg_wind(ud, +1.0) + bdry.set_explicit_boundary_data(mem.sol, mem.elem, ud, mem.th, mem.mpv) def euler_backward_non_advective_impl_part( Sol, @@ -123,7 +125,7 @@ def euler_backward_non_advective_impl_part( th, t, dt, - alpha_diff, + mem, Sol0=None, writer=None, label=None, @@ -151,21 +153,21 @@ def euler_backward_non_advective_impl_part( bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt) - i0 = node.ndim * [(slice(0, -1))] + i0 = node.ndim * [slice(0, -1)] i0 = tuple(i0) if writer != None: writer.populate(str(label), "hcenter", mpv.wcenter) writer.populate(str(label), "wplusx", mpv.wplus[0]) writer.populate(str(label), "wplusy", mpv.wplus[1]) - writer.populate( - str(label), "wplusz", mpv.wplus[2] - ) if elem.ndim == 3 else writer.populate( - str(label), "wplusz", np.zeros_like(mpv.wplus[0]) + ( + writer.populate(str(label), "wplusz", mpv.wplus[2]) + if elem.ndim == 3 + else writer.populate(str(label), "wplusz", np.zeros_like(mpv.wplus[0])) ) bdry.set_ghostnodes_p2(mpv.p2_nodes, node, ud) - correction_nodes(Sol, elem, node, mpv, mpv.p2_nodes, dt, ud, th, 0) + correction_nodes(mem, ud, dt, mem.mpv.p2_nodes, 0) bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) rhs[...] = divergence_nodes(rhs, elem, node, Sol, ud) @@ -200,7 +202,7 @@ def euler_backward_non_advective_impl_part( if elem.ndim == 2: Vec = mpv coriolis_params = multiply_inverse_coriolis( - Vec, Sol, mpv, ud, elem, node, dt, attrs=("u", "v", "w"), get_coeffs=True + Vec, mem, ud, dt, attrs=("u", "v", "w"), get_coeffs=True ) # lap = stencil_9pt(elem,node,mpv,Sol,ud,diag_inv,dt,coriolis_params) # sh = (ud.inx)*(ud.iny) @@ -274,14 +276,14 @@ def euler_backward_non_advective_impl_part( # p2,info = bicgstab(lap,rhs.ravel(),x0=p2.ravel(),tol=1e-16,maxiter=6000,callback=counter) # print("Convergence info = %i, no. of iterations = %i" %(info,counter.niter)) - global total_calls, total_iter - total_iter += counter.niter - total_calls += 1 - logging.info(counter.niter) - logging.info( - "Total calls to BiCGStab routine = %i, total iterations = %i" - % (total_calls, total_iter) - ) + # global total_calls, total_iter + # total_iter += counter.niter + # total_calls += 1 + # logging.info(counter.niter) + # logging.info( + # "Total calls to BiCGStab routine = %i, total iterations = %i" + # % (total_calls, total_iter) + # ) p2_full = np.zeros(nc).squeeze() if elem.ndim == 2: @@ -314,38 +316,36 @@ def euler_backward_non_advective_impl_part( writer.populate(str(label), "p2_full", p2_full) bdry.set_ghostnodes_p2(p2_full, node, ud) - correction_nodes(Sol, elem, node, mpv, p2_full, dt, ud, th, 1) + correction_nodes(mem, ud, dt, p2_full, 1) mpv.p2_nodes[...] += p2_full bdry.set_ghostnodes_p2(mpv.p2_nodes, node, ud) bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) -def correction_nodes(Sol, elem, node, mpv, p, dt, ud, th, updt_chi): - ndim = node.ndim - Gammainv = th.Gammainv +def correction_nodes(mem, ud, dt, p, updt_chi): + ndim = mem.node.ndim + Gammainv = mem.th.Gammainv - dSdy = mpv.HydroState_n.get_dSdy(elem, node) + dSdy = mem.mpv.HydroState_n.get_dSdy(mem.elem, mem.node) - Dpx, Dpy, Dpz = grad_nodes(p, elem.ndim, node.dxyz) + Dpx, Dpy, Dpz = operators.compute_gradient_nodes(p, mem.elem.ndim, mem.node.dxyz) - thinv = Sol.rho / Sol.rhoY + thinv = mem.sol.rho / mem.sol.rhoY - Y = Sol.rhoY / Sol.rho - coeff = Gammainv * Sol.rhoY * Y + Y = mem.sol.rhoY / mem.sol.rho + coeff = Gammainv * mem.sol.rhoY * Y - mpv.u = -dt * coeff * Dpx - mpv.v = -dt * coeff * Dpy - mpv.w = -dt * coeff * Dpz + mem.mpv.u = -dt * coeff * Dpx + mem.mpv.v = -dt * coeff * Dpy + mem.mpv.w = -dt * coeff * Dpz - multiply_inverse_coriolis(mpv, Sol, mpv, ud, elem, elem, dt, attrs=["u", "v", "w"]) + multiply_inverse_coriolis(mem.mpv, mem, ud, dt, attrs=["u", "v", "w"]) - Sol.rhou += thinv * mpv.u - Sol.rhov += thinv * mpv.v - Sol.rhow += thinv * mpv.w if ndim == 3 else 0.0 - Sol.rhoX += -updt_chi * dt * dSdy * Sol.rhov - - # set_explicit_boundary_data(Sol, elem, ud, th, mpv) + mem.sol.rhou += thinv * mem.mpv.u + mem.sol.rhov += thinv * mem.mpv.v + mem.sol.rhow += thinv * mem.mpv.w if ndim == 3 else 0.0 + mem.sol.rhoX += -updt_chi * dt * dSdy * mem.sol.rhov assert True @@ -402,245 +402,7 @@ def operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt): assert True if not hasattr(ud, "ATMOSPHERIC_EXTENSION"): - scale_wall_node_values(mpv.wcenter, node, ud) - - -# def operator_coefficients_nodes(elem, node, Sol, mpv, ud, th, dt): -# g = ud.gravity_strength[1] -# Msq = ud.Msq -# Gammainv = th.Gammainv - -# ndim = node.ndim -# nonhydro = ud.nonhydrostasy -# dy = elem.dy - -# wh1, wv, wh2 = dt * ud.coriolis_strength - -# ccenter = - ud.Msq * th.gm1inv / (dt**2) -# cexp = 2.0 - th.gamm - -# igs = elem.igs - -# nindim = np.empty((ndim),dtype='object') -# innerdim = np.empty((ndim),dtype='object') -# innerdim1 = np.empty((ndim),dtype='object') -# eindim = np.empty((ndim),dtype='object') - -# for dim in range(ndim): -# is_periodic = ud.bdry_type[dim] == BdryType.PERIODIC -# nindim[dim] = slice(igs[dim]-is_periodic,-igs[dim]+is_periodic) -# innerdim[dim] = slice(igs[dim],-igs[dim]) -# eindim[dim] = slice(igs[dim]-is_periodic,-igs[dim]+is_periodic-1) - -# if dim == 1: -# y_idx = slice(igs[dim]-is_periodic,-igs[dim]+is_periodic-1) -# right_idx = None if -igs[dim]+is_periodic == 0 else -igs[dim]+is_periodic -# y_idx1 = slice(igs[dim]-is_periodic+1, right_idx) - -# innerdim1[dim] = slice(igs[dim]-1, (-igs[dim]+1)) - -# strat = (mpv.HydroState_n.S0[y_idx1] - mpv.HydroState_n.S0[y_idx]) / dy - -# nindim = tuple(nindim) -# eindim = tuple(eindim) -# innerdim = tuple(innerdim) -# innerdim1 = tuple(innerdim1) - -# for dim in range(0,elem.ndim,2): -# is_periodic = ud.bdry_type[dim] != BdryType.PERIODIC -# strat = np.expand_dims(strat, dim) -# strat = np.repeat(strat, elem.sc[dim]-int(2*is_periodic+igs[dim]), axis=dim) - -# Y = Sol.rhoY[nindim] / Sol.rho[nindim] -# coeff = Gammainv * Sol.rhoY[nindim] * Y - -# nu = np.zeros_like(mpv.wcenter) -# nu[eindim] = -dt**2 * (g / Msq) * strat * Y - -# setattr(mpv, 'nu_c', nu) -# nu = nu[eindim] - -# denom = 1.0 / (wh1**2 + wh2**2 + (nu + nonhydro) * (wv**2 + 1)) - -# fimp = denom -# gimp = denom - -# for dim in range(ndim): -# ## Assuming 2D vertical slice! -# if dim == 1: -# mpv.wplus[dim][eindim] = coeff #* gimp #* (wv**2 + 1.0) -# else: -# mpv.wplus[dim][eindim] = coeff #* fimp #* (wh1**2 + nu + nonhydro) - -# kernel = np.ones([2] * ndim) - -# mpv.wcenter[innerdim] = ccenter * signal.fftconvolve(Sol.rhoY[innerdim1]**cexp,kernel,mode='valid') / kernel.sum() - -# scale_wall_node_values(mpv.wcenter, node, ud) - - -def scale_wall_node_values(rhs, node, ud, factor=0.5): - # if factor < 1.0: - # if factor < 1.0: - # factor = 1.0 - # rhs[:,1] *= factor - # rhs[:,-2] *= factor - - # # rhs[:,0] *= 1.0# factor - # # rhs[:,-1] *= 1.0# factor - # rhs[:,0] = rhs[:,2] * factor - # rhs[:,-1] = rhs[:,-3] * factor - # else: - # factor = 1.0 - # rhs[:,:2] *= factor - # rhs[:,-2:] *= factor - - ndim = node.ndim - igs = node.igs - - wall_idx = np.empty((ndim), dtype=object) - for dim in range(ndim): - wall_idx[dim] = slice(igs[dim], -igs[dim]) - - for dim in range(ndim): - is_wall = ( - ud.bdry_type[dim] == opts.BdryType.WALL - or ud.bdry_type[dim] == opts.BdryType.RAYLEIGH - ) - if is_wall: - for direction in [-1, 1]: - wall_idx[dim] = (igs[dim] - 1) * direction - if direction == -1: - wall_idx[dim] -= 1 - wall_idx_tuple = tuple(wall_idx) - rhs[wall_idx_tuple] *= factor - - # is_rayleigh = ud.bdry_type[dim] == BdryType.RAYLEIGH - # if is_rayleigh: - # for direction in [1]: - # wall_idx[dim] = (igs[dim]-1) * direction - # # if direction == -1: - # # wall_idx[dim] -= 1 - # wall_idx_tuple = tuple(wall_idx) - # rhs[wall_idx_tuple] *= factor - - -def grad_nodes_fft(p2n, elem, node): - ndim = node.ndim - dx, dy, dz = node.dx, node.dy, node.dz - - kernels = [] - for dim in range(ndim): - kernel = np.ones([2] * ndim) - slc = [ - slice( - None, - ) - ] * ndim - slc[dim] = slice(0, 1) - kernel[tuple(slc)] *= -1.0 - kernels.append(kernel) - - dpdx = ( - -(0.5 ** (ndim - 1)) * sp.signal.fftconvolve(p2n, kernels[0], mode="valid") / dx - ) - dpdy = ( - -(0.5 ** (ndim - 1)) * sp.signal.fftconvolve(p2n, kernels[1], mode="valid") / dy - if elem.iicy > 1 - else 0.0 - ) - dpdz = ( - -(0.5 ** (ndim - 1)) * sp.signal.fftconvolve(p2n, kernels[2], mode="valid") / dz - if (ndim == 3) - else 0.0 - ) - - return dpdx, dpdy, dpdz - - -# @jit(nopython=True, nogil=False, cache=True) -def grad_nodes(p, ndim, dxy): - dx, dy, dz = dxy - - indices = [idx for idx in it.product([slice(0, -1), slice(1, None)], repeat=ndim)] - if ndim == 2: - signs_x = (-1.0, -1.0, +1.0, +1.0) - signs_y = (-1.0, +1.0, -1.0, +1.0) - signs_z = (0.0, 0.0, 0.0, 0.0) - elif ndim == 3: - signs_x = (-1.0, -1.0, -1.0, -1.0, +1.0, +1.0, +1.0, +1.0) - signs_y = (-1.0, -1.0, +1.0, +1.0, -1.0, -1.0, +1.0, +1.0) - signs_z = (-1.0, +1.0, -1.0, +1.0, -1.0, +1.0, -1.0, +1.0) - - Dpx, Dpy, Dpz = 0.0, 0.0, 0.0 - cnt = 0 - for index in indices: - Dpx += signs_x[cnt] * p[index] - Dpy += signs_y[cnt] * p[index] - Dpz += signs_z[cnt] * p[index] - cnt += 1 - - Dpx *= 0.5 ** (ndim - 1) / dx - Dpy *= 0.5 ** (ndim - 1) / dy - Dpz *= 0.5 ** (ndim - 1) / dz - - return Dpx, Dpy, Dpz - - -def divergence_nodes(rhs, elem, node, Sol, ud): - ndim = elem.ndim - igs = elem.igs - dxyz = node.dxyz - inner_idx = np.empty((ndim), dtype=object) - - for dim in range(ndim): - is_periodic = ud.bdry_type[dim] == opts.BdryType.PERIODIC - inner_idx[dim] = slice(igs[dim] - is_periodic, -igs[dim] + is_periodic) - inner_idx_p1y = np.copy(inner_idx) - inner_idx_p1y[1] = slice(1, -1) - - indices = [idx for idx in it.product([slice(0, -1), slice(1, None)], repeat=ndim)] - signs = [sgn for sgn in it.product([1, -1], repeat=ndim)] - inner_idx = tuple(inner_idx) - inner_idx_p1y = tuple(inner_idx_p1y) - - if not hasattr(ud, "ATMOSPHERIC_EXTENSION"): - if ( - ud.bdry_type[1] == opts.BdryType.WALL - or ud.bdry_type[1] == opts.BdryType.RAYLEIGH - ): - Sol.rhou[:, :2, ...] = 0.0 - Sol.rhov[:, :2, ...] = 0.0 - Sol.rhow[:, :2, ...] = 0.0 - - # if ud.bdry_type[1] == BdryType.WALL: - Sol.rhou[:, -2:, ...] = 0.0 - Sol.rhov[:, -2:, ...] = 0.0 - Sol.rhow[:, -2:, ...] = 0.0 - - Y = Sol.rhoY / Sol.rho - - Ux = np.diff(Sol.rhou * Y, axis=0) / elem.dx - Ux = 0.5 * (Ux[:, :-1, ...] + Ux[:, 1:, ...]) - - Vy = np.diff(Sol.rhov * Y, axis=1) / elem.dy - Vy = 0.5 * (Vy[:-1, ...] + Vy[1:, ...]) - - if ndim == 3: - Ux = -0.5 * (Ux[..., :-1] + Ux[..., 1:]) - Vy = 0.5 * (Vy[..., :-1] + Vy[..., 1:]) - - Wz = np.diff(Sol.rhow * Y, axis=2) / elem.dz - Wz = 0.5 * (Wz[:-1, ...] + Wz[1:, ...]) - Wz = 0.5 * (Wz[:, :-1, ...] + Wz[:, 1:, ...]) - - rhs[1:-1, 1:-1, 1:-1] = Ux + Vy + Wz - else: - rhs = Ux + Vy - - rhs_max = np.max(rhs[inner_idx]) if np.max(rhs[inner_idx]) > 0 else 0 - return rhs - + bdry.scale_wall_node_values(mpv.wcenter, node, ud) def rhs_from_p_old(rhs, node, mpv): igs = node.igs @@ -667,60 +429,152 @@ def rhs_from_p_old(rhs, node, mpv): rhs_n = rhs + 0.0 * rhs_hh return rhs_n +@njit(cache=True) +def _compute_coriolis_coefficients(wh1, wh2, wv, nu, nonhydro): + """Compute coefficients for the H^-1 matrix multiplication. + + This corresponds to equation (C11) in the mathematical formulation. + """ + # Common terms + wh1_sq = wh1 * wh1 + wh2_sq = wh2 * wh2 + wv_sq = wv * wv + nu_nh = nu + nonhydro + + # Denominator (det(H)) + denom = 1.0 / (wh1_sq + wh2_sq + nu_nh * (wv_sq + 1.0)) + + # H^-1 matrix elements (row-major order) + # Row 1: U equation coefficients + h11 = (wh1_sq + nu_nh) * denom + h12 = nonhydro * (wh1 * wv + wh2) * denom + h13 = (wh1 * wh2 - nu_nh * wv) * denom + + # Row 2: V equation coefficients + h21 = (wh1 * wv - wh2) * denom + h22 = nonhydro * (1.0 + wv_sq) * denom + h23 = (wh2 * wv + wh1) * denom + + # Row 3: W equation coefficients + h31 = (wh1 * wh2 + nu_nh * wv) * denom + h32 = nonhydro * (wh2 * wv - wh1) * denom + h33 = (nu_nh + wh2_sq) * denom + + return h11, h12, h13, h21, h22, h23, h31, h32, h33 + +@njit(cache=True) +def apply_coriolis_matrix_inplace(u_vec, v_vec, w_vec, U,V,W, wh1, wh2, wv, nu, nonhydro): + """Apply H^-1 matrix multiplication in-place. + + Corresponds to the equation: U^{n+1} = H^{-1}(U^{n*} - Δt_{cp}(Pθ)^* ∇π^{n+1}) + """ + # Get matrix coefficients + h11, h12, h13, h21, h22, h23, h31, h32, h33 = _compute_coriolis_coefficients( + wh1, wh2, wv, nu, nonhydro + ) + + U[...] = u_vec + V[...] = v_vec + W[...] = w_vec + + # Matrix multiplication: [U_new, V_new, W_new] = H^-1 @ [U_old, V_old, W_old] + u_vec[...] = h11 * U + h12 * V + h13 * W + v_vec[...] = h21 * U + h22 * V + h23 * W + w_vec[...] = h31 * U + h32 * V + h33 * W +# Refactored main function def multiply_inverse_coriolis( - Vec, Sol, mpv, ud, elem, node, dt, attrs=("rhou", "rhov", "rhow"), get_coeffs=False + Vec, mem, ud, dt, attrs=("rhou", "rhov", "rhow"), get_coeffs=False ): + """Coriolis matrix multiplication.""" nonhydro = ud.nonhydrostasy g = ud.gravity_strength[1] Msq = ud.Msq wh1, wv, wh2 = dt * ud.coriolis_strength + strat = mem.mpv.HydroState_n.get_dSdy(mem.elem, mem.node) + Y = mem.sol.rhoY / mem.sol.rho + nu = -(dt**2) * (g / Msq) * strat * Y + + # Get vector components + VecU = getattr(Vec, attrs[0]) + VecV = getattr(Vec, attrs[1]) + VecW = getattr(Vec, attrs[2]) - strat = mpv.HydroState_n.get_dSdy(elem, node) + U,V,W = mem.cache.get_velocity_array_views(VecU.shape) + + apply_coriolis_matrix_inplace(VecU, VecV, VecW, U,V,W,wh1, wh2, wv, nu, nonhydro) + + # Return coefficients + if get_coeffs: + h11, h12, _, h21, h22, h23, h31, h32, h33 = _compute_coriolis_coefficients(wh1, wh2, wv, nu, nonhydro) + return (h11.T, h22.T, h12.T, h21.T) + - Y = Sol.rhoY / Sol.rho - nu = -(dt**2) * (g / Msq) * strat * Y +def divergence_nodes(rhs, elem, node, Sol, ud): + """Main divergence function - handles boundary conditions and calls JIT-compiled core.""" + ndim = elem.ndim + + # Handle boundary conditions + if not hasattr(ud, "ATMOSPHERIC_EXTENSION"): + if (ud.bdry_type[1] == opts.BdryType.WALL or + ud.bdry_type[1] == opts.BdryType.RAYLEIGH): + Sol.rhou[:, :2, ...] = 0.0 + Sol.rhov[:, :2, ...] = 0.0 + Sol.rhow[:, :2, ...] = 0.0 + Sol.rhou[:, -2:, ...] = 0.0 + Sol.rhov[:, -2:, ...] = 0.0 + Sol.rhow[:, -2:, ...] = 0.0 + + # Call appropriate JIT-compiled function + if ndim == 2: + rhs[:] = _momentum_pot_temp_divergence_2d_jit( + Sol.rho, Sol.rhou, Sol.rhov, Sol.rhoY, + elem.dx, elem.dy + ) + else: + _momentum_pot_temp_divergence_3d_jit( + rhs, Sol.rho, Sol.rhou, Sol.rhov, Sol.rhow, Sol.rhoY, + elem.dx, elem.dy, elem.dz + ) + + return rhs - # get coefficients of the explicit terms - # common denominator - denom = 1.0 / (wh1**2 + wh2**2 + (nu + nonhydro) * (wv**2 + 1.0)) - # U update - coeff_uu = wh1**2 + nu + nonhydro - coeff_uv = nonhydro * (wh1 * wv + wh2) - coeff_uw = wh1 * wh2 - (nu + nonhydro) * wv +@njit(cache=True) +def _momentum_pot_temp_divergence_2d_jit(rho, rhou, rhov, rhoY, dx, dy): + """ + JIT-compiled 2D momentum-potential temperature divergence calculation. + Computes ∇·(ρu θ, ρv θ) where θ = ρY/ρ is the potential temperature. + """ + # Calculate potential temperature θ = ρY / ρ + theta = rhoY / rho - # V update - coeff_vu = wh1 * wv - wh2 - coeff_vv = nonhydro * (1 + wv**2) - coeff_vw = wh2 * wv + wh1 + # Compute momentum-potential temperature flux components + rhou_theta = rhou * theta # x-momentum flux weighted by potential temperature + rhov_theta = rhov * theta # y-momentum flux weighted by potential temperature + + # Use generic divergence operator + return operators.compute_divergence_2d(rhou_theta, rhov_theta, dx, dy) - # W update - coeff_wu = wh1 * wh2 + (nu + nonhydro) * wv - coeff_wv = nonhydro * (wh2 * wv - wh1) - coeff_ww = nu + nonhydro + wh2**2 - VecU = getattr(Vec, attrs[0]) - VecV = getattr(Vec, attrs[1]) - VecW = getattr(Vec, attrs[2]) +@njit(cache=True) +def _momentum_pot_temp_divergence_3d_jit(rhs, rho, rhou, rhov, rhow, rhoY, dx, dy, dz): + """ + JIT-compiled 3D momentum-potential temperature divergence calculation. + Computes ∇·(ρu θ, ρv θ, ρw θ) where θ = ρY/ρ is the potential temperature. + """ + # Calculate potential temperature θ = ρY / ρ + theta = rhoY / rho - # Do the updates - U = denom * (coeff_uu * VecU + coeff_uv * VecV + coeff_uw * VecW) - V = denom * (coeff_vu * VecU + coeff_vv * VecV + coeff_vw * VecW) - W = denom * (coeff_wu * VecU + coeff_wv * VecV + coeff_ww * VecW) + # Compute momentum-potential temperature flux components + rhou_theta = rhou * theta # x-momentum flux weighted by potential temperature + rhov_theta = rhov * theta # y-momentum flux weighted by potential temperature + rhow_theta = rhow * theta # z-momentum flux weighted by potential temperature - VecU[...] = U - VecV[...] = V - VecW[...] = W + # Use generic total divergence operator + total_div = operators.compute_divergence_3d_total(rhou_theta, rhov_theta, rhow_theta, dx, dy, dz) - i1 = node.i1 - if get_coeffs: - # coriolis_parameters = ((coeff_uu * denom)[i1].reshape(-1,), (coeff_vv * denom)[i1].reshape(-1,), (coeff_uv * denom)[i1].reshape(-1,), (coeff_vu * denom)[i1].reshape(-1,)) - coriolis_parameters = ( - (coeff_uu * denom).T, - (coeff_vv * denom).T, - (coeff_uv * denom).T, - (coeff_vu * denom).T, - ) - return coriolis_parameters + # Assign to inner region + rhs[1:-1, 1:-1, 1:-1] = total_div + diff --git a/src/pybella/flow_solver/utils/boundary.py b/src/pybella/flow_solver/utils/boundary.py index 74209bf9..a2dccff7 100644 --- a/src/pybella/flow_solver/utils/boundary.py +++ b/src/pybella/flow_solver/utils/boundary.py @@ -1,9 +1,10 @@ """ For more details on this module, refer to the write-up :ref:`boundary_handling`. """ + import copy import numpy as np -from . import options as opts +from ...utils import options as opts from ...utils import io @@ -285,7 +286,7 @@ def get_gravity_padding(ndim, cur_idx, direction, offset, elem, y_axs=None): """ cur_i = np.copy(cur_idx) cur_idx += offset * ((elem.icy - 1) - 2 * cur_idx) - gravity_padding = [(slice(None))] * ndim + gravity_padding = [slice(None)] * ndim if y_axs == None: # y_axs = ndim - 1 y_axs = 1 @@ -498,6 +499,7 @@ def get_bottom_tau_y(ud, elem, node, alpha, cutoff=0.5): return tauc_y, taun_y + def apply_rayleigh_forcing( Sol, mpv, @@ -520,9 +522,7 @@ def apply_rayleigh_forcing( t_offset = 0.5 * dt if half else dt if ud.rayleigh_forcing_type == "file": - reader = io.read_input( - ud.rayleigh_forcing_fn, ud.rayleigh_forcing_path - ) + reader = io.read_input(ud.rayleigh_forcing_fn, ud.rayleigh_forcing_path) if Sol_half_new is None or mpv_half_new is None: Sol_half_new = copy.deepcopy(Sol) @@ -536,9 +536,7 @@ def apply_rayleigh_forcing( Yp = Sol_half_new.rhoY / Sol_half_new.rho - mpv.HydroState.Y0.reshape(1, -1) pi = mpv_half_new.p2_nodes - bdry.rayleigh_damping( - Sol, mpv, ud, elem, node, [up, vp, Yp, pi, t + t_offset] - ) + bdry.rayleigh_damping(Sol, mpv, ud, elem, node, [up, vp, Yp, pi, t + t_offset]) elif ud.rayleigh_forcing_type == "func": s = 5.0e-3 + 1e-4 + 0e-5 @@ -554,6 +552,7 @@ def apply_rayleigh_forcing( bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) + def rayleigh_damping(Sol, mpv, ud, elem, node, forcing=None): u = Sol.rhou / Sol.rho # [elem.i2] v = Sol.rhov / Sol.rho # [elem.i2] @@ -657,3 +656,28 @@ def check_flux_bcs(Lefts, Rights, elem, split_step, ud): Rights.rhow[right_ghost] = Lefts.rhow[:, -igx - 2] Rights.rhoY[right_ghost] = Lefts.rhoY[:, -igx - 2] # print(Rights.rhoY[right_ghost]) + + +def scale_wall_node_values(rhs, node, ud, factor=0.5): + """Scale values at wall boundary nodes by a given factor.""" + ndim = node.ndim + igs = node.igs + + for dim in range(ndim): + # Check if this dimension has wall boundaries + is_wall = ( + ud.bdry_type[dim] == opts.BdryType.WALL or + ud.bdry_type[dim] == opts.BdryType.RAYLEIGH + ) + + if is_wall: + # Create index for all dimensions + idx = [slice(igs[d], -igs[d]) for d in range(ndim)] + + # Scale first and last interior nodes in this dimension + for boundary_idx in [igs[dim], -igs[dim] - 1]: + idx[dim] = boundary_idx + rhs[tuple(idx)] *= factor + # rhs = rhs.at[tuple(idx)].multiply(factor) + + return rhs \ No newline at end of file diff --git a/src/pybella/flow_solver/utils/options.py b/src/pybella/flow_solver/utils/options.py deleted file mode 100644 index 621bea7a..00000000 --- a/src/pybella/flow_solver/utils/options.py +++ /dev/null @@ -1,51 +0,0 @@ -from enum import Enum # ! Version > Python 3.4 - -# class RecoveryOrder(Enum): -# FIRST = 0 -# SECOND = 1 - -# class TimeIntegrator(Enum): -# OP_SPLIT = 0 -# OP_SPLIT_MD_UPDATE = 1 -# HUEN = 2 -# EXPL_MIDPT = 3 -# RK3_SKAMA = 4 -# RK3_TEST = 5 -# SI_MIDPT = 6 -# STRANG = 7 - -# class HillShapes(Enum): -# SCHULTOW = 0 -# AGNESI = 1 - -# class MolecularTransport(Enum): -# FULL_MOLECULAR_TRANSPORT = 0 -# STRAKA_DIFFUSION_MODEL = 1 -# NO_MOLECULAR_TRANSPORT = 2 - -# class BottomBC(Enum): -# ZERO_ORDER_EXTRAPOL = 0 -# BOTTOM_BC_DEFAULT = 1 - - -class LimiterType(Enum): - NONE = 0 - # MINMOD = 1 - # VANLEER = 2 - # VANLEERSmooth = 3 - # SUPERBEE = 4 - # MONOTONIZED_CENTRAL = 5 - # SWEBY_MUNZ = 6 - # RUPE = 7 - # NO_SLOPE = 8 - # NUMBER_OF_LIMITER = 9 - - -class BdryType(Enum): - """ - An enumeration class that defines the accepted boundary condition types. - """ - - WALL = "symmetric" - PERIODIC = "wrap" - RAYLEIGH = "radiation" diff --git a/src/pybella/utils/prepare.py b/src/pybella/flow_solver/utils/prepare.py similarity index 86% rename from src/pybella/utils/prepare.py rename to src/pybella/flow_solver/utils/prepare.py index 5d776d59..9178de84 100644 --- a/src/pybella/utils/prepare.py +++ b/src/pybella/flow_solver/utils/prepare.py @@ -1,23 +1,23 @@ import numpy as np -from . import user_data, io, data_structures +from ...utils import user_data, io, data_structures -from ..flow_solver.discretisation import grid as dis_grid -from ..flow_solver.utils import variable as var -from ..flow_solver.utils import boundary as bdry -from ..flow_solver.physics import hydrostatics -from ..flow_solver.physics.low_mach import mpv as lm_var -from ..flow_solver.physics.gas_dynamics import thermodynamics as gd_thermodynamics +from ..discretisation import grid as dis_grid +from . import variable as var +from . import boundary as bdry +from ..physics import hydrostatics +from ..physics.low_mach import mpv as lm_var +from ..physics.gas_dynamics import thermodynamics as gd_thermodynamics # test module -from ..tests import diagnostics as diag +from ...tests import diagnostics as diag def initialise(): #### # Initialise simulation state #### - from . import sim_params as params + from ...utils import sim_params as params np.set_printoptions(precision=params.print_precision) @@ -43,6 +43,7 @@ def initialise(): sol = var.Vars(elem.sc, ud) + # Move these to the FlowSolverCache flux = np.empty((3), dtype=object) flux[0] = var.States(elem.sfx, ud) if elem.ndim > 1: @@ -84,6 +85,10 @@ def initialise(): sol = sol_init(sol, mpv, elem, node, th, ud) + # Initialise cache and add to simulation state + flow_cache = var.FlowSolverCache() + + ensemble_state.update_member( elem=elem, node=node, @@ -91,6 +96,7 @@ def initialise(): flux=flux, mpv=mpv, th=th, + cache=flow_cache ) restart_params = data_structures.RestartParameters( @@ -112,6 +118,8 @@ def initialise(): interface_params=interface_params, ) + + return sim_st diff --git a/src/pybella/flow_solver/utils/variable.py b/src/pybella/flow_solver/utils/variable.py index 48811fb0..093dff1b 100644 --- a/src/pybella/flow_solver/utils/variable.py +++ b/src/pybella/flow_solver/utils/variable.py @@ -1,6 +1,6 @@ import numpy as np import scipy as sp - +import logging class Vars(object): """ @@ -37,6 +37,14 @@ def __init__(self, size, ud): self.rhow = np.zeros((size)) self.rhoY = np.zeros((size)) self.rhoX = np.zeros(([ud.nspec] + list(size))) + + self.u = np.zeros((size)) + self.v = np.zeros((size)) + self.w = np.zeros((size)) + self.Y = np.zeros((size)) + self.X = np.zeros(([ud.nspec] + list(size))) + self.p = np.zeros((size)) + self.squeezer() # will be a better way of doing this @@ -48,8 +56,6 @@ def squeezer(self): for key, value in vars(self).items(): setattr(self, key, value.squeeze()) - # method written for 2D - def primitives(self, th): """ Calculate the primitive quantities from the state variables and extend the data container to include these quantities. @@ -69,23 +75,15 @@ def primitives(self, th): p : ndarray(size_of_rhoY) """ - nonzero_idx = np.nonzero(self.rho) - - self.u = np.zeros_like(self.rhou) - self.v = np.zeros_like(self.rhov) - self.w = np.zeros_like(self.rhow) - self.Y = np.zeros_like(self.rhoY) - self.X = np.zeros_like(self.rhoX) - self.p = np.zeros_like(self.rhoY) - - # the non-zero indices are used here for the case where Lefts and - # Rights in the HLLE Solver has one column that is zeros. - self.u[nonzero_idx] = self.rhou[nonzero_idx] / self.rho[nonzero_idx] - self.v[nonzero_idx] = self.rhov[nonzero_idx] / self.rho[nonzero_idx] - self.w[nonzero_idx] = self.rhow[nonzero_idx] / self.rho[nonzero_idx] - self.Y[nonzero_idx] = self.rhoY[nonzero_idx] / self.rho[nonzero_idx] - self.X[nonzero_idx] = self.rhoX[nonzero_idx] / self.rho[nonzero_idx] - self.p[nonzero_idx] = self.rhoY[nonzero_idx] ** th.gamm + with np.errstate(divide='ignore', invalid='ignore'): + # Direct division without nonzero indexing + # We know that when this method is called in recovery, we always have one column of zeroes in self.rho. + self.u[...] = self.rhou / self.rho + self.v[...] = self.rhov / self.rho + self.w[...] = self.rhow / self.rho + self.Y[...] = self.rhoY / self.rho + self.X[...] = self.rhoX / self.rho + self.p[...] = self.rhoY ** th.gamm def flip(self): """ @@ -112,9 +110,9 @@ def mod_bg_wind(self, ud, fac): v0 = ud.v_wind_speed w0 = ud.w_wind_speed - self.rhou = self.rhou + fac * u0 * self.rho - self.rhov = self.rhov + fac * v0 * self.rho - self.rhow = self.rhow + fac * w0 * self.rho + self.rhou[...] = self.rhou + fac * u0 * self.rho + self.rhov[...] = self.rhov + fac * v0 * self.rho + self.rhow[...] = self.rhow + fac * w0 * self.rho class States(Vars): @@ -138,16 +136,6 @@ def __init__(self, size, ud): """ super().__init__(size, ud) - self.u = np.zeros((size)) - self.v = np.zeros((size)) - self.w = np.zeros((size)) - self.q = np.zeros((size)) - self.p = np.zeros((size)) - self.c = np.zeros((size)) - self.entro = np.zeros((size)) - self.H = np.zeros((size)) - self.Y = np.zeros((size)) - self.X = np.zeros(([ud.nspec] + list(size))) self.p0 = np.zeros((size)) self.p20 = np.zeros((size)) @@ -163,40 +151,34 @@ def __init__(self, size, ud): self.get_S0c = self.get_S0c self.init_dSdy = False - self.init_S0c = False + self.init_S0c = False def get_dSdy(self, elem, node): - if self.init_dSdy: - return self.dSdy - else: - ndim = node.ndim - dy = node.dy - - dSdy = self.S0 - dSdy = sp.signal.convolve(dSdy, [1.0, -1.0], mode="valid") / dy - - for dim in range(0, ndim, 2): - dSdy = np.expand_dims(dSdy, dim) - dSdy = np.repeat(dSdy, elem.sc[dim], axis=dim) - - self.dSdy = dSdy + if not self.init_dSdy: + logging.info("Computing dSdy") + self.dSdy = sp.signal.convolve(self.S0, [1.0, -1.0], mode="valid") / node.dy + + for dim in range(0, node.ndim, 2): + self.dSdy = np.expand_dims(self.dSdy, dim) + self.dSdy = np.repeat(self.dSdy, elem.sc[dim], axis=dim) + self.init_dSdy = True - return dSdy + + return self.dSdy def get_S0c(self, elem): - if self.init_S0c: - return self.S0c - else: - ndim = elem.ndim - S0c = self.S0 - - for dim in range(0, ndim, 2): - S0c = np.expand_dims(S0c, dim) - S0c = np.repeat(S0c, elem.sc[dim], axis=dim) - - self.S0c = S0c + if not self.init_S0c: + logging.info("Computing S0c") + S0c_result = self.S0 + + for dim in range(0, elem.ndim, 2): + S0c_result = np.expand_dims(S0c_result, dim) + S0c_result = np.repeat(S0c_result, elem.sc[dim], axis=dim) + + self.S0c = S0c_result self.init_S0c = True - return S0c + + return self.S0c class Characters(object): @@ -230,10 +212,6 @@ def __init__(self, size): self.X = np.zeros((size)) self.Y = np.zeros((size)) - self.plus = np.zeros((size)) - self.minus = np.zeros((size)) - self.entro = np.zeros((size)) - self.squeezer() def squeezer(self): @@ -243,3 +221,69 @@ def squeezer(self): """ for key, value in vars(self).items(): setattr(self, key, value.squeeze()) + + +class FlowSolverCache: + """Cache for flow solver specific computations.""" + + def __init__(self): + self._recovery_cache = {} + self._velocity_cache = {} + + def get_recovery_objects(self, shape, ud): + """Get cached recovery objects or create new ones.""" + + cache_key = (tuple(shape), id(ud)) + if cache_key not in self._recovery_cache: + logging.info("Cache: Creating new recovery objects with shape %s", shape) + self._recovery_cache[cache_key] = { + 'Diffs': Characters(shape), + 'Ampls': Characters(shape), + 'Lefts': States(shape, ud), + 'Rights': States(shape, ud), + 'Slopes': Characters(shape), + } + + # Reset objects if they have reset methods + cache_obj = self._recovery_cache[cache_key] + # for obj in cache_obj.values(): + # if hasattr(obj, 'zero'): + # obj.zero() + + return cache_obj + + def get_velocity_arrays(self, shape, dtype=np.float64): + """Get cached velocity arrays (U, V, W) or create new ones.""" + + cache_key = (tuple(shape), dtype) + + if cache_key not in self._velocity_cache: + logging.info("Cache: Creating new velocity arrays with shape %s", shape) + self._velocity_cache[cache_key] = { + 'U': np.zeros(shape, dtype=dtype), + 'V': np.zeros(shape, dtype=dtype), + 'W': np.zeros(shape, dtype=dtype) + } + + # Clear arrays for reuse + cache_obj = self._velocity_cache[cache_key] + # for arr in cache_obj.values(): + # arr.fill(0.0) + + return cache_obj + + def get_velocity_array_views(self, shape, dtype=np.float64): + """Get views of cached velocity arrays for in-place operations.""" + cache_obj = self.get_velocity_arrays(shape, dtype) + + return cache_obj['U'], cache_obj['V'], cache_obj['W'] + + def clear_velocity_cache(self): + """Clear velocity cache to free memory.""" + self._velocity_cache.clear() + + def clear_all(self): + """Clear all caches to free memory.""" + self._recovery_cache.clear() + self._velocity_cache.clear() + diff --git a/src/pybella/interfaces/dynamics_blending/schemes.py b/src/pybella/interfaces/dynamics_blending/schemes.py index bad652ce..4904372c 100644 --- a/src/pybella/interfaces/dynamics_blending/schemes.py +++ b/src/pybella/interfaces/dynamics_blending/schemes.py @@ -7,6 +7,7 @@ from ...flow_solver.physics.gas_dynamics import eos as gd_eos from ...utils import io + class Blend(object): """ Class that takes care of the blending interface. @@ -119,8 +120,8 @@ def do_comp_to_psinc_conv(mem, bld, ud, label, writer): def do_psinc_to_comp_conv( - sst, mem, + ud, bld, label, writer, @@ -134,16 +135,14 @@ def do_psinc_to_comp_conv( mpv_freeze = copy.deepcopy(mem.mpv) ret = time_update.do( - sst, mem, + ud, tout, bld=None, writer=None, debug_writer=io.NullDebugWriter(), ) - ud = sst.ud - fac_old = ud.blending_weight fac_new = 1.0 - fac_old dp2n_0 = fac_new * ret.mpv.p2_nodes_half + fac_old * mpv_freeze.p2_nodes_half @@ -390,147 +389,147 @@ def do_hydro_to_nonhydro_conv( ############################### # if c1 or c2: - # logging.info( - # termcolor.colored("hydrostatic to nonhydrostatic conversion...", "blue") - # ) - - # writer.write_all(mem, str(label) + "_half_full") - # writer.populate(str(label) + "_ic", "pwchi", Sol.pwchi) - - # if test_hydrob == False: - # Sol = copy.deepcopy(Sol_half_old) - # # mpv = copy.deepcopy(mpv_half_old) - - # logging.info(termcolor.colored("test_hydrob == False", "red")) - # writer.write_all(mem, str(label) + "_quarter") - - # writer.populate(str(label) + "_quarter", "pwchi", Sol.pwchi) - - # logging.info("quarter dt = %.8f" % (dt * 0.5)) - - # ret = do( - # Sol_half_old, - # flux_half_old, - # mpv_half_old, - # dt - 0.5 * dt, - # dt + 0.5 * dt, - # ud, - # elem, - # node, - # [0, 0], - # th, - # bld=None, - # writer=None, - # debug=False, - # ) - - # Sol_tu = copy.deepcopy(ret[0]) - # # mpv_tu = copy.deepcopy(ret[2]) - # Sol.rho[...] = Sol_tu.rho_half - # Sol.rhou[...] = Sol_tu.rhou_half - # Sol.rhov[...] = Sol_tu.rhov_half - # Sol.rhow[...] = Sol_tu.rhow_half - # Sol.rhoX[...] = Sol_tu.rhoX_half - # Sol.rhoY[...] = Sol_tu.rhoY_half - # Sol.pwchi[...] = Sol_tu.pwchi - - # # mpv.p2_nodes[...] = mpv_tu.p2_nodes_half - - # writer.write_all(mem, str(label) + "_half") - - # writer.populate(str(label) + "_half", "pwchi", Sol.pwchi) - - # ret = do( - # Sol, - # flux, - # mpv, - # dt, - # 2.0 * dt, - # ud, - # elem, - # node, - # [0, 0], - # th, - # bld=None, - # writer=None, - # debug=False, - # ) - - # Sol = copy.deepcopy(ret[0]) - # flux = copy.deepcopy(ret[1]) - # mpv = copy.deepcopy(ret[2]) - - # if test_hydrob == True: - # Sol = copy.deepcopy(Sol_half_old) - # # mpv = copy.deepcopy(mpv_half_old) - - # logging.info(termcolor.colored("test_hydrob == False", "red")) - # writer.write_all(mem, str(label) + "_quarter") - - # # writer.populate(str(label)+'_quarter', 'pwchi', Sol.pwchi) - - # logging.info("quarter dt = %.8f" % (dt * 0.5)) - - # ret = do( - # Sol_half_old, - # flux_half_old, - # mpv_half_old, - # dt - 0.5 * dt, - # dt + 0.5 * dt, - # ud, - # elem, - # node, - # [0, 0], - # th, - # bld=None, - # writer=None, - # debug=False, - # ) - - # Sol_tu = copy.deepcopy(ret[0]) - # # mpv_tu = copy.deepcopy(ret[2]) - # Sol.rho[...] = Sol_tu.rho_half - # Sol.rhou[...] = Sol_tu.rhou_half - # Sol.rhov[...] = Sol_tu.rhov_half - # Sol.rhow[...] = Sol_tu.rhow_half - # Sol.rhoX[...] = Sol_tu.rhoX_half - # Sol.rhoY[...] = Sol_tu.rhoY_half - # Sol.pwchi[...] = Sol_tu.pwchi - - # # mpv.p2_nodes[...] = mpv_tu.p2_nodes_half - - # # writer.write_all(Sol,mpv,elem,node,th,str(label)+'_half') - - # # writer.populate(str(label)+'_half', 'pwchi', Sol.pwchi) - - # ret = do( - # Sol, - # flux, - # mpv, - # dt, - # 2.0 * dt, - # ud, - # elem, - # node, - # [0, 0], - # th, - # bld=None, - # writer=None, - # debug=False, - # ) - - # Sol = copy.deepcopy(ret[0]) - # flux = copy.deepcopy(ret[1]) - # mpv = copy.deepcopy(ret[2]) - # # writer.write_all(Sol,mpv,elem,node,th,str(label)+'_half') - # # writer.populate(str(label)+'_half', 'pwchi', Sol.pwchi) - - # logging.info(termcolor.colored("test_hydrob == True", "red")) - - # if test_hydrob == False: - # dt *= 2.0 - # if c2: - # ud.is_nonhydrostatic = 1 + # logging.info( + # termcolor.colored("hydrostatic to nonhydrostatic conversion...", "blue") + # ) + + # writer.write_all(mem, str(label) + "_half_full") + # writer.populate(str(label) + "_ic", "pwchi", Sol.pwchi) + + # if test_hydrob == False: + # Sol = copy.deepcopy(Sol_half_old) + # # mpv = copy.deepcopy(mpv_half_old) + + # logging.info(termcolor.colored("test_hydrob == False", "red")) + # writer.write_all(mem, str(label) + "_quarter") + + # writer.populate(str(label) + "_quarter", "pwchi", Sol.pwchi) + + # logging.info("quarter dt = %.8f" % (dt * 0.5)) + + # ret = do( + # Sol_half_old, + # flux_half_old, + # mpv_half_old, + # dt - 0.5 * dt, + # dt + 0.5 * dt, + # ud, + # elem, + # node, + # [0, 0], + # th, + # bld=None, + # writer=None, + # debug=False, + # ) + + # Sol_tu = copy.deepcopy(ret[0]) + # # mpv_tu = copy.deepcopy(ret[2]) + # Sol.rho[...] = Sol_tu.rho_half + # Sol.rhou[...] = Sol_tu.rhou_half + # Sol.rhov[...] = Sol_tu.rhov_half + # Sol.rhow[...] = Sol_tu.rhow_half + # Sol.rhoX[...] = Sol_tu.rhoX_half + # Sol.rhoY[...] = Sol_tu.rhoY_half + # Sol.pwchi[...] = Sol_tu.pwchi + + # # mpv.p2_nodes[...] = mpv_tu.p2_nodes_half + + # writer.write_all(mem, str(label) + "_half") + + # writer.populate(str(label) + "_half", "pwchi", Sol.pwchi) + + # ret = do( + # Sol, + # flux, + # mpv, + # dt, + # 2.0 * dt, + # ud, + # elem, + # node, + # [0, 0], + # th, + # bld=None, + # writer=None, + # debug=False, + # ) + + # Sol = copy.deepcopy(ret[0]) + # flux = copy.deepcopy(ret[1]) + # mpv = copy.deepcopy(ret[2]) + + # if test_hydrob == True: + # Sol = copy.deepcopy(Sol_half_old) + # # mpv = copy.deepcopy(mpv_half_old) + + # logging.info(termcolor.colored("test_hydrob == False", "red")) + # writer.write_all(mem, str(label) + "_quarter") + + # # writer.populate(str(label)+'_quarter', 'pwchi', Sol.pwchi) + + # logging.info("quarter dt = %.8f" % (dt * 0.5)) + + # ret = do( + # Sol_half_old, + # flux_half_old, + # mpv_half_old, + # dt - 0.5 * dt, + # dt + 0.5 * dt, + # ud, + # elem, + # node, + # [0, 0], + # th, + # bld=None, + # writer=None, + # debug=False, + # ) + + # Sol_tu = copy.deepcopy(ret[0]) + # # mpv_tu = copy.deepcopy(ret[2]) + # Sol.rho[...] = Sol_tu.rho_half + # Sol.rhou[...] = Sol_tu.rhou_half + # Sol.rhov[...] = Sol_tu.rhov_half + # Sol.rhow[...] = Sol_tu.rhow_half + # Sol.rhoX[...] = Sol_tu.rhoX_half + # Sol.rhoY[...] = Sol_tu.rhoY_half + # Sol.pwchi[...] = Sol_tu.pwchi + + # # mpv.p2_nodes[...] = mpv_tu.p2_nodes_half + + # # writer.write_all(Sol,mpv,elem,node,th,str(label)+'_half') + + # # writer.populate(str(label)+'_half', 'pwchi', Sol.pwchi) + + # ret = do( + # Sol, + # flux, + # mpv, + # dt, + # 2.0 * dt, + # ud, + # elem, + # node, + # [0, 0], + # th, + # bld=None, + # writer=None, + # debug=False, + # ) + + # Sol = copy.deepcopy(ret[0]) + # flux = copy.deepcopy(ret[1]) + # mpv = copy.deepcopy(ret[2]) + # # writer.write_all(Sol,mpv,elem,node,th,str(label)+'_half') + # # writer.populate(str(label)+'_half', 'pwchi', Sol.pwchi) + + # logging.info(termcolor.colored("test_hydrob == True", "red")) + + # if test_hydrob == False: + # dt *= 2.0 + # if c2: + # ud.is_nonhydrostatic = 1 return Sol, mpv @@ -539,8 +538,8 @@ def do_hydro_to_nonhydro_conv( # Blending calls from data.py ###################################################### def blending_before_timestep( - sst, mem, + ud, bld, label, writer, @@ -555,8 +554,7 @@ def blending_before_timestep( # Blending : Do full regime to limit regime conversion ###################################################### # do unpacking - elem, node, Sol, flux, mpv, th, _ = mem - ud = sst.ud + elem, node, Sol, flux, mpv, th, _, _ = mem # these make sure that we are the correct window step if bld is not None and window_step == 0: @@ -585,8 +583,8 @@ def blending_before_timestep( # distinguish between Euler and SWE blending if ud.blending_conv != "swe": do_psinc_to_comp_conv( - sst, mem, + ud, bld, ud, label, @@ -639,8 +637,8 @@ def blending_before_timestep( # Distinguish between SWE and Euler blendings if ud.blending_conv != "swe": do_psinc_to_comp_conv( - sst, mem, + ud, bld, label, writer, @@ -742,8 +740,8 @@ def blending_after_timestep( def prepare_blending( - sst, mem, + ud, bld, label, writer, @@ -752,16 +750,15 @@ def prepare_blending( t, dt, swe_to_lake, - debug, - ): + debug, +): - ud = sst.ud if check_and_apply_initial_hydrostatic_conversion(step, ud, bld): ud.is_nonhydrostatic = 0 swe_to_lake, Sol, mpv, t = blending_before_timestep( - sst, mem, + ud, bld, label, writer, @@ -779,22 +776,19 @@ def prepare_blending( def check_and_apply_initial_hydrostatic_conversion(step, ud, bld): """ Check and apply initial blending conversion if needed. - + Returns: bool: True if conversion was applied, False otherwise """ if step != 0 or bld is None or "imbal" not in ud.aux: return False - + hydrostatic_case = ud.is_nonhydrostatic == 0 - nonhydrostatic_case = ( - ud.is_nonhydrostatic == 1 and ud.initial_blending == True - ) - + nonhydrostatic_case = ud.is_nonhydrostatic == 1 and ud.initial_blending == True + if hydrostatic_case or nonhydrostatic_case: logging.info("nonhydrostatic to hydrostatic conversion...") ud.is_nonhydrostatic = 0 return True - - return False + return False diff --git a/src/pybella/interfaces/ic_config.py b/src/pybella/interfaces/ic_config.py index 067e5af1..09d19be5 100644 --- a/src/pybella/interfaces/ic_config.py +++ b/src/pybella/interfaces/ic_config.py @@ -19,9 +19,9 @@ "swe": "inputs.shallow_water_3D", "swe_icshear": "inputs.shallow_water_3D_icshear", "swe_dvortex": "inputs.shallow_water_3D_dvortex", - "test_travelling_vortex": "pybella.tests.test_travelling_vortex", + "test_travelling_vortex": "pybella.tests.test_travelling_vortex", "test_internal_long_wave": "pybella.tests.test_internal_long_wave", "test_lamb_wave": "pybella.tests.test_lamb_wave", "test_blending_warm_bubble": "pybella.tests.test_blending_warm_bubble", "test_unstable_lamb": "pybella.tests.test_unstable_lamb", -} \ No newline at end of file +} diff --git a/src/pybella/interfaces/postprocessing/strip_target_file.py b/src/pybella/interfaces/postprocessing/strip_target_file.py index 8c9b5948..73fa9c9a 100644 --- a/src/pybella/interfaces/postprocessing/strip_target_file.py +++ b/src/pybella/interfaces/postprocessing/strip_target_file.py @@ -1,7 +1,10 @@ import h5py import re -def do(h5path: str, ud, output_intervals: bool = False, time_increment: bool = False) -> None: + +def do( + h5path: str, ud, output_intervals: bool = False, time_increment: bool = False +) -> None: if ud.diag_updt_targets: stepmax = ud.stepmax @@ -12,7 +15,9 @@ def do(h5path: str, ud, output_intervals: bool = False, time_increment: bool = F elif 101 < stepmax <= 1000: keep_steps = list(range(100, stepmax, 100)) else: - raise ValueError(f"stepmax {stepmax} > 1000 is too large as a diagnostic target.") + raise ValueError( + f"stepmax {stepmax} > 1000 is too large as a diagnostic target." + ) else: keep_steps = [stepmax - 1] @@ -32,7 +37,7 @@ def do(h5path: str, ud, output_intervals: bool = False, time_increment: bool = F for dset_name in src_group.keys(): # Identify timestep in dataset name, e.g. p2_nodes_010_after_full_step - step_match = re.search(r'_(\d+)_after_full_step$', dset_name) + step_match = re.search(r"_(\d+)_after_full_step$", dset_name) if step_match: step = int(step_match.group(1)) if step in keep_steps: @@ -41,4 +46,4 @@ def do(h5path: str, ud, output_intervals: bool = False, time_increment: bool = F # Copy non-timestep fields like "p2_nodes_ic" src_group.copy(dset_name, dst_group) - print(f"Stripped target file written to: {output_path}") \ No newline at end of file + print(f"Stripped target file written to: {output_path}") diff --git a/src/pybella/interfaces/time_stepper/prestep.py b/src/pybella/interfaces/time_stepper/prestep.py index 3206579a..f91bc819 100644 --- a/src/pybella/interfaces/time_stepper/prestep.py +++ b/src/pybella/interfaces/time_stepper/prestep.py @@ -2,6 +2,7 @@ Pre-step modifications for the time stepper. """ + def apply_modifcations(dt, ud, step): """Apply modifications to the timestep based on user-defined parameters.""" @@ -14,4 +15,4 @@ def _apply_cfl_override(dt, ud, step): """Apply CFL override for fixed timestep cases.""" if "CFLfixed" in ud.aux and step < 2: return 21.69 / ud.t_ref - return dt \ No newline at end of file + return dt diff --git a/src/pybella/tests/diagnostics.py b/src/pybella/tests/diagnostics.py index 35f8b6c2..e7ab3c3a 100644 --- a/src/pybella/tests/diagnostics.py +++ b/src/pybella/tests/diagnostics.py @@ -25,18 +25,25 @@ def update_targets(self): for tc_name, tc in self.tcs.items(): tp = self.tps[tc_name] - dump_name = tp.name.replace("target","test") + dump_name = tp.name.replace("target", "test") self.arr_dump[dump_name] = {} for attribute in tp.attributes: - arr = self.__get_ens(tc, tp, attribute, time_increment=self.time_increment, summed=False) - self.arr_dump[dump_name][attribute] = float( - arr.sum() + arr = self.__get_ens( + tc, tp, attribute, time_increment=self.time_increment, summed=False ) + self.arr_dump[dump_name][attribute] = float(arr.sum()) if self.plot: # vis_pt.plotter accepts a list of tuples with plot and panel title. - pl = vis_pt.plotter([(arr.T, "ref"), ], ncols=1, figsize=(4, 3), sharey=False) + pl = vis_pt.plotter( + [ + (arr.T, "ref"), + ], + ncols=1, + figsize=(4, 3), + sharey=False, + ) _ = pl.plot(method="contour", lvls=None, suptitle=attribute) pl.img.savefig(tp.dir + attribute + ".png") @@ -51,7 +58,9 @@ def test_do(self, mem, ud): ref_mem = copy.deepcopy(mem) for attribute in tp.attributes: try: - ref_data = self.__get_ens(tc, tp, attribute, time_increment=False, summed=False) + ref_data = self.__get_ens( + tc, tp, attribute, time_increment=False, summed=False + ) if attribute != "p2_nodes": setattr(ref_mem.sol, attribute, ref_data) else: @@ -63,12 +72,26 @@ def test_do(self, mem, ud): tc_test.base_fn = tc_test.base_fn.replace("target", "test") tc_test.base_fn = tc_test.base_fn.replace("_stripped.h5", ".h5") tc_test.py_dir = tc_test.py_dir.replace("target", "test") - data = self.__get_ens(tc_test, tp, attribute, time_increment=self.time_increment, summed=False) + data = self.__get_ens( + tc_test, + tp, + attribute, + time_increment=self.time_increment, + summed=False, + ) setattr(mem.mpv, attribute, data) - ref_data = self.__get_ens(tc, tp, attribute, time_increment=self.time_increment, summed=False) + ref_data = self.__get_ens( + tc, + tp, + attribute, + time_increment=self.time_increment, + summed=False, + ) setattr(ref_mem.mpv, attribute, ref_data) except Exception as e: - raise AssertionError(f"test {self.current_run} has no target for comparison: {e}") + raise AssertionError( + f"test {self.current_run} has no target for comparison: {e}" + ) if self.plot: self.__plot_comparison(mem, ref_mem, ud) @@ -113,7 +136,7 @@ def test_do(self, mem, ud): Test passed for {self.current_run} {'#' * 10} """.strip() - ) + ) def __init(self, ds: DiagnosticState): tp = test_params(ds) @@ -178,9 +201,11 @@ def __get_sol_for_comparison(mem, ud, attribute): return test_sol @staticmethod - def __get_ens(tc, params, attribute, time_increment=False, summed=True, normed=False): + def __get_ens( + tc, params, attribute, time_increment=False, summed=True, normed=False + ): if time_increment and attribute == "p2_nodes": - times = [params.times[0]-1, params.times[0]] + times = [params.times[0] - 1, params.times[0]] else: times = params.times l_typ = params.l_typ @@ -205,8 +230,8 @@ def __get_ens(tc, params, attribute, time_increment=False, summed=True, normed=F if time_increment and attribute == "p2_nodes": ens = ens[1] - ens[0] else: - ens = ens[0] # removes time axis - ens = ens[0] # removes ensemble axis + ens = ens[0] # removes time axis + ens = ens[0] # removes ensemble axis if summed: return ens.sum() diff --git a/src/pybella/tests/test_blending_warm_bubble.py b/src/pybella/tests/test_blending_warm_bubble.py index 67022857..104da5c9 100644 --- a/src/pybella/tests/test_blending_warm_bubble.py +++ b/src/pybella/tests/test_blending_warm_bubble.py @@ -65,7 +65,7 @@ def __init__(self): "rhoY": 1.0e-0, "rhoX": 1.0e-0, "p2_nodes": 1.0e-4, - } + }, ) self.autogen_fn = False @@ -90,9 +90,7 @@ def sol_init(Sol, mpv, elem, node, th, ud, seed=None): r = np.sqrt((x) ** 2 + (y - y0) ** 2) / r0 p = np.repeat(mpv.HydroState.p0.reshape(1, -1), elem.icx, axis=0) - rhoY = mpv.HydroState.rhoY0[ - np.newaxis, : - ] + rhoY = mpv.HydroState.rhoY0[np.newaxis, :] perturbation = (delth / 300.0) * (np.cos(0.5 * np.pi * r) ** 2) perturbation[np.where(r > 1.0)] = 0.0 diff --git a/src/pybella/tests/test_internal_long_wave.py b/src/pybella/tests/test_internal_long_wave.py index 3f79e0af..5510d348 100644 --- a/src/pybella/tests/test_internal_long_wave.py +++ b/src/pybella/tests/test_internal_long_wave.py @@ -1,80 +1,45 @@ import numpy as np -from ..flow_solver.utils import options as opts, boundary as bdry, variable as var +from ..utils import options as opts + +from ..flow_solver.utils import boundary as bdry, variable as var from ..flow_solver.physics import hydrostatics from ..utils.data_structures import DiagnosticState class UserData(object): - NSPEC = 1 - BUOY = 0 - - grav = 9.81 - omega = 7.292 * 1e-5 # [s^{-1}] - - R_gas = 287.4 - R_vap = 461.0 - Q_vap = 2.53e06 - gamma = 1.4 - - h_ref = 10000.0 - t_ref = 100.0 - T_ref = 300.00 - p_ref = 1e5 - u_ref = h_ref / t_ref - rho_ref = p_ref / (R_gas * T_ref) - - Nsq_ref = 1.0e-4 - # planetary -> 160.0; long-wave -> 20.0; standard -> 1.0; scale_factor = 20.0 - i_gravity = np.zeros((3)) - i_coriolis = np.zeros((3)) - - tout = np.zeros((2)) - def __init__(self): self.scale_factor = self.scale_factor - self.h_ref = self.h_ref - self.t_ref = self.t_ref - self.T_ref = self.T_ref - self.p_ref = self.p_ref - self.rho_ref = self.rho_ref - self.u_ref = self.u_ref - self.Nsq_ref = self.Nsq_ref - self.g_ref = self.grav - self.gamm = self.gamma - self.Rg_over_Rv = self.R_gas / self.R_vap - self.Q = self.Q_vap / (self.R_gas * self.T_ref) - - self.nspec = self.NSPEC - - self.is_nonhydrostatic = 1 - self.is_compressible = 1 - self.is_ArakawaKonor = 0 - - self.compressibility = 0.0 - self.acoustic_timestep = 0 - self.Msq = self.u_ref * self.u_ref / (self.R_gas * self.T_ref) + self.h_ref = 10000.0 # [m] + self.t_ref = 100.0 # [s] + self.T_ref = 300.00 # [K] + self.p_ref = 1e5 # [Pa] + self.omega = 7.292 * 1e-5 # [s^{-1}] + self.grav = 9.81 # [m/s^2] + self.R_gas = 287.4 # [J kg^{-1} K^{-1}] + self.u_ref = self.h_ref / self.t_ref # [m/s] + self.Nsq_ref = 1.0e-4 # [s^{-2}] + self.Msq = ( + self.u_ref * self.u_ref / (self.R_gas * self.T_ref) + ) # Mach number squared self.gravity_strength = np.zeros((3)) - self.coriolis_strength = np.zeros((3)) self.gravity_strength[1] = self.grav * self.h_ref / (self.R_gas * self.T_ref) - self.coriolis_strength[0] = self.omega * self.t_ref - # self.coriolis_strength[2] = self.omega * self.t_ref - # self.coriolis_strength[1] = self.omega * self.t_ref - gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | (np.arange(3) == 1) + gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | ( + np.arange(3) == 1 + ) self.i_gravity = gravity_mask.astype(int) if np.any(gravity_mask): - self.gravity_direction = np.where(gravity_mask)[0][-1] # Use last matching index - - coriolis_mask = self.coriolis_strength > np.finfo(np.float64).eps - self.i_coriolis = coriolis_mask.astype(int) + self.gravity_direction = np.where(gravity_mask)[0][ + -1 + ] # Use last matching index self.xmin = -15.0 * self.scale_factor self.xmax = 15.0 * self.scale_factor @@ -83,10 +48,6 @@ def __init__(self): self.zmin = -1.0 self.zmax = 1.0 - self.u_wind_speed = 0.0 * 20.0 / self.u_ref - self.v_wind_speed = 0.0 - self.w_wind_speed = 0.0 - self.bdry_type = np.empty((3), dtype=object) self.bdry_type[0] = opts.BdryType.PERIODIC self.bdry_type[1] = opts.BdryType.WALL @@ -110,37 +71,15 @@ def __init__(self): self.dtfixed = 1.0 self.inx = 301 + 1 - # self.inx = 1205+1 self.iny = 10 + 1 - # self.iny = 40+1 self.inz = 1 - self.limiter_type_scalars = opts.LimiterType.NONE - self.limiter_type_velocity = opts.LimiterType.NONE - - self.initial_projection = False - - self.do_advection = True - self.tout = [self.scale_factor * 1.0 * 3000.0 / self.t_ref] self.tol = 1.0e-12 self.stepmax = 31 self.max_iterations = 6000 - self.continuous_blending = False - self.no_of_pi_initial = 0 - self.no_of_pi_transition = 0 - self.no_of_hy_initial = 1 - self.no_of_hy_transition = 0 - - self.blending_weight = 0.0 / 16 - self.blending_mean = "rhoY" # 1.0, rhoY - self.blending_conv = "rho" # theta, rho - self.blending_type = "half" # half, full - - self.initial_blending = False - self.autogen_fn = False self.output_timesteps = True diff --git a/src/pybella/tests/test_lamb_wave.py b/src/pybella/tests/test_lamb_wave.py index 9d885688..d474213d 100644 --- a/src/pybella/tests/test_lamb_wave.py +++ b/src/pybella/tests/test_lamb_wave.py @@ -1,6 +1,8 @@ import numpy as np -from ..flow_solver.utils import options as opts, boundary as bdry +from ..utils import options as opts + +from ..flow_solver.utils import boundary as bdry from ..flow_solver.physics import hydrostatics @@ -8,46 +10,20 @@ class UserData(object): - NSPEC = 1 - grav = 9.81 # [m s^{-2}] - omega = 0.0 * 1e-5 # [s^{-1}] - - R_gas = 287.4 # [J kg^{-1} K^{-1}] - R_vap = 461.0 - Q_vap = 2.53e06 - gamma = 1.4 - cp_gas = gamma * R_gas / (gamma - 1.0) - - p_ref = 1e5 - T_ref = 300.00 # [K] - rho_ref = p_ref / (R_gas * T_ref) - N_ref = grav / np.sqrt(cp_gas * T_ref) - Cs = np.sqrt(gamma * R_gas * T_ref) - - h_ref = 10.0e3 # [m] - t_ref = 100.0 # [s] - u_ref = h_ref / t_ref - - i_gravity = np.zeros((3)) - i_coriolis = np.zeros((3)) def __init__(self): - self.h_ref = self.h_ref - self.t_ref = self.t_ref - self.T_ref = self.T_ref - self.p_ref = self.p_ref - self.rho_ref = self.rho_ref - self.u_ref = self.u_ref - self.Nsq_ref = self.N_ref * self.N_ref - self.g_ref = self.grav - self.gamm = self.gamma - self.Rg_over_Rv = self.R_gas / self.R_vap - self.Q = self.Q_vap / (self.R_gas * self.T_ref) - self.R_gas = self.R_gas + self.grav = 9.81 + self.h_ref = 10.0e3 # [m] + self.t_ref = 100.0 # [s] + self.T_ref = 300.00 # [K] + self.p_ref = 1e5 + self.u_ref = self.h_ref / self.t_ref + self.R_gas = 287.4 + self.gamm = 1.4 + self.Cs = np.sqrt(self.gamm * self.R_gas * self.T_ref) + self.cp_gas = self.gamm * self.R_gas / (self.gamm - 1.0) + self.N_ref = self.grav / np.sqrt(self.cp_gas * self.T_ref) self.Rg = self.R_gas / (self.h_ref**2 / self.t_ref**2 / self.T_ref) - self.cp_gas = self.cp_gas - - self.nspec = self.NSPEC self.is_nonhydrostatic = 1 self.is_compressible = 1 @@ -61,15 +37,15 @@ def __init__(self): self.coriolis_strength = np.zeros((3)) self.gravity_strength[1] = self.grav * self.h_ref / (self.R_gas * self.T_ref) - self.coriolis_strength[2] = 2.0 * self.omega * self.t_ref - gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | (np.arange(3) == 1) + gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | ( + np.arange(3) == 1 + ) self.i_gravity = gravity_mask.astype(int) if np.any(gravity_mask): - self.gravity_direction = np.where(gravity_mask)[0][-1] # Use last matching index - - coriolis_mask = self.coriolis_strength > np.finfo(np.float64).eps - self.i_coriolis = coriolis_mask.astype(int) + self.gravity_direction = np.where(gravity_mask)[0][ + -1 + ] # Use last matching index j = 4.0 Lx = 1.0 * np.pi * self.Cs / self.N_ref * j @@ -80,14 +56,6 @@ def __init__(self): self.zmin = -1.0 self.zmax = 1.0 - self.u_wind_speed = 0.0 * self.u_ref - self.v_wind_speed = 0.0 - self.w_wind_speed = 0.0 - - self.bdry_type = np.empty((3), dtype=object) - self.bdry_type[0] = opts.BdryType.PERIODIC - self.bdry_type[1] = opts.BdryType.WALL - self.bdry_type[2] = opts.BdryType.WALL self.ATMOSPHERIC_EXTENSION = True self.rayleigh_bdry_switch = False @@ -103,29 +71,9 @@ def __init__(self): self.dtfixed0 = 10.0 / self.t_ref self.dtfixed = self.dtfixed0 - self.do_advection = True - self.limiter_type_scalars = opts.LimiterType.NONE - self.limiter_type_velocity = opts.LimiterType.NONE - self.tol = 1.0e-30 self.max_iterations = 10000 - # blending parameters - self.perturb_type = "pos_perturb" - self.blending_mean = "rhoY" # 1.0, rhoY - self.blending_conv = "rho" # theta, rho - self.blending_type = "half" # half, full - - self.continuous_blending = False - self.no_of_pi_initial = 1 - self.no_of_pi_transition = 0 - self.no_of_hy_initial = 0 - self.no_of_hy_transition = 0 - - self.blending_weight = 0.0 / 16 - self.initial_blending = False - self.initial_projection = True - self.tout = [360.0] # self.tout = np.arange(0,361,1.0) # self.tout = np.append(self.tout, [720.0]) @@ -153,11 +101,6 @@ def __init__(self): self.stratification = self.stratification_wrapper self.init_forcing = self.forcing - self.rayleigh_forcing = False - self.rayleigh_forcing_type = "func" # func or file - self.rayleigh_forcing_fn = None - self.rayleigh_forcing_path = None - def stratification_wrapper(self, dy): return lambda y: self.stratification_function(y, dy) @@ -294,11 +237,6 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): if ud.bdry_type[1] == opts.BdryType.RAYLEIGH: ud.tcy, ud.tny = bdry.get_tau_y(ud, elem, node, 0.5) - if ud.rayleigh_forcing: - ud.forcing_tcy, ud.forcing_tny = bdry.get_bottom_tau_y( - ud, elem, node, 0.2, cutoff=0.3 - ) - A0 = 1.0e-1 / ud.u_ref Msq = ud.Msq g = ud.gravity_strength[1] * ud.Rg diff --git a/src/pybella/tests/test_travelling_vortex.py b/src/pybella/tests/test_travelling_vortex.py index 0487ae18..08b86ebb 100644 --- a/src/pybella/tests/test_travelling_vortex.py +++ b/src/pybella/tests/test_travelling_vortex.py @@ -1,7 +1,10 @@ import numpy as np -from ..flow_solver.utils import options as opts, boundary as bdry + +from ..utils import options as opts +from ..flow_solver.utils import boundary as bdry from ..flow_solver.physics import hydrostatics from ..flow_solver.physics.low_mach import second_projection as lm_sp +from ..flow_solver.utils import variable as var from ..utils.data_structures import DiagnosticState @@ -9,68 +12,19 @@ class UserData(object): - NSPEC = 1 - grav = 0.0 - omega = 0.0 - - R_gas = 287.4 - R_vap = 461.0 - Q_vap = 2.53e06 - gamma = 1.4 h_ref = 10000.0 t_ref = 100.0 T_ref = 300.00 p_ref = 1e5 - u_ref = h_ref / t_ref - rho_ref = p_ref / (R_gas * T_ref) - - Nsq_ref = 0.0 - - i_gravity = np.zeros((3)) - i_coriolis = np.zeros((3)) - - tout = np.zeros((2)) def __init__(self): self.h_ref = self.h_ref self.t_ref = self.t_ref self.T_ref = self.T_ref self.p_ref = self.p_ref - self.rho_ref = self.rho_ref - self.u_ref = self.u_ref - self.Nsq_ref = self.Nsq_ref - self.g_ref = self.grav - self.gamm = self.gamma - self.Rg_over_Rv = self.R_gas / self.R_vap - self.Q = self.Q_vap / (self.R_gas * self.T_ref) - - self.nspec = self.NSPEC - - self.is_nonhydrostatic = 1 - self.is_compressible = 1 - self.is_ArakawaKonor = 0 - - self.compressibility = 1.0 - self.acoustic_timestep = 0 - self.acoustic_order = 0 - self.Msq = self.u_ref * self.u_ref / (self.R_gas * self.T_ref) - - self.gravity_strength = np.zeros((3)) - self.coriolis_strength = np.zeros((3)) - - self.gravity_strength[1] = self.grav * self.h_ref / (self.R_gas * self.T_ref) - self.coriolis_strength[0] = self.omega * self.t_ref - self.coriolis_strength[2] = self.omega * self.t_ref - - gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | (np.arange(3) == 1) - self.i_gravity = gravity_mask.astype(int) - if np.any(gravity_mask): - self.gravity_direction = np.where(gravity_mask)[0][-1] # Use last matching index - - coriolis_mask = self.coriolis_strength > np.finfo(np.float64).eps - self.i_coriolis = coriolis_mask.astype(int) + self.grav = self.grav self.xmin = -0.5 self.xmax = 0.5 @@ -92,9 +46,6 @@ def __init__(self): # NUMERICS ########################################## self.CFL = 0.9 / 2.0 - # self.CFL = 0.95 - # self.dtfixed0 = 2.1 * 1.200930e-2 - # self.dtfixed = 2.1 * 1.200930e-2 self.dtfixed = 0.01 self.dtfixed0 = 0.01 @@ -102,31 +53,7 @@ def __init__(self): self.iny = 64 + 1 self.inz = 1 - self.limiter_type_scalars = opts.LimiterType.NONE - self.limiter_type_velocity = opts.LimiterType.NONE - - self.tol = 1.0e-8 - self.max_iterations = 6000 - - self.perturb_type = "pos_perturb" - self.blending_mean = "rhoY" # 1.0, rhoY - self.blending_conv = "rho" # theta, rho - self.blending_type = "half" # half, full - - self.do_advection = True - - self.continuous_blending = False - self.no_of_pi_initial = 1 - self.no_of_pi_transition = 0 - self.no_of_hy_initial = 0 - self.no_of_hy_transition = 0 - - self.blending_weight = 0.0 / 16 - - self.initial_blending = False - self.initial_projection = True - self.initial_impl_Euler = False self.tout = [1.0] # np.arange(0.0,10.1,0.1)[1:] self.stepmax = 101 @@ -374,8 +301,19 @@ def sol_init(Sol, mpv, elem, node, th, ud, seed=None): Sol.rhou -= u0 * Sol.rho Sol.rhov -= v0 * Sol.rho + mem = obj() + mem.sol = Sol + mem.mpv = mpv + mem.elem = elem + mem.node = node + mem.th = th + mem.time = obj() + mem.time.t = ud.dtfixed + mem.time.step = 0 + mem.cache = var.FlowSolverCache() + lm_sp.euler_backward_non_advective_impl_part( - Sol, mpv, elem, node, ud, th, 0.0, ud.dtfixed, 0.5 + Sol, mpv, elem, node, ud, th, 0.0, ud.dtfixed, mem ) mpv.p2_nodes[...] = p2aux @@ -392,3 +330,7 @@ def sol_init(Sol, mpv, elem, node, th, ud, seed=None): def T_from_p_rho(p, rho): return np.divide(p, rho) + + +class obj(object): + pass \ No newline at end of file diff --git a/src/pybella/tests/test_unstable_lamb.py b/src/pybella/tests/test_unstable_lamb.py index 34fe949a..54dfb078 100644 --- a/src/pybella/tests/test_unstable_lamb.py +++ b/src/pybella/tests/test_unstable_lamb.py @@ -5,7 +5,7 @@ import numpy as np from ..flow_solver.physics import hydrostatics from ..flow_solver.utils import boundary as bdry -from ..flow_solver.utils import options as opts +from ..utils import options as opts from ..utils.data_structures import DiagnosticState @@ -13,34 +13,27 @@ class UserData(object): def __init__(self): self.grav = 9.81 # [m/s^2] - self.omega = 7.292 * 1e-5 # [s^{-1}] - self.t_ref = 100.0 # [s] - - self.h_ref = 10.0e3 # [m] - self.u_ref = self.h_ref / self.t_ref + self.omega = 2.0 * 7.292 * 1e-5 # [s^{-1}] - self.T_ref = 300.00 # [K] - self.R_gas = 287.4 # [J kg^{-1} K^{-1}] + self.h_ref = 10.0e3 # [m] + self.t_ref = 100.0 # [s] + self.T_ref = 300.00 # [K] + self.R_gas = 287.4 # [J kg^{-1} K^{-1}] self.Rg = self.R_gas / (self.h_ref**2 / self.t_ref**2 / self.T_ref) self.gamma = 1.4 - self.cp_gas = self.gamma * self.R_gas / (self.gamma-1.0) - - self.Nsq = (self.grav / np.sqrt(self.cp_gas * self.T_ref))**2 - self.Msq = self.u_ref * self.u_ref / (self.R_gas * self.T_ref) self.gravity_strength = np.zeros((3)) - self.coriolis_strength = np.zeros((3)) self.gravity_strength[1] = self.grav * self.h_ref / (self.R_gas * self.T_ref) - self.coriolis_strength[2] = 2.0 * self.omega * self.t_ref - gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | (np.arange(3) == 1) + gravity_mask = (self.gravity_strength > np.finfo(np.float64).eps) | ( + np.arange(3) == 1 + ) self.i_gravity = gravity_mask.astype(int) if np.any(gravity_mask): - self.gravity_direction = np.where(gravity_mask)[0][-1] # Use last matching index - - coriolis_mask = self.coriolis_strength > np.finfo(np.float64).eps - self.i_coriolis = coriolis_mask.astype(int) + self.gravity_direction = np.where(gravity_mask)[0][ + -1 + ] # Use last matching index ########################################## # SPATIAL GRID @@ -55,10 +48,6 @@ def __init__(self): # BOUNDARY CONDITIONS ########################################## - self.bdry_type = np.empty((3), dtype=object) - self.bdry_type[0] = opts.BdryType.PERIODIC - self.bdry_type[1] = opts.BdryType.WALL - self.bdry_type[2] = opts.BdryType.WALL self.ATMOSPHERIC_EXTENSION = True ########################################## @@ -71,35 +60,6 @@ def __init__(self): self.tout = [36.0] self.stepmax = 11 - self.is_compressible = 1 - self.is_nonhydrostatic = 1 - self.is_ArakawaKonor = 0 - - self.compressibility = 1.0 - self.acoustic_timestep = 0 - - ########################################## - # PHYSICS AND BACKGROUND WIND - ########################################## - self.u_wind_speed = 0.0 - self.v_wind_speed = 0.0 - self.w_wind_speed = 0.0 - - ########################################## - # BLENDING - ########################################## - self.continuous_blending = False - self.no_of_pi_initial = 1 - self.no_of_pi_transition = 0 - self.no_of_hy_initial = 0 - self.no_of_hy_transition = 0 - - self.blending_weight = 0.0 / 16 - self.blending_mean = "rhoY" - self.blending_conv = "rho" - self.blending_type = "half" - self.initial_blending = False - ########################################## # STRATIFICATION ########################################## @@ -147,7 +107,7 @@ def __init__(self): "rhoY": tol, "rhoX": tol, "p2_nodes": tol, - } + }, ) self.autogen_fn = False @@ -177,7 +137,7 @@ def _setup_domain(self): self.zmax = 1.0 def stratification_wrapper(self, dy): - return lambda y : self.stratification_function(y, dy) + return lambda y: self.stratification_function(y, dy) def stratification_function(self, y, dy): g = self.gravity_strength[1] @@ -187,22 +147,40 @@ def stratification_function(self, y, dy): pi_p = np.exp(-(y + 0.5 * dy) / Hex) pi_m = np.exp(-(y - 0.5 * dy) / Hex) - Theta = - (Gamma * g * dy) / (pi_p - pi_m) + Theta = -(Gamma * g * dy) / (pi_p - pi_m) return Theta - + @staticmethod def rayleigh_bc_function(ud): if ud.bdry_type[1] == opts.BdryType.RAYLEIGH or ud.rayleigh_forcing == True: ud.inbcy = ud.iny - 1 ud.iny0 = np.copy(ud.iny) - ud.iny = ud.iny0 + int(3*ud.inbcy) + ud.iny = ud.iny0 + int(3 * ud.inbcy) # tentative workaround ud.bcy = ud.ymax ud.ymax += 3.0 * ud.bcy class forcing(object): - def __init__(self, k, mu, Cs, F, N, Gamma, ampl, g, rhobar, Ybar, rhobar_n, Ybar_n, X, Y, Xn, Yn): + def __init__( + self, + k, + mu, + Cs, + F, + N, + Gamma, + ampl, + g, + rhobar, + Ybar, + rhobar_n, + Ybar_n, + X, + Y, + Xn, + Yn, + ): self.k = k self.mu = mu self.Cs = Cs @@ -226,41 +204,52 @@ def __init__(self, k, mu, Cs, F, N, Gamma, ampl, g, rhobar, Ybar, rhobar_n, Ybar def get_T_matrix(self): # system matrix of linearized equations - matrix = -np.array([[0, self.F, 0, 1j*self.Cs*self.k], - [-self.F, 0, -self.N, self.Cs*(self.mu+self.Gamma)], - [0, self.N, 0, 0], - [1j*self.Cs*self.k, self.Cs*(self.mu-self.Gamma), 0, 0]]) - + matrix = -np.array( + [ + [0, self.F, 0, 1j * self.Cs * self.k], + [-self.F, 0, -self.N, self.Cs * (self.mu + self.Gamma)], + [0, self.N, 0, 0], + [1j * self.Cs * self.k, self.Cs * (self.mu - self.Gamma), 0, 0], + ] + ) + self.T_matrix = matrix - def eigenfunction(self, t, s, grid='c'): - if grid == 'c': + def eigenfunction(self, t, s, grid="c"): + if grid == "c": x, z = self.X, self.Y - elif grid == 'n': - x, z, = self.Xn, self.Yn - + elif grid == "n": + ( + x, + z, + ) = ( + self.Xn, + self.Yn, + ) + # Compute eigenvalues and eigenvectors - eigval, eigvec = np.linalg.eig( self.T_matrix ) + eigval, eigvec = np.linalg.eig(self.T_matrix) - # Find index of eigenvalue + # Find index of eigenvalue # with greatest real part aka the instability growth rate - ind = np.argmax( np.real( eigval ) ) + ind = np.argmax(np.real(eigval)) # construct solution according to eq. 2.27 and 2.19 - exponentials = np.exp( 1j * self.k * x + self.mu * z - + ( eigval[ind] ) * (t) + 1j * s * t ) - chi_u = self.ampl * np.real( eigvec[0,ind] * exponentials ) - chi_w = self.ampl * np.real( eigvec[1,ind] * exponentials ) - chi_th = self.ampl * np.real( eigvec[2,ind] * exponentials ) - chi_pi = self.ampl * np.real( eigvec[3,ind] * exponentials ) - - self.arrs = ( chi_u, chi_w, chi_th, chi_pi ) - - def dehatter(self, th, grid='c'): - if grid == 'n': + exponentials = np.exp( + 1j * self.k * x + self.mu * z + (eigval[ind]) * (t) + 1j * s * t + ) + chi_u = self.ampl * np.real(eigvec[0, ind] * exponentials) + chi_w = self.ampl * np.real(eigvec[1, ind] * exponentials) + chi_th = self.ampl * np.real(eigvec[2, ind] * exponentials) + chi_pi = self.ampl * np.real(eigvec[3, ind] * exponentials) + + self.arrs = (chi_u, chi_w, chi_th, chi_pi) + + def dehatter(self, th, grid="c"): + if grid == "n": Ybar = self.Ybar_n oorhobarsqrt = self.oorhobarsqrt_n - elif grid == 'c': + elif grid == "c": Ybar = self.Ybar oorhobarsqrt = self.oorhobarsqrt @@ -270,12 +259,12 @@ def dehatter(self, th, grid='c'): vp = oorhobarsqrt * chi_v Yp = oorhobarsqrt * self.N / self.g * Ybar * chi_Y pi_p = oorhobarsqrt * self.Cs / Ybar / th.Gammainv * chi_pi - + return up.T, vp.T, Yp.T, pi_p.T def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): - if hasattr(ud, 'rayleigh_bdry_switch'): + if hasattr(ud, "rayleigh_bdry_switch"): if ud.rayleigh_bdry_switch: ud.bdry_type[1] = opts.BdryType.RAYLEIGH @@ -283,8 +272,9 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): ud.tcy, ud.tny = bdry.get_tau_y(ud, elem, node, 0.5) if ud.rayleigh_forcing: - ud.forcing_tcy, ud.forcing_tny = bdry.get_bottom_tau_y(ud, elem, node, 0.2, cutoff=0.3) - + ud.forcing_tcy, ud.forcing_tny = bdry.get_bottom_tau_y( + ud, elem, node, 0.2, cutoff=0.3 + ) A0 = 1.0e-1 / ud.u_ref Msq = ud.Msq @@ -317,7 +307,7 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): ################################################## # dimensionless Brunt-Väisälä frequency - N = ud.t_ref * np.sqrt(ud.Nsq) + N = ud.t_ref * np.sqrt(ud.Nsq_ref) # dimensionless speed of sound Cs = np.sqrt(th.gamm / Msq) ud.Cs = Cs @@ -327,11 +317,13 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): ud.coriolis_strength[2] += 1e-15 F = ud.coriolis_strength[2] - G = np.sqrt(9. / 40.) + G = np.sqrt(9.0 / 40.0) Gamma = G * N / Cs k = N / Cs - ud.rf_bot = ud.init_forcing(k, -Gamma, Cs, F, N, Gamma, A0, g, rhobar, Ybar, rhobar_n, Ybar_n, X, Y, Xn, Yn) + ud.rf_bot = ud.init_forcing( + k, -Gamma, Cs, F, N, Gamma, A0, g, rhobar, Ybar, rhobar_n, Ybar_n, X, Y, Xn, Yn + ) ud.rf_bot.get_T_matrix() ud.u_wind_speed = 0.0 @@ -357,18 +349,18 @@ def sol_init(Sol, mpv, elem, node, th, ud, seeds=None): ################################################### # initialise nodal pi - ud.rf_bot.eigenfunction(0, 1, grid='n') - _, _, _, pi_n = ud.rf_bot.dehatter(th, grid='n') + ud.rf_bot.eigenfunction(0, 1, grid="n") + _, _, _, pi_n = ud.rf_bot.dehatter(th, grid="n") mpv.p2_nodes[...] = pi_n bdry.set_explicit_boundary_data(Sol, elem, ud, th, mpv) - if hasattr(ud, 'mixed_run'): + if hasattr(ud, "mixed_run"): if ud.mixed_run: ud.coriolis_strength[2] = 2.0 * 7.292 * 1e-5 * ud.t_ref - if hasattr(ud, 'trad_forcing'): + if hasattr(ud, "trad_forcing"): if ud.trad_forcing: ud.rf_bot.F = 0.0 ud.rf_bot.get_T_matrix() diff --git a/src/pybella/utils/data_structures.py b/src/pybella/utils/data_structures.py index 15b2aef5..e00b9f0a 100644 --- a/src/pybella/utils/data_structures.py +++ b/src/pybella/utils/data_structures.py @@ -2,7 +2,7 @@ from dataclasses import dataclass, field, fields from ..flow_solver.discretisation.grid import Grid -from ..flow_solver.utils.variable import Vars +from ..flow_solver.utils.variable import Vars, FlowSolverCache from ..flow_solver.physics.low_mach.mpv import MPV from ..flow_solver.physics.gas_dynamics.thermodynamics import ThermodynamicalQuantities @@ -22,6 +22,7 @@ class ModelState: flux: List[Vars] mpv: MPV th: ThermodynamicalQuantities + cache: FlowSolverCache = field(default_factory=FlowSolverCache) time: IntegrationTime = field(init=False) def __post_init__(self): @@ -71,8 +72,11 @@ def update_member( mpv: MPV, flux: List[Vars], th: ThermodynamicalQuantities, + cache: Optional[FlowSolverCache] = None ): - new_state = ModelState(elem, node, sol, flux, mpv, th) + if cache is None: + cache = FlowSolverCache() + new_state = ModelState(elem, node, sol, flux, mpv, th, cache) self.members.append(new_state) def set_members(self, members: List[ModelState]): @@ -116,18 +120,20 @@ class DiagnosticState: # plot the comparison? plot_compare: bool = True - tolerances: dict[str, float] = field(default_factory=lambda: { - # for most tests, we want to ensure that the - # max absolute error is within singular precision - # limit. - "rhou": 1e-5, - "rhov": 1e-5, - "rhow": 1e-5, - "rhoY": 1e-5, - "rhoX": 1e-5, - "rho": 1e-5, - "p2_nodes": 1e-5 - }) + tolerances: dict[str, float] = field( + default_factory=lambda: { + # for most tests, we want to ensure that the + # max absolute error is within singular precision + # limit. + "rhou": 1e-5, + "rhov": 1e-5, + "rhow": 1e-5, + "rhoY": 1e-5, + "rhoX": 1e-5, + "rho": 1e-5, + "p2_nodes": 1e-5, + } + ) time_increment: bool = False diff --git a/src/pybella/utils/debug_helpers.py b/src/pybella/utils/debug_helpers.py index d7fc86d6..fbdf43ff 100644 --- a/src/pybella/utils/debug_helpers.py +++ b/src/pybella/utils/debug_helpers.py @@ -1,7 +1,8 @@ import matplotlib.pyplot as plt + def pl_sol(arr): plt.figure() plt.imshow(arr, origin="lower", aspect="auto") plt.colorbar() - plt.show() \ No newline at end of file + plt.show() diff --git a/src/pybella/utils/io.py b/src/pybella/utils/io.py index ec8038de..68cdeba7 100644 --- a/src/pybella/utils/io.py +++ b/src/pybella/utils/io.py @@ -567,9 +567,9 @@ def get_args(): if ic in IC_MODULES: module_name = IC_MODULES[ic] try: - module = __import__(module_name, fromlist=['UserData', 'sol_init']) - UserData = getattr(module, 'UserData') - sol_init = getattr(module, 'sol_init') + module = __import__(module_name, fromlist=["UserData", "sol_init"]) + UserData = getattr(module, "UserData") + sol_init = getattr(module, "sol_init") except ImportError as e: raise ImportError(f"Failed to import {module_name}: {e}") else: @@ -734,41 +734,43 @@ def init_logger(ud): logging.getLogger("numba").setLevel(logging.WARNING) logging.getLogger("numba.core").setLevel(logging.WARNING) - logging.getLogger().setLevel(logging.WARNING) + logging.getLogger().setLevel(logging.INFO) logging.info("Input file is %s" % input_filename) + class NullDebugWriter: """Null object that does nothing but implements the DebugWriter interface.""" - + def populate(self, label, field_name, data): """No-op populate method.""" pass - + def write(self, label): """No-op write method.""" pass - + def populate_flux_components(self, label, flux, elem): """No-op populate flux components method.""" pass + class DebugWriter: """Enhanced debug writer with populate functionality.""" - + def __init__(self, debug, writer, mem): self.debug = debug self.writer = writer self.mem = mem self._pending_populations = {} - + def populate(self, label, field_name, data): """Populate field data for a given label.""" if self.debug and self.writer is not None: if label not in self._pending_populations: self._pending_populations[label] = [] self._pending_populations[label].append((field_name, data)) - + def write(self, label): """Write all data including any pending populations.""" if self.debug and self.writer is not None: @@ -777,10 +779,10 @@ def write(self, label): for field_name, data in self._pending_populations[label]: self.writer.populate(label, field_name, data) del self._pending_populations[label] - + # Write the data self.writer.write_all(self.mem, label) - + def populate_flux_components(self, label, flux, elem): """Helper method to populate flux components.""" if self.debug and self.writer is not None: @@ -789,9 +791,10 @@ def populate_flux_components(self, label, flux, elem): if elem.ndim == 3: self.populate(label, "rhoYw", flux[2].rhoY) + def create_debug_writer(debug, writer, mem): """Factory function to create appropriate debug writer.""" if debug and writer is not None: return DebugWriter(debug, writer, mem) else: - return NullDebugWriter() \ No newline at end of file + return NullDebugWriter() diff --git a/src/pybella/utils/operators.py b/src/pybella/utils/operators.py new file mode 100644 index 00000000..37e9a0a3 --- /dev/null +++ b/src/pybella/utils/operators.py @@ -0,0 +1,650 @@ +import numpy as np +import scipy as sp +from numba import njit +from functools import lru_cache + +@lru_cache(maxsize=2) +def get_flux_convolution_kernels(ndim): + """Create convolution kernels for advective flux computation. + + Parameters + ---------- + ndim : int + Number of dimensions (2 or 3) + + Returns + ------- + dict + Dictionary containing kernels for each direction + + Notes + ----- + Results are cached since kernels don't change during computation. + """ + if ndim == 2: + kernel_u = np.array([[0.5, 1.0, 0.5], [0.5, 1.0, 0.5]]) + return { + 'u': kernel_u, + 'v': kernel_u.T + } + elif ndim == 3: + kernel_u = np.array([ + [[1, 2, 1], [2, 4, 2], [1, 2, 1]], + [[1, 2, 1], [2, 4, 2], [1, 2, 1]] + ]) + return { + 'u': kernel_u, + 'v': np.swapaxes(kernel_u, 1, 0), + 'w': np.swapaxes(kernel_u, 2, 0) + } + else: + raise ValueError(f"Unsupported dimension: {ndim}") + + +@lru_cache(maxsize=4) +def get_averaging_kernel(ndim, width=3, normalize=True): + """ + Create a generic averaging kernel for arbitrary dimensions. + + Parameters + ---------- + ndim : int + Number of dimensions (e.g., 2 or 3) + width : int, default=3 + Size of the kernel along each axis (can be even or odd) + normalize : bool, default=True + Whether to normalize the kernel to sum to 1 + + Returns + ------- + np.ndarray + Averaging kernel of shape (width,) * ndim + + Notes + ----- + - Odd widths result in centered kernels. + - Even widths are useful for staggered/grid-face averaging. + """ + shape = (width,) * ndim + kernel = np.ones(shape, dtype=np.float64) + + if normalize: + kernel /= kernel.size + + return kernel + +@njit(cache=True) +def _numba_convolve_2d(data, kernel): + """Numba-compiled 2D convolution for better performance.""" + data_h, data_w = data.shape + kernel_h, kernel_w = kernel.shape + + result_h = data_h - kernel_h + 1 + result_w = data_w - kernel_w + 1 + result = np.zeros((result_h, result_w)) + + for i in range(result_h): + for j in range(result_w): + for ki in range(kernel_h): + for kj in range(kernel_w): + result[i, j] += data[i + ki, j + kj] * kernel[ki, kj] + + return result + + +@njit(cache=True) +def _numba_convolve_3d(data, kernel): + """Numba-compiled 3D convolution for better performance.""" + data_d, data_h, data_w = data.shape + kernel_d, kernel_h, kernel_w = kernel.shape + + result_d = data_d - kernel_d + 1 + result_h = data_h - kernel_h + 1 + result_w = data_w - kernel_w + 1 + result = np.zeros((result_d, result_h, result_w)) + + for i in range(result_d): + for j in range(result_h): + for k in range(result_w): + for ki in range(kernel_d): + for kj in range(kernel_h): + for kk in range(kernel_w): + result[i, j, k] += data[i + ki, j + kj, k + kk] * kernel[ki, kj, kk] + + return result + + +def apply_convolution_kernel(data, kernel, normalize=True, axis_swap=None, use_numba=True): + """Apply convolution kernel with optional normalization and axis swapping. + + Parameters + ---------- + data : np.ndarray + Input data array + kernel : np.ndarray + Convolution kernel + normalize : bool, default=True + Whether to normalize by kernel sum + axis_swap : tuple or None, default=None + Tuple of (from_axis, to_axis) for np.moveaxis + use_numba : bool, default=True + Whether to use Numba-compiled convolution (faster for repeated calls) + + Returns + ------- + np.ndarray + Convolved result + + Notes + ----- + For large arrays or single calls, scipy.signal.fftconvolve might be faster. + For repeated calls on smaller arrays, Numba convolution is typically faster. + """ + if use_numba and data.ndim in (2, 3): + if data.ndim == 2: + result = _numba_convolve_2d(data, kernel) + else: # 3D + result = _numba_convolve_3d(data, kernel) + else: + # Fallback to scipy for other dimensions or when requested + result = sp.signal.fftconvolve(data, kernel, mode="valid") + + if normalize: + result = result / kernel.sum() + + if axis_swap is not None: + result = np.moveaxis(result, axis_swap[0], axis_swap[1]) + + return result + + +# Configuration for directional convolutions +DIRECTION_CONFIG = { + 2: { + 'u': {'axis_swap': (0, -1)}, + 'v': {'axis_swap': None} + }, + 3: { + 'u': {'axis_swap': (0, -1)}, + 'v': {'axis_swap': (-1, 0)}, + 'w': {'axis_swap': None} + } +} + +def apply_directional_convolution(data, kernel, direction, ndim, normalize=True, use_numba=True): + """Apply convolution kernel for a specific direction with appropriate axis swapping. + + Parameters + ---------- + data : np.ndarray + Input data array + kernel : np.ndarray + Convolution kernel for the direction + direction : str + Direction ('u', 'v', or 'w') + ndim : int + Number of dimensions (2 or 3) + normalize : bool, default=True + Whether to normalize by kernel sum + use_numba : bool, default=True + Whether to use Numba-compiled convolution + + Returns + ------- + np.ndarray + Convolved result with appropriate axis swapping + """ + if direction not in DIRECTION_CONFIG[ndim]: + raise ValueError(f"Direction '{direction}' not supported for {ndim}D") + + config = DIRECTION_CONFIG[ndim][direction] + return apply_convolution_kernel( + data, kernel, + normalize=normalize, + axis_swap=config['axis_swap'], + use_numba=use_numba + ) + + +@njit(cache=True) +def compute_divergence_2d(u_field, v_field, dx, dy): + """ + Compute 2D divergence: ∇·F = ∂u/∂x + ∂v/∂y + + Parameters + ---------- + u_field : np.ndarray + Field component in x-direction + v_field : np.ndarray + Field component in y-direction + dx : float + Grid spacing in x-direction + dy : float + Grid spacing in y-direction + + Returns + ------- + np.ndarray + Divergence field averaged to cell centers + """ + # X-direction: ∂u/∂x + div_x = finite_difference_1d(u_field, dx, axis=0) + # Average to y-cell centers + div_x = 0.5 * (div_x[:, :-1] + div_x[:, 1:]) + + # Y-direction: ∂v/∂y + div_y = finite_difference_1d(v_field, dy, axis=1) + # Average to x-cell centers + div_y = 0.5 * (div_y[:-1, :] + div_y[1:, :]) + + return div_x + div_y + + +@njit(cache=True) +def compute_divergence_3d(u_field, v_field, w_field, dx, dy, dz): + """ + Compute 3D divergence: ∇·F = ∂u/∂x + ∂v/∂y + ∂w/∂z + + Parameters + ---------- + u_field : np.ndarray + Field component in x-direction + v_field : np.ndarray + Field component in y-direction + w_field : np.ndarray + Field component in z-direction + dx : float + Grid spacing in x-direction + dy : float + Grid spacing in y-direction + dz : float + Grid spacing in z-direction + + Returns + ------- + tuple + (div_x, div_y, div_z) - Individual divergence components + """ + # X-direction: ∂u/∂x + div_x = finite_difference_1d(u_field, dx, axis=0) + # Average to y-cell centers, then to z-faces + div_x = 0.5 * (div_x[:, :-1, :] + div_x[:, 1:, :]) + div_x = -0.5 * (div_x[:, :, :-1] + div_x[:, :, 1:]) # Note: negative from original + + # Y-direction: ∂v/∂y + div_y = finite_difference_1d(v_field, dy, axis=1) + # Average to x-cell centers, then to z-faces + div_y = 0.5 * (div_y[:-1, :, :] + div_y[1:, :, :]) + div_y = 0.5 * (div_y[:, :, :-1] + div_y[:, :, 1:]) + + # Z-direction: ∂w/∂z + div_z = finite_difference_1d(w_field, dz, axis=2) + # Average to cell centers + div_z = 0.5 * (div_z[:-1, :, :] + div_z[1:, :, :]) + div_z = 0.5 * (div_z[:, :-1, :] + div_z[:, 1:, :]) + + return div_x, div_y, div_z + + +@njit(cache=True) +def compute_divergence_3d_total(u_field, v_field, w_field, dx, dy, dz): + """ + Compute total 3D divergence: ∇·F = ∂u/∂x + ∂v/∂y + ∂w/∂z + + Parameters + ---------- + u_field : np.ndarray + Field component in x-direction + v_field : np.ndarray + Field component in y-direction + w_field : np.ndarray + Field component in z-direction + dx : float + Grid spacing in x-direction + dy : float + Grid spacing in y-direction + dz : float + Grid spacing in z-direction + + Returns + ------- + np.ndarray + Total divergence field + """ + div_x, div_y, div_z = compute_divergence_3d(u_field, v_field, w_field, dx, dy, dz) + return div_x + div_y + div_z + + +@njit(cache=True) +def finite_difference_1d(field, spacing, axis=0): + """ + Compute 1D finite difference along specified axis. + + Parameters + ---------- + field : np.ndarray + Input field + spacing : float + Grid spacing + axis : int, default=0 + Axis along which to compute difference + + Returns + ------- + np.ndarray + Finite difference result + """ + if field.ndim == 2: + if axis == 0: + return (field[1:, :] - field[:-1, :]) / spacing + elif axis == 1: + return (field[:, 1:] - field[:, :-1]) / spacing + else: + raise ValueError("axis must be 0 or 1 for 2D arrays") + elif field.ndim == 3: + if axis == 0: + return (field[1:, :, :] - field[:-1, :, :]) / spacing + elif axis == 1: + return (field[:, 1:, :] - field[:, :-1, :]) / spacing + elif axis == 2: + return (field[:, :, 1:] - field[:, :, :-1]) / spacing + else: + raise ValueError("axis must be 0, 1, or 2 for 3D arrays") + else: + raise ValueError("field must be 2D or 3D array") + +# @njit +def average_to_centers_2d(field, axis): + """ + Average field values to cell centers along specified axis. + + Parameters + ---------- + field : np.ndarray + Input field (2D) + axis : int + Axis along which to average (0 or 1) + + Returns + ------- + np.ndarray + Averaged field + """ + if axis == 0: + return 0.5 * (field[:-1, :] + field[1:, :]) + elif axis == 1: + return 0.5 * (field[:, :-1] + field[:, 1:]) + else: + raise ValueError("axis must be 0 or 1 for 2D arrays") + + +# @njit +def average_to_centers_3d(field, axis): + """ + Average field values to cell centers along specified axis. + + Parameters + ---------- + field : np.ndarray + Input field (3D) + axis : int + Axis along which to average (0, 1, or 2) + + Returns + ------- + np.ndarray + Averaged field + """ + if axis == 0: + return 0.5 * (field[:-1, :, :] + field[1:, :, :]) + elif axis == 1: + return 0.5 * (field[:, :-1, :] + field[:, 1:, :]) + elif axis == 2: + return 0.5 * (field[:, :, :-1] + field[:, :, 1:]) + else: + raise ValueError("axis must be 0, 1, or 2 for 3D arrays") + + +@njit(cache=True) +def compute_gradient_2d(field, dx, dy): + """ + Compute 2D gradient: ∇φ = (∂φ/∂x, ∂φ/∂y) + + Parameters + ---------- + field : np.ndarray + Scalar field + dx : float + Grid spacing in x-direction + dy : float + Grid spacing in y-direction + + Returns + ------- + tuple + (grad_x, grad_y) - Gradient components + """ + grad_x = finite_difference_1d(field, dx, axis=0) + grad_y = finite_difference_1d(field, dy, axis=1) + + return grad_x, grad_y + + +@njit(cache=True) +def compute_gradient_3d(field, dx, dy, dz): + """ + Compute 3D gradient: ∇φ = (∂φ/∂x, ∂φ/∂y, ∂φ/∂z) + + Parameters + ---------- + field : np.ndarray + Scalar field + dx : float + Grid spacing in x-direction + dy : float + Grid spacing in y-direction + dz : float + Grid spacing in z-direction + + Returns + ------- + tuple + (grad_x, grad_y, grad_z) - Gradient components + """ + grad_x = finite_difference_1d(field, dx, axis=0) + grad_y = finite_difference_1d(field, dy, axis=1) + grad_z = finite_difference_1d(field, dz, axis=2) + + return grad_x, grad_y, grad_z + + +####### +# Compute gradient at nodes +####### + +@njit(cache=True) +def _compute_grad_nodes_2d(p, dx, dy): + """Compute 2D gradient at nodes using corner averaging.""" + # Pre-computed signs for each corner + signs_x = np.array([-1.0, -1.0, +1.0, +1.0]) + signs_y = np.array([-1.0, +1.0, -1.0, +1.0]) + + # Initialize gradient components + Dpx = np.zeros((p.shape[0] - 1, p.shape[1] - 1)) + Dpy = np.zeros((p.shape[0] - 1, p.shape[1] - 1)) + + # Corner contributions + # Bottom-left + Dpx += signs_x[0] * p[0:-1, 0:-1] + Dpy += signs_y[0] * p[0:-1, 0:-1] + + # Bottom-right + Dpx += signs_x[1] * p[0:-1, 1:] + Dpy += signs_y[1] * p[0:-1, 1:] + + # Top-left + Dpx += signs_x[2] * p[1:, 0:-1] + Dpy += signs_y[2] * p[1:, 0:-1] + + # Top-right + Dpx += signs_x[3] * p[1:, 1:] + Dpy += signs_y[3] * p[1:, 1:] + + # Apply scaling factors + scale_factor = 0.5 ** (2 - 1) # 0.5^(ndim-1) + Dpx *= scale_factor / dx + Dpy *= scale_factor / dy + + return Dpx, Dpy + +@njit(cache=True) +def _compute_grad_nodes_3d(p, dx, dy, dz): + """Compute 3D gradient at nodes using corner averaging.""" + # Pre-computed signs for each corner + signs_x = np.array([-1.0, -1.0, -1.0, -1.0, +1.0, +1.0, +1.0, +1.0]) + signs_y = np.array([-1.0, -1.0, +1.0, +1.0, -1.0, -1.0, +1.0, +1.0]) + signs_z = np.array([-1.0, +1.0, -1.0, +1.0, -1.0, +1.0, -1.0, +1.0]) + + # Initialize gradient components + Dpx = np.zeros((p.shape[0] - 1, p.shape[1] - 1, p.shape[2] - 1)) + Dpy = np.zeros((p.shape[0] - 1, p.shape[1] - 1, p.shape[2] - 1)) + Dpz = np.zeros((p.shape[0] - 1, p.shape[1] - 1, p.shape[2] - 1)) + + # Corner contributions (8 corners for 3D) + # Bottom-left-back + Dpx += signs_x[0] * p[0:-1, 0:-1, 0:-1] + Dpy += signs_y[0] * p[0:-1, 0:-1, 0:-1] + Dpz += signs_z[0] * p[0:-1, 0:-1, 0:-1] + + # Bottom-left-front + Dpx += signs_x[1] * p[0:-1, 0:-1, 1:] + Dpy += signs_y[1] * p[0:-1, 0:-1, 1:] + Dpz += signs_z[1] * p[0:-1, 0:-1, 1:] + + # Bottom-right-back + Dpx += signs_x[2] * p[0:-1, 1:, 0:-1] + Dpy += signs_y[2] * p[0:-1, 1:, 0:-1] + Dpz += signs_z[2] * p[0:-1, 1:, 0:-1] + + # Bottom-right-front + Dpx += signs_x[3] * p[0:-1, 1:, 1:] + Dpy += signs_y[3] * p[0:-1, 1:, 1:] + Dpz += signs_z[3] * p[0:-1, 1:, 1:] + + # Top-left-back + Dpx += signs_x[4] * p[1:, 0:-1, 0:-1] + Dpy += signs_y[4] * p[1:, 0:-1, 0:-1] + Dpz += signs_z[4] * p[1:, 0:-1, 0:-1] + + # Top-left-front + Dpx += signs_x[5] * p[1:, 0:-1, 1:] + Dpy += signs_y[5] * p[1:, 0:-1, 1:] + Dpz += signs_z[5] * p[1:, 0:-1, 1:] + + # Top-right-back + Dpx += signs_x[6] * p[1:, 1:, 0:-1] + Dpy += signs_y[6] * p[1:, 1:, 0:-1] + Dpz += signs_z[6] * p[1:, 1:, 0:-1] + + # Top-right-front + Dpx += signs_x[7] * p[1:, 1:, 1:] + Dpy += signs_y[7] * p[1:, 1:, 1:] + Dpz += signs_z[7] * p[1:, 1:, 1:] + + # Apply scaling factors + scale_factor = 0.5 ** (3 - 1) # 0.5^(ndim-1) + Dpx *= scale_factor / dx + Dpy *= scale_factor / dy + Dpz *= scale_factor / dz + + return Dpx, Dpy, Dpz + +@njit(cache=True) +def compute_gradient_nodes_2d(p, dx, dy): + """ + Compute gradient at nodes using finite difference averaging. + + Parameters + ---------- + p : np.ndarray + Scalar field (2D) + dx : float + Grid spacing in x-direction + dy : float + Grid spacing in y-direction + + Returns + ------- + tuple + (grad_x, grad_y) - Gradient components at nodes + + Notes + ----- + The gradient is computed at nodes by averaging contributions from + all neighboring cells. Output arrays have shape (nx-1, ny-1). + """ + return _compute_grad_nodes_2d(p, dx, dy) + +@njit(cache=True) +def compute_gradient_nodes_3d(p, dx, dy, dz): + """ + Compute gradient at nodes using finite difference averaging. + + Parameters + ---------- + p : np.ndarray + Scalar field (3D) + dx : float + Grid spacing in x-direction + dy : float + Grid spacing in y-direction + dz : float + Grid spacing in z-direction + + Returns + ------- + tuple + (grad_x, grad_y, grad_z) - Gradient components at nodes + + Notes + ----- + The gradient is computed at nodes by averaging contributions from + all neighboring cells. Output arrays have shape (nx-1, ny-1, nz-1). + """ + return _compute_grad_nodes_3d(p, dx, dy, dz) + +def compute_gradient_nodes(p, ndim, dxy): + """ + Compute gradient at nodes using finite difference averaging. + + Parameters + ---------- + p : np.ndarray + Scalar field + ndim : int + Number of dimensions (2 or 3) + dxy : tuple + Grid spacings (dx, dy, dz) + + Returns + ------- + tuple + Gradient components at nodes. For 2D: (grad_x, grad_y, zeros) + For 3D: (grad_x, grad_y, grad_z) + + Notes + ----- + This is the main interface function that dispatches to the appropriate + dimension-specific implementation. Always returns 3 components for + consistency, with the z-component being zero for 2D cases. + """ + dx, dy, dz = dxy + + if ndim == 2: + grad_x, grad_y = compute_gradient_nodes_2d(p, dx, dy) + grad_z = np.zeros_like(grad_x) + return grad_x, grad_y, grad_z + elif ndim == 3: + return compute_gradient_nodes_3d(p, dx, dy, dz) + else: + raise ValueError(f"Unsupported dimension: {ndim}") \ No newline at end of file diff --git a/src/pybella/utils/options.py b/src/pybella/utils/options.py new file mode 100644 index 00000000..a28c065a --- /dev/null +++ b/src/pybella/utils/options.py @@ -0,0 +1,24 @@ +from enum import Enum # ! Version > Python 3.4 + + +class LimiterType(Enum): + NONE = 0 + # MINMOD = 1 + # VANLEER = 2 + # VANLEERSmooth = 3 + # SUPERBEE = 4 + # MONOTONIZED_CENTRAL = 5 + # SWEBY_MUNZ = 6 + # RUPE = 7 + # NO_SLOPE = 8 + # NUMBER_OF_LIMITER = 9 + + +class BdryType(Enum): + """ + An enumeration class that defines the accepted boundary condition types. + """ + + WALL = "symmetric" + PERIODIC = "wrap" + RAYLEIGH = "radiation" diff --git a/src/pybella/utils/slices.py b/src/pybella/utils/slices.py new file mode 100644 index 00000000..b53a51bc --- /dev/null +++ b/src/pybella/utils/slices.py @@ -0,0 +1,163 @@ + +def get_neighbor_indices(ndim): + """Create left and right neighbor indices for n-dimensional arrays.""" + lefts_idx = [slice(None)] * ndim + rights_idx = [slice(None)] * ndim + + lefts_idx[-1] = slice(0, -1) + rights_idx[-1] = slice(1, None) + + return tuple(lefts_idx), tuple(rights_idx) + + +def get_inner_slice(ndim): + """Get slice tuple for inner cells (excluding outermost ghost cells).""" + return tuple([slice(1, -1)] * ndim) + +def get_last_dim_inner_slice(ndim): + """Get slice for inner faces in the last dimension only.""" + idx = [slice(None)] * ndim + idx[-1] = slice(1, -1) + return tuple(idx) + + +def get_interface_indices(ndim): + """ + Get complete set of indices for interface calculations. + + Parameters + ---------- + ndim : int + Number of dimensions + + Returns + ------- + tuple + (lefts_idx, rights_idx, inner_idx) where: + - lefts_idx: indices for left neighbors + - rights_idx: indices for right neighbors + - inner_idx: indices for inner cells (face_inner_idx) + """ + lefts_idx, rights_idx = get_neighbor_indices(ndim) + inner_idx = get_inner_slice(ndim) + + return lefts_idx, rights_idx, inner_idx + + +def get_all_slice_indices(ndim): + """ + Get all commonly used slice indices for finite volume calculations. + + Parameters + ---------- + ndim : int + Number of dimensions + + Returns + ------- + dict + Dictionary containing all slice indices: + - 'lefts': left neighbor indices + - 'rights': right neighbor indices + - 'inner': inner cell indices + - 'face_inner': face inner indices (alias for inner) + """ + lefts_idx, rights_idx, inner_idx = get_interface_indices(ndim) + + return { + 'lefts': lefts_idx, + 'rights': rights_idx, + 'inner': inner_idx, + 'face_inner': inner_idx # alias for clarity in interface flux calculations + } + + + +def get_averaging_indices(ndim, axis): + """ + Get indices for averaging along a specific axis. + + Parameters + ---------- + ndim : int + Number of dimensions + axis : int + Axis along which to average + + Returns + ------- + tuple + (left_idx, right_idx) for averaging operation + """ + left_idx = [slice(None)] * ndim + right_idx = [slice(None)] * ndim + + left_idx[axis] = slice(0, -1) + right_idx[axis] = slice(1, None) + + return tuple(left_idx), tuple(right_idx) + + +def get_periodic_inner_slice(igs, is_periodic, ndim): + """ + Get inner slice accounting for periodic boundary conditions. + + Parameters + ---------- + igs : array-like + Number of ghost cells in each dimension + is_periodic : array-like + Boolean array indicating periodic boundaries + ndim : int + Number of dimensions + + Returns + ------- + tuple + Slice tuple for inner region with periodic boundaries + """ + inner_idx = [] + for dim in range(ndim): + start = igs[dim] - int(is_periodic[dim]) + end = -igs[dim] + int(is_periodic[dim]) if igs[dim] > int(is_periodic[dim]) else None + inner_idx.append(slice(start, end)) + + return tuple(inner_idx) + + +def get_face_center_averaging_indices(ndim): + """ + Get indices for averaging finite difference results to face centers. + + Parameters + ---------- + ndim : int + Number of dimensions + + Returns + ------- + dict + Dictionary with 'y_avg' and 'x_avg' index tuples for 2D, + or 'y_avg', 'x_avg', 'z_avg' for 3D averaging operations + """ + indices = {} + + if ndim >= 2: + # For averaging in y-direction (axis=1) + indices['y_avg'] = (slice(None, -1), slice(None)) + indices['y_avg_right'] = (slice(1, None), slice(None)) + + # For averaging in x-direction (axis=0) + indices['x_avg'] = (slice(None), slice(None, -1)) + indices['x_avg_right'] = (slice(None), slice(1, None)) + + if ndim == 3: + # For averaging in z-direction (axis=2) + indices['z_avg'] = (slice(None), slice(None), slice(None, -1)) + indices['z_avg_right'] = (slice(None), slice(None), slice(1, None)) + + # 3D specific averaging combinations + indices['xy_avg'] = (slice(None, -1), slice(None, -1), slice(None)) + indices['xy_avg_right'] = (slice(1, None), slice(1, None), slice(None)) + + return indices \ No newline at end of file diff --git a/src/pybella/utils/user_data.py b/src/pybella/utils/user_data.py index 4aea5e28..66af313f 100644 --- a/src/pybella/utils/user_data.py +++ b/src/pybella/utils/user_data.py @@ -1,28 +1,102 @@ import numpy as np - -from ..flow_solver.utils import options as opts -from . import sim_params as params - - -class UserDataInit(object): +from collections import defaultdict + +from . import sim_params +from . import options as opts + + +class DependencyManager: + """Manages computational dependencies between attributes.""" + + def __init__(self): + # dependency_graph[computed_attr] = [list of required attributes] + self.dependency_graph = { + "u_ref": ["h_ref", "t_ref"], + "Msq": ["u_ref", "R_gas", "T_ref"], + "gravity_strength": ["grav", "h_ref", "R_gas", "T_ref"], + "i_gravity": ["grav", "h_ref", "R_gas", "T_ref"], + "coriolis_strength": ["omega", "t_ref"], + "cp_gas": ["gamm", "R_gas"], + "N_ref": ["grav", "cp_gas", "T_ref"], + "Nsq_ref": ["grav", "cp_gas", "T_ref"], + "rho_ref": ["p_ref", "R_gas", "T_ref"], + "Cs": ["gamm", "R_gas", "T_ref"], + } + + # Reverse mapping: which computations depend on each attribute + self.reverse_deps = defaultdict(list) + for computed, deps in self.dependency_graph.items(): + for dep in deps: + self.reverse_deps[dep].append(computed) + + def get_computations_to_update(self, changed_attr): + """Get list of computations that need to be updated when an attribute changes.""" + return self.reverse_deps.get(changed_attr, []) + + def can_compute(self, obj, computation): + """Check if all dependencies are available for a computation.""" + required_attrs = self.dependency_graph[computation] + return all(hasattr(obj, attr) for attr in required_attrs) + + +class UserDataInit: """ - Loads user defined initial conditions. Specifically, all attributes of the class object defined in the initial condition is overwritten. - - Attributes - ---------- - **kwargs: class object - + Loads user defined initial conditions with automatic dependency management. """ def __init__(self, **kwargs): - gconsts = params.global_constants() + # Initialise dependency manager + self._dep_manager = DependencyManager() + self._updating = False # Prevent infinite recursion + + # Load global constants first + gconsts = sim_params.global_constants() for key, value in vars(gconsts).items(): setattr(self, key, value) - # else: - ########################################## - # SPATIAL GRID - ########################################## + # Initialise with default values + self._init_defaults() + + # Apply any user-provided kwargs + if kwargs: + for key, value in kwargs.items(): + setattr(self, key, value) + + # Initialise all computable attributes that can be computed + self._initialise_computed_attributes() + + def _initialise_computed_attributes(self): + """Initialise all computed attributes that have their dependencies available.""" + # Get all computed attributes from the dependency graph + computed_attrs = list(self._dep_manager.dependency_graph.keys()) + + for computation in computed_attrs: + # Only compute if not already set and dependencies are available + if not hasattr(self, computation) and self._dep_manager.can_compute( + self, computation + ): + method_name = f"compute_{computation}" + if hasattr(self, method_name): + getattr(self, method_name)() + + def _set_example_global_constants(self): + """Set example global constants for demonstration.""" + self.nspec = 1 + self.buoy = 0 + self.grav = 9.81 # [m s^{-2}] + self.omega = 0.0 # [s^{-1}] + self.R_gas = 287.4 # [J kg^{-1} K^{-1}] + self.R_vap = 461.0 + self.Q_vap = 2.53e06 + self.gamm = 1.4 + self.p_ref = 8.61 * 1e4 # [N/m^2] + self.T_ref = 300.00 # [K] + self.h_ref = 10000.0 # [m] + self.t_ref = 100.0 # [s] + + def _init_defaults(self): + """Initialise default values.""" + # Spatial grid self.inx = 64 + 1 self.iny = 64 + 1 self.inz = 1 @@ -34,102 +108,108 @@ def __init__(self, **kwargs): self.zmin = -1.0 self.zmax = 1.0 - ########################################## - # BOUNDARY CONDITIONS - ########################################## + # Blending choices + self.initial_blending = False + + self.continuous_blending = False + self.no_of_pi_initial = 1 + self.no_of_pi_transition = 0 + self.no_of_hy_initial = 0 + self.no_of_hy_transition = 0 + + self.perturb_type = "pos_perturb" + self.blending_mean = "rhoY" # 1.0, rhoY + self.blending_conv = "rho" # theta, rho + self.blending_type = "half" # half, full + self.blending_weight = 0.0 / 16 + + # Boundary conditions self.bdry_type = np.empty((3), dtype=object) self.bdry_type[0] = opts.BdryType.PERIODIC self.bdry_type[1] = opts.BdryType.WALL self.bdry_type[2] = opts.BdryType.WALL - ########################################## - # TEMPORAL - ########################################## + # Temporal self.CFL = 0.5 self.dtfixed0 = 100.0 self.dtfixed = 100.0 - self.acoustic_timestep = 0 - self.tout = np.arange(0.0, 1.01, 0.01)[10:] self.stepmax = 10000 - ########################################## - # MODEL REGIMES - ########################################## - self.is_ArakawaKonor = 0 - self.is_nonhydrostatic = 1 + # Model regimes self.is_compressible = 1 - + self.is_nonhydrostatic = 1 + self.is_ArakawaKonor = 0 self.compressibility = 1.0 - ########################################## - # PHYSICS AND BACKGROUND WIND - ########################################## + # Physics and background wind self.u_wind_speed = 0.0 self.v_wind_speed = 0.0 self.w_wind_speed = 0.0 - self.stratification = self.stratification_function - ########################################## - # NUMERICS - ########################################## - # Do we solve the left-hand side? + # Numerics self.do_advection = True - - # Advection limiter types self.limiter_type_scalars = opts.LimiterType.NONE self.limiter_type_velocity = opts.LimiterType.NONE - - # Iterative solver self.tol = 1.0e-8 self.max_iterations = 6000 - ########################################## - # BLENDING - ########################################## - self.blending_weight = 0.0 / 16 - self.blending_mean = "rhoY" # 1.0, rhoY - self.blending_conv = "rho" # theta, rho - self.blending_type = "half" - - self.continuous_blending = False - self.no_of_pi_initial = 1 - self.no_of_pi_transition = 0 - self.no_of_hy_initial = 0 - self.no_of_hy_transition = 0 - - self.initial_blending = False - - ########################################## - # DIAGNOSTICS - ########################################## + # Other attributes self.diag = False self.diag_state = None - - ########################################## - # OUTPUTS - ########################################## self.autogen_fn = False self.output_timesteps = False self.output_type = "output" self.output_suffix = "_%i_%i" % (self.inx - 1, self.iny - 1) - if len(kwargs) > 0: - for key, value in kwargs.items(): - setattr(self, key, value) - + def __setattr__(self, name, value): + """Override setattr to handle dependency updates.""" + # Always set the attribute first + super().__setattr__(name, value) + + # Skip dependency updates during initialisation or recursive updates + if name.startswith("_") or not hasattr(self, "_dep_manager") or self._updating: + return + + # Update dependent computations + self._update_dependencies(name) + + def _update_dependencies(self, changed_attr): + """Update all computations that depend on the changed attribute.""" + if self._updating: # Prevent infinite recursion + return + + self._updating = True + try: + computations_to_update = self._dep_manager.get_computations_to_update( + changed_attr + ) + + for computation in computations_to_update: + if self._dep_manager.can_compute(self, computation): + method_name = f"compute_{computation}" + if hasattr(self, method_name): + getattr(self, method_name)() + finally: + self._updating = False + + # Computation methods def compute_u_ref(self): + """Compute reference velocity.""" self.u_ref = self.h_ref / self.t_ref - self.compute_Msq() + # u_ref change triggers Msq update automatically def compute_Msq(self): - self.Msq = self.u_ref * self.u_ref / (self.R_gas * self.T_ref) + """Compute Mach number squared.""" + if hasattr(self, "u_ref"): + self.Msq = self.u_ref * self.u_ref / (self.R_gas * self.T_ref) - def compute_gravity(self): - self.i_gravity = np.zeros((3)) - self.gravity_strength = np.zeros((3)) + def compute_gravity_strength(self): + """Compute gravity-related parameters.""" + self.i_gravity = np.zeros(3) + self.gravity_strength = np.zeros(3) self.gravity_strength[1] = self.grav * self.h_ref / (self.R_gas * self.T_ref) @@ -138,166 +218,57 @@ def compute_gravity(self): self.i_gravity[i] = 1 self.gravity_direction = i - def compute_coriolis(self): - self.i_coriolis = np.zeros((3)) - self.coriolis_strength = np.zeros((3)) + # Alias for backward compatibility + compute_i_gravity = compute_gravity_strength + + def compute_coriolis_strength(self): + """Compute Coriolis parameters.""" + self.i_coriolis = np.zeros(3) + self.coriolis_strength = np.zeros(3) self.coriolis_strength[0] = self.omega * self.t_ref self.coriolis_strength[2] = self.omega * self.t_ref def compute_cp_gas(self): + """Compute specific heat at constant pressure.""" self.cp_gas = self.gamm * self.R_gas / (self.gamm - 1.0) - if all(hasattr(self, attr) for attr in ["grav", "cp_gas", "T_ref"]): - self.compute_N_ref() - def compute_rho_ref(self): + """Compute reference density.""" self.rho_ref = self.p_ref / (self.R_gas * self.T_ref) def compute_N_ref(self): - self.N_ref = self.grav / np.sqrt(self.cp_gas * self.T_ref) - self.Nsq_ref = self.N_ref * self.N_ref + """Compute Brunt-Väisälä frequency.""" + if hasattr(self, "cp_gas"): + self.N_ref = self.grav / np.sqrt(self.cp_gas * self.T_ref) + self.Nsq_ref = self.N_ref * self.N_ref + + # Alias for backward compatibility + compute_Nsq_ref = compute_N_ref def compute_Cs(self): + """Compute sound speed.""" self.Cs = np.sqrt(self.gamm * self.R_gas * self.T_ref) @staticmethod def stratification_function(y): + """Default stratification function.""" return 1.0 def update_ud(self, obj): - for key, value in obj.items(): - setattr(self, key, value) - - # ########################################## - # # SETTER FUNCTIONS - # ########################################## - - # gravity and Msq arguments - @property - def R_gas(self): - return self._R_gas - - @R_gas.setter - def R_gas(self, val): - self._R_gas = val - - if all(hasattr(self, attr) for attr in ["grav", "h_ref", "R_gas", "T_ref"]): - self.compute_gravity() - - if all(hasattr(self, attr) for attr in ["u_ref, R_gas", "T_ref"]): - self.compute_Msq() - - if all(hasattr(self, attr) for attr in ["gamm", "R_gas"]): - self.compute_cp_gas() - - if all(hasattr(self, attr) for attr in ["p_ref, R_gas", "T_ref"]): - self.compute_rho_ref() - - if all(hasattr(self, attr) for attr in ["gamm", "R_gas", "T_ref"]): - self.compute_Cs() - - @property - def T_ref(self): - return self._T_ref - - @T_ref.setter - def T_ref(self, val): - self._T_ref = val - - if all(hasattr(self, attr) for attr in ["grav", "h_ref", "R_gas", "T_ref"]): - self.compute_gravity() - - if all(hasattr(self, attr) for attr in ["u_ref", "R_gas", "T_ref"]): - self.compute_Msq() - - if all(hasattr(self, attr) for attr in ["p_ref", "R_gas", "T_ref"]): - self.compute_rho_ref() - - if all(hasattr(self, attr) for attr in ["grav", "cp_gas", "T_ref"]): - self.compute_N_ref() - - if all(hasattr(self, attr) for attr in ["gamm", "R_gas", "T_ref"]): - self.compute_Cs() - - # gravity arguments - @property - def grav(self): - return self._grav - - @grav.setter - def grav(self, val): - self._grav = val - - if all(hasattr(self, attr) for attr in ["grav", "h_ref", "R_gas", "T_ref"]): - self.compute_gravity() - - if all(hasattr(self, attr) for attr in ["grav", "cp_gas", "T_ref"]): - self.compute_N_ref() - - @property - def h_ref(self): - return self._h_ref - - @h_ref.setter - def h_ref(self, val): - self._h_ref = val - - if all(hasattr(self, attr) for attr in ["h_ref", "t_ref"]): - self.compute_u_ref() - - if all(hasattr(self, attr) for attr in ["grav", "h_ref", "R_gas", "T_ref"]): - self.compute_gravity() - - # coriolis arguments - @property - def t_ref(self): - return self._t_ref - - @t_ref.setter - def t_ref(self, val): - self._t_ref = val - - if all(hasattr(self, attr) for attr in ["h_ref", "t_ref"]): - self.compute_u_ref() - - if all(hasattr(self, attr) for attr in ["omega", "t_ref"]): - self.compute_coriolis() - - @property - def omega(self): - return self._omega - - @omega.setter - def omega(self, val): - self._omega = val - - if all(hasattr(self, attr) for attr in ["omega", "t_ref"]): - self.compute_coriolis() - - # Cs and cp_gas argument - @property - def gamm(self): - return self._gamm - - @gamm.setter - def gamm(self, val): - self._gamm = val - - if all(hasattr(self, attr) for attr in ["gamm", "R_gas"]): - self.compute_cp_gas() - - if all(hasattr(self, attr) for attr in ["gamm", "R_gas", "T_ref"]): - self.compute_Cs() - - # rho_ref argument - @property - def p_ref(self): - return self._p_ref - - @p_ref.setter - def p_ref(self, val): - self._p_ref = val - - if all(hasattr(self, attr) for attr in ["p_ref, R_gas", "T_ref"]): - self.compute_rho_ref() + """Update multiple attributes at once.""" + # Temporarily disable dependency updates + old_updating = self._updating + self._updating = True + + try: + # Set all attributes first + for key, value in obj.items(): + super(UserDataInit, self).__setattr__(key, value) + finally: + self._updating = old_updating + + # Now update all dependencies at once + if not self._updating: + for key in obj.keys(): + self._update_dependencies(key) diff --git a/test_scripts/test_blending.py b/test_scripts/test_blending.py index ad20cd26..ef94d151 100644 --- a/test_scripts/test_blending.py +++ b/test_scripts/test_blending.py @@ -1,6 +1,7 @@ import pytest import subprocess + @pytest.mark.parametrize( "ic", [ diff --git a/test_scripts/test_flow_solver.py b/test_scripts/test_flow_solver.py index 99053894..0f8a4d41 100644 --- a/test_scripts/test_flow_solver.py +++ b/test_scripts/test_flow_solver.py @@ -1,21 +1,23 @@ import pytest import subprocess -@pytest.mark.parametrize("ic", - ["test_travelling_vortex", - "test_internal_long_wave", - "test_lamb_wave", - "test_unstable_lamb",] - ) + +@pytest.mark.parametrize( + "ic", + [ + "test_travelling_vortex", + "test_internal_long_wave", + "test_lamb_wave", + "test_unstable_lamb", + ], +) def test_single_run(ic): result = subprocess.run( - ["pybella", "-ic", ic, "-N", "1"], - capture_output=True, - text=True + ["pybella", "-ic", ic, "-N", "1"], capture_output=True, text=True ) assert result.returncode == 0, ( f"Command failed with return code {result.returncode}\n" f"STDERR:\n{result.stderr.strip()}\n" f"STDOUT:\n{result.stdout.strip()}" - ) \ No newline at end of file + )