| 33 | } |
| 34 | |
| 35 | array mandelbrot(const array &in, int iter, float maxval) { |
| 36 | array C = in; |
| 37 | array Z = C; |
| 38 | array mag = constant(0, C.dims()); |
| 39 | |
| 40 | for (int ii = 1; ii < iter; ii++) { |
| 41 | // Do the calculation |
| 42 | Z = Z * Z + C; |
| 43 | |
| 44 | // Get indices where abs(Z) crosses maxval |
| 45 | array cond = (abs(Z) > maxval).as(f32); |
| 46 | mag = af::max(mag, cond * ii); |
| 47 | |
| 48 | // If abs(Z) cross maxval, turn off those locations |
| 49 | C = C * (1 - cond); |
| 50 | Z = Z * (1 - cond); |
| 51 | |
| 52 | // Ensuring the JIT does not become too large |
| 53 | af::eval(C, Z); |
| 54 | mag.eval(); |
| 55 | } |
| 56 | |
| 57 | // Normalize |
| 58 | return mag / maxval; |
| 59 | } |
| 60 | |
| 61 | array normalize(array a) { |
| 62 | float mx = af::max<float>(a); |