Skip to content

Optimize computation algorithm #4

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all 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
235 changes: 151 additions & 84 deletions scripts/math/fractal_calculation.ts
Original file line number Diff line number Diff line change
Expand Up @@ -7,63 +7,162 @@ type PlotScale = {
y_display_range: number
}

class Complex32 {
re: number;
im: number;
constructor(re: number, im: number) {
this.re = re;
this.im = im;
}

clone() {
return new Complex32(this.re, this.im);
}
function complexAbs(real: number, imag: number): number {
// This is more numerically stable than `sqrt(x^2 + y^2)`.
let ratio = real / imag;
return real * Math.sqrt(1 + ratio * ratio);
}

add(other: Complex32) {
this.re += other.re;
this.im += other.im;
function getRootId(roots: Float32Array, iterationsCount: number, xp: number, yp: number): number {
/*
This is the following pseudocode loop where `Im` is the imaginary unit and `real(x)` and
`imag(x)` get the real and imaginary parts of `x` respectively:

```js
let z = xp + yp*Im;

// Perform a Newton's method approximation using `iterationsCount` rounds to find the
// nearest root. If it gets close enough to a root, return it.
for (let i = 0; i < iterationsCount; i++) {
let sum = 0 + 0*Im;
for (let i = 0; i < roots.length; i += 2) {
let rootReal = roots[i + 0];
let rootImag = roots[i + 1];
let root = rootReal + rootImag*Im;
let diff = z - root;
if (real(diff) < 0.001 && imag(diff) < 0.001) {
return i;
}
sum += 1 / diff;
}
z -= 1 / sum;
}

subtract(other: Complex32) {
this.re -= other.re;
this.im -= other.im;
// Failing a close-enough match, find and return the root with the nearest distance.
let minDistance = Infinity;
let closestRootId = 0;

for (let i = 0; i < roots.length; i += 2) {
let rootReal = roots[i + 0];
let rootImag = roots[i + 1];
let dst = abs(z - root);
if (dst < minDistance) {
minDistance = dst;
closestRootId = i;
}
}

invert() {
const square_sum = this.normSqr();
this.re /= square_sum;
this.im /= -square_sum;
return closestRootId;
```

First, I split it apart into using real and imaginary scalars directly:

```js
let realZ = xp;
let imagZ = yp;

// Perform a Newton's method approximation using `iterationsCount` rounds to find the
// nearest root. If it gets close enough to a root, return it.
for (let i = 0; i < iterationsCount; i++) {
let sumReal = 0;
let sumImag = 0;
for (let i = 0; i < roots.length; i += 2) {
let rootReal = roots[i + 0];
let rootImag = roots[i + 1];
let diffReal = realZ - rootReal;
let diffImag = imagZ - rootImag;
if (diffReal < 0.001 && diffImag < 0.001) {
return i;
}
// 1/(x + yi) = x/(x^2 + y^2) - (y/(x^2 + y^2))i
// (a + bi) + 1/(x + yi) = (a + x/(x^2 + y^2)) + (b - y/(x^2 + y^2))i
let squareNorm = diffReal * diffReal + diffImag * diffImag;
sumReal += diffReal / squareNorm;
sumImag += diffImag / -squareNorm;
}
// 1/(x + yi) = x/(x^2 + y^2) - (y/(x^2 + y^2))i
// (a + bi) - 1/(x + yi) = (a - x/(x^2 + y^2)) + (b + y/(x^2 + y^2))i
let squareNorm = sumReal * sumReal + sumImag * sumImag;
realZ += sumReal / -squareNorm;
imagZ += sumImag / squareNorm;
}

normSqr(): number {
return this.re * this.re + this.im * this.im;
// Failing a close-enough match, find and return the root with the nearest distance.
let minDistance = Infinity;
let closestRootId = 0;

for (let i = 0; i < roots.length; i += 2) {
let rootReal = roots[i + 0];
let rootImag = roots[i + 1];
let dst = complexAbs(realZ - rootReal, imagZ - rootImag);
if (dst < minDistance) {
minDistance = dst;
closestRootId = i;
}
}

distance(): number {
return Math.sqrt(this.re * this.re + this.im * this.im);
return closestRootId;
```

Then, in the inner loop, I take advantage of a useful mathematical identity to avoid some
calls to `Math.abs` (cheap) as well as an extra floating point comparison (expensive).
`a < b and c < d` and `a + c < b + d` are equivalent provided all four variables are
non-negative. If I square both sides of each comparison, `a < b` and `c < d`, that condition
would be guaranteed, and that leaves the following expression: `a^2 + c^2 < b^2 * d^2`. When
I substitute `diffReal` for `a`, `diffImag` for `c`, and `0.001` for `b` and `d`, I get the
following expression:

```js
diffReal * diffReal + diffImag * diffImag < 0.001 * 0.001 + 0.001 * 0.001
```

The left side just so happens to be the `squareNorm` computed from earlier, and the right
side happens to be all constant (it evaluates to `0.000002`). Leveraging this, I can then
revise that code to what's now below.
*/

let realZ = xp;
let imagZ = yp;

// Perform a Newton's method approximation using `iterationsCount` rounds to find the
// nearest root. If it gets close enough to a root, return it.
for (let i = 0; i < iterationsCount; i++) {
let sumReal = 0;
let sumImag = 0;
for (let i = 0; i < roots.length; i += 2) {
let rootReal = roots[i + 0];
let rootImag = roots[i + 1];
let diffReal = realZ - rootReal;
let diffImag = imagZ - rootImag;
// 1/(x + yi) = x/(x^2 + y^2) - (y/(x^2 + y^2))i
// (a + bi) + 1/(x + yi) = (a + x/(x^2 + y^2)) + (b - y/(x^2 + y^2))i
let squareNorm = diffReal * diffReal + diffImag * diffImag;
if (squareNorm < 0.000002) {
return i;
}
sumReal += diffReal / squareNorm;
sumImag += diffImag / -squareNorm;
}
// 1/(x + yi) = x/(x^2 + y^2) - (y/(x^2 + y^2))i
// (a + bi) - 1/(x + yi) = (a - x/(x^2 + y^2)) + (b + y/(x^2 + y^2))i
let squareNorm = sumReal * sumReal + sumImag * sumImag;
realZ += sumReal / -squareNorm;
imagZ += sumImag / squareNorm;
}
}

function newtonMethodApprox(roots: Complex32[], z: Complex32): {
idx: number;
z: Complex32;
} {
let sum = new Complex32(0, 0);
let i = 0;
for (const root of roots) {
let diff = z.clone();
diff.subtract(root);
if (Math.abs(diff.re) < 0.001 && Math.abs(diff.im) < 0.001) {
return { idx: i, z };
// Failing a close-enough match, find and return the root with the nearest distance.
let minDistance = Infinity;
let closestRootId = 0;

for (let i = 0; i < roots.length; i += 2) {
let rootReal = roots[i + 0];
let rootImag = roots[i + 1];
let dst = complexAbs(realZ - rootReal, imagZ - rootImag);
if (dst < minDistance) {
minDistance = dst;
closestRootId = i;
}
diff.invert();
sum.add(diff);
i++;
}
sum.invert();
let result = z.clone();
result.subtract(sum);
return { idx: -1, z: result };

return closestRootId;
}

function transformPointToPlotScale(x: number, y: number, plotScale: PlotScale): number[] {
Expand All @@ -90,50 +189,18 @@ function fillPixelsJavascript(buffer: SharedArrayBuffer, plotScale: PlotScale, r
y_display_range: height
} = plotScale;
let [w_int, h_int] = [Math.round(width), Math.round(height)];

let flatColors = new Uint8Array(colors.flat());
let colorPacks = new Uint32Array(flatColors.buffer);

let complexRoots: Complex32[] = roots.map((pair) => new Complex32(pair[0], pair[1]));

const filler = (x: number, y: number) => {
let minDistance = Number.MAX_SAFE_INTEGER;
let closestRootId = 0;
let [xp, yp] = transformPointToPlotScale(x, y, plotScale);
let z = new Complex32(xp, yp);
let colorId = -1;
for (let i = 0; i < iterationsCount; i++) {
let { idx, z: zNew } = newtonMethodApprox(complexRoots, z);
if (idx != -1) {
colorId = idx;
break;
}
z = zNew;
}
if (colorId != -1) {
return colorPacks[colorId];
}
let i = 0;
for(const root of complexRoots) {
let d = z.clone();
d.subtract(root);
let dst = d.distance();
if (dst < minDistance) {
minDistance = dst;
closestRootId = i;
}
i++;
}
return colorPacks[closestRootId];
}
let flattenedRoots = Float32Array.from(roots.flat());
let u32BufferView = new Uint32Array(buffer, bufferPtr);

let totalSize = w_int * h_int;
let thisBorder = calculatePartSize(totalSize, partsCount, partOffset, 1);
let nextBorder = calculatePartSize(totalSize, partsCount, partOffset + 1, 1);

let u32BufferView = new Uint32Array(buffer, bufferPtr);

for (let i = thisBorder; i < nextBorder; i++) {
u32BufferView[i] = filler(i % w_int, i / w_int);
let [xp, yp] = transformPointToPlotScale(i % w_int, i / w_int, plotScale);
u32BufferView[i] = colorPacks[getRootId(flattenedRoots, iterationsCount, xp, yp)];
}
}
}