init_gno Subroutine

  subroutine init_gno(this, input_shape, verbose)
    !! Initialise the Graph Neural Operator layer
    !!
    !! input_shape(1) = num_inputs (F_in)
    !! input_shape(2) = num_vertices (set to 0 if variable)
    implicit none

    ! Arguments
    class(graph_nop_layer_type), intent(inout) :: this
    !! Layer instance to initialise
    integer, dimension(:), intent(in) :: input_shape
    !! Input feature/vertex shape
    integer, optional, intent(in) :: verbose
    !! Verbosity level

    ! Local variables
    integer :: num_inputs, H, F_out, F_in, d, F
    !! Effective input count and kernel dimensions
    integer :: kernel_size, off_U, off_bu, off_V, off_bv
    !! Packed-kernel size and section offsets
    integer :: verbose_ = 0
    !! Effective verbosity level

    if(present(verbose)) verbose_ = verbose

    !---------------------------------------------------------------------------
    ! Set shapes
    !---------------------------------------------------------------------------
    if(.not.allocated(this%input_shape)) call this%set_shape(input_shape)

    F_in  = input_shape(1)
    F_out = this%num_outputs
    d     = this%coord_dim
    H     = this%kernel_hidden
    F     = F_out * F_in

    !---------------------------------------------------------------------------
    ! Set msgpass fields
    !---------------------------------------------------------------------------
    if(allocated(this%num_vertex_features)) deallocate(this%num_vertex_features)
    allocate(this%num_vertex_features(0:1))
    this%num_vertex_features(0) = F_in
    this%num_vertex_features(1) = F_out

    if(allocated(this%num_edge_features)) deallocate(this%num_edge_features)
    allocate(this%num_edge_features(0:1), source=0)

    kernel_size = H * d + H + F * H + F

    if(allocated(this%num_params_msg)) deallocate(this%num_params_msg)
    allocate(this%num_params_msg(1))
    this%num_params_msg(1) = kernel_size + F_out * F_in
    if(this%use_bias) this%num_params_msg(1) = this%num_params_msg(1) + F_out
    this%num_params_readout = 0

    this%output_shape = [this%num_outputs, 0]
    this%num_params = this%get_num_params()

    !---------------------------------------------------------------------------
    ! Allocate learnable parameters
    !
    ! params(1): packed kernel MLP  [kernel_size, 1]
    !            Layout: U [H*d] | b_u [H] | V [F*H] | b_v [F]
    ! params(2): W   [F_out*F_in, 1]  - linear bypass weights
    ! params(3): b   [F_out, 1]       - output bias (optional)
    !---------------------------------------------------------------------------
    if(allocated(this%weight_shape)) deallocate(this%weight_shape)
    if(allocated(this%params)) deallocate(this%params)
    if(this%use_bias)then
       if(allocated(this%bias_shape)) deallocate(this%bias_shape)
       this%bias_shape = [ F_out ]
       allocate(this%weight_shape(2, 3))
       this%weight_shape(:,3) = [ F_out, 1 ]
       allocate(this%params(3))
    else
       allocate(this%weight_shape(2, 2))
       allocate(this%params(2))
    end if
    this%weight_shape(:,1) = [ kernel_size, 1 ]
    this%weight_shape(:,2) = [ F_out, F_in ]

    ! params(1): packed kernel MLP params
    call this%params(1)%allocate([kernel_size, 1])
    call this%params(1)%set_requires_grad(.true.)
    this%params(1)%fix_pointer = .true.
    this%params(1)%is_sample_dependent = .false.
    this%params(1)%is_temporary = .false.

    ! params(2): W [F_out x F_in]
    call this%params(2)%allocate([F_out, F_in, 1])
    call this%params(2)%set_requires_grad(.true.)
    this%params(2)%fix_pointer = .true.
    this%params(2)%is_sample_dependent = .false.
    this%params(2)%is_temporary = .false.

    if(this%use_bias)then
       ! params(3): b [F_out]
       call this%params(3)%allocate([F_out, 1])
       call this%params(3)%set_requires_grad(.true.)
       this%params(3)%fix_pointer = .true.
       this%params(3)%is_sample_dependent = .false.
       this%params(3)%is_temporary = .false.
    end if


    !---------------------------------------------------------------------------
    ! Initialise learnable parameters
    !---------------------------------------------------------------------------
    off_U  = 0
    off_bu = H * d
    off_V  = off_bu + H
    off_bv = off_V + F * H

    ! U [H x d] — kernel first-layer weights
    call this%kernel_init%initialise( &
         this%params(1)%val(off_U+1:off_bu, 1), &
         fan_in = d, fan_out = H, &
         spacing = [ H ] &
    )
    ! b_u [H] — kernel first-layer bias
    call this%bias_init%initialise( &
         this%params(1)%val(off_bu+1:off_V, 1), &
         fan_in = d, fan_out = H &
    )
    ! V [F x H] — kernel second-layer weights
    call this%kernel_init%initialise( &
         this%params(1)%val(off_V+1:off_bv, 1), &
         fan_in = H, fan_out = F, &
         spacing = [ F ] &
    )
    ! b_v [F] — kernel second-layer bias
    call this%bias_init%initialise( &
         this%params(1)%val(off_bv+1:, 1), &
         fan_in = H, fan_out = F &
    )
    ! W [F_out x F_in] — linear bypass
    num_inputs = F_in
    if(this%use_bias) num_inputs = F_in + 1
    call this%kernel_init%initialise( &
         this%params(2)%val(:,1), &
         fan_in = num_inputs, fan_out = F_out, &
         spacing = [ F_out ] &
    )
    if(this%use_bias)then
       call this%bias_init%initialise( &
            this%params(3)%val(:,1), &
            fan_in = num_inputs, fan_out = F_out &
       )
    end if


    !---------------------------------------------------------------------------
    ! Allocate output arrays
    !---------------------------------------------------------------------------
    if(allocated(this%output)) deallocate(this%output)

  end subroutine init_gno