arb complex plot

gistfile1.c

#include "acb_modular.h"
#include "profiler.h"

/* HLS algorithm from python's colorsys module */
static double vv(double m1, double m2, double hue)
{
    hue = hue - floor(hue);

    if (hue < 1/6.)
        return m1 + (m2-m1)*hue*6.0;
    if (hue < 0.5)
        return m2;
    if (hue < 2/3.)
        return m1 + (m2-m1)*(2/3.-hue)*6.0;
    return m1;
}

static void hls_to_rgb(double *R, double *G, double *B,
                double h, double l, double s)
{
    double m1, m2;

    if (s == 0.0)
    {
        *R = *G = *B = l;
        return;
    }

    if (l <= 0.5)
        m2 = l * (1.0+s);
    else
        m2 = l+s-(l*s);

    m1 = 2.0*l - m2;

    *R = vv(m1, m2, h + 1/3.);
    *G = vv(m1, m2, h);
    *B = vv(m1, m2, h - 1/3.);
}


#define PI 3.1415926535898

/* note: should make z exact first */
void
color_function(double * R, double * G, double * B, const acb_t z)
{
    if (!acb_is_finite(z))
    {
        *R = *G = *B = 0.5;
    }
    else
    {
        double H, L, S;
        long prec;
        arb_t t, u;

        arb_init(t);
        arb_init(u);

        prec = 32;

        arf_set_round(arb_midref(t), arb_midref(acb_realref(z)), prec, ARF_RND_DOWN);
        arf_set_round(arb_midref(u), arb_midref(acb_imagref(z)), prec, ARF_RND_DOWN);

        arb_atan2(t, u, t, prec);

        H = arf_get_d(arb_midref(t), ARF_RND_DOWN);

        H = (H + PI) / (2 * PI) + 0.5;
        H = H - floor(H);

        acb_abs(t, z, prec);

        if (arf_cmpabs_2exp_si(arb_midref(t), 200) > 0)
        {
            L = 1.0;
        }
        else if (arf_cmpabs_2exp_si(arb_midref(t), -200) < 0)
        {
            L = 0.0;
        }
        else
        {
            L = arf_get_d(arb_midref(t), ARF_RND_DOWN);
            L = 1.0 - 1.0/(1.0 + pow(L, 0.3));
        }

        S = 0.8;

        hls_to_rgb(R, G, B, H, L, S);

        arb_clear(t);
        arb_clear(u);
    }
}

int main()
{
    arf_t xa, xb, ya, yb;
    long x, y, xnum, ynum, prec;
    double R, G, B;
    FILE * fp;

    acb_t tau, z;

    acb_init(tau);
    acb_init(z);

    ynum = 128;
    xnum = 128 * 4;
    prec = 128;

    arf_init(xa);
    arf_init(xb);
    arf_init(ya);
    arf_init(yb);

    arf_set_d(xa, -2.0);
    arf_set_d(xb, 2.0);
    arf_set_d(ya, 0.0);
    arf_set_d(yb, 1.0);

    acb_zero(tau);

    fp = fopen("arbplot.ppm", "w");

    fprintf(fp, "P6\n%ld %ld 255\n", xnum, ynum);

    TIMEIT_ONCE_START

    for (y = ynum - 1; y >= 0; y--)
    {
        printf("row %ld\n", y);

        for (x = 0; x < xnum; x++)
        {
            arf_sub(arb_midref(acb_imagref(tau)), yb, ya, prec, ARF_RND_DOWN);
            arf_mul_ui(arb_midref(acb_imagref(tau)), arb_midref(acb_imagref(tau)), y, prec, ARF_RND_DOWN);
            arf_div_ui(arb_midref(acb_imagref(tau)), arb_midref(acb_imagref(tau)), ynum - 1, prec, ARF_RND_DOWN);
            arf_add(arb_midref(acb_imagref(tau)), arb_midref(acb_imagref(tau)), ya, prec, ARF_RND_DOWN);

            arf_sub(arb_midref(acb_realref(tau)), xb, xa, prec, ARF_RND_DOWN);
            arf_mul_ui(arb_midref(acb_realref(tau)), arb_midref(acb_realref(tau)), x, prec, ARF_RND_DOWN);
            arf_div_ui(arb_midref(acb_realref(tau)), arb_midref(acb_realref(tau)), xnum - 1, prec, ARF_RND_DOWN);
            arf_add(arb_midref(acb_realref(tau)), arb_midref(acb_realref(tau)), xa, prec, ARF_RND_DOWN);

            acb_modular_j(z, tau, prec);
            acb_div_ui(z, z, 1728, prec);

            if (acb_is_finite(z))
            {
                mag_zero(arb_radref(acb_realref(z)));
                mag_zero(arb_radref(acb_imagref(z)));
            }

            //acb_log(z, z, prec);
            //acb_printd(z, 10); printf("\n");

            color_function(&R, &G, &B, z);

            fputc(FLINT_MIN(255, floor(R * 255)), fp);
            fputc(FLINT_MIN(255, floor(G * 255)), fp);
            fputc(FLINT_MIN(255, floor(B * 255)), fp);
        }
    }

    TIMEIT_ONCE_STOP

    fclose(fp);

    arf_clear(xa);
    arf_clear(xb);
    arf_clear(ya);
    arf_clear(yb);

    acb_clear(tau);
    acb_clear(z);

    flint_cleanup();
}