Skip to content

Commit

Permalink
[SPARK-7329] [MLLIB] simplify ParamGridBuilder impl
Browse files Browse the repository at this point in the history
as suggested by justinuang on #5601.

Author: Xiangrui Meng <meng@databricks.com>

Closes #5873 from mengxr/SPARK-7329 and squashes the following commits:

d08f9cf [Xiangrui Meng] simplify tests
b7a7b9b [Xiangrui Meng] simplify grid build
  • Loading branch information
mengxr committed May 4, 2015
1 parent 9e25b09 commit 1ffa8cb
Showing 1 changed file with 9 additions and 19 deletions.
28 changes: 9 additions & 19 deletions python/pyspark/ml/tuning.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
# limitations under the License.
#

import itertools

__all__ = ['ParamGridBuilder']


Expand All @@ -37,14 +39,10 @@ class ParamGridBuilder(object):
{lr.regParam: 1.0, lr.featuresCol: 'f', lr.maxIter: 5, lr.labelCol: 'l', lr.predictionCol: 'p'}, \
{lr.regParam: 2.0, lr.featuresCol: 'f', lr.maxIter: 5, lr.labelCol: 'l', lr.predictionCol: 'p'}, \
{lr.regParam: 3.0, lr.featuresCol: 'f', lr.maxIter: 5, lr.labelCol: 'l', lr.predictionCol: 'p'}]
>>> fail_count = 0
>>> for e in expected:
... if e not in output:
... fail_count += 1
>>> if len(expected) != len(output):
... fail_count += 1
>>> fail_count
0
>>> len(output) == len(expected)
True
>>> all([m in expected for m in output])
True
"""

def __init__(self):
Expand Down Expand Up @@ -76,17 +74,9 @@ def build(self):
Builds and returns all combinations of parameters specified
by the param grid.
"""
param_maps = [{}]
for (param, values) in self._param_grid.items():
new_param_maps = []
for value in values:
for old_map in param_maps:
copied_map = old_map.copy()
copied_map[param] = value
new_param_maps.append(copied_map)
param_maps = new_param_maps

return param_maps
keys = self._param_grid.keys()
grid_values = self._param_grid.values()
return [dict(zip(keys, prod)) for prod in itertools.product(*grid_values)]


if __name__ == "__main__":
Expand Down

0 comments on commit 1ffa8cb

Please sign in to comment.