API Reference¶

def demo_frame(rows: int = 1_500, seed: int = 42) -> pl.DataFrame:
    """Daily NAVs for a handful of made-up funds, on deliberately different scales.

    They share a market factor, so the lines rhyme without being copies, and they
    start anywhere from 1 to 1,450 — which is exactly the case that needs indexing
    before two of them can be read on one axis.

    The dates come first, so the frame is ready for :func:`jointview.plot.line_chart`
    as it stands, and ``seed`` makes it the same frame every time:

    >>> frame = demo_frame(rows=10)
    >>> frame.columns[:2]
    ['date', 'world_equity']
    >>> frame.height
    10
    """
    rng = np.random.default_rng(seed)
    market = rng.normal(0.0004, 0.011, rows)

    navs = {
        name: start * np.cumprod(1.0 + drift + beta * market + rng.normal(0.0, wobble, rows))
        for name, (start, beta, drift, wobble) in FUNDS.items()
    }
    return pl.DataFrame({"date": _business_days(date(2020, 1, 1), rows), **navs})

def load_frame(path: str | Path | None) -> pl.DataFrame:
    """Read a DataFrame from ``path``, or build the demo frame when it is None.

    The suffix picks the reader, so no path means the generated frame rather than
    an error — which is what makes ``marimo run app.py`` work with no arguments:

    >>> load_frame(None).columns[0]
    'date'

    Whatever the format, a column of dates arrives as dates: the text formats have no
    way to record that one ever was, so they are parsed back on the way in. A column
    that does not read as dates all the way down is left as the text it is.

    A path that is not there is reported before a reader is chosen, so the message
    names the file rather than complaining about its extension:

    >>> load_frame("nowhere.parquet")
    Traceback (most recent call last):
        ...
    FileNotFoundError: no such file: nowhere.parquet

    A directory is named as one. This is also where an empty ``Path`` lands, since
    ``Path("")`` is ``Path(".")`` — the two are the same object by the time anything
    here sees them, so the current directory is what actually arrived:

    >>> load_frame(".")
    Traceback (most recent call last):
        ...
    IsADirectoryError: not a file: .
    """
    # Only the empty *string* is the "no path" sentinel: it is what marimo hands over
    # for a flag that was not passed. A Path cannot carry it — see the docstring.
    if path is None or path == "":
        return demo_frame()

    file = Path(path).expanduser()
    if not file.exists():
        raise FileNotFoundError(f"no such file: {file}")  # noqa: TRY003

    # Before the suffix lookup, which would otherwise reject a directory for having
    # the wrong extension — and an empty one for having no name to quote at all.
    if file.is_dir():
        raise IsADirectoryError(f"not a file: {file}")  # noqa: TRY003

    reader = READERS.get(file.suffix.lower())
    if reader is None:
        supported = ", ".join(sorted(READERS))
        raise ValueError(f"cannot read {file.suffix or file.name!r}; supported: {supported}")  # noqa: TRY003

    return reader(file)

def aligned(frame: pl.DataFrame, a: str, b: str) -> pl.DataFrame:
    """The two series on their common sample: ``period``, ``a``, ``b``, in order.

    Both the picture and the summary tables are built from this, so the numbers
    beside the chart always describe the lines in it. Renaming also sidesteps
    Vega-Lite's field-shorthand escaping and lets ``a`` and ``b`` be the same column.

    The three names are fixed whatever the columns were called, the rows come out
    sorted by period, and a date where either series is missing is not part of the
    sample — here the frame arrives unsorted and with a gap on the 2nd:

    >>> import datetime as dt, polars as pl
    >>> frame = pl.DataFrame(
    ...     {
    ...         "date": [dt.date(2024, 1, 3), dt.date(2024, 1, 1), dt.date(2024, 1, 2)],
    ...         "x": [3.0, 1.0, None],
    ...         "y": [30.0, 10.0, 20.0],
    ...     }
    ... )
    >>> pair = aligned(frame, "x", "y")
    >>> pair.columns
    ['period', 'a', 'b']
    >>> pair["a"].to_list()
    [1.0, 3.0]
    """
    for column in (a, b):
        if column not in frame.columns:
            raise KeyError(f"no column {column!r} in frame")  # noqa: TRY003
        if not frame.schema[column].is_numeric():
            raise TypeError(f"column {column!r} is {frame.schema[column]}, which cannot be drawn")  # noqa: TRY003

    date = date_column(frame)
    period = pl.col(date).alias(PERIOD) if date else pl.int_range(pl.len()).alias(PERIOD)
    data = frame.select(period, pl.col(a).alias("a"), pl.col(b).alias("b"))
    return data.drop_nulls().sort(PERIOD)

def date_column(frame: pl.DataFrame) -> str | None:
    """The first temporal column, which becomes the x-axis. None means row number.

    >>> import datetime as dt, polars as pl
    >>> date_column(pl.DataFrame({"when": [dt.date(2024, 1, 1)], "nav": [1.0]}))
    'when'
    >>> date_column(pl.DataFrame({"nav": [1.0]})) is None
    True
    """
    return next((name for name, dtype in frame.schema.items() if dtype.is_temporal()), None)

def default_pair(frame: pl.DataFrame) -> tuple[int, int]:
    """Indices into :func:`series_columns` to open on — the first two series.

    A frame with a single series opens on it twice, rather than refusing to draw.

    >>> import polars as pl
    >>> default_pair(pl.DataFrame({"a": [1.0], "b": [2.0]}))
    (0, 1)
    >>> default_pair(pl.DataFrame({"only": [1.0]}))
    (0, 0)
    """
    names = series_columns(frame)
    if not names:
        raise ValueError("frame has no numeric columns to plot")  # noqa: TRY003
    return 0, 1 if len(names) > 1 else 0

def series_columns(frame: pl.DataFrame) -> list[str]:
    """The columns that can be drawn: every numeric one.

    >>> import datetime as dt, polars as pl
    >>> frame = pl.DataFrame({"date": [dt.date(2024, 1, 1)], "nav": [1.0], "label": ["a"]})
    >>> series_columns(frame)
    ['nav']
    """
    return [name for name, dtype in frame.schema.items() if dtype.is_numeric()]

def line_chart(
    frame: pl.DataFrame,
    a: str,
    b: str,
    *,
    rebase: bool = True,
    base: float = BASE,
    width: int | str = "container",
    height: int | str = 700,
    max_points: int = MAX_POINTS,
) -> alt.LayerChart:
    """Draw columns ``a`` and ``b`` of ``frame`` as two lines against time.

    Four layers over one plotting area — the crosshair, the lines, the hover markers
    and the end labels — handed back as a plain Altair chart, so nothing here needs
    marimo to draw it:

    >>> import polars as pl
    >>> frame = pl.DataFrame({"cash": [1.0, 1.01, 1.02], "balanced": [1450.0, 1479.0, 1465.0]})
    >>> chart = line_chart(frame, "cash", "balanced")
    >>> type(chart).__name__
    'LayerChart'
    >>> len(chart.to_dict()["layer"])
    4

    Width defaults to ``"container"``: the plot takes whatever the column around it
    gives it, which is the point of a full-width app. Height stays a number, because
    nothing in the page has a height for a chart to follow — 700 fills a laptop window
    once the notebook margins are out of the way, without spilling off a short one.

    Asking to rebase a pair that cannot be indexed draws the raw levels, and the
    y-axis says ``level`` rather than claiming otherwise — see :func:`_rebasable`.
    The title is read back out of the compiled spec, because that is the only place it
    exists; layer 1 is :func:`_lines`, the series themselves:

    >>> chart.to_dict()["layer"][1]["encoding"]["y"]["title"]
    'indexed to 100'
    >>> pnl = pl.DataFrame({"strategy": [0.0, 5.0, 3.0], "benchmark": [0.0, 2.0, 4.0]})
    >>> line_chart(pnl, "strategy", "benchmark").to_dict()["layer"][1]["encoding"]["y"]["title"]
    'level'
    """
    wide, rebased = _wide(frame, a, b, rebase=rebase, base=base, max_points=max_points)
    names = [name for name in wide.columns if name != PERIOD]

    date = date_column(frame)
    x_type: XType = "temporal" if date else "quantitative"
    x_title = date or "row"
    # `rebased`, not `rebase`: the title names what the numbers underneath actually are.
    y_title = f"indexed to {base:g}" if rebased else "level"
    x = alt.X(PERIOD, type=x_type, title=x_title)

    # Made here rather than inside a layer because two of them share it: the crosshair
    # carries the parameter, the markers only read it.
    hover = _hover()
    lines = _lines(wide, names, x, y_title)
    return (
        alt.layer(
            _crosshair(wide, names, x, x_type, x_title, hover),
            lines,
            _markers(lines, hover),
            _end_labels(wide, names, x),
        )
        .resolve_scale(color="shared")
        .configure_axis(grid=True, gridOpacity=0.3, domain=False, labelPadding=4, tickSize=4)
        .configure_view(stroke=None)
        .configure_legend(labelFontSize=12)
        .properties(
            # One plotting area for all four layers, sized at the top level so a
            # container width is measured once rather than per layer.
            width=width,
            height=height,
            # Only the right margin earns its keep: it is where the end labels go.
            padding={"left": 0, "top": 0, "bottom": 0, "right": 76},
            # "pad", the default, grows the figure past the box it was given and puts
            # the gutter back; fitting spends the padding out of the size instead.
            autosize=_autosize(width, height),
        )
    )

def line_frame(
    frame: pl.DataFrame,
    a: str,
    b: str,
    *,
    rebase: bool = True,
    base: float = BASE,
    max_points: int = MAX_POINTS,
) -> pl.DataFrame:
    """What the chart draws: one ``period`` column and one column per named series.

    Wide rather than long, because the crosshair reads every series at the hovered
    period out of a single row.

    >>> import polars as pl
    >>> frame = pl.DataFrame({"cash": [1.0, 1.01, 1.02], "balanced": [1450.0, 1479.0, 1465.0]})
    >>> drawn = line_frame(frame, "cash", "balanced")
    >>> drawn.columns
    ['period', 'cash', 'balanced']

    Rebasing is what lets those two share a y-axis at all: both leave the first
    period at ``base``, whatever they were priced at.

    >>> round(drawn["cash"][0], 6), round(drawn["balanced"][0], 6)
    (100.0, 100.0)

    A column against itself is one line rather than two identical ones:

    >>> line_frame(frame, "cash", "cash").columns
    ['period', 'cash']

    A series starting at zero cannot be indexed — a cumulative P&L curve starts there
    by construction — so the pair keeps its own levels instead:

    >>> pnl = pl.DataFrame({"strategy": [0.0, 5.0, 3.0], "benchmark": [0.0, 2.0, 4.0]})
    >>> line_frame(pnl, "strategy", "benchmark")["strategy"].to_list()
    [0.0, 5.0, 3.0]
    """
    wide, _ = _wide(frame, a, b, rebase=rebase, base=base, max_points=max_points)
    return wide

def drawdown(frame: pl.DataFrame, column: str, *, date_col: str = PERIOD) -> pl.Series:
    """Distance below the running maximum, as a fraction — zero or negative.

    jQuantStats reports the same quantity as a positive depth, so the sign is flipped
    here: a drawdown is a fall, and every other rate in this module carries its
    direction in its sign.

    >>> import polars as pl
    >>> frame = pl.DataFrame({"period": [0, 1, 2], "nav": [100.0, 120.0, 60.0]})
    >>> round(drawdown(frame, "nav").min(), 4)
    -0.5
    """
    stats = _stats(frame, column, date_col=date_col)
    return -stats.drawdown()[column]

def metrics(frame: pl.DataFrame, column: str, *, date_col: str = PERIOD, rf: float = 0.0) -> dict[str, float]:
    """The raw numbers behind the summary table, in natural units (0.07 is 7%).

    A figure that cannot be formed — a Sharpe ratio for a flat series, a growth rate
    for a series that touches zero — comes back as NaN or infinity rather than
    raising, so one odd column never blanks the whole table.

    Raises:
        ValueError: if there are too few observations to derive anything.

    >>> import polars as pl
    >>> frame = pl.DataFrame({"period": [0, 1, 2, 3], "nav": [100.0, 110.0, 105.0, 120.0]})
    >>> numbers = metrics(frame, "nav")
    >>> numbers["Observations"]
    4.0
    >>> round(numbers["Total return"], 4)
    0.2
    """
    stats = _stats(frame, column, date_col=date_col, rf=rf)
    levels = frame.get_column(column).drop_nulls().cast(pl.Float64)

    numbers: dict[str, float] = {
        "Observations": float(levels.len()),
        "Start": float(levels[0]),
        "End": float(levels[-1]),
    }
    for label, name in STATISTICS.items():
        numbers[label] = _number(stats, name, column)
    return numbers

def returns(frame: pl.DataFrame, column: str, *, date_col: str = PERIOD) -> pl.Series:
    """Simple period-over-period returns of a level series.

    >>> import polars as pl
    >>> frame = pl.DataFrame({"period": [0, 1, 2], "nav": [100.0, 110.0, 99.0]})
    >>> [round(value, 4) for value in returns(frame, "nav")]
    [0.1, -0.1]
    """
    stats = _stats(frame, column, date_col=date_col)
    return stats.returns[column]

def summary(frame: pl.DataFrame, column: str, *, date_col: str = PERIOD, rf: float = 0.0) -> pl.DataFrame:
    """The same numbers as :func:`metrics`, formatted for display.

    >>> import polars as pl
    >>> frame = pl.DataFrame({"period": [0, 1, 2, 3], "nav": [100.0, 110.0, 105.0, 120.0]})
    >>> table = summary(frame, "nav")
    >>> table.columns
    ['metric', 'value']
    >>> table.row(by_predicate=pl.col("metric") == "Total return")[1]
    '+20.00%'
    """
    numbers = metrics(frame, column, date_col=date_col, rf=rf)
    return pl.DataFrame(
        {
            "metric": list(numbers),
            "value": [_format(name, value) for name, value in numbers.items()],
        }
    )

def summary_markdown(
    frame: pl.DataFrame,
    column: str,
    *,
    title: str | None = None,
    date_col: str = PERIOD,
    rf: float = 0.0,
) -> str:
    """A two-column markdown table, ready for ``mo.md``.

    >>> import polars as pl
    >>> frame = pl.DataFrame({"period": [0, 1, 2, 3], "nav": [100.0, 110.0, 105.0, 120.0]})
    >>> print(summary_markdown(frame, "nav", title="fund").splitlines()[0])
    | fund | |
    """
    table = summary(frame, column, date_col=date_col, rf=rf)
    header = f"| {title or column} | |", "|:---|---:|"
    body = (f"| {row['metric']} | {row['value']} |" for row in table.iter_rows(named=True))
    return "\n".join((*header, *body))

def main(argv: list[str] | None = None) -> int:
    """Run `marimo run` (or `edit`) on the packaged notebook, and return its exit code."""
    ours, marimo_args = _split(list(sys.argv[1:] if argv is None else argv))

    parser = _parser()
    args = parser.parse_args(ours)
    app_args = _app_args(args, parser)

    # `python -m marimo`, not a bare `marimo`: under uvx the two need not be the same
    # interpreter, and only this one is sure to have jointview importable — which the
    # notebook needs on its first cell.
    command = [
        sys.executable,
        "-m",
        "marimo",
        "edit" if args.edit else "run",
        # In front of the notebook path, where `marimo run [OPTIONS] NAME` wants them;
        # after it they would be read as arguments to the notebook.
        *marimo_args,
        str(APP),
    ]
    if app_args:
        command += ["--", *app_args]

    try:
        # No shell, and argv is a list, so nothing here is word-split or glob-expanded:
        # the executable is this interpreter, the notebook path is the installed
        # wheel's, and the rest are the user's own arguments on their own machine —
        # the same ones they would have typed after `marimo run`.
        return subprocess.call(command)  # noqa: S603 # nosec B603
    except KeyboardInterrupt:
        # Ctrl-C reached the child too; it has already said its goodbyes.
        return 130