Public Interface of RegressionTree and TreeEnsemble (#2243)

* Immutable RegressionTree and TreeEnsemble wrappers

* Address comments and internalize RegressionTree and TreeEsnmble.
A lot of things become internal and possibly best-friend because of
RegressionTree and TreeEnsemble.

* Create new files for public wrapper classes

* private protected is good

* Doc strings

* Update entry point

* Add a dynamic and static API tests

* Address comments.
1. Rename classes.
2. Use IReadOnlyList instead of ReadOnlySpan.
3. Remove new namespace, Representation.

* Strongly-typed public surface

Seperate TreeEnsemble

* Rename files

Author: wschin

src/Microsoft.ML.Data/Scorers/PredictionTransformer.cs

@@ -171,8 +171,6 @@ public abstract class SingleFeaturePredictionTransformerBase<TModel, TScorer> :
         public SingleFeaturePredictionTransformerBase(IHost host, TModel model, Schema trainSchema, string featureColumn)
             : base(host, model, trainSchema)
         {
-            FeatureColumn = featureColumn;
-
             FeatureColumn = featureColumn;
             if (featureColumn == null)
                 FeatureColumnType = null;

src/Microsoft.ML.FastTree/BoostingFastTree.cs

@@ -62,7 +62,7 @@ protected override void CheckArgs(IChannel ch)
             base.CheckArgs(ch);
         }
 
-        protected override TreeLearner ConstructTreeLearner(IChannel ch)
+        private protected override TreeLearner ConstructTreeLearner(IChannel ch)
         {
             return new LeastSquaresRegressionTreeLearner(
                 TrainSet, Args.NumLeaves, Args.MinDocumentsInLeafs, Args.EntropyCoefficient,
@@ -73,7 +73,7 @@ protected override TreeLearner ConstructTreeLearner(IChannel ch)
                 Args.MinDocsPercentageForCategoricalSplit, Args.Bundling, Args.MinDocsForCategoricalSplit, Args.Bias);
         }
 
-        protected override OptimizationAlgorithm ConstructOptimizationAlgorithm(IChannel ch)
+        private protected override OptimizationAlgorithm ConstructOptimizationAlgorithm(IChannel ch)
         {
             Contracts.CheckValue(ch, nameof(ch));
             OptimizationAlgorithm optimizationAlgorithm;

src/Microsoft.ML.FastTree/Dataset/FeatureFlock.cs

@@ -51,7 +51,7 @@ public PerBinStats(Double sumTargets, Double sumWeights, int count)
     /// per flock. Note that feature indices, whenever present, refer to the feature within the
     /// particular flock the same as they do with <see cref="FeatureFlockBase"/>.
     /// </summary>
-    public abstract class SufficientStatsBase
+    internal abstract class SufficientStatsBase
     {
         // REVIEW: Holdover from histogram. I really don't like this. Figure out if
         // there's a better way.
@@ -929,7 +929,7 @@ public void FillSplitCandidatesCategoricalNeighborBundling(LeastSquaresRegressio
     /// </summary>
     /// <typeparam name="TSuffStats">The type of sufficient stats that we will be able to do
     /// "peer" operations against, like subtract. This will always be the derived class itself.</typeparam>
-    public abstract class SufficientStatsBase<TSuffStats> : SufficientStatsBase
+    internal abstract class SufficientStatsBase<TSuffStats> : SufficientStatsBase
         where TSuffStats : SufficientStatsBase<TSuffStats>
     {
         protected SufficientStatsBase(int features)
@@ -1005,7 +1005,7 @@ protected FeatureFlockBase(int count, bool categorical = false)
         /// <param name="hasWeights">Whether structures related to tracking
         /// example weights should be allocated</param>
         /// <returns>A sufficient statistics object</returns>
-        public abstract SufficientStatsBase CreateSufficientStats(bool hasWeights);
+        internal abstract SufficientStatsBase CreateSufficientStats(bool hasWeights);
 
         /// <summary>
         /// Returns a forward indexer for a single feature. This has a default implementation that
@@ -1207,7 +1207,7 @@ public override long SizeInBytes()
                    + sizeof(int) * HotFeatureStarts.Length;
         }
 
-        public override SufficientStatsBase CreateSufficientStats(bool hasWeights)
+        internal override SufficientStatsBase CreateSufficientStats(bool hasWeights)
         {
             return new SufficientStats(this, hasWeights);
         }

src/Microsoft.ML.FastTree/Dataset/SingletonFeatureFlock.cs

@@ -34,7 +34,7 @@ public override long SizeInBytes()
             return _bins.SizeInBytes() + sizeof(double) * _binUpperBounds.Length;
         }
 
-        public override SufficientStatsBase CreateSufficientStats(bool hasWeights)
+        internal override SufficientStatsBase CreateSufficientStats(bool hasWeights)
         {
             return new SufficientStats(this, hasWeights);
         }

src/Microsoft.ML.FastTree/FastTree.cs

@@ -17,7 +17,7 @@
 using Microsoft.ML.Data;
 using Microsoft.ML.Data.Conversion;
 using Microsoft.ML.EntryPoints;
-using Microsoft.ML.Internal.Calibration;
+using Microsoft.ML.FastTree;
 using Microsoft.ML.Internal.Internallearn;
 using Microsoft.ML.Internal.Utilities;
 using Microsoft.ML.Model;
@@ -56,16 +56,16 @@ public abstract class FastTreeTrainerBase<TArgs, TTransformer, TModel> :
     {
         protected readonly TArgs Args;
         protected readonly bool AllowGC;
-        protected TreeEnsemble TrainedEnsemble;
         protected int FeatureCount;
+        private protected InternalTreeEnsemble TrainedEnsemble;
         private protected RoleMappedData ValidData;
         /// <summary>
         /// If not null, it's a test data set passed in from training context. It will be converted to one element in
         /// <see cref="Tests"/> by calling <see cref="ExamplesToFastTreeBins.GetCompatibleDataset"/> in <see cref="InitializeTests"/>.
         /// </summary>
         private protected RoleMappedData TestData;
-        protected IParallelTraining ParallelTraining;
-        protected OptimizationAlgorithm OptimizationAlgorithm;
+        private protected IParallelTraining ParallelTraining;
+        private protected OptimizationAlgorithm OptimizationAlgorithm;
         protected Dataset TrainSet;
         protected Dataset ValidSet;
         /// <summary>
@@ -89,7 +89,7 @@ public abstract class FastTreeTrainerBase<TArgs, TTransformer, TModel> :
         protected double[] InitValidScores;
         protected double[][] InitTestScores;
         //protected int Iteration;
-        protected TreeEnsemble Ensemble;
+        private protected InternalTreeEnsemble Ensemble;
 
         protected bool HasValidSet => ValidSet != null;
 
@@ -175,8 +175,9 @@ private protected FastTreeTrainerBase(IHostEnvironment env, TArgs args, SchemaSh
 
         protected abstract Test ConstructTestForTrainingData();
 
-        protected abstract OptimizationAlgorithm ConstructOptimizationAlgorithm(IChannel ch);
-        protected abstract TreeLearner ConstructTreeLearner(IChannel ch);
+        private protected abstract OptimizationAlgorithm ConstructOptimizationAlgorithm(IChannel ch);
+
+        private protected abstract TreeLearner ConstructTreeLearner(IChannel ch);
 
         protected abstract ObjectiveFunctionBase ConstructObjFunc(IChannel ch);
 
@@ -488,7 +489,7 @@ protected bool AreSamplesWeighted(IChannel ch)
 
         private void InitializeEnsemble()
         {
-            Ensemble = new TreeEnsemble();
+            Ensemble = new InternalTreeEnsemble();
         }
 
         /// <summary>
@@ -793,7 +794,7 @@ private float GetMachineAvailableBytes()
 
         // This method is called at the end of each training iteration, with the tree that was learnt on that iteration.
         // Note that this tree can be null if no tree was learnt this iteration.
-        protected virtual void CustomizedTrainingIteration(RegressionTree tree)
+        private protected virtual void CustomizedTrainingIteration(InternalRegressionTree tree)
         {
         }
 
@@ -924,7 +925,7 @@ internal abstract class DataConverter
         /// of features we actually trained on. This can be null in the event that no filtering
         /// occurred.
         /// </summary>
-        /// <seealso cref="TreeEnsemble.RemapFeatures"/>
+        /// <seealso cref="InternalTreeEnsemble.RemapFeatures"/>
         public int[] FeatureMap;
 
         protected readonly IHost Host;
@@ -2810,9 +2811,10 @@ public abstract class TreeEnsembleModelParameters :
         ISingleCanSavePfa,
         ISingleCanSaveOnnx
     {
-        //The below two properties are necessary for tree Visualizer
+        // The below two properties are necessary for tree Visualizer
         [BestFriend]
-        internal TreeEnsemble TrainedEnsemble { get; }
+        internal InternalTreeEnsemble TrainedEnsemble { get; }
+
         int ITreeEnsemble.NumTrees => TrainedEnsemble.NumTrees;
 
         // Inner args is used only for documentation purposes when saving comments to INI files.
@@ -2854,7 +2856,9 @@ public abstract class TreeEnsembleModelParameters :
         /// </summary>
         public FeatureContributionCalculator FeatureContributionCalculator => new FeatureContributionCalculator(this);
 
-        public TreeEnsembleModelParameters(IHostEnvironment env, string name, TreeEnsemble trainedEnsemble, int numFeatures, string innerArgs)
+        /// The following function is used in both FastTree and LightGBM so <see cref="BestFriendAttribute"/> is required.
+        [BestFriend]
+        internal TreeEnsembleModelParameters(IHostEnvironment env, string name, InternalTreeEnsemble trainedEnsemble, int numFeatures, string innerArgs)
             : base(env, name)
         {
             Host.CheckValue(trainedEnsemble, nameof(trainedEnsemble));
@@ -2868,7 +2872,7 @@ public TreeEnsembleModelParameters(IHostEnvironment env, string name, TreeEnsemb
             InnerArgs = innerArgs;
             NumFeatures = numFeatures;
 
-            MaxSplitFeatIdx = FindMaxFeatureIndex(trainedEnsemble);
+            MaxSplitFeatIdx = trainedEnsemble.GetMaxFeatureIndex();
             Contracts.Assert(NumFeatures > MaxSplitFeatIdx);
 
             InputType = new VectorType(NumberType.Float, NumFeatures);
@@ -2892,8 +2896,8 @@ protected TreeEnsembleModelParameters(IHostEnvironment env, string name, ModelLo
             if (ctx.Header.ModelVerWritten >= VerCategoricalSplitSerialized)
                 categoricalSplits = true;
 
-            TrainedEnsemble = new TreeEnsemble(ctx, usingDefaultValues, categoricalSplits);
-            MaxSplitFeatIdx = FindMaxFeatureIndex(TrainedEnsemble);
+            TrainedEnsemble = new InternalTreeEnsemble(ctx, usingDefaultValues, categoricalSplits);
+            MaxSplitFeatIdx = TrainedEnsemble.GetMaxFeatureIndex();
 
             InnerArgs = ctx.LoadStringOrNull();
             if (ctx.Header.ModelVerWritten >= VerNumFeaturesSerialized)
@@ -3195,7 +3199,7 @@ private void SaveEnsembleAsCode(TextWriter writer, RoleMappedSchema schema)
             MetadataUtils.GetSlotNames(schema, RoleMappedSchema.ColumnRole.Feature, NumFeatures, ref names);
 
             int i = 0;
-            foreach (RegressionTree tree in TrainedEnsemble.Trees)
+            foreach (InternalRegressionTree tree in TrainedEnsemble.Trees)
             {
                 writer.Write("double treeOutput{0}=", i);
                 SaveTreeAsCode(tree, writer, in names);
@@ -3211,13 +3215,13 @@ private void SaveEnsembleAsCode(TextWriter writer, RoleMappedSchema schema)
         /// <summary>
         /// Convert a single tree to code, called recursively
         /// </summary>
-        private void SaveTreeAsCode(RegressionTree tree, TextWriter writer, in VBuffer<ReadOnlyMemory<char>> names)
+        private void SaveTreeAsCode(InternalRegressionTree tree, TextWriter writer, in VBuffer<ReadOnlyMemory<char>> names)
         {
             ToCSharp(tree, writer, 0, in names);
         }
 
         // converts a subtree into a C# expression
-        private void ToCSharp(RegressionTree tree, TextWriter writer, int node, in VBuffer<ReadOnlyMemory<char>> names)
+        private void ToCSharp(InternalRegressionTree tree, TextWriter writer, int node, in VBuffer<ReadOnlyMemory<char>> names)
         {
             if (node < 0)
             {
@@ -3259,23 +3263,6 @@ public void GetFeatureWeights(ref VBuffer<Float> weights)
             bldr.GetResult(ref weights);
         }
 
-        private static int FindMaxFeatureIndex(TreeEnsemble ensemble)
-        {
-            int ifeatMax = 0;
-            for (int i = 0; i < ensemble.NumTrees; i++)
-            {
-                var tree = ensemble.GetTreeAt(i);
-                for (int n = 0; n < tree.NumNodes; n++)
-                {
-                    int ifeat = tree.SplitFeature(n);
-                    if (ifeat > ifeatMax)
-                        ifeatMax = ifeat;
-                }
-            }
-
-            return ifeatMax;
-        }
-
         ITree[] ITreeEnsemble.GetTrees()
         {
             return TrainedEnsemble.Trees.Select(k => new Tree(k)).ToArray();
@@ -3318,9 +3305,9 @@ Row ICanGetSummaryAsIRow.GetStatsIRowOrNull(RoleMappedSchema schema)
 
         private sealed class Tree : ITree<VBuffer<Float>>
         {
-            private readonly RegressionTree _regTree;
+            private readonly InternalRegressionTree _regTree;
 
-            public Tree(RegressionTree regTree)
+            public Tree(InternalRegressionTree regTree)
             {
                 _regTree = regTree;
             }
@@ -3390,4 +3377,82 @@ public TreeNode(Dictionary<string, object> keyValues)
             public Dictionary<string, object> KeyValues { get; }
         }
     }
+
+    /// <summary>
+    /// <see cref="TreeEnsembleModelParametersBasedOnRegressionTree"/> is derived from
+    /// <see cref="TreeEnsembleModelParameters"/> plus a strongly-typed public attribute,
+    /// <see cref="TrainedTreeEnsemble"/>, for exposing trained model's details to users.
+    /// Its function, <see cref="CreateTreeEnsembleFromInternalDataStructure"/>, is
+    /// called to create <see cref="TrainedTreeEnsemble"/> inside <see cref="TreeEnsembleModelParameters"/>.
+    /// Note that the major difference between <see cref="TreeEnsembleModelParametersBasedOnQuantileRegressionTree"/>
+    /// and <see cref="TreeEnsembleModelParametersBasedOnRegressionTree"/> is the type of
+    /// <see cref="TrainedTreeEnsemble"/>.
+    /// </summary>
+    public abstract class TreeEnsembleModelParametersBasedOnRegressionTree : TreeEnsembleModelParameters
+    {
+        /// <summary>
+        /// An ensemble of trees exposed to users. It is a wrapper on the <see langword="internal"/>
+        /// <see cref="InternalTreeEnsemble"/> in <see cref="ML.FastTree.TreeEnsemble{T}"/>.
+        /// </summary>
+        public RegressionTreeEnsemble TrainedTreeEnsemble { get; }
+
+        [BestFriend]
+        internal TreeEnsembleModelParametersBasedOnRegressionTree(IHostEnvironment env, string name, InternalTreeEnsemble trainedEnsemble, int numFeatures, string innerArgs)
+            : base(env, name, trainedEnsemble, numFeatures, innerArgs)
+        {
+            TrainedTreeEnsemble = CreateTreeEnsembleFromInternalDataStructure();
+        }
+
+        protected TreeEnsembleModelParametersBasedOnRegressionTree(IHostEnvironment env, string name, ModelLoadContext ctx, VersionInfo ver)
+            : base(env, name, ctx, ver)
+        {
+            TrainedTreeEnsemble = CreateTreeEnsembleFromInternalDataStructure();
+        }
+
+        private RegressionTreeEnsemble CreateTreeEnsembleFromInternalDataStructure()
+        {
+            var trees = TrainedEnsemble.Trees.Select(tree => new RegressionTree(tree));
+            var treeWeights = TrainedEnsemble.Trees.Select(tree => tree.Weight);
+            return new RegressionTreeEnsemble(trees, treeWeights, TrainedEnsemble.Bias);
+        }
+    }
+
+    /// <summary>
+    /// <see cref="TreeEnsembleModelParametersBasedOnQuantileRegressionTree"/> is derived from
+    /// <see cref="TreeEnsembleModelParameters"/> plus a strongly-typed public attribute,
+    /// <see cref="TrainedTreeEnsemble"/>, for exposing trained model's details to users.
+    /// Its function, <see cref="CreateTreeEnsembleFromInternalDataStructure"/>, is
+    /// called to create <see cref="TrainedTreeEnsemble"/> inside <see cref="TreeEnsembleModelParameters"/>.
+    /// Note that the major difference between <see cref="TreeEnsembleModelParametersBasedOnQuantileRegressionTree"/>
+    /// and <see cref="TreeEnsembleModelParametersBasedOnRegressionTree"/> is the type of
+    /// <see cref="TrainedTreeEnsemble"/>.
+    /// </summary>
+    public abstract class TreeEnsembleModelParametersBasedOnQuantileRegressionTree : TreeEnsembleModelParameters
+    {
+        /// <summary>
+        /// An ensemble of trees exposed to users. It is a wrapper on the <see langword="internal"/>
+        /// <see cref="InternalTreeEnsemble"/> in <see cref="ML.FastTree.TreeEnsemble{T}"/>.
+        /// </summary>
+        public QuantileRegressionTreeEnsemble TrainedTreeEnsemble { get; }
+
+        [BestFriend]
+        internal TreeEnsembleModelParametersBasedOnQuantileRegressionTree(IHostEnvironment env, string name, InternalTreeEnsemble trainedEnsemble, int numFeatures, string innerArgs)
+            : base(env, name, trainedEnsemble, numFeatures, innerArgs)
+        {
+            TrainedTreeEnsemble = CreateTreeEnsembleFromInternalDataStructure();
+        }
+
+        protected TreeEnsembleModelParametersBasedOnQuantileRegressionTree(IHostEnvironment env, string name, ModelLoadContext ctx, VersionInfo ver)
+            : base(env, name, ctx, ver)
+        {
+            TrainedTreeEnsemble = CreateTreeEnsembleFromInternalDataStructure();
+        }
+
+        private QuantileRegressionTreeEnsemble CreateTreeEnsembleFromInternalDataStructure()
+        {
+            var trees = TrainedEnsemble.Trees.Select(tree => new QuantileRegressionTree((InternalQuantileRegressionTree)tree));
+            var treeWeights = TrainedEnsemble.Trees.Select(tree => tree.Weight);
+            return new QuantileRegressionTreeEnsemble(trees, treeWeights, TrainedEnsemble.Bias);
+        }
+    }
 }

src/Microsoft.ML.FastTree/FastTreeArguments.cs

@@ -156,7 +156,7 @@ public abstract class TreeArgs : LearnerInputBaseWithGroupId
         /// Allows to choose Parallel FastTree Learning Algorithm.
         /// </summary>
         [Argument(ArgumentType.Multiple, HelpText = "Allows to choose Parallel FastTree Learning Algorithm", ShortName = "parag")]
-        public ISupportParallelTraining ParallelTrainer = new SingleTrainerFactory();
+        internal ISupportParallelTraining ParallelTrainer = new SingleTrainerFactory();
 
         /// <summary>
         /// The number of threads to use.

src/Microsoft.ML.FastTree/FastTreeClassification.cs

@@ -44,7 +44,7 @@
 namespace Microsoft.ML.Trainers.FastTree
 {
     public sealed class FastTreeBinaryModelParameters :
-        TreeEnsembleModelParameters
+        TreeEnsembleModelParametersBasedOnRegressionTree
     {
         internal const string LoaderSignature = "FastTreeBinaryExec";
         internal const string RegistrationName = "FastTreeBinaryPredictor";
@@ -70,7 +70,7 @@ private static VersionInfo GetVersionInfo()
 
         protected override uint VerCategoricalSplitSerialized => 0x00010005;
 
-        public FastTreeBinaryModelParameters(IHostEnvironment env, TreeEnsemble trainedEnsemble, int featureCount, string innerArgs)
+        internal FastTreeBinaryModelParameters(IHostEnvironment env, InternalTreeEnsemble trainedEnsemble, int featureCount, string innerArgs)
             : base(env, RegistrationName, trainedEnsemble, featureCount, innerArgs)
         {
         }
@@ -204,7 +204,7 @@ protected override ObjectiveFunctionBase ConstructObjFunc(IChannel ch)
                 ParallelTraining);
         }
 
-        protected override OptimizationAlgorithm ConstructOptimizationAlgorithm(IChannel ch)
+        private protected override OptimizationAlgorithm ConstructOptimizationAlgorithm(IChannel ch)
         {
             OptimizationAlgorithm optimizationAlgorithm = base.ConstructOptimizationAlgorithm(ch);
             if (Args.UseLineSearch)
@@ -365,7 +365,7 @@ protected override void GetGradientInOneQuery(int query, int threadIndex)
                 }
             }
 
-            public void AdjustTreeOutputs(IChannel ch, RegressionTree tree,
+            public void AdjustTreeOutputs(IChannel ch, InternalRegressionTree tree,
                 DocumentPartitioning partitioning, ScoreTracker trainingScores)
             {
                 const double epsilon = 1.4e-45;