InitSolution Subroutine

SUBROUTINE InitSolution(U_init,which_init_in)
! MODULES
  USE MODgvec_Globals,       ONLY:ProgressBar,getTime
  USE MODgvec_MHD3D_Vars   , ONLY:init_fromBConly,init_BC,init_average_axis,average_axis_move
  USE MODgvec_MHD3D_Vars   , ONLY:X1_base,X1_BC_Type,X1_a,X1_b
  USE MODgvec_MHD3D_Vars   , ONLY:X2_base,X2_BC_Type,X2_a,X2_b
  USE MODgvec_MHD3D_Vars   , ONLY:LA_base,LA_BC_Type,init_LA
  USE MODgvec_sol_var_MHD3D, ONLY:t_sol_var_mhd3d
  USE MODgvec_lambda_solve,  ONLY:lambda_solve
  USE MODgvec_VMEC_Vars,     ONLY:Rmnc_spl,Rmns_spl,Zmnc_spl,Zmns_spl
  USE MODgvec_VMEC_Vars,     ONLY:lmnc_spl,lmns_spl
  USE MODgvec_VMEC_Readin,   ONLY:lasym
  USE MODgvec_VMEC,          ONLY:VMEC_EvalSplMode
!$ USE omp_lib
  IMPLICIT NONE
!-----------------------------------------------------------------------------------------------------------------------------------
! INPUT VARIABLES
  INTEGER, INTENT(IN) :: which_init_in
!-----------------------------------------------------------------------------------------------------------------------------------
! OUTPUT VARIABLES
  CLASS(t_sol_var_MHD3D), INTENT(INOUT) :: U_init
!-----------------------------------------------------------------------------------------------------------------------------------
! LOCAL VARIABLES
  INTEGER  :: iMode,i_m
  REAL(wp) :: BC_val(2)
  REAL(wp) :: rhopos
  REAL(wp) :: X1_gIP(1:X1_base%s%nBase,1:X1_base%f%modes)
  REAL(wp) :: X2_gIP(1:X2_base%s%nBase,1:X2_base%f%modes)
  REAL(wp) :: LA_gIP(1:LA_base%s%nBase,1:LA_base%f%modes)
!===================================================================================================================================
  IF(.NOT.MPIroot) CALL abort(__STAMP__, &
                       "InitSolution should only be called by MPIroot!")
  X1_gIP=0.0_wp; X2_gIP=0.0_wp; LA_gIP=0.0_wp

  SELECT CASE(which_init_in)
  CASE(-1) !restart
    X1_a(:)=U_init%X1(1,:)
    X2_a(:)=U_init%X2(1,:)
    X1_b(:)=U_init%X1(X1_base%s%nBase,:)
    X2_b(:)=U_init%X2(X2_base%s%nBase,:)
    IF (init_average_axis)THEN
      WRITE(UNIT_stdOut,'(A)') "WARNING: init_average_axis ignored due to restart!"
    END IF
  CASE(0)
    !X1_a,X2_a and X1_b,X2_b already filled from parameter file readin...
    IF(init_average_axis)THEN
      CALL InitAverageAxis()
    END IF !init_average_axis
  CASE(1) !VMEC
    IF((init_BC.EQ.-1).OR.(init_BC.EQ.1))THEN ! compute  axis from VMEC, else use the one defined in paramterfile
      rhopos=0.0_wp
      ASSOCIATE(sin_range    => X1_base%f%sin_range, cos_range    => X1_base%f%cos_range )
      DO imode=cos_range(1)+1,cos_range(2)
        X1_a(iMode:iMode)= VMEC_EvalSplMode(X1_base%f%Xmn(:,iMode),0,(/rhopos/),Rmnc_Spl)
      END DO
      IF(lasym)THEN
        DO imode=sin_range(1)+1,sin_range(2)
          X1_a(iMode:iMode)=X1_a(iMode:imode)  +VMEC_EvalSplMode(X1_base%f%Xmn(:,iMode),0,(/rhopos/),Rmns_Spl)
        END DO
      END IF !lasym
      END ASSOCIATE !X1
      ASSOCIATE(sin_range    => X2_base%f%sin_range, cos_range    => X2_base%f%cos_range )
      DO imode=sin_range(1)+1,sin_range(2)
        X2_a(iMode:iMode)=VMEC_EvalSplMode(X2_base%f%Xmn(:,iMode),0,(/rhopos/),Zmns_Spl)
      END DO
      IF(lasym)THEN
        DO imode=cos_range(1)+1,cos_range(2)
          X2_a(iMode:iMode)=X2_a(iMode:iMode)  +VMEC_EvalSplMode(X2_base%f%Xmn(:,iMode),0,(/rhopos/),Zmnc_Spl)
        END DO
      END IF !lasym
      END ASSOCIATE !X2
    END IF
    IF((init_BC.EQ.-1).OR.(init_BC.EQ.0))THEN ! compute edge from VMEC, else use the one defined in paramterfile
      rhopos=1.0_wp
      ASSOCIATE(sin_range    => X1_base%f%sin_range, cos_range    => X1_base%f%cos_range )
      DO imode=cos_range(1)+1,cos_range(2)
        X1_b(iMode:iMode)=VMEC_EvalSplMode(X1_base%f%Xmn(:,iMode),0,(/rhopos/),Rmnc_Spl)
      END DO
      IF(lasym)THEN
        DO imode=sin_range(1)+1,sin_range(2)
          X1_b(iMode:iMode)=X1_b(iMode:iMode)  +VMEC_EvalSplMode(X1_base%f%Xmn(:,iMode),0,(/rhopos/),Rmns_Spl)
        END DO
      END IF !lasym
      END ASSOCIATE !X1
      ASSOCIATE(sin_range    => X2_base%f%sin_range, cos_range    => X2_base%f%cos_range )
      DO imode=sin_range(1)+1,sin_range(2)
        X2_b(iMode:iMode)=VMEC_EvalSplMode(X2_base%f%Xmn(:,iMode),0,(/rhopos/),Zmns_Spl)
      END DO
      IF(lasym)THEN
        DO imode=cos_range(1)+1,cos_range(2)
          X2_b(iMode:iMode)=X2_b(iMode:iMode)  +VMEC_EvalSplMode(X2_base%f%Xmn(:,iMode),0,(/rhopos/),Zmnc_Spl)
        END DO
      END IF !lasym
      END ASSOCIATE !X2
    END IF
    IF(init_average_axis)THEN
      CALL InitAverageAxis()
    END IF !init_average_axis
    IF(.NOT.init_fromBConly)THEN !only boundary and axis from VMEC
      ASSOCIATE(s_IP         => X1_base%s%s_IP, &
                nBase        => X1_base%s%nBase, &
                sin_range    => X1_base%f%sin_range,&
                cos_range    => X1_base%f%cos_range )
      DO imode=cos_range(1)+1,cos_range(2)
        X1_gIP(:,iMode)  =VMEC_EvalSplMode(X1_base%f%Xmn(:,iMode),0,s_IP,Rmnc_Spl)
      END DO !imode=cos_range
      IF(lasym)THEN
        DO imode=sin_range(1)+1,sin_range(2)
          X1_gIP(:,iMode)=X1_gIP(:,iMode)  +VMEC_EvalSplMode(X1_base%f%Xmn(:,iMode),0,s_IP,Rmns_Spl)
        END DO !imode= sin_range
      END IF !lasym
      END ASSOCIATE !X1
      ASSOCIATE(s_IP         => X2_base%s%s_IP, &
                nBase        => X2_base%s%nBase, &
                sin_range    => X2_base%f%sin_range,&
                cos_range    => X2_base%f%cos_range )
      DO imode=sin_range(1)+1,sin_range(2)
        X2_gIP(:,iMode)=VMEC_EvalSplMode(X2_base%f%Xmn(:,iMode),0,s_IP,Zmns_Spl)
      END DO !imode=sin_range
      IF(lasym)THEN
        DO imode=cos_range(1)+1,cos_range(2)
          X2_gIP(:,iMode)=X2_gIP(:,iMode)  +VMEC_EvalSplMode(X2_base%f%Xmn(:,iMode),0,s_IP,Zmnc_Spl)
        END DO !imode= sin_range
      END IF !lasym
      END ASSOCIATE !X2
      ASSOCIATE(s_IP         => LA_base%s%s_IP, &
                nBase        => LA_base%s%nBase, &
                sin_range    => LA_base%f%sin_range,&
                cos_range    => LA_base%f%cos_range )
      DO imode=sin_range(1)+1,sin_range(2)
        LA_gIP(:,iMode)=VMEC_EvalSplMode(LA_base%f%Xmn(:,iMode),0,s_IP,lmns_Spl)
      END DO !imode= sin_range
      IF(lasym)THEN
        DO imode=cos_range(1)+1,cos_range(2)
          LA_gIP(:,iMode)=LA_gIP(:,iMode)  +VMEC_EvalSplMode(LA_base%f%Xmn(:,iMode),0,s_IP,lmnc_Spl)
        END DO !imode=cos_range
      END IF !lasym
      END ASSOCIATE !LA
     END IF !fullIntVmec
  END SELECT !which_init


  IF((which_init_in.NE.-1).AND.(init_fromBConly))THEN
    !no restart(=-1) and initialization only
    !smoothly interpolate between  edge and axis data
    ASSOCIATE(s_IP         =>X1_base%s%s_IP, &
              modes        =>X1_base%f%modes, &
              zero_odd_even=>X1_base%f%zero_odd_even)
    DO imode=1,modes
      SELECT CASE(zero_odd_even(iMode))
      CASE(MN_ZERO,M_ZERO) !X1_a only used here!!
        X1_gIP(:,iMode)=(1.0_wp-(s_IP(:)**2))*X1_a(iMode)+(s_IP(:)**2)*X1_b(iMode)  ! meet edge and axis, ~(1-s^2)
      CASE(M_ODD_FIRST)
        X1_gIP(:,iMode)=s_IP(:)*X1_b(iMode)      ! first odd mode ~s
      CASE DEFAULT ! ~s^m
        X1_gIP(:,iMode)=s_IP(:)*X1_b(iMode)
        DO i_m=2,X1_base%f%Xmn(1,iMode)
          X1_gIP(:,iMode)=X1_gIP(:,iMode)*s_IP(:)
        END DO
      END SELECT !X1(:,iMode) zero odd even
    END DO !iMode
    END ASSOCIATE

    ASSOCIATE(s_IP         =>X2_base%s%s_IP, &
              modes        =>X2_base%f%modes, &
              zero_odd_even=>X2_base%f%zero_odd_even)
    DO imode=1,modes
      SELECT CASE(zero_odd_even(iMode))
      CASE(MN_ZERO,M_ZERO) !X2_a only used here!!!
        X2_gIP(:,iMode)=(1.0_wp-(s_IP(:)**2))*X2_a(iMode)+(s_IP(:)**2)*X2_b(iMode) ! meet edge and axis, ~(1-s^2)
      CASE(M_ODD_FIRST)
        X2_gIP(:,iMode)=s_IP(:)*X2_b(iMode)      ! first odd mode ~s
      CASE DEFAULT ! ~s^m
        X2_gIP(:,iMode)=s_IP(:)*X2_b(iMode)
        DO i_m=2,X2_base%f%Xmn(1,iMode)
          X2_gIP(:,iMode)=X2_gIP(:,iMode)*s_IP(:)
        END DO
      END SELECT !X2(:,iMode) zero odd even
    END DO
    END ASSOCIATE
  END IF !init_fromBConly

  !apply strong boundary conditions
  ASSOCIATE(modes        =>X1_base%f%modes, &
            zero_odd_even=>X1_base%f%zero_odd_even)
  DO imode=1,modes
    SELECT CASE(zero_odd_even(iMode))
    CASE(MN_ZERO,M_ZERO)
      BC_val =(/ X1_a(iMode)    ,      X1_b(iMode)/)
    !CASE(M_ODD_FIRST,M_ODD,M_EVEN)
    CASE DEFAULT
      BC_val =(/          0.0_wp,      X1_b(iMode)/)
    END SELECT !X1(:,iMode) zero odd even
    IF(which_init_in.NE.-1) U_init%X1(:,iMode)=X1_base%s%initDOF( X1_gIP(:,iMode) )
    CALL X1_base%s%applyBCtoDOF(U_init%X1(:,iMode),X1_BC_type(:,iMode),BC_val)
  END DO
  END ASSOCIATE !X1

  ASSOCIATE(modes        =>X2_base%f%modes, &
            zero_odd_even=>X2_base%f%zero_odd_even)
  DO imode=1,modes
    SELECT CASE(zero_odd_even(iMode))
    CASE(MN_ZERO,M_ZERO)
      BC_val =(/     X2_a(iMode),      X2_b(iMode)/)
    !CASE(M_ODD_FIRST,M_ODD,M_EVEN)
    CASE DEFAULT
      BC_val =(/          0.0_wp,      X2_b(iMode)/)
    END SELECT !X1(:,iMode) zero odd even
    IF(which_init_in.NE.-1) U_init%X2(:,iMode)=X2_base%s%initDOF( X2_gIP(:,iMode) )
    CALL X2_base%s%applyBCtoDOF(U_init%X2(:,iMode),X2_BC_type(:,iMode),BC_val)
  END DO
  END ASSOCIATE !X2

  ASSOCIATE(modes        =>LA_base%f%modes, &
            zero_odd_even=>LA_base%f%zero_odd_even)
  IF((which_init_in.NE.-1).AND.(.NOT.init_LA)) THEN
    !lambda init might not be needed since it has no boundary condition and changes anyway after the update of the mapping...
    IF(.NOT.init_fromBConly)THEN
      WRITE(UNIT_stdOut,'(4X,A)') "... lambda initialized with VMEC ..."
    ELSE
      WRITE(UNIT_stdOut,'(4X,A)') "... initialize lambda =0 ..."
      LA_gIP=0.0_wp
    END IF
  END IF !!which_init_in>-1 and not init_LA
  !always apply strong BC
  DO imode=1,modes
    IF(zero_odd_even(iMode).EQ.MN_ZERO)THEN
      U_init%LA(:,iMode)=0.0_wp ! (0,0) mode should not be here, but must be zero if its used.
    ELSE
      BC_val =(/ 0.0_wp, 0.0_wp/)
      IF(which_init_in.NE.-1) U_init%LA(:,iMode)=LA_base%s%initDOF( LA_gIP(:,iMode) )
      CALL LA_base%s%applyBCtoDOF(U_init%LA(:,iMode),LA_BC_type(:,iMode),BC_val)
    END IF!iMode ~ MN_ZERO
  END DO !iMode

  END ASSOCIATE !LA
END SUBROUTINE InitSolution