mp_obj_t mp_obj_complex_binary_op(int op, mp_float_t lhs_real, mp_float_t lhs_imag, mp_obj_t rhs_in) { mp_float_t rhs_real, rhs_imag; mp_obj_get_complex(rhs_in, &rhs_real, &rhs_imag); // can be any type, this function will convert to float (if possible) switch (op) { case RT_BINARY_OP_ADD: case RT_BINARY_OP_INPLACE_ADD: lhs_real += rhs_real; lhs_imag += rhs_imag; break; case RT_BINARY_OP_SUBTRACT: case RT_BINARY_OP_INPLACE_SUBTRACT: lhs_real -= rhs_real; lhs_imag -= rhs_imag; break; case RT_BINARY_OP_MULTIPLY: case RT_BINARY_OP_INPLACE_MULTIPLY: { mp_float_t real = lhs_real * rhs_real - lhs_imag * rhs_imag; lhs_imag = lhs_real * rhs_imag + lhs_imag * rhs_real; lhs_real = real; break; } /* TODO floor(?) the value case RT_BINARY_OP_FLOOR_DIVIDE: case RT_BINARY_OP_INPLACE_FLOOR_DIVIDE: val = lhs_val / rhs_val; break; */ /* TODO case RT_BINARY_OP_TRUE_DIVIDE: case RT_BINARY_OP_INPLACE_TRUE_DIVIDE: val = lhs_val / rhs_val; break; */ return NULL; // op not supported } return mp_obj_new_complex(lhs_real, lhs_imag); }
mp_obj_t mp_cmath_sqrt(mp_obj_t z_obj) { mp_float_t real, imag; mp_obj_get_complex(z_obj, &real, &imag); mp_float_t sqrt_abs = MICROPY_FLOAT_C_FUN(pow)(real*real + imag*imag, 0.25); mp_float_t theta = 0.5 * MICROPY_FLOAT_C_FUN(atan2)(imag, real); return mp_obj_new_complex(sqrt_abs * cos(theta), sqrt_abs * sin(theta)); }
mp_obj_t mp_cmath_polar(mp_obj_t z_obj) { mp_float_t real, imag; mp_obj_get_complex(z_obj, &real, &imag); mp_obj_t tuple[2] = { mp_obj_new_float(MICROPY_FLOAT_C_FUN(sqrt)(real*real + imag*imag)), mp_obj_new_float(MICROPY_FLOAT_C_FUN(atan2)(imag, real)), }; return mp_obj_new_tuple(2, tuple); }
mp_obj_t mp_obj_complex_binary_op(int op, mp_float_t lhs_real, mp_float_t lhs_imag, mp_obj_t rhs_in) { mp_float_t rhs_real, rhs_imag; mp_obj_get_complex(rhs_in, &rhs_real, &rhs_imag); // can be any type, this function will convert to float (if possible) switch (op) { case MP_BINARY_OP_ADD: case MP_BINARY_OP_INPLACE_ADD: lhs_real += rhs_real; lhs_imag += rhs_imag; break; case MP_BINARY_OP_SUBTRACT: case MP_BINARY_OP_INPLACE_SUBTRACT: lhs_real -= rhs_real; lhs_imag -= rhs_imag; break; case MP_BINARY_OP_MULTIPLY: case MP_BINARY_OP_INPLACE_MULTIPLY: { mp_float_t real; multiply: real = lhs_real * rhs_real - lhs_imag * rhs_imag; lhs_imag = lhs_real * rhs_imag + lhs_imag * rhs_real; lhs_real = real; break; } case MP_BINARY_OP_FLOOR_DIVIDE: case MP_BINARY_OP_INPLACE_FLOOR_DIVIDE: nlr_raise(mp_obj_new_exception_msg_varg(&mp_type_TypeError, "can't do truncated division of a complex number")); case MP_BINARY_OP_TRUE_DIVIDE: case MP_BINARY_OP_INPLACE_TRUE_DIVIDE: if (rhs_imag == 0) { if (rhs_real == 0) { nlr_raise(mp_obj_new_exception_msg(&mp_type_ZeroDivisionError, "complex division by zero")); } lhs_real /= rhs_real; lhs_imag /= rhs_real; } else if (rhs_real == 0) { mp_float_t real = lhs_imag / rhs_imag; lhs_imag = -lhs_real / rhs_imag; lhs_real = real; } else { mp_float_t rhs_len_sq = rhs_real*rhs_real + rhs_imag*rhs_imag; rhs_real /= rhs_len_sq; rhs_imag /= -rhs_len_sq; goto multiply; } break; case MP_BINARY_OP_POWER: case MP_BINARY_OP_INPLACE_POWER: { // z1**z2 = exp(z2*ln(z1)) // = exp(z2*(ln(|z1|)+i*arg(z1))) // = exp( (x2*ln1 - y2*arg1) + i*(y2*ln1 + x2*arg1) ) // = exp(x3 + i*y3) // = exp(x3)*(cos(y3) + i*sin(y3)) mp_float_t abs1 = MICROPY_FLOAT_C_FUN(sqrt)(lhs_real*lhs_real + lhs_imag*lhs_imag); if (abs1 == 0) { if (rhs_imag == 0) { lhs_real = 1; rhs_real = 0; } else { nlr_raise(mp_obj_new_exception_msg(&mp_type_ZeroDivisionError, "0.0 to a complex power")); } } else { mp_float_t ln1 = MICROPY_FLOAT_C_FUN(log)(abs1); mp_float_t arg1 = MICROPY_FLOAT_C_FUN(atan2)(lhs_imag, lhs_real); mp_float_t x3 = rhs_real * ln1 - rhs_imag * arg1; mp_float_t y3 = rhs_imag * ln1 + rhs_real * arg1; mp_float_t exp_x3 = MICROPY_FLOAT_C_FUN(exp)(x3); lhs_real = exp_x3 * MICROPY_FLOAT_C_FUN(cos)(y3); lhs_imag = exp_x3 * MICROPY_FLOAT_C_FUN(sin)(y3); } break; } case MP_BINARY_OP_EQUAL: return MP_BOOL(lhs_real == rhs_real && lhs_imag == rhs_imag); default: return MP_OBJ_NULL; // op not supported } return mp_obj_new_complex(lhs_real, lhs_imag); }
mp_obj_t mp_cmath_sin(mp_obj_t z_obj) { mp_float_t real, imag; mp_obj_get_complex(z_obj, &real, &imag); return mp_obj_new_complex(MICROPY_FLOAT_C_FUN(sin)(real) * MICROPY_FLOAT_C_FUN(cosh)(imag), MICROPY_FLOAT_C_FUN(cos)(real) * MICROPY_FLOAT_C_FUN(sinh)(imag)); }
mp_obj_t mp_cmath_log10(mp_obj_t z_obj) { mp_float_t real, imag; mp_obj_get_complex(z_obj, &real, &imag); return mp_obj_new_complex(0.5 * MICROPY_FLOAT_C_FUN(log10)(real*real + imag*imag), MICROPY_FLOAT_C_FUN(atan2)(imag, real)); }
mp_obj_t mp_cmath_exp(mp_obj_t z_obj) { mp_float_t real, imag; mp_obj_get_complex(z_obj, &real, &imag); mp_float_t exp_real = MICROPY_FLOAT_C_FUN(exp)(real); return mp_obj_new_complex(exp_real * MICROPY_FLOAT_C_FUN(cos)(imag), exp_real * MICROPY_FLOAT_C_FUN(sin)(imag)); }
mp_obj_t mp_cmath_phase(mp_obj_t z_obj) { mp_float_t real, imag; mp_obj_get_complex(z_obj, &real, &imag); return mp_obj_new_float(MICROPY_FLOAT_C_FUN(atan2)(imag, real)); }