## Copyright (C) 2006 Peter V. Lanspeary <pvl@mecheng.adelaide.edu.au>
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program; see the file COPYING. If not, see
## <https://www.gnu.org/licenses/>.

## -*- texinfo -*-
## @deftypefn  {Function File} {} ar_psd (@var{a}, @var{v})
## @deftypefnx {Function File} {} ar_psd (@var{a}, @var{v}, @var{freq})
## @deftypefnx {Function File} {} ar_psd (@var{a}, @var{v}, @var{freq}, @var{Fs})
## @deftypefnx {Function File} {} ar_psd (@dots{}, @var{range})
## @deftypefnx {Function File} {} ar_psd (@dots{}, @var{method})
## @deftypefnx {Function File} {} ar_psd (@dots{}, @var{plottype})
## @deftypefnx {Function File} {[@var{psd}, @var{f_out}] =} ar_psd (@dots{})
##
## Calculate the power spectrum of the autoregressive model
##
## @example
## @group
##                        M
## x(n) = sqrt(v).e(n) + SUM a(k).x(n-k)
##                       k=1
## @end group
## @end example
##
## where @math{x(n)} is the output of the model and @math{e(n)} is white noise.
## This function is intended for use with
## @code{[a, v, k] = arburg (x, poles, criterion)}
## which use the Burg (1968) method to calculate a "maximum entropy"
## autoregressive model of @var{x}.
##
## If the @var{freq} argument is a vector (of frequencies) the spectrum is
## calculated using the polynomial method and the @var{method} argument is
## ignored.  For scalar @var{freq}, an integer power of 2, or @var{method} =
## "FFT", causes the spectrum to be calculated by FFT.  Otherwise, the spectrum
## is calculated as a polynomial.  It may be computationally more
## efficient to use the FFT method if length of the model is not much
## smaller than the number of frequency values.  The spectrum is scaled so
## that spectral energy (area under spectrum) is the same as the time-domain
## energy (mean square of the signal).
##
## ARGUMENTS:
## All but the first two arguments are optional and may be empty.
## @itemize
## @item
## @var{a}
## list of M=(order+1) autoregressive model
## coefficients.  The first element of "ar_coeffs" is the
## zero-lag coefficient, which always has a value of 1.
## @item
## @var{v}
## square of the moving-average coefficient of the AR model.
## @item
## @var{freq}
## frequencies at which power spectral density is calculated, or a scalar
## indicating the number of uniformly distributed frequency
## values at which spectral density is calculated.
## (default = 256)
## @item
## @var{Fs}
## sampling frequency (Hertz) (default=1)
## @end itemize
##
## CONTROL-STRING ARGUMENTS -- each of these arguments is a character string.
## Control-string arguments can be in any order after the other arguments.
##
## Range:
##
## 'half',  'onesided' : frequency range of the spectrum is
## from zero up to but not including sample_f/2.  Power
## from negative frequencies is added to the positive
## side of the spectrum.
## 'whole', 'twosided' : frequency range of the spectrum is
## -sample_f/2 to sample_f/2, with negative frequencies
## stored in "wrap around" order after the positive
## frequencies; e.g. frequencies for a 10-point 'twosided'
## spectrum are 0 0.1 0.2 0.3 0.4 0.5 -0.4 -0.3 -0.2 -0.1
## 'shift', 'centerdc' : same as 'whole' but with the first half
## of the spectrum swapped with second half to put the
## zero-frequency value in the middle. (See "help
## fftshift". If "freq" is vector, 'shift' is ignored.
## If model coefficients "ar_coeffs" are real, the default
## range is 'half', otherwise default range is 'whole'.
##
## Method:
##
## 'fft':  use FFT to calculate power spectrum.
## 'poly': calculate power spectrum as a polynomial of 1/z
## N.B. this argument is ignored if the "freq" argument is a
## vector.  The default is 'poly' unless the "freq"
## argument is an integer power of 2.
##
## Plot type:
##
## 'plot', 'semilogx', 'semilogy', 'loglog', 'squared' or 'db':
## specifies the type of plot.  The default is 'plot', which
## means linear-linear axes. 'squared' is the same as 'plot'.
## 'dB' plots "10*log10(psd)".  This argument is ignored and a
## spectrum is not plotted if the caller requires a returned
## value.
##
## RETURNED VALUES:
## If returned values are not required by the caller, the spectrum
## is plotted and nothing is returned.
## @itemize
## @item
## @var{psd}
## estimate of power-spectral density
## @item
## @var{f_out}
## frequency values
## @end itemize
##
## REFERENCE
## [1] Equation 2.28 from Steven M. Kay and Stanley Lawrence Marple Jr.:
##   "Spectrum analysis -- a modern perspective",
##   Proceedings of the IEEE, Vol 69, pp 1380-1419, Nov., 1981
##
## @end deftypefn

function varargout = ar_psd(a,v,varargin)

  ##
  ## Check fixed arguments
  if ( nargin < 2 )
    error( 'ar_psd: needs at least 2 args. Use help ar_psd.' );
  elseif ( ~isvector(a) || length(a)<2 )
    error( 'ar_psd: arg 1 (a) must be vector, length>=2.' );
  elseif ( ~isscalar(v) )
    error( 'ar_psd: arg 2 (v) must be real scalar >0.' );
  else
    real_model = isreal(a);
  ##
  ##  default values for optional arguments
    freq = 256;
    user_freqs = 0;    ## boolean: true for user-specified frequencies
    Fs   = 1.0;
    ##  FFT padding factor (is also frequency range divisor): 1=whole, 2=half.
    pad_fact = 1 + real_model;
    do_shift   = 0;
    force_FFT  = 0;
    force_poly = 0;
    plot_type  = 1;
  ##
  ##  decode and check optional arguments
  ##  end_numeric_args is boolean; becomes true at 1st string arg
    end_numeric_args = 0;
    for iarg = 1:length(varargin)
      arg = varargin{iarg};
      end_numeric_args = end_numeric_args || ischar(arg);
      ## skip empty arguments
      if ( isempty(arg) )
        1;
      ## numeric optional arguments must be first, cannot follow string args
      ## N.B. older versions of matlab may not have the function "error" so
      ## the user writes "function error(msg); disp(msg); end" and we need
      ## a "return" here.
      elseif ( ~ischar(arg) )
        if ( end_numeric_args )
          error( 'ar_psd: control arg must be string.' );
        ##
        ## first optional numeric arg is "freq"
        elseif ( iarg == 1 )
          user_freqs = isvector(arg) && length(arg)>1;
          if ( ~isscalar(arg) && ~user_freqs )
            error( 'ar_psd: arg 3 (freq) must be vector or scalar.' );
          elseif ( ~user_freqs && ( ~isreal(arg) || ...
                   fix(arg)~=arg || arg <= 2 || arg >= 1048576 ) )
            error('ar_psd: arg 3 (freq) must be integer >=2, <=1048576' );
          elseif ( user_freqs && ~isreal(arg) )
            error( 'ar_psd: arg 3 (freq) vector must be real.' );
          endif
          freq = arg(:); # -> column vector
        ##
        ## second optional numeric arg is  "Fs" - sampling frequency
        elseif ( iarg == 2 )
          if ( ~isscalar(arg) || ~isreal(arg) || arg<=0 )
            error( 'ar_psd: arg 4 (Fs) must be real positive scalar.' );
          endif
          Fs = arg;
        ##
        else
          error( 'ar_psd: control arg must be string.' );
        endif
    ##
    ## decode control-string arguments
      elseif ( strcmp(arg,'plot') || strcmp(arg,'squared') )
        plot_type = 1;
      elseif ( strcmp(arg,'semilogx') )
        plot_type = 2;
      elseif ( strcmp(arg,'semilogy') )
        plot_type = 3;
      elseif ( strcmp(arg,'loglog') )
        plot_type = 4;
      elseif ( strcmp(arg,'dB') )
        plot_type = 5;
      elseif ( strcmp(arg,'fft') )
        force_FFT  = 1;
        force_poly = 0;
      elseif ( strcmp(arg,'poly') )
        force_FFT  = 0;
        force_poly = 1;
      elseif ( strcmp(arg,'half') || strcmp(arg,'onesided') )
        pad_fact = 2;    # FFT zero-padding factor (pad FFT to double length)
        do_shift = 0;
      elseif ( strcmp(arg,'whole') || strcmp(arg,'twosided') )
        pad_fact = 1;    # FFT zero-padding factor (do not pad)
        do_shift = 0;
      elseif ( strcmp(arg,'shift') || strcmp(arg,'centerdc') )
        pad_fact = 1;
        do_shift = 1;
      else
        error( 'ar_psd: string arg: illegal value: %s', arg );
      endif
    endfor
  ##  end of decoding and checking args
  ##
    if ( user_freqs )
      ## user provides (column) vector of frequencies
      if ( any(abs(freq)>Fs/2) )
        error( 'ar_psd: arg 3 (freq) cannot exceed half sampling frequency.' );
      elseif ( pad_fact==2 && any(freq<0) )
        error( 'ar_psd: arg 3 (freq) must be positive in onesided spectrum' );
      endif
      freq_len = length(freq);
      fft_len  = freq_len;
      use_FFT  = 0;
      do_shift = 0;
    else
      ## internally generated frequencies
      freq_len = freq;
      freq = (Fs/pad_fact/freq_len) * [0:freq_len-1]';
      ## decide which method to use (poly or FFT)
      is_power_of_2 = rem(log(freq_len),log(2))<10.*eps;
      use_FFT = ( ~ force_poly && is_power_of_2 ) || force_FFT;
      fft_len = freq_len * pad_fact;
    endif
    ##
    ## calculate denominator of Equation 2.28, Kay and Marple, ref [1]Jr.:
    len_coeffs = length(a);
    if ( use_FFT )
      ## FFT method
      fft_out = fft( [ a(:); zeros(fft_len-len_coeffs,1) ] );
    else
      ## polynomial method
      ## complex data on "half" frequency range needs -ve frequency values
      if ( pad_fact==2 && ~real_model )
        freq = [freq; -freq(freq_len:-1:1)];
        fft_len = 2*freq_len;
      endif
      fft_out = polyval( a(len_coeffs:-1:1), exp( (-i*2*pi/Fs) * freq ) );
    endif
    ##
    ## The power spectrum (PSD) is the scaled squared reciprocal of amplitude
    ## of the FFT/polynomial. This is NOT the reciprocal of the periodogram.
    ## The PSD is a continuous function of frequency.  For uniformly
    ## distributed frequency values, the FFT algorithm might be the most
    ## efficient way of calculating it.
    ##
    psd = ( v / Fs ) ./ ( fft_out .* conj(fft_out) );
    ##
    ## range='half' or 'onesided',
    ##   add PSD at -ve frequencies to PSD at +ve frequencies
    ## N.B. unlike periodogram, PSD at zero frequency _is_ doubled.
    if ( pad_fact==2 )
      freq = freq(1:freq_len);
      if ( real_model )
        ## real data, double the psd
        psd = 2 * psd(1:freq_len);
      elseif ( use_FFT )
        ## complex data, FFT method, internally-generated frequencies
        psd = psd(1:freq_len)+[psd(1); psd(fft_len:-1:freq_len+2)];
      else
        ## complex data, polynomial method
        ##  user-defined and internally-generated frequencies
        psd = psd(1:freq_len)+psd(fft_len:-1:freq_len+1);
      endif
    ##
    ## range='shift'
    ##   disabled for user-supplied frequencies
    ##   Shift zero-frequency to the middle (pad_fact==1)
    elseif ( do_shift )
      len2 = fix((fft_len+1)/2);
      psd  = [psd(len2+1:fft_len); psd(1:len2)];
      freq = [freq(len2+1:fft_len)-Fs; freq(1:len2)];
    endif
    ##
    ## Plot the spectrum if there are no return variables.
    if ( nargout >= 2 )
       varargout{1} = psd;
       varargout{2} = freq;
    elseif ( nargout == 1 )
       varargout{1} = psd;
    else
      if ( plot_type == 1 )
        plot(freq,psd);
      elseif ( plot_type == 2 )
        semilogx(freq,psd);
      elseif ( plot_type == 3 )
        semilogy(freq,psd);
      elseif ( plot_type == 4 )
        loglog(freq,psd);
      elseif ( plot_type == 5 )
        plot(freq,10*log10(psd));
      endif
    endif
  endif

endfunction
