libs/math/example/color_maps_example.cpp
// (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); }