Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[FEAT] Implements trigonometry expressions: arctanh arccosh arcsinh #2476

Merged
Merged
Show file tree
Hide file tree
Changes from 6 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions daft/daft.pyi
Original file line number Diff line number Diff line change
Expand Up @@ -1036,6 +1036,9 @@ class PyExpr:
def arccos(self) -> PyExpr: ...
def arctan(self) -> PyExpr: ...
def arctan2(self, other: PyExpr) -> PyExpr: ...
def arctanh(self) -> PyExpr: ...
def arccosh(self) -> PyExpr: ...
def arcsinh(self) -> PyExpr: ...
def degrees(self) -> PyExpr: ...
def radians(self) -> PyExpr: ...
def log2(self) -> PyExpr: ...
Expand Down Expand Up @@ -1225,6 +1228,9 @@ class PySeries:
def arccos(self) -> PySeries: ...
def arctan(self) -> PySeries: ...
def arctan2(self, other: PySeries) -> PySeries: ...
def arctanh(self) -> PySeries: ...
def arccosh(self) -> PySeries: ...
def arcsinh(self) -> PySeries: ...
def degrees(self) -> PySeries: ...
def radians(self) -> PySeries: ...
def log2(self) -> PySeries: ...
Expand Down
15 changes: 15 additions & 0 deletions daft/expressions/expressions.py
Original file line number Diff line number Diff line change
Expand Up @@ -434,6 +434,21 @@ def arctan2(self, other: Expression) -> Expression:
expr = Expression._to_expression(other)
return Expression._from_pyexpr(self._expr.arctan2(expr._expr))

def arctanh(self) -> Expression:
"""The elementwise inverse hyperbolic tangent of a numeric expression (``expr.arctanh()``)"""
expr = self._expr.arctanh()
return Expression._from_pyexpr(expr)

def arccosh(self) -> Expression:
"""The elementwise inverse hyperbolic cosine of a numeric expression (``expr.arccosh()``)"""
expr = self._expr.arccosh()
return Expression._from_pyexpr(expr)

def arcsinh(self) -> Expression:
"""The elementwise inverse hyperbolic sine of a numeric expression (``expr.arcsinh()``)"""
expr = self._expr.arcsinh()
return Expression._from_pyexpr(expr)

def radians(self) -> Expression:
"""The elementwise radians of a numeric expression (``expr.radians()``)"""
expr = self._expr.radians()
Expand Down
12 changes: 12 additions & 0 deletions daft/series.py
Original file line number Diff line number Diff line change
Expand Up @@ -402,6 +402,18 @@ def arctan2(self, other: Series) -> Series:
raise TypeError(f"expected another Series but got {type(other)}")
return Series._from_pyseries(self._series.arctan2(other._series))

def arctanh(self) -> Series:
"""The elementwise inverse hyperbolic tangent of a numeric series"""
return Series._from_pyseries(self._series.arctanh())

def arccosh(self) -> Series:
"""The elementwise inverse hyperbolic cosine of a numeric series"""
return Series._from_pyseries(self._series.arccosh())

def arcsinh(self) -> Series:
"""The elementwise inverse hyperbolic sine of a numeric series"""
return Series._from_pyseries(self._series.arcsinh())

def radians(self) -> Series:
"""The elementwise radians of a numeric series"""
return Series._from_pyseries(self._series.radians())
Expand Down
3 changes: 3 additions & 0 deletions docs/source/api_docs/expressions.rst
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,9 @@ Numeric
Expression.arccos
Expression.arctan
Expression.arctan2
Expression.arctanh
Expression.arccosh
Expression.arcsinh
Expression.radians
Expression.degrees
Expression.log2
Expand Down
9 changes: 9 additions & 0 deletions src/daft-core/src/array/ops/trigonometry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,9 @@ pub enum TrigonometricFunction {
ArcTan,
Radians,
Degrees,
ArcTanh,
ArcCosh,
ArcSinh,
}

impl TrigonometricFunction {
Expand All @@ -34,6 +37,9 @@ impl TrigonometricFunction {
ArcTan => "arctan",
Radians => "radians",
Degrees => "degrees",
ArcTanh => "arctanh",
ArcCosh => "arccosh",
ArcSinh => "arcsinh",
}
}
}
Expand All @@ -55,6 +61,9 @@ where
ArcTan => self.apply(|v| v.atan()),
Radians => self.apply(|v| v.to_radians()),
Degrees => self.apply(|v| v.to_degrees()),
ArcTanh => self.apply(|v| v.atanh()),
ArcCosh => self.apply(|v| v.acosh()),
ArcSinh => self.apply(|v| v.asinh()),
}
}
}
Expand Down
21 changes: 21 additions & 0 deletions src/daft-core/src/python/series.rs
Original file line number Diff line number Diff line change
Expand Up @@ -205,6 +205,27 @@ impl PySeries {
.into())
}

pub fn arctanh(&self) -> PyResult<Self> {
Ok(self
.series
.trigonometry(&TrigonometricFunction::ArcTanh)?
.into())
}

pub fn arccosh(&self) -> PyResult<Self> {
Ok(self
.series
.trigonometry(&TrigonometricFunction::ArcCosh)?
.into())
}

pub fn arcsinh(&self) -> PyResult<Self> {
Ok(self
.series
.trigonometry(&TrigonometricFunction::ArcSinh)?
.into())
}

pub fn log2(&self) -> PyResult<Self> {
Ok(self.series.log2()?.into())
}
Expand Down
30 changes: 30 additions & 0 deletions src/daft-dsl/src/functions/numeric/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,9 @@ pub enum NumericExpr {
Log(FloatWrapper<f64>),
Ln,
Exp,
ArcTanh,
ArcCosh,
ArcSinh,
}

impl NumericExpr {
Expand Down Expand Up @@ -77,6 +80,9 @@ impl NumericExpr {
Log(_) => &LogEvaluator(log::LogFunction::Log),
Ln => &LogEvaluator(log::LogFunction::Ln),
Exp => &ExpEvaluator {},
ArcTanh => &TrigonometryEvaluator(TrigonometricFunction::ArcTanh),
ArcCosh => &TrigonometryEvaluator(TrigonometricFunction::ArcCosh),
ArcSinh => &TrigonometryEvaluator(TrigonometricFunction::ArcSinh),
}
}
}
Expand Down Expand Up @@ -209,6 +215,30 @@ pub fn degrees(input: ExprRef) -> ExprRef {
.into()
}

pub fn arctanh(input: ExprRef) -> ExprRef {
Expr::Function {
func: super::FunctionExpr::Numeric(NumericExpr::ArcTanh),
inputs: vec![input],
}
.into()
}

pub fn arccosh(input: ExprRef) -> ExprRef {
Expr::Function {
func: super::FunctionExpr::Numeric(NumericExpr::ArcCosh),
inputs: vec![input],
}
.into()
}

pub fn arcsinh(input: ExprRef) -> ExprRef {
Expr::Function {
func: super::FunctionExpr::Numeric(NumericExpr::ArcSinh),
inputs: vec![input],
}
.into()
}

pub fn log2(input: ExprRef) -> ExprRef {
Expr::Function {
func: super::FunctionExpr::Numeric(NumericExpr::Log2),
Expand Down
15 changes: 15 additions & 0 deletions src/daft-dsl/src/python.rs
Original file line number Diff line number Diff line change
Expand Up @@ -282,6 +282,21 @@ impl PyExpr {
Ok(degrees(self.into()).into())
}

pub fn arctanh(&self) -> PyResult<Self> {
use functions::numeric::arctanh;
Ok(arctanh(self.into()).into())
}

pub fn arccosh(&self) -> PyResult<Self> {
use functions::numeric::arccosh;
Ok(arccosh(self.into()).into())
}

pub fn arcsinh(&self) -> PyResult<Self> {
use functions::numeric::arcsinh;
Ok(arcsinh(self.into()).into())
}

pub fn log2(&self) -> PyResult<Self> {
use functions::numeric::log2;
Ok(log2(self.into()).into())
Expand Down
27 changes: 27 additions & 0 deletions tests/expressions/test_expressions.py
Original file line number Diff line number Diff line change
Expand Up @@ -195,6 +195,33 @@ def test_repr_functions_atan2() -> None:
assert repr_out == repr(copied)


def test_repr_functions_arctanh() -> None:
a = col("a")
y = a.arctanh()
repr_out = repr(y)
assert repr_out == "arctanh(col(a))"
copied = copy.deepcopy(y)
assert repr_out == repr(copied)


def test_repr_functions_arccosh() -> None:
a = col("a")
y = a.arccosh()
repr_out = repr(y)
assert repr_out == "arccosh(col(a))"
copied = copy.deepcopy(y)
assert repr_out == repr(copied)


def test_repr_functions_arcsinh() -> None:
a = col("a")
y = a.arcsinh()
repr_out = repr(y)
assert repr_out == "arcsinh(col(a))"
copied = copy.deepcopy(y)
assert repr_out == repr(copied)


def test_repr_functions_day() -> None:
a = col("a")
y = a.dt.day()
Expand Down
30 changes: 30 additions & 0 deletions tests/expressions/typing/test_arithmetic.py
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,36 @@ def test_atan2(binary_data_fixture):
)


def test_atanh(unary_data_fixture):
arg = unary_data_fixture
assert_typing_resolve_vs_runtime_behavior(
data=(unary_data_fixture,),
expr=col(arg.name()).arctanh(),
run_kernel=lambda: arg.arctanh(),
resolvable=is_numeric(arg.datatype()),
)


def test_asinh(unary_data_fixture):
arg = unary_data_fixture
assert_typing_resolve_vs_runtime_behavior(
data=(unary_data_fixture,),
expr=col(arg.name()).arcsinh(),
run_kernel=lambda: arg.arcsinh(),
resolvable=is_numeric(arg.datatype()),
)


def test_acosh(unary_data_fixture):
arg = unary_data_fixture
assert_typing_resolve_vs_runtime_behavior(
data=(unary_data_fixture,),
expr=col(arg.name()).arccosh(),
run_kernel=lambda: arg.arccosh(),
resolvable=is_numeric(arg.datatype()),
)


def test_exp(unary_data_fixture):
arg = unary_data_fixture
assert_typing_resolve_vs_runtime_behavior(
Expand Down
45 changes: 45 additions & 0 deletions tests/table/test_eval.py
Original file line number Diff line number Diff line change
Expand Up @@ -367,6 +367,51 @@ def test_table_numeric_atan2_literals() -> None:
)


def test_table_numeric_arctanh() -> None:
table = MicroPartition.from_pydict({"a": [0.0, 0.5, 0.9, -0.9, -0.5, -0.0, math.nan]})
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you also test that for values |x| >= 1, the arctanh expression returns Nan?

s = table.to_pandas()["a"]
np_result = np.arctanh(s)

arct = table.eval_expression_list([col("a").arctanh()])
assert (
all(
x - y < 1.0e-10 or (x is None and y is None) or (math.isnan(x) and math.isnan(y))
for x, y in zip(arct.get_column("a").to_pylist(), np_result.to_list())
)
is True
)


def test_table_numeric_arcsinh() -> None:
table = MicroPartition.from_pydict({"a": [0.0, 1.0, 0.5, -0.5, -0.0, math.nan]})
s = table.to_pandas()["a"]
np_result = np.arcsinh(s)

arcs = table.eval_expression_list([col("a").arcsinh()])
assert (
all(
x - y < 1.0e-10 or (x is None and y is None) or (math.isnan(x) and math.isnan(y))
for x, y in zip(arcs.get_column("a").to_pylist(), np_result.to_list())
)
is True
)


def test_table_numeric_arccosh() -> None:
table = MicroPartition.from_pydict({"a": [1.0, 2.0, 1.5, math.nan]})
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

And here, if you could also test that values x < 1 evaluate to NaN

s = table.to_pandas()["a"]
np_result = np.arccosh(s)

arcc = table.eval_expression_list([col("a").arccosh()])
assert (
all(
x - y < 1.0e-10 or (x is None and y is None) or (math.isnan(x) and math.isnan(y))
for x, y in zip(arcc.get_column("a").to_pylist(), np_result.to_list())
)
is True
)


def test_table_numeric_round() -> None:
from decimal import ROUND_HALF_UP, Decimal

Expand Down
Loading