gray/src/splines.f90

! This module provides a high-level interface for creating and evaluating
! several kind of splines:
!
!   `spline_simple` is a simple interpolating cubic spline,
!   `spline_1d` and `spline_2d` are wrapper around the DIERCKX cubic splines,
!   `linear_1d` is a linear interpolation.
module splines
  use const_and_precisions, only : wp_

  implicit none

  ! A 1D interpolating cubic spline
  type spline_simple
    integer :: ndata                      ! Number of data points
    real(wp_), allocatable :: data(:)     ! Data points         (ndata)
    real(wp_), allocatable :: coeffs(:,:) ! Spline coefficients (ndata, 4)
    contains
      procedure :: init     => spline_simple_init
      procedure :: deinit   => spline_simple_deinit
      procedure :: eval     => spline_simple_eval
      procedure :: raw_eval => spline_simple_raw_eval
      procedure :: deriv    => spline_simple_deriv
  end type

  ! A 1D smoothing/interpolating cubic spline
  type spline_1d
    integer :: nknots                   ! Number of spline knots
    real(wp_), allocatable :: knots(:)  ! Knots positions
    real(wp_), allocatable :: coeffs(:) ! B-spline coefficients
    contains
      procedure :: init   => spline_1d_init
      procedure :: deinit => spline_1d_deinit
      procedure :: eval   => spline_1d_eval
      procedure :: deriv  => spline_1d_deriv
  end type

  ! A 2D smoothing/interpolating cubic spline s(x, y)
  type spline_2d
    integer :: nknots_x                   ! Number of x knots
    integer :: nknots_y                   ! Number of y knots
    real(wp_), allocatable :: knots_x(:)  ! Knots x positions
    real(wp_), allocatable :: knots_y(:)  ! Knots y positions
    real(wp_), allocatable :: coeffs(:)   ! B-spline coefficients

    ! B-spline coefficients of the partial derivatives
    type(pointer), allocatable :: partial(:,:)

    contains
      procedure :: init       => spline_2d_init
      procedure :: deinit     => spline_2d_deinit
      procedure :: eval       => spline_2d_eval
      procedure :: init_deriv => spline_2d_init_deriv
      procedure :: deriv      => spline_2d_deriv
      procedure :: transform  => spline_2d_transform
  end type

  ! Wrapper to store pointers in an array
  type pointer
    real(wp_), pointer :: ptr(:) => null()
  end type

  ! A simple 1D linear interpolation
  type linear_1d
    integer :: ndata                   ! Number of data points
    real(wp_), allocatable :: xdata(:) ! X data (ndata)
    real(wp_), allocatable :: ydata(:) ! Y data (ndata)
    contains
      procedure :: init     => linear_1d_init
      procedure :: deinit   => linear_1d_deinit
      procedure :: eval     => linear_1d_eval
      procedure :: raw_eval => linear_1d_raw_eval
  end type

  private
  public spline_simple, spline_1d, spline_2d, linear_1d

contains

  subroutine spline_simple_init(self, x, y, n)
    ! Initialises the spline

    ! subroutine arguments
    class(spline_simple),    intent(inout) :: self
    integer,                 intent(in)    :: n
    real(wp_), dimension(n), intent(in)    :: x, y

    call self%deinit
    self%ndata = n
    allocate(self%data(n))
    allocate(self%coeffs(n, 4))

    self%data = x
    call spline_simple_coeffs(x, y, n, self%coeffs)
  end subroutine spline_simple_init


  subroutine spline_simple_deinit(self)
    ! Deinitialises a simple_spline

    ! subroutine arguments
    class(spline_simple), intent(inout) :: self

    if (allocated(self%data))   deallocate(self%data)
    if (allocated(self%coeffs)) deallocate(self%coeffs)
    self%ndata = 0
  end subroutine spline_simple_deinit


  function spline_simple_eval(self, x) result(y)
    ! Evaluates the spline at x
    use utils, only : locate

    ! subroutine arguments
    class(spline_simple), intent(in) :: self
    real(wp_),            intent(in) :: x
    real(wp_)                        :: y

    ! local variables
    integer   :: i
    real(wp_) :: dx

    call locate(self%data, self%ndata, x, i)
    i = min(max(1, i), self%ndata - 1)
    dx = x - self%data(i)
    y = self%raw_eval(i, dx)
  end function spline_simple_eval


  function spline_simple_raw_eval(self, i, dx) result(y)
    ! Evaluates the i-th polynomial of the spline at dx

    ! subroutine arguments
    class(spline_simple), intent(in) :: self
    integer,              intent(in) :: i
    real(wp_),            intent(in) :: dx
    real(wp_)                        :: y

    y = self%coeffs(i,1) + dx*(self%coeffs(i,2) &
                         + dx*(self%coeffs(i,3) + dx*self%coeffs(i,4)))
  end function spline_simple_raw_eval


  function spline_simple_deriv(self, x) result(y)
    ! Computes the derivative of the spline at x
    use utils, only : locate


    ! subroutine arguments
    class(spline_simple), intent(in) :: self
    real(wp_),            intent(in) :: x
    real(wp_)                        :: y

    ! local variables
    integer   :: i
    real(wp_) :: dx

    call locate(self%data, self%ndata, x, i)
    i = min(max(1, i), self%ndata - 1)
    dx = x - self%data(i)
    y = self%coeffs(i,2) + dx*(2*self%coeffs(i,3) + 3*dx*self%coeffs(i,4))
  end function spline_simple_deriv


  subroutine spline_simple_coeffs(x, y, n, c)
    ! Computes the cubic coefficients of all n polynomials

    ! subroutine arguments
    integer,   intent(in)    :: n
    real(wp_), intent(in)    :: x(n), y(n)
    real(wp_), intent(inout) :: c(n*4)

    ! local variables
    integer :: jmp, i, ii, j, j1, j2, j3, k
    real(wp_) :: xb, xc, ya, yb, h, a, r, dya, dyb, dy2
    jmp = 1

    if (n <= 1) return

    ! initialisation
    xc = x(1)
    yb = y(1)
    h = 0
    a = 0
    r = 0
    dyb = 0
    dy2 = c(2)

    ! form the system of linear equations
    ! and eliminate subsequentially
    j = 1
    do i = 1, n
      j2 = n + i
      j3 = j2 + n
      a = h*(2 - a)
      dya = dyb + h*r
      if (i >= n) then
        ! set derivative dy2 at last point
        dyb = dy2
        h = 0
        dyb = dya
        goto 13
      else
        j  = j+jmp
        xb = xc
        xc = x(j)
        h  = xc-xb

        ! ii=0 - increasing abscissae
        ! ii=1 - decreasing abscissae
        ii = 0
        if (h == 0) return
        if (h < 0) ii = 1
        ya = yb
        yb = y(j)
        dyb = (yb - ya)/h

        if (i <= 1) then
          j1 = ii
          goto 13
        end if
      end if

      if (j1-ii /= 0) return
      a = 1 / (2*h + a)
      13 continue
      r = a*(dyb - dya)
      c(j3) = r
      a = h*a
      c(j2) = a
      c(i)  = dyb
    end do

    ! back substitution of the system of linear equations
    ! and computation of the other coefficients
    a  = 1
    j1 = j3+n+ii-ii*n
    i  = n
    do k = 1, n
      xb = x(j)
      h  = xc - xb
      xc = xb
      a  = a+h
      yb = r
      r  = c(j3)-r*c(j2)
      ya = 2*r
      c(j3) = ya + r
      c(j2) = c(i) - h*(ya+yb)
      c(j1) = (yb - r)/a
      c(i)  = y(j)
      a = 0
      j = j-jmp
      i = i-1
      j2 = j2-1
      j3 = j3-1
      j1 = j3+n+ii
    end do
  end subroutine spline_simple_coeffs


  subroutine spline_1d_init(self, x, y, n, range, weights, tension, err)
    ! Initialises a spline_1d.
    ! Takes:
    !   x: x data points
    !   y: y data points
    !   n: number of data points
    !   range: interpolation range as [x_min, x_max]
    !   weights: factors weighting the data points             (default: all 1)
    !   tension: parameter controlling the amount of smoothing (default: 0)
    ! Returns:
    !   err: error code of `curfit`
    use dierckx, only : curfit

    ! subroutine arguments
    class(spline_1d), intent(inout)         :: self
    real(wp_),        intent(in)            :: x(n)
    real(wp_),        intent(in)            :: y(n)
    integer,          intent(in)            :: n
    real(wp_),        intent(in)            :: range(2)
    real(wp_),        intent(in),  optional :: weights(n)
    real(wp_),        intent(in),  optional :: tension
    integer,          intent(out), optional :: err

    ! local variables
    integer   :: nknots_est  ! over-estimate of the number of knots
    real(wp_) :: residuals   ! sum of the residuals
    integer   :: work_int(n + 4)              ! integer working space
    real(wp_) :: work_real(4*n + 16*(n + 4))  ! real working space

    ! default values
    integer   :: err_def
    real(wp_) :: weights_def(n), tension_def

    weights_def = 1
    tension_def = 0
    if (present(weights)) weights_def = weights
    if (present(tension)) tension_def = tension

    ! clear memory, if necessary
    call self%deinit

    ! allocate the spline arrays
    nknots_est = n + 4
    allocate(self%knots(nknots_est), self%coeffs(nknots_est))

    call curfit(0, n, x, y, weights_def, range(1), range(2), 3, tension_def, &
                nknots_est, self%nknots, self%knots, self%coeffs, residuals, &
                work_real, size(work_real), work_int, err_def)

    if (present(err)) err = err_def
  end subroutine spline_1d_init


  subroutine spline_1d_deinit(self)
    ! Deinitialises a spline_1d

    class(spline_1d), intent(inout) :: self

    if (allocated(self%knots))  deallocate(self%knots)
    if (allocated(self%coeffs)) deallocate(self%coeffs)
    self%nknots = 0
  end subroutine spline_1d_deinit


  function spline_1d_eval(self, x) result(y)
    ! Evaluates the spline at x
    use dierckx, only : splev

    ! subroutine arguments
    class(spline_1d), intent(in) :: self
    real(wp_),        intent(in) :: x
    real(wp_)                    :: y

    ! local variables
    integer   :: err
    real(wp_) :: yv(1) ! because splev returns a vector

    call splev(self%knots, self%nknots, self%coeffs, 3, [x], yv, 1, err)
    y = yv(1)
  end function spline_1d_eval


  function spline_1d_deriv(self, x, order) result(y)
    ! Evaluates the spline n-th order derivative at x
    use dierckx, only : splder

    ! subroutine arguments
    class(spline_1d), intent(in)           :: self
    real(wp_),        intent(in)           :: x
    integer,          intent(in), optional :: order
    real(wp_)                              :: y

    ! local variables
    integer   :: err, n
    real(wp_) :: yv(1)              ! because splev returns a vector
    real(wp_) :: work(self%nknots)  ! working space array

    n = 1
    if (present(order)) n = order

    call splder(self%knots, self%nknots, self%coeffs, &
                3, n, [x], yv, 1, work, err)
    y = yv(1)
  end function spline_1d_deriv


  subroutine spline_2d_init(self, x, y, z, nx, ny, range, tension, err)
    ! Initialises a spline_2d.
    ! Takes:
    !   x, y: data points on a regular grid
    !   z: data points of z(x, y)
    !   n: number of data points
    !   range: interpolation range as [x_min, x_max, y_min, y_max]
    !   weights: factors weighting the data points             (default: all 1)
    !   tension: parameter controlling the amount of smoothing (default: 0)
    ! Returns:
    !   err: error code of `curfit`
    use dierckx, only : regrid

    ! subroutine arguments
    class(spline_2d), intent(inout)         :: self
    real(wp_),        intent(in)            :: x(nx)
    real(wp_),        intent(in)            :: y(ny)
    real(wp_),        intent(in)            :: z(nx * ny)
    integer,          intent(in)            :: nx, ny
    real(wp_),        intent(in)            :: range(4)
    real(wp_),        intent(in),  optional :: tension
    integer,          intent(out), optional :: err

    ! local variables
    integer   :: nknots_x_est  ! over-estimate of the number of knots
    integer   :: nknots_y_est  !
    real(wp_) :: residuals     ! sum of the residuals

    ! working space arrays
    integer   :: work_int(2*(ny + nx) + 11)
    real(wp_) :: work_real(15*(nx + ny) + ny*(nx + 4) + 92 + max(ny, nx + 4))

    ! default values
    integer   :: err_def
    real(wp_) :: tension_def

    tension_def = 0
    if (present(tension)) tension_def = tension

    ! clear memory, if necessary
    call self%deinit

    ! allocate the spline arrays
    nknots_x_est = nx + 4
    nknots_y_est = ny + 4
    allocate(self%knots_x(nknots_x_est), self%knots_y(nknots_y_est))
    allocate(self%coeffs(nx * ny))

    call regrid(0, nx, x, ny, y, z, range(1), range(2), range(3), range(4),  &
                3, 3, tension_def, nknots_x_est, nknots_y_est,               &
                self%nknots_x, self%knots_x, self%nknots_y, self%knots_y,    &
                self%coeffs, residuals, work_real, size(work_real),          &
                work_int, size(work_int), err_def)

    if (present(err)) err = err_def
  end subroutine spline_2d_init


  subroutine spline_2d_deinit(self)
    ! Deinitialises a spline_2d

    class(spline_2d), intent(inout) :: self

    if (allocated(self%knots_x)) deallocate(self%knots_x)
    if (allocated(self%knots_y)) deallocate(self%knots_y)
    if (allocated(self%coeffs))  deallocate(self%coeffs)

    ! Note: partial derivatives coeff. are pointers
    if (allocated(self%partial)) then
      deallocate(self%partial(1, 0)%ptr)
      deallocate(self%partial(0, 1)%ptr)
      deallocate(self%partial(1, 1)%ptr)
      deallocate(self%partial(2, 0)%ptr)
      deallocate(self%partial(0, 2)%ptr)
      deallocate(self%partial)
    end if

    self%nknots_x = 0
    self%nknots_y = 0
  end subroutine spline_2d_deinit


  function spline_2d_eval(self, x, y) result(z)
    ! Evaluates the spline at (x, y)
    use dierckx, only : fpbisp

    ! subroutine arguments
    class(spline_2d), intent(in) :: self
    real(wp_),        intent(in) :: x, y
    real(wp_)                    :: z

    ! local variables
    real(wp_) :: zv(1)  ! because fpbisp returns a vector

    integer   :: work_int(8)   ! integer working space
    real(wp_) :: work_real(8)  ! real working space

    ! Note: see https://netlib.org/dierckx/bispev.f for
    ! this apparently nonsensical invocation
    call fpbisp(self%knots_x, self%nknots_x,       &
                self%knots_y, self%nknots_y,       &
                self%coeffs, 3, 3, [x], 1, [y], 1, &
                zv, work_real(1), work_real(5),    &
                work_int(1), work_int(2))
    z = zv(1)
  end function spline_2d_eval


  subroutine spline_2d_init_deriv(self, p, q, n, m)
    ! Computes the spline coefficients of n-th partial derivative
    ! w.r.t x and m-th partial derivative w.r.t y on a grid of
    ! p×q points.
    !
    ! Note: for simplicity, only up to second-order is supported.
    use dierckx, only : coeff_parder

    ! subroutine arguments
    class(spline_2d), intent(inout) :: self
    integer,          intent(in)    :: p, q  ! grid dimensions
    integer,          intent(in)    :: n, m  ! derivative order

    ! local variables
    integer :: coeff_size
    integer :: err

    ! coeff. array (actually, the working space) size
    coeff_size = p*(4 - n) + q*(4 - m) + p*q

    ! allocate slots for storing the derivatives (first call only)
    if (.not. allocated(self%partial)) allocate(self%partial(0:2, 0:2))

    ! allocate the coefficients array
    allocate(self%partial(n, m)%ptr(coeff_size))

    ! compute the coefficients
    call coeff_parder(self%knots_x, self%nknots_x, &
                      self%knots_y, self%nknots_y, &
                      self%coeffs, 3, 3, n, m,     &
                      self%partial(n, m)%ptr, coeff_size, err)
  end subroutine spline_2d_init_deriv


  function spline_2d_deriv(self, x, y, n, m) result(z)
    ! Evaluates the spline n-th partial derivative w.r.t x
    ! and m-th partial derivative w.r.t y at (x, y)
    !
    ! Note: the coefficients of the derivative must have been
    ! initialised with init_deriv before calling this method.
    use dierckx, only : fpbisp

    ! subroutine arguments
    class(spline_2d), intent(in) :: self
    real(wp_), intent(in)        :: x, y
    integer,   intent(in)        :: n, m
    real(wp_)                    :: z

    ! local variables
    real(wp_), dimension(1)   :: zv   ! because splev returns a vector
    integer,   dimension(1)   :: work_int_x, work_int_y    ! integer working space
    real(wp_), dimension(1,4) :: work_real_x, work_real_y  ! real working space

    call fpbisp(self%knots_x(1 + n), self%nknots_x - 2*n,  &
                self%knots_y(1 + m), self%nknots_y - 2*m,  &
                self%partial(n, m)%ptr,                    &
                3 - n, 3 - m, [x], 1, [y], 1, zv,          &
                work_real_x, work_real_y, work_int_x, work_int_y)
    z = zv(1)
  end function spline_2d_deriv


  subroutine spline_2d_transform(self, a, b)
    ! Applies an affine transformation to the spline by
    ! a manipulation of its B-spline coefficients.
    ! If the spline is s(x,y), this maps s(x,y) → a⋅s(x,y) + b

    ! subroutine arguments
    class(spline_2d), intent(inout) :: self
    real(wp_),        intent(in)    :: a, b

    ! local variables
    integer :: i, j, m

    ! Note: since s(x,y) = Σ_ij c_ij B_i(x)B_j(y), the affine
    ! transformation simply acts on the spline coefficients as
    ! c_ij → a⋅c_ij + b

    ! Transform the spline
    self%coeffs = a * self%coeffs + b

    ! Transform the derivative splines too
    if (allocated(self%partial)) then

      ! Note: self%partial is a workspace array, only the
      ! first m elements contain the spline coefficients
      m = (self%nknots_x - 3 - 1) * (self%nknots_y - 3 - 1)

      do i = 0, size(self%partial, dim=1) - 1
      do j = 0, size(self%partial, dim=2) - 1
        if (associated(self%partial(i, j)%ptr)) then
          self%partial(i, j)%ptr(1:m) = a * self%partial(i, j)%ptr(1:m)
        end if
      end do
      end do

    end if
  end subroutine spline_2d_transform


  subroutine linear_1d_init(self, x, y, n)
    ! Initialises the spline

    ! subroutine arguments
    class(linear_1d),        intent(inout) :: self
    integer,                 intent(in)    :: n
    real(wp_), dimension(n), intent(in)    :: x, y

    call self%deinit
    self%ndata = n
    allocate(self%xdata(n))
    allocate(self%ydata(n))

    self%xdata = x
    self%ydata = y
  end subroutine linear_1d_init


  subroutine linear_1d_deinit(self)
    ! Deinitialises a linear_1d

    ! subroutine arguments
    class(linear_1d), intent(inout) :: self

    if (allocated(self%xdata)) deallocate(self%xdata)
    if (allocated(self%ydata)) deallocate(self%ydata)
    self%ndata = 0
  end subroutine linear_1d_deinit


  function linear_1d_eval(self, x) result(y)
    ! Evaluates the linear interpolated data at x
    use utils, only : locate

    ! subroutine arguments
    class(linear_1d), intent(in) :: self
    real(wp_),        intent(in) :: x
    real(wp_)                    :: y

    ! local variables
    integer :: i

    call locate(self%xdata, self%ndata, x, i)
    i = min(max(1, i), self%ndata - 1)
    y = self%raw_eval(i, x)
  end function linear_1d_eval


  function linear_1d_raw_eval(self, i, x) result(y)
    ! Performs the linear interpolation in the i-th interval

    ! subroutine arguments
    class(linear_1d), intent(in) :: self
    integer,          intent(in) :: i
    real(wp_),        intent(in) :: x
    real(wp_)                    :: y

    ! local variables
    real(wp_) :: dx, dy

    dx = self%xdata(i+1) - self%xdata(i)
    dy = self%ydata(i+1) - self%ydata(i)
    y = self%ydata(i) + dy/dx * (x - self%xdata(i))
  end function linear_1d_raw_eval

end module splines