main.cpp

// LibPNG example
// A.Greensted
// http://www.labbookpages.co.uk

// Version 2.0
// With some minor corrections to Mandlebrot code (thanks to Jan-Oliver)

// Version 1.0 - Initial release

#include <stdio.h>
#include <math.h>
#include <png.h>
#include <memory>
#include <vector>
#include <functional>
#include <complex>
#include <thread>
#include <algorithm>
#include <iterator>
#include <deque>
#include <mutex>
#include <condition_variable>

#include "argstream.h"



void HSVtoRGB(float& fR, float& fG, float& fB, float fH, float fS, float fV) {
  float fC = fV * fS; // Chroma
  float fHPrime = fmod(fH / 60.0, 6);
  float fX = fC * (1 - fabs(fmod(fHPrime, 2) - 1));
  float fM = fV - fC;
  
  if(0 <= fHPrime && fHPrime < 1) {fR = fC;fG = fX;fB = 0; }
  else if(1 <= fHPrime && fHPrime < 2) { fR = fX; fG = fC; fB = 0; }
  else if(2 <= fHPrime && fHPrime < 3) { fR = 0; fG = fC; fB = fX; } 
  else if(3 <= fHPrime && fHPrime < 4) { fR = 0; fG = fX; fB = fC; } 
  else if(4 <= fHPrime && fHPrime < 5) { fR = fX; fG = 0; fB = fC; } 
  else if(5 <= fHPrime && fHPrime < 6) { fR = fC; fG = 0; fB = fX; } 
  else { fR = 0; fG = 0; fB = 0; }
  
  fR += fM;
  fG += fM;
  fB += fM;
}

template<class _Tnumber, class _Titerator >
bool next_combination(_Titerator const& _First , _Titerator const& _Last, _Tnumber const& _Max )
{
    _Titerator _Current = _First;
    if( _Current  == _Last ) return false;
    *_Current += 1;
    if( *_Current < _Max ) return true;
    _Titerator _Next = _Current + 1;
    if( _Next == _Last ) return false;
    if( ! next_combination( _Next, _Last, _Max - 1 ) ) return false;
    *_Current = *_Next + 1; 
    return *_Current < _Max;
}


int iHeight = 1024;
int iWidth = 1024;
int poly_degree = 3;
double scale = 1.0;
std::complex<double> offset(0,0);
double spiral_scale = 1.0 , spiral_scale_factor = 1.0;
std::string base_filename;
double real_ellipse_factor = 1.0, imag_ellipse_factor = 1.0;
int num_cpus = std::thread::hardware_concurrency();
int max_iters = 20;
double rotation = 0.0;
double overlap = 1.0;
bool edges = false;
double saturation = 1.0;
double value = 1.0;
double hue_multiplier = 1.0;
double start_hue = 0.0;
int frames = 1;

int writeImage(const char* filename, int width, int height, int *buffer, const char* title);

class range_string
{
    public:
    std::string str;
    double& value;
    double min_range, max_range;
    std::string name;
    range_string(std::string n, double& v) : name(n), value(v) {}
    void parse()
    {        
        if (str.empty())
        {
            min_range = max_range = value;
            return;
        }
        if (sscanf(str.c_str(), "%lf:%lf", &min_range, &max_range) != 2)
        {
            if (sscanf(str.c_str(), "%lf", &value) != 1)
            {
                printf("Can't parse \"%s\" for parameter --%s\n", str.c_str(), name.c_str());
                throw std::runtime_error("Error parsing parameter");
            }
            min_range = max_range = value;
        }
    }
    void interp(int frame_num, int frames)
    {
        if (frame_num == 0 && frames == 0)
        {
            frames = 1;
        }
        value = min_range + (max_range - min_range) * (double(frame_num) / double(frames - 1));
        if (max_range == min_range)
        {
            printf("%s = %lf\n", name.c_str(), value);
        }
        else
        {
            printf("%s = %lf (%lf->%lf) frame %d of %d\n", name.c_str(), value, min_range, max_range, frame_num+1, frames);
        }
        
    }
};

int main(int argc, char *argv[])
{

    argstream::argstream<char> as(argc, argv);

    range_string scale_str("scale",scale);
    range_string real_ellipse_factor_str("real_ellipse_factor",real_ellipse_factor);
    range_string imag_ellipse_factor_str("imag_ellipse_factor", imag_ellipse_factor);
    range_string spiral_scale_str("spiral_scale", spiral_scale);
    range_string spiral_scale_factor_str("spiral_scale_factor", spiral_scale_factor);
    range_string rotation_str("rotation", rotation);
    range_string saturation_str("saturation", saturation);
    range_string value_str("value", value);
    range_string start_hue_str("start_hue", start_hue);
    range_string hue_multiplier_str("hue_multiplier", hue_multiplier);
    range_string overlap_str("overlap", overlap);

    std::vector<range_string*> range_strings{
        & scale_str,
        & real_ellipse_factor_str,
        & imag_ellipse_factor_str,
        & spiral_scale_str,
        & spiral_scale_factor_str,
        & rotation_str,
        & saturation_str,
        & value_str,
        & start_hue_str,
        & hue_multiplier_str,
        & overlap_str,
    };    

    as >> argstream::copyright("Newtons Fractal. Parameters marked with an asterisk may be specified as a range (a:b), or just a value.\nFrames are generated by interpolating ranged values.");
    as >> argstream::parameter('w', "width", iWidth, "Image Width", false);
    as >> argstream::parameter('h', "height", iHeight, "Image Height", false);
    as >> argstream::parameter('o', "file", base_filename, "Output filename (without extension)", true);
    as >> argstream::parameter('d', "degree", poly_degree, "Degree of polynomial",false);
    as >> argstream::parameter('r', "real_ellipse_factor", real_ellipse_factor_str.str, "Flatten to ellipse along real axis *", false);
    as >> argstream::parameter('i', "imag_ellipse_factor", imag_ellipse_factor_str.str, "Flatten to ellipse along imaginary axis *", false);
    as >> argstream::parameter('s', "scale", scale_str.str, "Scale image *", false);
    as >> argstream::parameter('p', "spiral_scale", spiral_scale_str.str, "Spiral scale start *", false);
    as >> argstream::parameter('f', "spiral_factor", spiral_scale_factor_str.str, "Spiral scale factor *", false);
    as >> argstream::parameter('n', "n_threads", num_cpus, "Number of threads to use", false);
    as >> argstream::parameter('m', "max_iters", max_iters, "Max iterations", false);
    as >> argstream::parameter('t', "rotation", rotation_str.str, "Radians to rotate roots by *", false);
    as >> argstream::parameter('v', "overlap", overlap_str.str, "Root circle overlap factor *", false);
    as >> argstream::parameter('e', "edges", edges, "Colour only the edges", false);
    as >> argstream::parameter('a', "saturation", saturation_str.str, "Colour saturation *", false);
    as >> argstream::parameter('l', "value", value_str.str, "Colour value *", false);
    as >> argstream::parameter('u', "start_hue", start_hue_str.str, "Start Hue *", false);
    as >> argstream::parameter('y', "hue_multiplier", hue_multiplier_str.str, "Hue multiplier *", false);
    as >> argstream::parameter('x', "frames", frames, "Number of frames to generate", false);

    for(auto rp : range_strings) rp->parse();

    as.defaultErrorHandling();

	// Make sure that the output filename argument has been provided
	if (base_filename.empty()) {		
        std::cout << as.usage();
        fprintf(stderr, "Please specify output file\n");
		return 1;
	}
    
    double aspect = double(iWidth) / double(iHeight);

    for(int frame_num = 0; frame_num < frames; frame_num++)
    {
        for(auto rp : range_strings) rp->interp(frame_num, frames);

        std::string filename = 
            frames == 1 ? base_filename + ".png" :
            base_filename + "-" + std::to_string(frame_num) + ".png";

        printf("Creating Image %s using %d threads\n", filename.c_str(), num_cpus);
        std::vector<int> buffer(iWidth * iHeight);

        // nth roots of unity, modified in various ways
        std::vector<std::complex<double>> roots(poly_degree);
        
        for(int i = 0; i < poly_degree; i++)
        {
            // Modify root according to parameters
            roots[i] = std::exp(std::complex<double>(0, rotation + (overlap*2.0*double(i)*3.141592654) / double(poly_degree)));
            roots[i].real( roots[i].real() * real_ellipse_factor );
            roots[i].imag( roots[i].imag() * imag_ellipse_factor );
            roots[i] *= spiral_scale;
            spiral_scale = spiral_scale_factor;
        }

        printf("Polynomial Roots\n");
        for(auto c : roots)
        {
            printf("  (%0.6f%+0.6fi)\n", c.real(), c.imag());
        }

        // from the roots, find the polynomal coeffs
        std::vector<std::complex<double>> coeffs(poly_degree+1);
        
        for(int power = 0; power < poly_degree; power++)
        {
            std::vector<int> term_indexes;
            for(int i = poly_degree - (1 + power); i >= 0; i--)
            {
                term_indexes.emplace_back(i);
            }
            do
            {
                std::complex<double> coeff(1,0);
                for(auto term_index : term_indexes)
                {   
                    coeff *= -roots[term_index];
                }
                coeffs[power] += coeff;
            } while (next_combination(term_indexes.begin(), term_indexes.end(), poly_degree));
        }    
        coeffs[poly_degree].real(1.0);
        int ipower = 0;
        printf("Polynomial\n");
        for(auto c : coeffs)
        {
            printf("  x^%d*(%0.6f%+0.6fi)\n", ipower, c.real(), c.imag());
            ipower++;
        }

        printf("Calculating...\n");

        // threads
        std::vector<std::thread> threads;
        std::deque<bool> lines_done;
        std::mutex lines_lock;
        std::condition_variable cv;
        for(int div = 0; div < num_cpus; div++)
        {
            threads.emplace_back([&, div]()
            {
                std::vector<double> distances(poly_degree);
                for(int x = div*(iWidth / num_cpus); x < (div+1)*(iWidth / num_cpus) ; x++)
                {
                    int last_value = 0;
                    for(int y = 0; y < iHeight ; y++)
                    {
                        std::complex<double> xn( double(x - iWidth/2) / double(iWidth / 2), double(y - iHeight/2) / double(iHeight / 2) );
                        xn.real(xn.real() * aspect);
                        xn *= scale;
                        xn += offset;
                        int i = 0;
                        do
                        {
                            // evaluate poly at xn
                            
                            std::complex<double> poly_value{};
                            for(int power = 0; power < poly_degree + 1; power++)
                            {
                                poly_value += coeffs[power] * std::pow(xn, power);
                            }

                            // derivative
                            std::complex<double> derv_value{};
                            for(int power = 1; power < poly_degree + 1; power++)
                            {
                                derv_value += coeffs[power] * std::pow(xn, power - 1) * double(power);
                            }
                            auto delta = -(poly_value / derv_value);
                            
                            if (std::norm(delta) < 0.00000001)
                            {
                            break;
                            }
                            xn += delta;
                            i++;
                            // break out toooo many
                        } while (i < max_iters);
                        // which root we closest to?
                        for(int root = 0; root < poly_degree; root++)
                        {
                            distances[root] = std::norm(xn - roots[root]);
                        }
                        auto this_value = std::min_element(distances.begin(), distances.end()) - distances.begin();
                        if (edges)
                        {
                            if (this_value != last_value)
                            {
                                buffer[y*iWidth + x] = 0;
                            }
                            else
                            {
                                buffer[y*iWidth + x] = 1;
                            }
                            last_value = this_value;
                        }
                        else
                        {
                            buffer[y*iWidth + x] = this_value;
                        }                    
                    }
                    std::lock_guard<std::mutex> _lk(lines_lock);
                    lines_done.push_front(true);
                    cv.notify_one();
                }
                std::lock_guard<std::mutex> _lk(lines_lock);
                lines_done.push_front(false);
                cv.notify_one();
            });
        }
        // Monitor progress
        {
        int done = num_cpus;
        int completed = 0;
        while(done > 0)
        {
            std::unique_lock<std::mutex> ul(lines_lock);
            cv.wait(ul, [&](){ return !lines_done.empty(); });
            if (lines_done.back())
            {
            completed++;
            printf("%0.01f%%\r  ",100.0*(double(completed) / (double(iWidth))));
            fflush(stdout);
            }
            else
            {
            done--;
            }
            lines_done.pop_back();
        };
        }

        // Join threads (they should all be complete anyway now)
        for(auto& t : threads)
        {
            t.join();
        }

        // Save the image to a PNG file
        // The 'title' string is stored as part of the PNG file
        printf("\nDone - Saving PNG\n");
        if ( writeImage(filename.c_str(), iWidth, iHeight, buffer.data(), "This is my test image") != 0)
        {
            return 1;
        }
    }
    return 0;
}

inline void setRGB(png_byte *ptr, int val)
{
    float r,g,b;
    HSVtoRGB(r,g,b,start_hue + hue_multiplier*((360.0f * float(val)) / float(poly_degree)), saturation, value);
    ptr[0] = png_byte(r*255.0f);
    ptr[1] = png_byte(g*255.0f);
    ptr[2] = png_byte(b*255.0f);
}

int writeImage(const char* filename, int width, int height, int *buffer, const char* title)
{
	// Open file for writing (binary mode)
	std::unique_ptr<FILE, decltype(&fclose)> fp(fopen(filename, "wb"), &fclose);
	if (!fp) {
		fprintf(stderr, "Could not open file %s for writing\n", filename);
        return 1;
	}

	// Initialize write structure
	std::unique_ptr<png_struct, std::function<void(png_struct*)>> png_ptr(png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL), [](png_struct*s)
    {
        png_destroy_write_struct(&s, nullptr);
    });
	if (!png_ptr) return 1;

	// Initialize info structure
	std::unique_ptr<png_info, std::function<void(png_info*)>> info_ptr(png_create_info_struct(png_ptr.get()), [&](png_info*p)
    {
        png_free_data(png_ptr.get(), p, PNG_FREE_ALL, -1);
    });
	if (!info_ptr) return 1;

	png_init_io(png_ptr.get(), fp.get());

	// Write header (8 bit colour depth)
	png_set_IHDR(png_ptr.get(), info_ptr.get(), width, height,
			8, PNG_COLOR_TYPE_RGB, PNG_INTERLACE_NONE,
			PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

	// Set title
	if (title != nullptr) {
		png_text title_text;

		title_text.compression = PNG_TEXT_COMPRESSION_NONE;
		title_text.key = (char*)"Title";
		title_text.text = (char*)title;
		png_set_text(png_ptr.get(), info_ptr.get(), &title_text, 1);
	}

	png_write_info(png_ptr.get(), info_ptr.get());

	// Allocate memory for one row (3 bytes per pixel - RGB)
	auto row = std::vector<png_byte>(3 * width);

	// Write image data
	int x, y;
	for (y=0 ; y<height ; y++) {
		for (x=0 ; x<width ; x++) {
			setRGB(&(row[x*3]), buffer[y*width + x]);
		}
		png_write_row(png_ptr.get(), row.data());
	}

	// End write
	png_write_end(png_ptr.get(), NULL);

	return 0;
}