compute_cce Function

  function compute_cce(this, predicted, expected) result(output)
    !! Compute the categorical cross entropy loss of a model
    implicit none

    ! Arguments
    class(cce_loss_type), intent(in), target :: this
    !! Instance of the physics-informed neural network loss function
    type(array_type), dimension(:,:), intent(inout), target :: predicted
    !! Predicted values
    type(array_type), dimension(size(predicted,1),size(predicted,2)), intent(in) :: &
         expected
    !! Expected values
    type(array_type), pointer :: output
    !! Categorical cross entropy loss

    ! Local variables
    integer :: s, i
    !! Loop indices
    type(array_type), pointer :: ptr
    !! Temporary pointer for calculations

    output => -mean( sum( &
         expected(1,1) * log(predicted(1,1) + this%epsilon), &
         dim=1 ), dim=2)
    if(any(shape(predicted).gt.1))then
       do s = 1, size(predicted,2)
          do i = 1, size(predicted,1)
             if(i.eq.1 .and. s.eq.1) cycle
             if(.not.predicted(i,s)%allocated .or. &
                  .not.expected(i,s)%allocated) cycle
             ptr => mean( sum( &
                  expected(i,s) * log(predicted(i,s) + this%epsilon), &
                  dim=1 ), dim=2)

             output => output - ptr
          end do
       end do
    end if

  end function compute_cce