sovai 0.2.30__tar.gz → 0.2.31__tar.gz

{sovai-0.2.30 → sovai-0.2.31}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: sovai
-Version: 0.2.30
+Version: 0.2.31
 Summary: python for asset management
 License: Proprietary License
 Author: Derek Snow

{sovai-0.2.30 → sovai-0.2.31}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "sovai"
-version = "0.2.30"
+version = "0.2.31"
 description = "python for asset management"
 authors = ["Derek Snow <d.snow@outlook.com>"]
 license = "Proprietary License"

{sovai-0.2.30 → sovai-0.2.31}/sovai/extensions/pandas_extensions.py

@@ -101,7 +101,10 @@ import pandas as pd
 from typing import Union, List
 from sovai.extensions.ask_df_llm import Ask
 from sovai.extensions.filter_df import Filter
-
+import pandas as pd
+from functools import lru_cache
+import re
+from typing import Union, Tuple
 
 class CustomDataFrame(pd.DataFrame):
     @property
@@ -125,6 +128,82 @@ class CustomDataFrame(pd.DataFrame):
         filtered_df = Filter.filter(conditions, df_multi=self, verbose=verbose)
         return CustomDataFrame(filtered_df)
     
+    @staticmethod
+    @lru_cache(maxsize=1)
+    def _load_combined_df():
+        df = pd.read_parquet("https://storage.googleapis.com/sovai-public/concats/filters/latest.parquet")
+        # print("Columns in combined DataFrame:", df.columns)  # Debug print
+        # print("Index name:", df.index.name)  # Debug print for index name
+        return df
+    
+    def merge_data(self, column: str) -> 'CustomDataFrame':
+        """
+        Merge the current DataFrame with the combined DataFrame based on ticker and a specified column.
+        
+        :param column: The column from the combined DataFrame to merge.
+        :return: A new CustomDataFrame with the merged data.
+        """
+        # Load the combined DataFrame
+        df_combined = self._load_combined_df()
+        
+        # Ensure the specified column exists in the combined DataFrame
+        if column not in df_combined.columns:
+            raise ValueError(f"Column '{column}' not found in the combined DataFrame.")
+        
+        # Identify the ticker information in the current DataFrame
+        ticker_info = self._identify_ticker_info()
+        
+        # Prepare the current DataFrame for merging
+        df_to_merge = self.copy()
+        if ticker_info[0] != 'column':
+            df_to_merge = df_to_merge.reset_index()
+        
+        # Standardize ticker format in the current DataFrame
+        
+        # Prepare the combined DataFrame for merging
+        df_combined_merged = df_combined.reset_index()
+        
+        # Perform the merge
+        merged_df = pd.merge(df_to_merge, df_combined_merged[['ticker', column]], 
+                             left_on=ticker_info[1], right_on='ticker', 
+                             how='left')
+        
+        # Drop the redundant 'ticker' column from the merge if it's not the original ticker column
+        if 'ticker' in merged_df.columns and 'ticker' != ticker_info[1]:
+            merged_df = merged_df.drop('ticker', axis=1)
+        
+        # Restore the original index structure
+        if ticker_info[0] == 'index':
+            merged_df = merged_df.set_index(ticker_info[1])
+        elif ticker_info[0] == 'multi_index':
+            index_cols = [col for col in df_to_merge.columns if col not in self.columns]
+            merged_df = merged_df.set_index(index_cols)
+        
+        return CustomDataFrame(merged_df)
+    
+    def _identify_ticker_info(self) -> Tuple[str, str]:
+        """
+        Identify the location and name of ticker information in the DataFrame.
+        
+        :return: A tuple of (location, column_name).
+                 Location can be 'index', 'multi_index', or 'column'.
+        """
+        # Check if ticker is in the index
+        if self.index.name and 'ticker' in self.index.name.lower():
+            return ('index', self.index.name)
+        elif isinstance(self.index, pd.MultiIndex):
+            for name in self.index.names:
+                if name and 'ticker' in name.lower():
+                    return ('multi_index', name)
+        
+        # If ticker is not in the index, check the columns
+        for col in self.columns:
+            if 'ticker' in col.lower():
+                return ('column', col)
+        
+        raise ValueError("Couldn't identify ticker information in the DataFrame.")
+    
+    
     def get_latest(self, column=None):
         # Your get_latest implementation
         if isinstance(self.index, pd.MultiIndex) and "date" in self.index.names:

{sovai-0.2.30 → sovai-0.2.31}/sovai/get_plots.py

@@ -95,6 +95,7 @@ PLOT_FUNCTION_MAPPER = {
     ("allocation", "line", "local", True): create_line_plot_allocation,
     ("allocation", "stacked", "local", True): create_stacked_bar_plot_allocation,
     ("earnings", "line", "local", True): create_earnings_surprise_plot,
+    ("earnings", "tree", "local", True): earnings_tree,
     ("bankruptcy", "compare", "local", True): plot_bankruptcy_monthly_line,
     ("bankruptcy", "pca_clusters", "local", True): plot_pca_clusters,
     ("bankruptcy", "predictions", "local", True): plot_ticker_widget,

{sovai-0.2.30 → sovai-0.2.31}/sovai/plots/earnings_surprise/earnings_surprise_plots.py

@@ -114,3 +114,73 @@ def create_earnings_surprise_plot():
     return app.run_server(
         debug=False, port=get_unique_port(app_name)
     )  # Apply exponential moving average to smooth the data
+
+
+import plotly.express as px
+import pandas as pd
+import numpy as np
+
+def earnings_tree():
+    df_earn = data("earnings/surprise")
+    df_earn = df_earn.merge_data("market_cap")
+    # Ensure 'ticker' is a column, not an index
+    if 'ticker' not in df_earn.columns:
+        df_earn = df_earn.reset_index()
+    
+    # Calculate market cap deciles and handle NaN values
+    df_earn['market_cap_decile'] = pd.qcut(df_earn['market_cap'].rank(method='first'), q=10, labels=range(1, 11))
+    df_earn['market_cap_decile'] = df_earn['market_cap_decile'].astype(str).replace('nan', 'Unknown')
+    
+    # Update decile labels
+    df_earn['market_cap_decile'] = df_earn['market_cap_decile'].apply(lambda x: f"{x} Cap Decile" if x != 'Unknown' else x)
+    
+    # Create the main treemap
+    fig = px.treemap(
+        df_earn,
+        path=[px.Constant("All"), 'market_cap_decile', "ticker"],
+        values="market_cap",
+        color="surprise_probability",
+        color_continuous_scale="RdYlBu",
+        color_continuous_midpoint=0,
+        hover_data=["surprise_probability", "market_cap", "date"],
+        title="Earnings Surprise Probability and Market Cap Treemap",
+    )
+    
+    # Update layout for better appearance
+    fig.update_layout(
+        font=dict(family="Arial", size=14),
+        title=dict(font=dict(size=20, color="lightgrey")),
+        coloraxis_colorbar=dict(
+            title="Surprise Probability",
+            tickformat=".2%",
+            len=0.75,
+            thickness=20,
+            title_font=dict(size=16),
+            tickfont=dict(size=14),
+        ),
+        height=900,
+        # width=1200,
+        # plot_bgcolor='rgba(240,240,240,0.8)',
+        # paper_bgcolor='rgba(240,240,240,0.8)',
+        margin=dict(t=50, l=25, r=25, b=25),
+    )
+    
+    # Update traces for better text visibility and style
+    fig.update_traces(
+        textposition="middle center",
+        textfont=dict(family="Arial", size=12, color="black"),
+        hovertemplate=(
+            "<b>%{label}</b><br>"
+            "Date: %{customdata[2]}<br>"
+            "Surprise Probability: %{customdata[0]:.2%}<br>"
+            "Market Cap: $%{customdata[1]:,.0f}M"
+            "<extra></extra>"
+        ),
+        marker=dict(
+            cornerradius=5,
+            line=dict(width=1, color="white"),
+        ),
+    )
+    
+    # Show the plot
+    fig.show()