#lang scheme
;;; PLT Scheme Science Collection
;;; chebyshev.ss
;;; Copyright (c) 2004-2009 M. Douglas Williams
;;;
;;; This file is part of the PLT Scheme Science Collection.
;;;
;;; The PLT Scheme Science Collection is free software: you can redistribute it
;;; and/or modify it under the terms of the GNU Lesser General Public License as
;;; published by the Free Software Foundation, either version 3 of the License,
;;; or (at your option) any later version.
;;;
;;; The PLT Scheme Science Collection 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 Lesser
;;; General Public License for more details.
;;;
;;; You should have received a copy of the GNU Lesser General Public License
;;; along with the PLT Scheme Science Collection.  If not, see
;;; <http://www.gnu.org/licenses/>.
;;;
;;; -----------------------------------------------------------------------------
;;;
;;; This code in based on the Chebyshev Approximations in the GNU Scientific
;;; Library (GSL), which is licensed under the GPL.
;;;
;;; Version  Date      Description
;;; 0.1.0    08/07/04  This is the initial release of the Chebyshev approx-
;;;                    imations routines ported from GSL. (Doug Williams)
;;; 1.0.0    09/29/04  Added contracts for functions.  Marked as ready fo
;;;                    Release 1.0.  (Doug Williams)
;;; 2.0.0    11/17/07  Added unchecked versions of functions and getting ready
;;;                    for PLT Scheme V4.0.  (Doug Williams)
;;; 3.0.0    06/09/08  Global changes for V4.0.  (Doug Williams)
;;; 4.0.0    09/13/09  Rewritten for better clarity and efficiency; changes
;;;                    make-chebyshev-series to make it more general; added
;;;                    derivative and integral series construction. (Doug
;;;                    Williams)

;;; (struct chebyshev-series (coefficients order lower upper))
;;;   coefficients : (vectorof real?)
;;;   order : exact-positive-integer?
;;;   lower : real?
;;;   upper : real?
(define-values (struct:chebyshev-series
                construct-chebyshev-series
                chebyshev-series?
                chebyshev-series-ref
                chebyshev-series-set!)
  (make-struct-type 'chebyshev-series #f 4 0))

(define chebyshev-series-coefficients
  (make-struct-field-accessor chebyshev-series-ref 0 'coefficients))
(define set-chebyshev-series-coefficients!
  (make-struct-field-mutator chebyshev-series-set! 0 'coefficients))

(define chebyshev-series-order
  (make-struct-field-accessor chebyshev-series-ref 1 'order))

(define chebyshev-series-lower
  (make-struct-field-accessor chebyshev-series-ref 2 'lower))
(define set-chebyshev-series-lower!
  (make-struct-field-mutator chebyshev-series-set! 2 'lower))
  
(define chebyshev-series-upper
  (make-struct-field-accessor chebyshev-series-ref 3 'upper))
(define set-chebyshev-series-upper!
  (make-struct-field-mutator chebyshev-series-set! 3 'upper))

;;; (make-chebyshev-series order) -> chebyshev-series?
;;;   order : exact-positive-integer?
;;; (make-chebyshev-series coeffs-or-func order lower upper) -> chebyshev-series?
;;;   coeffs-of-func : (or/c (vectorof real?) (-> real? real?))
;;;   order : exact-positive-integer?
;;;   lower : real?
;;;   upper : real?
(define make-chebyshev-series
  (case-lambda
    ((order)
     (construct-chebyshev-series (make-vector (add1 order) 0.0) order 0.0 0.0))
    ((coeffs-or-func order lower upper)
     (unless (< lower upper)
       (error 'make-chebyshev-series
              "null interval [~a, ~a]" lower upper))
     (cond ((vector? coeffs-or-func)
            (unless (= (vector-length coeffs-or-func) (add1 order))
              (error 'make-chebyshev-series
                     "expected vector of length ~a, given vector of length ~a"
                     (add1 order) (vector-length coeffs-or-func)))
            (construct-chebyshev-series coeffs-or-func order lower upper))
           ((procedure? coeffs-or-func)
            (construct-chebyshev-series
             (compute-coefficients coeffs-or-func order lower upper)
             order lower upper))
           (else
            (error 'make-chebyshev-series
                   "expected vector or procedure, given ~a" coeffs-or-func))))))

;;; (compute-coefficients func order a b) -> (vectorof real?)
;;;   func : (-> real? real?)
;;;   order : exact-positive-integer?
;;;   lower : real?
;;;   upper : real?
(define (compute-coefficients func order a b)
  (let* ((bma (* 0.5 (- b a)))
         (bpa (* 0.5 (+ b a)))
         (fac (/ 2.0 (add1 order)))
         (f (build-vector
             (add1 order)
             (lambda (i)
               (let ((y (cos (* pi (/ (+ i 0.5) (add1 order))))))
                 (func (+ (* y bma) bpa)))))))
    (build-vector
     (add1 order)
     (lambda (j)
       (* fac (for/fold ((sum 0.0))
                        ((k (in-naturals))
                         (x (in-vector f)))
                (+ sum (* x (cos (/ (* pi j (+ k 0.5)) (add1 order)))))))))))

;;; (make-chebyshev-series-order order) -> chebyshev-series?
;;;   order : exact-positive-integer?
;;; Legacy function to create a new Chebyshev series with the specified order.
;;; Use (make-chebyshev-series order) instead.
(define (make-chebyshev-series-order order)
  (make-chebyshev-series order))

;;; (chebyshev-series-init series func a b) -> void?
;;;   series : chebyshev-series?
;;;   order : exact-positive-integer?
;;;   a : real?
;;;   b : real?
;;; Legacy function to initialize a Chebyshev series to estimate func on the
;;; interval [a, b]. Use (make-chebyshev-series func order a b) instead.
(define (chebyshev-series-init series func a b)
     (when (>= a b)
       (error 'make-chebyshev-series
              "null interval [~a, ~a]" a b))
  (let ((order (chebyshev-series-order series)))
    (set-chebyshev-series-coefficients!
     series (compute-coefficients func order a b))
    (set-chebyshev-series-lower! series a)
    (set-chebyshev-series-upper! series b)))

;;; (chebyshev-eval series x) -> real?
;;;   series : chebyshev-series?
;;;   x : real?
(define (chebyshev-eval series x)
  (let* ((coefficients (chebyshev-series-coefficients series))
         (order (chebyshev-series-order series))
         (lower (chebyshev-series-lower series))
         (upper (chebyshev-series-upper series))
         (c0 (vector-ref coefficients 0))
         (y (/ (- (* 2.0 x) lower upper) (- upper lower)))
         (y2 (* 2.0 y)))
    (let-values (((d dd)
                  (for/fold ((d 0.0)
                             (dd 0.0))
                            ((c (in-vector coefficients order 0 -1)))
                    (values (+ (* y2 d) (- dd) c) d))))
      (+ (* y d) (- dd) (* 0.5 c0)))))

;;; (chebyshev-eval-n series n x) -> real?
;;;   series : chebyshev-series?
;;;   n : exact-positive-integer?
;;;   x : real?
(define (chebyshev-eval-n series n x)
  (let* ((coefficients (chebyshev-series-coefficients series))
         (order (chebyshev-series-order series))
         (lower (chebyshev-series-lower series))
         (upper (chebyshev-series-upper series))
         (c0 (vector-ref coefficients 0))
         (y (/ (- (* 2.0 x) lower upper) (- upper lower)))
         (y2 (* 2.0 y)))
    (let-values (((d dd)
                  (for/fold ((d 0.0)
                             (dd 0.0))
                            ((c (in-vector coefficients (min n order) 0 -1)))
                    (values (+ (* y2 d) (- dd) c) d))))
      (+ (* y d) (- dd) (* 0.5 c0)))))

;;; (make-chebyshev-series-derivative series) -> chebyshev-series?
;;;   series : chebyshev-series?
(define (make-chebyshev-series-derivative series)
  (let* ((series-coefficients (chebyshev-series-coefficients series))
         (order (chebyshev-series-order series))
         (lower (chebyshev-series-lower series))
         (upper (chebyshev-series-upper series))
         (n (add1 order))
         (deriv-coefficients (make-vector n))
         (con (/ 2.0 (- upper lower))))
    (vector-set! deriv-coefficients (sub1 n) 0.0)
    (when (> n 1)
      (vector-set! deriv-coefficients (- n 2)
                   (* 2.0 (- n 1) (vector-ref series-coefficients (- n 1))))
      (for ((i (in-range (- n 3) 0 -1)))
        (vector-set! deriv-coefficients i
                     (+ (vector-ref deriv-coefficients (+ i 2))
                        (* 2.0 (+ i 1) (vector-ref series-coefficients (+ i 1))))))
      (vector-set! deriv-coefficients 0
                   (+ (vector-ref deriv-coefficients 2)
                      (* 2.0 (vector-ref series-coefficients 1))))
      (for ((i (in-range n)))
        (vector-set! deriv-coefficients i
                     (* (vector-ref deriv-coefficients i) con))))
    (make-chebyshev-series deriv-coefficients order lower upper)))

;;; (make-chebyshev-series-integral series) -> chebyshev-series?
;;;   series : chebyshev-series?
(define (make-chebyshev-series-integral series)
  (let* ((series-coefficients (chebyshev-series-coefficients series))
         (order (chebyshev-series-order series))
         (lower (chebyshev-series-lower series))
         (upper (chebyshev-series-upper series))
         (n (add1 order))
         (integ-coefficients (make-vector n))
         (con (* 0.25 (- upper lower))))
    (cond ((= n 1)
           (vector-set! integ-coefficients 0 0.0))
          ((= n 2)
           (vector-set! integ-coefficients 1
                        (* con (vector-ref series-coefficients 0)))
           (vector-set! integ-coefficients 0
                        (* 2.0 (vector-ref integ-coefficients 1))))
          (else
           (let-values (((sum fac)
                         (for/fold ((sum 0.0)
                                    (fac 1.0))
                                   ((i (in-range 1 (- n 2))))
                           (vector-set! integ-coefficients i
                                        (* con (- (vector-ref series-coefficients (- i 1))
                                                  (/ (vector-ref series-coefficients (+ 1))
                                                     i))))
                           (values (+ sum (* fac (vector-ref integ-coefficients i)))
                                   (- fac)))))
             (vector-set! integ-coefficients (- n 1)
                          (* con (/ (vector-ref series-coefficients (- n 2))
                                    (- n 1))))
             (set! sum (+ sum (* fac (vector-ref integ-coefficients (- n 1)))))
             (vector-set! integ-coefficients 0
                          (* 2.0 sum)))))
    (make-chebyshev-series integ-coefficients order lower upper)))

;;; Module Contracts

(provide
 (rename-out (chebyshev-eval unchecked-chebyshev-eval)
             (chebyshev-eval-n unchecked-chebyshev-eval-n)))

(provide/contract
 (chebyshev-series?
  (-> any/c boolean?))
 (chebyshev-series-coefficients
  (-> chebyshev-series? (vectorof real?)))
 (chebyshev-series-order
  (-> chebyshev-series? exact-positive-integer?))
 (chebyshev-series-lower
  (-> chebyshev-series? real?))
 (chebyshev-series-upper
  (-> chebyshev-series? real?))
 (make-chebyshev-series
  (case-> (-> exact-positive-integer? chebyshev-series?)
          (-> (or/c (vectorof real?) (-> real? real?)) exact-positive-integer?
              real? real? chebyshev-series?)))
 (make-chebyshev-series-order
  (-> exact-positive-integer? chebyshev-series?))
 (chebyshev-series-init
  (-> chebyshev-series? (-> real? real?) real? real? void?))
 (chebyshev-eval
  (-> chebyshev-series? real? real?))
 (chebyshev-eval-n
  (-> chebyshev-series? exact-positive-integer? real? real?))
 (make-chebyshev-series-derivative
  (-> chebyshev-series? chebyshev-series?))
 (make-chebyshev-series-integral
  (-> chebyshev-series? chebyshev-series?)))