dgx-spark-playbooks/nvidia/portfolio-optimization/assets/setup/src/portfolio.py
2026-01-02 22:21:53 +00:00

641 lines
22 KiB
Python

# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # noqa
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Portfolio class for managing and analyzing investment portfolios."""
import json
import os
import matplotlib.colors as mcolors
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
class Portfolio:
"""
Portfolio class for managing asset allocations and cash holdings.
Stores portfolio weights, cash allocation, and provides methods for
portfolio analysis and visualization.
Attributes
----------
name : str
Portfolio identifier
tickers : list
Asset symbols/tickers
weights : np.ndarray
Portfolio weights for each asset
cash : float
Cash allocation (typically 0-1)
time_range : tuple, optional
(start_date, end_date) for portfolio period
"""
def __init__(self, name="", tickers=None, weights=None, cash=0.0, time_range=None):
"""
Initialize Portfolio with assets, weights, and cash allocation.
Parameters
----------
name : str, default ""
Portfolio identifier/name
tickers : list, optional
Asset symbols (e.g., ['AAPL', 'MSFT'])
weights : array-like, optional
Portfolio weights for each asset
cash : float, default 0.0
Cash allocation
time_range : tuple, optional
(start_date, end_date) for portfolio period
"""
self.name = name
self.tickers = tickers if tickers is not None else []
self._n_assets = len(self.tickers)
self.weights = weights if weights is not None else []
self.cash = float(cash)
self.time_range = time_range
def __eq__(self, other_portfolio, atol=1e-3):
"""Check portfolio equality based on weights within tolerance."""
if isinstance(other_portfolio, Portfolio):
return np.allclose(self.weights, other_portfolio.weights, atol=atol)
return False
def _check_self_financing(self, weights=None, cash=None):
"""Verify that portfolio weights and cash sum to 1.0
(self-financing constraint)."""
if weights is None or weights.size == 0:
weights = self.weights
if cash is None:
cash = self.cash
self_finance = np.sum(weights) + cash
if np.abs(self_finance - 1) > 1e-3:
print(f"weights: {np.sum(weights)}; cash: {cash}")
raise ValueError("Portfolio weights and cash do not sum to 1!")
def portfolio_from_dict(self, portfolio_name, user_portfolio_dict, cash):
"""
Create portfolio from user-specified weights dictionary.
Parameters
----------
portfolio_name : str
Name for the portfolio
user_portfolio_dict : dict
Asset weights as {ticker: weight}
cash : float
Cash allocation
"""
weights = pd.Series(dtype=np.float64, index=self.tickers)
for ticker, weight in user_portfolio_dict.items():
ticker = ticker.upper()
if ticker in self.tickers:
weights[ticker] = weight
else:
raise ValueError("Selected ticker is not available in the dataset!")
weights = weights.fillna(0).T
weights = weights.to_numpy()
self._check_self_financing(weights, cash)
self.weights = np.array(weights)
self.cash = float(cash)
self.name = portfolio_name
def print_clean(self, cutoff=1e-3, min_percentage=0.0, rounding=3, verbose=False):
"""
Display clean portfolio allocation with formatting.
Filters positions based on both cutoff threshold and minimum percentage.
Only positions meeting both criteria are displayed/returned.
Parameters
----------
cutoff : float, default 1e-3
Minimum absolute weight to display (positions below this are
considered zero).
min_percentage : float, default 0.0
Minimum percentage threshold (0-100) for displaying/returning assets.
Only assets with absolute allocation >= min_percentage% will be included.
Example: min_percentage=2.0 shows only assets with ≥2% allocation.
rounding : int, default 3
Number of decimal places for display.
verbose : bool, default False
Whether to print portfolio breakdown.
Returns
-------
tuple
(clean_portfolio_dict, cash) - Dictionary of significant positions
and cash amount.
"""
residual = 0
clean_ptf_dict = {}
long_positions = {}
short_positions = {}
# Process each position
for idx, ticker in enumerate(self.tickers):
value = self.weights[idx]
if value > cutoff:
clean_ptf_dict[ticker] = value
long_positions[ticker] = value
elif value < -cutoff:
clean_ptf_dict[ticker] = value
short_positions[ticker] = value
else:
residual += value
cash = round(self.cash, rounding)
# Filter positions by minimum percentage threshold
min_threshold = min_percentage / 100.0 # Convert percentage to decimal
if min_threshold > 0:
# Filter clean_ptf_dict to only include positions >= min_percentage%
clean_ptf_dict = {
k: v for k, v in clean_ptf_dict.items() if abs(v) >= min_threshold
}
# Re-separate long and short positions after filtering
long_positions = {k: v for k, v in clean_ptf_dict.items() if v > 0}
short_positions = {k: v for k, v in clean_ptf_dict.items() if v < 0}
# Filter cash if below threshold
if abs(cash) < min_threshold:
cash = 0.0
if verbose:
# Portfolio header
print(f"\nPORTFOLIO: {self.name.upper()}")
print(f"{'-' * 40}")
if self.time_range:
print(f"Period: {self.time_range[0]} to {self.time_range[1]}")
# Long positions section
if long_positions:
print(f"\nLONG POSITIONS ({len(long_positions)} assets)")
print(f"{'-' * 25}")
total_long = 0
for ticker, weight in sorted(
long_positions.items(), key=lambda x: x[1], reverse=True
):
print(f"{ticker:8} {weight:>8.{rounding}f} ({weight * 100:>6.2f}%)")
total_long += weight
print(
f"{'Total Long':8} {total_long:>8.{rounding}f} "
f"({total_long * 100:>6.2f}%)"
)
# Short positions section
if short_positions:
print(f"\nSHORT POSITIONS ({len(short_positions)} assets)")
print(f"{'-' * 26}")
total_short = 0
for ticker, weight in sorted(
short_positions.items(), key=lambda x: x[1]
):
print(f"{ticker:8} {weight:>8.{rounding}f} ({weight * 100:>6.2f}%)")
total_short += weight
print(
f"{'Total Short':8} {total_short:>8.{rounding}f} "
f"({total_short * 100:>6.2f}%)"
)
# Cash and summary section
print("\nCASH & SUMMARY")
print(f"{'-' * 20}")
print(f"{'Cash':8} {cash:>8.{rounding}f} ({cash * 100:>6.2f}%)")
if abs(residual) > 1e-6:
print(
f"{'Residual':8} {residual:>8.{rounding}f} "
f"({residual * 100:>6.2f}%)"
)
# Portfolio totals
net_equity = sum(clean_ptf_dict.values())
total_allocation = net_equity + cash + residual
gross_exposure = sum(abs(w) for w in clean_ptf_dict.values())
print(
f"\n{'Net Equity':15} {net_equity:>8.{rounding}f} "
f"({net_equity * 100:>6.2f}%)"
)
print(
f"{'Total Portfolio':15} {total_allocation:>8.{rounding}f} "
f"({total_allocation * 100:>6.2f}%)"
)
print(
f"{'Gross Exposure':15} {gross_exposure:>8.{rounding}f} "
f"({gross_exposure * 100:>6.2f}%)"
)
print(f"{'-' * 40}")
return (clean_ptf_dict, float(cash))
def calculate_portfolio_expected_return(self, mean):
"""
Calculate portfolio expected return from asset mean returns.
Parameters
----------
mean : np.ndarray
Mean returns for each asset (shape: n_assets)
Returns
-------
float
Portfolio expected return
"""
assert (
mean.shape[0] == self._n_assets
), f"Incorrect mean vector size! Expecting: {self._n_assets}."
return mean @ self.weights
def calculate_portfolio_variance(self, covariance):
"""
Calculate portfolio variance from asset covariance matrix.
Parameters
----------
covariance : np.ndarray
Asset covariance matrix (shape: n_assets x n_assets)
Returns
-------
float
Portfolio variance
"""
assert (
covariance.shape[0] == self._n_assets
or covariance.shape[1] != self._n_assets
), (
f"Incorrect covariance size! Expecting: {self._n_assets} by "
+ f"{self._n_assets}."
)
return self.weights.T @ covariance @ self.weights
def plot_portfolio(
self,
show_plot=False,
ax=None,
title=None,
figsize=(12, 8),
style="modern",
cutoff=1e-3,
min_percentage=0.0,
sort_by_weight=True,
save_path=None,
dpi=300,
):
"""
Create a portfolio allocation visualization with gradient colors.
Uses color gradients based on position weights:
- Long positions: Blue gradient (light to dark based on weight)
- Short positions: Red gradient (light to dark based on absolute weight)
- Cash: Yellow (separate category at bottom)
Only displays assets above the minimum percentage threshold for cleaner plots.
Parameters
----------
show_plot : bool, default False
Whether to display the plot immediately.
ax : matplotlib.axes.Axes, optional
Existing axes to plot on. If None, creates new figure.
title : str, optional
Custom plot title. If None, auto-generates from portfolio info.
figsize : tuple, default (12, 8)
Figure size (width, height) in inches.
style : str, default "modern"
Visual style ("modern", "classic", "minimal").
cutoff : float, default 1e-3
Minimum weight to display (smaller positions grouped as "Other").
min_percentage : float, default 0.0
Minimum percentage threshold (0-100) for displaying assets.
Only assets with absolute allocation >= min_percentage% will be shown.
Example: min_percentage=1.0 shows only assets with ≥1% allocation.
sort_by_weight : bool, default True
Whether to sort positions by absolute weight (largest first).
save_path : str, optional
Path to save the plot. If None, plot is not saved.
dpi : int, default 300
Resolution for saved figure.
Returns
-------
matplotlib.axes.Axes
The axes object containing the plot.
"""
# Color schemes consistent with rebalance plot
color_schemes = {
"modern": {
"long": "#7cd7fe", # Light blue (frontier)
"short": "#ff8181", # Pink/Red (custom)
"cash": "#fcde7b", # Purple (assets)
"background": "#ffffff",
"grid": "#E0E0E0",
"text": "#000000",
}
}
colors = color_schemes.get(style, color_schemes["modern"])
plt.style.use("seaborn-v0_8-whitegrid")
# Create figure if needed
if ax is None:
fig, ax = plt.subplots(
figsize=figsize, dpi=dpi, facecolor=colors["background"]
)
ax.set_facecolor(colors["background"])
else:
_ = ax.get_figure() # fig (unused)
# Get portfolio data (filtering handled by print_clean)
portfolio_data, filtered_cash = self.print_clean(
cutoff=cutoff, min_percentage=min_percentage
)
cash = filtered_cash
# Separate long and short positions
long_positions = {k: v for k, v in portfolio_data.items() if v > 0}
short_positions = {k: v for k, v in portfolio_data.items() if v < 0}
# Prepare data for plotting
all_tickers = []
all_weights = []
all_colors = []
# Sort positions if requested
if sort_by_weight:
long_sorted = sorted(
long_positions.items(), key=lambda x: x[1], reverse=True
)
short_sorted = sorted(
short_positions.items(), key=lambda x: abs(x[1]), reverse=True
)
else:
long_sorted = list(long_positions.items())
short_sorted = list(short_positions.items())
# Create color gradients based on weights
# For long positions: gradient from light to dark blue
if long_positions:
max_long = max(long_positions.values())
min_long = (
min(long_positions.values())
if len(long_positions) > 1
else max_long * 0.1
)
long_range = max_long - min_long if max_long != min_long else max_long
# Create blue gradient colormap
long_cmap = mcolors.LinearSegmentedColormap.from_list(
"long_gradient", ["#7cd7fe", "#0046a4"], N=100
)
# For short positions: gradient from light to dark red
if short_positions:
max_short = max(abs(w) for w in short_positions.values())
min_short = (
min(abs(w) for w in short_positions.values())
if len(short_positions) > 1
else max_short * 0.1
)
short_range = max_short - min_short if max_short != min_short else max_short
# Create red gradient colormap
short_cmap = mcolors.LinearSegmentedColormap.from_list(
"short_gradient", ["#ff8181", "#961515"], N=100
)
# Add long positions with gradient colors
for ticker, weight in long_sorted:
all_tickers.append(ticker)
all_weights.append(weight)
if long_range > 0:
# Normalize weight to [0, 1] for colormap
intensity = (weight - min_long) / long_range
color = long_cmap(intensity)
else:
color = long_cmap(0.8) # Default intensity
all_colors.append(color)
# Add short positions with gradient colors
for ticker, weight in short_sorted:
all_tickers.append(ticker)
all_weights.append(weight)
if short_range > 0:
# Normalize absolute weight to [0, 1] for colormap
intensity = (abs(weight) - min_short) / short_range
color = short_cmap(intensity)
else:
color = short_cmap(0.8) # Default intensity
all_colors.append(color)
# Add cash at the bottom as separate category (yellow)
if abs(cash) > cutoff:
all_tickers.append("CASH")
all_weights.append(cash)
all_colors.append(colors["cash"]) # Gold/Yellow color for cash
# Create horizontal bar chart with extra space before cash
cash_gap = 0.8 # Extra space between equity positions and cash
y_positions = []
# Calculate positions with gap before cash
for i, ticker in enumerate(all_tickers):
if (
ticker == "CASH" and i > 0
): # Add gap before cash if it's not the only item
y_positions.append(i + cash_gap)
else:
y_positions.append(i)
_ = ax.barh( # bars (unused)
y_positions,
all_weights,
color=all_colors,
edgecolor="white",
linewidth=0.8,
alpha=0.8,
)
# Customize appearance - reverse y-axis so first items appear at top
ax.set_yticks(y_positions)
ax.set_yticklabels(all_tickers, fontsize=9, color=colors["text"])
ax.invert_yaxis() # Reverse y-axis so long positions appear at top,
# cash at bottom
ax.set_xlabel("Portfolio Weight", fontsize=10, color=colors["text"])
ax.set_ylabel("Assets", fontsize=10, color=colors["text"])
# Set title
if title is None:
portfolio_name = self.name if self.name else "Portfolio"
if self.time_range:
title = (
f"{portfolio_name} Allocation\n{self.time_range[0]} to "
f"{self.time_range[1]}"
)
else:
title = f"{portfolio_name} Allocation"
ax.set_title(title, fontsize=11, pad=15, color=colors["text"])
# Percentage labels removed for cleaner appearance
# Set x-axis limits with padding
max_abs_weight = max(abs(w) for w in all_weights) if all_weights else 0.1
padding = max_abs_weight * 0.15
ax.set_xlim(-max_abs_weight - padding, max_abs_weight + padding)
# Add horizontal separator line between cash and risky assets
if abs(cash) > cutoff and len(all_tickers) > 1:
# Find cash position and position separator in the middle of the gap
cash_ticker_index = next(
(i for i, ticker in enumerate(all_tickers) if ticker == "CASH"), -1
)
if cash_ticker_index >= 0:
cash_y_pos = y_positions[cash_ticker_index]
# Find the last risky asset position (just before cash)
last_risky_y_pos = (
y_positions[cash_ticker_index - 1] if cash_ticker_index > 0 else 0
)
# Position separator in the middle of the gap
separator_y = (cash_y_pos + last_risky_y_pos) / 2
ax.axhline(
separator_y,
color=colors["text"],
linewidth=1,
alpha=0.6,
linestyle="--",
)
# Grid and styling - subtle grid similar to backtest
ax.grid(True, alpha=0.3, color=colors["grid"])
ax.spines["top"].set_visible(False)
ax.spines["right"].set_visible(False)
ax.spines["left"].set_color(colors["grid"])
ax.spines["bottom"].set_color(colors["grid"])
# Tick styling
ax.tick_params(axis="both", colors=colors["text"], labelsize=9)
ax.tick_params(axis="x", which="both", bottom=True, top=False)
ax.tick_params(axis="y", which="both", left=False, right=False)
# Legend with gradient representation
legend_elements = []
if long_positions:
# Use darkest blue for legend
num_long = len(long_positions)
legend_elements.append(
plt.Rectangle(
(0, 0),
1,
1,
facecolor="#1F5F8B",
label=f"Long Positions ({num_long})",
)
)
if short_positions:
# Use darkest red for legend
num_short = len(short_positions)
legend_elements.append(
plt.Rectangle(
(0, 0),
1,
1,
facecolor="#8B0000",
label=f"Short Positions ({num_short})",
)
)
if abs(cash) > cutoff:
legend_elements.append(
plt.Rectangle((0, 0), 1, 1, facecolor="#FFD700", label="Cash")
)
if legend_elements:
ax.legend(
handles=legend_elements,
loc="upper left",
frameon=True,
fancybox=True,
shadow=True,
framealpha=0.9,
fontsize=8,
)
plt.tight_layout()
# Save if requested
if save_path:
save_name = f"{self.name.lower()}_allocation.png"
plt.savefig(
os.path.join(save_path, save_name),
dpi=dpi,
bbox_inches="tight",
facecolor=colors["background"],
edgecolor="none",
)
print(f"Portfolio plot saved: {save_path}")
# Show if requested
if show_plot:
plt.show()
return ax
def save_portfolio(self, save_path):
"""Save portfolio to JSON file."""
save_weights = self.weights.tolist()
portfolio = {
"name": self.name,
"weights": save_weights,
"cash": self.cash,
"tickers": self.tickers,
"time_range": self.time_range,
}
with open(save_path, "w") as json_file:
json.dump(portfolio, json_file, indent=4)
def load_portfolio_from_json(self, load_path):
"""Load portfolio from JSON file and update current instance."""
with open(load_path, "r") as json_file:
data = json.load(json_file)
self.name = data["name"]
self.tickers = data["tickers"]
self._n_assets = len(self.tickers)
weights = np.array(data["weights"])
self._check_self_financing(weights, data["cash"])
self.weights = weights
self.cash = data["cash"]
self.time_range = data["time_range"]