Qumo/JPMC-quant
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

saving plots to subdirectories grouped by filter

Browse Source
This commit is contained in:
Luigi Maiorano
2026-01-28 15:58:38 +01:00
parent 365e70b834
commit 62e75fe899
3 changed files with 88 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
utils.py
View File

@@ -210,6 +210,13 @@ class JPMCSurvey(JPMCPlotsMixin):
# Rename columns with the extracted ImportIds
df.columns = new_columns
# Store unique values for filters (ignoring nulls) to detect "all selected" state
self.options_age = sorted(df['QID1'].drop_nulls().unique().to_list()) if 'QID1' in df.columns else []
self.options_gender = sorted(df['QID2'].drop_nulls().unique().to_list()) if 'QID2' in df.columns else []
self.options_consumer = sorted(df['Consumer'].drop_nulls().unique().to_list()) if 'Consumer' in df.columns else []
self.options_ethnicity = sorted(df['QID3'].drop_nulls().unique().to_list()) if 'QID3' in df.columns else []
self.options_income = sorted(df['QID15'].drop_nulls().unique().to_list()) if 'QID15' in df.columns else []
return df.lazy()
def _get_subset(self, q: pl.LazyFrame, QIDs, rename_cols=True, include_record_id=True) -> pl.LazyFrame:
@@ -239,24 +246,24 @@ class JPMCSurvey(JPMCPlotsMixin):
"""
# Apply filters
self.filter_age = age
if age is not None:
self.filter_age = age
q = q.filter(pl.col('QID1').is_in(age))
self.filter_gender = gender
if gender is not None:
self.filter_gender = gender
q = q.filter(pl.col('QID2').is_in(gender))
self.filter_consumer = consumer
if consumer is not None:
self.filter_consumer = consumer
q = q.filter(pl.col('Consumer').is_in(consumer))
self.filter_ethnicity = ethnicity
if ethnicity is not None:
self.filter_ethnicity = ethnicity
q = q.filter(pl.col('QID3').is_in(ethnicity))
self.filter_income = income
if income is not None:
self.filter_income = income
q = q.filter(pl.col('QID15').is_in(income))
self.data_filtered = q