diff --git a/tests/test_regression.py b/tests/test_regression.py index 1220395..cd6ea23 100644 --- a/tests/test_regression.py +++ b/tests/test_regression.py @@ -64,3 +64,24 @@ def test_simulation(orca_session): assert orca.get_table('obs').to_frame()['a_predicted'].equals(m.predicted_values) + +def test_out_transform(orca_session): + """ + Test transformation of the predicted values. + + """ + modelmanager.initialize() + + m = OLSRegressionStep() + m.tables = 'obs' + m.model_expression = 'a ~ b' + m.fit() + + m.out_column = 'a_predicted' + m.out_transform = 'np.exp' + m.run() + + predictions = m.predicted_values.apply(np.exp) + + assert orca.get_table('obs').to_frame()['a_predicted'].equals(predictions) + diff --git a/urbansim_templates/models/regression.py b/urbansim_templates/models/regression.py index 7683498..9d6b843 100644 --- a/urbansim_templates/models/regression.py +++ b/urbansim_templates/models/regression.py @@ -1,5 +1,6 @@ from __future__ import print_function +import math import numpy as np import pandas as pd from datetime import datetime as dt @@ -69,10 +70,12 @@ class OLSRegressionStep(TemplateStep): side variable from the model expression will be used. Replaces the `out_fname` argument in UrbanSim. - out_transform : callable, optional - Transformation to apply to the predicted values, for example to reverse a - transformation of the left-hand-side variable in the model expression. Replaces - the `ytransform` argument in UrbanSim. + out_transform : str, optional + Element-wise transformation to apply to the predicted values, for example to + reverse a transformation of the left-hand-side variable in the model expression. + This should be provided as a string containing a function name. Supports anything + from NumPy or Python's built-in math library, for example 'np.exp' or + 'math.floor'. Replaces the `ytransform` argument in UrbanSim. out_filters : str or list of str, optional Filters to apply to the data before simulation. If not provided, no filters will @@ -168,7 +171,7 @@ def fit(self): """ self.model = RegressionModel(model_expression=self.model_expression, fit_filters=self.filters, predict_filters=self.out_filters, - ytransform=self.out_transform, name=self.name) + ytransform=None, name=self.name) df = get_data(tables = self.tables, filters = self.filters, @@ -207,7 +210,7 @@ def run(self): self.predicted_values = values if self.out_transform is not None: - values = self.out_transform(values) + values = values.apply(eval(self.out_transform)) colname = self._get_out_column() tabname = self._get_out_table()