guile-ml/tree.scm

(define-module (tree)
  ;; The exported procedures are created from the macro which defines record
  ;; types, where one cannot use define-public.
  #:export (Split
            make-split
            split?
            split-feature-index
            split-value
            set-split-value
            split-subsets
            set-split-subsets
            split-cost
            set-split-cost

            Node
            make-node
            node?
            node-data
            set-node-data
            node-split-feature-index
            set-node-split-feature-index
            node-split-value
            set-node-split-value
            node-left
            set-node-left
            node-right
            set-node-right))

(use-modules
 ;; SRFI-9 for standardized records
 (srfi srfi-9)
 ;; for functional structs (not part of srfi-9 directly)
 (srfi srfi-9 gnu)
 (utils display-utils)
 (utils string-utils)
 (dataset))

;; ===============
;; DATA STRUCTURES
;; ===============

;; A Split record is simply a structure to contain all information about a
;; split. It is not a node containing data.
(define-immutable-record-type Split
  ;; define constructor
  (make-split index value subsets cost)
  ;; define predicate
  split?
  ;; define accessors and functional setters

  ;; Note: It should never be required to update the index of the feature, which
  ;; was used to create the split. If that feature changed, the whole split
  ;; record would become invalid and a new instance should be created.
  (index split-feature-index)
  (value split-value set-split-value)
  (subsets split-subsets set-split-subsets)
  (cost split-cost set-split-cost))


;; A Node record is a node of the decision tree. It contains data and
;; information about the feature, at which it splits and the value for that
;; split. It also contains references to the child nodes.
(define-immutable-record-type Node
  ;; define constructor
  (make-node data split-feature-index split-value left right)
  ;; define predicate
  node?
  ;; define accessors and functional setters
  (data node-data set-node-data)

  (split-feature-index
   node-split-feature-index
   set-node-split-feature-index)

  (split-value
   node-split-value
   set-node-split-value)

  (left node-left set-node-left)
  (right node-right set-node-right))


(define-public make-leaf-node
  (lambda (data)
    (make-node data
               'none
               'none
               empty-dataset
               empty-dataset)))


(define-public make-leaf-node-from-split-node
  (lambda (split-node)
    (make-leaf-node (node-data split-node))))


(define-public leaf-node?
  (lambda (node)
    (and (dataset-empty? (node-left node))
         (dataset-empty? (node-right node)))))


(define-public last-split-node?
  (lambda (node)
    (cond [(leaf-node? node) #f]
          [else
           (and (leaf-node? (node-left node))
                (leaf-node? (node-right node)))])))


;; displays a string representation of a learned decision tree
(define-public print-tree
  (lambda (tree label-column-index)
    (define tree->string
      (lambda (tree depth)
        (cond
         [(leaf-node? tree)
          (string-append (n-times-string "  " depth)
                         "["
                         (number->string
                          (dataset-majority-prediction (node-data tree)
                                                       label-column-index))
                         "]\n")]
         [else
          (string-append
           (string-append (n-times-string "  " depth)
                          "[feature:"
                          (number->string (node-split-feature-index tree))
                          " < "
                          (number->string (node-split-value tree))
                          "]\n")
           (tree->string (node-left tree) (+ depth 1))
           (tree->string (node-right tree) (+ depth 1)))])))
    (displayln (tree->string tree 0))))