lf-pollywog 0.1.1__py3-none-any.whl

pollywog/display.py

@@ -0,0 +1,149 @@
+_DISPLAY_THEME = "light"
+
+def set_theme(theme):
+    """
+    Set the global theme for display_calcset ("light" or "dark").
+    """
+    global _DISPLAY_THEME
+    if theme not in ("light", "dark"):
+        raise ValueError("Theme must be 'light' or 'dark'.")
+    _DISPLAY_THEME = theme
+
+def display_calcset(calcset, theme=None, colors=None, display_output=True):
+    """
+    Display a CalcSet in a Jupyter notebook with a visual style similar to Leapfrog, rendering equations and logic blocks visually.
+    Supports 'theme' ("light" or "dark") and custom color palettes via 'colors' dict. If theme is None, uses global setting from set_theme().
+    """
+    from IPython.display import display, HTML
+    import html
+
+    # Default color palettes
+    default_colors = {
+        "light": {
+            "background": "#eee",
+            "text": "#222",
+            "variable": "#0057b7",
+            "label": "#222",
+            "if": "#0057b7",
+            "arrow": "#222",
+            "comment": "#999",
+            "var_ref": "#b77",
+        },
+        "dark": {
+            "background": "#222",
+            "text": "#eee",
+            "variable": "#7abaff",
+            "label": "#eee",
+            "if": "#7abaff",
+            "arrow": "#eee",
+            "comment": "#bbb",
+            "var_ref": "#ffb77a",
+        },
+    }
+    use_theme = theme if theme is not None else _DISPLAY_THEME
+    palette = default_colors.get(use_theme, default_colors["light"]).copy()
+    if colors:
+        palette.update(colors)
+
+    def render_expression(expr, indent=0):
+        pad = " " * (indent * 4)
+        if isinstance(expr, str):
+            expr = html.escape(expr)
+            expr = expr.replace('[', f'<span style="color:{palette["var_ref"]};">[').replace(']', ']</span>')
+            return pad + f'<span style="color:{palette["text"]};">{expr}</span>'
+        elif isinstance(expr, list):
+            return '<br>'.join(render_expression(e, indent) for e in expr)
+        if isinstance(expr, dict):
+            typ = expr.get("type")
+            if typ == "if":
+                rows = expr.get("rows", [])
+                otherwise = expr.get("otherwise", {}).get("children", [])
+                html_rows = []
+                for row in rows:
+                    cond = row.get("test", {}).get("children", [])
+                    res = row.get("result", {}).get("children", [])
+                    html_rows.append(
+                        f'<div style="margin-left:{indent*24}px;border-left:2px solid {palette["if"]};padding-left:8px;">'
+                        f'<span style="color:{palette["if"]};">if</span> '
+                        f'{render_expression(cond, 0)} '
+                        f'<span style="color:{palette["arrow"]};">&rarr;</span> '
+                        f'{render_expression(res, indent+1)}'
+                        f'</div>'
+                    )
+                if otherwise:
+                    html_rows.append(
+                        f'<div style="margin-left:{indent*24}px;border-left:2px solid {palette["if"]};padding-left:8px;">'
+                        f'<span style="color:{palette["if"]};">otherwise</span> '
+                        f'<span style="color:{palette["arrow"]};">&rarr;</span> '
+                        f'{render_expression(otherwise, indent+1)}'
+                        f'</div>'
+                    )
+                return ''.join(html_rows)
+            elif typ == "if_row":
+                cond = expr.get("test", {}).get("children", [])
+                res = expr.get("result", {}).get("children", [])
+                return (
+                    f'<div style="margin-left:{indent*24}px;border-left:2px solid {palette["if"]};padding-left:8px;">'
+                    f'<span style="color:{palette["if"]};">if</span> '
+                    f'{render_expression(cond, 0)} '
+                    f'<span style="color:{palette["arrow"]};">&rarr;</span> '
+                    f'{render_expression(res, indent+1)}'
+                    f'</div>'
+                )
+            elif typ == "list":
+                children = expr.get("children", [])
+                return render_expression(children, indent)
+            else:
+                return pad + f'<span style="color:{palette["text"]};">{html.escape(str(expr))}</span>'
+        else:
+            return pad + f'<span style="color:{palette["text"]};">{html.escape(str(expr))}</span>'
+
+    def render_equation(eq):
+        if isinstance(eq, dict) and eq.get("type") == "equation":
+            statement = eq["statement"]
+            comment = eq.get("comment", "")
+            expr_html = render_expression(statement)
+            comment_html = f'<span style="color:#999;">{html.escape(comment)}</span>' if comment else ''
+            return f'<div style="margin-left:1em;color:#555;">{expr_html} {comment_html}</div>'
+        return html.escape(str(eq))
+
+    def render_item(item):
+        d = item.to_dict() if hasattr(item, "to_dict") else item
+        name = d.get("name", "")
+        typ = d.get("type", "")
+        eq = d.get("equation", None)
+        comment = d.get("comment", "")
+        html_block = f'<div style="margin-bottom:0.5em;">'
+        html_block += f'<b style="color:#0057b7;">{html.escape(name)}</b> '
+        # Show calculation_type for calculation items
+        calc_type = d.get("calculation_type")
+        label = typ
+        if typ == "calculation" and calc_type:
+            label = calc_type
+        html_block += f'<span style="background:#eee;border-radius:4px;padding:2px 6px;color:#222;">{html.escape(label)}</span>'
+        if eq:
+            html_block += render_equation(eq)
+        if comment:
+            html_block += f'<div style="color:#999;margin-left:1em;">{html.escape(comment)}</div>'
+        html_block += '</div>'
+        return html_block
+
+    def section(title, items):
+        if not items:
+            return ''
+        html_items = ''.join(render_item(item) for item in items)
+        return f'<details open><summary style="font-size:1.1em;font-weight:bold;color:{palette["variable"]};">{title}</summary>{html_items}</details>'
+
+    variables = [i for i in calcset.items if getattr(i, 'item_type', None) == 'variable']
+    calculations = [i for i in calcset.items if getattr(i, 'item_type', None) == 'calculation']
+    filters = [i for i in calcset.items if getattr(i, 'item_type', None) == 'filter']
+
+    html_out = f'<div style="font-family:sans-serif;max-width:900px;color:{palette["text"]};">'
+    html_out += section('Variables', variables)
+    html_out += section('Calculations', calculations)
+    html_out += section('Filters', filters)
+    html_out += '</div>'
+    if display_output:
+        display(HTML(html_out))
+    else:
+        return html_out

pollywog/helpers.py

@@ -0,0 +1,246 @@
+from .core import If, IfRow, Number, Category
+from .utils import ensure_variables
+
+
+def Sum(*variables, name=None, comment=None):
+    """
+    Create a Number representing the sum of the given variables.
+
+    Args:
+        *variables: Variable names (as strings) to sum, e.g. "Au", "Ag", or a single list of variable names as strings.
+        name (str, optional): Name for the output variable. If None, defaults to "sum_<var1>_<var2>_..."
+        comment (str, optional): Optional comment for the calculation.
+
+    Returns:
+        Number: A pollywog Number representing the sum calculation.
+
+    Example:
+        >>> Sum("Au", "Ag", name="sum_Au_Ag")
+    """
+    if not variables:
+        raise ValueError("At least one variable must be provided.")
+    if len(variables) == 1 and isinstance(variables[0], (list, tuple)):
+        variables = variables[0]
+    if name is None:
+        name = "sum_" + "_".join(variables)
+    expr = f"({' + '.join(f'[{v}]' for v in variables)})"
+    return Number(
+        name, [expr], comment_equation=comment or f"Sum of {', '.join(variables)}"
+    )
+
+
+def Product(*variables, name=None, comment=None):
+    """
+    Create a Number representing the product of the given variables.
+
+    Args:
+        *variables: Variable names (as strings) to multiply, e.g. "Au", "Ag", or a single list of variable names as strings.
+        name (str, optional): Name for the output variable. If None, defaults to "prod_<var1>_<var2>_..."
+        comment (str, optional): Optional comment for the calculation.
+
+    Returns:
+        Number: A pollywog Number representing the product calculation.
+
+    Example:
+        >>> Product("Au", "Ag", name="prod_Au_Ag")
+    """
+    if not variables:
+        raise ValueError("At least one variable must be provided.")
+    if len(variables) == 1 and isinstance(variables[0], (list, tuple)):
+        variables = variables[0]
+    if name is None:
+        name = "prod_" + "_".join(variables)
+    expr = f"({' * '.join(f'[{v}]' for v in variables)})"
+    return Number(
+        name, [expr], comment_equation=comment or f"Product of {', '.join(variables)}"
+    )
+
+
+def Normalize(variable, min_value, max_value, name=None, comment=None):
+    """
+    Create a Number that normalizes a variable to [0, 1] given min and max values.
+
+    Args:
+        variable (str): Variable name to normalize.
+        min_value (float): Minimum value for normalization.
+        max_value (float): Maximum value for normalization.
+        name (str, optional): Name for the output variable. If None, defaults to "norm_<variable>".
+        comment (str, optional): Optional comment for the calculation.
+
+    Returns:
+        Number: A pollywog Number representing the normalization calculation.
+
+    Example:
+        >>> Normalize("Au", 0, 10, name="norm_Au")
+    """
+    if name is None:
+        name = f"norm_{variable}"
+    expr = f"([{variable}] - {min_value}) / ({max_value} - {min_value})"
+    return Number(
+        name,
+        [expr],
+        comment_equation=comment
+        or f"Normalize {variable} to [0, 1] using min={min_value}, max={max_value}",
+    )
+
+
+def Average(*variables, name=None, comment=None):
+    """
+    Create a Number representing the average of the given variables.
+
+    Args:
+        *variables: Variable names (as strings) to average, e.g. "Au", "Ag", or a single list of variable names as strings.
+        name (str, optional): Name for the output variable. If None, defaults to "avg_<var1>_<var2>_..."
+        comment (str, optional): Optional comment for the calculation.
+
+    Returns:
+        Number: A pollywog Number representing the average calculation.
+
+    Example:
+        >>> Average("Au", "Ag", name="avg_Au_Ag")
+    """
+    if not variables:
+        raise ValueError("At least one variable must be provided.")
+    if len(variables) == 1 and isinstance(variables[0], list):
+        variables = variables[0]
+    if name is None:
+        name = "avg_" + "_".join(variables)
+    expr = f"({' + '.join(f'[{v}]' for v in variables)}) / {len(variables)}"
+    return Number(
+        name, [expr], comment_equation=comment or f"Average of {', '.join(variables)}"
+    )
+
+
+def WeightedAverage(variables, weights, name=None, comment=None):
+    """
+    Create a Number representing the weighted average of variables.
+
+    Args:
+        variables (list of str): Variable names to average, e.g. ["Au", "Ag", "Cu"]
+        weights (list of float or string): Corresponding weights for each variable, either constant values ( e.g. [0.5, 0.3, 0.2]) or variable names (e.g. ["w1", "w2", "w3"]).
+        name (str, optional): Name for the output variable. If None, defaults to "wavg_<var1>_<var2>_..."
+        comment (str, optional): Optional comment for the calculation.
+
+    Returns:
+        Number: A pollywog Number representing the weighted average calculation.
+
+    Example:
+        >>> WeightedAverage(["Au", "Ag"], [0.7, 0.3], name="wavg_Au_Ag")
+    """
+    if not variables or not weights or len(variables) != len(weights):
+        raise ValueError("variables and weights must be non-empty and of equal length.")
+    if name is None:
+        name = "wavg_" + "_".join(variables)
+    weights = ensure_variables(weights)
+    sum_weights = " + ".join(weights)
+    weighted_terms = [f"[{v}] * {w}" for v, w in zip(variables, weights)]
+    expr = f"({' + '.join(weighted_terms)}) / ({sum_weights})"
+    return Number(
+        name,
+        [expr],
+        comment_equation=comment
+        or f"Weighted average of {', '.join(variables)} with weights {weights}",
+    )
+
+
+def Scale(variable, factor, name=None, comment=None):
+    """
+    Create a Number that multiplies a variable by a factor.
+
+    Args:
+        variable (str): Variable name to scale.
+        factor (float or str): Scaling factor (can be a constant or another variable).
+        name (str, optional): Name for the output variable. If None, defaults to "scale_<variable>".
+        comment (str, optional): Optional comment for the calculation.
+
+    Returns:
+        Number: A pollywog Number representing the scaled variable.
+
+    Example:
+        >>> Scale("Au", 2, name="Au_scaled")
+    """
+    if name is None:
+        name = f"scale_{variable}"
+    factor_expr = f"[{factor}]" if isinstance(factor, str) else str(factor)
+    expr = f"[{variable}] * {factor_expr}"
+    return Number(
+        name, [expr], comment_equation=comment or f"Scale {variable} by {factor}"
+    )
+
+
+def IfElse(condition, then, else_, name=None, comment=None, output_type=None):
+    """
+    Create a Number or Category with conditional logic (if/else).
+
+    Args:
+        condition (str): Condition expression (e.g., "[Au] > 1").
+        then: Value or expression if condition is true.
+        else_: Value or expression if condition is false.
+        name (str, optional): Name for the output variable.
+        comment (str, optional): Optional comment for the calculation.
+        output_type (type, optional): Number or Category (default: Number).
+
+    Returns:
+        Number or Category: A pollywog item with conditional logic.
+
+    Example:
+        >>> IfElse("[Au] > 1", "High", "Low", name="Au_class")
+    """
+    if output_type is None:
+        output_type = Number
+    if name is None:
+        name = "ifelse"
+    if isinstance(then, (int, float)):
+        then_expr = str(then)
+    else:
+        then_expr = then
+    if isinstance(else_, (int, float)):
+        else_expr = str(else_)
+    else:
+        else_expr = else_
+    if_block = If([IfRow([condition], [then_expr])], otherwise=[else_expr])
+    return output_type(
+        name,
+        [if_block],
+        comment_equation=comment or f"If {condition} then {then_expr} else {else_expr}",
+    )
+
+
+def CategoryFromThresholds(variable, thresholds, categories, name=None, comment=None):
+    """
+    Create a Category assigning labels based on value thresholds.
+
+    Args:
+        variable (str): Variable to threshold.
+        thresholds (list of float): Threshold values (must be sorted ascending).
+        categories (list of str): Category labels (len(categories) == len(thresholds) + 1).
+        name (str, optional): Name for the output category.
+        comment (str, optional): Optional comment for the calculation.
+
+    Returns:
+        Category: A pollywog Category assigning labels based on thresholds.
+
+    Example:
+        >>> CategoryFromThresholds("Au", [0.5, 1.0], ["Low", "Medium", "High"], name="Au_class")
+    """
+    if len(categories) != len(thresholds) + 1:
+        raise ValueError("categories must have one more element than thresholds")
+    rows = []
+    prev = None
+    for i, threshold in enumerate(thresholds):
+        if prev is None:
+            cond = f"[{variable}] <= {threshold}"
+        else:
+            cond = f"([{variable}] > {prev}) and ([{variable}] <= {threshold})"
+        rows.append(([cond], [categories[i]]))
+        prev = threshold
+    # Otherwise case
+    otherwise = [categories[-1]]
+    if_block = If([IfRow(cond, val) for cond, val in rows], otherwise=otherwise)
+    if name is None:
+        name = f"class_{variable}"
+    return Category(
+        name,
+        [if_block],
+        comment_equation=comment or f"Classify {variable} by thresholds {thresholds}",
+    )

pollywog/run.py

@@ -0,0 +1,116 @@
+import copy
+from pollywog.core import CalcSet, If, IfRow
+import re
+
+
+def run_calcset(
+    calcset, inputs=None, dataframe=None, assign_results=True, output_variables=False
+):
+    """
+    Evaluate a CalcSet with external inputs or a pandas DataFrame.
+    Returns a dict of results (if inputs provided) or a DataFrame (if dataframe provided).
+    Pandas is only required if using DataFrame input/output.
+    By default, only calculations, categories, and filters are output (Leapfrog-like).
+    Set output_variables=True to include variables in output (for debugging).
+    """
+
+    # Helper to evaluate an expression or If object
+    def eval_expr(expr, context):
+        if isinstance(expr, str):
+            if not expr.strip():
+                return None
+
+            # Replace [var] with context["var"] using regex
+            def repl(m):
+                var = m.group(1)
+                return f"context[{repr(var)}]"
+
+            expr_eval = re.sub(r"\[([^\]]+)\]", repl, expr)
+            try:
+                return eval(expr_eval, {"context": context}, context)
+            except Exception:
+                return None
+        elif isinstance(expr, If):
+            for row in expr.rows:
+                cond = eval_expr(row.condition[0], context) if row.condition else True
+                if cond:
+                    return eval_expr(row.value[0], context)
+            if expr.otherwise:
+                return eval_expr(expr.otherwise[0], context)
+            return None
+        elif isinstance(expr, IfRow):
+            # Should not be evaluated directly, only as part of If
+            return None
+        else:
+            return expr
+
+    # Dependency resolution
+    sorted_items = calcset.topological_sort().items
+
+    def run_single(context):
+        results = {}
+        for item in sorted_items:
+            # If item is a Variable, assign its value from context or inputs directly
+            if getattr(item, "item_type", None) == "variable":
+                results[item.name] = context.get(item.name, None)
+                continue
+            child_results = []
+            for child in item.children:
+                child_results.append(eval_expr(child, {**context, **results}))
+            results[item.name] = child_results[0] if child_results else None
+        # Filter output according to output_variables flag
+        item_type_map = {
+            item.name: getattr(item, "item_type", None) for item in sorted_items
+        }
+        if not output_variables:
+            return {
+                k: v for k, v in results.items() if item_type_map.get(k) != "variable"
+            }
+        return results
+
+    if dataframe is not None:
+        try:
+            import pandas as pd
+        except ImportError:
+            raise ImportError(
+                "pandas is required for DataFrame input/output. Please install pandas or use dict inputs."
+            )
+        df = dataframe.copy()
+        for idx, row in df.iterrows():
+            context = dict(row)
+            results = run_single(context)
+            for k, v in results.items():
+                df.at[idx, k] = v
+        # Remove variable columns if output_variables is False
+        if not output_variables:
+            variable_names = [
+                item.name
+                for item in sorted_items
+                if getattr(item, "item_type", None) == "variable"
+            ]
+            df = df.drop(columns=variable_names, errors="ignore")
+        return df
+    else:
+        context = inputs if inputs is not None else {}
+        return run_single(context)
+
+
+# Pandas DataFrame extension accessor
+try:
+    import pandas as pd
+
+    @pd.api.extensions.register_dataframe_accessor("pw")
+    class PollywogAccessor:
+        def __init__(self, pandas_obj):
+            self._obj = pandas_obj
+
+        def run(self, calcset, assign_results=True):
+            """
+            Run a CalcSet on this DataFrame, returning a copy with results assigned.
+            """
+            return run_calcset(
+                calcset, dataframe=self._obj, assign_results=assign_results
+            )
+
+except ImportError:
+    pass

pollywog/utils.py

@@ -0,0 +1,95 @@
+def ensure_list(x):
+    """
+    Ensure the input is a list. If not, wrap it in a list.
+
+    Args:
+        x (Any): Input value or list.
+    Returns:
+        list: The input as a list.
+    """
+    if not isinstance(x, list):
+        x = [x]
+    return x
+
+
+def ensure_str_list(x):
+    """
+    Ensure the input is a list of strings, padding with empty strings if needed.
+
+    Args:
+        x (Any): Input value or list.
+    Returns:
+        list: List of strings, padded with empty strings if necessary.
+    """
+    if not isinstance(x, list):
+        x = [x]
+    if isinstance(x, list):
+        if not isinstance(x[0], str):
+            x = [""] + x
+        if not isinstance(x[-1], str):
+            x = x + [""]
+    return x
+
+
+def to_dict(items, guard_strings=False):
+    """
+    Convert a list of items to their dictionary representations if possible.
+
+    Args:
+        items (list): List of items or single item.
+        guard_strings (bool): If True, pad with empty strings if first/last are not strings.
+    Returns:
+        list: List of dicts or items, possibly padded with empty strings.
+    """
+    out = [
+        item.to_dict() if hasattr(item, "to_dict") else item
+        for item in ensure_list(items)
+    ]
+    if guard_strings:
+        if not isinstance(out[0], str):
+            out = [""] + out
+        if not isinstance(out[-1], str):
+            out = out + [""]
+    return out
+
+def is_number(v):
+    """
+    Check if the input can be converted to a float (i.e., is a number).
+
+    Args:
+        v (Any): Input value.
+    Returns:
+        bool: True if input is a number, False otherwise.
+    """
+    try:
+        float(v)
+        return True
+    except (ValueError, TypeError):
+        return False
+    
+def ensure_brackets(var):
+    """
+    Ensure the variable name is wrapped in brackets [var].
+
+    Args:
+        var (str): Variable name.
+    Returns:
+        str: Variable name wrapped in brackets.
+    """
+    var = var.strip()
+    if not (var.startswith("[") and var.endswith("]")):
+        var = f"[{var}]"
+    return var
+
+def ensure_variables(variables):
+    """
+    Ensures that each item in the input is formatted as a variable.
+    For each item in `variables`, if the item is a number, it is converted to a string.
+    Otherwise, it is passed to `ensure_brackets` to ensure proper bracket formatting.
+    Args:
+        variables (Any): A single variable or a list of variables to be processed.
+    Returns:
+        list: A list of formatted variable strings.
+    """
+
+    return [f"{v}" if is_number(v) else ensure_brackets(v) for v in ensure_list(variables)]

tests/conftest.py

@@ -0,0 +1,4 @@
+import sys
+import os
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

tests/test_conversion.py

@@ -0,0 +1,72 @@
+import pytest
+from pollywog.conversion.sklearn import convert_tree, convert_linear_model
+from pollywog.core import Number, Category
+from sklearn.tree import DecisionTreeRegressor, DecisionTreeClassifier
+from sklearn.linear_model import LinearRegression
+import numpy as np
+
+
+def make_regressor():
+    X = np.array([[0, 0], [1, 1], [2, 2]])
+    y = np.array([0, 1, 2])
+    return DecisionTreeRegressor().fit(X, y)
+
+
+def make_linear():
+    X = np.array([[0, 0], [1, 0], [0, 1]])
+    y = np.array([1, 3, 1])
+    lm = LinearRegression().fit(X, y)
+    lm.coef_ = np.array([2.0, 0.0])
+    lm.intercept_ = 1.0
+    return lm
+
+
+class DummyTree:
+    class tree_:
+        feature = [0, -2, -2]
+        threshold = [1.5, 0, 0]
+        children_left = [1, -1, -1]
+        children_right = [2, -1, -1]
+        value = [[[0]], [[1]], [[2]]]
+
+    class _tree:
+        class Dummy:
+            pass
+
+
+def test_convert_tree_regressor():
+    tree = make_regressor()
+    result = convert_tree(tree, ["x1", "x2"], "target")
+    assert isinstance(result, Number)
+    assert result.name == "target"
+    assert "Converted from DecisionTreeRegressor" in result.comment_equation
+
+
+def test_convert_tree_classifier():
+    X = np.array([[0, 0], [1, 1], [2, 2]])
+    y = np.array([0, 1, 2])
+    tree = DecisionTreeClassifier().fit(X, y)
+    result = convert_tree(tree, ["x1", "x2"], "target")
+    assert isinstance(result, Category)
+    assert result.name == "target"
+    assert "Converted from DecisionTreeClassifier" in result.comment_equation
+
+
+def test_convert_tree_invalid():
+    class Dummy:
+        pass
+
+    dummy = Dummy()
+    dummy.__class__ = type("UnknownTree", (), {})
+    with pytest.raises(Exception):
+        convert_tree(dummy, ["x1"], "target")
+
+
+def test_convert_linear_model():
+    lm = make_linear()
+    result = convert_linear_model(lm, ["x1", "x2"], "target")
+    assert isinstance(result, Number)
+    assert result.name == "target"
+    assert "Converted from LinearRegression" in result.comment_equation
+    assert "1.000000" in result.children[0]
+    assert "2.000000 * [x1]" in result.children[0]