Decision Tree: save tree as svg image (#212)

Author: gyrdym

CHANGELOG.md

@@ -1,5 +1,9 @@
 # Changelog
 
+## 16.7.0
+- DecisionTreeClassifier:
+    - Added `saveAsSvg` method which returns '.svg' file with a graphical representation of a tree
+
 ## 16.6.3
 - KDTree:
     - `fromIterable` constructor added

benchmark/knn_regressor.dart

@@ -1,4 +1,4 @@
-// 10.0 sec (MacBook Air mid 2017)
+// 6.0 sec (MacBook Air mid 2017)
 import 'package:benchmark_harness/benchmark_harness.dart';
 import 'package:ml_algo/ml_algo.dart';
 import 'package:ml_dataframe/ml_dataframe.dart';

e2e/decision_tree_classifier/decision_tree_classifier_save_as_svg_test.dart

@@ -0,0 +1,48 @@
+import 'package:ml_algo/ml_algo.dart';
+import 'package:ml_dataframe/ml_dataframe.dart';
+import 'package:ml_preprocessing/ml_preprocessing.dart';
+import 'package:test/test.dart';
+
+void main() async {
+  group('DecisionTreeClassifier', () {
+    test('should save graphical representation as svg image, iris dataset',
+        () async {
+      final samples = (await fromCsv('e2e/_datasets/iris.csv'))
+          .shuffle()
+          .dropSeries(seriesNames: ['Id']);
+      final pipeline = Pipeline(samples, [
+        encodeAsIntegerLabels(
+          featureNames: ['Species'],
+        ),
+      ]);
+      final processed = pipeline.process(samples);
+      final classifier = DecisionTreeClassifier(
+        processed,
+        'Species',
+        minError: 0.3,
+        minSamplesCount: 5,
+        maxDepth: 4,
+      );
+
+      await classifier.saveAsSvg('e2e/decision_tree_classifier/iris_tree.svg');
+    });
+
+    test(
+        'should save graphical representation as svg image, pima indians diabetes dataset',
+        () async {
+      final samples =
+          (await fromCsv('e2e/_datasets/pima_indians_diabetes_database.csv'))
+              .shuffle();
+      final classifier = DecisionTreeClassifier(
+        samples,
+        'class variable (0 or 1)',
+        minError: 0.1,
+        minSamplesCount: 2,
+        maxDepth: 4,
+      );
+
+      await classifier
+          .saveAsSvg('e2e/decision_tree_classifier/pima_indians_tree.svg');
+    });
+  });
+}

e2e/decision_tree_classifier/iris_tree.svg

@@ -1,3 +1,5 @@
+import 'dart:io';
+
 import 'package:ml_algo/src/classifier/classifier.dart';
 import 'package:ml_algo/src/classifier/decision_tree_classifier/_init_module.dart';
 import 'package:ml_algo/src/classifier/decision_tree_classifier/decision_tree_classifier_factory.dart';
@@ -122,4 +124,22 @@ abstract class DecisionTreeClassifier
   ///
   /// The value is read-only, it's a hyperparameter of the model
   int get maxDepth;
+
+  /// Saves tree as SVG-image. Example:
+  ///
+  /// ```dart
+  /// final samples = (await fromCsv('path/to/dataset.csv'));
+  /// final classifier = DecisionTreeClassifier(
+  ///   samples,
+  ///   'target',
+  ///   minError: 0.3,
+  ///   minSamplesCount: 5,
+  ///   maxDepth: 4,
+  /// );
+  //
+  //  await classifier.saveAsSvg('tree.svg');
+  /// ```
+  ///
+  /// The file 'tree.svg' now contains a graphical representation of the tree
+  Future<File> saveAsSvg(String filePath);
 }

e2e/decision_tree_classifier/pima_indians_tree.svg

@@ -1,10 +1,13 @@
+import 'dart:io';
+
 import 'package:json_annotation/json_annotation.dart';
 import 'package:ml_algo/src/classifier/_mixins/assessable_classifier_mixin.dart';
 import 'package:ml_algo/src/classifier/decision_tree_classifier/_injector.dart';
 import 'package:ml_algo/src/classifier/decision_tree_classifier/decision_tree_classifier.dart';
 import 'package:ml_algo/src/classifier/decision_tree_classifier/decision_tree_classifier_constants.dart';
 import 'package:ml_algo/src/classifier/decision_tree_classifier/decision_tree_classifier_factory.dart';
 import 'package:ml_algo/src/classifier/decision_tree_classifier/decision_tree_json_keys.dart';
+import 'package:ml_algo/src/classifier/decision_tree_classifier/helpers/create_tree_svg_markup/create_tree_svg_markup.dart';
 import 'package:ml_algo/src/common/constants/common_json_keys.dart';
 import 'package:ml_algo/src/common/json_converter/dtype_json_converter.dart';
 import 'package:ml_algo/src/common/serializable/serializable_mixin.dart';
@@ -143,4 +146,12 @@ class DecisionTreeClassifierImpl
           maxDepth,
         );
   }
+
+  @override
+  Future<File> saveAsSvg(String filePath) async {
+    final markup = createTreeSvgMarkup(treeRootNode);
+    final file = await File(filePath).create(recursive: true);
+
+    return file.writeAsString(markup);
+  }
 }

lib/src/classifier/decision_tree_classifier/decision_tree_classifier.dart

@@ -0,0 +1,52 @@
+import 'package:ml_algo/src/classifier/decision_tree_classifier/helpers/create_tree_svg_markup/create_tree_svg_markup_constants.dart';
+import 'package:ml_algo/src/classifier/decision_tree_classifier/helpers/create_tree_svg_markup/get_tree_levels.dart';
+import 'package:ml_algo/src/classifier/decision_tree_classifier/helpers/create_tree_svg_markup/get_tree_node_distance_by_level.dart';
+import 'package:ml_algo/src/classifier/decision_tree_classifier/helpers/create_tree_svg_markup/get_tree_node_markup.dart';
+import 'package:ml_algo/src/classifier/decision_tree_classifier/helpers/create_tree_svg_markup/get_tree_width.dart';
+import 'package:ml_algo/src/tree_trainer/tree_node/tree_node.dart';
+
+class _NodesMarkupData {
+  _NodesMarkupData(
+      {required this.markup, required this.level, required this.y});
+
+  final String markup;
+  final int level;
+  final num y;
+}
+
+String createTreeSvgMarkup(TreeNode node) {
+  final shape = node.shape;
+  final levels = getTreeLevels(node, shape.length);
+  final nodeDistanceByLevel =
+      getTreeNodeDistanceByLevel(levels, nodeWidth, minNodeHorizontalDistance);
+  final totalWidth = getTreeWidth(levels, nodeWidth, minNodeHorizontalDistance);
+  final totalHeight = shape.length * (nodeHeight + nodeVerticalDistance);
+  final markup = _generateMarkup(levels, node, nodeDistanceByLevel);
+
+  return '<svg xmlns="http://www.w3.org/2000/svg" width="$totalWidth" height="$totalHeight">$textStyles$markup</svg>';
+}
+
+String _generateMarkup(
+    List<List<TreeNode>> levels, TreeNode root, Map<int, num> distByLevel) {
+  return levels.fold<_NodesMarkupData>(
+      _NodesMarkupData(markup: '', level: 0, y: 20), (data, nodes) {
+    final spacing = distByLevel[data.level]!;
+    final childSpacing = distByLevel.containsKey(data.level + 1)
+        ? distByLevel[data.level + 1]!
+        : null;
+    final getX =
+        (int idx) => spacing / 2 + (idx == 0 ? 0 : idx * (nodeWidth + spacing));
+
+    var nodeIdx = 0;
+    final nodesMarkup = nodes
+        .map((node) =>
+            getTreeNodeMarkup(node, getX(nodeIdx++), data.y, childSpacing))
+        .join();
+
+    return _NodesMarkupData(
+      markup: '${data.markup}$nodesMarkup',
+      level: data.level + 1,
+      y: data.y + nodeHeight + nodeVerticalDistance,
+    );
+  }).markup;
+}

lib/src/classifier/decision_tree_classifier/decision_tree_classifier_impl.dart

@@ -0,0 +1,19 @@
+const nodeWidth = 200;
+const nodeHeight = 80;
+const minNodeHorizontalDistance = 20;
+const nodeVerticalDistance = 150;
+
+const textStyles = '<style>'
+    '.label { font: 14px sans-serif; }'
+    '.value { font: bold 14px sans-serif; }'
+    '.root-node-label { font: 24px sans-serif; }'
+    '</style>';
+
+const nodeStyle =
+    'fill:blue;fill-opacity:.25;stroke:cornflowerblue;stroke-width:2';
+const nodeLineStyle = 'stroke:cornflowerblue;stroke-width:1';
+
+const labelMargin = 20;
+const labelWidth = 100;
+const labelHeight = 20;
+const noValue = '-';