ceq.f90 Source File


Contents

Source Code


Source Code

!> Reads in the geometry using a CHEASE output file.
!! The CHEASE output file is read using the helper module
!! read_chease.
module ceq

  implicit none

  public :: ceq_init, ceqin, gradient, eqitem, bgradient, Hahm_Burrell, ceq_finish
  public :: invR,     initialize_invR
  public :: Rpos
  public :: Zpos
  public :: rcenter,  initialize_rcenter 
  public :: diameter, initialize_diameter
  public :: btori,    initialize_btori
  public :: dbtori,   initialize_dbtori
  public :: qfun,     initialize_q
  public :: pfun,     initialize_pressure
  public :: dpfun,    initialize_dpressure
  public :: betafun,  initialize_beta
  public :: psi,      initialize_psi
  public :: eqdbish

  public :: ceq_testing_type

  private

  integer :: nr, nt
  
  real, allocatable, dimension (:)     :: rho_d, eqpsi, psi_bar, fp, beta
  real, allocatable, dimension (:)     :: pressure, diam, rc, qsf, rho_b
  real, allocatable, dimension (:,:)   :: R_psi, Z_psi !, B_psi
  real, allocatable, dimension (:,:,:) :: drm, dzm, dbtm, dpm, dtm  !, dbm
  real, allocatable, dimension (:,:,:) :: dpcart, dbcart, dtcart, dbtcart
  real, allocatable, dimension (:,:,:) :: dpbish, dtbish, dbtbish

  real :: psi_0, psi_a, B_T, beta_0
  real :: R_mag, Z_mag, aminor
  
  logical :: init_rcenter = .true.
  logical :: init_diameter = .true.
  logical :: init_btori = .true.
  logical :: init_dbtori = .true.
  logical :: init_q = .true.
  logical :: init_pressure = .true.
  logical :: init_dpressure = .true.
  logical :: init_beta = .true.
  logical :: init_psi = .true.
  logical :: init_invR = .true.

  logical, parameter :: debug = .true.

  !> Expose private routines for testing only!
  type :: ceq_testing_type
  contains
    procedure, nopass :: alloc_arrays
    procedure, nopass :: chease_chi_index
    procedure, nopass :: dealloc_arrays
    procedure, nopass :: derm
    procedure, nopass :: eqdbish
    procedure, nopass :: eqdcart
    procedure, nopass :: mod2pi
    procedure, nopass :: set_ceq_from_chease
    procedure, nopass :: set_dpcart
    procedure, nopass :: set_drm
    procedure, nopass :: set_dzm
    procedure, nopass :: set_eqpsi
    procedure, nopass :: set_nrnt
    procedure, nopass :: set_R_psi
    procedure, nopass :: set_Z_psi
    procedure, nopass :: set_cart_arrays
    procedure, nopass :: set_bish_arrays
    procedure, nopass :: get_cart_arrays
    procedure, nopass :: get_bish_arrays
  end type ceq_testing_type

contains

  !> Set module-level `nt`, `nr`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine set_nrnt(nradius, ntheta)
    integer, intent(in) :: nradius, ntheta
    nr = nradius
    nt = ntheta
  end subroutine set_nrnt

  !> Set module-level `drm`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine set_drm(drm_in)
    real, dimension(nr,nt,2), intent(in) :: drm_in
    drm = drm_in
  end subroutine set_drm

  !> Set module-level `dzm`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine set_dzm(dzm_in)
    real, dimension(nr,nt,2), intent(in) :: dzm_in
    dzm = dzm_in
  end subroutine set_dzm

  !> Set module-level `dpcart`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine set_dpcart(dpcart_in)
    real, dimension(nr,nt,2), intent(in) :: dpcart_in
    dpcart = dpcart_in
  end subroutine set_dpcart

  !> Set module-level `eqpsi`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine set_eqpsi(eqpsi_in)
    real, dimension(nr), intent(in) :: eqpsi_in
    eqpsi = eqpsi_in
  end subroutine set_eqpsi

  !> Set module-level `R_psi`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine set_R_psi(R_psi_in)
    real, dimension(nr, nt), intent(in) :: R_psi_in
    R_psi = R_psi_in
  end subroutine set_R_psi

  !> Set module-level `Z_psi`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine set_Z_psi(Z_psi_in)
    real, dimension(nr, nt), intent(in) :: Z_psi_in
    Z_psi = Z_psi_in
  end subroutine set_Z_psi

  !> Set module-level `dpcart`, `dbtcart` and `dtcart`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine set_cart_arrays(dpcart_in, dbtcart_in, dtcart_in)
    real, dimension(nr, nt, 2), intent(in) :: dpcart_in
    real, dimension(nr, nt, 2), intent(in) :: dbtcart_in
    real, dimension(nr, nt, 2), intent(in) :: dtcart_in
    dpcart = dpcart_in
    dbtcart = dbtcart_in
    dtcart = dtcart_in
  end subroutine set_cart_arrays

  !> Set module-level `dpbish`, `dbtbish` and `dtbish`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine set_bish_arrays(dpbish_in, dbtbish_in, dtbish_in)
    real, dimension(nr, nt, 2), intent(in) :: dpbish_in
    real, dimension(nr, nt, 2), intent(in) :: dbtbish_in
    real, dimension(nr, nt, 2), intent(in) :: dtbish_in
    dpbish = dpbish_in
    dbtbish = dbtbish_in
    dtbish = dtbish_in
  end subroutine set_bish_arrays

  !> Get module-level `dpcart`, `dbtcart` and `dtcart`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine get_cart_arrays(dpcart_out, dbtcart_out, dtcart_out)
    real, dimension(nr, nt, 2), intent(out) :: dpcart_out
    real, dimension(nr, nt, 2), intent(out) :: dbtcart_out
    real, dimension(nr, nt, 2), intent(out) :: dtcart_out
    dpcart_out = dpcart
    dbtcart_out = dbtcart
    dtcart_out = dtcart
  end subroutine get_cart_arrays

  !> Get module-level `dpbish`, `dbtbish` and `dtbish`
  !! PURELY FOR TESTING
  !! FIXME: Remove when we have more control over module-level variables for testing
  subroutine get_bish_arrays(dpbish_out, dbtbish_out, dtbish_out)
    real, dimension(nr, nt, 2), intent(out) :: dpbish_out
    real, dimension(nr, nt, 2), intent(out) :: dbtbish_out
    real, dimension(nr, nt, 2), intent(out) :: dtbish_out
    dpbish_out = dpbish
    dbtbish_out = dbtbish
    dtbish_out = dtbish
  end subroutine get_bish_arrays

  !> Convert from theta-index which is `-pi` to `pi`, to chi-index, which
  !> is `0` to `2*pi.`
  !> Assumes `nchi` is even?
  function chease_chi_index(nchi,itheta)
    integer, intent(in) :: nchi, itheta
    integer :: chease_chi_index
    ! when ichi = 0, itheta = (ntheta-1)/2 + 1
    ! when itheta = 0, ichi = ntheta / 2
    ! itheta  1 2 3 4 5 6 7 8 
    ! ichi    5 6 7 8 1 2 3 4         
    if (itheta > nchi/2) then 
       chease_chi_index = itheta - nchi/2
    else
       chease_chi_index = itheta + nchi/2
    end if
  end function chease_chi_index

  !> FIXME : Add documentation  
  subroutine ceqin(eqfile, psi_0_out, psi_a_out, rmaj, B_T0, &
       avgrmid, initeq, in_nt, nthg) 
    use read_chease, only: read_infile
    implicit none
    character(len=*), intent(in) :: eqfile
    real, intent(out) :: psi_0_out, psi_a_out, rmaj, B_T0, avgrmid
    integer, intent(in) :: initeq
    logical, intent(in) :: in_nt 
    integer, intent(out) :: nthg
    real :: f_N, psi_N

    ! Exit early if already initialised
    if(initeq == 0) then
      ! Ensure `nthg` gets set if exiting early or it may be uninitialised in
      ! the calling procedure. The value it gets set to must be consistent
      ! with the value set at the end of this subroutine when not exiting
      ! early or the calling procedure may end up with inconsistently sized
      ! arrays.
      nthg = nt
      return
    endif
    
    if (.not.in_nt) then
       
       write (*,*)  'Reading CHEASE input file: ', trim(eqfile)
       !    Read the data
       call read_infile(trim(eqfile))

       call set_ceq_from_chease
    endif
                       
    !
    !     Normalize, rename quantities 
    !

    avgrmid = aminor
    ! rmaj changes depending on iflux.. FIX!!
    ! see geometr... for Miller is the rgeo for the flux surface of 
    ! interest... rgeo is rgeo of LCFS... check!!!
    ! actually when iflux = 1, it is the major radius of the magnetic axis,
    ! so below is OK.
    rmaj = R_mag 
    B_T0 = abs(B_T)

    psi_N = B_T0 * avgrmid**2
    psi_a = psi_a / psi_N
    psi_0 = psi_0 / psi_N
    psi_a_out = psi_a 
    psi_0_out = psi_0 
    eqpsi = eqpsi / psi_N

    f_N = B_T0*avgrmid
    fp=fp/f_N

    nthg=nt

    if (debug) then
       write (*,*) "Finished ceqin... imported CHEASE equilibrium"
       write (*,*) 'Some important quantities:'
       write (*,*) "aminor", aminor
       write (*,*) 'R_mag', R_mag
       write (*,*) 'B_T0', B_T0
       write (*,*) 'f_N', f_N
       write (*,*) 'nthg', nthg
       write (*,*) 'beta', beta_0
    end if
  end subroutine ceqin

  !> Set the ceq module-level variables from the read_chease module
  !> variables.
  !>
  !> This involves changing the theta grid from 0 to 2pi => -pi to pi,
  !> adding a point at theta == 0, and some normalisations
  subroutine set_ceq_from_chease
    use constants, only: pi
    use read_chease, only: npsi_chease, nchi_chease, b0exp_chease
    use read_chease, only: psi_chease, f_chease, q_chease, p_chease
    use read_chease, only: r_chease, z_chease

    real :: d, R_geo
    integer :: i, j

    !     nr == number of actual grid points in radial array
    nr = npsi_chease

    !     nt == number of theta grid points in theta eq grid
    nt = nchi_chease + 1

    call alloc_arrays(nr, nt)

    !     eqpsi(1:nr) == values of psi on the radial grid
    eqpsi = psi_chease

    !     fp(1:nr) == the function that satisfies
    !              B = fp grad zeta + grad zeta x grad psi
    fp = f_chease

    qsf = q_chease

    !     pressure(1:nr) == pressure profile on the radial grid,
    !     normalized to the value at the magnetic axis.
    pressure = p_chease

    !     psi_0 == value of psi at the magnetic axis
    !     psi_a == value of psi at the boundary of the plasma
    psi_0 = eqpsi(1)
    psi_a = eqpsi(nr)
    psi_bar = (eqpsi-psi_0)/(psi_a-psi_0)

    do j=1,nt
      do i=1,nr
        R_psi(i,j) = r_chease(i,chease_chi_index(nchi_chease, j))
      enddo
    enddo

    do j=1,nt
      do i=1,nr
        Z_psi(i,j) = z_chease(i,chease_chi_index(nchi_chease, j))
      enddo
    enddo

    !     aminor    == half diameter of last flux surface at elevation of Z_mag
    d = R_psi(nr,nt/2+1) - R_psi(nr,1)
    R_geo = (R_psi(nr,nt/2+1) + R_psi(nr, 1))/2.
    aminor = d/2.

    B_T = fp(nr)/R_geo

    R_psi = R_psi / aminor
    Z_psi = Z_psi / aminor

    !     R_mag == R position of magnetic axis
    !     Z_mag == Z position of magnetic axis
    R_mag = R_psi(1,1)
    Z_mag = Z_psi(1,1)

    beta = 8. * pi * 1.e-7 * pressure / b0exp_chease**2  ! not correct definition yet
    beta_0 = beta(1)
    pressure = pressure/pressure(1)
  end subroutine set_ceq_from_chease

  !> Allocate module-level arrays
  subroutine alloc_arrays(nr, nt)
    implicit none
    !> Number of points in radial direction
    integer, intent(in) :: nr
    !> Number of points in theta direction
    integer, intent(in) :: nt

    allocate(rho_d(nr), eqpsi(nr), psi_bar(nr), fp(nr), beta(nr), pressure(nr), &
         rc(nr), diam(nr), qsf(nr), rho_b(nr))
    rho_d = 0. ; eqpsi = 0. ; psi_bar = 0. ; fp = 0. ; beta = 0. ; pressure = 0. 
    rc = 0. ; diam = 0. ; qsf = 0. ; rho_b = 0.
    allocate(R_psi(nr, nt), Z_psi(nr, nt))
    R_psi = 0.  ; Z_psi = 0.
    allocate(drm(nr, nt, 2), dzm(nr, nt, 2), dbtm(nr, nt, 2), &
         dpm(nr, nt, 2), dtm(nr, nt, 2))
    drm = 0. ; dzm = 0. ; dbtm = 0. ; dpm = 0. ; dtm = 0.
    allocate(dpcart(nr, nt, 2), dtcart(nr, nt, 2), dbtcart(nr, nt, 2))
    allocate(dpbish(nr, nt, 2), dtbish(nr, nt, 2), dbtbish(nr, nt, 2))

    dpcart = 0. ; dtcart = 0. ; dbtcart = 0. 
    dpbish = 0. ; dtbish = 0. ; dbtbish = 0.

  end subroutine alloc_arrays

  !> Deallocate module-level arrays
  subroutine dealloc_arrays
    implicit none
    if(allocated(rho_d)) deallocate(rho_d,eqpsi,psi_bar,fp,beta,pressure,rc,diam,qsf,rho_b)
    if(allocated(R_psi)) deallocate(R_psi,Z_psi)
    if(allocated(drm)) deallocate(drm,dzm,dbtm,dpm,dtm)
    if(allocated(dpcart)) deallocate(dpcart,dtcart,dbtcart)
    if(allocated(dpbish)) deallocate(dpbish,dtbish,dbtbish)
  end subroutine dealloc_arrays

  !> Finish this module; deallocate arrays
  subroutine ceq_finish
    use read_chease, only : finish
    implicit none
    call dealloc_arrays
    call finish
  end subroutine ceq_finish

  !> Initialise CHEASE equilibrium module
  !!
  !! Assumes data has already been read in
  subroutine ceq_init
    use constants, only: pi
    implicit none
    real, dimension(nr,nt) :: eqpsi1, eqth, eqbtor
    integer i, j
   
    do j=1,nt
       do i=1,nr
          eqbtor(i,j) = fp(i)/R_psi(i,j)
          eqpsi1(i,j) = eqpsi(i)
       enddo
    enddo
    
    do j=1,nt
       eqth(:,j) = (j-1)*2.*pi/real(nt-1)-pi
    enddo

    call derm(eqth,   dtm,  'T')
    call derm(R_psi,  drm,  'E')
    call derm(Z_psi,  dzm,  'O')
    call derm(eqbtor, dbtm, 'E')
    call derm(eqpsi1, dpm,  'E')

    ! grad(psi) in cartesian form 
    call eqdcart(dpm, dpcart)
    ! grad(psi) in Bishop form 
    call eqdbish(dpcart, dpbish)

    ! grad(BT) in cartesian form
    call eqdcart(dbtm, dbtcart)
    ! grad(BT) in Bishop form
    call eqdbish(dbtcart, dbtbish)

    ! grad(theta) in cartesian form
    call eqdcart(dtm, dtcart)
    ! grad(theta) in Bishop form
    call eqdbish(dtcart, dtbish)

  end subroutine ceq_init

  !> Calculate the derivative of f w.r.t to the radial
  !! and poloidal indexes (i.e. calculate the finite 
  !! differences without dividing by 
  !! delta psi and delta theta).
  !! - dfm(:,:,1) is the psi derivative
  !! - dfm(:,:,2) is the theta derivative
  !! - char specifies the periodicity in theta
  !!   - 'E', 'O' mean continuous at theta = +/- pi
  !!   - 'T' means a jump of 2 pi at theta = +/- pi
  subroutine derm(f, dfm, char)
    use constants, only: pi
    implicit none
    integer i, j
    character(1), intent(in) :: char
    real, intent(in) :: f(:,:)
    real, intent(out) :: dfm(:,:,:)
    
    i=1
    ! Radial derivative
    dfm(i,:,1) = -0.5*(3*f(i,:)-4*f(i+1,:)+f(i+2,:))         
    
    i=nr
    ! Radial derivative
    dfm(i,:,1) = 0.5*(3*f(i,:)-4*f(i-1,:)+f(i-2,:))
   
    select case (char)
    case ('E') 
       j=1
       ! theta derivative, periodic
       dfm(:,j,2) = 0.5*(f(:,j+1)-f(:,nt-1))
       
       j=nt      
       dfm(:,j,2) = 0.5*(f(:,2)-f(:,j-1))
    case ('O')
       j=1
       dfm(:,j,2) = 0.5*(f(:,j+1)-f(:,nt-1))
       
       j=nt      
       dfm(:,j,2) = 0.5*(f(:,2)-f(:,j-1))
    case ('T')
       j=1
       dfm(:,j,2) = 0.5*(f(:,j+1)-f(:,nt-1)+2.*pi)
          
       j=nt
       dfm(:,j,2) = 0.5*(2.*pi + f(:,2) - f(:,j-1))
    end select
    
    do i=2,nr-1
       dfm(i,:,1)=0.5*(f(i+1,:)-f(i-1,:))
    enddo
    
    do j=2,nt-1
       dfm(:,j,2)=0.5*(f(:,j+1)-f(:,j-1))
    enddo

  end subroutine derm

  !> Calculate the derivative of a selected field w.r.t. R and Z
  subroutine gradient(rgrid, theta, grad, char, rp, nth_used, ntm)
    use splines, only: inter_d_cspl
    implicit none
    !> Number of theta points
    integer, intent(in) :: nth_used
    !> Lower index of theta array
    integer, intent(in) :: ntm
    !> Select field to take gradient of:
    !>   - 'D': Bt
    !>   - 'P': psi
    !>   - 'R': pressure
    !>   - 'T': theta
    !>   - 'B': for debugging with `bishop==0` only
    character(1), intent(in) :: char
    !> Minor radius grid of flux surface (?)
    real, intent(in) :: rgrid(-ntm:)
    !> Theta grid
    real, intent(in) :: theta(-ntm:)
    !> grad(:,1): derivative w.r.t. R as function of theta
    !> grad(:,2): derivative w.r.t  Z as function of theta
    real, intent(out) :: grad(-ntm:,:)
    !> Radius for taking pressure derivative (FIXME: double-check)
    real, intent(in) :: rp
    real :: aa(1), daa(1), rpt(1)
    real, dimension(nr,nt,2) :: dcart
    integer :: i
    
    select case(char)
    case('B') 
       dcart = dbcart
    case('D')  ! diagnostic 
       dcart = dbtcart
    case('P') 
       dcart = dpcart
    case('R') 
       dcart = dpcart  ! dpcart is correct for 'R'
    case('T')
       dcart = dtcart
    end select
    
    do i=-nth_used,nth_used
       call eqitem(rgrid(i), theta(i), dcart(:,:,1), grad(i,1))
       call eqitem(rgrid(i), theta(i), dcart(:,:,2), grad(i,2))
    enddo

    !     to get grad(pressure), multiply grad(psi) by dpressure/dpsi
    if(char == 'R') then
       rpt(1) = rp
       call inter_d_cspl(nr, eqpsi, pressure, 1, rpt, aa, daa)
       do i=-nth_used, nth_used
          grad(i,1)=grad(i,1)*daa(1) * 0.5*beta_0
          grad(i,2)=grad(i,2)*daa(1) * 0.5*beta_0
       enddo
    endif

  end subroutine gradient

  !> Calculate the derivative of a selected field w.r.t. the Bishop
  !> coordinate system
  subroutine bgradient(rgrid, theta, grad, char, rp, nth_used, ntm)
    use mp, only: mp_abort
    use splines, only: inter_d_cspl
    implicit none
    !> Number of theta points
    integer, intent(in) :: nth_used
    !> Lower index of theta array
    integer, intent(in) :: ntm
    !> Select field to take gradient of:
    !>   - 'D': Bt
    !>   - 'P': psi
    !>   - 'R': pressure
    !>   - 'T': theta
    !>   - 'B': for debugging with `bishop==0` only
    character(1), intent(in) :: char
    !> Minor radius grid of flux surface (?)
    real, intent(in) :: rgrid(-ntm:)
    !> Theta grid
    real, intent(in) :: theta(-ntm:)
    !> Computed derivative, second dimension is the Bishop coordinate
    real, intent(out) :: grad(-ntm:,:)
    !> Radius for taking pressure derivative (?)
    real, intent(in) :: rp
    real :: aa(1), daa(1), rpt(1)
    real, dimension(nr,nt,2) ::  dbish
    integer :: i

    select case(char)
    case('B') 
       call mp_abort('error: bishop = 1 not allowed with ceq. (2)',.true.)
    case('D')  ! diagnostic
       dbish = dbtbish
    case('P') 
       dbish = dpbish
    case('R') 
       dbish = dpbish  ! dpcart is correct for 'R'
    case('T')
       dbish = dtbish
    end select

    do i=-nth_used,nth_used
       call eqitem(rgrid(i), theta(i), dbish(:,:,1), grad(i,1))
       call eqitem(rgrid(i), theta(i), dbish(:,:,2), grad(i,2))
    enddo

    !     to get grad(pressure), multiply grad(psi) by dpressure/dpsi
    if(char == 'R') then
       rpt(1) = rp
       call inter_d_cspl(nr, eqpsi, pressure, 1, rpt, aa, daa)
       do i=-nth_used, nth_used
          grad(i,1)=grad(i,1)*daa(1) * 0.5*beta_0
          grad(i,2)=grad(i,2)*daa(1) * 0.5*beta_0
       enddo
    endif

  end subroutine bgradient

  !> Calculates fstar which is f interpolated at the location (r,theta).
  !! Here r is the normalised poloidal flux coordinate rp (= psi_pN + psi_0N)
  !! and theta_in is theta. f is a grid of values of f as a function of psi_p,theta
  subroutine eqitem(r, theta_in, f, fstar)
    use mp, only: mp_abort
    use constants, only: pi
    real, intent(in) :: r, theta_in
    real, intent(out) :: fstar
    real, dimension(:,:), intent(inout) :: f
    integer :: i, j, istar, jstar
    real :: thet, sign
    real :: st, dt, sr, dr
    real, dimension(size(f,2)) :: mtheta

    ! check for axis evaluation
    if(r == eqpsi(1)) then
       write(*,*) 'no evaluation at axis allowed in eqitem'
       write(*,*) r, theta_in, eqpsi(1)
       call mp_abort('no evaluation at axis allowed in eqitem')
    endif
    
    ! allow psi(r) to be a decreasing function
    sign=1.
    if(eqpsi(2) < eqpsi(1)) sign=-1.
    
    if(r < sign*eqpsi(1)) then
       write(*,*) 'r < Psi_0 in eqitem'
       write(*,*) r,sign,eqpsi(1)
       call mp_abort('r < Psi_0 in eqitem')
    endif
      
    ! find r on psi mesh
    
    ! disallow evaluations on or outside the surface for now
    if(r >= eqpsi(nr)) then
       write(*,*) 'No evaluation of eqitem allowed on or outside surface'
       write(*,*) '(Could this be relaxed a bit?)'
       write(*,*) r, theta_in, eqpsi(nr), sign
       call mp_abort('No evaluation of eqitem allowed on or outside surface')
    endif
    
    istar=0
    do i=2,nr
       if(r < sign*eqpsi(i)) then
          dr = r - sign*eqpsi(i-1)
          sr = sign*eqpsi(i) - r
          istar=i-1
          exit
       endif
    enddo
    
    ! no gradients available at axis, so do not get too close
    if(istar == 1) then
       write(*,*) 'Too close to axis in eqitem'
       write(*,*) r, theta_in, eqpsi(1), eqpsi(2)
       call mp_abort('Too close to axis in eqitem')
    endif
  
    ! Now do theta direction
    thet = mod2pi(theta_in)

    ! get thet on theta mesh
    do j=1,nt
       mtheta(j)=(j-1)*2.*pi/real(nt-1)-pi
    enddo
       
    ! note that theta(1)=0 for gen_eq theta 
    jstar=-1
    do j=1,nt
       if(jstar /= -1) cycle
       if(thet < mtheta(j)) then
          dt = thet - mtheta(j-1)
          st = mtheta(j) - thet
          jstar=j-1
       endif
    enddo
      
    ! treat pi separately
    
    if(jstar == -1) then
       jstar=nt-1
       dt=mtheta(jstar+1)-mtheta(jstar)
       st=0.
    endif

    ! use opposite area stencil to interpolate
    fstar=f(istar    , jstar    ) * sr * st &
         +f(istar + 1, jstar    ) * dr * st &
         +f(istar    , jstar + 1) * sr * dt &
         +f(istar + 1, jstar + 1) * dr * dt
    fstar=fstar &
         /abs(eqpsi(istar+1)-eqpsi(istar)) &
         /(mtheta(jstar+1)-mtheta(jstar))
  end subroutine eqitem

  !> Converts derivatives w.r.t. (psi_index,theta_index) to derivatives
  !! w.r.t. (R,Z).
  !! Note that `dfcart(1, :, :)` is not valid
  subroutine eqdcart(dfm, dfcart)    
    implicit none
    !> dfm(:,:,1) is deriv w.r.t. psi_index (i.e. (df/dpsi)_theta * delta psi);
    !> dfm(:,:,2) is deriv w.r.t. theta_index
    real, dimension (:,:,:), intent(in)  :: dfm
    !> dfcart(:,:,1) is deriv w.r.t. R;
    !> dfcart(:,:,2) is deriv w.r.t. Z
    real, dimension (:,:,:), intent(out) :: dfcart
    real, dimension (size(dfm,1),size(dfm,2)) :: denom
    integer :: i, j
      
    denom(:,:) = drm(:,:,1)*dzm(:,:,2) - drm(:,:,2)*dzm(:,:,1)
    
    dfcart(:,:,1) =   dfm(:,:,1)*dzm(:,:,2) - dzm(:,:,1)*dfm(:,:,2)
    dfcart(:,:,2) = - dfm(:,:,1)*drm(:,:,2) + drm(:,:,1)*dfm(:,:,2)

    ! Inner loop skips first point because it's likely exactly on the
    ! magnetic axis, which means the Jacobian is zero. Do we need
    ! need a better way to either calculate dfcart there, or to ensure
    ! we ignore it in other places?
    do j=1,nt
       do i=2,nr
          dfcart(i,j,:)=dfcart(i,j,:)/denom(i,j)
       enddo
    enddo    

  end subroutine eqdcart

  !> Convert gradients of a function f w.r.t. R,Z into bishop form
  !! Note that `dbish(1, :, :)` is not valid
  subroutine eqdbish(dcart, dbish)
    implicit none
    !> dcart(:,:,1) is gradient of f w.r.t. R;
    !> dcart(:,:,2) is gradient of f w.r.t. Z
    real, dimension(:, :, :), intent (in) :: dcart
    !> dbish(:,:,1) is set to (df/dR dpsi/dR + df/dZ dpsi/dZ)/|grad psi|;
    !> dbish(:,:,2) is set to (-df/dR dpsi/dZ + df/dZ dpsi/dR)/|grad psi|.
    !> Note that in the special case where f=psi:
    !> dbish(:,:,1) is |grad psi|;
    !> dbish(:,:,2) is 0
    real, dimension(:, :, :), intent(out) :: dbish
    real, dimension(size(dcart,1),size(dcart,2)) :: denom
    integer :: i, j

    ! denom is |grad psi|
    denom(:,:) = sqrt(dpcart(:,:,1)**2 + dpcart(:,:,2)**2)

    dbish(:,:,1) = dcart(:,:,1)*dpcart(:,:,1) + dcart(:,:,2)*dpcart(:,:,2)
    dbish(:,:,2) =-dcart(:,:,1)*dpcart(:,:,2) + dcart(:,:,2)*dpcart(:,:,1)
    
    ! Inner loop skips first point because it's likely exactly on the
    ! magnetic axis, which means the Jacobian is zero. Do we need
    ! need a better way to either calculate dfcart there, or to ensure
    ! we ignore it in other places?
    do j=1,nt
       do i=2,nr
          dbish(i,j,:) = dbish(i,j,:)/denom(i,j)
       enddo
    enddo

  end subroutine eqdbish

  !> If `init == 1`, set flag to initialize invR spline when
  !> `invR` is called
  !! FIXME: this _always_ returns 1, change to subroutine
  function initialize_invR (init) 
    implicit none
    integer :: init, initialize_invR
    
    init_invR = (init == 1)
    initialize_invR = 1
  end function initialize_invR

  !> Return the inverse of the major radius at \((r, \theta)\)
  function invR (r, theta)
    implicit none
    real, intent (in) :: r, theta
    real :: invR
    
    invR = 1./Rpos(r, theta)

  end function invR

  !> Return the major radius at \((r, \theta)\)
  function Rpos (r, theta)
    implicit none
    real, intent (in) :: r, theta
    real :: f, Rpos
    real :: th
    
    th = mod2pi( theta)
    
    call eqitem(r, th, R_psi, f)
    Rpos=f
    
  end function Rpos

  !> Return the vertical coordinate at \((r, \theta)\)
  function Zpos (r, theta)
    implicit none
    real, intent (in) :: r, theta
    real :: f, Zpos
    real :: th
    
    th = mod2pi( theta)
    
    call eqitem(r, th, Z_psi, f)
    Zpos=f
    
  end function Zpos

  !> If `init == 1`, set flag to initialize psi spline when
  !> `psi` is called
  !! FIXME: this _always_ returns 1, change to subroutine
  function initialize_psi (init) 
    implicit none
    integer, intent(in) :: init
    integer :: initialize_psi
    
    init_psi = (init == 1)
    initialize_psi = 1
  end function initialize_psi

  !> Return the flux surface label `psi`
  function psi (r)
    implicit none
    real, intent (in) :: r
    real :: psi

    psi = r

  end function psi

  !> Put `theta` into the range \([-\pi, \pi]\)
  function mod2pi (theta)
    use constants, only: pi
    implicit none
    real, intent(in) :: theta
    real :: th, mod2pi
    logical :: out
    
    if(theta <= pi .and. theta >= -pi) then
       mod2pi = theta
       return
    endif
    
    th=theta
    out=.true.
    do while(out)
       if(th > pi) th = th - 2.*pi
       if(th <-pi) th = th + 2.*pi
       if(th <= pi .and. th >= -pi) out=.false.
    enddo
    mod2pi=th
  end function mod2pi

  !> If `init == 1`, set flag to initialize diameter spline when
  !> `diameter` is called
  !! FIXME: this _always_ returns 1, change to subroutine
  function initialize_diameter (init)
    implicit none
    integer, intent(in) :: init
    integer :: initialize_diameter

    init_diameter = (init == 1)
    initialize_diameter = 1

  end function initialize_diameter

  !> Return the diameter of the flux surface at a given major radius
  function diameter (rp)
    use splines, only: new_spline, splint, spline
    implicit none
    real, intent(in) :: rp
    real :: diameter
    type (spline), save :: spl

    if(init_diameter) then
       diam(:) = abs(R_psi(:,nt/2 + 1) - R_psi(:,1))
       call new_spline(nr, eqpsi, diam, spl)
       init_diameter = .false.
    endif

    diameter = splint(rp, spl)

  end function diameter

  !> If `init == 1`, set flag to initialize rcenter spline when
  !> `rcenter` is called
  !! FIXME: this _always_ returns 1, change to subroutine
  function initialize_rcenter (init)
    implicit none
    integer, intent(in) :: init
    integer :: initialize_rcenter

    init_rcenter = (init == 1)
    initialize_rcenter = 1
  end function initialize_rcenter

  !> Return the major radius of the centre of a given flux surface
  function rcenter (rp)
    use splines, only: new_spline, splint, spline
    implicit none
    real, intent(in) :: rp
    real :: rcenter
    type (spline), save :: spl

    if(init_rcenter) then
       rc(:) = 0.5*(R_psi(:,1) + R_psi(:,nt/2 + 1))
       call new_spline(nr, eqpsi, rc, spl)
       init_rcenter = .false.
    endif

    rcenter = splint(rp, spl)

  end function rcenter

  !> If `init == 1`, set flag to initialize dbtori spline when
  !> `dbtori` is called
  !! FIXME: this _always_ returns 1, change to subroutine
  function initialize_dbtori (init) 
    implicit none
    integer, intent(in) :: init
    integer :: initialize_dbtori

    init_dbtori = (init == 1)
    initialize_dbtori = 1
  end function initialize_dbtori

  !> \(\frac{dI}{d\psi}\) at the given normalised \(\psi\)
  function dbtori (pbar)
    use splines, only: new_spline, dsplint, spline
    implicit none
    real, intent(in) :: pbar
    real :: dbtori
    type (spline), save :: spl

    if(init_dbtori) call new_spline(nr, psi_bar, fp, spl)
    init_dbtori=.false.

    dbtori = dsplint(pbar, spl)/(psi_a-psi_0)

  end function dbtori

  !> If `init == 1`, set flag to initialize btori spline when
  !> `btori` is called
  !! FIXME: this _always_ returns 1, change to subroutine
  function initialize_btori (init)
    implicit none
    integer, intent(in) :: init
    integer :: initialize_btori

    init_btori = (init == 1)
    initialize_btori = 1
  end function initialize_btori
  
  !> FIXME : Add documentation
  function btori (pbar)
    use splines, only: new_spline, splint, spline
    implicit none
    real, intent(in) :: pbar
    real :: btori
    type (spline), save :: spl

    if(init_btori) call new_spline(nr, psi_bar, fp, spl)
    init_btori=.false.

    btori = splint(pbar, spl)
  end function btori

  !> If `init == 1`, set flag to initialize `q` spline when
  !> `qfun` is called
  !! FIXME: this _always_ returns 1, change to subroutine
  function initialize_q (init) 
    implicit none
    integer, intent(in) :: init
    integer :: initialize_q

    init_q = (init == 1)
    initialize_q = 1
  end function initialize_q

  !> FIXME : Add documentation
  function qfun (pbar)
    use splines, only: new_spline, splint, spline
    implicit none
    real, intent(in) :: pbar
    real :: qfun
    type (spline), save :: spl

    if(init_q) call new_spline(nr, psi_bar, qsf, spl)
    init_q = .false.

    qfun = splint(pbar, spl)

  end function qfun

  !> If `init == 1`, set flag to initialize pressure spline when
  !> `pfun` is called
  !! FIXME: this _always_ returns 1, change to subroutine
  function initialize_pressure (init) 
    implicit none
    integer, intent(in) :: init
    integer :: initialize_pressure

    init_pressure = (init == 1)
    initialize_pressure = 1
  end function initialize_pressure

  !> FIXME : Add documentation
  function pfun (pbar) 
    use splines, only: new_spline, splint, spline
    implicit none
    real, intent(in) :: pbar
    real :: pfun
    type (spline), save :: spl

    if(init_pressure) call new_spline(nr, psi_bar, pressure, spl)
    init_pressure = .false.
    !
    ! p_N would be B**2/mu_0 => p = beta/2 in our units
    !
    pfun = 0.5*beta_0*splint(pbar, spl)

  end function pfun

  !> If `init == 1`, set flag to initialize dpressure spline when
  !> `dpfun` is called
  !! FIXME: this _always_ returns 1, change to subroutine
  function initialize_dpressure (init) 
    implicit none
    integer, intent(in) :: init
    integer :: initialize_dpressure

    init_dpressure = (init == 1)
    initialize_dpressure = 1
  end function initialize_dpressure

  !> FIXME : Add documentation  
  function dpfun (pbar)   
    use splines, only: new_spline, dsplint, spline
    implicit none
    real, intent(in) :: pbar
    real :: dpfun
    type (spline), save :: spl
    !
    ! p_N would be B**2/mu_0 => p = beta/2 in our units
    !
    if(init_dpressure) then
       call new_spline(nr, psi_bar, pressure, spl)
       init_dpressure = .false.
    endif

    dpfun = dsplint(pbar, spl)/(psi_a-psi_0) * beta_0/2.

  end function dpfun

  !> If `init == 1`, set flag to initialize beta spline when
  !> `betafun` is called
  !! FIXME: this _always_ returns 1, change to subroutine
  function initialize_beta (init)
    implicit none
    integer, intent(in) :: init
    integer :: initialize_beta

    init_beta = (init == 1)
    initialize_beta = 1
  end function initialize_beta

  !> FIXME : Add documentation  
  function betafun (pbar)
    use splines, only: new_spline, splint, spline
    implicit none
    real, intent(in) :: pbar
    real :: betafun
    type (spline), save :: spl

    if(pbar == 0.) then
       betafun=beta(1)
       return
    endif

    if(init_beta) call new_spline(nr, psi_bar, beta, spl)
    init_beta = .false.

    betafun = splint(pbar, spl)

  end function betafun

  !> FIXME : Add documentation  
  subroutine Hahm_Burrell(irho, a) 
    use file_utils, only: open_output_file, close_output_file, get_unused_unit
    implicit none
    real, intent(in) :: a
    integer, intent(in) :: irho
    integer :: i
    real :: gradpsi, mag_B, rho_eq, rp1, rp2, rho1, rho2, drhodpsiq
    real, dimension(nr) :: gamma, pbar, dp, d2p, pres
    integer :: fort24_unit

    gamma = 0.

    pbar = (eqpsi-eqpsi(1))/(eqpsi(nr)-eqpsi(1))

    do i=2, nr-1
       dp(i)  = dpfun(pbar(i))
    enddo

    do i=3,nr-2
       d2p(i) = (dp(i+1)-dp(i-1))/(eqpsi(i+1)-eqpsi(i-1))
    enddo

    pres=0.
    do i=3, nr-2
       rp1=eqpsi(i+1)
       rp2=eqpsi(i-1)
       rho1=0.5*diameter(rp1)
       rho2=0.5*diameter(rp2)
       drhodpsiq=(rho1-rho2)/(rp1-rp2)
       
       call eqitem(eqpsi(i), 0., dpbish(:,:,1), gradpsi)
       
       mag_B=sqrt( (btori(pbar(i))**2 + gradpsi**2))/Rpos(eqpsi(i),0.)

       pres(i) = pfun(pbar(i))

       gamma(i) = 0.01*gradpsi**2*(d2p(i)/pres(i)-a*(dp(i)/pres(i))**2) &
            /mag_B*(2.*pres(i)/beta_0)**((1-a)/2.) &
            *(-pres(i)/(dp(i)/drhodpsiq))            
    enddo
    
    call get_unused_unit(fort24_unit)
    call open_output_file(fort24_unit,".fort.24")
    do i=3,nr-2
       if(irho == 1) then
          rho_eq = eqpsi(i)
       else if(irho == 2) then
          rho_eq = 0.5 * diameter(eqpsi(i))
       else if(irho == 3) then
          rho_eq = pbar(i)
       endif
       write(fort24_unit,1000) i, pbar(i), rho_eq, pres(i), gamma(i)
    enddo

    if(irho == 1) write(fort24_unit,*) '# irho = 1 produces psi instead of rho_eq'

    call close_output_file(fort24_unit)

1000 format(i5,11(1x,e16.9))

  end subroutine Hahm_Burrell
end module ceq