upgini 1.2.156.dev5__tar.gz → 1.2.157.dev2__tar.gz

{upgini-1.2.156.dev5 → upgini-1.2.157.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: upgini
-Version: 1.2.156.dev5
+Version: 1.2.157.dev2
 Summary: Intelligent data search & enrichment for Machine Learning
 Project-URL: Bug Reports, https://github.com/upgini/upgini/issues
 Project-URL: Homepage, https://upgini.com/

upgini-1.2.157.dev2/src/upgini/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.2.157.dev2"

{upgini-1.2.156.dev5 → upgini-1.2.157.dev2}/src/upgini/autofe/binary.py

@@ -6,6 +6,7 @@ import numpy as np
 import pandas as pd
 from jarowinkler import jarowinkler_similarity
 
+from upgini.autofe.operand import OperandKind, OperandValue
 from upgini.autofe.operator import PandasOperator, VectorizableMixin
 
 
@@ -15,7 +16,9 @@ class Min(PandasOperator):
     is_symmetrical: bool = True
     has_symmetry_importance: bool = True
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         return np.minimum(left, right)
 
 
@@ -25,7 +28,9 @@ class Max(PandasOperator):
     is_symmetrical: bool = True
     has_symmetry_importance: bool = True
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         return np.maximum(left, right)
 
 
@@ -37,7 +42,9 @@ class Add(PandasOperator, VectorizableMixin):
     has_symmetry_importance: bool = True
     is_vectorizable: bool = True
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         return left + right
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -56,7 +63,9 @@ class Subtract(PandasOperator, VectorizableMixin):
     has_symmetry_importance: bool = True
     is_vectorizable: bool = True
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         return left - right
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -75,7 +84,9 @@ class Multiply(PandasOperator, VectorizableMixin):
     has_symmetry_importance: bool = True
     is_vectorizable: bool = True
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         return left * right
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -94,7 +105,9 @@ class Divide(PandasOperator, VectorizableMixin):
     is_vectorizable: bool = True
     output_type: Optional[str] = "float"
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         return left / right.replace(0, np.nan)
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -111,7 +124,9 @@ class Combine(PandasOperator):
     has_symmetry_importance: bool = True
     output_type: Optional[str] = "object"
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         temp = left.astype(str) + "_" + right.astype(str)
         temp[left.isna() | right.isna()] = np.nan
         return pd.Series(temp, index=left.index)
@@ -126,7 +141,9 @@ class CombineThenFreq(PandasOperator):
     is_distribution_dependent: bool = True
     input_type: Optional[str] = "discrete"
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         temp = left.astype(str) + "_" + right.astype(str)
         temp[left.isna() | right.isna()] = np.nan
         value_counts = temp.value_counts(normalize=True)
@@ -140,120 +157,40 @@ class Distance(PandasOperator):
     is_symmetrical: bool = True
     has_symmetry_importance: bool = True
 
-    def _to_float_vector(self, value) -> Optional[np.ndarray]:
-        if value is None or (isinstance(value, float) and np.isnan(value)):
-            return None
-        return np.atleast_1d(np.asarray(value, dtype=np.float64))
-
-    def _stack_vectors(self, arr: np.ndarray, indices: np.ndarray) -> Optional[np.ndarray]:
-        try:
-            stacked = np.vstack(arr[indices])
-        except (ValueError, TypeError):
-            return None
-        if stacked.ndim != 2 or stacked.shape[1] == 0:
-            return None
-        if stacked.dtype != np.float64:
-            stacked = stacked.astype(np.float64, copy=False)
-        return stacked
-
-    def _cosine_distance_stacked(self, left_mat: np.ndarray, right_mat: np.ndarray) -> np.ndarray:
-        if left_mat.shape != right_mat.shape:
-            raise ValueError("Left and right matrices must have the same shape")
-        dots = np.sum(left_mat * right_mat, axis=1)
-        denom = np.sqrt(np.sum(left_mat * left_mat, axis=1)) * np.sqrt(np.sum(right_mat * right_mat, axis=1))
-        results = np.full(len(left_mat), np.nan, dtype=np.float64)
-        valid = denom > 0
-        results[valid] = 1 - dots[valid] / denom[valid]
-        return results
-
-    def _try_fast_distance(
-        self, left_arr: np.ndarray, right_arr: np.ndarray, row_present: np.ndarray, n: int
-    ) -> Optional[np.ndarray]:
-        valid_idx = np.flatnonzero(row_present)
-        if len(valid_idx) == 0:
-            return None
-
-        left_mat = self._stack_vectors(left_arr, valid_idx)
-        right_mat = self._stack_vectors(right_arr, valid_idx)
-        if left_mat is None or right_mat is None:
-            return None
-
-        stacked_results = self._cosine_distance_stacked(left_mat, right_mat)
-        if row_present.all():
-            return stacked_results
-
-        results = np.full(n, np.nan, dtype=np.float64)
-        results[valid_idx] = stacked_results
-        return results
-
-    def _parse_vectors(self, left_arr: np.ndarray, right_arr: np.ndarray) -> tuple[list, list, np.ndarray]:
-        n = len(left_arr)
-        left_vecs: list[Optional[np.ndarray]] = [None] * n
-        right_vecs: list[Optional[np.ndarray]] = [None] * n
-        row_valid = np.zeros(n, dtype=bool)
-        for i, (left_value, right_value) in enumerate(zip(left_arr, right_arr)):
-            left_vec = self._to_float_vector(left_value)
-            right_vec = self._to_float_vector(right_value)
-            left_vecs[i] = left_vec
-            right_vecs[i] = right_vec
-            if left_vec is not None and right_vec is not None and left_vec.size > 0 and right_vec.size > 0:
-                row_valid[i] = True
-        return left_vecs, right_vecs, row_valid
-
-    def _masked_cosine_distance(
-        self, left_mat: np.ndarray, right_mat: np.ndarray, row_valid: np.ndarray
-    ) -> np.ndarray:
-        results = self._cosine_distance_stacked(left_mat, right_mat)
-        if len(results) == len(row_valid):
-            results[~row_valid] = np.nan
-        return results
-
-    def _calculate_rowwise(
-        self, results: np.ndarray, indices: np.ndarray, left_vecs: list, right_vecs: list
-    ) -> None:
-        for index in indices:
-            left_vec = left_vecs[index]
-            right_vec = right_vecs[index]
-            dot = np.sum(left_vec * right_vec)
-            denom = np.sqrt(np.sum(left_vec * left_vec)) * np.sqrt(np.sum(right_vec * right_vec))
-            if denom != 0:
-                results[index] = 1 - dot / denom
-
-    def _calculate_fallback_distance(
-        self, left_arr: np.ndarray, right_arr: np.ndarray, n: int, index: pd.Index
-    ) -> pd.Series:
-        results = np.full(n, np.nan, dtype=np.float64)
-        left_vecs, right_vecs, row_valid = self._parse_vectors(left_arr, right_arr)
-        if not row_valid.any():
-            return pd.Series(results, index=index, dtype=np.float64)
-
-        valid_idx = np.flatnonzero(row_valid)
-        dim = left_vecs[valid_idx[0]].size
-        if not all(left_vecs[i].size == dim and right_vecs[i].size == dim for i in valid_idx):
-            self._calculate_rowwise(results, valid_idx, left_vecs, right_vecs)
-            return pd.Series(results, index=index, dtype=np.float64)
-
-        left_mat = np.zeros((n, dim), dtype=np.float64)
-        right_mat = np.zeros((n, dim), dtype=np.float64)
-        left_mat[valid_idx] = np.stack([left_vecs[i] for i in valid_idx])
-        right_mat[valid_idx] = np.stack([right_vecs[i] for i in valid_idx])
-
-        return pd.Series(self._masked_cosine_distance(left_mat, right_mat, row_valid), index=index, dtype=np.float64)
-
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
-        n = len(left)
-        if n == 0:
-            return pd.Series(np.full(n, np.nan), index=left.index, dtype=np.float64)
-
-        row_present = left.notna().to_numpy() & right.notna().to_numpy()
-        left_arr = left.to_numpy()
-        right_arr = right.to_numpy()
-
-        fast_results = self._try_fast_distance(left_arr, right_arr, row_present, n)
-        if fast_results is not None:
-            return pd.Series(fast_results, index=left.index, dtype=np.float64)
-
-        return self._calculate_fallback_distance(left_arr, right_arr, n, left.index)
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        if left.kind == OperandKind.MATRIX and right.kind == OperandKind.MATRIX:
+            return self._calculate_binary_matrices(left, right)
+        return self._calculate_binary_series(left.as_series(), right.as_series())
+
+    def _calculate_binary_matrices(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left_mat = left.as_matrix()
+        right_mat = right.as_matrix()
+        dot = np.sum(left_mat * right_mat, axis=1)
+        left_norm = np.sqrt(np.sum(left_mat * left_mat, axis=1))
+        right_norm = np.sqrt(np.sum(right_mat * right_mat, axis=1))
+        denom = left_norm * right_norm
+        with np.errstate(invalid="ignore", divide="ignore"):
+            result = 1 - dot / denom
+        result[denom == 0] = np.nan
+        return pd.Series(result, index=left.index, dtype=np.float64)
+
+    def _calculate_binary_series(self, left: pd.Series, right: pd.Series) -> pd.Series:
+        return pd.Series(
+            1 - self.__dot(left, right) / (self.__norm(left) * self.__norm(right)), index=left.index
+        ).astype(np.float64)
+
+    # row-wise dot product, handling None values
+    def __dot(self, left: pd.Series, right: pd.Series) -> pd.Series:
+        left = left.apply(lambda x: np.array(x).astype(np.float64))
+        right = right.apply(lambda x: np.array(x).astype(np.float64))
+        res = (left.dropna() * right.dropna()).apply(np.sum)
+        res = res.reindex(left.index.union(right.index))
+        return res
+
+    # Calculate the norm of a vector, handling None values
+    def __norm(self, vector: pd.Series) -> pd.Series:
+        vector = vector.fillna(np.nan)
+        return np.sqrt(self.__dot(vector, vector))
 
 
 # Left for backward compatibility
@@ -264,12 +201,14 @@ class Sim(Distance):
     is_symmetrical: bool = True
     has_symmetry_importance: bool = True
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
         return 1 - super().calculate_binary(left, right)
 
 
 class StringSim(PandasOperator, abc.ABC):
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         sims = []
         for i in left.index:
             left_i = self._prepare_value(left.get(i))

{upgini-1.2.156.dev5 → upgini-1.2.157.dev2}/src/upgini/autofe/date.py

@@ -7,6 +7,7 @@ import pandas as pd
 from pandas.core.arrays.timedeltas import TimedeltaArray
 from pydantic import BaseModel, __version__ as pydantic_version
 
+from upgini.autofe.operand import OperandValue
 from upgini.autofe.operator import PandasOperator, ParametrizedOperator
 from upgini.autofe.utils import bin_index, bin_index_many, bin_index_vectorized, pydantic_validator
 
@@ -64,7 +65,9 @@ class DateDiff(PandasOperator, DateDiffMixin):
         )
         return res
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         if left.isna().all() or right.isna().all():
             return pd.Series([None] * len(left), index=left.index, dtype=np.float64)
 
@@ -95,7 +98,9 @@ class DateDiffType2(PandasOperator, DateDiffMixin):
         )
         return res
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         left = self._convert_to_date(left, self.left_unit)
         right = self._convert_to_date(right, self.right_unit)
         future = right + (left.dt.year - right.dt.year).apply(
@@ -187,7 +192,9 @@ class DateListDiffLists(PandasOperator, DateDiffMixin, ParametrizedOperator):
             diffs = diffs[diffs > 0]
         return np.atleast_1d(np.asarray(diffs, dtype=np.float64)).tolist()
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         if left.isna().all() or right.isna().all():
             return pd.Series([None] * len(left), index=left.index, dtype=object)
 
@@ -271,7 +278,8 @@ class DateListDiffAggWithinBounds(PandasOperator, ParametrizedOperator):
             return agg_res / orig_len
         return agg_res
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         results = np.empty(len(data), dtype=np.float64)
         results[:] = np.nan
         for i, diffs in enumerate(data.to_numpy()):
@@ -332,13 +340,15 @@ class DateListDiff(PandasOperator, DateDiffMixin, ParametrizedOperator):
             return pd.Series([None] * len(left), index=left.index, dtype=np.float64)
 
         right_mask = DateListDiffLists._non_empty_list_mask(right)
-        diff_lists = self._lists_op().calculate_binary(left, right)
-        result = self._agg_op().calculate_unary(diff_lists)
+        diff_lists = self._lists_op().calculate(left=left, right=right)
+        result = self._agg_op().calculate(data=diff_lists).as_series()
         if self.aggregation in _count_aggregations:
             result[~right_mask] = 0.0
         return result.astype(np.float64)
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         return self._compose_list_diff(left, right)
 
 
@@ -404,7 +414,9 @@ class DatePercentileBase(PandasOperator, abc.ABC):
 
     date_unit: Optional[str] = None
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         # Assuming that left is a date column, right is a feature column
         left = pd.to_datetime(left, unit=self.date_unit)
 

{upgini-1.2.156.dev5 → upgini-1.2.157.dev2}/src/upgini/autofe/feature.py

@@ -8,13 +8,14 @@ import pandas as pd
 from pandas._typing import DtypeObj
 
 from upgini.autofe.all_operators import find_op
+from upgini.autofe.operand import CalculationContext, OperandValue
 from upgini.autofe.operator import Operator, PandasOperator
 from upgini.autofe.timeseries.base import TimeSeriesBase
 from upgini.autofe.utils import pydantic_dump_method, pydantic_parse_method
 
 
 class Column:
-    def __init__(self, name: str, data: Optional[pd.Series] = None, calculate_all=False):
+    def __init__(self, name: str, data: Optional[OperandValue] = None, calculate_all=False):
         self.name = name
         self.data = data
         self.calculate_all = calculate_all
@@ -61,8 +62,12 @@ class Column:
     def infer_type(self, data: pd.DataFrame) -> DtypeObj:
         return data[self.name].dtype
 
-    def calculate(self, data: pd.DataFrame) -> pd.Series:
-        self.data = data[self.name]
+    def calculate(self, data: Union[pd.DataFrame, CalculationContext]) -> pd.Series:
+        ctx = data if isinstance(data, CalculationContext) else CalculationContext.from_dataframe(data)
+        return self._eval(ctx).as_series()
+
+    def _eval(self, ctx: CalculationContext) -> OperandValue:
+        self.data = ctx.resolve_operand(self.name)
         return self.data
 
     def to_formula(self, **kwargs) -> str:
@@ -236,27 +241,38 @@ class Feature:
             # either a symmetrical operator or group by
             return self.children[0].infer_type(data)
 
-    def calculate(self, data: pd.DataFrame, is_root=False) -> Union[pd.Series, pd.DataFrame]:
-        if isinstance(self.op, PandasOperator):
-            if self.op.is_vector:
-                ds = [child.calculate(data) for child in self.children]
-                new_data = self.op.calculate(data=ds)
-            else:
-                d1 = self.children[0].calculate(data)
-                d2 = None if len(self.children) < 2 else self.children[1].calculate(data)
-                new_data = self.op.calculate(data=d1, left=d1, right=d2)
-        else:
-            raise NotImplementedError(f"Unrecognized operator {self.op.name}.")
-
-        if (str(new_data.dtype) == "category") | (str(new_data.dtype) == "object"):
-            pass
-        else:
-            new_data = new_data.replace([-np.inf, np.inf], np.nan)
+    def calculate(
+        self,
+        data: Union[pd.DataFrame, CalculationContext],
+        is_root: bool = False,
+    ) -> Union[pd.Series, pd.DataFrame]:
+        ctx = data if isinstance(data, CalculationContext) else CalculationContext.from_dataframe(data)
+        result = self._eval(ctx)
+        new_data = self._operand_to_output(result)
 
         if is_root:
             self.data = new_data
         return new_data
 
+    def _eval(self, ctx: CalculationContext) -> OperandValue:
+        if isinstance(self.op, PandasOperator):
+            if self.op.is_vector:
+                operands = [child._eval(ctx) for child in self.children]
+                return self.op.calculate(data=operands)
+            left = self.children[0]._eval(ctx)
+            if len(self.children) < 2:
+                return self.op.calculate(data=left)
+            right = self.children[1]._eval(ctx)
+            return self.op.calculate(left=left, right=right)
+        raise NotImplementedError(f"Unrecognized operator {self.op.name}.")
+
+    @staticmethod
+    def _operand_to_output(operand: OperandValue) -> Union[pd.Series, pd.DataFrame]:
+        new_data = operand.as_series()
+        if (str(new_data.dtype) == "category") | (str(new_data.dtype) == "object"):
+            return new_data
+        return new_data.replace([-np.inf, np.inf], np.nan)
+
     @staticmethod
     def check_xor(left: Union[Column, "Feature"], right: Union[Column, "Feature"]) -> bool:
         def _get_all_columns(feature: Union[Column, "Feature"]) -> List[str]:

{upgini-1.2.156.dev5 → upgini-1.2.157.dev2}/src/upgini/autofe/groupby.py

@@ -2,6 +2,7 @@ from typing import Optional
 
 import pandas as pd
 
+from upgini.autofe.operand import OperandValue
 from upgini.autofe.operator import PandasOperator, ParametrizedOperator, VectorizableMixin
 
 
@@ -27,7 +28,9 @@ class GroupByThenAgg(
             return None
         return cls(agg=agg)
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         temp = left.groupby(right).agg(self.agg.lower())
         return self._loc(right, temp)
 
@@ -46,7 +49,9 @@ class GroupByThenRank(PandasOperator, VectorizableMixin):
     output_type: Optional[str] = "float"
     is_distribution_dependent: bool = True
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         temp = pd.DataFrame(left[~right.isna()].groupby(right).rank(ascending=True, pct=True)).reset_index()
         return temp.merge(pd.DataFrame(right).reset_index(), how="right", on=["index"])[left.name]
 
@@ -66,7 +71,9 @@ class GroupByThenNUnique(PandasOperator, VectorizableMixin):
     is_distribution_dependent: bool = True
     input_type: Optional[str] = "discrete"
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         nunique = left.groupby(right).nunique()
         return self._loc(right, nunique)
 
@@ -85,7 +92,9 @@ class GroupByThenFreq(PandasOperator):
     is_distribution_dependent: bool = True
     input_type: Optional[str] = "discrete"
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(self, left: OperandValue, right: OperandValue) -> pd.Series:
+        left = left.as_series()
+        right = right.as_series()
         def _f(x):
             value_counts = x.value_counts(normalize=True)
             return self._loc(x, value_counts)

upgini-1.2.157.dev2/src/upgini/autofe/operand.py

@@ -0,0 +1,112 @@
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, Optional, Union
+
+import numpy as np
+import pandas as pd
+
+
+class OperandKind(str, Enum):
+    SERIES = "series"
+    MATRIX = "matrix"
+    ARRAY = "array"
+
+
+@dataclass
+class OperandValue:
+    kind: OperandKind
+    data: Union[pd.Series, np.ndarray]
+    index: pd.Index
+    source: Optional[str] = None
+
+    @classmethod
+    def from_series(cls, series: pd.Series, source: Optional[str] = None) -> "OperandValue":
+        return cls(kind=OperandKind.SERIES, data=series, index=series.index, source=source)
+
+    @classmethod
+    def from_matrix(cls, matrix: np.ndarray, index: pd.Index, source: Optional[str] = None) -> "OperandValue":
+        return cls(kind=OperandKind.MATRIX, data=matrix, index=index, source=source)
+
+    @classmethod
+    def from_array(cls, array: np.ndarray, index: pd.Index, source: Optional[str] = None) -> "OperandValue":
+        return cls(kind=OperandKind.ARRAY, data=array, index=index, source=source)
+
+    def as_series(self) -> pd.Series:
+        if self.kind == OperandKind.SERIES:
+            return self.data  # type: ignore[return-value]
+        if self.kind == OperandKind.ARRAY:
+            return pd.Series(self.data, index=self.index, dtype=np.float64)
+        if self.kind == OperandKind.MATRIX:
+            return pd.Series(list(self.as_matrix()), index=self.index)
+        raise TypeError(f"Cannot convert operand kind {self.kind.value} to series")
+
+    def as_matrix(self) -> np.ndarray:
+        if self.kind != OperandKind.MATRIX:
+            raise TypeError(f"Cannot convert operand kind {self.kind.value} to matrix")
+        return self.data  # type: ignore[return-value]
+
+
+@dataclass
+class CalculationContext:
+    data: Optional[pd.DataFrame] = None
+    matrices: Dict[str, np.ndarray] = field(default_factory=dict)
+    row_index: Optional[pd.Index] = None
+
+    def __post_init__(self) -> None:
+        if self.data is None and not self.matrices:
+            raise ValueError("CalculationContext requires data and/or matrices")
+        if self.matrices:
+            lengths = {name: matrix.shape[0] for name, matrix in self.matrices.items()}
+            if len(set(lengths.values())) > 1:
+                raise ValueError(f"Matrix row counts mismatch: {lengths}")
+
+    @property
+    def index(self) -> pd.Index:
+        if self.data is not None:
+            return self.data.index
+        if self.row_index is not None:
+            return self.row_index
+        if self.matrices:
+            nrows = next(iter(self.matrices.values())).shape[0]
+            return pd.RangeIndex(nrows)
+        raise ValueError("CalculationContext has no index source")
+
+    @classmethod
+    def from_dataframe(cls, data: pd.DataFrame) -> "CalculationContext":
+        return cls(data=data)
+
+    @classmethod
+    def from_matrices(
+        cls,
+        matrices: Dict[str, np.ndarray],
+        index: Optional[pd.Index] = None,
+    ) -> "CalculationContext":
+        return cls(matrices=matrices, row_index=index)
+
+    def resolve_operand(self, name: str) -> OperandValue:
+        matrix = self.matrices.get(name)
+        if matrix is not None:
+            return OperandValue.from_matrix(matrix, index=self.index, source=name)
+        if self.data is not None and name in self.data:
+            return OperandValue.from_series(self.data[name], source=name)
+        raise KeyError(name)
+
+
+def wrap_operand(
+    result: Union[pd.Series, np.ndarray, "OperandValue"],
+    source: Optional[str] = None,
+    index: Optional[pd.Index] = None,
+) -> "OperandValue":
+    if isinstance(result, OperandValue):
+        return result
+    if isinstance(result, np.ndarray):
+        if index is None:
+            raise TypeError("index is required when wrapping ndarray result")
+        if result.ndim == 1:
+            return OperandValue.from_array(result, index=index, source=source)
+        if result.ndim == 2:
+            return OperandValue.from_matrix(result, index=index, source=source)
+        raise TypeError(f"Cannot wrap ndarray with {result.ndim} dimensions")
+    return OperandValue.from_series(result, source=source)

{upgini-1.2.156.dev5 → upgini-1.2.157.dev2}/src/upgini/autofe/operator.py

@@ -6,6 +6,8 @@ import numpy as np
 import pandas as pd
 from pydantic import BaseModel, PrivateAttr
 
+from upgini.autofe.operand import OperandValue, wrap_operand
+
 
 class OperatorRegistry(type(BaseModel)):
     _registry = {}
@@ -132,21 +134,35 @@ MAIN_COLUMN = "main_column"
 
 
 class PandasOperator(Operator, abc.ABC):
-    def calculate(self, **kwargs) -> pd.Series:
-        if self.is_unary:
-            return self.calculate_unary(kwargs["data"])
-        elif self.is_binary or self.is_grouping:
-            return self.calculate_binary(kwargs["left"], kwargs["right"])
-        else:
-            return self.calculate_vector(kwargs["data"])
+    @staticmethod
+    def _operand(value: Union[pd.Series, OperandValue]) -> OperandValue:
+        if isinstance(value, OperandValue):
+            return value
+        return OperandValue.from_series(value)
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate(self, **kwargs) -> OperandValue:
+        if self.is_unary:
+            data = self._operand(kwargs["data"])
+            return wrap_operand(self.calculate_unary(data), source=data.source, index=data.index)
+        if self.is_binary or self.is_grouping:
+            left = self._operand(kwargs["left"])
+            right = self._operand(kwargs["right"])
+            return wrap_operand(self.calculate_binary(left, right), source=left.source, index=left.index)
+        operands = [self._operand(value) for value in kwargs["data"]]
+        index = operands[0].index if operands else None
+        return wrap_operand(self.calculate_vector(operands), index=index)
+
+    def calculate_unary(self, data: OperandValue) -> Union[pd.Series, np.ndarray, OperandValue]:
         pass
 
-    def calculate_binary(self, left: pd.Series, right: pd.Series) -> pd.Series:
+    def calculate_binary(
+        self, left: OperandValue, right: OperandValue
+    ) -> Union[pd.Series, np.ndarray, OperandValue]:
         pass
 
-    def calculate_vector(self, data: List[pd.Series]) -> pd.Series:
+    def calculate_vector(
+        self, data: List[OperandValue]
+    ) -> Union[pd.Series, np.ndarray, OperandValue]:
         pass
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:

{upgini-1.2.156.dev5 → upgini-1.2.157.dev2}/src/upgini/autofe/timeseries/base.py

@@ -2,6 +2,7 @@ import abc
 from typing import Dict, List, Optional, Tuple
 
 import pandas as pd
+from upgini.autofe.operand import OperandValue
 from upgini.autofe.operator import PandasOperator
 
 
@@ -22,7 +23,8 @@ class TimeSeriesBase(PandasOperator, abc.ABC):
         )
         return res
 
-    def calculate_vector(self, data: List[pd.Series]) -> pd.Series:
+    def calculate_vector(self, data: List[OperandValue]) -> pd.Series:
+        data = [operand.as_series() for operand in data]
         # assuming first is date, last is value, rest is group columns
         date = pd.to_datetime(data[0], unit=self.date_unit, errors="coerce")
         ts = pd.concat([date] + data[1:], axis=1)

{upgini-1.2.156.dev5 → upgini-1.2.157.dev2}/src/upgini/autofe/timeseries/cross.py

@@ -5,6 +5,7 @@ import numpy as np
 import pandas as pd
 
 from upgini.autofe.all_operators import find_op
+from upgini.autofe.operand import OperandValue
 from upgini.autofe.operator import PandasOperator, ParametrizedOperator
 from upgini.autofe.timeseries.base import TimeSeriesBase
 from upgini.autofe.utils import pydantic_validator
@@ -117,14 +118,15 @@ class CrossSeriesInteraction(TimeSeriesBase, ParametrizedOperator):
     def _get_interaction_op_name(self) -> str:
         return self.interaction_op.alias or self.interaction_op.to_formula()
 
-    def calculate_vector(self, data: List[pd.Series]) -> pd.Series:
+    def calculate_vector(self, data: List[OperandValue]) -> pd.Series:
+        data = [operand.as_series() for operand in data]
         left_mask = self._get_mask(data, self.left_descriptor)
         left = self._extract_series(data, left_mask)
 
         right_mask = self._get_mask(data, self.right_descriptor)
         right = self._extract_series(data, right_mask)
 
-        interaction: pd.Series = self.interaction_op.calculate_binary(left, right)
+        interaction: pd.Series = self.interaction_op.calculate(left=left, right=right).as_series()
         interaction = interaction.reindex(self._get_index(data))
         res = pd.Series(np.nan, index=data[-1].index, name=data[-1].name)
         res.loc[left_mask] = interaction[left_mask].values
@@ -137,7 +139,7 @@ class CrossSeriesInteraction(TimeSeriesBase, ParametrizedOperator):
 
     def _extract_series(self, data: List[pd.Series], mask: pd.Series) -> pd.Series:
         masked_data = [d[mask] for d in data]
-        shifted = super().calculate_vector(masked_data)
+        shifted = super().calculate_vector([OperandValue.from_series(d) for d in masked_data])
         shifted.index = self._get_index(masked_data)
         return shifted
 

{upgini-1.2.156.dev5 → upgini-1.2.157.dev2}/src/upgini/autofe/unary.py

@@ -3,6 +3,7 @@ from typing import Dict, List, Optional
 import numpy as np
 import pandas as pd
 
+from upgini.autofe.operand import OperandValue
 from upgini.autofe.operator import PandasOperator, ParametrizedOperator, VectorizableMixin
 from upgini.autofe.utils import bin_index, bin_index_vectorized, pydantic_validator
 
@@ -13,7 +14,8 @@ class Abs(PandasOperator, VectorizableMixin):
     is_vectorizable: bool = True
     group_index: int = 0
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         return data.astype(np.float64).abs()
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -27,7 +29,8 @@ class Log(PandasOperator, VectorizableMixin):
     output_type: Optional[str] = "float"
     group_index: int = 0
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         return self._round_value(np.log(np.abs(data.replace(0, np.nan))), 10)
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -41,7 +44,8 @@ class Sqrt(PandasOperator, VectorizableMixin):
     output_type: Optional[str] = "float"
     group_index: int = 0
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         return self._round_value(np.sqrt(np.abs(data)))
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -54,7 +58,8 @@ class Square(PandasOperator, VectorizableMixin):
     is_vectorizable: bool = True
     group_index: int = 0
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         return np.square(data)
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -68,7 +73,8 @@ class Sigmoid(PandasOperator, VectorizableMixin):
     output_type: Optional[str] = "float"
     group_index: int = 0
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         return self._round_value(1 / (1 + np.exp(-data)))
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -83,7 +89,8 @@ class Floor(PandasOperator, VectorizableMixin):
     input_type: Optional[str] = "continuous"
     group_index: int = 0
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         return np.floor(data)
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -97,7 +104,8 @@ class Residual(PandasOperator, VectorizableMixin):
     input_type: Optional[str] = "continuous"
     group_index: int = 0
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         return data - np.floor(data)
 
     def calculate_group(self, data: pd.DataFrame, **kwargs) -> pd.DataFrame:
@@ -111,7 +119,8 @@ class Freq(PandasOperator):
     is_distribution_dependent: bool = True
     input_type: Optional[str] = "discrete"
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         value_counts = data.value_counts(normalize=True)
         return self._loc(data, value_counts)
 
@@ -122,7 +131,8 @@ class Norm(PandasOperator):
     output_type: Optional[str] = "float"
     norm: Optional[float] = None
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         data_dropna = data.dropna()
         if data_dropna.empty:
             return data
@@ -152,7 +162,7 @@ class Embeddings(PandasOperator):
     name: str = "emb"
     is_unary: bool = True
     input_type: Optional[str] = "string"
-    output_type: Optional[str] = "vector"
+    output_type: Optional[str] = "matrix"
 
 
 class Bin(PandasOperator):
@@ -162,7 +172,8 @@ class Bin(PandasOperator):
     bin_bounds: List[int] = []
     is_categorical: bool = True
 
-    def calculate_unary(self, data: pd.Series) -> pd.Series:
+    def calculate_unary(self, data: OperandValue) -> pd.Series:
+        data = data.as_series()
         bounds_arr = np.asarray(self.bin_bounds, dtype=np.float64)
         values = pd.to_numeric(data, errors="coerce").to_numpy(dtype=np.float64, copy=False)
         result = bin_index_vectorized(values, bounds_arr)

{upgini-1.2.156.dev5 → upgini-1.2.157.dev2}/src/upgini/autofe/vector.py

@@ -2,6 +2,7 @@ from typing import List, Optional
 
 import pandas as pd
 
+from upgini.autofe.operand import OperandValue
 from upgini.autofe.operator import OperatorRegistry, PandasOperator, ParametrizedOperator, VectorizableMixin
 
 
@@ -11,8 +12,9 @@ class Mean(PandasOperator, VectorizableMixin):
     is_vector: bool = True
     group_index: int = 0
 
-    def calculate_vector(self, data: List[pd.Series]) -> pd.Series:
-        return pd.DataFrame(data).T.fillna(0).mean(axis=1)
+    def calculate_vector(self, data: List[OperandValue]) -> pd.Series:
+        series = [operand.as_series() for operand in data]
+        return pd.DataFrame(series).T.fillna(0).mean(axis=1)
 
 
 class Sum(PandasOperator, VectorizableMixin):
@@ -20,8 +22,9 @@ class Sum(PandasOperator, VectorizableMixin):
     is_vector: bool = True
     group_index: int = 0
 
-    def calculate_vector(self, data: List[pd.Series]) -> pd.Series:
-        return pd.DataFrame(data).T.fillna(0).sum(axis=1)
+    def calculate_vector(self, data: List[OperandValue]) -> pd.Series:
+        series = [operand.as_series() for operand in data]
+        return pd.DataFrame(series).T.fillna(0).sum(axis=1)
 
 
 class Vectorize(PandasOperator, VectorizableMixin):
@@ -29,8 +32,9 @@ class Vectorize(PandasOperator, VectorizableMixin):
     is_vector: bool = True
     group_index: int = 0
 
-    def calculate_vector(self, data: List[pd.Series]) -> pd.Series:
-        return pd.DataFrame(data).T.apply(lambda x: x.to_list(), axis=1)
+    def calculate_vector(self, data: List[OperandValue]) -> pd.Series:
+        series = [operand.as_series() for operand in data]
+        return pd.DataFrame(series).T.apply(lambda x: x.to_list(), axis=1)
 
 
 class OnnxModel(PandasOperator, ParametrizedOperator, metaclass=OperatorRegistry):

upgini-1.2.156.dev5/src/upgini/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.2.156.dev5"