File: color_maps_example.cpp

package info (click to toggle)
boost1.83 1.83.0-5
  • links: PTS, VCS
  • area: main
  • in suites: forky, sid
  • size: 545,632 kB
  • sloc: cpp: 3,857,086; xml: 125,552; ansic: 34,414; python: 25,887; asm: 5,276; sh: 4,799; ada: 1,681; makefile: 1,629; perl: 1,212; pascal: 1,139; sql: 810; yacc: 478; ruby: 102; lisp: 24; csh: 6
file content (221 lines) | stat: -rw-r--r-- 7,068 bytes parent folder | download | duplicates (9)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
//  (C) Copyright Nick Thompson 2021.
//  (C) Copyright Matt Borland 2022.
//  Use, modification and distribution are subject to the
//  Boost Software License, Version 1.0. (See accompanying file
//  LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

#include <cmath>
#include <cstdint>
#include <array>
#include <complex>
#include <tuple>
#include <iostream>
#include <vector>
#include <limits>
#include <boost/math/tools/color_maps.hpp>

#if !__has_include("lodepng.h")
 #error "lodepng.h is required to run this example."
#endif
#include "lodepng.h"
#include <iostream>
#include <string>
#include <vector>


// In lodepng, the vector is expected to be row major, with the top row
// specified first. Note that this is a bit confusing sometimes as it's more
// natural to let y increase moving *up*.
unsigned write_png(const std::string &filename,
                   const std::vector<std::uint8_t> &img, std::size_t width,
                   std::size_t height) {
  unsigned error = lodepng::encode(filename, img, width, height,
                                   LodePNGColorType::LCT_RGBA, 8);
  if (error) {
    std::cerr << "Error encoding png: " << lodepng_error_text(error) << "\n";
  }
  return error;
}


// Computes ab - cd.
// See: https://pharr.org/matt/blog/2019/11/03/difference-of-floats.html
template <typename Real>
inline Real difference_of_products(Real a, Real b, Real c, Real d)
{
    Real cd = c * d;
    Real err = std::fma(-c, d, cd);
    Real dop = std::fma(a, b, -cd);
    return dop + err;
}

template<typename Real>
auto fifth_roots(std::complex<Real> z)
{
    std::complex<Real> v = std::pow(z,4);
    std::complex<Real> dw = Real(5)*v;
    std::complex<Real> w = v*z - Real(1);
    return std::make_pair(w, dw);
}

template<typename Real>
auto g(std::complex<Real> z)
{
    std::complex<Real> z2 = z*z;
    std::complex<Real> z3 = z*z2;
    std::complex<Real> z4 = z2*z2;
    std::complex<Real> w = z4*(z4 + Real(15)) - Real(16);
    std::complex<Real> dw = Real(4)*z3*(Real(2)*z4 + Real(15));
    return std::make_pair(w, dw);
}

template<typename Real>
std::complex<Real> complex_newton(std::function<std::pair<std::complex<Real>,std::complex<Real>>(std::complex<Real>)> f, std::complex<Real> z)
{
    // f(x(1+e)) = f(x) + exf'(x)
    bool close = false;
    do
    {
        auto [y, dy] = f(z);
        z -= y/dy;
        close = (abs(y) <= 1.4*std::numeric_limits<Real>::epsilon()*abs(z*dy));
    } while(!close);
    return z;
}

template<typename Real>
class plane_pixel_map
{
public:
    plane_pixel_map(int64_t image_width, int64_t image_height, Real xmin, Real ymin)
    {
        image_width_ = image_width;
        image_height_ = image_height;
        xmin_ = xmin;
        ymin_ = ymin;
    }

    std::complex<Real> to_complex(int64_t i, int64_t j) const {
        Real x = xmin_ + 2*abs(xmin_)*Real(i)/Real(image_width_ - 1);
        Real y = ymin_ + 2*abs(ymin_)*Real(j)/Real(image_height_ - 1);
        return std::complex<Real>(x,y);
    }

    std::pair<int64_t, int64_t> to_pixel(std::complex<Real> z) const {
        Real x = z.real();
        Real y = z.imag();
        Real ii = (image_width_ - 1)*(x - xmin_)/(2*abs(xmin_));
        Real jj = (image_height_ - 1)*(y - ymin_)/(2*abs(ymin_));

        return std::make_pair(std::round(ii), std::round(jj));
    }

private:
    int64_t image_width_;
    int64_t image_height_;
    Real xmin_;
    Real ymin_;
};

int main(int argc, char** argv)
{
    using Real = double;
    using boost::math::tools::viridis;
    using std::sqrt;

    std::function<std::array<Real, 3>(Real)> color_map = viridis<Real>;
    std::string requested_color_map = "viridis";
    if (argc == 2) {
       requested_color_map = std::string(argv[1]);
       if (requested_color_map == "smooth_cool_warm") {
          color_map = boost::math::tools::smooth_cool_warm<Real>;
       }
       else if (requested_color_map == "plasma") {
          color_map = boost::math::tools::plasma<Real>;
       }
       else if (requested_color_map == "black_body") {
          color_map = boost::math::tools::black_body<Real>;
       }
       else if (requested_color_map == "inferno") {
          color_map = boost::math::tools::inferno<Real>;
       }
       else if (requested_color_map == "kindlmann") {
          color_map = boost::math::tools::kindlmann<Real>;
       }
       else if (requested_color_map == "extended_kindlmann") {
          color_map = boost::math::tools::extended_kindlmann<Real>;
       }
       else {
          std::cerr << "Could not recognize color map " << argv[1] << ".";
          return 1;
       }
    }
    constexpr int64_t image_width = 1024;
    constexpr int64_t image_height = 1024;
    constexpr const Real two_pi = 6.28318530718;

    std::vector<std::uint8_t> img(4*image_width*image_height, 0);
    plane_pixel_map<Real> map(image_width, image_height, Real(-2), Real(-2));

    for (int64_t j = 0; j < image_height; ++j)
    {
        std::cout << "j = " << j << "\n";
        for (int64_t i = 0; i < image_width; ++i)
        {
            std::complex<Real> z0 = map.to_complex(i,j);
            auto rt = complex_newton<Real>(g<Real>, z0);
            // The root is one of exp(2*pi*ij/5). Therefore, it can be classified by angle.
            Real theta = std::atan2(rt.imag(), rt.real());
            // Now theta in [-pi,pi]. Get it into [0,2pi]:
            if (theta < 0) {
                theta += two_pi;
            }
            theta /= two_pi;
            if (std::isnan(theta)) {
                std::cerr << "Theta is a nan!\n";
            }
            auto c = boost::math::tools::to_8bit_rgba(color_map(theta));
            int64_t idx = 4 * image_width * (image_height - 1 - j) + 4 * i;
            img[idx + 0] = c[0];
            img[idx + 1] = c[1];
            img[idx + 2] = c[2];
            img[idx + 3] = c[3];
        }
    }

    std::array<std::complex<Real>, 8> roots;
    roots[0] = -Real(1);
    roots[1] = Real(1);
    roots[2] = {Real(0), Real(1)};
    roots[3] = {Real(0), -Real(1)};
    roots[4] = {sqrt(Real(2)), sqrt(Real(2))};
    roots[5] = {sqrt(Real(2)), -sqrt(Real(2))};
    roots[6] = {-sqrt(Real(2)), -sqrt(Real(2))};
    roots[7] = {-sqrt(Real(2)), sqrt(Real(2))};

    for (int64_t k = 0; k < 8; ++k)
    {
        auto [ic, jc] = map.to_pixel(roots[k]);

        int64_t r = 7;
        for (int64_t i = ic - r; i < ic + r; ++i)
        {
            for (int64_t j = jc - r; j < jc + r; ++j)
            {
                if ((i-ic)*(i-ic) + (j-jc)*(j-jc) > r*r)
                {
                    continue;
                }
                int64_t idx = 4 * image_width * (image_height - 1 - j) + 4 * i;
                img[idx + 0] = 0;
                img[idx + 1] = 0;
                img[idx + 2] = 0;
                img[idx + 3] = 0xff;
            }
        }
    }

    // Requires lodepng.h
    // See: https://github.com/lvandeve/lodepng for download and compilation instructions
    write_png(requested_color_map + "_newton_fractal.png", img, image_width, image_height);
}