reintegrate_column.py

# Reintegrate column based on subroutine in MOM6: MOM_diag_vkernels.F90 translated into python 
# Redo a reintegration
def reintegrate_column(h_src, uh_src, h_dest):
  missing_value = 0.
  nsrc = h_src.size
  ndest = h_dest.size
  uh_dest = np.zeros(h_dest.shape)
  print(uh_dest.shape)
  k_src = -1
  k_dest = -1
  h_dest_rem = 0.
  h_src_rem = 0.
  src_ran_out = False
  src_exists = False

  while True:
    if (h_src_rem==0. and k_src<nsrc-1):
        
        # Supply is empty so move to the next source cell
        k_src = k_src + 1
        h_src_rem = h_src[k_src]
        uh_src_rem = uh_src[k_src]         
        if (h_src_rem==0.):
            continue       
        src_exists = False # This stops us masking out the entire column    
        
    if (h_dest_rem==0. and k_dest<ndest-1):
        # Sink has no capacity so move to the next destination cell
        k_dest = k_dest + 1
        h_dest_rem = h_dest[k_dest]
        uh_dest[k_dest] = 0.
        if (h_dest_rem==0.):
            continue    
            
    if (k_src==nsrc-1 and h_src_rem==0.):                
        if (src_ran_out):
            break
        src_ran_out = True
        continue
    
    duh = 0.
    if (h_src_rem<h_dest_rem):        
        # The source cell is fully within the destination cell
        dh = h_src_rem
        if (dh>0.):
            duh = uh_src_rem
        h_src_rem = 0.
        uh_src_rem = 0.
        h_dest_rem = max(0., h_dest_rem - dh)
    elif (h_src_rem>h_dest_rem):
        # Only part of the source cell can be used up
        dh = h_dest_rem
        duh = (dh / h_src_rem) * uh_src_rem        
        h_src_rem = max(0., h_src_rem - dh)
        uh_src_rem = uh_src_rem - duh
        h_dest_rem = 0.
    else: # h_src_rem==h_dest_rem
        # The source cell exactly fits the destination cell
        duh = uh_src_rem
        h_src_rem = 0.
        uh_src_rem = 0.
        h_dest_rem = 0.    
    uh_dest[k_dest] = uh_dest[k_dest] + duh
    if (k_dest==ndest-1 and (k_src==nsrc-1 or h_dest_rem==0.)):
        break  

  return uh_dest