pearcey.c

#include <string.h>
#include "acb_calc.h"
#include "flint/profiler.h"

/* the integrand */
int
f_pearcey(acb_ptr res, const acb_t z, void * param, slong order, slong prec)
{
    acb_t t, u;

    if (order > 1)
        flint_abort();  /* Would be needed for Taylor method. */

    acb_init(t);
    acb_init(u);

    acb_mul(t, z, z, prec);
    acb_mul(u, t, t, prec);
    acb_neg(u, u);
    acb_addmul(u, t, ((acb_ptr) param) + 0, prec);
    acb_exp(u, u, prec);

    acb_mul(t, ((acb_ptr) param) + 1, z, prec);
    acb_cosh(t, t, prec);
    acb_mul(res, t, u, prec);

    acb_clear(t);
    acb_clear(u);

    return 0;
}

void
acb_pearcey(acb_t res, const acb_t x, const acb_t y, slong prec)
{
    acb_t a, b;
    arb_t A, B, C, t, err;
    slong N;

    acb_init(a);
    acb_init(b);
    arb_init(A);
    arb_init(B);
    arb_init(C);
    arb_init(t);
    arb_init(err);

    /* a = exp(pi i (3/4)) */
    arb_sqrt_ui(acb_imagref(a), 2, prec);
    arb_neg(acb_realref(a), acb_imagref(a));
    acb_mul_2exp_si(a, a, -1);

    /* b = exp(pi i (5/4)) */
    acb_sqrt(b, a, prec);
    arb_neg(acb_realref(b), acb_realref(b));

    acb_mul(a, a, x, prec);
    acb_mul(b, b, y, prec);

    arb_nonnegative_part(A, acb_realref(a));
    arb_abs(B, acb_realref(b));

    arb_pos_inf(err);

    for (N = 2; N < prec; N++)
    {
        arb_one(C);
        arb_div_ui(t, A, N * N, prec);
        arb_sub(C, C, t, prec);
        arb_div_ui(t, B, N * N * N, prec);
        arb_sub(C, C, t, prec);

        if (arb_is_positive(C))
        {
            arb_mul_ui(err, C, N * N, prec);
            arb_mul_ui(err, err, N * N, prec);
            arb_neg(err, err);
            arb_exp(err, err, prec);
            arb_div(err, err, C, prec);

            arb_one(t);
            arb_mul_2exp_si(t, t, -prec);
            if (arb_lt(err, t))
                break;
        }
        else
        {
            arb_pos_inf(err);
        }
    }

    if (arb_is_finite(err))
    {
        acb_struct param[2];
        acb_t xa, xb;
        mag_t tol;

        acb_init(xa);
        acb_init(xb);
        mag_init(tol);

        param[0] = *a;
        param[1] = *b;

        acb_zero(xa);
        acb_set_si(xb, N);

        mag_set_ui_2exp_si(tol, 1, -prec);

        acb_calc_integrate(res, f_pearcey, param, xa, xb, prec, tol, NULL, prec);

        arb_add_error(acb_realref(res), err);
        arb_add_error(acb_imagref(res), err);

        arb_sqrt_ui(acb_imagref(a), 2, prec);
        arb_set(acb_realref(a), acb_imagref(a));
        acb_mul_2exp_si(a, a, -1);
        acb_sqrt(a, a, prec);

        acb_mul(res, res, a, prec);
        acb_mul_2exp_si(res, res, 1);

        acb_clear(xa);
        acb_clear(xb);
        mag_clear(tol);
    }
    else
    {
        acb_indeterminate(res);
    }

    acb_clear(a);
    acb_clear(b);
    arb_clear(A);
    arb_clear(B);
    arb_clear(C);
    arb_clear(t);
    arb_clear(err);
}


/* ---- modified contents of complex_plot.c below ---- */


/* some useful color operations */
#define CLAMP(y) FLINT_MAX(0.0, FLINT_MIN((y), 1.0))
#define BLEND(x,y) (0.5*(x) + 0.5*(y))
#define DODGE(a,b) ((a) / ((1.0 - (b)) + 1/256.0))

/* 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.);
}

static void
rgb_to_hls(double *H, double *L, double *S, double R, double G, double B)
{
    double h, l, s, hi, lo, d;

    hi = FLINT_MAX(FLINT_MAX(R, G), B);
    lo = FLINT_MIN(FLINT_MIN(R, G), B);

    l = 0.5 * (lo + hi);
    d = hi - lo;

    if (hi == lo)
    {
        s = 0.0;
        h = 0.0;
    }
    else
    {
        if (l <= 0.5)
            s = d / (hi + lo);
        else
            s = d / (2.0 - hi - lo);

        if (d == 0.0)
            d = 1.0;

        if (R == hi)
            h = (G - B) / d;
        else if (G == hi)
            h = (B - R) / d + 2.0;
        else
            h = (R - G) / d + 4.0;

        h = h / 6.0;
        if (h < 0.0)
            h += 1.0;
    }

    *H = h;
    *L = l;
    *S = s;
}

/* color balance algorithm from gimp */
static double balance_channel(double value, double l,
    double shadows, double midtones, double highlights)
{
    double a = 0.25, b = 0.333, scale = 0.7;

    shadows    *= CLAMP((l - b) / (-a) + 0.5) * scale;
    midtones   *= CLAMP((l - b) / ( a) + 0.5) *
                  CLAMP((l + b - 1.0) / (-a) + 0.5) * scale;
    highlights *= CLAMP((l + b - 1.0) / ( a) + 0.5) * scale;

    value += shadows;
    value += midtones;
    value += highlights;
    return CLAMP(value);
}

static void balance(double * R, double * G, double * B,
    double ra, double rb, double rc,  /* red shadows, midtones, highlights */
    double ga, double gb, double gc,  /* green */
    double ba, double bb, double bc)  /* blue */
{
    double h, l, s;
    double h2, l2, s2;

    rgb_to_hls(&h, &l, &s, *R, *G, *B);

    *R = balance_channel(*R, *R, ra, rb, rc);
    *G = balance_channel(*G, *G, ga, gb, gc);
    *B = balance_channel(*B, *B, ba, bb, bc);

    /* preserve lightness */
    rgb_to_hls(&h2, &l2, &s2, *R, *G, *B);
    hls_to_rgb(R, G, B, h2, l, s2);
}

#define PI 3.1415926535898

const double blue_orange_colors[][4] = {
  {-1.0,  0.0, 0.0, 0.0},
  {-0.95, 0.1, 0.2, 0.5},
  {-0.5,  0.0, 0.5, 1.0},
  {-0.05, 0.4, 0.8, 0.8},
  { 0.0,  1.0, 1.0, 1.0},
  { 0.05, 1.0, 0.9, 0.3},
  { 0.5,  0.9, 0.5, 0.0},
  { 0.95, 0.7, 0.1, 0.0},
  { 1.0,  0.0, 0.0, 0.0},
  { 2.0,  0.0, 0.0, 0.0},
};

void
color_function(double * R, double * G, double * B, const acb_t z, int mode)
{
    double H, L, S;
    slong prec, i;
    arb_t t, u;

    if (!acb_is_finite(z) || acb_rel_accuracy_bits(z) < 4)
    {
        *R = *G = *B = 0.5;
        return;
    }

    if (mode >= 2)
    {
        double R1, G1, B1;
        double R2, G2, B2;

        /* combine both color functions */
        color_function(&R1, &G1, &B1, z, 0);
        color_function(&R2, &G2, &B2, z, 1);

        *R = BLEND(R1, CLAMP(DODGE(R1, R2)));
        *G = BLEND(G1, CLAMP(DODGE(G1, G2)));
        *B = BLEND(B1, CLAMP(DODGE(B1, B2)));

        /* then play with the levels */
        if (mode == 3)
            balance(R, G, B, 0.0, -0.5, 0.2, 0.0, 0.0, -0.1, 0.0, -1.0, -0.2);
        else if (mode == 4)
            balance(R, G, B, 0.0, -0.5, 0.2, 0.0, 0.5, -0.1, 0.0, -0.3, -1.0);
        else if (mode == 5)
            balance(R, G, B, 0.0, -0.5, -1.0, 0.0, -0.1, -0.67, 0.0, -0.55, -0.12);
        else if (mode == 6)
            balance(R, G, B, 0.86, 0.0, 0.13, 0.57, 0.19, -0.52, 0.31, -0.30, -0.94);

        return;
    }

    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);

    if (mode == 0)
    {
        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.2));
        }

        S = 0.8;

        hls_to_rgb(R, G, B, H, L, S);
    }
    else
    {
        H = H / PI;
        H = FLINT_MAX(FLINT_MIN(H, 1.0), -1.0);

        for (i = 1; ; i++)
        {
            if (blue_orange_colors[i][0] > H)
            {
                double a, ra, ga, ba, b, rb, gb, bb, s;

                a  = blue_orange_colors[i-1][0];
                ra = blue_orange_colors[i-1][1];
                ga = blue_orange_colors[i-1][2];
                ba = blue_orange_colors[i-1][3];
                b  = blue_orange_colors[i][0];
                rb = blue_orange_colors[i][1];
                gb = blue_orange_colors[i][2];
                bb = blue_orange_colors[i][3];

                s = (H - a) / (b - a);
                *R = ra + (rb - ra) * s;
                *G = ga + (gb - ga) * s;
                *B = ba + (bb - ba) * s;
                break;
            }
        }
    }

    arb_clear(t);
    arb_clear(u);
}

typedef void (*func_ptr)(acb_t, const acb_t, slong);

void
pearcey(acb_t res, const acb_t z, slong prec)
{
    acb_t x, y;
    acb_init(x);
    acb_init(y);
    acb_set_arb(x, acb_imagref(z));
    acb_set_arb(y, acb_realref(z));
    acb_pearcey(res, x, y, prec);
    acb_clear(x);
    acb_clear(y);
}

int main(int argc, char *argv[])
{
    slong x, y, xnum, ynum, prec, i;
    double R, G, B, dxa, dxb, dya, dyb;
    FILE * fp;
    arf_t xa, xb, ya, yb;
    acb_t z, w;
    func_ptr func;
    int color_mode;

    if (argc < 2)
    {
        printf("complex_plot [-range xa xb ya yb] [-size xn yn] [-color n] <func>\n\n");

        printf("Plots one of the predefined functions on [xa,xb] + [ya,yb]i\n");
        printf("using domain coloring, at a resolution of xn by yn pixels.\n\n");

        printf("Defaults parameters are [-10,10] + [-10,10]i and xn = yn = 512.\n\n");

        printf("A color function can be selected with -color. The choices are:\n");
        printf("0   phase=hue, magnitude=brightness\n");
        printf("1   phase only, white-gold-black-blue-white counterclockwise\n");
        printf("2   0+1 (dodge filter)\n");
        printf("3   0+1, shiny\n");
        printf("4   0+1, warm\n");
        printf("5   0+1, cold\n");
        printf("6   0+1, tomato\n\n");

        printf("The output is written to arbplot.ppm. If you have ImageMagick,\n");
        printf("run [convert arbplot.ppm arbplot.png] to get a PNG.\n\n");

        return 1;
    }

    xnum = 512;
    ynum = 512;
    dxa = dya = -10;
    dxb = dyb = 10;
    func = pearcey;
    color_mode = 0;

    for (i = 1; i < argc; i++)
    {
        if (!strcmp(argv[i], "-size"))
        {
            xnum = atol(argv[i+1]);
            ynum = atol(argv[i+2]);
            i += 2;
        }
        else if (!strcmp(argv[i], "-range"))
        {
            dxa = atof(argv[i+1]);
            dxb = atof(argv[i+2]);
            dya = atof(argv[i+3]);
            dyb = atof(argv[i+4]);
            i += 4;
        }
        else if (!strcmp(argv[i], "-color"))
        {
            color_mode = atoi(argv[i+1]);
            i++;
        }
        else
        {
            printf("unknown option: %s\n", argv[i]);
            return 1;
        }
    }

    acb_init(z);
    acb_init(w);

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

    arf_set_d(xa, dxa);
    arf_set_d(xb, dxb);
    arf_set_d(ya, dya);
    arf_set_d(yb, dyb);

    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--)
    {
        if (y % (ynum / 16) == 0 || 1)
            printf("row %ld\n", y);

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

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

                func(w, z, prec);

                if (acb_rel_accuracy_bits(w) > 4)
                    break;
            }

            color_function(&R, &G, &B, w, color_mode);

            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(z);
    acb_clear(w);

    flint_cleanup();
    return 0;
}