Curve Fit: Picklable model

Author: VesnaT

Orange/data/util.py

@@ -268,3 +268,19 @@ def get_unique_names_domain(attributes, class_vars=(), metas=()):
                                  for old, new in zip(all_names, unique_names)
                                  if new != old))
     return (attributes, class_vars, metas), renamed
+
+
+def sanitized_name(name: str) -> str:
+    """
+    Replace non-alphanumeric characters and leading zero with `_`.
+
+    Args:
+        name (str): proposed name
+
+    Returns:
+        name (str): new name
+    """
+    sanitized = re.sub(r"\W", "_", name)
+    if sanitized[0].isdigit():
+        sanitized = "_" + sanitized
+    return sanitized

Orange/regression/curvefit.py

@@ -1,24 +1,40 @@
-from typing import Iterable, Callable, List, Optional
+import ast
+from typing import Callable, List, Optional, Union, Dict, Tuple, Any
 
 import numpy as np
 from scipy.optimize import curve_fit
 
 from Orange.data import Table, Domain, ContinuousVariable, StringVariable
 from Orange.data.filter import HasClass
+from Orange.data.util import sanitized_name, get_unique_names
 from Orange.preprocess import RemoveNaNColumns, Impute
 from Orange.regression import Learner, Model
 
 __all__ = ["CurveFitLearner"]
 
 
 class CurveFitModel(Model):
-    def __init__(self, domain: Domain, original_domain: Domain,
-                 parameters_names: List[str],
-                 parameters: np.ndarray, function: Callable):
+    def __init__(
+            self,
+            domain: Domain,
+            original_domain: Domain,
+            parameters_names: List[str],
+            parameters: np.ndarray,
+            function: Optional[Callable],
+            create_lambda_args: Optional[Tuple]
+    ):
         super().__init__(domain, original_domain)
         self.__parameters_names = parameters_names
         self.__parameters = parameters
+
+        if function is None and create_lambda_args is not None:
+            function, names, _ = _create_lambda(**create_lambda_args)
+            assert parameters_names == names
+
+        assert function
+
         self.__function = function
+        self.__create_lambda_args = create_lambda_args
 
     @property
     def coefficients(self) -> Table:
@@ -34,41 +50,219 @@ def predict(self, X: np.ndarray) -> np.ndarray:
             return np.full(len(X), predicted)
         return predicted.flatten()
 
+    def __getstate__(self) -> Dict:
+        if not self.__create_lambda_args:
+            raise AttributeError(
+                "Can't pickle/copy callable. Use str expression instead."
+            )
+        return {
+            "domain": self.domain,
+            "original_domain": self.original_domain,
+            "parameters_names": self.__parameters_names,
+            "parameters": self.__parameters,
+            "function": None,
+            "args": self.__create_lambda_args,
+        }
+
+    def __setstate__(self, state: Dict):
+        self.__init__(*state.values())
+
 
 class CurveFitLearner(Learner):
     preprocessors = [HasClass(), RemoveNaNColumns(), Impute()]
     __returns__ = CurveFitModel
     name = "Curve Fit"
 
-    def __init__(self, function: Callable, parameters_names: List[str],
-                 feature_names: List[str], p0: Optional[Iterable] = None,
-                 bounds: Iterable = (-np.inf, np.inf), preprocessors=None):
+    def __init__(
+            self,
+            expression: Union[Callable, ast.Expression, str],
+            parameters_names: Optional[List[str]] = None,
+            features_names: Optional[List[str]] = None,
+            available_feature_names: Optional[List[str]] = None,
+            functions: Optional[List[str]] = None,
+            globals_: Optional[Dict[str, Any]] = None,
+            p0: Union[List, Dict, None] = None,
+            bounds: Union[Tuple, Dict] = (-np.inf, np.inf),
+            preprocessors=None
+    ):
         super().__init__(preprocessors)
 
-        if not callable(function):
-            raise TypeError("Function is not callable.")
+        if callable(expression):
+            if parameters_names is None:
+                raise TypeError("Provide 'parameters_names' parameter.")
+            if features_names is None:
+                raise TypeError("Provide 'features_names' parameter.")
+
+            args = None
+            function = expression
+        else:
+            if available_feature_names is None:
+                raise TypeError("Provide 'available_feature_names' parameter.")
+            if functions is None:
+                raise TypeError("Provide 'functions' parameter.")
+
+            args = dict(expression=expression,
+                        available_feature_names=available_feature_names,
+                        functions=functions, globals_=globals_)
+            function, parameters_names, features_names = _create_lambda(**args)
+
+        if isinstance(p0, dict):
+            p0 = [p0.get(p, 1) for p in parameters_names]
+        if isinstance(bounds, dict):
+            d = [-np.inf, np.inf]
+            lower_bounds = [bounds.get(p, d)[0] for p in parameters_names]
+            upper_bounds = [bounds.get(p, d)[1] for p in parameters_names]
+            bounds = lower_bounds, upper_bounds
 
         self.__function = function
         self.__parameters_names = parameters_names
-        self.__feature_names = feature_names
+        self.__features_names = features_names
         self.__p0 = p0
         self.__bounds = bounds
 
+        # learner is not picklable, if expression is lambda
+        # properties, needed for pickling
+        self.__create_lambda_args = args
+
+    @property
+    def parameters_names(self) -> List[str]:
+        return self.__parameters_names
+
     def fit_storage(self, data: Table) -> CurveFitModel:
-        domain = data.domain
+        domain: Domain = data.domain
         attributes = []
         for attr in domain.attributes:
-            if attr.name in self.__feature_names:
+            if attr.name in self.__features_names:
                 if not attr.is_continuous:
                     raise ValueError("Numeric feature expected.")
                 attributes.append(attr)
 
         new_domain = Domain(attributes, domain.class_vars, domain.metas)
         transformed = data.transform(new_domain)
-        params, _ = curve_fit(self.__function, transformed.X, transformed.Y,
-                              p0=self.__p0, bounds=self.__bounds)
+        params = curve_fit(self.__function, transformed.X, transformed.Y,
+                           p0=self.__p0, bounds=self.__bounds)[0]
         return CurveFitModel(new_domain, domain,
-                             self.__parameters_names, params, self.__function)
+                             self.__parameters_names, params, self.__function,
+                             self.__create_lambda_args)
+
+    def __getstate__(self) -> Dict:
+        if not self.__create_lambda_args:
+            raise AttributeError(
+                "Can't pickle/copy callable. Use str expression instead."
+            )
+        state = self.__create_lambda_args.copy()
+        state["parameters_names"] = None
+        state["features_names"] = None
+        state["p0"] = self.__p0
+        state["bounds"] = self.__bounds
+        state["preprocessors"] = self.preprocessors
+        return state
+
+    def __setstate__(self, state: Dict):
+        expression = state.pop("expression")
+        self.__init__(expression, **state)
+
+
+def _create_lambda(
+        expression: Union[str, ast.Expression] = "",
+        available_feature_names: List[str] = None,
+        functions: List[str] = None,
+        globals_: Optional[Dict[str, Any]] = None
+) -> Tuple[Callable, List[str], List[str]]:
+    if globals_ is None:
+        globals_ = {name: getattr(np, name) for name in functions}
+
+    sanitized_vars_names = [sanitized_name(name) for name
+                            in available_feature_names]
+
+    exp = ast.parse(expression, mode="eval")
+    search = _ParametersSearch(sanitized_vars_names, functions)
+    search.visit(exp)
+    params = search.parameters
+    used_sanitized_vars_names = search.variables
+
+    name = sanitized_name(get_unique_names(params, "x"))
+    vars_mapper = {var: i for i, var in enumerate(used_sanitized_vars_names)}
+    exp = _ReplaceVars(name, vars_mapper, functions).visit(exp)
+
+    lambda_ = ast.Lambda(
+        args=ast.arguments(
+            posonlyargs=[],
+            args=[ast.arg(arg=arg) for arg in [name] + params],
+            varargs=None,
+            kwonlyargs=[],
+            kw_defaults=[],
+            defaults=[],
+        ),
+        body=exp.body
+    )
+    exp = ast.Expression(body=lambda_)
+    ast.fix_missing_locations(exp)
+    vars_ = [name for name in available_feature_names
+             if sanitized_name(name) in used_sanitized_vars_names]
+
+    # pylint: disable=eval-used
+    return eval(compile(exp, "<lambda>", mode="eval"), globals_), params, vars_
+
+
+class _ParametersSearch(ast.NodeVisitor):
+    """
+    Find features and parameters:
+     - feature: if node is instance of ast.Name and is included in vars_names
+     - parameters: if node is instance of ast.Name and is not included
+     in functions
+    """
+
+    def __init__(self, vars_names: List[str], functions: List[str]):
+        super().__init__()
+        self.__vars_names = vars_names
+        self.__functions = functions
+        self.__parameters: List[str] = []
+        self.__variables: List[str] = []
+
+    @property
+    def parameters(self) -> List[str]:
+        return self.__parameters
+
+    @property
+    def variables(self) -> List[str]:
+        return self.__variables
+
+    def visit_Name(self, node: ast.Name) -> ast.Name:
+        if node.id in self.__vars_names:
+            # don't use Set in order to preserve parameters order
+            if node.id not in self.__variables:
+                self.__variables.append(node.id)
+        elif node.id not in self.__functions:
+            # don't use Set in order to preserve parameters order
+            if node.id not in self.__parameters:
+                self.__parameters.append(node.id)
+        return node
+
+
+class _ReplaceVars(ast.NodeTransformer):
+    """
+    Replace feature names with X[:, i], where i is index of feature.
+    """
+
+    def __init__(self, name: str, vars_mapper: Dict, functions: List):
+        super().__init__()
+        self.__name = name
+        self.__vars_mapper = vars_mapper
+        self.__functions = functions
+
+    def visit_Name(self, node: ast.Name) -> Union[ast.Name, ast.Subscript]:
+        if node.id not in self.__vars_mapper or node.id in self.__functions:
+            return node
+        else:
+            n = self.__vars_mapper[node.id]
+            return ast.Subscript(
+                value=ast.Name(id=self.__name, ctx=ast.Load()),
+                slice=ast.ExtSlice(
+                    dims=[ast.Slice(lower=None, upper=None, step=None),
+                          ast.Index(value=ast.Num(n=n))]),
+                ctx=node.ctx
+            )
 
 
 if __name__ == "__main__":