From 73f2f78745da6c623aa680534b008c7329a94f4b Mon Sep 17 00:00:00 2001 From: Mateo Date: Wed, 25 May 2022 13:30:47 -0400 Subject: [PATCH] Replace list(map()) with list comprehensions --- HARK/distribution.py | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/HARK/distribution.py b/HARK/distribution.py index c747de119..ca53d2c6c 100644 --- a/HARK/distribution.py +++ b/HARK/distribution.py @@ -1305,11 +1305,10 @@ def calc_expectation(dstn, func=lambda x: x, *args): Scalar if only one value. """ - f_query = list( - map( - lambda x: func(x, *args), [dstn.X[..., i] for i in range(len(dstn.pmf))] - ) - ) + f_query = [ + func(dstn.X[...,i], *args) for i in range(len(dstn.pmf)) + ] + f_query = np.stack(f_query, axis=-1) f_exp = np.dot(f_query, np.vstack(dstn.pmf)) @@ -1349,11 +1348,10 @@ def distr_of_function(dstn, func=lambda x: x, *args): The distribution of func(dstn). """ # Apply function to every event realization - f_query = list( - map( - lambda x: func(x, *args), [dstn.X[..., i] for i in range(len(dstn.pmf))] - ) - ) + f_query = [ + func(dstn.X[...,i], *args) for i in range(len(dstn.pmf)) + ] + # Stack results along their last (new) axis f_query = np.stack(f_query, axis=-1)