kumoai 2.14.0.dev202601011731__cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

kumoai/utils/display.py

@@ -0,0 +1,51 @@
+from collections.abc import Sequence
+
+import pandas as pd
+
+from kumoai import in_notebook, in_snowflake_notebook
+
+
+def message(msg: str) -> None:
+    msg = msg.replace("`", "'") if not in_notebook() else msg
+
+    if in_snowflake_notebook():
+        import streamlit as st
+        st.markdown(msg)
+    elif in_notebook():
+        from IPython.display import Markdown, display
+        display(Markdown(msg))
+    else:
+        print(msg)
+
+
+def title(msg: str) -> None:
+    message(f"### {msg}" if in_notebook() else f"{msg}:")
+
+
+def italic(msg: str) -> None:
+    message(f"*{msg}*" if in_notebook() else msg)
+
+
+def unordered_list(items: Sequence[str]) -> None:
+    if in_notebook():
+        msg = '\n'.join([f"- {item}" for item in items])
+    else:
+        msg = '\n'.join([f"• {item.replace('`', '')}" for item in items])
+    message(msg)
+
+
+def dataframe(df: pd.DataFrame) -> None:
+    if in_snowflake_notebook():
+        import streamlit as st
+        st.dataframe(df, hide_index=True)
+    elif in_notebook():
+        from IPython.display import display
+        try:
+            if hasattr(df.style, 'hide'):
+                display(df.style.hide(axis='index'))  # pandas=2
+            else:
+                display(df.style.hide_index())  # pandas<1.3
+        except ImportError:
+            print(df.to_string(index=False))  # missing jinja2
+    else:
+        print(df.to_string(index=False))

kumoai/utils/forecasting.py

@@ -0,0 +1,209 @@
+from typing import Dict, List
+
+import pandas as pd
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots
+
+
+class ForecastVisualizer:
+    r"""A tool for quickly visualizing forecast results from the holdout
+    dataframe of a forecasting job.
+
+    .. code-block:: python
+
+        import kumoai
+
+        # Retrieve job results from a training training job. Note
+        # that the job ID passed here must be in a completed state:
+        job_result = kumoai.TrainingJob("trainingjob-...").result()
+
+        # Read the holdout table as a Pandas DataFrame:
+        holdout_df = job_result.holdout_df()
+
+        # Pass holdout table to ForecastVisualizer and visualize results
+        holdout_forecast = kumoai.utils.ForecastVisualizer(holdout_df)
+        holdout_forecast.visualize()
+    """
+    def __init__(self, holdout_df: pd.DataFrame) -> None:
+        # Sort the holdout dataframe and extract unique entities:
+        self.forecast = holdout_df.sort_values(['ENTITY', 'TIMESTAMP'])
+        self.entities = holdout_df['ENTITY'].unique().tolist()
+
+        self.fig = self._initialize_subplot()
+        self.buttons: List[Dict] = []
+        self.plot_config = {
+            'target': {
+                'color': 'blue',
+                'name': 'TARGET'
+            },
+            'prediction': {
+                'color': 'red',
+                'name': 'TARGET_PRED'
+            },
+            'residuals': {
+                'color': 'green',
+                'name': 'Residuals'
+            },
+            'residuals_time': {
+                'color': 'orange',
+                'name': 'Residuals Over Time'
+            }
+        }
+
+    @staticmethod
+    def _initialize_subplot() -> go.Figure:
+        r"""Initializes the subplot structure with three rows:
+        Row 1: Line plot of actual forecast vs predicted
+        Row 2: Line plot of residuals overtime
+        Row 3: Histogram distribution of residuals
+        """
+        return make_subplots(
+            rows=3,
+            cols=1,
+            specs=[[{
+                "type": "scatter"
+            }], [{
+                "type": "scatter"
+            }], [{
+                "type": "xy"
+            }]],
+            row_heights=[0.5, 0.25, 0.25],
+            subplot_titles=('Forecast vs Actual', 'Residuals Over Time',
+                            'Residuals Distribution'),
+        )
+
+    def _create_time_series_trace(
+        self,
+        data: pd.Series,
+        entity: str,
+        trace_type: str,
+        visibility: bool,
+    ) -> go.Scatter:
+        r"""Create a time series trace for either target or prediction"""
+        config = self.plot_config[trace_type]
+        return go.Scatter(
+            x=data["TIMESTAMP"],
+            y=data[config['name']],
+            name=f"{entity} - {config['name']}",
+            mode="lines",
+            line=dict(color=config['color']),
+            visible=visibility,
+            opacity=0.75,
+        )
+
+    def _create_residuals_time_trace(
+        self,
+        data: pd.Series,
+        entity: str,
+        visibility: bool,
+    ) -> go.Scatter:
+        r"""Create a time series trace for residuals over time"""
+        residuals = data["TARGET"] - data["TARGET_PRED"]
+        return go.Scatter(
+            x=data["TIMESTAMP"],
+            y=residuals,
+            name=f"{entity} - Residuals Over Time",
+            mode="lines+markers",
+            line=dict(color=self.plot_config['residuals_time']['color']),
+            visible=visibility,
+            opacity=0.75,
+        )
+
+    def _create_residuals_hist_trace(
+        self,
+        data: pd.Series,
+        entity: str,
+        visibility: bool,
+    ) -> go.Histogram:
+        r"""Create a histogram trace for residuals distribution."""
+        residuals = data["TARGET"] - data["TARGET_PRED"]
+        return go.Histogram(
+            x=residuals,
+            name=f"{entity} - Residuals Distribution",
+            marker=dict(color=self.plot_config['residuals']['color']),
+            visible=visibility,
+            opacity=0.75,
+            nbinsx=30,
+        )
+
+    def _create_button(self, index: int, entity: str) -> None:
+        r"""Create visibility toggle button for an entity."""
+        # target, prediction, residuals time, and residuals hist:
+        num_traces_per_entity = 4
+        total_traces = len(self.entities) * num_traces_per_entity
+
+        button = dict(label=entity, method="update", args=[{
+            "visible": [False] * total_traces
+        }])
+
+        # Set visibility for the entity's traces:
+        base_index = index * num_traces_per_entity
+        for i in range(num_traces_per_entity):
+            button["args"][0]["visible"][base_index + i] = True  # type: ignore
+
+        self.buttons.append(button)
+
+    def _create_traces(self) -> None:
+        """Create all traces for the visualization."""
+        for i, entity in enumerate(self.entities):
+            entity_data = self.forecast.loc[self.forecast.ENTITY == entity]
+
+            # First entity's traces are visible by default:
+            visibility = (i == 0)
+
+            # Create traces
+            trace_target = self._create_time_series_trace(
+                entity_data, entity, 'target', visibility)
+            trace_pred = self._create_time_series_trace(
+                entity_data, entity, 'prediction', visibility)
+            trace_residuals_time = self._create_residuals_time_trace(
+                entity_data, entity, visibility)
+            trace_residuals_hist = self._create_residuals_hist_trace(
+                entity_data, entity, visibility)
+
+            # Add traces to appropriate subplots
+            self.fig.add_trace(trace_target, row=1, col=1)
+            self.fig.add_trace(trace_pred, row=1, col=1)
+            self.fig.add_trace(trace_residuals_time, row=2, col=1)
+            self.fig.add_trace(trace_residuals_hist, row=3, col=1)
+
+            self._create_button(i, entity)
+
+    def _update_layout(self) -> None:
+        r"""Update the figure layout with all necessary configurations."""
+        self.fig.update_layout(
+            updatemenus=[
+                dict(active=0, buttons=self.buttons, direction="down", pad={
+                    "r": 10,
+                    "t": 10
+                }, showactive=True, x=1, xanchor="left", y=1.07, yanchor="top")
+            ],
+            title="Forecast Results by Department",
+            height=1000,  # Increased height to accommodate third plot
+            width=1300,
+            showlegend=True,
+            hovermode='x unified')
+
+        # Update axis labels and add zero reference line for residuals
+        self.fig.update_xaxes(title_text="Timestamp", row=1, col=1)
+        self.fig.update_xaxes(title_text="Timestamp", row=2, col=1)
+        self.fig.update_xaxes(title_text="Residual Value", row=3, col=1)
+
+        self.fig.update_yaxes(title_text="Patient Volume", row=1, col=1)
+        self.fig.update_yaxes(title_text="Residual Value", row=2, col=1)
+        self.fig.update_yaxes(title_text="Frequency", row=3, col=1)
+
+        # Add zero reference line for residuals time series
+        self.fig.add_hline(
+            y=0,
+            line_dash="dash",
+            line_color="gray",
+            row=2,
+            col=1,
+        )
+
+    def visualize(self) -> None:
+        r"""Generate and display the complete visualization."""
+        self._create_traces()
+        self._update_layout()
+        self.fig.show()

kumoai/utils/progress_logger.py

@@ -0,0 +1,343 @@
+import re
+import sys
+import time
+from typing import Any
+
+from rich.console import Console, ConsoleOptions, RenderResult
+from rich.live import Live
+from rich.padding import Padding
+from rich.progress import (
+    BarColumn,
+    MofNCompleteColumn,
+    Progress,
+    Task,
+    TextColumn,
+    TimeRemainingColumn,
+)
+from rich.spinner import Spinner
+from rich.table import Table
+from rich.text import Text
+from typing_extensions import Self
+
+
+class ProgressLogger:
+    def __init__(self, msg: str, verbose: bool = True) -> None:
+        self.msg = msg
+        self.verbose = verbose
+
+        self.logs: list[str] = []
+
+        self.start_time: float | None = None
+        self.end_time: float | None = None
+
+    @classmethod
+    def default(cls, msg: str, verbose: bool = True) -> 'ProgressLogger':
+        from kumoai import in_snowflake_notebook
+
+        if in_snowflake_notebook():
+            return StreamlitProgressLogger(msg, verbose)
+        return RichProgressLogger(msg, verbose)
+
+    @property
+    def duration(self) -> float:
+        assert self.start_time is not None
+        if self.end_time is not None:
+            return self.end_time - self.start_time
+        return time.perf_counter() - self.start_time
+
+    def log(self, msg: str) -> None:
+        self.logs.append(msg)
+
+    def init_progress(self, total: int, description: str) -> None:
+        pass
+
+    def step(self) -> None:
+        pass
+
+    def __enter__(self) -> Self:
+        self.start_time = time.perf_counter()
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.end_time = time.perf_counter()
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}({self.msg})'
+
+
+class ColoredMofNCompleteColumn(MofNCompleteColumn):
+    def __init__(self, style: str = 'green') -> None:
+        super().__init__()
+        self.style = style
+
+    def render(self, task: Task) -> Text:
+        return Text(str(super().render(task)), style=self.style)
+
+
+class ColoredTimeRemainingColumn(TimeRemainingColumn):
+    def __init__(self, style: str = 'cyan') -> None:
+        super().__init__()
+        self.style = style
+
+    def render(self, task: Task) -> Text:
+        return Text(str(super().render(task)), style=self.style)
+
+
+class RichProgressLogger(ProgressLogger):
+    def __init__(
+        self,
+        msg: str,
+        verbose: bool = True,
+        refresh_per_second: int = 10,
+    ) -> None:
+        super().__init__(msg=msg, verbose=verbose)
+
+        self.refresh_per_second = refresh_per_second
+
+        self._progress: Progress | None = None
+        self._task: int | None = None
+
+        self._live: Live | None = None
+        self._exception: bool = False
+
+    def init_progress(self, total: int, description: str) -> None:
+        assert self._progress is None
+        if self.verbose:
+            self._progress = Progress(
+                TextColumn(f'   ↳ {description}', style='dim'),
+                BarColumn(bar_width=None),
+                ColoredMofNCompleteColumn(style='dim'),
+                TextColumn('•', style='dim'),
+                ColoredTimeRemainingColumn(style='dim'),
+            )
+            self._task = self._progress.add_task("Progress", total=total)
+
+    def step(self) -> None:
+        if self.verbose:
+            assert self._progress is not None
+            assert self._task is not None
+            self._progress.update(self._task, advance=1)  # type: ignore
+
+    def __enter__(self) -> Self:
+        from kumoai import in_notebook
+
+        super().__enter__()
+
+        if not in_notebook():  # Render progress bar in TUI.
+            sys.stdout.write("\x1b]9;4;3\x07")
+            sys.stdout.flush()
+
+        if self.verbose:
+            self._live = Live(
+                self,
+                refresh_per_second=self.refresh_per_second,
+                vertical_overflow='visible',
+            )
+            self._live.start()
+
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        from kumoai import in_notebook
+
+        super().__exit__(exc_type, exc_val, exc_tb)
+
+        if exc_type is not None:
+            self._exception = True
+
+        if self._progress is not None:
+            self._progress.stop()
+            self._progress = None
+            self._task = None
+
+        if self._live is not None:
+            self._live.update(self, refresh=True)
+            self._live.stop()
+            self._live = None
+
+        if not in_notebook():
+            sys.stdout.write("\x1b]9;4;0\x07")
+            sys.stdout.flush()
+
+    def __rich_console__(
+        self,
+        console: Console,
+        options: ConsoleOptions,
+    ) -> RenderResult:
+
+        table = Table.grid(padding=(0, 1))
+
+        icon: Text | Padding
+        if self._exception:
+            style = 'red'
+            icon = Text('❌', style=style)
+        elif self.end_time is not None:
+            style = 'green'
+            icon = Text('✅', style=style)
+        else:
+            style = 'cyan'
+            icon = Padding(Spinner('dots', style=style), (0, 1, 0, 0))
+
+        title = Text.from_markup(
+            f'{self.msg} ({self.duration:.2f}s)',
+            style=style,
+        )
+        table.add_row(icon, title)
+
+        for log in self.logs:
+            table.add_row('', Text(f'↳ {log}', style='dim'))
+
+        yield table
+
+        if self.verbose and self._progress is not None:
+            yield self._progress.get_renderable()
+
+
+class StreamlitProgressLogger(ProgressLogger):
+    def __init__(
+        self,
+        msg: str,
+        verbose: bool = True,
+    ) -> None:
+        super().__init__(msg=msg, verbose=verbose)
+
+        self._status: Any = None
+
+        self._total = 0
+        self._current = 0
+        self._description: str = ''
+        self._progress: Any = None
+
+    def __enter__(self) -> Self:
+        super().__enter__()
+
+        import streamlit as st
+
+        # Adjust layout for prettier output:
+        st.markdown(STREAMLIT_CSS, unsafe_allow_html=True)
+
+        if self.verbose:
+            self._status = st.status(
+                f':blue[{self._sanitize_text(self.msg)}]',
+                expanded=True,
+            )
+
+        return self
+
+    def log(self, msg: str) -> None:
+        super().log(msg)
+        if self.verbose and self._status is not None:
+            self._status.write(self._sanitize_text(msg))
+
+    def init_progress(self, total: int, description: str) -> None:
+        if self.verbose and self._status is not None:
+            self._total = total
+            self._current = 0
+            self._description = self._sanitize_text(description)
+            percent = min(self._current / self._total, 1.0)
+            self._progress = self._status.progress(
+                value=percent,
+                text=f'{self._description} [{self._current}/{self._total}]',
+            )
+
+    def step(self) -> None:
+        self._current += 1
+
+        if self.verbose and self._progress is not None:
+            percent = min(self._current / self._total, 1.0)
+            self._progress.progress(
+                value=percent,
+                text=f'{self._description} [{self._current}/{self._total}]',
+            )
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        super().__exit__(exc_type, exc_val, exc_tb)
+
+        if not self.verbose or self._status is None:
+            return
+
+        label = f'{self._sanitize_text(self.msg)} ({self.duration:.2f}s)'
+
+        if exc_type is not None:
+            self._status.update(
+                label=f':red[{label}]',
+                state='error',
+                expanded=True,
+            )
+        else:
+            self._status.update(
+                label=f':green[{label}]',
+                state='complete',
+                expanded=True,
+            )
+
+    @staticmethod
+    def _sanitize_text(msg: str) -> str:
+        return re.sub(r'\[/?bold\]', '**', msg)
+
+
+STREAMLIT_CSS = """
+<style>
+/* Fix horizontal scrollbar */
+.stExpander summary {
+    width: auto;
+}
+
+/* Fix paddings/margins */
+.stExpander summary {
+    padding: 0.75rem 1rem 0.5rem;
+}
+.stExpander p {
+    margin: 0px 0px 0.2rem;
+}
+.stExpander [data-testid="stExpanderDetails"] {
+    padding-bottom: 1.45rem;
+}
+.stExpander .stProgress div:first-child {
+    padding-bottom: 4px;
+}
+
+/* Fix expand icon position */
+.stExpander summary svg {
+    height: 1.5rem;
+}
+
+/* Fix summary icons */
+.stExpander summary [data-testid="stExpanderIconCheck"] {
+    font-size: 1.8rem;
+    margin-top: -3px;
+    color: rgb(21, 130, 55);
+}
+.stExpander summary [data-testid="stExpanderIconError"] {
+    font-size: 1.8rem;
+    margin-top: -3px;
+    color: rgb(255, 43, 43);
+}
+.stExpander summary span:first-child span:first-child {
+    width: 1.6rem;
+}
+
+/* Add border between title and content */
+.stExpander [data-testid="stExpanderDetails"] {
+    border-top: 1px solid rgba(30, 37, 47, 0.2);
+    padding-top: 0.5rem;
+}
+
+/* Fix title font size */
+.stExpander summary p {
+    font-size: 1rem;
+}
+
+/* Gray out content */
+.stExpander [data-testid="stExpanderDetails"] {
+    color: rgba(30, 37, 47, 0.5);
+}
+
+/* Fix progress bar font size */
+.stExpander .stProgress p {
+    line-height: 1.6;
+    font-size: 1rem;
+    color: rgba(30, 37, 47, 0.5);
+}
+</style>
+"""

kumoai/utils/sql.py

@@ -0,0 +1,3 @@
+def quote_ident(name: str) -> str:
+    r"""Quotes a SQL identifier."""
+    return '"' + name.replace('"', '""') + '"'

kumoai-2.14.0.dev202601011731.dist-info/METADATA

@@ -0,0 +1,71 @@
+Metadata-Version: 2.4
+Name: kumoai
+Version: 2.14.0.dev202601011731
+Summary: AI on the Modern Data Stack
+Author-email: "Kumo.AI" <hello@kumo.ai>
+License-Expression: MIT
+Project-URL: homepage, https://kumo.ai
+Project-URL: documentation, https://kumo.ai/docs
+Keywords: deep-learning,graph-neural-networks,cloud-data-warehouse
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3 :: Only
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: pandas
+Requires-Dist: pyarrow<19.0.0,>=8.0.0
+Requires-Dist: requests>=2.28.2
+Requires-Dist: urllib3
+Requires-Dist: plotly
+Requires-Dist: typing_extensions>=4.5.0
+Requires-Dist: kumo-api==0.49.0
+Requires-Dist: tqdm>=4.66.0
+Requires-Dist: aiohttp>=3.10.0
+Requires-Dist: pydantic>=1.10.21
+Requires-Dist: rich>=9.0.0
+Provides-Extra: doc
+Requires-Dist: sphinx; extra == "doc"
+Requires-Dist: sphinx-book-theme; extra == "doc"
+Requires-Dist: sphinx-copybutton; extra == "doc"
+Requires-Dist: sphinx-autodoc-typehints; extra == "doc"
+Requires-Dist: graphviz; extra == "doc"
+Provides-Extra: test
+Requires-Dist: pytest; extra == "test"
+Requires-Dist: pytest-mock; extra == "test"
+Requires-Dist: requests-mock; extra == "test"
+Provides-Extra: sqlite
+Requires-Dist: adbc_driver_sqlite; extra == "sqlite"
+Provides-Extra: snowflake
+Requires-Dist: numpy<2.0; extra == "snowflake"
+Requires-Dist: snowflake-connector-python; extra == "snowflake"
+Requires-Dist: pyyaml; extra == "snowflake"
+Provides-Extra: sagemaker
+Requires-Dist: boto3<2.0,>=1.30.0; extra == "sagemaker"
+Requires-Dist: mypy-boto3-sagemaker-runtime<2.0,>=1.34.0; extra == "sagemaker"
+Provides-Extra: test-sagemaker
+Requires-Dist: sagemaker<3.0; extra == "test-sagemaker"
+Dynamic: license-file
+Dynamic: requires-dist
+
+<p align="center">
+  <img height="180" src="https://s3.us-west-1.amazonaws.com/data.kumo.ai/img/kumo_pink_md.svg" />
+</p>
+
+______________________________________________________________________
+
+The Kumo SDK implements a pythonic interface for users to programmatically
+interact with the Kumo machine learning platform
+([documentation](https://kumo-ai.github.io/kumo-sdk/docs/#)).
+
+## Installation
+
+The Kumo SDK is available for Python 3.10 to Python 3.13. To install, simply run
+
+```
+pip install kumoai
+```