From bb208e1d93e16ac68fb3c68fa2ec0ec81890da41 Mon Sep 17 00:00:00 2001 From: nanglo123 Date: Mon, 30 Sep 2024 11:05:38 -0400 Subject: [PATCH] Propogate param_renaming_uses_strata_names argument from stratify endpoint to stratify method --- mira/dkg/model.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/mira/dkg/model.py b/mira/dkg/model.py index 5af4eb6b..bdc9c20b 100644 --- a/mira/dkg/model.py +++ b/mira/dkg/model.py @@ -285,6 +285,13 @@ class StratificationQuery(BaseModel): "will stratify all concepts.", examples=[["recovered"]], ) + param_renaming_uses_strata_names: Optional[bool] = Field( + False, + description="If true, the strata names will be used in the parameter " + "renaming. If false, the strata indices will be used. " + "Default: False", + examples=[False] + ) def get_conversion_cls(self) -> Type[Template]: if self.conversion_cls == "natural_conversion": @@ -339,6 +346,7 @@ def model_stratification( params_to_preserve=stratification_query.params_to_preserve, concepts_to_stratify=stratification_query.concepts_to_stratify, concepts_to_preserve=stratification_query.concepts_to_preserve, + param_renaming_uses_strata_names=stratification_query.param_renaming_uses_strata_names ) return template_model