fitvis.py

gistfile1.txt

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.11"
# dependencies = [
#   "fire",
#   "altair",
#   "polars",
#   "pycrux",
#   "ipython"
# ]
# ///

# Example usage:
# uv run \
#         fitvis.py                                                       \
#         --crux_common="/Users/thomascamminadywahoo/Dev/crux_common/"    \
#         --folder="/Users/thomascamminadywahoo/Dev/fitvis/data/set (5)"  ;

import json
import logging
import os
from typing import Literal

import altair as alt  # ty: ignore[unresolved-import]
import fire  # ty: ignore[unresolved-import]
import polars as pl  # ty: ignore[unresolved-import]
import pycrux  # ty: ignore[unresolved-import]


class AltairConfig:
    @staticmethod
    def setup(
        where: Literal["browser", "html", "jupyter"] = "browser",
        name: str = "my_theme",
        enable: bool = True,
        dark_color: str = "gray",
        font_weight: str = "normal",
        small_font_size: int = 12,
        medium_font_size: int = 14,
        large_font_size: int = 18,
        height: int = 500,
        width: int = 1500,
        spacing_facet: int = 60,
        spacing_concat: int = 60,
    ):
        """Setup Altair."""
        AltairConfig.disable_max_rows()
        AltairConfig.renderers_enable(where=where)
        AltairConfig.register_theme(
            name=name,
            enable=enable,
            dark_color=dark_color,
            font_weight=font_weight,
            small_font_size=small_font_size,
            medium_font_size=medium_font_size,
            large_font_size=large_font_size,
            height=height,
            width=width,
            spacing_facet=spacing_facet,
            spacing_concat=spacing_concat,
        )

    @staticmethod
    def disable_max_rows():
        """Disables the max rows limit for Altair."""
        alt.data_transformers.disable_max_rows()

    @staticmethod
    def renderers_enable(where: Literal["browser", "html", "jupyter"]):
        """Enables Altair renderers."""
        alt.renderers.enable(where)

    @staticmethod
    def register_theme(
        name: str,
        enable: bool,
        dark_color: str,
        font_weight: str,
        small_font_size: int,
        medium_font_size: int,
        large_font_size: int,
        height: int,
        width: int,
        spacing_facet: int,
        spacing_concat: int,
    ):
        """Registers my custom Altair theme."""

        @alt.theme.register(name, enable=enable)
        def loader():
            return json.loads(
                json.dumps(
                    {
                        "config": {
                            "text": {
                                "color": dark_color,
                                "fontSize": small_font_size,
                            },
                            "header": {
                                "titleFontSize": large_font_size,
                                "labelFontSize": medium_font_size,
                                "color": dark_color,
                                "titleColor": dark_color,
                                "labelColor": dark_color,
                                "fontWeight": font_weight,
                                "titleFontWeight": font_weight,
                                "labelFontWeight": font_weight,
                            },
                            "view": {
                                "height": height,
                                "width": width,
                                "strokeWidth": 0,
                                "fill": "white",
                            },
                            "facet": {
                                "spacing": spacing_facet,
                            },
                            "concat": {
                                "spacing": spacing_concat,
                            },
                            "axis": {
                                "domain": True,
                                "domainColor": dark_color,
                                "domainWidth": 1,
                                "gridWidth": 1,
                                "labelAngle": 0,
                                "tickSize": 5,
                                "gridCap": "round",
                                "gridDash": [2, 4],
                                "fontWeight": font_weight,
                                "titleFontWeight": font_weight,
                                "labelFontWeight": font_weight,
                                "titleFontSize": medium_font_size,
                                "labelFontSize": small_font_size,
                                "color": dark_color,
                                "titleColor": dark_color,
                                "labelColor": dark_color,
                                "tickColor": dark_color,
                            },
                            "axisX": {
                                "titleAnchor": "end",
                                "titleAlign": "center",
                            },
                            "axisY": {
                                "titleAnchor": "end",
                                "titleAngle": 0,
                                "titleAlign": "center",
                                "titleY": -15,
                                "titleX": 0,
                            },
                            "legend": {
                                "fontWeight": font_weight,
                                "titleFontWeight": font_weight,
                                "labelFontWeight": font_weight,
                                "titleFontSize": medium_font_size,
                                "labelFontSize": small_font_size,
                                "labelLimit": 0,
                                "color": dark_color,
                                "titleColor": dark_color,
                                "labelColor": dark_color,
                                "tickColor": dark_color,
                                "domainColor": dark_color,
                            },
                            "title": {
                                "anchor": "middle",
                                "fontWeight": font_weight,
                                "titleFontWeight": font_weight,
                                "labelFontWeight": font_weight,
                                "titleFontSize": large_font_size,
                                "labelFontSize": small_font_size,
                                "color": dark_color,
                                "titleColor": dark_color,
                                "labelColor": dark_color,
                                "tickColor": dark_color,
                                "domainColor": dark_color,
                            },
                        }
                    }
                )
            )


def io(folder: str, crux_common: str) -> pl.DataFrame:
    import glob

    _df_list: list[pl.DataFrame] = []
    crux = pycrux.Crux(crux_common)

    files = (
        glob.glob(os.path.join(folder, "*.fit"))
        + glob.glob(os.path.join(folder, "*.FIT"))
        + glob.glob(os.path.join(folder, "*.Fit"))
    )
    for file in files:
        try:
            _df_list.append(
                pl.DataFrame(crux.read_fit(file)).with_columns(
                    file=pl.lit(os.path.basename(file)),
                    folder=pl.lit(folder),
                )
            )
        except Exception as e:
            logging.warning(f"Failed to read {file}: {e}")

    try:
        df = pl.concat(_df_list, how="diagonal_relaxed")
    except Exception as e:
        raise ValueError(f"Failed to create df: {e}") from e

    try:
        # Try to convert "timestamp" which looks like this: 2025-09-04T08:00:46.000Z to datetime
        df = df.with_columns(
            pl.col("timestamp").str.strptime(
                pl.Datetime, format="%Y-%m-%dT%H:%M:%S%.3fZ"
            )
        )
    except Exception as e:
        logging.warning(f"Failed to convert timestamp: {e}")

    try:
        # drop columns that are all null

        def drop_columns_that_are_all_null(_df: pl.DataFrame) -> pl.DataFrame:
            return _df[
                [s.name for s in _df if not (s.null_count() == _df.height)]
            ]

        df = df.pipe(drop_columns_that_are_all_null)
    except Exception as e:
        logging.warning(f"Failed to drop columns that are all null: {e}")
    return df


def plot(df: pl.DataFrame, height: int = 500, width: int = 1500) -> None:
    brush = alt.selection_interval(encodings=["x"], empty=True)
    base = (
        alt.Chart(df)
        .encode(
            color=alt.Color(
                "file:N",
                legend=alt.Legend(
                    orient="none",
                    legendX=100,
                    legendY=height + 40,
                    direction="horizontal",
                    titleAnchor="middle",
                ),
            ).title(None),
        )
        .add_params(brush)
        .properties(width=(width), height=(height), title=df["folder"].unique())
    )

    def _chart(y) -> alt.Chart:
        return alt.hconcat(
            alt.layer(
                base.mark_line().encode(
                    x=alt.X("sec:Q", title="Time (Seconds)"),
                    y=alt.Y(f"{y}:Q", title=y),
                    opacity=alt.when(brush)
                    .then(alt.value(1))
                    .otherwise(alt.value(0.7)),
                ),
                base.mark_rule()
                .encode(y=f"mean({y}):Q")
                .transform_filter(brush),
            ),
            alt.vconcat(
                base.mark_text(align="left", dx=0, fontSize=14)
                .encode(
                    x=alt.X("file:O").title(None).axis(None),
                    y=alt.value(0),
                    text=alt.Text(f"mean({y}):Q", format=".2f"),
                )
                .properties(
                    title=alt.Title(f"Mean ({y})", align="center"),
                    width=300,
                    height=10,
                )
                .transform_filter(brush),
                base.mark_text(align="left", dx=0, fontSize=14)
                .encode(
                    x=alt.X("file:O").title(None).axis(None),
                    y=alt.value(0),
                    text=alt.Text(f"stdev({y}):Q", format=".2f"),
                )
                .properties(
                    title=alt.Title(
                        f"Standard Deviation ({y})", align="center"
                    ),
                    width=300,
                    height=10,
                )
                .transform_filter(brush),
                base.mark_text(align="left", dx=0, fontSize=14)
                .encode(
                    x=alt.X("file:O").title(None).axis(None),
                    y=alt.value(0),
                    text=alt.Text(f"min({y}):Q", format=".2f"),
                )
                .properties(
                    title=alt.Title(f"Minimum ({y})", align="center"),
                    width=300,
                    height=10,
                )
                .transform_filter(brush),
                base.mark_text(align="left", dx=0, fontSize=14)
                .encode(
                    x=alt.X("file:O").title(None).axis(None),
                    y=alt.value(0),
                    text=alt.Text(f"max({y}):Q", format=".2f"),
                )
                .properties(
                    title=alt.Title(f"Maximum ({y})", align="center"),
                    width=300,
                    height=10,
                )
                .transform_filter(brush),
                spacing=20,
            ),
            spacing=100,
        )

    alt.vconcat(
        _chart("pwr_watts"),
        _chart("cad_rpm"),
        spacing=250,
    ).resolve_scale(color="independent").show()


def main(folder: str, crux_common: str):
    AltairConfig.setup()

    df = io(folder, crux_common)
    plot(df)


if __name__ == "__main__":
    fire.Fire(main)