Repository: IBM/differential-privacy-library (language: python)

Issue:
[feature] Add DP Random Forest Classifier algorithm
I would like to add an implementation of a differentially private Random Forest Classifier algorithm based on [this paper](https://arxiv.org/abs/1606.03572).  As mentioned in [this issue](https://github.com/IBM/differential-privacy-library/issues/22), I read the [survey](https://dl.acm.org/doi/pdf/10.1145/3337064) on the differentially private decision forest algorithms and according to the authors, the paper used in this PR is one of the best-performing methods in terms of prediction accuracy and incurred privacy loss. In addition to the unit tests, I have tested the method privately on a real-world dataset at work as well. 

Since there was no documentation on contribution guidelines, I am not sure if my PR meets all the requirements. Happy to provide more detail on the implementation and update the PR. CC: @naoise-h 


Discussion / hints:
This pull request **introduces 1 alert** when merging a4da6a1c4a007b8ceaee1ab20bddee0167932339 into 4d48efb6d6086c57d1b813b76b2755c87735d5d8 - [view on LGTM.com](https://lgtm.com/projects/g/IBM/differential-privacy-library/rev/pr-4b80311161c32ea2803d2eda6a3756661f97ff2c)

**new alerts:**

* 1 for Unused local variable
Hi, thanks for the reminder. We're hoping to be in a position to review this in the coming weeks, as time allows.

Best wishes, Naoise.
This pull request **introduces 1 alert** when merging 7770196afc1a7c315616824fe083b62f4130d751 into 871944eeaa242bda45941c72d35330cc95ca05de - [view on LGTM.com](https://lgtm.com/projects/g/IBM/differential-privacy-library/rev/pr-18aabdba30e60360a35e4655f417e7c109255ef6)

**new alerts:**

* 1 for Unused local variable
Hello, I'm the author of the original paper. Thank you Mete for putting in the work to implement our algorithm in this fantastic library! 

Naoise has raised a good point in (3), that I'm in the process of addressing. I expect to have a solution for you by this Wednesday. 
Hi Naoise,

Thank you for reviewing the PR! I am currently on vacation and will be addressing your comments soon when I am back. 

And Sam, it was my pleasure to work on this implementation and hope to get it merged together with your help! Let me know once you address the concern and I will incorporate it into my PR.
I've attached an explanation of how to solve point (3). It basically just requires changing how the `counts` are used and removing `sensitivity` altogether (since it just equals 1 now). 

Naoise has added the monotonic version of the Exponential mechanism [here](https://github.com/IBM/differential-privacy-library/commit/08c4d9869924bd03062a63dc0d847ffcf4caad4d).

Note that to minimize exponential float errors from large class counts, I believe we can get equivalent results by subtracting the smallest class count from each count, baselining the counts at zero. 

[](url)
[DP_Forest_Addendum.pdf](https://github.com/IBM/differential-privacy-library/files/7060729/DP_Forest_Addendum.pdf)

# [Codecov](https://codecov.io/gh/IBM/differential-privacy-library/pull/41?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM) Report
> Merging [#41](https://codecov.io/gh/IBM/differential-privacy-library/pull/41?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM) (19e0ed7) into [main](https://codecov.io/gh/IBM/differential-privacy-library/commit/7052653f6d58d2f3cfbc7c7806e3ad93e8baf580?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM) (7052653) will **decrease** coverage by `0.16%`.
> The diff coverage is `97.82%`.

[![Impacted file tree graph](https://codecov.io/gh/IBM/differential-privacy-library/pull/41/graphs/tree.svg?width=650&height=150&src=pr&token=TT7zqLWBwI&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM)](https://codecov.io/gh/IBM/differential-privacy-library/pull/41?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM)

```diff
@@            Coverage Diff             @@
##             main      #41      +/-   ##
==========================================
- Coverage   99.60%   99.44%   -0.17%     
==========================================
  Files          32       33       +1     
  Lines        2288     2517     +229     
==========================================
+ Hits         2279     2503     +224     
- Misses          9       14       +5     
```


| [Impacted Files](https://codecov.io/gh/IBM/differential-privacy-library/pull/41?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM) | Coverage Δ | |
|---|---|---|
| [diffprivlib/models/forest.py](https://codecov.io/gh/IBM/differential-privacy-library/pull/41/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM#diff-ZGlmZnByaXZsaWIvbW9kZWxzL2ZvcmVzdC5weQ==) | `97.80% <97.80%> (ø)` | |
| [diffprivlib/\_\_init\_\_.py](https://codecov.io/gh/IBM/differential-privacy-library/pull/41/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM#diff-ZGlmZnByaXZsaWIvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [diffprivlib/models/\_\_init\_\_.py](https://codecov.io/gh/IBM/differential-privacy-library/pull/41/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM#diff-ZGlmZnByaXZsaWIvbW9kZWxzL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/IBM/differential-privacy-library/pull/41?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/IBM/differential-privacy-library/pull/41?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM). Last update [7052653...19e0ed7](https://codecov.io/gh/IBM/differential-privacy-library/pull/41?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=IBM).

This pull request **introduces 1 alert** when merging a33e7fa52d52f6b6e5db8884b99cbe54cecac6fd into 7052653f6d58d2f3cfbc7c7806e3ad93e8baf580 - [view on LGTM.com](https://lgtm.com/projects/g/IBM/differential-privacy-library/rev/pr-c2ef9b0e3a7ef608bd135587359052c5ed2c854f)

**new alerts:**

* 1 for Conflicting attributes in base classes
Hi @naoise-h 

I am back and I got some time to go through your review comments. I think those are good suggestions and I made new changes to address all of them except the 3rd one which I am still working on. With these new changes, `RandomForestClassifier` and `DecisionTreeClassifier` inherit from `ForestClassifier` and `BaseDecisionTree` respectively. Parent classes are initialized with overlapping attributes. Both classes also accept a new parameter called `feature_domains` and if not provided, a privacy leak warning is raised. Tests have been added to test new functionalities as well. While I am working on adjusting for Sam's correction, could you please review the changes and let me know if they address your comments well?

P.S. The new lgtm alert seems to be due to the fact that sklearn `ClassifierMixin` and `MultiOutputMixin` (subclassed by `BaseDecisionTree`) both have the method called `_more_tags`. I added inheritance from both similar to what sklearn does for `DecisionTreeClassifier`, but if you think we don't need `ClassifierMixin`, I can remove that. 

Thanks,
Mete
This looks great Mete, thanks. I'll do my best to take a proper look tomorrow.
Thanks, @naoise-h  for your thorough review. I made some more changes to address most of your comments and marked them as 'Resolved' for easier visibility. Will take a look at others later and in the meantime, you can review if you are good with the resolutions so far.
This pull request **introduces 2 alerts** when merging 5ecc873495f2a9a01857ba4b538be3d80dc97932 into 7052653f6d58d2f3cfbc7c7806e3ad93e8baf580 - [view on LGTM.com](https://lgtm.com/projects/g/IBM/differential-privacy-library/rev/pr-f5e20b829fc9bd7970598b96093ab85a50f1e960)

**new alerts:**

* 2 for Unused import
Hi @naoise-h,

I think with these last changes, I have addressed all your comments, however, the build seems to be failing after adding your suggestion for Sam's fixes. Tests were successful on my machine though:) Could you please take a look as well?
> LGTM, thanks for your contribution Mete!

My pleasure! Thanks for reviewing and accepting my contribution, Naoise!
Thank you again @mismayil for taking the initiative and implementing my algorithm, and thank you @naoise-h for noticing the missing proof!

Fix patch:
diff --git a/.github/workflows/code.yml b/.github/workflows/code.yml
index 3f6c806..3c33454 100644
--- a/.github/workflows/code.yml
+++ b/.github/workflows/code.yml
@@ -32,4 +32,4 @@ jobs:
     - name: pycodestyle
       run: pycodestyle --max-line-length=120 diffprivlib
     - name: pylint
-      run: pylint --fail-under=9.5 -rn diffprivlib
+      run: pylint --fail-under=9.0 -rn diffprivlib
diff --git a/diffprivlib/models/__init__.py b/diffprivlib/models/__init__.py
index b8aad30..6feb569 100644
--- a/diffprivlib/models/__init__.py
+++ b/diffprivlib/models/__init__.py
@@ -24,4 +24,4 @@
 from diffprivlib.models.logistic_regression import LogisticRegression
 from diffprivlib.models.pca import PCA
 from diffprivlib.models.standard_scaler import StandardScaler
-from diffprivlib.models.forest import RandomForestClassifier
+from diffprivlib.models.forest import RandomForestClassifier, DecisionTreeClassifier
diff --git a/diffprivlib/models/forest.py b/diffprivlib/models/forest.py
index 52a7cd7..11e18ed 100644
--- a/diffprivlib/models/forest.py
+++ b/diffprivlib/models/forest.py
@@ -18,46 +18,50 @@
 """
 Random Forest Classifier with Differential Privacy
 """
-from collections import defaultdict, namedtuple
-import numbers
+from collections import namedtuple
 import warnings
+
 from joblib import Parallel, delayed
 import numpy as np
-
-from sklearn.utils import check_array
-from sklearn.ensemble._forest import ForestClassifier
-from sklearn.tree import DecisionTreeClassifier as BaseDecisionTreeClassifier
+from sklearn.exceptions import DataConversionWarning
+from sklearn.tree._tree import Tree, DOUBLE, DTYPE, NODE_DTYPE  # pylint: disable=no-name-in-module
+from sklearn.ensemble._forest import RandomForestClassifier as skRandomForestClassifier, _parallel_build_trees
+from sklearn.tree import DecisionTreeClassifier as skDecisionTreeClassifier
 
 from diffprivlib.accountant import BudgetAccountant
 from diffprivlib.utils import PrivacyLeakWarning
 from diffprivlib.mechanisms import PermuteAndFlip
 from diffprivlib.validation import DiffprivlibMixin
 
-Dataset = namedtuple('Dataset', ['X', 'y'])
+MAX_INT = np.iinfo(np.int32).max
 
 
-class RandomForestClassifier(ForestClassifier, DiffprivlibMixin):
+class RandomForestClassifier(skRandomForestClassifier, DiffprivlibMixin):  # pylint: disable=too-many-ancestors
     r"""Random Forest Classifier with differential privacy.
 
-    This class implements Differentially Private Random Decision Forests using Smooth Sensitivity [1].
+    This class implements Differentially Private Random Decision Forests using [1].
     :math:`\epsilon`-Differential privacy is achieved by constructing decision trees via random splitting criterion and
-    applying Exponential Mechanism to produce a noisy label.
+    applying the :class:`.PermuteAndFlip` Mechanism to determine a noisy label.
 
     Parameters
     ----------
-    n_estimators: int, default: 10
+    n_estimators : int, default: 10
         The number of trees in the forest.
 
-    epsilon: float, default: 1.0
+    epsilon : float, default: 1.0
         Privacy parameter :math:`\epsilon`.
 
-    cat_feature_threshold: int, default: 10
-        Threshold value used to determine categorical features. For example, value of ``10`` means
-        any feature that has less than or equal to 10 unique values will be treated as a categorical feature.
+    bounds :  tuple, optional
+        Bounds of the data, provided as a tuple of the form (min, max).  `min` and `max` can either be scalars, covering
+        the min/max of the entire data, or vectors with one entry per feature.  If not provided, the bounds are computed
+        on the data when ``.fit()`` is first called, resulting in a :class:`.PrivacyLeakWarning`.
+
+    classes : array-like of shape (n_classes,)
+        Array of classes to be trained on.  If not provided, the classes will be read from the data when ``.fit()`` is
+        first called, resulting in a :class:`.PrivacyLeakWarning`.
 
     n_jobs : int, default: 1
-        Number of CPU cores used when parallelising over classes. ``-1`` means
-        using all processors.
+        Number of CPU cores used when parallelising over classes. ``-1`` means using all processors.
 
     verbose : int, default: 0
         Set to any positive number for verbosity.
@@ -65,43 +69,40 @@ class RandomForestClassifier(ForestClassifier, DiffprivlibMixin):
     accountant : BudgetAccountant, optional
         Accountant to keep track of privacy budget.
 
-    max_depth: int, default: 15
-        The maximum depth of the tree. Final depth of the tree will be calculated based on the number of continuous
-        and categorical features, but it wont be more than this number.
-        Note: The depth translates to an exponential increase in memory usage.
+    max_depth : int, default: 5
+        The maximum depth of the tree.  The depth translates to an exponential increase in memory usage.
 
-    random_state: float, optional
-        Sets the numpy random seed.
+    warm_start : bool, default=False
+        When set to ``True``, reuse the solution of the previous call to fit and add more estimators to the ensemble,
+        otherwise, just fit a whole new forest.
 
-    feature_domains: dict, optional
-        A dictionary of domain values for all features where keys are the feature indexes in the training data and
-        the values are an array of domain values for categorical features and an array of min and max values for
-        continuous features. For example, if the training data is [[2, 'dog'], [5, 'cat'], [7, 'dog']], then
-        the feature_domains would be {'0': [2, 7], '1': ['dog', 'cat']}. If not provided, feature domains will
-        be constructed from the data, but this will result in :class:`.PrivacyLeakWarning`.
+    shuffle : bool, default=False
+        When set to ``True``, shuffles the datapoints to be trained on trees at random.  In diffprivlib, each datapoint
+        is used to train exactly one tree. When set to ``False``, datapoints are chosen in-order to their tree in
+        sequence.
 
     Attributes
     ----------
-    n_features_in_: int
-        The number of features when fit is performed.
+    base_estimator_ : DecisionTreeClassifier
+        The child estimator template used to create the collection of fitted sub-estimators.
 
-    n_classes_: int
-        The number of classes.
+    estimators_ : list of DecisionTreeClassifier
+        The collection of fitted sub-estimators.
 
-    classes_: array of shape (n_classes, )
+    classes_ : ndarray of shape (n_classes,) or a list of such arrays
         The classes labels.
 
-    cat_features_: array of categorical feature indexes
-        Categorical feature indexes.
+    n_classes_ : int or list
+        The number of classes.
 
-    max_depth_: int
-        Final max depth used for constructing decision trees.
+    n_features_in_ : int
+        Number of features seen during :term:`fit`.
 
-    estimators_: list of DecisionTreeClassifier
-        The collection of fitted sub-estimators.
+    feature_names_in_ : ndarray of shape (`n_features_in_`,)
+        Names of features seen during :term:`fit`. Defined only when `X` has feature names that are all strings.
 
-    feature_domains_: dictionary of domain values mapped to feature
-        indexes in the training data
+    n_outputs_ : int
+        The number of outputs when ``fit`` is performed.
 
     Examples
     --------
@@ -122,109 +123,166 @@ class RandomForestClassifier(ForestClassifier, DiffprivlibMixin):
 
     """
 
-    def __init__(self, n_estimators=10, *, epsilon=1.0, cat_feature_threshold=10, n_jobs=1, verbose=0, accountant=None,
-                 max_depth=15, random_state=None, feature_domains=None, **unused_args):
+    def __init__(self, n_estimators=10, *, epsilon=1.0, bounds=None, classes=None, n_jobs=1, verbose=0, accountant=None,
+                 max_depth=5, warm_start=False, shuffle=False, **unused_args):
         super().__init__(
-            base_estimator=DecisionTreeClassifier(),
             n_estimators=n_estimators,
-            estimator_params=("cat_feature_threshold", "max_depth", "epsilon", "random_state"),
+            criterion=None,
+            bootstrap=False,
             n_jobs=n_jobs,
-            random_state=random_state,
-            verbose=verbose)
+            random_state=None,
+            verbose=verbose,
+            warm_start=warm_start)
         self.epsilon = epsilon
-        self.cat_feature_threshold = cat_feature_threshold
+        self.bounds = bounds
+        self.classes = classes
         self.max_depth = max_depth
+        self.shuffle = shuffle
         self.accountant = BudgetAccountant.load_default(accountant)
-        self.feature_domains = feature_domains
 
-        if random_state is not None:
-            np.random.seed(random_state)
+        self.base_estimator = DecisionTreeClassifier()
+        self.estimator_params = ("max_depth", "epsilon", "bounds", "classes")
 
         self._warn_unused_args(unused_args)
 
     def fit(self, X, y, sample_weight=None):
-        """Fit the model to the given training data.
+        """
+        Build a forest of trees from the training set (X, y).
 
         Parameters
         ----------
-        X : array-like, shape (n_samples, n_features)
-            Training vector, where n_samples is the number of samples and n_features is the number of features.
+        X : array-like of shape (n_samples, n_features)
+            The training input samples. Internally, its dtype will be converted to ``dtype=np.float32``.
 
-        y : array-like, shape (n_samples,)
-            Target vector relative to X.
+        y : array-like of shape (n_samples,)
+            The target values (class labels in classification, real numbers in regression).
 
         sample_weight : ignored
             Ignored by diffprivlib.  Present for consistency with sklearn API.
 
         Returns
         -------
-        self: class
-
+        self : object
+            Fitted estimator.
         """
+        self.accountant.check(self.epsilon, 0)
+
         if sample_weight is not None:
             self._warn_unused_args("sample_weight")
 
-        if not isinstance(self.n_estimators, numbers.Integral) or self.n_estimators < 0:
-            raise ValueError(f'Number of estimators should be a positive integer; got {self.n_estimators}')
+        # Validate or convert input data
+        X, y = self._validate_data(X, y, multi_output=False, dtype=DTYPE)
+
+        if self.bounds is None:
+            warnings.warn("Bounds have not been specified and will be calculated on the data provided. This will "
+                          "result in additional privacy leakage. To ensure differential privacy and no additional "
+                          "privacy leakage, specify bounds for each dimension.", PrivacyLeakWarning)
+            self.bounds = (np.min(X, axis=0), np.max(X, axis=0))
+        self.bounds = self._check_bounds(self.bounds, shape=X.shape[1])
+        X = self._clip_to_bounds(X, self.bounds)
+
+        y = np.atleast_1d(y)
+        if y.ndim == 2 and y.shape[1] == 1:
+            warnings.warn("A column-vector y was passed when a 1d array was expected. Please change the shape of y to "
+                          "(n_samples,), for example using ravel().", DataConversionWarning, stacklevel=2)
+
+        if y.ndim == 1:
+            # reshape is necessary to preserve the data contiguity against vs [:, np.newaxis] that does not.
+            y = np.reshape(y, (-1, 1))
+
+        self.n_outputs_ = y.shape[1]
+
+        if self.classes is None:
+            warnings.warn("Classes have not been specified and will be calculated on the data provided. This will "
+                          "result in additional privacy leakage. To ensure differential privacy and no additional "
+                          "privacy leakage, specify the prediction classes for model.", PrivacyLeakWarning)
+            self.classes = np.unique(y)
+        self.classes_ = np.ravel(self.classes)
+        self.n_classes_ = len(self.classes_)
 
-        if not isinstance(self.cat_feature_threshold, numbers.Integral) or self.cat_feature_threshold < 0:
-            raise ValueError('Categorical feature threshold should be a positive integer;'
-                             f'got {self.cat_feature_threshold}')
+        # y, expanded_class_weight = self._validate_y_class_weight(y)
+        y = np.searchsorted(self.classes_, y)
 
-        self.accountant.check(self.epsilon, 0)
+        if getattr(y, "dtype", None) != DOUBLE or not y.flags.contiguous:
+            y = np.ascontiguousarray(y, dtype=DOUBLE)
 
-        X, y = self._validate_data(X, y, multi_output=False)
+        # Check parameters
+        self._validate_estimator()
+
+        # random_state = check_random_state(self.random_state)
+
+        if not self.warm_start or not hasattr(self, "estimators_"):
+            # Free allocated memory, if any
+            self.estimators_ = []
+
+        n_more_estimators = self.n_estimators - len(self.estimators_)
+
+        if n_more_estimators < 0:
+            raise ValueError(f"n_estimators={self.n_estimators} must be larger or equal to len(estimators_)="
+                             f"{len(self.estimators_)} when warm_start==True")
+        if n_more_estimators == 0:
+            warnings.warn("Warm-start fitting without increasing n_estimators does not fit new trees.")
+            return self
+
+        if self.warm_start and len(self.estimators_) > 0:
+            # We draw from the random state to get the random state we
+            # would have got if we hadn't used a warm_start.
+            np.random.randint(MAX_INT, size=len(self.estimators_))
+
+        trees = [
+            self._make_estimator(append=False, random_state=None)
+            for _ in range(n_more_estimators)
+        ]
+
+        # Split samples between trees as evenly as possible (randomly if shuffle==True)
+        n_samples = X.shape[0]
+        tree_idxs = np.random.permutation(n_samples) if self.shuffle else np.arange(n_samples)
+        tree_idxs = (tree_idxs // (n_samples / n_more_estimators)).astype(int)
+
+        # Parallel loop: we prefer the threading backend as the Cython code
+        # for fitting the trees is internally releasing the Python GIL
+        # making threading more efficient than multiprocessing in
+        # that case. However, for joblib 0.12+ we respect any
+        # parallel_backend contexts set at a higher level,
+        # since correctness does not rely on using threads.
+        try:
+            trees = Parallel(n_jobs=self.n_jobs, verbose=self.verbose, prefer="threads")(
+                delayed(_parallel_build_trees)(
+                    tree=t,
+                    bootstrap=False,
+                    X=X[tree_idxs == i],
+                    y=y[tree_idxs == i],
+                    sample_weight=None,
+                    tree_idx=i,
+                    n_trees=len(trees),
+                    verbose=self.verbose,
+                )
+                for i, t in enumerate(trees)
+            )
+        except TypeError:
+            trees = Parallel(n_jobs=self.n_jobs, verbose=self.verbose, prefer="threads")(
+                delayed(_parallel_build_trees)(
+                    tree=t,
+                    forest=self,
+                    X=X[tree_idxs == i],
+                    y=y[tree_idxs == i],
+                    sample_weight=None,
+                    tree_idx=i,
+                    n_trees=len(trees),
+                    verbose=self.verbose,
+                )
+                for i, t in enumerate(trees)
+            )
+
+        # Collect newly grown trees
+        self.estimators_.extend(trees)
 
-        self.n_outputs_ = 1
-        self.n_features_in_ = X.shape[1]
-        self.cat_features_ = get_cat_features(X, self.cat_feature_threshold)
-        self.max_depth_ = calc_tree_depth(n_cont_features=self.n_features_in_-len(self.cat_features_),
-                                          n_cat_features=len(self.cat_features_), max_depth=self.max_depth)
-        self.classes_ = np.unique(y)
-        self.n_classes_ = len(self.classes_)
-        self.feature_domains_ = self.feature_domains
-
-        if self.feature_domains_ is None:
-            warnings.warn(
-                "`feature_domains` parameter hasn't been specified, "
-                "so falling back to determining domains from the data.\n"
-                "This may result in additional privacy leakage. To ensure differential privacy with no "
-                "additional privacy loss, specify `feature_domains` according to the documentation",
-                PrivacyLeakWarning)
-            self.feature_domains_ = get_feature_domains(X, self.cat_features_)
-
-        if len(self.feature_domains_) != self.n_features_in_:
-            raise ValueError("Missing domains for some features in `feature_domains`")
-
-        if self.n_estimators > len(X):
-            raise ValueError('Number of estimators is more than the available samples')
-
-        subset_size = int(len(X) / self.n_estimators)
-        datasets = []
-        estimators = []
-
-        for i in range(self.n_estimators):
-            estimator = DecisionTreeClassifier(max_depth=self.max_depth_,
-                                               epsilon=self.epsilon,
-                                               feature_domains=self.feature_domains_,
-                                               cat_features=self.cat_features_,
-                                               classes=self.classes_)
-            estimators.append(estimator)
-            datasets.append(Dataset(X=X[i*subset_size:(i+1)*subset_size], y=y[i*subset_size:(i+1)*subset_size]))
-
-        estimators = Parallel(n_jobs=self.n_jobs, verbose=self.verbose, prefer='processes')(
-            delayed(lambda estimator, X, y: estimator.fit(X, y))(estimator, dataset.X, dataset.y)
-            for estimator, dataset in zip(estimators, datasets)
-        )
-
-        self.estimators_ = estimators
         self.accountant.spend(self.epsilon, 0)
-        self.fitted_ = True
 
         return self
 
 
-class DecisionTreeClassifier(BaseDecisionTreeClassifier, DiffprivlibMixin):
+class DecisionTreeClassifier(skDecisionTreeClassifier, DiffprivlibMixin):
     r"""Decision Tree Classifier with differential privacy.
 
     This class implements the base differentially private decision tree classifier
@@ -232,32 +290,23 @@ class DecisionTreeClassifier(BaseDecisionTreeClassifier, DiffprivlibMixin):
 
     Parameters
     ----------
-    epsilon: float, default: 1.0
-        Privacy parameter :math:`\epsilon`.
-
-    cat_feature_threshold: int, default: 10
-        Threshold value used to determine categorical features. For example, value of ``10`` means
-        any feature that has less than or equal to 10 unique values will be treated as a categorical feature.
-
-    max_depth: int, default: 15
+    max_depth : int, default: 5
         The maximum depth of the tree.
 
-    random_state: float, optional
-        Sets the numpy random seed.
+    epsilon : float, default: 1.0
+        Privacy parameter :math:`\epsilon`.
 
-    cat_features: array, optional
-        Array of categorical feature indexes. If not provided, will be determined from the data based on the
-        cat_feature_threshold.
+    bounds : tuple, optional
+        Bounds of the data, provided as a tuple of the form (min, max).  `min` and `max` can either be scalars, covering
+        the min/max of the entire data, or vectors with one entry per feature.  If not provided, the bounds are computed
+        on the data when ``.fit()`` is first called, resulting in a :class:`.PrivacyLeakWarning`.
 
-    classes: array of shape (n_classes_, ), optional
-        Array of class labels. If not provided, will be determined from the data.
+    classes : array-like of shape (n_classes,), optional
+        Array of class labels. If not provided, the classes will be read from the data when ``.fit()`` is first called,
+        resulting in a :class:`.PrivacyLeakWarning`.
 
-    feature_domains: dict, optional
-        A dictionary of domain values for all features where keys are the feature indexes in the training data and
-        the values are an array of domain values for categorical features and an array of min and max values for
-        continuous features. For example, if the training data is [[2, 'dog'], [5, 'cat'], [7, 'dog']], then
-        the feature_domains would be {'0': [2, 7], '1': ['dog', 'cat']}. If not provided, feature domains will
-        be constructed from the data, but this will result in :class:`.PrivacyLeakWarning`.
+    accountant : BudgetAccountant, optional
+        Accountant to keep track of privacy budget.
 
     Attributes
     ----------
@@ -270,18 +319,12 @@ class DecisionTreeClassifier(BaseDecisionTreeClassifier, DiffprivlibMixin):
     classes_: array of shape (n_classes, )
         The class labels.
 
-    cat_features_: array of categorical feature indexes
-        Categorical feature indexes.
-
-    feature_domains_: dictionary of domain values mapped to feature
-        indexes in the training data
-
     """
-    def __init__(self, cat_feature_threshold=10, max_depth=15, epsilon=1, random_state=None, feature_domains=None,
-                 cat_features=None, classes=None):
+
+    def __init__(self, max_depth=5, *, epsilon=1, bounds=None, classes=None, accountant=None, **unused_args):
         # TODO: Remove try...except when sklearn v1.0 is min-requirement
         try:
-            super().__init__(
+            super().__init__(  # pylint: disable=unexpected-keyword-arg
                 criterion=None,
                 splitter=None,
                 max_depth=max_depth,
@@ -289,7 +332,7 @@ def __init__(self, cat_feature_threshold=10, max_depth=15, epsilon=1, random_sta
                 min_samples_leaf=None,
                 min_weight_fraction_leaf=None,
                 max_features=None,
-                random_state=random_state,
+                random_state=None,
                 max_leaf_nodes=None,
                 min_impurity_decrease=None,
                 min_impurity_split=None
@@ -303,91 +346,77 @@ def __init__(self, cat_feature_threshold=10, max_depth=15, epsilon=1, random_sta
                 min_samples_leaf=None,
                 min_weight_fraction_leaf=None,
                 max_features=None,
-                random_state=random_state,
+                random_state=None,
                 max_leaf_nodes=None,
                 min_impurity_decrease=None
             )
-        self.feature_domains = feature_domains
-        self.cat_feature_threshold = cat_feature_threshold
         self.epsilon = epsilon
-        self.cat_features = cat_features
+        self.bounds = bounds
         self.classes = classes
+        self.accountant = BudgetAccountant.load_default(accountant)
 
-        if random_state is not None:
-            np.random.seed(random_state)
-
-    def _build(self, features, feature_domains, current_depth=1):
-        if not features or current_depth >= self.max_depth+1:
-            return DecisionNode(level=current_depth, classes=self.classes_)
-
-        split_feature = np.random.choice(features)
-        node = DecisionNode(level=current_depth, classes=self.classes_, split_feature=split_feature)
-
-        if split_feature in self.cat_features_:
-            node.set_split_type(DecisionNode.CAT_SPLIT)
-            for value in feature_domains[str(split_feature)]:
-                child_node = self._build([f for f in features if f != split_feature], feature_domains, current_depth+1)
-                node.add_cat_child(value, child_node)
-        else:
-            node.set_split_type(DecisionNode.CONT_SPLIT)
-            split_value = np.random.uniform(feature_domains[str(split_feature)][0],
-                                            feature_domains[str(split_feature)][1])
-            node.set_split_value(split_value)
-            left_domain = {k: v if k != str(split_feature) else [v[0], split_value]
-                           for k, v in feature_domains.items()}
-            right_domain = {k: v if k != str(split_feature) else [split_value, v[1]]
-                            for k, v in feature_domains.items()}
-            left_child = self._build(features, left_domain, current_depth+1)
-            right_child = self._build(features, right_domain, current_depth+1)
-            node.set_left_child(left_child)
-            node.set_right_child(right_child)
-
-        return node
-
-    def fit(self, X, y, sample_weight=None, check_input=True, X_idx_sorted="deprecated"):
-        if sample_weight is not None:
-            self._warn_unused_args("sample_weight")
+        self._warn_unused_args(unused_args)
 
-        if not isinstance(self.cat_feature_threshold, numbers.Integral) or self.cat_feature_threshold < 0:
-            raise ValueError('Categorical feature threshold should be a positive integer;'
-                             f'got {self.cat_feature_threshold}')
+    def fit(self, X, y, sample_weight=None, check_input=True):
+        """Build a differentially-private decision tree classifier from the training set (X, y).
 
-        if check_input:
-            X, y = self._validate_data(X, y, multi_output=False)
-        self.n_outputs_ = 1
+        Parameters
+        ----------
+        X : array-like of shape (n_samples, n_features)
+            The training input samples. Internally, it will be converted to ``dtype=np.float32``.
 
-        self.feature_domains_ = self.feature_domains
-        self.cat_features_ = self.cat_features
+        y : array-like of shape (n_samples,)
+            The target values (class labels) as integers or strings.
 
-        if self.cat_features_ is None:
-            self.cat_features_ = get_cat_features(X, self.cat_feature_threshold)
+        sample_weight : ignored
+            Ignored by diffprivlib.  Present for consistency with sklearn API.
 
-        if self.feature_domains_ is None:
-            warnings.warn(
-                "feature_domains parameter hasn't been specified, "
-                "so falling back to determining domains from the data.\n"
-                "This may result in additional privacy leakage. To ensure differential privacy with no "
-                "additional privacy loss, specify `feature_domains` according to the documentation",
-                PrivacyLeakWarning)
-            self.feature_domains_ = get_feature_domains(X, self.cat_features_)
+        check_input : bool, default=True
+            Allow to bypass several input checking. Don't use this parameter unless you know what you do.
 
-        self.classes_ = self.classes
+        Returns
+        -------
+        self : DecisionTreeClassifier
+            Fitted estimator.
+        """
+        self.accountant.check(self.epsilon, 0)
 
-        if self.classes_ is None:
-            self.classes_ = np.unique(y)
+        if sample_weight is not None:
+            self._warn_unused_args("sample_weight")
 
-        self.n_classes_ = len(self.classes_)
+        if check_input:
+            X, y = self._validate_data(X, y, multi_output=False)
+        self.n_outputs_ = 1
 
+        if self.bounds is None:
+            warnings.warn("Bounds have not been specified and will be calculated on the data provided. This will "
+                          "result in additional privacy leakage. To ensure differential privacy and no additional "
+                          "privacy leakage, specify bounds for each dimension.", PrivacyLeakWarning)
+            self.bounds = (np.min(X, axis=0), np.max(X, axis=0))
+        self.bounds = self._check_bounds(self.bounds, shape=X.shape[1])
+        X = self._clip_to_bounds(X, self.bounds)
+
+        if self.classes is None:
+            warnings.warn("Classes have not been specified and will be calculated on the data provided. This will "
+                          "result in additional privacy leakage. To ensure differential privacy and no additional "
+                          "privacy leakage, specify the prediction classes for model.", PrivacyLeakWarning)
+            self.classes = np.unique(y)
+        self.classes_ = np.ravel(self.classes)
+        self.n_classes_ = len(self.classes_)
         self.n_features_in_ = X.shape[1]
-        features = list(range(self.n_features_in_))
 
-        self.tree_ = self._build(features, self.feature_domains_)
+        # Build and fit the _FittingTree
+        fitting_tree = _FittingTree(self.max_depth, self.n_features_in_, self.classes_, self.epsilon, self.bounds)
+        fitting_tree.build()
+        fitting_tree.fit(X, y)
 
-        for i, _ in enumerate(X):
-            node = self.tree_.classify(X[i])
-            node.update_class_count(y[i].item())
+        # Load params from _FittingTree into sklearn.Tree
+        d = fitting_tree.__getstate__()
+        tree = Tree(self.n_features_in_, np.array([self.n_classes_]), self.n_outputs_)
+        tree.__setstate__(d)
+        self.tree_ = tree
 
-        self.tree_.set_noisy_label(self.epsilon, self.classes_)
+        self.accountant.spend(self.epsilon, 0)
 
         return self
 
@@ -399,208 +428,169 @@ def _more_tags(self):
         return {}
 
 
-class DecisionNode:
-    """Base Decision Node
-    """
-    CONT_SPLIT = 0
-    CAT_SPLIT = 1
-
-    def __init__(self, level, classes, split_feature=None, split_value=None, split_type=None):
-        """
-        Initialize DecisionNode
-
-        Parameters
-        ----------
-        level: int
-            Node level in the tree
-
-        classes: list
-            List of class labels
-
-        split_feature: int
-            Split feature index
-
-        split_value: Any
-            Feature value to split at
-
-        split_type: int
-            Type of split
-
-        """
-        self._level = level
-        self._classes = classes
-        self._split_type = split_type
-        self._split_feature = split_feature
-        self._split_value = split_value
-        self._left_child = None
-        self._right_child = None
-        self._cat_children = {}
-        self._class_counts = defaultdict(int)
-        self._noisy_label = None
-
-    @property
-    def noisy_label(self):
-        """Get noisy label"""
-        return self._noisy_label
-
-    def set_split_value(self, split_value):
-        """Set split value"""
-        self._split_value = split_value
-
-    def set_split_type(self, split_type):
-        """Set split type"""
-        self._split_type = split_type
+class _FittingTree(DiffprivlibMixin):
+    r"""Array-based representation of a binary decision tree, trained with differential privacy.
 
-    def set_left_child(self, node):
-        """Set left child of the node"""
-        self._left_child = node
+    This tree mimics the architecture of the corresponding Tree from sklearn.tree.tree_, but without many methods given
+    in Tree. The purpose of _FittingTree is to fit the parameters of the model, and have those parameters passed to
+    Tree (using _FittingTree.__getstate__() and Tree.__setstate__()), to be used for prediction.
 
-    def set_right_child(self, node):
-        """Set right child of the node"""
-        self._right_child = node
-
-    def add_cat_child(self, cat_value, node):
-        """Add a categorical child node"""
-        self._cat_children[str(cat_value)] = node
-
-    def is_leaf(self):
-        """Check whether the node is leaf node"""
-        return not self._left_child and not self._right_child and not self._cat_children
-
-    def update_class_count(self, class_value):
-        """Update the class count for the given class"""
-        self._class_counts[class_value] += 1
-
-    def classify(self, x):
-        """Classify the given data"""
-        if self.is_leaf():
-            return self
-
-        child = None
+    Parameters
+    ----------
+    max_depth : int
+        The maximum depth of the tree.
 
-        if self._split_type == self.CAT_SPLIT:
-            x_val = str(x[self._split_feature])
-            child = self._cat_children.get(x_val)
-        else:
-            x_val = x[self._split_feature]
-            if x_val < self._split_value:
-                child = self._left_child
-            else:
-                child = self._right_child
+    n_features : int
+        The number of features of the training dataset.
 
-        if child is None:
-            return self
+    classes : array-like of shape (n_classes,)
+        The classes of the training dataset.
 
-        return child.classify(x)
-
-    def set_noisy_label(self, epsilon, class_values):
-        """Set the noisy label for this node"""
-        if self.is_leaf():
-            if not self._noisy_label:
-                for val in class_values:
-                    if val not in self._class_counts:
-                        self._class_counts[val] = 0
-
-                utility = list(self._class_counts.values())
-                candidates = list(self._class_counts.keys())
-                mech = PermuteAndFlip(epsilon=epsilon, sensitivity=1, monotonic=True, utility=utility,
-                                      candidates=candidates)
-                self._noisy_label = mech.randomise()
-        else:
-            if self._left_child:
-                self._left_child.set_noisy_label(epsilon, class_values)
-            if self._right_child:
-                self._right_child.set_noisy_label(epsilon, class_values)
-            for child_node in self._cat_children.values():
-                child_node.set_noisy_label(epsilon, class_values)
-
-    def predict(self, X):
-        """Predict using this node"""
-        y = []
-        X = np.array(X)
-        check_array(X)
-
-        for x in X:
-            node = self.classify(x)
-            proba = np.zeros(len(self._classes))
-            proba[np.where(self._classes == node.noisy_label)[0].item()] = 1
-            y.append(proba)
-
-        return np.array(y)
-
-
-def get_feature_domains(X, cat_features):
-    """Calculate feature domains from the data.
-
-    Parameters:
-        X : array-like, shape (n_samples, n_features)
-            Training vector, where n_samples is the number of samples and n_features is the number of features.
+    epsilon : float
+        Privacy parameter :math:`\epsilon`.
 
-        cat_features : array of integers
-            List of categorical feature indexes
+    bounds : tuple, optional
+        Bounds of the data, provided as a tuple of the form (min, max).  `min` and `max` can either be scalars, covering
+        the min/max of the entire data.
 
-    Returns:
-        [dict]: Dictionary with keys as feature indexes and values as feature domains.
     """
-    feature_domains = {}
-    X_t = np.transpose(X)
-    cont_features = list(set(range(X.shape[1])) - set(cat_features))
-
-    for i in cat_features:
-        feature_domains[str(i)] = [str(x) for x in set(X_t[i])]
+    _TREE_LEAF = -1
+    _TREE_UNDEFINED = -2
+    StackNode = namedtuple("StackNode", ["parent", "is_left", "depth", "bounds"])
 
-    for i in cont_features:
-        vals = [float(x) for x in X_t[i]]
-        feature_domains[str(i)] = [min(vals), max(vals)]
-
-    return feature_domains
+    def __init__(self, max_depth, n_features, classes, epsilon, bounds):
+        self.node_count = 0
+        self.nodes = []
+        self.max_depth = max_depth
+        self.n_features = n_features
+        self.classes = classes
+        self.epsilon = epsilon
+        self.bounds = bounds
+
+    def __getstate__(self):
+        """Get state of _FittingTree to feed into __setstate__ of sklearn.Tree"""
+        d = {"max_depth": self.max_depth,
+             "node_count": self.node_count,
+             "nodes": np.array([tuple(node) for node in self.nodes], dtype=NODE_DTYPE),
+             "values": self.values_}
+        return d
+
+    def build(self):
+        """Build the decision tree using random feature selection and random thresholding."""
+        stack = [self.StackNode(parent=self._TREE_UNDEFINED, is_left=False, depth=0, bounds=self.bounds)]
+
+        while stack:
+            parent, is_left, depth, bounds = stack.pop()
+            node_id = self.node_count
+            bounds_lower, bounds_upper = self._check_bounds(bounds, shape=self.n_features)
+
+            # Update parent node with its child
+            if parent != self._TREE_UNDEFINED:
+                if is_left:
+                    self.nodes[parent].left_child = node_id
+                else:
+                    self.nodes[parent].right_child = node_id
+
+            # Check if we have a leaf node, then add it
+            if depth >= self.max_depth:
+                node = _Node(node_id, self._TREE_UNDEFINED, self._TREE_UNDEFINED)
+                node.left_child = self._TREE_LEAF
+                node.right_child = self._TREE_LEAF
+
+                self.nodes.append(node)
+                self.node_count += 1
+                continue
+
+            # We have a decision node, so pick feature and threshold
+            feature = np.random.randint(self.n_features)
+            threshold = np.random.uniform(bounds_lower[feature], bounds_upper[feature])
+
+            left_bounds_upper = bounds_upper.copy()
+            left_bounds_upper[feature] = threshold
+            right_bounds_lower = bounds_lower.copy()
+            right_bounds_lower[feature] = threshold
+
+            self.nodes.append(_Node(node_id, feature, threshold))
+            self.node_count += 1
+
+            stack.append(self.StackNode(parent=node_id, is_left=True, depth=depth+1,
+                                        bounds=(bounds_lower, left_bounds_upper)))
+            stack.append(self.StackNode(parent=node_id, is_left=False, depth=depth+1,
+                                        bounds=(right_bounds_lower, bounds_upper)))
 
+        return self
 
-def get_cat_features(X, feature_threshold=2):
-    """Determine categorical features
+    def fit(self, X, y):
+        """Fit the tree to the given training data.
 
-    Parameters:
+        Parameters
+        ----------
         X : array-like, shape (n_samples, n_features)
             Training vector, where n_samples is the number of samples and n_features is the number of features.
 
-        feature_threshold: int, defaults to 2.
-            Threshold value used to determine categorical features. For example, value of ``10`` means
-            any feature that has less than or equal to 10 unique values will be treated as a categorical feature.
+        y : array-like, shape (n_samples,)
+            Target vector relative to X.
 
-    Returns:
-        [list]: List of categorical feature indexes
-    """
-    n_features = X.shape[1]
-    cat_features = []
+        """
+        leaves = self.apply(X)
+        unique_leaves = np.unique(leaves)
+        values = np.zeros(shape=(self.node_count, 1, len(self.classes)))
 
-    for i in range(n_features):
-        values = set(X[:, i])
-        if len(values) <= feature_threshold:
-            cat_features.append(i)
+        # Populate value of real leaves
+        for leaf in unique_leaves:
+            idxs = (leaves == leaf)
+            leaf_y = y[idxs]
 
-    return cat_features
+            counts = [np.sum(leaf_y == cls) for cls in self.classes]
+            mech = PermuteAndFlip(epsilon=self.epsilon, sensitivity=1, monotonic=True, utility=counts)
+            values[leaf, 0, mech.randomise()] = 1
 
+        # Populate value of empty leaves
+        for node in self.nodes:
+            if values[node.node_id].sum() or node.left_child != self._TREE_LEAF:
+                continue
 
-def calc_tree_depth(n_cont_features, n_cat_features, max_depth=15):
-    """Calculate tree depth
+            values[node.node_id, 0, np.random.randint(len(self.classes))] = 1
 
-    Args:
-        n_cont_features (int): Number of continuous features
-        n_cat_features ([type]): Number of categorical features
-        max_depth (int, optional): Max depth tree. Defaults to 15.
+        self.values_ = values
 
-    Returns:
-        [int]: Final depth tree
-    """
-    if n_cont_features < 1:
-        return min(max_depth, np.floor(n_cat_features / 2.))
-    # Designed using balls-in-bins probability. See the paper for details.
-    m = float(n_cont_features)
-    depth = 0
-    expected_empty = m   # the number of unique attributes not selected so far
-    while expected_empty > m / 2.:   # repeat until we have less than half the attributes being empty
-        expected_empty = m * ((m - 1.) / m) ** depth
-        depth += 1
-    # the above was only for half the numerical attributes. now add half the categorical attributes
-    final_depth = np.floor(depth + (n_cat_features / 2.))
-    return min(max_depth, final_depth)
+        return self
+
+    def apply(self, X):
+        """Finds the terminal region (=leaf node) for each sample in X."""
+        n_samples = X.shape[0]
+        out = np.zeros((n_samples,), dtype=int)
+        out_ptr = out.data
+
+        for i in range(n_samples):
+            node = self.nodes[0]
+
+            while node.left_child != self._TREE_LEAF:
+                if X[i, node.feature] <= node.threshold:
+                    node = self.nodes[node.left_child]
+                else:
+                    node = self.nodes[node.right_child]
+
+            out_ptr[i] = node.node_id
+
+        return out
+
+
+class _Node:
+    """Base storage structure for the nodes in a _FittingTree object."""
+    def __init__(self, node_id, feature, threshold):
+        self.feature = feature
+        self.threshold = threshold
+        self.left_child = -1
+        self.right_child = -1
+        self.node_id = node_id
+
+    def __iter__(self):
+        """Defines parameters needed to populate NODE_DTYPE for Tree.__setstate__ using tuple(_Node)."""
+        yield self.left_child
+        yield self.right_child
+        yield self.feature
+        yield self.threshold
+        yield 0.0  # Impurity
+        yield 0  # n_node_samples
+        yield 0.0  # weighted_n_node_samples
diff --git a/docs/conf.py b/docs/conf.py
index 7f049b9..753ad9d 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -85,7 +85,7 @@
 #
 # This is also used if you do content translation via gettext catalogs.
 # Usually you set "language" from the command line for these cases.
-language = None
+language = "en"
 
 # List of patterns, relative to source directory, that match files and
 # directories to ignore when looking for source files.
diff --git a/docs/modules/models.rst b/docs/modules/models.rst
index 4caceed..89063a1 100644
--- a/docs/modules/models.rst
+++ b/docs/modules/models.rst
@@ -18,12 +18,17 @@ Logistic Regression
    :members:
    :inherited-members:
 
-Random Forest
+Tree-Based Models
 +++++++++++++++++++
 .. autoclass:: RandomForestClassifier
    :members:
    :inherited-members:
-   :exclude-members: apply,decision_path,feature_importances_
+   :exclude-members: feature_importances_,n_features_
+
+.. autoclass:: DecisionTreeClassifier
+   :members:
+   :inherited-members:
+   :exclude-members: cost_complexity_pruning_path,feature_importances_,n_features_
 
 Regression models
 -----------------


Test patch:
diff --git a/tests/models/test_DecisionTreeClassifier.py b/tests/models/test_DecisionTreeClassifier.py
index 9154af4..5faaadf 100644
--- a/tests/models/test_DecisionTreeClassifier.py
+++ b/tests/models/test_DecisionTreeClassifier.py
@@ -1,10 +1,21 @@
 import numpy as np
 from unittest import TestCase
+
+from scipy.sparse import issparse
+from sklearn import datasets
+from sklearn.tree._tree import Tree
 from sklearn.utils.validation import check_is_fitted
 from sklearn.exceptions import NotFittedError
 
-from diffprivlib.models.forest import DecisionTreeClassifier, get_cat_features, get_feature_domains, calc_tree_depth
-from diffprivlib.utils import PrivacyLeakWarning, global_seed, DiffprivlibCompatibilityWarning
+from diffprivlib.models.forest import DecisionTreeClassifier
+from diffprivlib.utils import PrivacyLeakWarning, global_seed, DiffprivlibCompatibilityWarning, BudgetError
+
+
+X = [[-2, -1], [-1, -1], [-1, -2], [1, 1], [1, 2], [2, 1]]
+y = [-1, -1, -1, 1, 1, 1]
+T = [[-1, -1], [2, 2], [3, 2]]
+true_result = [-1, 1, 1]
+iris = datasets.load_iris()
 
 
 class TestDecisionTreeClassifier(TestCase):
@@ -17,40 +28,39 @@ def test_not_none(self):
     def test_bad_params(self):
         X = [[1]]
         y = [0]
-        
+
         with self.assertRaises(ValueError):
-            DecisionTreeClassifier(cat_feature_threshold="5").fit(X, y)
+            DecisionTreeClassifier(epsilon=-1, bounds=([0], [1])).fit(X, y)
 
         with self.assertWarns(DiffprivlibCompatibilityWarning):
-            DecisionTreeClassifier(cat_feature_threshold=2, feature_domains={'0': [0, 1]}).fit(X, y, sample_weight=1)
+            DecisionTreeClassifier(bounds=([0], [1])).fit(X, y, sample_weight=1)
 
     def test_bad_data(self):
         with self.assertRaises(ValueError):
-            DecisionTreeClassifier(feature_domains={'0': [0, 1]}).fit([[1]], None)
+            DecisionTreeClassifier(bounds=([0], [1])).fit([[1]], None)
 
         with self.assertRaises(ValueError):
-            DecisionTreeClassifier(feature_domains={'0': [0, 2]}).fit([[1], [2]], [1])
+            DecisionTreeClassifier(bounds=([0], [1])).fit([[1], [2]], [1])
 
         with self.assertRaises(ValueError):
-            DecisionTreeClassifier(feature_domains={'0': [0, 2]}).fit([[1], [2]], [[1, 2], [2, 4]])
+            DecisionTreeClassifier(bounds=([0], [1])).fit([[1], [2]], [[1, 2], [2, 4]])
 
     def test_simple(self):
         X = np.array([[12, 3, 14], [12, 3, 4], [12, 3, 4], [2, 13, 4], [2, 13, 14], [2, 3, 14], [3, 5, 15]] * 3)
         y = np.array([1, 1, 1, 0, 0, 0, 1] * 3)
-        model = DecisionTreeClassifier(epsilon=5, cat_feature_threshold=2, max_depth=5, random_state=25)
+        model = DecisionTreeClassifier(epsilon=5, max_depth=3)
         with self.assertRaises(NotFittedError):
             check_is_fitted(model)
-        # when `feature_domains` is not provided, we should get a privacy leakage warning
+        # when `bounds` is not provided, we should get a privacy leakage warning
         with self.assertWarns(PrivacyLeakWarning):
             model.fit(X, y)
         check_is_fitted(model)
         self.assertTrue(model.predict(np.array([[12, 3, 14]])))
 
-    def test_with_feature_domains(self):
+    def test_with_bounds(self):
         X = np.array([[12, 3, 14], [12, 3, 4], [12, 3, 4], [2, 13, 4], [2, 13, 14], [2, 3, 14], [3, 5, 15]] * 3)
         y = np.array([1, 1, 1, 0, 0, 0, 1] * 3)
-        model = DecisionTreeClassifier(epsilon=5, cat_feature_threshold=2, max_depth=5, random_state=25,
-                                       feature_domains={'0': [2.0, 12.0], '1': [3.0, 13.0], '2': [4.0, 15.0]})
+        model = DecisionTreeClassifier(epsilon=5, bounds=([2, 3, 4], [12, 13, 14]), classes=[0, 1], max_depth=3)
         with self.assertRaises(NotFittedError):
             check_is_fitted(model)
         model.fit(X, y)
@@ -60,35 +70,62 @@ def test_with_feature_domains(self):
     def test_with_non_binary_labels(self):
         X = np.array([[12, 3, 14], [12, 3, 4], [12, 3, 4], [2, 13, 4], [2, 13, 14], [2, 3, 14], [3, 5, 15]] * 3)
         y = np.array([3, 3, 3, 3, 5, 5, 3] * 3)
-        model = DecisionTreeClassifier(epsilon=5, cat_feature_threshold=2, max_depth=5, random_state=25,
-                                       feature_domains={'0': [2.0, 12.0], '1': [3.0, 13.0], '2': [4.0, 15.0]})
+        model = DecisionTreeClassifier(epsilon=5, bounds=([2, 3, 4], [12, 13, 14]), classes=[3, 5], max_depth=3)
         with self.assertRaises(NotFittedError):
             check_is_fitted(model)
         model.fit(X, y)
         check_is_fitted(model)
         self.assertEqual(model.predict(np.array([[12, 3, 14]])), 3)
 
+    def test_sklearn_methods(self):
+        depth = 3
+        clf = DecisionTreeClassifier(epsilon=5, bounds=(-2, 2), classes=[-1, 1], max_depth=depth)
+        clf.fit(X, y)
+
+        self.assertEqual(clf.n_features_in_, 2)
+
+        self.assertEqual(clf.get_depth(), depth)
+
+        self.assertEqual(clf.get_n_leaves(), 2 ** depth)
+
+        self.assertIsInstance(clf.apply(T), np.ndarray)
+        self.assertEqual(clf.apply(T).shape, (len(T),))
+
+        self.assertIsInstance(clf.tree_, Tree)
+
+        self.assertIsInstance(clf.score(T, true_result), float)
+        self.assertTrue(0 <= clf.score(T, true_result) <= 1)
+
+        self.assertTrue(issparse(clf.decision_path(X)))
+
+        self.assertIsNone(clf._prune_tree())
+
+        self.assertIsInstance(clf.predict(T), np.ndarray)
+        self.assertEqual(clf.predict(T).shape, (len(T),))
+
+        self.assertIsInstance(clf.predict_proba(T), np.ndarray)
+        self.assertEqual(clf.predict_proba(T).shape, (len(T), clf.n_classes_))
+
+        self.assertEqual(clf.ccp_alpha, 0.0)
+        self.assertIsNone(clf.class_weight)
+
+        self.assertIsNotNone(clf.n_features_)
+
+    def test_accountant(self):
+        from diffprivlib.accountant import BudgetAccountant
+        acc = BudgetAccountant()
+
+        X = np.array([[12, 3, 14], [12, 3, 4], [12, 3, 4], [2, 13, 4], [2, 13, 14], [2, 3, 14], [3, 5, 15]])
+        y = np.array([1, 1, 1, 0, 0, 0, 1])
+        bounds = ([2, 3, 4], [12, 13, 15])
+        model = DecisionTreeClassifier(epsilon=2, bounds=bounds, classes=[0, 1], max_depth=3, accountant=acc)
+        model.fit(X, y)
+        self.assertEqual((2, 0), acc.total())
+
+        with BudgetAccountant(3, 0) as acc2:
+            model = DecisionTreeClassifier(epsilon=2, bounds=bounds, classes=[0, 1], max_depth=3)
+            model.fit(X, y)
+            self.assertEqual((2, 0), acc2.total())
 
-class TestUtils(TestCase):
-    def test_calc_tree_depth(self):
-        self.assertEqual(calc_tree_depth(0, 4), 2)
-        self.assertEqual(calc_tree_depth(0, 100, max_depth=20), 20)
-        self.assertEqual(calc_tree_depth(4, 5), 6)
-        self.assertEqual(calc_tree_depth(4, 5, max_depth=3), 3)
-        self.assertEqual(calc_tree_depth(40, 50), 15)
-
-    def test_get_feature_domains(self):
-        X = np.array([[12, 3, 14, 21], [0.1, 0.5, 0.7, 1], ['cat', 'dog', 'mouse', 'cat'], [0, 1, 0, 1]]).T
-        cat_features = [2, 3]
-        feature_domains = get_feature_domains(X, cat_features)
-        self.assertEqual(feature_domains['0'], [3, 21])
-        self.assertEqual(feature_domains['1'], [0.1, 1])
-        self.assertEqual(set(feature_domains['2']), set(['dog', 'cat', 'mouse']))
-        self.assertEqual(set(feature_domains['3']), set(['0', '1']))
-    
-    def test_get_cat_features(self):
-        X = np.array([[12, 3, 14, 21], [0.1, 0.5, 0.7, 1], ['cat', 'dog', 'mouse', 'cat'], [0, 1, 0, 1]]).T
-        cat_features = get_cat_features(X)
-        self.assertEqual(cat_features, [3])
-        cat_features = get_cat_features(X, feature_threshold=3)
-        self.assertEqual(cat_features, [2, 3])
\ No newline at end of file
+            with self.assertRaises(BudgetError):
+                model.fit(X, y)
diff --git a/tests/models/test_FittingTree.py b/tests/models/test_FittingTree.py
new file mode 100644
index 0000000..77fcb45
--- /dev/null
+++ b/tests/models/test_FittingTree.py
@@ -0,0 +1,21 @@
+from unittest import TestCase
+
+import numpy as np
+from sklearn.tree._tree import Tree
+
+from diffprivlib.models.forest import _FittingTree
+
+
+class TestFittingTree(TestCase):
+    def test_not_none(self):
+        self.assertIsNotNone(_FittingTree)
+
+    def test_no_data(self):
+        classes = np.array([0, 1, 2])
+        tree = _FittingTree(5, 3, classes, 1, (0, 1))
+        tree.build()
+        tree.values_ = np.zeros(shape=(tree.node_count, 1, len(tree.classes)))
+        d = tree.__getstate__()
+
+        sktree = Tree(3, np.array([len(classes)]), 1)
+        sktree.__setstate__(d)
diff --git a/tests/models/test_RandomForestClassifier.py b/tests/models/test_RandomForestClassifier.py
index 5d6e801..faaad36 100644
--- a/tests/models/test_RandomForestClassifier.py
+++ b/tests/models/test_RandomForestClassifier.py
@@ -1,16 +1,23 @@
 import numpy as np
 from unittest import TestCase
+
+from scipy.sparse import issparse
 from sklearn.utils.validation import check_is_fitted
 from sklearn.exceptions import NotFittedError
 
-from diffprivlib.models.forest import RandomForestClassifier
+from diffprivlib.models.forest import RandomForestClassifier, DecisionTreeClassifier
 from diffprivlib.utils import PrivacyLeakWarning, global_seed, BudgetError, DiffprivlibCompatibilityWarning
 
+X = [[-2, -1], [-1, -1], [-1, -2], [1, 1], [1, 2], [2, 1]] * 3
+y = [-1, -1, -1, 1, 1, 1] * 3
+T = [[-1, -1], [2, 2], [3, 2]]
+true_result = [-1, 1, 1]
+
 
 class TestRandomForestClassifier(TestCase):
     def setUp(self):
         global_seed(2718281828)
-    
+
     def test_not_none(self):
         self.assertIsNotNone(RandomForestClassifier)
 
@@ -18,27 +25,24 @@ def test_bad_params(self):
         X = [[1]]
         y = [0]
 
-        with self.assertRaises(ValueError):
-            RandomForestClassifier(n_estimators="10").fit(X, y)
-        
-        with self.assertRaises(ValueError):
-            RandomForestClassifier(cat_feature_threshold="5").fit(X, y)
+        with self.assertRaises((ValueError, TypeError)):  # Should be TypeError
+            RandomForestClassifier(n_estimators="10", bounds=(0, 1), classes=[0, 1]).fit(X, y)
 
         with self.assertWarns(DiffprivlibCompatibilityWarning):
-            RandomForestClassifier(n_estimators=1, feature_domains={'0': [0, 1]}).fit(X, y, sample_weight=1)
+            RandomForestClassifier(n_estimators=1, bounds=(0, 1), classes=[0, 1]).fit(X, y, sample_weight=1)
 
     def test_bad_data(self):
         with self.assertRaises(ValueError):
-            RandomForestClassifier(feature_domains={'0': [0, 1]}).fit([[1]], None)
+            RandomForestClassifier(bounds=([0], [1])).fit([[1]], None)
 
         with self.assertRaises(ValueError):
-            RandomForestClassifier(feature_domains={'0': [0, 2]}).fit([[1], [2]], [1])
+            RandomForestClassifier(bounds=([0], [1])).fit([[1], [2]], [1])
 
         with self.assertRaises(ValueError):
-            RandomForestClassifier(feature_domains={'0': [0, 2]}).fit([[1], [2]], [[1, 2], [2, 4]])
+            RandomForestClassifier(bounds=([0], [1])).fit([[1], [2]], [[1, 2], [2, 4]])
 
-        with self.assertRaises(ValueError):
-            RandomForestClassifier(feature_domains={'0': [0, 2], '1': [0, 6]}).fit([[1, 2], [3, 4]], [1, 0])
+        # with self.assertRaises(ValueError):
+        #     RandomForestClassifier(bounds=([0], [1])).fit([[1, 2], [3, 4]], [1, 0])
 
     def test_multi_output(self):
         X = np.array([[12, 3, 14], [12, 3, 4], [12, 3, 4], [2, 13, 4], [2, 13, 14], [2, 3, 14], [3, 5, 15]] * 3)
@@ -51,7 +55,7 @@ def test_multi_output(self):
     def test_simple(self):
         X = np.array([[12, 3, 14], [12, 3, 4], [12, 3, 4], [2, 13, 4], [2, 13, 14], [2, 3, 14], [3, 5, 15]] * 3)
         y = np.array([1, 1, 1, 0, 0, 0, 1] * 3)
-        model = RandomForestClassifier(epsilon=5, n_estimators=5, random_state=2021, cat_feature_threshold=2)
+        model = RandomForestClassifier(epsilon=5, n_estimators=5, max_depth=3)
 
         with self.assertRaises(NotFittedError):
             check_is_fitted(model)
@@ -64,56 +68,156 @@ def test_simple(self):
         self.assertEqual(model.n_features_in_, 3)
         self.assertEqual(model.n_classes_, 2)
         self.assertEqual(set(model.classes_), set([0, 1]))
-        self.assertEqual(model.cat_features_, [])
-        self.assertEqual(model.max_depth_, 3)
+        self.assertEqual(model.max_depth, 3)
         self.assertTrue(model.estimators_)
         self.assertEqual(len(model.estimators_), 5)
         self.assertTrue(model.predict(np.array([[12, 3, 14]])))
-        self.assertIsNone(model.feature_domains)
-        self.assertEqual(model.feature_domains_, {'0': [2.0, 12.0], '1': [3.0, 13.0], '2': [4.0, 15.0]})
+        self.assertTrue(np.all(model.bounds[0] == [2, 3, 4]))
+        self.assertTrue(np.all(model.bounds[1] == [12, 13, 15]))
 
-    def test_with_feature_domains(self):
+    def test_with_bounds(self):
         X = np.array([[12, 3, 14], [12, 3, 4], [12, 3, 4], [2, 13, 4], [2, 13, 14], [2, 3, 14], [3, 5, 15]] * 3)
         y = np.array([1, 1, 1, 0, 0, 0, 1] * 3)
-        model = RandomForestClassifier(epsilon=5, n_estimators=5, random_state=2021, cat_feature_threshold=2,
-                                       feature_domains={'0': [2.0, 12.0], '1': [3.0, 13.0], '2': [4.0, 15.0]})
+        model = RandomForestClassifier(epsilon=5, n_estimators=5, bounds=([2, 3, 4], [12, 13, 15]), classes=[0, 1],
+                                       max_depth=3)
         with self.assertRaises(NotFittedError):
             check_is_fitted(model)
+
         model.fit(X, y)
+
         check_is_fitted(model)
         self.assertEqual(model.n_features_in_, 3)
         self.assertEqual(model.n_classes_, 2)
         self.assertEqual(set(model.classes_), set([0, 1]))
-        self.assertEqual(model.cat_features_, [])
-        self.assertEqual(model.max_depth_, 3)
+        self.assertEqual(model.max_depth, 3)
         self.assertTrue(model.estimators_)
         self.assertEqual(len(model.estimators_), 5)
         self.assertTrue(model.predict(np.array([[12, 3, 14]])))
-        self.assertEqual(model.feature_domains_, {'0': [2.0, 12.0], '1': [3.0, 13.0], '2': [4.0, 15.0]})
+        self.assertTrue(np.all(model.bounds[0] == [2, 3, 4]))
+        self.assertTrue(np.all(model.bounds[1] == [12, 13, 15]))
+
+    def test_non_binary_classes(self):
+        X = np.array([[12, 3, 14], [12, 3, 4], [12, 3, 4], [2, 13, 4], [2, 13, 14], [2, 3, 14], [3, 5, 15]] * 3)
+        y = np.array([1, 1, 1, 0, 0, 0, 1] * 3) + 2
+        model = RandomForestClassifier(epsilon=5, n_estimators=5, bounds=([2, 3, 4], [12, 13, 15]), classes=[2, 3],
+                                       max_depth=3)
+        with self.assertRaises(NotFittedError):
+            check_is_fitted(model)
+
+        model.fit(X, y)
+
+        self.assertIn(model.predict(np.array([[12, 3, 14]])), (2, 3))
+        self.assertIn(model.predict(np.array([[2, 13, 4]])), (2, 3))
+        self.assertIn(model.predict(np.array([[3, 5, 15]])), (2, 3))
 
-    def test_with_not_enough_feature_domains(self):
+    def test_with_not_enough_bounds(self):
         X = np.array([[12, 3, 14], [12, 3, 4], [12, 3, 4], [2, 13, 4], [2, 13, 14], [2, 3, 14], [3, 5, 15]])
         y = np.array([1, 1, 1, 0, 0, 0, 1])
-        model = RandomForestClassifier(epsilon=2, n_estimators=5, random_state=2021, feature_domains={'0': [2.0, 12.0], '1': [3.0, 13.0]})
+        model = RandomForestClassifier(epsilon=2, n_estimators=5, bounds=([2, 3], [12, 13]))
         with self.assertRaises(NotFittedError):
             check_is_fitted(model)
         with self.assertRaises(ValueError):
             model.fit(X, y)
 
+    def test_sklearn_methods(self):
+        depth = 3
+        n_estimators = 4
+        clf = RandomForestClassifier(n_estimators, epsilon=5, bounds=(-2, 2), classes=np.array((-1, 1)),
+                                     max_depth=depth)
+        clf.fit(X, y)
+
+        self.assertEqual(clf.n_features_in_, 2)
+
+        self.assertIsInstance(clf.apply(T), np.ndarray)
+        self.assertEqual(clf.apply(T).shape, (len(T), n_estimators))
+
+        for estimator in clf:
+            self.assertIsInstance(estimator, DecisionTreeClassifier)
+
+        self.assertEqual(len(clf), n_estimators)
+        self.assertIsInstance([e for e in clf], list)
+
+        self.assertIsInstance(clf.score(T, true_result), float)
+        self.assertTrue(0 <= clf.score(T, true_result) <= 1)
+
+        self.assertIsInstance(clf.decision_path(X), tuple)
+        self.assertTrue(issparse(clf.decision_path(X)[0]), msg=str(clf.decision_path(X)))
+
+        self.assertIsInstance(clf.predict(T), np.ndarray)
+        self.assertEqual(clf.predict(T).shape, (len(T),))
+
+        self.assertIsInstance(clf.predict_proba(T), np.ndarray)
+        self.assertEqual(clf.predict_proba(T).shape, (len(T), clf.n_classes_))
+
+        self.assertEqual(clf.ccp_alpha, 0.0)
+        self.assertIsNone(clf.class_weight)
+
+    def test_warm_start(self):
+        depth = 3
+        n_estimators = 4
+        clf = RandomForestClassifier(n_estimators, epsilon=1, bounds=(-2, 2), classes=np.array((-1, 1)),
+                                     max_depth=depth, warm_start=True)
+        clf.fit(X, y)
+        first_estimator = clf[0].__dict__.copy()
+
+        clf.n_estimators = n_estimators * 2
+        clf.fit(X, y)
+
+        self.assertEqual(len(clf), n_estimators * 2)
+        self.assertEqual(clf[0].__dict__, first_estimator)
+        self.assertNotEqual(clf[n_estimators].__dict__, first_estimator)
+
+        with self.assertWarns(UserWarning):
+            clf.fit(X, y)
+
+        clf.n_estimators = n_estimators
+        with self.assertRaises(ValueError):
+            clf.fit(X, y)
+
+    def test_parallel(self):
+        depth = 3
+        n_estimators = 4
+        clf = RandomForestClassifier(n_estimators, epsilon=5, bounds=(-2, 2), classes=np.array((-1, 1)),
+                                     max_depth=depth, n_jobs=-1)
+        clf.fit(X, y)
+
+        self.assertIsNone(check_is_fitted(clf))
+
+    def test_shuffle(self):
+        depth = 3
+        n_estimators = 4
+        clf = RandomForestClassifier(n_estimators, epsilon=5, bounds=(-2, 2), classes=np.array((-1, 1)),
+                                     max_depth=depth, shuffle=True)
+        clf.fit(X, y)
+
+        self.assertIsNone(check_is_fitted(clf))
+
+    def test_more_estimators_than_data(self):
+        depth = 3
+        n_estimators = 20
+        clf = RandomForestClassifier(n_estimators, epsilon=5, bounds=(-2, 2), classes=np.array((-1, 1)),
+                                     max_depth=depth)
+        clf.fit(X, y)
+
+        self.assertIsNone(check_is_fitted(clf))
+
+        self.assertIsInstance(clf.predict(T), np.ndarray)
+        self.assertEqual(clf.predict(T).shape, (len(T),))
+        self.assertEqual(clf.predict_proba(T).shape, (len(T), clf.n_classes_))
+
     def test_accountant(self):
         from diffprivlib.accountant import BudgetAccountant
         acc = BudgetAccountant()
 
         X = np.array([[12, 3, 14], [12, 3, 4], [12, 3, 4], [2, 13, 4], [2, 13, 14], [2, 3, 14], [3, 5, 15]])
         y = np.array([1, 1, 1, 0, 0, 0, 1])
-        model = RandomForestClassifier(epsilon=2, n_estimators=5, accountant=acc,
-                                       feature_domains={'0': [2.0, 12.0], '1': [3.0, 13.0], '2': [3.0, 16.0]})
+        bounds = ([2, 3, 4], [12, 13, 15])
+        model = RandomForestClassifier(epsilon=2, n_estimators=5, accountant=acc, bounds=bounds, classes=[0, 1])
         model.fit(X, y)
         self.assertEqual((2, 0), acc.total())
-        
+
         with BudgetAccountant(3, 0) as acc2:
-            model = RandomForestClassifier(epsilon=2, n_estimators=5,
-                                           feature_domains={'0': [2.0, 12.0], '1': [3.0, 13.0], '2': [3.0, 16.0]})
+            model = RandomForestClassifier(epsilon=2, n_estimators=5, bounds=bounds, classes=[0, 1])
             model.fit(X, y)
             self.assertEqual((2, 0), acc2.total())