Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Reduce parameter space for velocity tendency kernels #3924

Merged
merged 18 commits into from
Nov 19, 2024
Merged

Reduce parameter space for velocity tendency kernels #3924

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
18 commits
Select commit Hold shift + click to select a range
feda572
reduce parameter space
simone-silvestri Nov 14, 2024
3f8b609
change other kernels for symmetry
simone-silvestri Nov 14, 2024
06468b0
only free surface forcing
simone-silvestri Nov 14, 2024
d04c680
pass specific forcing in kernels
simone-silvestri Nov 14, 2024
41bad88
bugfix
simone-silvestri Nov 14, 2024
fa51e82
bugfix
simone-silvestri Nov 14, 2024
eca101b
Merge branch 'main' into ss/reduce-parameter-space
simone-silvestri Nov 14, 2024
24591f6
Merge branch 'main' into ss/reduce-parameter-space
simone-silvestri Nov 14, 2024
c44258b
Merge branch 'main' into ss/reduce-parameter-space
simone-silvestri Nov 15, 2024
2e09e59
reduce parameter space
simone-silvestri Nov 15, 2024
b138c36
import on_architecture
simone-silvestri Nov 15, 2024
4995636
some changes
simone-silvestri Nov 15, 2024
2061f57
put model_fields before the tracer loop
simone-silvestri Nov 15, 2024
1940b6e
remove tracers from here
simone-silvestri Nov 15, 2024
d2cc93f
remove auxiliary fields
simone-silvestri Nov 15, 2024
7b4997c
Merge branch 'main' into ss/reduce-parameter-space
simone-silvestri Nov 15, 2024
bf08dc4
Merge branch 'main' into ss/reduce-parameter-space
simone-silvestri Nov 18, 2024
32a711e
revert to previous code
simone-silvestri Nov 18, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions ...els/HydrostaticFreeSurfaceModels/SplitExplicitFreeSurfaces/split_explicit_free_surface.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
using Oceananigans.BuoyancyModels: g_Earth
using Oceananigans.Grids: with_halo
import Oceananigans.Grids: on_architecture

struct SplitExplicitFreeSurface{H, U, M, FT, K , S, T} <: AbstractFreeSurface{H, FT}
η :: H
Expand Down
25 changes: 12 additions & 13 deletions src/Models/HydrostaticFreeSurfaceModels/compute_hydrostatic_free_surface_tendencies.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,8 +91,8 @@ function compute_hydrostatic_free_surface_tendency_contributions!(model, kernel_
model.tracers,
model.diffusivity_fields,
model.auxiliary_fields,
c_forcing,
model.clock)
model.clock,
c_forcing)

launch!(arch, grid, kernel_parameters,
compute_hydrostatic_free_surface_Gc!,
Expand Down Expand Up @@ -137,20 +137,19 @@ function compute_hydrostatic_momentum_tendencies!(model, velocities, kernel_para
model.diffusivity_fields,
model.pressure.pHY′,
model.auxiliary_fields,
model.forcing,
model.clock)

u_kernel_args = tuple(start_momentum_kernel_args..., u_immersed_bc, end_momentum_kernel_args...)
v_kernel_args = tuple(start_momentum_kernel_args..., v_immersed_bc, end_momentum_kernel_args...)

launch!(arch, grid, kernel_parameters,
compute_hydrostatic_free_surface_Gu!, model.timestepper.Gⁿ.u, grid,
active_cells_map, u_kernel_args;
active_cells_map, u_kernel_args, model.forcing.u;
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
active_cells_map, u_kernel_args, model.forcing.u;
active_cells_map, tuple(u_kernel_args..., model.forcing.u);

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

but it doesnt matter

active_cells_map)

launch!(arch, grid, kernel_parameters,
compute_hydrostatic_free_surface_Gv!, model.timestepper.Gⁿ.v, grid,
active_cells_map, v_kernel_args;
active_cells_map, v_kernel_args, model.forcing.v;
active_cells_map)

return nothing
Expand Down Expand Up @@ -180,27 +179,27 @@ end
#####

""" Calculate the right-hand-side of the u-velocity equation. """
@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args)
@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, ::Nothing, args, forcing)
i, j, k = @index(Global, NTuple)
@inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...)
@inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args..., forcing)
end

@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args)
@kernel function compute_hydrostatic_free_surface_Gu!(Gu, grid, active_cells_map, args, forcing)
idx = @index(Global, Linear)
i, j, k = active_linear_index_to_tuple(idx, active_cells_map)
@inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args...)
@inbounds Gu[i, j, k] = hydrostatic_free_surface_u_velocity_tendency(i, j, k, grid, args..., forcing)
end

""" Calculate the right-hand-side of the v-velocity equation. """
@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args)
@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, ::Nothing, args, forcing)
i, j, k = @index(Global, NTuple)
@inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
@inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args..., forcing)
end

@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args)
@kernel function compute_hydrostatic_free_surface_Gv!(Gv, grid, active_cells_map, args, forcing)
idx = @index(Global, Linear)
i, j, k = active_linear_index_to_tuple(idx, active_cells_map)
@inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args...)
@inbounds Gv[i, j, k] = hydrostatic_free_surface_v_velocity_tendency(i, j, k, grid, args..., forcing)
end

#####
Expand Down
18 changes: 9 additions & 9 deletions ...Models/HydrostaticFreeSurfaceModels/hydrostatic_free_surface_tendency_kernel_functions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,8 +38,8 @@ implicitly during time-stepping.
diffusivities,
hydrostatic_pressure_anomaly,
auxiliary_fields,
forcings,
clock)
clock,
forcing)

model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)

Expand All @@ -49,7 +49,7 @@ implicitly during time-stepping.
- ∂xᶠᶜᶜ(i, j, k, grid, hydrostatic_pressure_anomaly)
- ∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy)
- immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields)
+ forcings.u(i, j, k, grid, clock, hydrostatic_prognostic_fields(velocities, free_surface, tracers)))
+ forcing(i, j, k, grid, clock, hydrostatic_prognostic_fields(velocities, free_surface, tracers)))
end

"""
Expand Down Expand Up @@ -77,8 +77,8 @@ implicitly during time-stepping.
diffusivities,
hydrostatic_pressure_anomaly,
auxiliary_fields,
forcings,
clock)
clock,
forcing)

model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)

Expand All @@ -88,7 +88,7 @@ implicitly during time-stepping.
- ∂yᶜᶠᶜ(i, j, k, grid, hydrostatic_pressure_anomaly)
- ∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, buoyancy)
- immersed_∂ⱼ_τ₂ⱼ(i, j, k, grid, velocities, v_immersed_bc, closure, diffusivities, clock, model_fields)
+ forcings.v(i, j, k, grid, clock, model_fields))
+ forcing(i, j, k, grid, clock, model_fields))
end

"""
Expand Down Expand Up @@ -116,8 +116,8 @@ where `c = C[tracer_index]`.
tracers,
diffusivities,
auxiliary_fields,
forcing,
clock) where tracer_index
clock,
forcing) where tracer_index

@inbounds c = tracers[tracer_index]
model_fields = merge(hydrostatic_fields(velocities, free_surface, tracers), auxiliary_fields)
Expand All @@ -135,4 +135,4 @@ where `c = C[tracer_index]`.
- immersed_∇_dot_qᶜ(i, j, k, grid, c, c_immersed_bc, closure, diffusivities, val_tracer_index, clock, model_fields)
+ biogeochemical_transition(i, j, k, grid, biogeochemistry, val_tracer_name, clock, model_fields)
+ forcing(i, j, k, grid, clock, model_fields))
end
end
8 changes: 4 additions & 4 deletions src/Models/NonhydrostaticModels/compute_nonhydrostatic_tendencies.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -91,15 +91,15 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti

u_kernel_args = tuple(start_momentum_kernel_args...,
u_immersed_bc, end_momentum_kernel_args...,
forcings, hydrostatic_pressure, clock)
hydrostatic_pressure, clock, forcings.u)

v_kernel_args = tuple(start_momentum_kernel_args...,
v_immersed_bc, end_momentum_kernel_args...,
forcings, hydrostatic_pressure, clock)
hydrostatic_pressure, clock, forcings.v)

w_kernel_args = tuple(start_momentum_kernel_args...,
w_immersed_bc, end_momentum_kernel_args...,
forcings, hydrostatic_pressure, clock)
hydrostatic_pressure, clock, forcings.w)

exclude_periphery = true
launch!(arch, grid, kernel_parameters, compute_Gu!,
Expand Down Expand Up @@ -128,7 +128,7 @@ function compute_interior_tendency_contributions!(model, kernel_parameters; acti
start_tracer_kernel_args...,
c_immersed_bc,
end_tracer_kernel_args...,
forcing, clock)
clock, forcing)

launch!(arch, grid, kernel_parameters, compute_Gc!,
c_tendency, grid, active_cells_map, args;
Expand Down
28 changes: 14 additions & 14 deletions src/Models/NonhydrostaticModels/nonhydrostatic_tendency_kernel_functions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,17 +56,17 @@ pressure anomaly.
tracers,
auxiliary_fields,
diffusivities,
forcings,
hydrostatic_pressure,
clock)
clock,
forcing)

model_fields = merge(velocities, tracers, auxiliary_fields)

total_velocities = (u = SumOfArrays{2}(velocities.u, background_fields.velocities.u),
v = SumOfArrays{2}(velocities.v, background_fields.velocities.v),
w = SumOfArrays{2}(velocities.w, background_fields.velocities.w))

total_velocities = with_advective_forcing(forcings.u, total_velocities)
total_velocities = with_advective_forcing(forcing, total_velocities)

return ( - div_𝐯u(i, j, k, grid, advection, total_velocities, velocities.u)
- div_𝐯u(i, j, k, grid, advection, velocities, background_fields.velocities.u)
Expand All @@ -77,7 +77,7 @@ pressure anomaly.
- immersed_∂ⱼ_τ₁ⱼ(i, j, k, grid, velocities, u_immersed_bc, closure, diffusivities, clock, model_fields)
+ x_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time)
+ ∂t_uˢ(i, j, k, grid, stokes_drift, clock.time)
+ forcings.u(i, j, k, grid, clock, model_fields))
+ forcing(i, j, k, grid, clock, model_fields))
end

"""
Expand Down Expand Up @@ -119,17 +119,17 @@ pressure anomaly.
tracers,
auxiliary_fields,
diffusivities,
forcings,
hydrostatic_pressure,
clock)
clock,
forcing)

model_fields = merge(velocities, tracers, auxiliary_fields)

total_velocities = (u = SumOfArrays{2}(velocities.u, background_fields.velocities.u),
v = SumOfArrays{2}(velocities.v, background_fields.velocities.v),
w = SumOfArrays{2}(velocities.w, background_fields.velocities.w))

total_velocities = with_advective_forcing(forcings.v, total_velocities)
total_velocities = with_advective_forcing(forcing, total_velocities)

return ( - div_𝐯v(i, j, k, grid, advection, total_velocities, velocities.v)
- div_𝐯v(i, j, k, grid, advection, velocities, background_fields.velocities.v)
Expand All @@ -140,7 +140,7 @@ pressure anomaly.
- immersed_∂ⱼ_τ₂ⱼ(i, j, k, grid, velocities, v_immersed_bc, closure, diffusivities, clock, model_fields)
+ y_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time)
+ ∂t_vˢ(i, j, k, grid, stokes_drift, clock.time)
+ forcings.v(i, j, k, grid, clock, model_fields))
+ forcing(i, j, k, grid, clock, model_fields))
end

# Only add buoyancy if the hydrostatic pressure isa Nothing
Expand Down Expand Up @@ -185,17 +185,17 @@ velocity components, tracer fields, and precalculated diffusivities where applic
tracers,
auxiliary_fields,
diffusivities,
forcings,
hydrostatic_pressure,
clock)
clock,
forcing)

model_fields = merge(velocities, tracers, auxiliary_fields)

total_velocities = (u = SumOfArrays{2}(velocities.u, background_fields.velocities.u),
v = SumOfArrays{2}(velocities.v, background_fields.velocities.v),
w = SumOfArrays{2}(velocities.w, background_fields.velocities.w))

total_velocities = with_advective_forcing(forcings.w, total_velocities)
total_velocities = with_advective_forcing(forcing, total_velocities)

return ( - div_𝐯w(i, j, k, grid, advection, total_velocities, velocities.w)
- div_𝐯w(i, j, k, grid, advection, velocities, background_fields.velocities.w)
Expand All @@ -205,7 +205,7 @@ velocity components, tracer fields, and precalculated diffusivities where applic
- immersed_∂ⱼ_τ₃ⱼ(i, j, k, grid, velocities, w_immersed_bc, closure, diffusivities, clock, model_fields)
+ z_curl_Uˢ_cross_U(i, j, k, grid, stokes_drift, velocities, clock.time)
+ ∂t_wˢ(i, j, k, grid, stokes_drift, clock.time)
+ forcings.w(i, j, k, grid, clock, model_fields))
+ forcing(i, j, k, grid, clock, model_fields))
end

"""
Expand Down Expand Up @@ -247,8 +247,8 @@ velocity components, tracer fields, and precalculated diffusivities where applic
tracers,
auxiliary_fields,
diffusivities,
forcing,
clock) where tracer_index
clock,
forcing) where tracer_index

@inbounds c = tracers[tracer_index]
@inbounds background_fields_c = background_fields.tracers[tracer_index]
Expand Down
12 changes: 6 additions & 6 deletions src/Models/ShallowWaterModels/compute_shallow_water_tendencies.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,22 +75,22 @@ function compute_interior_tendency_contributions!(tendencies,
formulation)

transport_args = (grid, gravitational_acceleration, advection.momentum, velocities, coriolis, closure,
bathymetry, solution, tracers, diffusivities, forcings, clock, formulation)
bathymetry, solution, tracers, diffusivities, clock, formulation)

h_args = (grid, gravitational_acceleration, advection.mass, coriolis, closure,
solution, tracers, diffusivities, forcings, clock, formulation)
solution, tracers, diffusivities, clock, formulation)

launch!(arch, grid, :xyz, compute_Guh!, tendencies[1], transport_args...; exclude_periphery=true)
launch!(arch, grid, :xyz, compute_Gvh!, tendencies[2], transport_args...; exclude_periphery=true)
launch!(arch, grid, :xyz, compute_Gh!, tendencies[3], h_args...)
launch!(arch, grid, :xyz, compute_Guh!, tendencies[1], transport_args..., forcings[1]; exclude_periphery=true)
launch!(arch, grid, :xyz, compute_Gvh!, tendencies[2], transport_args..., forcings[2]; exclude_periphery=true)
launch!(arch, grid, :xyz, compute_Gh!, tendencies[3], h_args..., forcings.h)

for (tracer_index, tracer_name) in enumerate(propertynames(tracers))
@inbounds Gc = tendencies[tracer_index+3]
@inbounds forcing = forcings[tracer_index+3]
@inbounds c_advection = advection[tracer_name]

launch!(arch, grid, :xyz, compute_Gc!, Gc, grid, Val(tracer_index),
c_advection, closure, solution, tracers, diffusivities, forcing, clock, formulation)
c_advection, closure, solution, tracers, diffusivities, clock, formulation, forcing)
end

return nothing
Expand Down
22 changes: 11 additions & 11 deletions src/Models/ShallowWaterModels/solution_and_tracer_tendencies.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,9 +32,9 @@ Compute the tendency for the x-directional transport, uh
solution,
tracers,
diffusivities,
forcings,
clock,
formulation)
formulation,
forcing)

g = gravitational_acceleration

Expand All @@ -45,7 +45,7 @@ Compute the tendency for the x-directional transport, uh
- x_f_cross_U(i, j, k, grid, coriolis, solution)
- bathymetry_contribution_x(i, j, k, grid, g, solution.h, bathymetry, formulation)
- sw_∂ⱼ_τ₁ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, formulation)
+ forcings[1](i, j, k, grid, clock, merge(solution, tracers)))
+ forcing(i, j, k, grid, clock, merge(solution, tracers)))
end

"""
Expand All @@ -61,9 +61,9 @@ Compute the tendency for the y-directional transport, vh.
solution,
tracers,
diffusivities,
forcings,
clock,
formulation)
formulation,
forcing)

g = gravitational_acceleration

Expand All @@ -74,7 +74,7 @@ Compute the tendency for the y-directional transport, vh.
- y_f_cross_U(i, j, k, grid, coriolis, solution)
- bathymetry_contribution_y(i, j, k, grid, g, solution.h, bathymetry, formulation)
- sw_∂ⱼ_τ₂ⱼ(i, j, k, grid, closure, diffusivities, clock, model_fields, formulation)
+ forcings[2](i, j, k, grid, clock, merge(solution, tracers)))
+ forcing(i, j, k, grid, clock, merge(solution, tracers)))
end

"""
Expand All @@ -88,12 +88,12 @@ Compute the tendency for the height, h.
solution,
tracers,
diffusivities,
forcings,
clock,
formulation)
formulation,
forcing)

return ( - div_Uh(i, j, k, grid, advection, solution, formulation)
+ forcings.h(i, j, k, grid, clock, merge(solution, tracers)))
+ forcing(i, j, k, grid, clock, merge(solution, tracers)))
end

@inline function tracer_tendency(i, j, k, grid,
Expand All @@ -103,9 +103,9 @@ end
solution,
tracers,
diffusivities,
forcing,
clock,
formulation) where tracer_index
formulation,
forcing) where tracer_index

@inbounds c = tracers[tracer_index]

Expand Down