portfolio rebalance

portfolio_rebalance.py

# I have 2 dataframes:
# a) My positions with columns: ticker, purchase_date, quantity, price_per_share
# b) My ideal portfolio: ticker, size (in percent), current_price
# Emit a dataframe with trades to execute to rebalance my portfolio: 
# ticker, quantity (positive means buy and negative means sell)
# State your assumptions

import polars as pl

def rebalance(positions: pl.DataFrame, ideal_portfolio: pl.DataFrame) -> pl.DataFrame:
  # Assumption: ideal_portfolio contains prices for all tickers in positions - even those that we need to fully exit
  return (positions
    .group_by("ticker")
    .agg(pl.col("quantity").sum().alias("num_shares"))
    .join(ideal_portfolio, on="ticker", how="outer", coalesce=True)
    .fill_null(0)
    .with_columns((pl.col("num_shares") * pl.col("current_price")).alias("value"))
    .with_columns(pl.col("value").sum().over().alias("total_value"))
    .with_columns(((pl.col("size")/100 * pl.col("total_value") - pl.col("value"))/pl.col("current_price")).alias("quantity"))
    .filter(pl.col("quantity") != 0)
    .select(["ticker", "quantity"]))