@@ -804,8 +804,8 @@ indices_data = read_data(
804
804
Then, calculate the yearly returns using polars:
805
805
806
806
``` {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 = []
809
809
810
810
for index_col in indices_data.columns:
811
811
if index_col != 'Date':
@@ -817,16 +817,18 @@ for index_col in indices_data.columns:
817
817
pl.col(index_col).last().alias('last_price')
818
818
])
819
819
.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' )
821
821
)
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']))
823
824
824
- yearly_returns_list .append(yearly_data)
825
+ all_yearly_data .append(yearly_data)
825
826
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')
830
832
831
833
yearly_returns
832
834
```
0 commit comments