Repository: earth-chris/elapid

Issue:
add tau scaling
Maxent uses a `tau` parameter to scale relative suitability predictions into detection probability scores. This is not currently implemented, as the translation from the R [maxnet predict function](https://github.com/mrmaxent/maxnet/blob/master/R/predict.maxnet.R) didn't include this parameter, effectively fixing `tau = 0.5`. 

We should look at the underlying math and apply tau scaling in the elapid.MaxentModel.predict() function ([link here](https://github.com/earth-chris/elapid/blob/e975cc64666b266b5cafdb8b524fc165383b6a84/elapid/models.py#L116)).

You can read more about `tau` in the [maxent for ecologists paper](https://onlinelibrary.wiley.com/doi/full/10.1111/j.1472-4642.2010.00725.x) and the [practical guide to maxent paper](https://onlinelibrary.wiley.com/doi/full/10.1111/j.1600-0587.2013.07872.x)

Discussion:

Patch:
diff --git a/elapid/config.py b/elapid/config.py
new file mode 100644
index 0000000..f4743a5
--- /dev/null
+++ b/elapid/config.py
@@ -0,0 +1,36 @@
+"""SDM model configuration parameters."""
+
+
+class MaxentConfig:
+    # constrain feature min/max ranges
+    clamp: bool = True
+
+    # beta regularization drops feature cofficients if they don't improve model
+    # performance. these scaler modify feature beta regularization behavior.
+    # increase this number to drop more coefficients, decrease to keep more
+    # coefficients. `beta_multiplier` scales regularization for all features,
+    # but regularization can be tuned by feature type
+    beta_multiplier: float = 1.0
+    beta_hinge: float = 1.0
+    beta_lqp: float = 1.0
+    beta_threshold: float = 1.0
+    beta_categorical: float = 1.0
+
+    # default maxent feature types to compute
+    feature_types: list = ["linear", "hinge", "product"]
+
+    # the number of hinge and threshold features to compute per-covariate
+    n_hinge_features: int = 30
+    n_threshold_features: int = 20
+
+    # model training metric (from `sklearn.metrics`)
+    scorer: str = "roc_auc"
+
+    # species prevalence scalar
+    tau: float = 0.5
+
+    # model convergence tolerance threshold
+    tolerance: float = 1e-7
+
+    # elasticnet lambda type to use ('best' or 'last')
+    use_lambdas: str = "best"
diff --git a/elapid/features.py b/elapid/features.py
index 26692f6..97fdb80 100644
--- a/elapid/features.py
+++ b/elapid/features.py
@@ -6,37 +6,37 @@ import numpy as np
 import pandas as pd
 from sklearn.preprocessing import OneHotEncoder
 
-from elapid.utils import MAXENT_DEFAULTS, ArrayLike, repeat_array
+from elapid.config import MaxentConfig
+from elapid.types import ArrayLike, validate_boolean, validate_feature_types, validate_numeric_scalar
+from elapid.utils import repeat_array
 
 
 class MaxentFeatureTransformer(object):
-    """Transforms covariate data into maxent-format feature data.
-
-    Attributes:
-        feature_types_: List of the feature types (quadratic, etc.) to compute
-        clamp_: Boolean to clamp feature to the min/max of the fit data
-        n_hinge_features_: Int of the linear hinge features to fit
-        n_threshold_features_: Int of the number of thresholds to fit
-        initialized_: Boolean flag identifying the transformer's fit status
-        hinge_ranges_: List of the hinge ranges fit
-        threshold_ranges_: List of the thresholds fit
-        categorical_encoders_: List of `sklearn` one-hot class encoders
-        categorical_: Array-like of column indices for which columns are categorical
-        feature_mins_: Array-like of the feature minimum values for clamping
-        feature_maxs_: Array-like of the feature maximum values for clamping
-        labels_: List of covariate labels
-    """
+    """Transforms covariate data into maxent-format feature data."""
+
+    feature_types_: list = MaxentConfig.feature_types
+    clamp_: bool = MaxentConfig.clamp
+    n_hinge_features_: int = MaxentConfig.n_hinge_features
+    n_threshold_features_: int = MaxentConfig.n_threshold_features
+    initialized_: bool = False
+    hinge_ranges_: dict = {}
+    threshold_ranges_: dict = {}
+    categorical_encoders_: dict = {}
+    categorical_: ArrayLike = None
+    feature_mins_: ArrayLike = None
+    feature_maxs_: ArrayLike = None
+    labels_: list = None
 
     def __init__(
         self,
-        feature_types: Union[str, list] = MAXENT_DEFAULTS["feature_types"],
-        clamp: bool = MAXENT_DEFAULTS["clamp"],
-        n_hinge_features: int = MAXENT_DEFAULTS["n_hinge_features"],
-        n_threshold_features: int = MAXENT_DEFAULTS["n_threshold_features"],
+        feature_types: Union[str, list] = MaxentConfig.feature_types,
+        clamp: bool = MaxentConfig.clamp,
+        n_hinge_features: int = MaxentConfig.n_hinge_features,
+        n_threshold_features: int = MaxentConfig.n_threshold_features,
     ):
         """Computes features based on the maxent feature types specified (like linear, quadratic, hinge).
 
-        Implemented using sklearn conventions (usng `.fit()` and `.transform()` functions.
+        Implemented using sklearn conventions (with `.fit()` and `.transform()` functions.
 
         Args:
             feature_types: list of maxent features to generate.
@@ -44,22 +44,11 @@ class MaxentFeatureTransformer(object):
             n_hinge_features: number of hinge knots to generate
             n_threshold_features: nuber of threshold features to generate
         """
-        # user-specified parameters
         self.feature_types_ = validate_feature_types(feature_types)
         self.clamp_ = validate_boolean(clamp)
         self.n_hinge_features_ = validate_numeric_scalar(n_hinge_features)
         self.n_threshold_features_ = validate_numeric_scalar(n_threshold_features)
 
-        # data-driven parameters
-        self.initialized_ = False
-        self.hinge_ranges_ = dict()
-        self.threshold_ranges_ = dict()
-        self.categorical_encoders_ = dict()
-        self.categorical_ = None
-        self.feature_mins_ = None
-        self.feature_maxs_ = None
-        self.labels_ = None
-
     def fit(self, x: ArrayLike, categorical: list = None, labels: list = None) -> None:
         """Fits features to covariates.
 
@@ -320,9 +309,7 @@ def hingeval(x: ArrayLike, mn: float, mx: float) -> ArrayLike:
     return np.minimum(1, np.maximum(0, (x - mn) / (mx - mn)))
 
 
-def hinge(
-    x: ArrayLike, n_hinges: int = MAXENT_DEFAULTS["n_hinge_features"], range: Tuple[float, float] = None
-) -> ArrayLike:
+def hinge(x: ArrayLike, n_hinges: int = MaxentConfig.n_hinge_features, range: Tuple[float, float] = None) -> ArrayLike:
     """Fits hinge transformations to an array of covariates.
 
     Args:
@@ -347,9 +334,7 @@ def hinge(
     return np.concatenate((lh, rh), axis=1)
 
 
-def threshold(
-    x, n_thresholds: int = MAXENT_DEFAULTS["n_threshold_features"], range: Tuple[float, float] = None
-) -> ArrayLike:
+def threshold(x, n_thresholds: int = MaxentConfig.n_threshold_features, range: Tuple[float, float] = None) -> ArrayLike:
     """Fits arbitrary threshold transformations to an array of covariates.
 
     Args:
@@ -423,11 +408,11 @@ def compute_weights(y: ArrayLike, pbr: int = 100) -> ArrayLike:
 def compute_regularization(
     f: pd.DataFrame,
     y: ArrayLike,
-    beta_multiplier: float = MAXENT_DEFAULTS["beta_multiplier"],
-    beta_lqp: float = MAXENT_DEFAULTS["beta_lqp"],
-    beta_threshold: float = MAXENT_DEFAULTS["beta_threshold"],
-    beta_hinge: float = ["beta_hinge"],
-    beta_categorical: float = MAXENT_DEFAULTS["beta_hinge"],
+    beta_multiplier: float = MaxentConfig.beta_multiplier,
+    beta_lqp: float = MaxentConfig.beta_lqp,
+    beta_threshold: float = MaxentConfig.beta_threshold,
+    beta_hinge: float = MaxentConfig.beta_hinge,
+    beta_categorical: float = MaxentConfig.beta_hinge,
 ) -> ArrayLike:
     """Computes variable regularization values for all feature data.
 
@@ -505,84 +490,3 @@ def compute_regularization(
     max_reg *= beta_multiplier
 
     return max_reg
-
-
-def validate_feature_types(features: Union[str, list]) -> list:
-    """Ensures the feature classes passed are maxent-legible
-
-    Args:
-        features: List or string that must be in ["linear", "quadratic", "product",
-            "hinge", "threshold", "auto"] or string "lqphta"
-
-    Returns:
-        valid_features: List of formatted valid feature values
-    """
-    valid_list = ["linear", "quadratic", "product", "hinge", "threshold"]
-    valid_string = "lqpht"
-    valid_features = list()
-
-    assert features is not None, "Features cannot be empty"
-
-    # ensure the string features are valid, and translate to a standard feature list
-    if type(features) is str:
-        for feature in features:
-
-            if feature == "a":
-                return valid_list
-
-            assert feature in valid_string, "Invalid feature passed: {}".format(feature)
-
-            if feature == "l":
-                valid_features.append("linear")
-            elif feature == "q":
-                valid_features.append("quadratic")
-            elif feature == "p":
-                valid_features.append("product")
-            elif feature == "h":
-                valid_features.append("hinge")
-            elif feature == "t":
-                valid_features.append("threshold")
-
-    # or ensure the list features are valid
-    elif type(features) is list:
-        for feature in features:
-            if feature == "auto":
-                return valid_list
-
-            assert feature in valid_list, "Invalid feature passed: {}".format(feature)
-
-            valid_features.append(feature)
-
-    return valid_features
-
-
-def validate_boolean(var: Any) -> bool:
-    """Evaluates whether an argument is boolean True/False
-
-    Args:
-        var: the input argument to validate
-
-    Returns:
-        var: the value if it passes validation
-
-    Raises:
-        AssertionError: `var` was not boolean
-    """
-    assert isinstance(var, bool), "Argument must be True/False"
-    return var
-
-
-def validate_numeric_scalar(var: Any) -> bool:
-    """Evaluates whether an argument is a single numeric value.
-
-    Args:
-        var: the input argument to validate
-
-    Returns:
-        var: the value if it passes validation
-
-    Raises:
-        AssertionError: `var` was not numeric.
-    """
-    assert isinstance(var, (int, float)), "Argument must be single numeric value"
-    return var
diff --git a/elapid/geo.py b/elapid/geo.py
index 8afac77..d8d7e44 100644
--- a/elapid/geo.py
+++ b/elapid/geo.py
@@ -12,6 +12,7 @@ import rasterio as rio
 from shapely.geometry import MultiPoint, Point
 from sklearn.base import BaseEstimator
 
+from elapid.types import CRSType
 from elapid.utils import (
     NoDataException,
     check_raster_alignment,
@@ -27,8 +28,6 @@ tqdm = get_tqdm()
 tqdm_opts = {"desc": "Geometry", "leave": False}
 tqdm.pandas(**tqdm_opts)
 
-CRSType = Union[pyproj.CRS, str]
-
 
 def xy_to_geoseries(
     x: Union[float, list, np.ndarray], y: Union[float, list, np.ndarray], crs: CRSType = "epsg:4236"
@@ -323,7 +322,7 @@ def raster_values_from_geoseries(
             if drop_na and src.nodata is not None:
                 values.replace(src.nodata, np.NaN, inplace=True)
 
-            raster_values.append(values.astype(src.profile["dtype"]))
+        raster_values.append(values)
 
     # format the values as a geodataframe
     df = pd.concat(raster_values, axis=1)
diff --git a/elapid/models.py b/elapid/models.py
index 18e6b82..ed6142c 100644
--- a/elapid/models.py
+++ b/elapid/models.py
@@ -1,57 +1,56 @@
 """Classes for training species distribution models."""
-
 from typing import List, Union
 
 import numpy as np
 import pandas as pd
-from glmnet.logistic import LogitNet
+from glmnet.linear import ElasticNet
 from sklearn.base import BaseEstimator
 
 from elapid import features as _features
-from elapid.utils import MAXENT_DEFAULTS, ArrayLike, Number, n_cpus
+from elapid.config import MaxentConfig
+from elapid.types import ArrayLike, Number, validate_feature_types
+from elapid.utils import n_cpus
 
 
 class MaxentModel(BaseEstimator):
-    """Creates a model estimator for Maxent-style species distribution models.
-
-    Attributes:
-        feature_types_: List of maxent feature types to compute
-        tau_: Float species prevalence scalar
-        clamp_: Boolean for whether to constrain feature min/max ranges
-        convergence_tolerance_: Float model convergence tolerance threshold
-        beta_multiplier_: Float scalar for all feature beta values
-        beta_hinge_: Float scalar for hinge beta values
-        beta_lqp_: Float scalar for linear/quadratic/product beta values
-        beta_threshold_: Float scalar for threshold beta values
-        beta_categorical_: Float scalar for categorical beta values
-        n_hinge_features_: Int number of hinge features to fit
-        n_threshold_features_: Int number of threshold features to fit
-        n_cpus_: Int of cpu threads ot use during model training
-        scorer_: Str of model training metric (from `sklearn.metrics`)
-        use_lambdas_: Str of lambda type to use (in `['best', 'last']`)
-        initialized_: Bool tracking model fit status
-        beta_scores_: Array of trained model betas
-        entropy_: Float of trained model entropy score
-        alpha_: Float of trained model alpha score 0.0
-        estimator: Object of trained model estimator
-        transformer: MaxentFeatureTransformer object
-    """
+    """Creates a model estimator for Maxent-style species distribution models."""
+
+    feature_types_: list = MaxentConfig.feature_types
+    tau_: float = MaxentConfig.tau
+    clamp_: bool = MaxentConfig.clamp
+    convergence_tolerance_: float = MaxentConfig.tolerance
+    beta_multiplier_: float = MaxentConfig.beta_multiplier
+    beta_hinge_: float = MaxentConfig.beta_hinge
+    beta_lqp_: float = MaxentConfig.beta_lqp
+    beta_threshold_: float = MaxentConfig.beta_threshold
+    beta_categorical_: float = MaxentConfig.beta_categorical
+    n_hinge_features_: int = MaxentConfig.n_hinge_features
+    n_threshold_features_: int = MaxentConfig.n_threshold_features
+    scorer_: str = MaxentConfig.scorer
+    use_lambdas_: str = MaxentConfig.use_lambdas
+    initialized_: bool = False
+    beta_scores_: np.array = None
+    entropy_: float = 0.0
+    alpha_: float = 0.0
+    estimator: BaseEstimator = None
+    transformer: _features.MaxentFeatureTransformer = None
+    n_cpus_ = n_cpus
 
     def __init__(
         self,
-        feature_types: Union[list, str] = MAXENT_DEFAULTS["feature_types"],
-        tau: float = MAXENT_DEFAULTS["tau"],
-        clamp: bool = MAXENT_DEFAULTS["clamp"],
-        scorer: str = MAXENT_DEFAULTS["scorer"],
-        beta_multiplier: float = MAXENT_DEFAULTS["beta_multiplier"],
-        beta_lqp: float = MAXENT_DEFAULTS["beta_lqp"],
-        beta_hinge: float = MAXENT_DEFAULTS["beta_hinge"],
-        beta_threshold: float = MAXENT_DEFAULTS["beta_lqp"],
-        beta_categorical: float = MAXENT_DEFAULTS["beta_categorical"],
-        n_hinge_features: int = MAXENT_DEFAULTS["n_hinge_features"],
-        n_threshold_features: int = MAXENT_DEFAULTS["n_threshold_features"],
-        convergence_tolerance: float = MAXENT_DEFAULTS["tolerance"],
-        use_lambdas: str = MAXENT_DEFAULTS["use_lambdas"],
+        feature_types: Union[list, str] = MaxentConfig.feature_types,
+        tau: float = MaxentConfig.tau,
+        clamp: bool = MaxentConfig.clamp,
+        scorer: str = MaxentConfig.scorer,
+        beta_multiplier: float = MaxentConfig.beta_multiplier,
+        beta_lqp: float = MaxentConfig.beta_lqp,
+        beta_hinge: float = MaxentConfig.beta_hinge,
+        beta_threshold: float = MaxentConfig.beta_lqp,
+        beta_categorical: float = MaxentConfig.beta_categorical,
+        n_hinge_features: int = MaxentConfig.n_hinge_features,
+        n_threshold_features: int = MaxentConfig.n_threshold_features,
+        convergence_tolerance: float = MaxentConfig.tolerance,
+        use_lambdas: str = MaxentConfig.use_lambdas,
         n_cpus: int = n_cpus,
     ):
         """Create a maxent model object.
@@ -62,16 +61,17 @@ class MaxentModel(BaseEstimator):
             tau: maxent prevalence value for scaling logistic output
             clamp: set features to min/max range from training during prediction
             scorer: sklearn scoring function for model training
-            beta_multiplier: scaler for all regularization parameters. higher values exclude more features
+            beta_multiplier: scaler for all regularization parameters.
+                higher values drop more coeffiecients
             beta_lqp: linear, quadratic and product feature regularization scaler
             beta_hinge: hinge feature regularization scaler
             beta_threshold: threshold feature regularization scaler
             beta_categorical: categorical feature regularization scaler
             convergence_tolerance: model convergence tolerance level
-            use_lambdas: guide for which model lambdas to select, from ["best", "last"]
+            use_lambdas: guide for which model lambdas to select (either "best" or "last")
             n_cpus: threads to use during model training
         """
-        self.feature_types_ = _features.validate_feature_types(feature_types)
+        self.feature_types_ = validate_feature_types(feature_types)
         self.tau_ = tau
         self.clamp_ = clamp
         self.convergence_tolerance_ = convergence_tolerance
@@ -86,13 +86,6 @@ class MaxentModel(BaseEstimator):
         self.scorer_ = scorer
         self.use_lambdas_ = use_lambdas
 
-        self.initialized_ = False
-        self.beta_scores_ = None
-        self.entropy_ = 0.0
-        self.alpha_ = 0.0
-        self.estimator = None
-        self.transformer = None
-
     def fit(
         self, x: ArrayLike, y: ArrayLike, categorical: List[int] = None, labels: list = None, is_features: bool = False
     ) -> None:
@@ -134,15 +127,19 @@ class MaxentModel(BaseEstimator):
         )
 
         if self.use_lambdas_ == "last":
-            self.beta_scores_ = self.estimator.coef_path_[0, :, -1]
+            self.beta_scores_ = self.estimator.coef_path_[:, -1]
         elif self.use_lambdas_ == "best":
-            self.beta_scores_ = self.estimator.coef_path_[0, :, self.estimator.lambda_best_inx_]
+            self.beta_scores_ = self.estimator.coef_path_[:, self.estimator.lambda_max_inx_]
+
+        # maxent-specific transformations
+        raw = self.predict(features[y == 0], transform="raw", is_features=True)
 
-        # maxent specific transformations
-        rr = self.predict(features[y == 0], transform="exponential", is_features=True)
-        raw = rr / np.sum(rr)
-        self.entropy_ = -np.sum(raw * np.log(raw))
-        self.alpha_ = -np.log(np.sum(rr))
+        # alpha is a normalizing constant that ensures that f1(z) integrates (sums) to 1
+        self.alpha_ = -np.log(np.sum(raw))
+
+        # the distance from f(z) is considered the relative entropy of f1(z) WRT f(z)
+        scaled = raw / np.sum(raw)
+        self.entropy_ = -np.sum(scaled * np.log(scaled))
 
     def predict(self, x: ArrayLike, transform: str = "logistic", is_features: bool = False) -> ArrayLike:
         """Applies a model to a set of covariates or features. Requires that a model has been fit.
@@ -165,16 +162,24 @@ class MaxentModel(BaseEstimator):
             features = self.transformer.transform(x)
 
         # apply the model
-        link = np.matmul(features, self.beta_scores_) + self.alpha_
+        engma = np.matmul(features, self.beta_scores_) + self.alpha_
 
+        # scale based on the transform type
         if transform == "raw":
-            return link
-        elif transform == "exponential":
-            return np.exp(link)
+            return np.exp(engma)
+
         elif transform == "logistic":
-            return 1 / (1 + np.exp(-self.entropy_ - link))
+            # below is R's maxnet (tau-free) logistic formulation
+            # return 1 / (1 + np.exp(-self.entropy_ - raw))
+            # use the java formulation instead
+            logratio = np.exp(engma) * np.exp(self.entropy_)
+            return (self.tau_ * logratio) / ((1 - self.tau_) + (self.tau_ * logratio))
+
         elif transform == "cloglog":
-            return 1 - np.exp(0 - np.exp(self.entropy_ + link))
+            # below is R's maxent cloglog formula
+            # return 1 - np.exp(0 - np.exp(self.entropy_ - raw))
+            # use java again
+            return 1 - np.exp(-np.exp(engma) * np.exp(self.entropy_))
 
     def fit_predict(
         self,
@@ -222,7 +227,7 @@ class MaxentModel(BaseEstimator):
         Returns:
             None. updates the self.estimator with an sklearn-style model estimator
         """
-        self.estimator = LogitNet(
+        self.estimator = ElasticNet(
             alpha=alpha,
             lambda_path=lambdas,
             standardize=standardize,
diff --git a/elapid/types.py b/elapid/types.py
new file mode 100644
index 0000000..41113a2
--- /dev/null
+++ b/elapid/types.py
@@ -0,0 +1,92 @@
+"""Custom maxent and typing data types."""
+from typing import Any, Union
+
+import numpy as np
+import pandas as pd
+import pyproj
+
+# typing
+Number = Union[int, float]
+ArrayLike = Union[np.array, pd.DataFrame]
+CRSType = Union[pyproj.CRS, str]
+
+
+# maxent feature transformations
+def get_feature_types(return_string: bool = False) -> Union[list, str]:
+    feature_types = "lqpht" if return_string else ["linear", "quadratic", "product", "hinge", "threshold"]
+    return feature_types
+
+
+def validate_feature_types(features: Union[str, list]) -> list:
+    """Ensures the feature classes passed are maxent-legible
+
+    Args:
+        features: List or string that must be in ["linear", "quadratic", "product",
+            "hinge", "threshold", "auto"] or string "lqphta"
+
+    Returns:
+        valid_features: List of formatted valid feature values
+    """
+    valid_list = get_feature_types(return_string=False)
+    valid_string = get_feature_types(return_string=True)
+    valid_features = list()
+
+    # ensure the string features are valid, and translate to a standard feature list
+    if type(features) is str:
+        for feature in features:
+            if feature == "a":
+                return valid_list
+            assert feature in valid_string, "Invalid feature passed: {}".format(feature)
+            if feature == "l":
+                valid_features.append("linear")
+            elif feature == "q":
+                valid_features.append("quadratic")
+            elif feature == "p":
+                valid_features.append("product")
+            elif feature == "h":
+                valid_features.append("hinge")
+            elif feature == "t":
+                valid_features.append("threshold")
+
+    # or ensure the list features are valid
+    elif type(features) is list:
+        for feature in features:
+            if feature == "auto":
+                return valid_list
+            assert feature in valid_list, "Invalid feature passed: {}".format(feature)
+            valid_features.append(feature)
+
+    return valid_features
+
+
+# data type checking
+def validate_boolean(var: Any) -> bool:
+    """Evaluates whether an argument is boolean True/False
+
+    Args:
+        var: the input argument to validate
+
+    Returns:
+        var: the value if it passes validation
+
+    Raises:
+        AssertionError: `var` was not boolean
+    """
+    assert isinstance(var, bool), "Argument must be True/False"
+    return var
+
+
+def validate_numeric_scalar(var: Any) -> bool:
+    """Evaluates whether an argument is a single numeric value.
+
+    Args:
+        var: the input argument to validate
+
+    Returns:
+        var: the value if it passes validation
+
+    Raises:
+        AssertionError: `var` was not numeric.
+    """
+    assert isinstance(var, (int, float)), "Argument must be single numeric value"
+    return var
diff --git a/elapid/utils.py b/elapid/utils.py
index 5515d93..afdc5c3 100644
--- a/elapid/utils.py
+++ b/elapid/utils.py
@@ -11,26 +11,9 @@ import numpy as np
 import pandas as pd
 import rasterio as rio
 
-n_cpus = mp.cpu_count()
+from elapid.types import Number
 
-MAXENT_DEFAULTS = {
-    "clamp": True,
-    "beta_multiplier": 1.0,
-    "beta_hinge": 1.0,
-    "beta_lqp": 1.0,
-    "beta_threshold": 1.0,
-    "beta_categorical": 1.0,
-    "feature_types": ["linear", "hinge", "product"],
-    "n_hinge_features": 30,
-    "n_threshold_features": 20,
-    "scorer": "roc_auc",
-    "tau": 0.5,
-    "tolerance": 1e-7,
-    "use_lambdas": "last",
-}
-
-Number = Union[int, float]
-ArrayLike = Union[np.array, pd.DataFrame]
+n_cpus = mp.cpu_count()
 
 
 class NoDataException(Exception):


Test patch:
diff --git a/.github/workflows/run-pytest.yml b/.github/workflows/run-pytest.yml
index 2edf826..a132023 100644
--- a/.github/workflows/run-pytest.yml
+++ b/.github/workflows/run-pytest.yml
@@ -11,17 +11,18 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.6", "3.7", "3.8"]
+        python-version: ["3.7", "3.8"]
     steps:
     - uses: actions/checkout@v2
     - name: Set up Python ${{ matrix.python-version }}
       uses: actions/setup-python@v2
       with:
         python-version: ${{ matrix.python-version }}
+        cache: 'pip'
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        python -m pip install pytest pytest-cov
+        pip install pytest pytest-cov -r requirements.txt
         pip install .
     - name: Run pytest
       run: |
diff --git a/tests/test_features.py b/tests/test_features.py
new file mode 100644
index 0000000..fd0bfe6
--- /dev/null
+++ b/tests/test_features.py
@@ -0,0 +1,8 @@
+from elapid import features
+from elapid.utils import load_sample_data
+
+x, y = load_sample_data()
+
+
+def test_MaxentFeatureTransformer_types():
+    pass
diff --git a/tests/test_models.py b/tests/test_models.py
new file mode 100644
index 0000000..5c8d2f2
--- /dev/null
+++ b/tests/test_models.py
@@ -0,0 +1,76 @@
+from sklearn import metrics
+
+from elapid import models
+from elapid.utils import load_sample_data
+
+x, y = load_sample_data()
+
+
+def test_MaxentModel_flow():
+    model = models.MaxentModel()
+    model.fit(x, y)
+    ypred = model.predict(x)
+    assert len(ypred) == len(y)
+    print(ypred)
+    print(ypred.max())
+    assert ypred.max() <= 1.0
+    assert ypred.min() >= 0.0
+
+
+def test_MaxentModel_performance():
+    model = models.MaxentModel(use_lambdas="last", scorer="roc_auc")
+    model.fit(x, y)
+    ypred = model.predict(x)
+    auc_score = metrics.roc_auc_score(y, ypred)
+
+    # check that model is not nonsense
+    assert 0.5 <= auc_score <= 1.0
+
+    # check that full model close to cross-val performance (+- stderr)
+    assert abs(auc_score - model.estimator.cv_mean_score_[-1]) < 0.1 + model.estimator.cv_standard_error_[-1]
+
+
+def test_MaxentModel_best_lambdas():
+    model = models.MaxentModel(use_lambdas="best")
+    model.fit(x, y)
+    ypred = model.predict(x, transform="logistic")
+    auc_score = metrics.roc_auc_score(y, ypred)
+    assert 0.5 <= auc_score <= 1.0
+    assert 0.48 < ypred[y == 1].mean() < 0.52
+
+
+def test_tau_scaler():
+    model = models.MaxentModel(tau=0.5)
+    model.fit(x, y)
+    ypred = model.predict(x, transform="logistic")
+    assert 0.48 < ypred[y == 1].mean() < 0.52
+
+    model = models.MaxentModel(tau=0.25)
+    model.fit(x, y)
+    ypred = model.predict(x, transform="logistic")
+    assert 0.23 < ypred[y == 1].mean() < 0.27
+
+
+def test_MaxentModel_feature_types():
+    all = models.MaxentModel(feature_types="lqpht", beta_multiplier=1.5)
+    auto = models.MaxentModel(feature_types=["auto"])
+    linear = models.MaxentModel(feature_types=["linear"])
+    qp = models.MaxentModel(feature_types=["quadratic", "product"], convergence_tolerance=2e-7, beta_lqp=1.5)
+    ht = models.MaxentModel(
+        feature_types=["hinge", "threshold"],
+        beta_threshold=1.5,
+        beta_hinge=0.75,
+        n_threshold_features=10,
+        n_hinge_features=10,
+    )
+
+    all.fit(x, y)
+    auto.fit(x, y)
+    linear.fit(x, y)
+    qp.fit(x, y)
+    ht.fit(x, y)
+
+    for model in [all, linear, qp, ht, auto]:
+        ypred = model.predict(x)
+        auc_score = metrics.roc_auc_score(y, ypred)
+        assert 0.5 <= auc_score <= 1.0