tabmat 4.3.0__cp315-cp315-win_amd64.whl

tabmat/__init__.py

@@ -0,0 +1,40 @@
+"""""" # start delvewheel patch
+def _delvewheel_patch_1_13_1():
+    import os
+    if os.path.isdir(libs_dir := os.path.abspath(os.path.join(os.path.dirname(__file__), os.pardir, 'tabmat.libs'))):
+        os.add_dll_directory(libs_dir)
+
+
+_delvewheel_patch_1_13_1()
+del _delvewheel_patch_1_13_1
+# end delvewheel patch
+
+import importlib.metadata
+
+from .categorical_matrix import CategoricalMatrix
+from .constructor import from_csc, from_df, from_formula, from_pandas
+from .dense_matrix import DenseMatrix
+from .matrix_base import MatrixBase
+from .sparse_matrix import SparseMatrix
+from .split_matrix import SplitMatrix, as_tabmat, hstack
+from .standardized_mat import StandardizedMatrix
+
+try:
+    __version__ = importlib.metadata.distribution(__name__).version
+except Exception:
+    __version__ = "unknown"
+
+__all__ = [
+    "DenseMatrix",
+    "MatrixBase",
+    "StandardizedMatrix",
+    "SparseMatrix",
+    "SplitMatrix",
+    "CategoricalMatrix",
+    "from_csc",
+    "from_formula",
+    "from_pandas",
+    "from_df",
+    "as_tabmat",
+    "hstack",
+]

tabmat/benchmark/generate_matrices.py

@@ -0,0 +1,144 @@
+import os
+import pickle
+
+import click
+import numpy as np
+from scipy import sparse as sps
+
+import tabmat as tm
+
+
+def make_dense_matrices(n_rows: int, n_cols: int) -> dict:
+    """Make dense matrices for benchmarks."""
+    dense_matrices = {"numpy_C": np.random.random((n_rows, n_cols))}
+    dense_matrices["numpy_F"] = dense_matrices["numpy_C"].copy(order="F")
+    assert dense_matrices["numpy_F"].flags["F_CONTIGUOUS"]
+    dense_matrices["tabmat"] = tm.DenseMatrix(dense_matrices["numpy_C"])  # type: ignore
+    return dense_matrices
+
+
+def make_cat_matrix(n_rows: int, n_cats: int) -> tm.CategoricalMatrix:
+    """Make categorical matrix for benchmarks."""
+    mat = tm.CategoricalMatrix(np.random.choice(np.arange(n_cats, dtype=int), n_rows))
+    return mat
+
+
+def make_cat_matrix_all_formats(n_rows: int, n_cats: int) -> dict:
+    """Make categorical matrix with all formats for benchmarks."""
+    mat = make_cat_matrix(n_rows, n_cats)
+    d = {
+        "tabmat": mat,
+        "scipy.sparse csr": mat.tocsr(),
+    }
+    d["scipy.sparse csc"] = d["scipy.sparse csr"].tocsc()
+    return d
+
+
+def make_cat_matrices(n_rows: int, n_cat_cols_1: int, n_cat_cols_2: int) -> dict:
+    """Make two categorical matrices for benchmarks."""
+    two_cat_matrices = {
+        "tabmat": tm.SplitMatrix(
+            [
+                make_cat_matrix(n_rows, n_cat_cols_1),
+                make_cat_matrix(n_rows, n_cat_cols_2),
+            ]
+        )
+    }
+    two_cat_matrices["scipy.sparse csr"] = sps.hstack(
+        [elt.tocsr() for elt in two_cat_matrices["tabmat"].matrices]
+    )
+    two_cat_matrices["scipy.sparse csc"] = two_cat_matrices["scipy.sparse csr"].tocsc()  # type: ignore
+    return two_cat_matrices
+
+
+def make_dense_cat_matrices(
+    n_rows: int, n_dense_cols: int, n_cats_1: int, n_cats_2: int
+) -> dict:
+    """Make dense categorical matrices for benchmarks."""
+    dense_block = np.random.random((n_rows, n_dense_cols))
+    two_cat_matrices = [
+        make_cat_matrix(n_rows, n_cats_1),
+        make_cat_matrix(n_rows, n_cats_2),
+    ]
+    dense_cat_matrices = {
+        "tabmat": tm.SplitMatrix(two_cat_matrices + [tm.DenseMatrix(dense_block)]),
+        "scipy.sparse csr": sps.hstack(
+            [elt.tocsr() for elt in two_cat_matrices] + [sps.csr_matrix(dense_block)]
+        ),
+    }
+    dense_cat_matrices["scipy.sparse csc"] = dense_cat_matrices[
+        "scipy.sparse csr"
+    ].tocsc()
+    return dense_cat_matrices
+
+
+def make_sparse_matrices(n_rows: int, n_cols: int) -> dict:
+    """Make sparse matrices for benchmarks."""
+    mat = sps.random(n_rows, n_cols).tocsc()
+    matrices = {
+        "scipy.sparse csc": mat,
+        "scipy.sparse csr": mat.tocsr(),
+        "tabmat": tm.SparseMatrix(mat),
+    }
+    return matrices
+
+
+def _get_matrix_path(name):
+    return f"benchmark/data/{name}_data.pkl"
+
+
+def get_all_benchmark_matrices():
+    """Get all matrices used in benchmarks."""
+    return {
+        "dense": lambda: make_dense_matrices(int(4e6), 10),
+        "sparse": lambda: make_sparse_matrices(int(4e5), int(1e2)),
+        "sparse_narrow": lambda: make_sparse_matrices(int(3e6), 3),
+        "sparse_wide": lambda: make_sparse_matrices(int(4e4), int(1e4)),
+        "one_cat": lambda: make_cat_matrix_all_formats(int(1e6), int(1e5)),
+        "two_cat": lambda: make_cat_matrices(int(1e6), int(1e3), int(1e3)),
+        "dense_cat": lambda: make_dense_cat_matrices(int(3e6), 5, int(1e3), int(1e3)),
+        "dense_smallcat": lambda: make_dense_cat_matrices(int(3e6), 5, 10, int(1e3)),
+    }
+
+
+# TODO: duplication with glm_benchmarks
+def get_comma_sep_names(xs: str):
+    """Return comma separated names from names in input string."""
+    return [x.strip() for x in xs.split(",")]
+
+
+def get_matrix_names():
+    """Return names for benchmark_matrices."""
+    return ",".join(get_all_benchmark_matrices().keys())
+
+
+@click.command()
+@click.option(
+    "--matrix_name",
+    type=str,
+    help=(
+        f"Specify a comma-separated list of matrices you want to build. "
+        f"Leaving this blank will default to building all matrices. "
+        f"Matrix options: {get_matrix_names()}"
+    ),
+)
+def generate_matrices(matrix_name: str) -> None:
+    """Generate example matrices for benchmarks."""
+    all_benchmark_matrices = get_all_benchmark_matrices()
+
+    if matrix_name is None:
+        benchmark_matrices = list(all_benchmark_matrices.keys())
+    else:
+        benchmark_matrices = get_comma_sep_names(matrix_name)
+
+    for name in benchmark_matrices:
+        f = all_benchmark_matrices[name]
+        mats = f()
+        save_path = _get_matrix_path(name)
+        os.makedirs(os.path.dirname(save_path), exist_ok=True)
+        with open(save_path, "wb") as fname:
+            pickle.dump(mats, fname)
+
+
+if __name__ == "__main__":
+    generate_matrices()

tabmat/benchmark/main.py

@@ -0,0 +1,383 @@
+import pickle
+import time
+from typing import Union
+
+import click
+import numpy as np
+import pandas as pd
+from scipy import sparse as sps
+
+import tabmat as tm
+from tabmat.benchmark.generate_matrices import (
+    get_all_benchmark_matrices,
+    get_comma_sep_names,
+    get_matrix_names,
+    make_cat_matrices,
+    make_cat_matrix_all_formats,
+    make_dense_cat_matrices,
+    make_dense_matrices,
+    make_sparse_matrices,
+)
+from tabmat.benchmark.memory_tools import track_peak_mem
+
+
+def _sandwich(mat: Union[tm.MatrixBase, np.ndarray, sps.csc_matrix], vec: np.ndarray):
+    if isinstance(mat, (tm.MatrixBase, tm.StandardizedMatrix)):
+        mat.sandwich(vec)
+    elif isinstance(mat, np.ndarray):
+        (mat * vec[:, None]).T @ mat
+    else:
+        mat.T @ sps.diags(vec) @ mat
+    return
+
+
+def _transpose_matvec(
+    mat: Union[tm.MatrixBase, np.ndarray, sps.csc_matrix], vec: np.ndarray
+):
+    if isinstance(mat, (tm.MatrixBase, tm.StandardizedMatrix)):
+        return mat.transpose_matvec(vec)
+    else:
+        return mat.T.dot(vec)
+
+
+def _matvec(mat, vec: np.ndarray) -> np.ndarray:
+    if isinstance(mat, (tm.MatrixBase, tm.StandardizedMatrix)):
+        return mat.matvec(vec)  # type: ignore
+    else:
+        return mat.dot(vec)
+
+
+def _lvec_setup(matrices):
+    return (np.random.random(next(iter(matrices.values())).shape[0]),)
+
+
+def _rvec_setup(matrices):
+    return (np.random.random(next(iter(matrices.values())).shape[1]),)
+
+
+ops = {
+    "matvec": (_rvec_setup, _matvec),
+    "transpose-matvec": (_lvec_setup, _transpose_matvec),
+    "sandwich": (_lvec_setup, _sandwich),
+}
+
+
+def get_op_names():
+    """Get names of operations."""
+    return ",".join(ops.keys())
+
+
+def run_one_benchmark_set(
+    matrices: dict[
+        str, Union[tm.MatrixBase, tm.StandardizedMatrix, np.ndarray, sps.spmatrix]
+    ],
+    include_baseline: bool,
+    name: str,
+    standardized: bool,
+    ops_to_run,
+    n_iterations: int,
+    bench_memory: bool,
+) -> pd.DataFrame:
+    """Run a single round of benchmarks."""
+    if not include_baseline:
+        for k in list(matrices.keys()):
+            if k != "tabmat":
+                del matrices[k]
+
+    if standardized:
+
+        def _to_standardized_mat(mat):
+            if isinstance(mat, tm.MatrixBase):
+                return tm.StandardizedMatrix(mat, np.zeros(mat.shape[1]))
+            print(
+                f"""For benchmarking a {type(mat)}, the baseline matrix will not
+                be standardized."""
+            )
+            return mat
+
+        matrices = {k: _to_standardized_mat(v) for k, v in matrices.items()}
+
+    times = pd.DataFrame(
+        index=pd.MultiIndex.from_product(
+            [ops_to_run, matrices.keys()],
+            names=["operation", "storage"],
+        ),
+        columns=["memory", "time"],
+    ).reset_index()
+
+    for i, row in times.iterrows():
+        mat_ = matrices[row["storage"]]
+        setup_fnc, op_fnc = ops[row["operation"]]
+        setup_data = setup_fnc(matrices)
+        runtimes = []
+        peak_mems = []
+        for _ in range(n_iterations):
+            start = time.time()
+            if bench_memory:
+                peak_mem = track_peak_mem(op_fnc, mat_, *setup_data)
+            else:
+                op_fnc(mat_, *setup_data)
+                peak_mem = 0
+            end = time.time()
+            peak_mems.append(peak_mem)
+            runtimes.append(end - start)
+
+        # We want to get a consistent measure of runtime so we take the
+        # minimum. Any increase in runtime is due to warmup or other
+        # processes running at the same time.
+        times.loc[i, "time"] = np.min(runtimes)
+
+        # On the other hand, we want the maximum memory usage because this
+        # metric is isolated to our current python process. Any lower
+        # values will be because the highest memory usage was "missed" by
+        # the tracker
+        times.loc[i, "memory"] = np.max(peak_mems)
+
+    times["design"] = name
+    return times
+
+
+@click.command()
+@click.option(
+    "--operation_name",
+    type=str,
+    help=(
+        f"Specify a comma-separated list of operations you want to run. Leaving this "
+        f"blank will default to running all operations. Operation options: "
+        f"{get_op_names()}"
+    ),
+)
+@click.option(
+    "--matrix_name",
+    type=str,
+    help=(
+        f"Specify a comma-separated list of matrices you want to run or specify. "
+        f"Leaving this blank will default to running all predefined matrices. "
+        f"Matrix options: {get_matrix_names()} OR custom. If custom, specify details "
+        f"using additional custom matrix options. See --dense, --sparse, --one_cat, "
+        f"--two_cat, and --dense_cat options for more details"
+    ),
+)
+@click.option(
+    "--dense",
+    nargs=2,
+    multiple=True,
+    help=(
+        "Specify n_rows, n_cols for custom dense matrix. "
+        "Only used if 'custom' included in matrix_name."
+    ),
+    default=None,
+)
+@click.option(
+    "--sparse",
+    nargs=2,
+    multiple=True,
+    help=(
+        "Specify n_rows, n_cols for custom sparse matrix. "
+        "Only used if 'custom' included in matrix_name."
+    ),
+    default=None,
+)
+@click.option(
+    "--one_cat",
+    nargs=2,
+    multiple=True,
+    help=(
+        "Specify n_rows, n_cols for custom one_cat matrix. "
+        "Only used if 'custom' included in matrix_name."
+    ),
+    default=None,
+)
+@click.option(
+    "--two_cat",
+    nargs=3,
+    multiple=True,
+    help=(
+        "Specify n_rows, n_cols for custom two_cat matrix. "
+        "Only used if 'custom' included in matrix_name."
+    ),
+    default=None,
+)
+@click.option(
+    "--dense_cat",
+    nargs=4,
+    multiple=True,
+    help=(
+        "Specify n_rows, n_cols for custom dense_cat matrix. "
+        "Only used if 'custom' included in matrix_name."
+    ),
+    default=None,
+)
+@click.option(
+    "--bench_memory",
+    type=bool,
+    is_flag=True,
+    help=(
+        "Should we benchmark memory usage with tracemalloc. Turning this on will make "
+        "the runtime benchmarks less useful due to memory benchmarking overhead. "
+        "Also, when memory benchmarking is on, debuggers like pdb and ipdb seem to "
+        "fail."
+    ),
+    default=False,
+)
+@click.option(
+    "--n_iterations",
+    type=int,
+    help=(
+        "How many times to re-run the benchmark. The maximum memory usage and minimum "
+        "runtime will be reported. Higher numbers of iterations reduce noise. This "
+        "defaults to 100 unless memory benchmarking is turned on in which case it will "
+        "be 1."
+    ),
+    default=None,
+)
+@click.option(
+    "--include_baseline",
+    type=bool,
+    is_flag=True,
+    help="Should we include a numpy/scipy baseline performance benchmark.",
+    default=False,
+)
+@click.option(
+    "--standardized",
+    type=bool,
+    is_flag=True,
+    help="Should we test with a tabmat.StandardizedMatrix?",
+    default=False,
+)
+def run_all_benchmarks(
+    operation_name: str,
+    matrix_name: str,
+    dense: list,
+    sparse: list,
+    one_cat: list,
+    two_cat: list,
+    dense_cat: list,
+    bench_memory: bool,
+    n_iterations: int,
+    include_baseline: bool,
+    standardized: bool,
+):
+    """
+    Usage examples.
+
+    python benchmark/main.py --operation_name matvec,transpose-matvec --matrix_name sparse --include_baseline\n
+              operation           storage memory         time
+
+    0            matvec  scipy.sparse csc      0   0.00129819\n
+    1            matvec  scipy.sparse csr      0   0.00266385\n
+    2            matvec  tabmat      0   0.00199628\n
+    3  transpose-matvec  scipy.sparse csc      0  0.000838518\n
+    4  transpose-matvec  scipy.sparse csr      0   0.00239468\n
+    5  transpose-matvec  tabmat      0  0.000296116\n
+
+    python benchmark/main.py --operation_name sandwich --matrix_name dense_cat --bench_memory\n
+
+      operation           storage    memory      time\n
+    0  sandwich  tabmat  52244505  0.159682\n
+
+    python benchmark/main.py --operation_name matvec --matrix_name custom --sparse 3e6 1 --sparse 3e6 10 --dense 10 10\n
+    operation           storage memory      time                            design \n
+    0    matvec  tabmat      0  0.000006  dense, #rows:10, #cols:10      \n
+    operation           storage memory      time                            design \n
+    0    matvec  tabmat      0  0.046355  sparse, #rows:3000000, #cols:1 \n
+    operation           storage memory      time                            design \n
+    0    matvec  tabmat      0  0.048141  sparse, #rows:3000000, #cols:10\n
+    """  # noqa
+    if n_iterations is None:
+        if bench_memory:
+            n_iterations = 1
+        else:
+            n_iterations = 100
+
+    if operation_name is None:
+        ops_to_run = list(ops.keys())
+    else:
+        ops_to_run = get_comma_sep_names(operation_name)
+
+    all_benchmark_matrices = get_all_benchmark_matrices()
+
+    benchmark_matrices = {}
+    if matrix_name is None:
+        for k in all_benchmark_matrices.keys():
+            with open(f"benchmark/data/{k}_data.pkl", "rb") as f:
+                benchmark_matrices[k] = pickle.load(f)
+
+    elif "custom" in matrix_name:
+        if dense:
+            for params in dense:
+                n_rows, n_cols = (int(float(x)) for x in params)
+                benchmark_matrices[f"dense, #rows:{n_rows}, #cols:{n_cols}"] = (
+                    make_dense_matrices(n_rows, n_cols)
+                )
+        if sparse:
+            for params in sparse:
+                n_rows, n_cols = (int(float(x)) for x in params)
+                benchmark_matrices[f"sparse, #rows:{n_rows}, #cols:{n_cols}"] = (
+                    make_sparse_matrices(n_rows, n_cols)
+                )
+        if one_cat:
+            for params in one_cat:
+                n_rows, n_cats = (int(float(x)) for x in params)
+                benchmark_matrices[f"one_cat, #rows:{n_rows}, #cats:{n_cats}"] = (
+                    make_cat_matrix_all_formats(n_rows, n_cats)
+                )
+        if two_cat:
+            for params in two_cat:
+                n_rows, n_cat_cols_1, n_cat_cols_2 = (int(float(x)) for x in params)
+                benchmark_matrices[
+                    f"two_cat #rows:{n_rows}, #cats_1:{n_cat_cols_1}, "
+                    f"#cats_2:{n_cat_cols_2}"
+                ] = make_cat_matrices(n_rows, n_cat_cols_1, n_cat_cols_2)
+        if dense_cat:
+            for params in dense_cat:
+                n_rows, n_dense_cols, n_cat_cols_1, n_cat_cols_2 = (
+                    int(float(x)) for x in params
+                )
+                benchmark_matrices[
+                    f"dense_cat #rows:{n_rows}, #dense:{n_dense_cols}, "
+                    f" cats_1:{n_cat_cols_1}, #cats_2:{n_cat_cols_2}"
+                ] = make_dense_cat_matrices(
+                    n_rows, n_dense_cols, n_cat_cols_1, n_cat_cols_2
+                )
+    else:
+        for k in get_comma_sep_names(matrix_name):
+            with open(f"benchmark/data/{k}_data.pkl", "rb") as f:
+                benchmark_matrices[k] = pickle.load(f)
+
+    for name, matrices in benchmark_matrices.items():
+        time_bench = run_one_benchmark_set(
+            matrices,
+            include_baseline,
+            name,
+            standardized,
+            ops_to_run,
+            n_iterations,
+            False,
+        )
+
+        if bench_memory:
+            memory_bench = run_one_benchmark_set(
+                matrices,
+                include_baseline,
+                name,
+                standardized,
+                ops_to_run,
+                1,
+                True,
+            )
+
+            full_bench = pd.merge(
+                memory_bench[["operation", "storage", "memory"]],
+                time_bench[["operation", "storage", "time", "design"]],
+                on=["operation", "storage"],
+            )
+        else:
+            full_bench = time_bench[["operation", "storage", "time", "design"]]
+        print(full_bench)
+
+        full_bench.to_csv(f"benchmark/data/{name}_bench.csv", index=False)
+
+
+if __name__ == "__main__":
+    run_all_benchmarks()

tabmat/benchmark/memory_tools.py

@@ -0,0 +1,48 @@
+import time
+import tracemalloc
+from threading import Thread
+
+
+class MemoryPoller:
+    """
+    Example usage.
+
+    with MemoryPoller() as mp:
+        do some stuff here
+        print('initial memory usage', mp.initial_memory)
+        print('max memory usage', mp.max_memory)
+        excess_memory_used = mp.max_memory - mp.initial_memory
+    """
+
+    def _poll_max_memory_usage(self):
+        while not self.stop_polling:
+            self.snapshots.append(tracemalloc.take_snapshot())
+            time.sleep(1e-3)
+
+    def __enter__(self):
+        tracemalloc.start()
+        self.stop_polling = False
+        self.snapshots = [tracemalloc.take_snapshot()]
+        self.t = Thread(target=self._poll_max_memory_usage)
+        self.t.start()
+        return self
+
+    def __exit__(self, *excargs):
+        self.stop_polling = True
+        self.t.join()
+        self.final_usage, self.peak_usage = tracemalloc.get_traced_memory()
+        tracemalloc.stop()
+
+
+def track_peak_mem(f, *args, **kwargs):
+    """Track peak memory. Used in benchmarks to track memory used during matrix
+    operations.
+    """
+    with MemoryPoller() as mp:
+        f(*args, **kwargs)
+    for s in mp.snapshots:
+        top_stats = s.statistics("lineno")
+        print("[ Top 4 ]")
+        for stat in top_stats[:4]:
+            print(stat)
+    return mp.peak_usage