s_lin_ols Subroutine

subroutine s_lin_ols(x, y, nd, nv, b0, b, r2, y_hat, se, cov_b)

! ==== Description
!! Multiple Linear Ordinary Least Squares (OLS) Regression with intercept.
!! NOTE: OLS could be wrapper for ridge (with lambda = 0 or presence checks
!! if mande an optional argument). However, it would increase computation
!! slightly and make code less readable. OLS is often used in teaching and
!! therefore, an easily readable standalone is kept.
!!
!! Computes:
!!  - Intercept b0 (scalar)
!!  - Predictor coefficients b(nv)
!!  - Coefficient of determination R²
!!  - Standard errors se(nv)
!!  - Covariance matrix of predictors covb(nv,nv)

! ==== Declarations
  integer(i4), intent(in)            :: nd                !! number of datapoints
  integer(i4), intent(in)            :: nv                !! number of predictors (excluding intercept)
  real(wp)   , intent(in)            :: x(nd,nv)          !! predictor data matrix (no intercept column)
  real(wp)   , intent(in)            :: y(nd)             !! response vector
  real(wp)   , intent(out)           :: b0                !! intercept coefficient
  real(wp)   , intent(out)           :: b(nv)             !! predictor coefficients
  real(wp)   , intent(out)           :: r2                !! coefficient of determination R²
  real(wp)   , intent(out), optional :: y_hat(nd)         !! predicted y values
  real(wp)   , intent(out), optional :: se(nv)            !! standard errors of predictor coefficients
  real(wp)   , intent(out), optional :: cov_b(nv,nv)      !! covariance matrix of predictor coefficients
  real(wp)                           :: x1(nd,nv+1)       !! matrix with intercept column
  real(wp)                           :: xt(nv+1,nd)       !! transpose of x1
  real(wp)                           :: xtx(nv+1,nv+1)    !! XᵗX matrix
  real(wp)                           :: xtx_i(nv+1,nv+1)  !! inverse of XᵗX
  real(wp)                           :: xty(nv+1)         !! Xᵗy
  real(wp)                           :: res(nd)           !! residuals
  real(wp)                           :: sse               !! sum of squared errors
  real(wp)                           :: sst               !! total sum of squares
  real(wp)                           :: y_bar             !! mean of y
  real(wp)                           :: ew(nv+1)          !! eigenvalues
  real(wp)                           :: ev(nv+1,nv+1)     !! eigenvectors
  real(wp)                           :: ew_diag(nv+1,nv+1)!! diag matrix of 1/λ
  integer(i4)                        :: i, j              !! loop counters

! ==== Instructions

  ! ---- validate input
  if (nd .le. nv + 1) then
     call s_err_print("[fsml error] OLS: Number of observations must&
                    & exceed number of predictors + intercept.")
     error stop
  endif

  ! ---- construct matrix with intercept column (first column = 1.0)
  do i = 1, nd
     x1(i,1) = 1.0_wp
     do j = 1, nv
        x1(i,j+1) = x(i,j)
     enddo
  enddo

  ! ---- eigen-decomposition for inversion of XᵗX

  ! compute transposed matrix and XᵗX
  xt   = transpose(x1)
  xtx  = matmul(xt, x1)

  ! eigendecomposition of s_pool
  call eigh(xtx, ew, vectors=ev)

  ! check result and return sentinel if needed
  do i = 1, nv+1
     if (ew(i) .le. 0.0_wp) then
        b0        = c_sentinel_r
        b(:)      = c_sentinel_r
        r2        = c_sentinel_r
        if (present(y_hat)) y_hat = c_sentinel_r
        if (present(se))       se = c_sentinel_r
        if (present(cov_b)) cov_b = c_sentinel_r
        return
     endif
  enddo

  ! construct diagonal matrix from inverted eigenvalues
  ew_diag = 0.0_wp
  do i = 1, nv+1
     ew_diag(i,i) = 1.0_wp / ew(i)
  enddo

  ! invert XᵗX
  xtx_i = matmul(ev, matmul(ew_diag, transpose(ev)))

  ! ---- compute coefficients: full vector including intercept
  xty = matmul(xt, y)
  b0  = 0.0_wp
  b   = 0.0_wp

  ! get intercept coefficient
  b0 = dot_product(xtx_i(1,:), xty)

  ! get predictor coefficients
  do j = 1, nv
     b(j) = dot_product(xtx_i(j+1, :), xty)
  enddo

  ! ---- predicted values and residuals
  y_hat = matmul(x1, [b0, b])  ! [b0, b] concatenates b0 and b
  res   = y - y_hat

  ! ---- R² calculation

  ! sum of squared errors
  sse = sum(res**2)

  ! mean of y
  y_bar = f_sts_mean_core(y)

  ! total sum of squares
  sst = sum((y - y_bar)**2)

  ! R²
  r2 = 1.0_wp - sse / sst

  ! ---- covariance matrix of full coefficients
  if (present(cov_b)) then
     cov_b = xtx_i(2:nv+1, 2:nv+1) * (sse / real(nd - nv - 1, kind=wp))
  endif

  ! ---- standard errors for predictors (exclude intercept)
  if (present(se)) then
     do i = 1, nv
        se(i) = sqrt(cov_b(i,i))
     enddo
  endif

end subroutine s_lin_ols