use num_bigint::BigInt;
use num_rational::Ratio;
use num_traits::{One, Signed};
use rustc_hash::FxHashMap;
use crate::api::expr::{Expr, Sort};
use crate::base::arena::Arena;
use crate::base::errors::SymplexError;
use crate::base::node::{ExprId, ExprNode};
use crate::base::walk;
use crate::output::common::display_sort_key;
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub(crate) enum Target {
Python,
NumPy,
Julia,
}
impl Target {
fn name(self) -> &'static str {
match self {
Target::Python => "to_python",
Target::NumPy => "to_numpy",
Target::Julia => "to_julia",
}
}
fn pow_op(self) -> &'static str {
match self {
Target::Python | Target::NumPy => "**",
Target::Julia => "^",
}
}
fn and_op(self) -> &'static str {
match self {
Target::Python => " and ",
Target::NumPy => "",
Target::Julia => " && ",
}
}
fn or_op(self) -> &'static str {
match self {
Target::Python => " or ",
Target::NumPy => "",
Target::Julia => " || ",
}
}
fn unary(self, node: &ExprNode) -> Option<&'static str> {
use ExprNode as N;
let py = |m: &'static str, np: &'static str, jl: &'static str| match self {
Target::Python => m,
Target::NumPy => np,
Target::Julia => jl,
};
Some(match node {
N::Sin(_) => py("math.sin", "numpy.sin", "sin"),
N::Cos(_) => py("math.cos", "numpy.cos", "cos"),
N::Tan(_) => py("math.tan", "numpy.tan", "tan"),
N::Asin(_) => py("math.asin", "numpy.arcsin", "asin"),
N::Acos(_) => py("math.acos", "numpy.arccos", "acos"),
N::Atan(_) => py("math.atan", "numpy.arctan", "atan"),
N::Sinh(_) => py("math.sinh", "numpy.sinh", "sinh"),
N::Cosh(_) => py("math.cosh", "numpy.cosh", "cosh"),
N::Tanh(_) => py("math.tanh", "numpy.tanh", "tanh"),
N::Asinh(_) => py("math.asinh", "numpy.arcsinh", "asinh"),
N::Acosh(_) => py("math.acosh", "numpy.arccosh", "acosh"),
N::Atanh(_) => py("math.atanh", "numpy.arctanh", "atanh"),
N::Exp(_) => py("math.exp", "numpy.exp", "exp"),
N::Ln(_) => py("math.log", "numpy.log", "log"),
N::Abs(_) => py("abs", "abs", "abs"),
N::Floor(_) => py("math.floor", "numpy.floor", "floor"),
N::Ceiling(_) => py("math.ceil", "numpy.ceil", "ceil"),
N::Sign(_) => match self {
Target::Python => return None,
Target::NumPy => "numpy.sign",
Target::Julia => "sign",
},
N::Gamma(_) if self == Target::Python => "math.gamma",
N::LogGamma(_) if self == Target::Python => "math.lgamma",
N::Erf(_) if self == Target::Python => "math.erf",
N::Erfc(_) if self == Target::Python => "math.erfc",
N::Factorial(_) if self == Target::Python => "math.factorial",
_ => return None,
})
}
fn constant(self, node: &ExprNode) -> Option<&'static str> {
let py = |m: &'static str, np: &'static str, jl: &'static str| match self {
Target::Python => m,
Target::NumPy => np,
Target::Julia => jl,
};
Some(match node {
ExprNode::Pi => py("math.pi", "numpy.pi", "pi"),
ExprNode::E => py("math.e", "numpy.e", "ℯ"),
ExprNode::Infinity => py("math.inf", "numpy.inf", "Inf"),
ExprNode::NegInfinity => py("(-math.inf)", "(-numpy.inf)", "(-Inf)"),
ExprNode::NaN => py("math.nan", "numpy.nan", "NaN"),
ExprNode::EulerGamma => "0.5772156649015329",
ExprNode::Catalan => "0.915965594177219",
ExprNode::GoldenRatio => "1.618033988749895",
_ => return None,
})
}
}
const PREC_OR: u8 = 1;
const PREC_AND: u8 = 2;
const PREC_NOT: u8 = 3;
const PREC_REL: u8 = 4;
const PREC_ADD: u8 = 5;
const PREC_MUL: u8 = 6;
const PREC_NEG: u8 = 7;
const PREC_POW: u8 = 8;
const PREC_ATOM: u8 = 9;
#[derive(Clone)]
struct Rendered {
text: String,
prec: u8,
}
impl Rendered {
fn atom(text: impl Into<String>) -> Self {
Rendered {
text: text.into(),
prec: PREC_ATOM,
}
}
fn new(text: String, prec: u8) -> Self {
Rendered { text, prec }
}
fn at(&self, min_prec: u8) -> String {
if self.prec < min_prec {
format!("({})", self.text)
} else {
self.text.clone()
}
}
}
type Cache = FxHashMap<ExprId, Rendered>;
struct Emitter<'a> {
arena: &'a Arena,
target: Target,
params: Option<&'a [&'a str]>,
cse_slots: &'a FxHashMap<ExprId, usize>,
}
impl Emitter<'_> {
fn unsupported(&self, what: &str) -> SymplexError {
SymplexError::NotImplemented(format!(
"{}: cannot generate code for `{what}`",
self.target.name()
))
}
fn cached(&self, cache: &Cache, id: ExprId) -> Result<Rendered, SymplexError> {
cache
.get(&id)
.cloned()
.ok_or_else(|| self.unsupported("an unrendered sub-expression"))
}
fn number(&self, r: &Ratio<BigInt>) -> Rendered {
if r.is_integer() {
if r.is_negative() {
Rendered::new(r.numer().to_string(), PREC_NEG)
} else {
Rendered::atom(r.numer().to_string())
}
} else {
Rendered::atom(format!("({}/{})", r.numer(), r.denom()))
}
}
fn call(&self, name: &str, args: &[&Rendered]) -> Rendered {
let list: Vec<String> = args.iter().map(|a| a.at(0)).collect();
Rendered::atom(format!("{name}({})", list.join(", ")))
}
fn product(
&self,
coeff: Option<&Ratio<BigInt>>,
factors: &[ExprId],
cache: &Cache,
) -> Result<Rendered, SymplexError> {
let mut numer: Vec<String> = Vec::new();
let mut denom: Vec<String> = Vec::new();
if let Some(c) = coeff {
if !c.numer().is_one() {
numer.push(c.numer().to_string());
}
if !c.denom().is_one() {
denom.push(c.denom().to_string());
}
}
for &f in factors {
if let ExprNode::Pow(base, exp) = self.arena.node(f)
&& let Some(r) = self.arena.as_num(*exp)
&& r.is_negative()
{
let base_r = self.cached(cache, *base)?;
let pos = -r.clone();
if pos.is_one() {
denom.push(base_r.at(PREC_MUL + 1));
} else {
let e = self.number(&pos);
denom.push(self.pow_text(&base_r, &e));
}
continue;
}
numer.push(self.cached(cache, f)?.at(PREC_MUL + 1));
}
let numer_text = if numer.is_empty() {
"1".to_string()
} else {
numer.join("*")
};
if denom.is_empty() {
return Ok(if numer.len() <= 1 {
Rendered::atom(numer_text)
} else {
Rendered::new(numer_text, PREC_MUL)
});
}
let denom_text = if denom.len() == 1 {
denom.remove(0)
} else {
format!("({})", denom.join("*"))
};
Ok(Rendered::new(
format!("{numer_text}/{denom_text}"),
PREC_MUL,
))
}
fn pow_text(&self, base: &Rendered, exp: &Rendered) -> String {
format!(
"{}{}{}",
base.at(PREC_ATOM),
self.target.pow_op(),
exp.at(PREC_POW)
)
}
fn render_pow(
&self,
base: ExprId,
exp: ExprId,
cache: &Cache,
) -> Result<Rendered, SymplexError> {
let b = self.cached(cache, base)?;
if let Some(r) = self.arena.as_num(exp) {
let one = BigInt::one();
if *r.numer() == one && *r.denom() == BigInt::from(2) {
let sqrt = match self.target {
Target::Python => "math.sqrt",
Target::NumPy => "numpy.sqrt",
Target::Julia => "sqrt",
};
return Ok(self.call(sqrt, &[&b]));
}
if *r.numer() == one
&& *r.denom() == BigInt::from(3)
&& let Some(cbrt) = match self.target {
Target::Python => None,
Target::NumPy => Some("numpy.cbrt"),
Target::Julia => Some("cbrt"),
}
{
return Ok(self.call(cbrt, &[&b]));
}
if r.is_integer() && *r.numer() == -one {
return Ok(Rendered::new(format!("1/{}", b.at(PREC_MUL + 1)), PREC_MUL));
}
let e = self.number(r);
return Ok(Rendered::new(self.pow_text(&b, &e), PREC_POW));
}
let e = self.cached(cache, exp)?;
Ok(Rendered::new(self.pow_text(&b, &e), PREC_POW))
}
fn signed_term(&self, id: ExprId, cache: &Cache) -> Result<(bool, Rendered), SymplexError> {
match self.arena.node(id) {
ExprNode::Num(nid) if self.arena.num(*nid).is_negative() => {
Ok((true, self.number(&-self.arena.num(*nid).clone())))
}
ExprNode::Neg(inner) => Ok((true, self.cached(cache, *inner)?)),
ExprNode::Mul(children) => {
if let Some(&first) = children.first()
&& let Some(c) = self.arena.as_num(first)
&& c.is_negative()
&& children.len() > 1
{
let body = self.product(Some(&(-c.clone())), &children[1..], cache)?;
return Ok((true, body));
}
Ok((false, self.cached(cache, id)?))
}
_ => Ok((false, self.cached(cache, id)?)),
}
}
fn render_add(&self, children: &[ExprId], cache: &Cache) -> Result<Rendered, SymplexError> {
if children.is_empty() {
return Ok(Rendered::atom("0"));
}
let mut ordered: Vec<ExprId> = children.to_vec();
ordered.sort_by_key(|a| display_sort_key(self.arena, *a));
let mut text = String::new();
for (i, &c) in ordered.iter().enumerate() {
let (negative, body) = self.signed_term(c, cache)?;
if i == 0 {
if negative {
text.push('-');
text.push_str(&body.at(PREC_MUL));
} else {
text.push_str(&body.at(PREC_ADD));
}
} else {
text.push_str(if negative { " - " } else { " + " });
text.push_str(&body.at(PREC_ADD + 1));
}
}
let prec = if ordered.len() == 1 {
let (negative, body) = self.signed_term(ordered[0], cache)?;
if negative { PREC_NEG } else { body.prec }
} else {
PREC_ADD
};
Ok(Rendered::new(text, prec))
}
fn render_mul(&self, children: &[ExprId], cache: &Cache) -> Result<Rendered, SymplexError> {
if children.is_empty() {
return Ok(Rendered::atom("1"));
}
if let Some(&first) = children.first()
&& let Some(c) = self.arena.as_num(first)
&& children.len() > 1
{
if c.is_negative() {
let body = self.product(Some(&(-c.clone())), &children[1..], cache)?;
return Ok(Rendered::new(format!("-{}", body.at(PREC_MUL)), PREC_NEG));
}
return self.product(Some(c), &children[1..], cache);
}
self.product(None, children, cache)
}
fn relation(&self, op: &str, lhs: &Rendered, rhs: &Rendered) -> Rendered {
Rendered::new(
format!("{} {op} {}", lhs.at(PREC_ADD), rhs.at(PREC_ADD)),
PREC_REL,
)
}
fn np_relation(&self, func: &str, lhs: &Rendered, rhs: &Rendered) -> Rendered {
self.call(&format!("numpy.{func}"), &[lhs, rhs])
}
fn connective(
&self,
op: &str,
np_func: &str,
prec: u8,
children: &[ExprId],
cache: &Cache,
) -> Result<Rendered, SymplexError> {
let parts: Vec<Rendered> = children
.iter()
.map(|c| self.cached(cache, *c))
.collect::<Result<_, _>>()?;
if self.target == Target::NumPy {
let mut acc = parts
.first()
.cloned()
.unwrap_or_else(|| Rendered::atom("True"));
for p in &parts[1.min(parts.len())..] {
acc = self.call(&format!("numpy.{np_func}"), &[&acc, p]);
}
return Ok(acc);
}
let text: Vec<String> = parts.iter().map(|p| p.at(prec + 1)).collect();
Ok(Rendered::new(text.join(op), prec))
}
fn render_piecewise(
&self,
pieces: &[(ExprId, ExprId)],
cache: &Cache,
) -> Result<Rendered, SymplexError> {
let nan = match self.target {
Target::Python => "math.nan",
Target::NumPy => "numpy.nan",
Target::Julia => "NaN",
};
if self.target == Target::NumPy {
let mut conds = Vec::new();
let mut vals = Vec::new();
for &(v, c) in pieces {
let c_r = if c == self.arena.bool_true() {
Rendered::atom("True")
} else {
self.cached(cache, c)?
};
conds.push(c_r.at(0));
vals.push(self.cached(cache, v)?.at(0));
}
return Ok(Rendered::atom(format!(
"numpy.select([{}], [{}], default={nan})",
conds.join(", "),
vals.join(", ")
)));
}
let mut acc = nan.to_string();
let mut acc_is_cond = false;
for (i, &(v, c)) in pieces.iter().enumerate().rev() {
let v_r = self.cached(cache, v)?;
if i + 1 == pieces.len() && c == self.arena.bool_true() {
acc = v_r.at(PREC_OR);
continue;
}
let c_r = self.cached(cache, c)?;
let rest = if acc_is_cond { format!("({acc})") } else { acc };
acc = match self.target {
Target::Julia => format!("{} ? {} : {rest}", c_r.at(PREC_OR), v_r.at(PREC_OR)),
_ => format!("{} if {} else {rest}", v_r.at(PREC_OR), c_r.at(PREC_OR)),
};
acc_is_cond = true;
}
Ok(Rendered::atom(format!("({acc})")))
}
fn render(&self, root: ExprId) -> Result<Rendered, SymplexError> {
let arena = self.arena;
let order = walk::post_order_ids(arena, root);
let mut cache: Cache = FxHashMap::default();
for &id in &order {
let node = arena.node(id);
let child = |c: &ExprId| self.cached(&cache, *c);
let rendered = match node {
ExprNode::Num(nid) => self.number(arena.num(*nid)),
ExprNode::Symbol(sid) => {
if let Some(&slot) = self.cse_slots.get(&id) {
Rendered::atom(format!("t{slot}"))
} else {
let name = arena.symbol_name(*sid);
if let Some(params) = self.params
&& !params.contains(&name)
{
return Err(SymplexError::FreeSymbol {
name: name.to_string(),
});
}
Rendered::atom(name)
}
}
ExprNode::PhysicalConstant(_, value) => child(value)?,
ExprNode::Pi
| ExprNode::E
| ExprNode::Infinity
| ExprNode::NegInfinity
| ExprNode::NaN
| ExprNode::EulerGamma
| ExprNode::Catalan
| ExprNode::GoldenRatio => {
let text = self
.target
.constant(node)
.ok_or_else(|| self.unsupported("a named constant"))?;
Rendered::atom(text)
}
ExprNode::BoolTrue => Rendered::atom(match self.target {
Target::Julia => "true",
_ => "True",
}),
ExprNode::BoolFalse => Rendered::atom(match self.target {
Target::Julia => "false",
_ => "False",
}),
ExprNode::Add(children) => self.render_add(children, &cache)?,
ExprNode::Mul(children) => self.render_mul(children, &cache)?,
ExprNode::Neg(inner) => {
let r = child(inner)?;
Rendered::new(format!("-{}", r.at(PREC_MUL)), PREC_NEG)
}
ExprNode::Pow(base, exp) => self.render_pow(*base, *exp, &cache)?,
ExprNode::Sign(a) if self.target == Target::Python => {
let r = child(a)?;
Rendered::atom(format!(
"(0.0 if {} == 0 else math.copysign(1, {}))",
r.at(PREC_ADD),
r.at(0)
))
}
ExprNode::Heaviside(a) => {
let r = child(a)?;
match self.target {
Target::Python => Rendered::atom(format!(
"(0.0 if {a} < 0 else (0.5 if {a} == 0 else 1.0))",
a = r.at(PREC_ADD)
)),
Target::NumPy => {
self.call("numpy.heaviside", &[&r, &Rendered::atom("0.5")])
}
Target::Julia => Rendered::atom(format!(
"({a} < 0 ? 0.0 : ({a} == 0 ? 0.5 : 1.0))",
a = r.at(PREC_ADD)
)),
}
}
ExprNode::Atan2(y, x) => {
let name = match self.target {
Target::Python => "math.atan2",
Target::NumPy => "numpy.arctan2",
Target::Julia => "atan",
};
self.call(name, &[&child(y)?, &child(x)?])
}
ExprNode::Min(args) | ExprNode::Max(args) => {
let is_min = matches!(node, ExprNode::Min(_));
let parts: Vec<Rendered> = args.iter().map(child).collect::<Result<_, _>>()?;
if parts.is_empty() {
return Err(self.unsupported("an empty min/max"));
}
if self.target == Target::NumPy {
let f = if is_min {
"numpy.minimum"
} else {
"numpy.maximum"
};
let mut acc = parts[0].clone();
for p in &parts[1..] {
acc = self.call(f, &[&acc, p]);
}
acc
} else {
let refs: Vec<&Rendered> = parts.iter().collect();
self.call(if is_min { "min" } else { "max" }, &refs)
}
}
ExprNode::Piecewise(pieces) => self.render_piecewise(pieces, &cache)?,
ExprNode::Gt(a, b) => match self.target {
Target::NumPy => self.np_relation("greater", &child(a)?, &child(b)?),
_ => self.relation(">", &child(a)?, &child(b)?),
},
ExprNode::Ge(a, b) => match self.target {
Target::NumPy => self.np_relation("greater_equal", &child(a)?, &child(b)?),
_ => self.relation(">=", &child(a)?, &child(b)?),
},
ExprNode::Eq_(a, b) => match self.target {
Target::NumPy => self.np_relation("equal", &child(a)?, &child(b)?),
_ => self.relation("==", &child(a)?, &child(b)?),
},
ExprNode::Ne(a, b) => match self.target {
Target::NumPy => self.np_relation("not_equal", &child(a)?, &child(b)?),
_ => self.relation("!=", &child(a)?, &child(b)?),
},
ExprNode::And(children) => self.connective(
self.target.and_op(),
"logical_and",
PREC_AND,
children,
&cache,
)?,
ExprNode::Or(children) => {
self.connective(self.target.or_op(), "logical_or", PREC_OR, children, &cache)?
}
ExprNode::Not(a) => {
let r = child(a)?;
match self.target {
Target::Python => {
Rendered::new(format!("not {}", r.at(PREC_REL + 1)), PREC_NOT)
}
Target::NumPy => self.call("numpy.logical_not", &[&r]),
Target::Julia => Rendered::new(format!("!{}", r.at(PREC_ATOM)), PREC_NEG),
}
}
ExprNode::DiracDelta(_) => Rendered::atom("0.0"),
ExprNode::Sin(a)
| ExprNode::Cos(a)
| ExprNode::Tan(a)
| ExprNode::Asin(a)
| ExprNode::Acos(a)
| ExprNode::Atan(a)
| ExprNode::Sinh(a)
| ExprNode::Cosh(a)
| ExprNode::Tanh(a)
| ExprNode::Asinh(a)
| ExprNode::Acosh(a)
| ExprNode::Atanh(a)
| ExprNode::Exp(a)
| ExprNode::Ln(a)
| ExprNode::Abs(a)
| ExprNode::Floor(a)
| ExprNode::Ceiling(a)
| ExprNode::Sign(a)
| ExprNode::Gamma(a)
| ExprNode::LogGamma(a)
| ExprNode::Erf(a)
| ExprNode::Erfc(a)
| ExprNode::Factorial(a) => {
let name = self
.target
.unary(node)
.ok_or_else(|| self.unsupported(&describe(node)))?;
self.call(name, &[&child(a)?])
}
ExprNode::Apply(sid, _) => {
return Err(
self.unsupported(&format!("the function `{}`", arena.symbol_name(*sid)))
);
}
other => return Err(self.unsupported(&describe(other))),
};
cache.insert(id, rendered);
}
self.cached(&cache, root)
}
}
fn describe(node: &ExprNode) -> String {
let dbg = format!("{node:?}");
dbg.split(['(', ' ']).next().unwrap_or("node").to_string()
}
pub(crate) fn to_expr_code(
arena: &Arena,
expr: ExprId,
target: Target,
) -> Result<String, SymplexError> {
let slots = FxHashMap::default();
let em = Emitter {
arena,
target,
params: None,
cse_slots: &slots,
};
Ok(em.render(expr)?.text)
}
pub(crate) fn to_fn_code(
arena: &mut Arena,
expr: ExprId,
name: &str,
args: &[&str],
target: Target,
) -> Result<String, SymplexError> {
let cse = crate::output::cse::cse(arena, expr);
let mut cse_slots: FxHashMap<ExprId, usize> = FxHashMap::default();
for (i, (name_id, _)) in cse.bindings.iter().enumerate() {
cse_slots.insert(*name_id, i);
}
let em = Emitter {
arena,
target,
params: Some(args),
cse_slots: &cse_slots,
};
let mut lines: Vec<String> = Vec::new();
for (i, (_, value)) in cse.bindings.iter().enumerate() {
lines.push(format!(" t{i} = {}", em.render(*value)?.text));
}
let result = em.render(cse.expr)?.text;
let params = args.join(", ");
let mut out = String::new();
match target {
Target::Python | Target::NumPy => {
out.push_str(&format!("def {name}({params}):\n"));
for l in &lines {
out.push_str(l);
out.push('\n');
}
out.push_str(&format!(" return {result}\n"));
}
Target::Julia => {
out.push_str(&format!("function {name}({params})\n"));
for l in &lines {
out.push_str(l);
out.push('\n');
}
out.push_str(&format!(" return {result}\nend\n"));
}
}
Ok(out)
}
impl<S: Sort> Expr<S> {
pub fn to_python(&self) -> Result<String, SymplexError> {
let inner = self.inner.read();
to_expr_code(&inner.arena, self.raw_id(), Target::Python)
}
pub fn to_numpy(&self) -> Result<String, SymplexError> {
let inner = self.inner.read();
to_expr_code(&inner.arena, self.raw_id(), Target::NumPy)
}
pub fn to_julia(&self) -> Result<String, SymplexError> {
let inner = self.inner.read();
to_expr_code(&inner.arena, self.raw_id(), Target::Julia)
}
pub fn to_python_fn(&self, name: &str, args: &[&str]) -> Result<String, SymplexError> {
let mut guard = self.inner.write();
to_fn_code(&mut guard.arena, self.raw_id(), name, args, Target::Python)
}
pub fn to_numpy_fn(&self, name: &str, args: &[&str]) -> Result<String, SymplexError> {
let mut guard = self.inner.write();
to_fn_code(&mut guard.arena, self.raw_id(), name, args, Target::NumPy)
}
pub fn to_julia_fn(&self, name: &str, args: &[&str]) -> Result<String, SymplexError> {
let mut guard = self.inner.write();
to_fn_code(&mut guard.arena, self.raw_id(), name, args, Target::Julia)
}
}
#[cfg(test)]
mod tests {
use crate::prelude::*;
fn py(s: &str) -> String {
let ctx = Context::new();
ctx.parse(s).unwrap().to_python().unwrap()
}
fn np(s: &str) -> String {
let ctx = Context::new();
ctx.parse(s).unwrap().to_numpy().unwrap()
}
fn jl(s: &str) -> String {
let ctx = Context::new();
ctx.parse(s).unwrap().to_julia().unwrap()
}
#[test]
fn python_arithmetic() {
assert_eq!(py("x^2 + 1"), "x**2 + 1");
assert_eq!(py("x/2"), "x/2");
assert_eq!(py("1/x"), "1/x");
assert_eq!(py("2/3"), "(2/3)");
assert_eq!(py("x/y"), "x/y");
assert_eq!(py("x/(y*z)"), "x/(y*z)");
assert_eq!(py("x/(y+1)"), "x/(y + 1)");
assert_eq!(py("-x"), "-x");
assert_eq!(py("-x^2"), "-x**2");
assert_eq!(py("(-x)^y"), "(-x)**y");
assert_eq!(py("x - y"), "x - y");
assert_eq!(py("1 - x/2"), "-x/2 + 1");
assert_eq!(py("x^(-2)"), "x**(-2)");
assert_eq!(py("x^(3/2)"), "x**(3/2)");
assert_eq!(py("x^y"), "x**y");
assert_eq!(py("(x+1)^2"), "(x + 1)**2");
assert_eq!(py("2^x"), "2**x");
assert_eq!(py("x^y^z"), "x**y**z");
assert_eq!(py("(x^y)^z"), "(x**y)**z");
assert_eq!(py("x^(y+1)"), "x**(y + 1)");
assert_eq!(py("e^x"), "math.exp(x)");
assert_eq!(py("pi*e"), "math.pi*math.e");
assert_eq!(py("inf"), "math.inf");
}
#[test]
fn python_functions() {
assert_eq!(py("sin(x)^2 + exp(x)"), "math.sin(x)**2 + math.exp(x)");
assert_eq!(py("sqrt(x)"), "math.sqrt(x)");
assert_eq!(py("cbrt(x)"), "x**(1/3)");
assert_eq!(py("abs(x)"), "abs(x)");
assert_eq!(py("floor(x) + ceil(y)"), "math.floor(x) + math.ceil(y)");
assert_eq!(py("gamma(x) + erf(x)"), "math.gamma(x) + math.erf(x)");
assert_eq!(py("loggamma(x)"), "math.lgamma(x)");
assert_eq!(py("atan2(y, x)"), "math.atan2(y, x)");
assert_eq!(py("min(x, y)"), "min(x, y)");
assert_eq!(py("x!"), "math.factorial(x)");
assert_eq!(py("ln(x)"), "math.log(x)");
assert_eq!(py("sign(x)"), "(0.0 if x == 0 else math.copysign(1, x))");
}
#[test]
fn python_logic_and_piecewise() {
let ctx = Context::new();
let (x, y) = (ctx.symbol("x"), ctx.symbol("y"));
let zero = ctx.int(0);
assert_eq!(x.gt(&zero).to_python().unwrap(), "x > 0");
assert_eq!(
x.gt(&zero).and(&y.ge(&zero)).to_python().unwrap(),
"x > 0 and y >= 0"
);
assert_eq!(
x.gt(&zero)
.and(&y.ge(&zero))
.or(&x.eq_expr(&y))
.to_python()
.unwrap(),
"x > 0 and y >= 0 or x == y"
);
assert_eq!(
x.gt(&zero)
.or(&y.ge(&zero))
.and(&x.ne_expr(&y))
.to_python()
.unwrap(),
"(x > 0 or y >= 0) and x != y"
);
assert_eq!(x.gt(&zero).not().to_python().unwrap(), "not (x > 0)");
let pw = Ex::piecewise(&[(&x, &x.gt(&zero)), (&(-&x), &x.le(&zero))]);
assert_eq!(
pw.to_python().unwrap(),
"(x if x > 0 else (-x if 0 >= x else math.nan))"
);
let pw2 = Ex::piecewise(&[(&x.powi(2), &x.lt(&zero)), (&x, &x.ge(&zero))]);
assert_eq!(
pw2.to_python().unwrap(),
"(x**2 if 0 > x else (x if x >= 0 else math.nan))"
);
}
#[test]
fn numpy_and_julia() {
assert_eq!(np("sin(x)"), "numpy.sin(x)");
assert_eq!(np("x^2 + 1"), "x**2 + 1");
assert_eq!(np("pi"), "numpy.pi");
assert_eq!(np("cbrt(x)"), "numpy.cbrt(x)");
assert_eq!(np("min(x, y, z)"), "numpy.minimum(numpy.minimum(x, y), z)");
assert_eq!(np("sign(x)"), "numpy.sign(x)");
assert_eq!(np("asin(x)"), "numpy.arcsin(x)");
let ctx = Context::new();
let (x, y) = (ctx.symbol("x"), ctx.symbol("y"));
let zero = ctx.int(0);
assert_eq!(
x.gt(&zero).and(&y.ge(&zero)).to_numpy().unwrap(),
"numpy.logical_and(numpy.greater(x, 0), numpy.greater_equal(y, 0))"
);
let pw = Ex::piecewise(&[(&x, &x.gt(&zero)), (&(-&x), &x.le(&zero))]);
assert_eq!(
pw.to_numpy().unwrap(),
"numpy.select([numpy.greater(x, 0), numpy.greater_equal(0, x)], [x, -x], default=numpy.nan)"
);
assert!(ctx.parse("gamma(x)").unwrap().to_numpy().is_err());
assert_eq!(jl("sin(x)^2 + exp(x)"), "sin(x)^2 + exp(x)");
assert_eq!(jl("x^2 + pi"), "x^2 + pi");
assert_eq!(jl("x/2"), "x/2");
assert_eq!(jl("2/3*x"), "2*x/3");
assert_eq!(jl("atan2(y, x)"), "atan(y, x)");
assert_eq!(
x.gt(&zero).and(&y.ge(&zero)).to_julia().unwrap(),
"x > 0 && y >= 0"
);
assert_eq!(x.gt(&zero).not().to_julia().unwrap(), "!(x > 0)");
assert_eq!(pw.to_julia().unwrap(), "(x > 0 ? x : (0 >= x ? -x : NaN))");
assert!(ctx.parse("gamma(x)").unwrap().to_julia().is_err());
}
#[test]
fn functions_with_cse_and_errors() {
let ctx = Context::new();
let (x, y) = (ctx.symbol("x"), ctx.symbol("y"));
let f = &x.sin().powi(2) + &x.sin() * &y;
assert_eq!(
f.to_python_fn("f", &["x", "y"]).unwrap(),
"def f(x, y):\n t0 = math.sin(x)\n return t0**2 + t0*y\n"
);
assert_eq!(
f.to_julia_fn("f", &["x", "y"]).unwrap(),
"function f(x, y)\n t0 = sin(x)\n return t0^2 + t0*y\nend\n"
);
assert!(matches!(
f.to_python_fn("f", &["x"]),
Err(SymplexError::FreeSymbol { .. })
));
assert!(matches!(
x.bessel_j(&ctx.int(0)).to_python(),
Err(SymplexError::NotImplemented(_))
));
assert!(matches!(
ctx.parse("Integral(x, x)").unwrap().to_python(),
Err(SymplexError::NotImplemented(_))
));
assert!(matches!(
ctx.parse("I").unwrap().to_python(),
Err(SymplexError::NotImplemented(_))
));
}
}