@@ -804,8 +804,8 @@ indices_data = read_data(
804804Then, calculate the yearly returns using polars:
805805
806806``` {code-cell} ipython3
807- # Add year column and calculate yearly returns
808- yearly_returns_list = []
807+ # Combine all yearly returns using concat and pivot approach
808+ all_yearly_data = []
809809
810810for index_col in indices_data.columns:
811811 if index_col != 'Date':
@@ -817,16 +817,18 @@ for index_col in indices_data.columns:
817817 pl.col(index_col).last().alias('last_price')
818818 ])
819819 .with_columns(
820- ((pl.col('last_price') - pl.col('first_price')) / pl.col('first_price')) .alias(indices_list[index_col] )
820+ ((pl.col('last_price') - pl.col('first_price') + 1e-10 ) / ( pl.col('first_price') + 1e-10)) .alias('return' )
821821 )
822- .select(['year', indices_list[index_col]]))
822+ .with_columns(pl.lit(indices_list[index_col]).alias('index_name'))
823+ .select(['year', 'index_name', 'return']))
823824
824- yearly_returns_list .append(yearly_data)
825+ all_yearly_data .append(yearly_data)
825826
826- # Join all yearly returns
827- yearly_returns = yearly_returns_list[0]
828- for df in yearly_returns_list[1:]:
829- yearly_returns = yearly_returns.join(df, on='year', how='outer')
827+ # Concatenate all data
828+ combined_data = pl.concat(all_yearly_data)
829+
830+ # Pivot to get indices as columns
831+ yearly_returns = combined_data.pivot(values='return', index='year', on='index_name')
830832
831833yearly_returns
832834```
0 commit comments