eddy_graph.f90

program eddy_graph

  use, intrinsic:: ISO_FORTRAN_ENV

  ! Libraries:
  use ezmpi, only: ezmpi_bcast
  use jumble, only: get_command_arg_dyn, new_unit, ediff1d, assert, deg_to_rad
  use mpi_f08, only: mpi_init, mpi_finalize, MPI_Comm_rank, MPI_Comm_world, &
       MPI_Comm_size

  use derived_types, only: snapshot, shpc
  use dispatch_snapshot_m, only: dispatch_snapshot
  use get_snapshot_m, only: get_snapshot
  use overlap_m, only: overlap
  use unit_edge_m, only: open_edge_file, unit_edge
  use shpc_open_m, only: shpc_open
  use shpc_close_m, only: shpc_close
  
  implicit none

  integer rank, n_proc, copy, delta, i, j
  integer:: n_dates = huge(0) ! number of dates to read
  integer k ! date index
  integer k_begin, k_end, k_end_main_loop
  character(len = :), allocatable:: shpc_dir
  integer:: max_delta = 1
  
  real:: corner_deg(2) = [huge(0.), huge(0.)]
  ! longitude and latitude of the corner of the whole grid in input
  ! NetCDF, in degrees

  real corner(2)
  ! longitude and latitude of the corner of the whole grid in input
  ! NetCDF, in rad

  real:: step_deg(2) = [huge(0.), huge(0.)]
  ! longitude and latitude steps, in degrees
  
  real step(2) ! longitude and latitude steps, in rad

  integer:: nlon = - 1, nlat = - 1
  ! size of ssh array in input NetCDF, assuming no repeated point if
  ! the grid is global

  integer:: dist_lim = 12
  ! We look for an overlapping eddy at dist_lim (in grid points) of
  ! the first extremum.

  logical periodic ! grid is periodic in longitude
  integer unit_isolated, unit
  integer n_shpc ! number of input SHPC directories
  integer e_overestim ! over-estimation of the number of eddies at each date
  TYPE(shpc), allocatable:: hshp(:) ! (n_shpc)
  type(snapshot), allocatable:: flow(:)
  character(len = 30) file

  namelist /grid_nml/ corner_deg, step_deg, nlon, nlat
  namelist /main_nml/ max_delta, dist_lim, n_dates

  !-------------------------------------------------------------------------

  ! 1. Front matter

  call mpi_init
  call MPI_Comm_rank(MPI_Comm_world, rank)
  call MPI_Comm_size(MPI_Comm_world, n_proc)

  if (rank == 0) then
     n_shpc = COMMAND_ARGUMENT_COUNT()
     call assert(n_shpc /= 0, &
          "Required arguments: SHPC-directory [SHPC-directory] ...")
  end if

  call ezmpi_bcast(n_shpc, root = 0)
  allocate(hshp(n_shpc))

  do i = 1, n_shpc
     call get_command_arg_dyn(i, shpc_dir)
     call shpc_open(hshp(i), trim(shpc_dir), pszaccess = "rb")
  end do

  if (rank == 0) then
     ! Assuming grid_nml.txt is the same in all input SHPC:
     call new_unit(unit)
     open(unit, file = hshp(1)%dir // "/grid_nml.txt", status = "old", &
          action = "read", position = "rewind")
     read(unit, nml = grid_nml)
     close(unit)

     write(unit = error_unit, nml = main_nml)
     write(unit = error_unit, fmt = *) "Enter namelist main_nml."
     read(unit = *, nml = main_nml)
     write(unit = *, nml = main_nml)
     call assert(max_delta >= 1, "eddy_graph max_delta")
     call assert(n_dates == huge(0) .or. n_shpc == 1, "If we have several " &
          // "SHPC directories, then we read all the dates from them")

     ! As we are requiring the grid spacing to be uniform, the value of
     ! "periodic" may be deduced from the values of step_deg(1) and nlon:
     periodic = nint(360. / step_deg(1)) == nlon
     print *, "periodic = ", periodic
     if (periodic) call assert(2 * dist_lim * step_deg(1) < 180., &
          "eddy_graph dist_lim")
     copy = merge(dist_lim, 0, periodic)

     if (n_dates == huge(0)) &
          n_dates = sum([(size(hshp(i)%ishape_last), i = 1, n_shpc)])
     print *, "n_dates = ", n_dates
     call assert(n_dates >= max_delta + 1, &
          "eddy_graph: n_dates should be >= max_delta + 1")
     call assert(n_proc <= n_dates / (max_delta + 1), &
          "eddy_graph: n_proc should be <= n_dates / (max_delta + 1)")
     e_overestim = maxval([(hshp(i)%ishape_last(hshp(i)%d0) + 1, &
          ediff1d(hshp(i)%ishape_last), i = 1, n_shpc)])
     open(unit, file = "node_id_param.json", status = "replace", &
          action = "write")
     write(unit, fmt = *) '{"e_overestim": ', e_overestim, '}'
     close(unit)
     corner = corner_deg * deg_to_rad
     step = step_deg * deg_to_rad
  end if

  call ezmpi_bcast(max_delta, root = 0)
  call ezmpi_bcast(corner, root = 0)
  call ezmpi_bcast(step, root = 0)
  call ezmpi_bcast(nlon, root = 0)
  call ezmpi_bcast(nlat, root = 0)
  call ezmpi_bcast(dist_lim, root = 0)
  call ezmpi_bcast(periodic, root = 0)
  call ezmpi_bcast(copy, root = 0)
  call ezmpi_bcast(n_dates, root = 0)
  call ezmpi_bcast(e_overestim, root = 0)

  ! Open output files:
  
  call new_unit(unit_isolated)
  write(unit = file, fmt = '("isolated_nodes_", i0, ".txt")') rank
  open(unit_isolated, file = trim(file), status = "replace", action = "write")
  
  call open_edge_file(rank)
  ! We do not write the title line. That will be handled by eddy_graph.sh.

  k_begin = hshp(1)%d0 + (rank * n_dates) / n_proc

  if (rank < n_proc - 1) then
     k_end = hshp(1)%d0 + ((rank + 1) * n_dates) / n_proc + max_delta - 1
     k_end_main_loop = k_end - max_delta + 1
  else
     k_end = hshp(1)%d0 + n_dates - 1
     k_end_main_loop = k_end
  end if

  allocate(flow(max_delta + 1))

  ! 2. Prologue:

  do k = k_begin, k_begin + max_delta
     call get_snapshot(flow(k - k_begin + 1), nlon, nlat, corner, step, copy, &
          hshp, k, k_end, rank, n_proc, max_delta) ! (read)
  end do

  do delta = 1, max_delta
     do k = k_begin + delta, k_begin + max_delta
        call overlap(flow, nlon, nlat, periodic, dist_lim, e_overestim, k, &
             delta, j = k - k_begin + 1)
     end do
  end do

  ! 3. Main loop

  do k = k_begin + max_delta + 1, k_end_main_loop
     call dispatch_snapshot(flow(1), unit_isolated, rank, k_begin, max_delta, &
          k = k - max_delta - 1)
     flow(:max_delta) = flow(2:)
     call get_snapshot(flow(max_delta + 1), nlon, nlat, corner, step, copy, &
          hshp, k, k_end, rank, n_proc, max_delta)
     
     do delta = 1, max_delta
        call overlap(flow, nlon, nlat, periodic, dist_lim, e_overestim, k, &
             delta, j = max_delta + 1)
     end do
  end do

  ! 4. Epilogue

  do k = k_end_main_loop + 1, k_end
     ! {rank < n_proc - 1 and k >= k_end - max_delta + 2}
     call dispatch_snapshot(flow(1), unit_isolated, rank, k_begin, max_delta, &
          k = k - max_delta - 1)
     flow(:max_delta) = flow(2:)
     call get_snapshot(flow(max_delta + 1), nlon, nlat, corner, step, copy, &
          hshp, k, k_end, rank, n_proc, max_delta)
     ! (reception)

     ! Stitching:
     do delta = k - k_end + max_delta, max_delta
        call overlap(flow, nlon, nlat, periodic, dist_lim, e_overestim, k, &
             delta, j = max_delta + 1)
     end do
  end do

  do j = 1, max_delta + 1
     call dispatch_snapshot(flow(j), unit_isolated, rank, k_begin, max_delta, &
          k = k_end - max_delta - 1 + j)
  end do

  ! 5. Back matter

  do i = 1, n_shpc
     CALL shpc_close(hshp(i))
  end do

  close(unit_isolated)
  close(unit_edge)
  call mpi_finalize 

end program eddy_graph