A tree-based (non-overlapping) solver

tests/test_fit.py

@@ -14,7 +14,7 @@
     explain_timeseries,
 )
 from wise_pizza.segment_data import SegmentData
-from wise_pizza.solver import solve_lasso, solve_lp
+from wise_pizza.solve.solver import solve_lasso, solve_lp
 from wise_pizza.time import create_time_basis
 from wise_pizza.plotting_time import plot_time
 
@@ -33,7 +33,7 @@
 # Too long, delete some values for quick starts, e.g. by deleting the parameters in nan_percent, size_one_percent
 deltas_test_values = [
     ("totals", "split_fits", "force_dim", "extra_dim"),  # how
-    ("lp", "lasso"),  # solver
+    ("lp", "lasso", "tree"),  # solver
     (True,),  # plot_is_static
     (explain_changes_in_average, explain_changes_in_totals),  # function
     (0.0, 90.0),  # nan_percent
@@ -44,7 +44,7 @@
 
 # possible values for explain_levels
 levels_test_values = [
-    ("lp", "lasso"),  # solver
+    ("lp", "lasso", "tree"),  # solver
     (0.0, 90.0),  # nan_percent
     (0.0, 90.0),  # size_one_percent
 ]
@@ -136,9 +136,9 @@ def test_categorical():
     print("yay!")
 
 
-@pytest.mark.parametrize("nan_percent", [0.0, 1.0])
-def test_synthetic_template(nan_percent: float):
-    all_data = synthetic_data(init_len=1000)
+@pytest.mark.parametrize("nan_percent, clustering", [[0.0, False], [1.0, False]])
+def test_synthetic_template(nan_percent: float, clustering: bool):
+    all_data = synthetic_data(init_len=10000, dim_values=5)
     data = all_data.data
 
     data.loc[(data["dim0"] == 0) & (data["dim1"] == 1), "totals"] += 100
@@ -155,6 +155,7 @@ def test_synthetic_template(nan_percent: float):
         min_segments=5,
         verbose=1,
         solver="lp",
+        cluster_values=clustering,
     )
     print("***")
     for s in sf.segments:
@@ -167,6 +168,38 @@ def test_synthetic_template(nan_percent: float):
     print("yay!")
 
 
+@pytest.mark.parametrize("nan_percent", [0.0, 1.0])
+def test_synthetic_template_tree(nan_percent: float):
+    all_data = synthetic_data(init_len=1000)
+    data = all_data.data
+
+    data.loc[(data["dim0"] == 0) & (data["dim1"] == 1), "totals"] += 200
+    data.loc[(data["dim1"] == 0) & (data["dim2"] == 1), "totals"] += 300
+
+    if nan_percent > 0:
+        data = values_to_nan(data, nan_percent)
+    sf = explain_levels(
+        data,
+        dims=all_data.dimensions,
+        total_name=all_data.segment_total,
+        size_name=all_data.segment_size,
+        max_depth=2,
+        min_segments=5,
+        verbose=1,
+        solver="tree",
+    )
+    print("***")
+    for s in sf.segments:
+        print(s)
+
+    # TODO: insert approppriate asserts
+    # assert abs(sf.segments[0]["coef"] - 300) < 2
+    # assert abs(sf.segments[1]["coef"] - 100) < 2
+
+    # sf.plot()
+    print("yay!")
+
+
 @pytest.mark.parametrize("nan_percent", [0.0, 1.0])
 def test_synthetic_ts_template(nan_percent: float):
     all_data = synthetic_ts_data(init_len=10000)

wise_pizza/cluster.py

@@ -1,3 +1,6 @@
+from typing import List, Dict, Tuple
+from collections import defaultdict
+
 import numpy as np
 import pandas as pd
 from sklearn.preprocessing import PowerTransformer
@@ -18,17 +21,27 @@ def guided_kmeans(X: np.ndarray, power_transform: bool = True) -> np.ndarray:
         X = X.values
 
     if power_transform:
-        if len(X[X > 0] > 1):
-            X[X > 0] = PowerTransformer(standardize=False).fit_transform(X[X > 0].reshape(-1, 1)).reshape(-1)
-        if len(X[X < 0] > 1):
-            X[X < 0] = -PowerTransformer(standardize=False).fit_transform(-X[X < 0].reshape(-1, 1)).reshape(-1)
+        if len(X[X > 0]) > 1:
+            X[X > 0] = (
+                PowerTransformer(standardize=False)
+                .fit_transform(X[X > 0].reshape(-1, 1))
+                .reshape(-1)
+            )
+        if len(X[X < 0]) > 1:
+            X[X < 0] = (
+                -PowerTransformer(standardize=False)
+                .fit_transform(-X[X < 0].reshape(-1, 1))
+                .reshape(-1)
+            )
 
     best_score = -1
     best_labels = None
     best_n = -1
     # If we allow 2 clusters, it almost always just splits positive vs negative - boring!
     for n_clusters in range(3, int(len(X) / 2) + 1):
-        cluster_labels = KMeans(n_clusters=n_clusters, init="k-means++", n_init=10).fit_predict(X)
+        cluster_labels = KMeans(
+            n_clusters=n_clusters, init="k-means++", n_init=10
+        ).fit_predict(X)
         score = silhouette_score(X, cluster_labels)
         # print(n_clusters, score)
         if score > best_score:
@@ -45,3 +58,55 @@ def to_matrix(labels: np.ndarray) -> np.ndarray:
     for i in labels.unique():
         out[labels == i, i] = 1.0
     return out
+
+
+def make_clusters(dim_df: pd.DataFrame, dims: List[str]):
+    cluster_names = {}
+    for dim in dims:
+        if len(dim_df[dim].unique()) >= 6:  # otherwise what's the point in clustering?
+            grouped_df = (
+                dim_df[[dim, "totals", "weights"]].groupby(dim, as_index=False).sum()
+            )
+            grouped_df["avg"] = grouped_df["totals"] / grouped_df["weights"]
+            grouped_df["cluster"], _ = guided_kmeans(grouped_df["avg"])
+            pre_clusters = (
+                grouped_df[["cluster", dim]]
+                .groupby("cluster")
+                .agg({dim: lambda x: "@@".join(x)})
+                .values
+            )
+            # filter out clusters with only one element
+            these_clusters = [c for c in pre_clusters.reshape(-1) if "@@" in c]
+            # create short cluster names
+            for i, c in enumerate(these_clusters):
+                cluster_names[f"{dim}_cluster_{i + 1}"] = c
+    return cluster_names
+
+
+def nice_cluster_names(x: List[Dict[str, List[str]]]) -> Tuple[List[Dict], Dict]:
+    # first pass just populate cluster names
+    cluster_strings = defaultdict(set)
+    for xx in x:
+        for dim, v in xx.items():
+            if len(v) > 1:
+                cluster_strings[dim].add("@@".join(v))
+
+    cluster_names = {}
+    reverse_cluster_names = {}
+    for dim, clusters in cluster_strings.items():
+        reverse_cluster_names[dim] = {}
+        for i, c in enumerate(clusters):
+            cluster_names[f"{dim}_cluster_{i + 1}"] = c
+            reverse_cluster_names[dim][c] = f"{dim}_cluster_{i + 1}"
+
+    col_defs = []
+    for xx in x:
+        this_def = {}
+        for dim, v in xx.items():
+            if len(v) > 1:
+                this_def[dim] = reverse_cluster_names[dim]["@@".join(v)]
+            else:
+                this_def[dim] = v[0]
+        col_defs.append(this_def)
+
+    return col_defs, cluster_names

wise_pizza/explain.py

@@ -361,6 +361,7 @@ def explain_timeseries(
     max_depth: int = 2,
     solver: str = "omp",
     verbose: bool = False,
+    constrain_signs: bool = False,
     cluster_values: bool = False,
     time_basis: Optional[pd.DataFrame] = None,
     fit_log_space: bool = False,
@@ -388,7 +389,10 @@ def explain_timeseries(
             fit_sizes = True
 
     if fit_log_space:
-        tf = LogTransform(offset=1, weight_pow_sc=log_space_weight_sc)
+        tf = LogTransform(
+            offset=1,
+            weight_pow_sc=log_space_weight_sc,
+        )
     else:
         tf = IdentityTransform()
 
@@ -415,6 +419,7 @@ def explain_timeseries(
             max_depth=max_depth,
             solver=solver,
             verbose=verbose,
+            constrain_signs=constrain_signs,
             cluster_values=cluster_values,
             time_basis=time_basis,
         )
@@ -441,6 +446,7 @@ def explain_timeseries(
         max_depth=max_depth,
         solver=solver,
         verbose=verbose,
+        constrain_signs=constrain_signs,
         cluster_values=cluster_values,
         time_basis=time_basis,
     )
@@ -477,6 +483,7 @@ def explain_timeseries(
         max_depth=max_depth,
         solver=solver,
         verbose=verbose,
+        constrain_signs=constrain_signs,
         cluster_values=cluster_values,
         time_basis=time_basis,
     )