blob: bf6fb51dc558ef5efca65d912f5aace2564afcb9 [file] [log] [blame]
Damien George04b91472014-05-03 23:27:38 +01001/*
Alexander Steffen55f33242017-06-30 09:22:17 +02002 * This file is part of the MicroPython project, http://micropython.org/
Damien George04b91472014-05-03 23:27:38 +01003 *
4 * The MIT License (MIT)
5 *
6 * Copyright (c) 2013, 2014 Damien P. George
7 *
8 * Permission is hereby granted, free of charge, to any person obtaining a copy
9 * of this software and associated documentation files (the "Software"), to deal
10 * in the Software without restriction, including without limitation the rights
11 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
12 * copies of the Software, and to permit persons to whom the Software is
13 * furnished to do so, subject to the following conditions:
14 *
15 * The above copyright notice and this permission notice shall be included in
16 * all copies or substantial portions of the Software.
17 *
18 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
24 * THE SOFTWARE.
25 */
26
Damiend99b0522013-12-21 18:17:45 +000027#include <stdlib.h>
Damien George20beff92014-09-11 22:24:45 +010028#include <stdio.h>
Damiend99b0522013-12-21 18:17:45 +000029#include <assert.h>
30
Damien George51dfcb42015-01-01 20:27:54 +000031#include "py/parsenum.h"
Damien George51dfcb42015-01-01 20:27:54 +000032#include "py/runtime.h"
Damiend99b0522013-12-21 18:17:45 +000033
Paul Sokolovsky3b6f7b92014-06-20 01:48:35 +030034#if MICROPY_PY_BUILTINS_COMPLEX
Damiend99b0522013-12-21 18:17:45 +000035
Damien Georgeae491052014-04-10 20:08:11 +010036#include <math.h>
Damien George51dfcb42015-01-01 20:27:54 +000037#include "py/formatfloat.h"
Damien George8767d072014-03-27 22:17:49 +000038
Damiend99b0522013-12-21 18:17:45 +000039typedef struct _mp_obj_complex_t {
40 mp_obj_base_t base;
41 mp_float_t real;
42 mp_float_t imag;
43} mp_obj_complex_t;
44
Damien George7f9d1d62015-04-09 23:56:15 +010045STATIC void complex_print(const mp_print_t *print, mp_obj_t o_in, mp_print_kind_t kind) {
Damien Georgeff8dd3f2015-01-20 12:47:20 +000046 (void)kind;
Damien George999cedb2015-11-27 17:01:44 +000047 mp_obj_complex_t *o = MP_OBJ_TO_PTR(o_in);
Damien George8767d072014-03-27 22:17:49 +000048#if MICROPY_FLOAT_IMPL == MICROPY_FLOAT_IMPL_FLOAT
Damien George20beff92014-09-11 22:24:45 +010049 char buf[16];
Damien George7a72c0d2017-04-19 18:46:53 +100050 #if MICROPY_OBJ_REPR == MICROPY_OBJ_REPR_C
51 const int precision = 6;
52 #else
stijn861670b2015-05-16 10:54:19 +020053 const int precision = 7;
Damien George7a72c0d2017-04-19 18:46:53 +100054 #endif
Damien George8767d072014-03-27 22:17:49 +000055#else
Damien George20beff92014-09-11 22:24:45 +010056 char buf[32];
stijn861670b2015-05-16 10:54:19 +020057 const int precision = 16;
58#endif
Damiend99b0522013-12-21 18:17:45 +000059 if (o->real == 0) {
stijn861670b2015-05-16 10:54:19 +020060 mp_format_float(o->imag, buf, sizeof(buf), 'g', precision, '\0');
Damien George7f9d1d62015-04-09 23:56:15 +010061 mp_printf(print, "%sj", buf);
Damiend99b0522013-12-21 18:17:45 +000062 } else {
stijn861670b2015-05-16 10:54:19 +020063 mp_format_float(o->real, buf, sizeof(buf), 'g', precision, '\0');
Damien George7f9d1d62015-04-09 23:56:15 +010064 mp_printf(print, "(%s", buf);
stijn709955b2015-05-13 21:35:58 +020065 if (o->imag >= 0 || isnan(o->imag)) {
Damien George7f9d1d62015-04-09 23:56:15 +010066 mp_print_str(print, "+");
Damien Georgef49782f2015-02-02 12:52:14 +000067 }
stijn861670b2015-05-16 10:54:19 +020068 mp_format_float(o->imag, buf, sizeof(buf), 'g', precision, '\0');
Damien George7f9d1d62015-04-09 23:56:15 +010069 mp_printf(print, "%sj)", buf);
Damiend99b0522013-12-21 18:17:45 +000070 }
71}
72
Damien George5b3f0b72016-01-03 15:55:55 +000073STATIC mp_obj_t complex_make_new(const mp_obj_type_t *type_in, size_t n_args, size_t n_kw, const mp_obj_t *args) {
Damien Georgeff8dd3f2015-01-20 12:47:20 +000074 (void)type_in;
Damien Georgeee7a8802014-05-11 18:37:21 +010075 mp_arg_check_num(n_args, n_kw, 0, 2, false);
Damien George20006db2014-01-18 14:10:48 +000076
Damien George71c51812014-01-04 20:21:15 +000077 switch (n_args) {
78 case 0:
79 return mp_obj_new_complex(0, 0);
80
81 case 1:
Damien Georgeeee1e882019-01-30 18:49:52 +110082 if (mp_obj_is_str(args[0])) {
Damien George6e48f7f2014-03-21 11:45:46 +000083 // a string, parse it
Damien George6b341072017-03-25 19:48:18 +110084 size_t l;
Damien George6e48f7f2014-03-21 11:45:46 +000085 const char *s = mp_obj_str_get_data(args[0], &l);
Damien George7d414a12015-02-08 01:57:40 +000086 return mp_parse_num_decimal(s, l, true, true, NULL);
Damien Georgeeee1e882019-01-30 18:49:52 +110087 } else if (mp_obj_is_type(args[0], &mp_type_complex)) {
Damien George6e48f7f2014-03-21 11:45:46 +000088 // a complex, just return it
Damien George71c51812014-01-04 20:21:15 +000089 return args[0];
90 } else {
Damien George6e48f7f2014-03-21 11:45:46 +000091 // something else, try to cast it to a complex
Damien George71c51812014-01-04 20:21:15 +000092 return mp_obj_new_complex(mp_obj_get_float(args[0]), 0);
93 }
94
Damien Georgeee7a8802014-05-11 18:37:21 +010095 case 2:
96 default: {
Damien George71c51812014-01-04 20:21:15 +000097 mp_float_t real, imag;
Damien Georgeeee1e882019-01-30 18:49:52 +110098 if (mp_obj_is_type(args[0], &mp_type_complex)) {
Damien George20006db2014-01-18 14:10:48 +000099 mp_obj_complex_get(args[0], &real, &imag);
Damien George71c51812014-01-04 20:21:15 +0000100 } else {
Damien George20006db2014-01-18 14:10:48 +0000101 real = mp_obj_get_float(args[0]);
Damien George71c51812014-01-04 20:21:15 +0000102 imag = 0;
103 }
Damien Georgeeee1e882019-01-30 18:49:52 +1100104 if (mp_obj_is_type(args[1], &mp_type_complex)) {
Damien George71c51812014-01-04 20:21:15 +0000105 mp_float_t real2, imag2;
Damien George20006db2014-01-18 14:10:48 +0000106 mp_obj_complex_get(args[1], &real2, &imag2);
Damien George71c51812014-01-04 20:21:15 +0000107 real -= imag2;
108 imag += real2;
109 } else {
Damien George20006db2014-01-18 14:10:48 +0000110 imag += mp_obj_get_float(args[1]);
Damien George71c51812014-01-04 20:21:15 +0000111 }
112 return mp_obj_new_complex(real, imag);
113 }
Damien George71c51812014-01-04 20:21:15 +0000114 }
115}
116
Damien George58321dd2017-08-29 13:04:01 +1000117STATIC mp_obj_t complex_unary_op(mp_unary_op_t op, mp_obj_t o_in) {
Damien George999cedb2015-11-27 17:01:44 +0000118 mp_obj_complex_t *o = MP_OBJ_TO_PTR(o_in);
Damiend99b0522013-12-21 18:17:45 +0000119 switch (op) {
Paul Sokolovsky1b586f32015-10-11 12:09:43 +0300120 case MP_UNARY_OP_BOOL: return mp_obj_new_bool(o->real != 0 || o->imag != 0);
Damien George19f2e472017-04-04 11:57:21 +1000121 case MP_UNARY_OP_HASH: return MP_OBJ_NEW_SMALL_INT(mp_float_hash(o->real) ^ mp_float_hash(o->imag));
Damien Georged17926d2014-03-30 13:35:08 +0100122 case MP_UNARY_OP_POSITIVE: return o_in;
123 case MP_UNARY_OP_NEGATIVE: return mp_obj_new_complex(-o->real, -o->imag);
Damien Georgefdb2aa82017-09-18 14:31:03 +1000124 case MP_UNARY_OP_ABS:
125 return mp_obj_new_float(MICROPY_FLOAT_C_FUN(sqrt)(o->real*o->real + o->imag*o->imag));
Damien George6ac5dce2014-05-21 19:42:43 +0100126 default: return MP_OBJ_NULL; // op not supported
Damiend99b0522013-12-21 18:17:45 +0000127 }
128}
129
Damien George58321dd2017-08-29 13:04:01 +1000130STATIC mp_obj_t complex_binary_op(mp_binary_op_t op, mp_obj_t lhs_in, mp_obj_t rhs_in) {
Damien George999cedb2015-11-27 17:01:44 +0000131 mp_obj_complex_t *lhs = MP_OBJ_TO_PTR(lhs_in);
Damien Georgee2e3d112014-01-06 22:13:00 +0000132 return mp_obj_complex_binary_op(op, lhs->real, lhs->imag, rhs_in);
133}
134
Damien Georgeb1bbe962015-04-01 14:10:50 +0000135STATIC void complex_attr(mp_obj_t self_in, qstr attr, mp_obj_t *dest) {
136 if (dest[0] != MP_OBJ_NULL) {
137 // not load attribute
138 return;
139 }
Damien George999cedb2015-11-27 17:01:44 +0000140 mp_obj_complex_t *self = MP_OBJ_TO_PTR(self_in);
Damien Georgef20375e2014-08-12 19:57:52 +0100141 if (attr == MP_QSTR_real) {
142 dest[0] = mp_obj_new_float(self->real);
143 } else if (attr == MP_QSTR_imag) {
144 dest[0] = mp_obj_new_float(self->imag);
145 }
146}
147
Damien George0c36da02014-03-08 15:24:39 +0000148const mp_obj_type_t mp_type_complex = {
Damien Georgec5966122014-02-15 16:10:44 +0000149 { &mp_type_type },
Damien Georgea71c83a2014-02-15 11:34:50 +0000150 .name = MP_QSTR_complex,
Damien George97209d32014-01-07 15:58:30 +0000151 .print = complex_print,
152 .make_new = complex_make_new,
153 .unary_op = complex_unary_op,
154 .binary_op = complex_binary_op,
Damien Georgeb1bbe962015-04-01 14:10:50 +0000155 .attr = complex_attr,
Damien Georgee2e3d112014-01-06 22:13:00 +0000156};
157
158mp_obj_t mp_obj_new_complex(mp_float_t real, mp_float_t imag) {
159 mp_obj_complex_t *o = m_new_obj(mp_obj_complex_t);
Damien George0c36da02014-03-08 15:24:39 +0000160 o->base.type = &mp_type_complex;
Damien Georgee2e3d112014-01-06 22:13:00 +0000161 o->real = real;
162 o->imag = imag;
Damien George999cedb2015-11-27 17:01:44 +0000163 return MP_OBJ_FROM_PTR(o);
Damien Georgee2e3d112014-01-06 22:13:00 +0000164}
165
166void mp_obj_complex_get(mp_obj_t self_in, mp_float_t *real, mp_float_t *imag) {
Damien Georgeeee1e882019-01-30 18:49:52 +1100167 assert(mp_obj_is_type(self_in, &mp_type_complex));
Damien George999cedb2015-11-27 17:01:44 +0000168 mp_obj_complex_t *self = MP_OBJ_TO_PTR(self_in);
Damien Georgee2e3d112014-01-06 22:13:00 +0000169 *real = self->real;
170 *imag = self->imag;
171}
172
Damien George58321dd2017-08-29 13:04:01 +1000173mp_obj_t mp_obj_complex_binary_op(mp_binary_op_t op, mp_float_t lhs_real, mp_float_t lhs_imag, mp_obj_t rhs_in) {
Damien Georgee2e3d112014-01-06 22:13:00 +0000174 mp_float_t rhs_real, rhs_imag;
175 mp_obj_get_complex(rhs_in, &rhs_real, &rhs_imag); // can be any type, this function will convert to float (if possible)
Damiend99b0522013-12-21 18:17:45 +0000176 switch (op) {
Damien Georged17926d2014-03-30 13:35:08 +0100177 case MP_BINARY_OP_ADD:
178 case MP_BINARY_OP_INPLACE_ADD:
Damiend99b0522013-12-21 18:17:45 +0000179 lhs_real += rhs_real;
180 lhs_imag += rhs_imag;
181 break;
Damien Georged17926d2014-03-30 13:35:08 +0100182 case MP_BINARY_OP_SUBTRACT:
183 case MP_BINARY_OP_INPLACE_SUBTRACT:
Damiend99b0522013-12-21 18:17:45 +0000184 lhs_real -= rhs_real;
185 lhs_imag -= rhs_imag;
186 break;
Damien Georged17926d2014-03-30 13:35:08 +0100187 case MP_BINARY_OP_MULTIPLY:
Damien George5589db82014-04-09 20:21:00 +0100188 case MP_BINARY_OP_INPLACE_MULTIPLY: {
189 mp_float_t real;
190 multiply:
191 real = lhs_real * rhs_real - lhs_imag * rhs_imag;
Damiend99b0522013-12-21 18:17:45 +0000192 lhs_imag = lhs_real * rhs_imag + lhs_imag * rhs_real;
193 lhs_real = real;
194 break;
195 }
Damien Georged17926d2014-03-30 13:35:08 +0100196 case MP_BINARY_OP_FLOOR_DIVIDE:
Damien George5589db82014-04-09 20:21:00 +0100197 case MP_BINARY_OP_INPLACE_FLOOR_DIVIDE:
Damien Georgeb01f66c2018-06-20 21:02:11 +1000198 mp_raise_TypeError("can't truncate-divide a complex number");
Damien George5589db82014-04-09 20:21:00 +0100199
Damien Georged17926d2014-03-30 13:35:08 +0100200 case MP_BINARY_OP_TRUE_DIVIDE:
Damien George5589db82014-04-09 20:21:00 +0100201 case MP_BINARY_OP_INPLACE_TRUE_DIVIDE:
202 if (rhs_imag == 0) {
203 if (rhs_real == 0) {
Damien Georgeb01f66c2018-06-20 21:02:11 +1000204 mp_raise_msg(&mp_type_ZeroDivisionError, "complex divide by zero");
Damien George5589db82014-04-09 20:21:00 +0100205 }
206 lhs_real /= rhs_real;
207 lhs_imag /= rhs_real;
208 } else if (rhs_real == 0) {
209 mp_float_t real = lhs_imag / rhs_imag;
210 lhs_imag = -lhs_real / rhs_imag;
211 lhs_real = real;
212 } else {
213 mp_float_t rhs_len_sq = rhs_real*rhs_real + rhs_imag*rhs_imag;
214 rhs_real /= rhs_len_sq;
215 rhs_imag /= -rhs_len_sq;
216 goto multiply;
217 }
218 break;
219
Damien Georgeae491052014-04-10 20:08:11 +0100220 case MP_BINARY_OP_POWER:
221 case MP_BINARY_OP_INPLACE_POWER: {
222 // z1**z2 = exp(z2*ln(z1))
223 // = exp(z2*(ln(|z1|)+i*arg(z1)))
224 // = exp( (x2*ln1 - y2*arg1) + i*(y2*ln1 + x2*arg1) )
225 // = exp(x3 + i*y3)
226 // = exp(x3)*(cos(y3) + i*sin(y3))
227 mp_float_t abs1 = MICROPY_FLOAT_C_FUN(sqrt)(lhs_real*lhs_real + lhs_imag*lhs_imag);
228 if (abs1 == 0) {
Damien George3ed0e5e2017-02-03 00:01:45 +1100229 if (rhs_imag == 0 && rhs_real >= 0) {
Damien Georgebd04ed32017-02-04 00:23:56 +1100230 lhs_real = (rhs_real == 0);
Damien Georgeae491052014-04-10 20:08:11 +0100231 } else {
Damien George7d0d7212016-10-17 12:17:37 +1100232 mp_raise_msg(&mp_type_ZeroDivisionError, "0.0 to a complex power");
Damien Georgeae491052014-04-10 20:08:11 +0100233 }
234 } else {
235 mp_float_t ln1 = MICROPY_FLOAT_C_FUN(log)(abs1);
236 mp_float_t arg1 = MICROPY_FLOAT_C_FUN(atan2)(lhs_imag, lhs_real);
237 mp_float_t x3 = rhs_real * ln1 - rhs_imag * arg1;
238 mp_float_t y3 = rhs_imag * ln1 + rhs_real * arg1;
239 mp_float_t exp_x3 = MICROPY_FLOAT_C_FUN(exp)(x3);
240 lhs_real = exp_x3 * MICROPY_FLOAT_C_FUN(cos)(y3);
241 lhs_imag = exp_x3 * MICROPY_FLOAT_C_FUN(sin)(y3);
242 }
243 break;
244 }
245
Paul Sokolovsky1b586f32015-10-11 12:09:43 +0300246 case MP_BINARY_OP_EQUAL: return mp_obj_new_bool(lhs_real == rhs_real && lhs_imag == rhs_imag);
Damien Georgeb8a053a2014-04-11 10:10:37 +0100247
Damien George5589db82014-04-09 20:21:00 +0100248 default:
Damien George6ac5dce2014-05-21 19:42:43 +0100249 return MP_OBJ_NULL; // op not supported
Damiend99b0522013-12-21 18:17:45 +0000250 }
251 return mp_obj_new_complex(lhs_real, lhs_imag);
252}
253
Damiend99b0522013-12-21 18:17:45 +0000254#endif