Extracting values of fields used for flux boundary conditions at an immersed boundary

[bethanlily]

Hi,

I'm trying to adapt functions I already have for calculating flux boundary conditions at an ice base to be able to output and inspect intermediate values which are used to calculate the fluxes.

I have a module with the rough structure below which creates functions FluxModel and FluxModelBoundaryConditions. I've removed some of the details for clarity.

module FluxModelFunctions

# Import necessary packages
using Oceananigans #...etc.
  
# Exports
export FluxModel, FluxModelBoundaryConditions

struct FluxModel{FT} <: Function
  parameter1 :: FT
  parameter1 :: FT
  #... etc
end

adapt_structure(to, fm::FluxModel) = FluxModel(adapt(to, fm.parameter1),
adapt(to, fm.parameter2))

function FluxModel(;parameter1::FT = 1.0,
  parameter2::FT = 1.0) where FT
  return FluxModel{FT}(parameter1, parameter2)
end

function (fluxes::FluxModel)(x,y,z,t,u,v,T,S,::Val{field}) where field
  p1 = fluxes.parameter1
  p2 = fluxes.parameter2
  
  T₁ = T .+ p1 
  S₁ = S .+ p2 
  U₁ = sqrt(u .^ 2 .+ v .^ 2)
  
  if field == :u
	  return # a function of u, U₁ and parameters
  elseif field == :v
	  return # a function of v, U₁ and parameters
  elseif field == :T
	  return # a function of T₁, S₁, U₁  and parameters
  elseif field == :S
	  return # a function of T₁, S₁, U₁ and parameters
  else
	  error("Invalid field")
  end
end
  

function FluxModelBoundaryConditions(; parameter1::FT = 1.0,
  parameter2::FT = 1.0) where FT
  
  flux_model_instance = FluxModel(; parameter1, parameter2)
  
  u = FluxBoundaryCondition(flux_model_instance, field_dependencies=(:u, :v, :T, :S), parameters = Val(:u))
  v = FluxBoundaryCondition(flux_model_instance, field_dependencies=(:u, :v, :T, :S), parameters = Val(:v))
  T = FluxBoundaryCondition(flux_model_instance, field_dependencies=(:u, :v, :T, :S), parameters = Val(:T))
  S = FluxBoundaryCondition(flux_model_instance, field_dependencies=(:u, :v, :T, :S), parameters = Val(:S))
  
  return (; u, v, T, S)
end

end # module

These functions are then used to construct the boundary conditions for my model:

flux_boundary_conditions = FluxModelBoundaryConditions()

boundary_conditions = (
u = FieldBoundaryConditions(immersed = flux_boundary_conditions.u),
v = FieldBoundaryConditions(immersed = flux_boundary_conditions.v),
T = FieldBoundaryConditions(immersed = flux_boundary_conditions.T),
S = FieldBoundaryConditions(immersed = flux_boundary_conditions.S)) 

This has been working well and allows me to prescribe dynamic boundary conditions for ice melt at an immersed boundary (at the top of my domain). However, I would like to inspect intermediate values used in calculating the fluxes such as the tracers and velocities in the boundary-adjacent grid cells and functions of these such as freezing point and melt rate.

I have tried calculating these after the fact and extracting the boundary-adjacent values of the fields using a version of BottomCellValue (from Oceanostics) which works for an immersed boundary at the top of the domain, but am finding some inconsistencies and so would like to make sure I am getting the exact values which are actually being used in the flux calculation for my boundary conditions.

I would like to
(a) extract the values of the fields being used for the boundary condition at the immersed boundary and intermediate functions of them that are part of my flux calculations
(b) save these as fields so I can plot them later.
I should possibly say that I need this is be able to run on GPU.

Any guidance on how to adapt my code would be much appreciated. If it is not possible to do this together with making the boundary conditions, then I'd welcome any ideas on how to be as consistent as possible after the fact. Thanks!

[glwagner]

This looks like a cool problem!

I am not sure I completely understand the question so I will provide a cloud of information first and hopefully some of it is useful:

1. If you would like to output a boundary condition, you may want to look into Oceananigans.Models.BoundaryConditionOperation (the docstring has some examples)
2. You may want to avoid if statements in your function. Especially for immersed boundary conditions, which are evaluated everywhere -- this can cause slowdown. You'll need independent functions for each kernel operation. You can reuse code by breaking things up into multiple functions (eg call function 1 from functions 2-5, etc)
3. For a complex production application like this I recommend using discrete_form=true rather than the continuous form.
4. Do not use broadcasting inside a boundary condition or forcing function. If you are using the continuous form, all of the inputs to the function are scalars (you can test this using @show on CPU). If you are using the discrete form, you receive array input but you still cannot use broadcasting --- the function must be evaluated at the index i, j, k. For example, you do not want to broadcast over the entire array at every grid point.

Next, I think you are saying that you would like to evaluate a field on the boundary, is that correct? You can save the entire 3D field for this. If you are using a 2D bathymetry representation, then we should also be able to build a diagnostic that outputs a field on the cells that lie just above the immersed boundary. Unfortunately we don't have that yet but I think it is good to add.

I think we need to take into account the staggered grid to understand what you are trying to output. Converting your boundary conditions to the discrete form will help. The continuous form boundary conditions do interpolation under the hood.

This last point does give me an idea, that one way to recover the field involved in your boundary condition would be to produce a similar function to the one you are using, and then plug this into the same machinery that is used for BoundaryConditionOperation. But we might be able to come up with a more elegant solution if we understand more precisely what we need to do.

[bethanlily]

Hi Greg,

Thank you very much for your reply.

Firstly, thanks for the tips for writing the functions - I will rewrite to remove the broadcasting and the if statements.

Secondly, yes I am just trying to evaluate the field involved in my boundary condition at the boundary.
I think I am currently doing something similar to what you suggest with saving the entire 3D field and outputting the field on the cells just above (or actually below in my case) the boundary.
To do that I am using:

@inline top_adjacent_node(i, j, k, grid, LX, LY, LZ) = !inactive_node(i, j, k, grid, LX, LY, LZ) && inactive_node(i, j, k + 1, grid, LX, LY, LZ) 

function top_cell_value(diagnostic)
    loc = location(diagnostic)
    instantiated_location = (L() for L in loc)
    condition_to_ban = KernelFunctionOperation{loc...}(top_adjacent_node, diagnostic.grid, instantiated_location...)
    if size(diagnostic.data, 2) == 1
        return condition_to_ban
    else
        return diagnostic * condition_to_ban
    end
end 

I'm currently using this on my outputted FieldTimeSeries rather than using it to create a diagnostic during the model run, which would perhaps be better.
But I think your point about using a discrete boundary condition is important here - it seems that the interpolation done for the continuous boundary condition would lead to different values for the fields than what I'm getting with the above function.
Would I use something like the function top_adjacent_node to find the correct z index within the discrete boundary condition function?

I wasn't aware of BoundaryConditionOperation, it seems very useful so thanks for pointing that out! It does seem like doing something similar could help but perhaps is overengineering it a bit. I'll try changing to a discrete boundary condition first!