augmented_lagrangian_objective.f90

! Copyright (c) 2024-25, The Neko Authors
! All rights reserved.
!
! Redistribution and use in source and binary forms, with or without
! modification, are permitted provided that the following conditions
! are met:
!
!   * Redistributions of source code must retain the above copyright
!     notice, this list of conditions and the following disclaimer.
!
!   * Redistributions in binary form must reproduce the above
!     copyright notice, this list of conditions and the following
!     disclaimer in the documentation and/or other materials provided
!     with the distribution.
!
!   * Neither the name of the authors nor the names of its
!     contributors may be used to endorse or promote products derived
!     from this software without specific prior written permission.
!
! THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
! "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
! LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
! FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
! COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
! INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
! BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
! LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
! CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT
! LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN
! ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
! POSSIBILITY OF SUCH DAMAGE.
!
module augmented_lagrangian_objective
  use num_types, only: rp
  use field, only: field_t
  use field_math, only: field_col3, field_addcol3, field_cmult
  use scratch_registry, only: scratch_registry_t, neko_scratch_registry
  use objective, only: objective_t
  use simulation_m, only: simulation_t
  use neko_config, only: neko_bcknd_device
  use math, only: copy, col2, invcol2
  use device_math, only: device_copy, device_col2, device_invcol2
  use design, only: design_t
  use json_module, only: json_file
  use json_utils, only: json_get_or_default
  use interpolation, only: interpolator_t
  use space, only: space_t, gl
  use coefs, only: coef_t
  implicit none
  private
 
  type, public, extends(objective_t) :: augmented_lagrangian_objective_t
     private
 
     type(field_t), pointer :: u => null()
     type(field_t), pointer :: v => null()
     type(field_t), pointer :: w => null()
 
     type(field_t), pointer :: adjoint_u => null()
     type(field_t), pointer :: adjoint_v => null()
     type(field_t), pointer :: adjoint_w => null()
 
     ! ---- everything GLL ----
     type(space_t), pointer :: xh_gll
     type(coef_t), pointer :: c_xh_gll
 
     ! ---- everything for GL ----
     type(space_t), pointer :: xh_gl
     type(coef_t), pointer :: c_xh_gl
     type(interpolator_t), pointer :: gll_to_gl
     logical :: dealias
     type(scratch_registry_t), pointer :: scratch_gl
 
   contains
     procedure, public, pass(this) :: init_json_sim => &
          augmented_lagrangian_init_json_sim
     procedure, public, pass(this) :: init_from_attributes => &
          augmented_lagrangian_init_attributes
     procedure, public, pass(this) :: free => augmented_lagrangian_free
     procedure, public, pass(this) :: update_value => &
          augmented_lagrangian_update_value
     procedure, public, pass(this) :: update_sensitivity => &
          augmented_lagrangian_update_sensitivity
 
  end type augmented_lagrangian_objective_t
 
contains
 
  subroutine augmented_lagrangian_init_json_sim(this, json, design, simulation)
    class(augmented_lagrangian_objective_t), intent(inout) :: this
    type(json_file), intent(inout) :: json
    class(design_t), intent(in) :: design
    type(simulation_t), target, intent(inout) :: simulation
 
    character(len=:), allocatable :: name
    character(len=:), allocatable :: mask_name
    real(kind=rp) :: weight
    logical :: dealias
 
    call json_get_or_default(json, "weight", weight, 1.0_rp)
    call json_get_or_default(json, "mask_name", mask_name, "")
    call json_get_or_default(json, "name", name, "Augmented Lagrangian")
    call json_get_or_default(json, "dealias", dealias, .true.)
 
    call this%init_from_attributes(design, simulation, weight, name, &
         mask_name, dealias)
  end subroutine augmented_lagrangian_init_json_sim
 
  subroutine augmented_lagrangian_init_attributes(this, design, simulation, &
       weight, name, mask_name, dealias)
    class(augmented_lagrangian_objective_t), intent(inout) :: this
    class(design_t), intent(in) :: design
    type(simulation_t), target, intent(inout) :: simulation
    real(kind=rp), intent(in) :: weight
    character(len=*), intent(in) :: name
    character(len=*), intent(in) :: mask_name
    logical, intent(in) :: dealias
 
    call this%init_base(name, design%size(), weight, mask_name)
 
    ! Save the simulation and design
    this%u => simulation%neko_case%fluid%u
    this%v => simulation%neko_case%fluid%v
    this%w => simulation%neko_case%fluid%w
    this%adjoint_u => simulation%adjoint_case%fluid_adj%u_adj
    this%adjoint_v => simulation%adjoint_case%fluid_adj%v_adj
    this%adjoint_w => simulation%adjoint_case%fluid_adj%w_adj
 
    this%dealias = dealias
    ! GLL
    this%c_Xh_GLL => simulation%neko_case%fluid%c_Xh
    this%Xh_GLL => this%c_Xh_GLL%Xh
 
    ! GL
    this%c_Xh_GL => simulation%adjoint_case%fluid_adj%c_Xh_GL
    this%Xh_GL => this%c_Xh_GL%Xh
 
    ! GLL to GL
    this%GLL_to_GL => simulation%adjoint_case%fluid_adj%GLL_to_GL
 
    this%scratch_GL => simulation%adjoint_case%fluid_adj%scratch_GL
 
  end subroutine augmented_lagrangian_init_attributes
 
  subroutine augmented_lagrangian_free(this)
    class(augmented_lagrangian_objective_t), intent(inout) :: this
    call this%free_base()
 
    if (associated(this%u)) nullify(this%u)
    if (associated(this%v)) nullify(this%v)
    if (associated(this%w)) nullify(this%w)
 
    if (associated(this%adjoint_u)) nullify(this%adjoint_u)
    if (associated(this%adjoint_v)) nullify(this%adjoint_v)
    if (associated(this%adjoint_w)) nullify(this%adjoint_w)
 
    nullify(this%scratch_GL)
 
 
  end subroutine augmented_lagrangian_free
 
  subroutine augmented_lagrangian_update_value(this, design)
    class(augmented_lagrangian_objective_t), intent(inout) :: this
    class(design_t), intent(in) :: design
 
  end subroutine augmented_lagrangian_update_value
 
  subroutine augmented_lagrangian_update_sensitivity(this, design)
    class(augmented_lagrangian_objective_t), intent(inout) :: this
    class(design_t), intent(in) :: design
    type(field_t), pointer :: work
    integer :: temp_indices(1)
    integer :: n_GL, nel
    type(field_t), pointer :: accumulate, fld_GL, adjoint_fld_GL
    integer :: temp_indices_GL(3)
 
    call neko_scratch_registry%request_field(work, temp_indices(1))
 
    if (this%dealias) then
 
       nel = this%c_Xh_GLL%msh%nelv
       n_gl = nel * this%Xh_GL%lxyz
       call this%scratch_GL%request_field(accumulate, temp_indices_gl(1))
       call this%scratch_GL%request_field(fld_gl, temp_indices_gl(2))
       call this%scratch_GL%request_field(adjoint_fld_gl, temp_indices_gl(3))
 
       call this%GLL_to_GL%map(fld_gl%x, this%u%x, nel, this%Xh_GL)
       call this%GLL_to_GL%map(adjoint_fld_gl%x, this%adjoint_u%x, nel, &
            this%Xh_GL)
       call field_col3(accumulate, fld_gl, adjoint_fld_gl)
 
       call this%GLL_to_GL%map(fld_gl%x, this%v%x, nel, this%Xh_GL)
       call this%GLL_to_GL%map(adjoint_fld_gl%x, this%adjoint_v%x, nel, &
            this%Xh_GL)
       call field_addcol3(accumulate, fld_gl, adjoint_fld_gl)
 
       call this%GLL_to_GL%map(fld_gl%x, this%w%x, nel, this%Xh_GL)
       call this%GLL_to_GL%map(adjoint_fld_gl%x, this%adjoint_w%x, nel, &
            this%Xh_GL)
       call field_addcol3(accumulate, fld_gl, adjoint_fld_gl)
 
       ! Evaluate term on GL and preempt the GLL premultiplication
       if (neko_bcknd_device .eq. 1) then
          call device_col2(accumulate%x_d, this%c_Xh_GL%B_d, n_gl)
          call this%GLL_to_GL%map(work%x, accumulate%x, nel, this%Xh_GLL)
          call device_invcol2(work%x_d, this%c_Xh_GLL%B_d, work%size())
       else
          call col2(accumulate%x, this%c_Xh_GL%B, n_gl)
          call this%GLL_to_GL%map(work%x, accumulate%x, nel, this%Xh_GLL)
          call invcol2(work%x, this%c_Xh_GLL%B, work%size())
       end if
 
       call this%scratch_GL%relinquish_field(temp_indices_gl)
    else
       call field_col3(work, this%u, this%adjoint_u)
       call field_addcol3(work, this%v, this%adjoint_v)
       call field_addcol3(work, this%w, this%adjoint_w)
    end if
    ! but negative
    call field_cmult(work, -1.0_rp)
 
    if (neko_bcknd_device .eq. 1) then
       call device_copy(this%sensitivity%x_d, work%x_d, this%sensitivity%size())
    else
       call copy(this%sensitivity%x, work%x, this%sensitivity%size())
    end if
 
    call neko_scratch_registry%relinquish_field(temp_indices)
 
  end subroutine augmented_lagrangian_update_sensitivity
 
end module augmented_lagrangian_objective