From 2ba7c42f1c109873a463c394df7615a2c1a51128 Mon Sep 17 00:00:00 2001 From: Mahdi Ben Jelloul Date: Fri, 11 Aug 2023 17:11:29 +0200 Subject: [PATCH] Fix filtering --- openfisca_survey_manager/simulations.py | 30 +++++++++++++++++++------ 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/openfisca_survey_manager/simulations.py b/openfisca_survey_manager/simulations.py index 7da4143c..191f16e9 100644 --- a/openfisca_survey_manager/simulations.py +++ b/openfisca_survey_manager/simulations.py @@ -164,9 +164,6 @@ def compute_quantiles(simulation = None, variable = None, nquantiles = None, per if weight_variable_by_entity: weight_variable = weight_variable_by_entity[entity_key] - if filter_by is not None: - NotImplementedError("filter_by option not implemented") - variable_values = simulation.adaptative_calculate_variable(variable, period) if weighted: assert (alternative_weights is not None) or (weight_variable is not None) @@ -178,6 +175,18 @@ def compute_quantiles(simulation = None, variable = None, nquantiles = None, per else: weight = np.ones(len(variable_values)) + if filtering_variable_by_entity is not None: + if filter_by is None: + filter_by = filtering_variable_by_entity.get(entity_key) + + if filter_by is not None: + filter_entity_key = simulation.tax_benefit_system.variables.get(filter_by).entity.key + assert filter_entity_key == entity_key + filter_dummy = simulation.adaptative_calculate_variable(filter_by, period = period) + + variable_values = variable_values[filter_dummy].copy() + weight = weight[filter_dummy].copy() + labels = np.arange(1, nquantiles + 1) method = 2 quantile, values = mark_weighted_percentiles(variable_values, labels, weight, method, return_quantiles = True) @@ -524,14 +533,21 @@ def compute_winners_loosers( ): entity_key = baseline_simulation.tax_benefit_system.variables[variable].entity.key + after = simulation.adaptative_calculate_variable(variable, period = period) before = baseline_simulation.adaptative_calculate_variable(variable, period = period) - if filter_by: - raise NotImplementedError("filter by not implemented for this function") + if filtering_variable_by_entity is not None: + if filter_by is None: + filter_by = filtering_variable_by_entity.get(entity_key) + + if filter_by is not None: + filter_entity_key = baseline_simulation.tax_benefit_system.variables.get(filter_by).entity.key + assert filter_entity_key == entity_key + filter_dummy = baseline_simulation.adaptative_calculate_variable(filter_by, period = period) - if filtering_variable_by_entity: - raise NotImplementedError("filtering_variable_by_entity not implemented for this function") + after = after[filter_dummy].copy() + before = before[filter_dummy].copy() weight = np.ones(len(after)) if weighted: