Program 9 - Basins of Attraction for Newton's Method.

Program Code

Off[General::"spell1", General::"spell"];

Needs["Graphics`Colors`"];

BeginPackage["NewtonPoly`"];

Unprotect[NewtonPoly, RandomFractal];

NewtonPoly::usage=
    "NewtonPoly[poly, n] applies Newton's method to the complex-valued \
polynomial in one variable given by poly to create an figure which shows the \
basins of attraction for the distinct roots of poly.  NewtonPoly[list, n] \
creates a polynomial with complex roots given in list, then generates the \
image.";

RandomFractal::usage=
    "RandomFractal[deg,n] randomly generates a polynomial of degree deg, then \
maps the basins of attration using NewtonPoly.";

Options[NewtonPoly] = {
	Coverage -> Filled,
	MaxIterations -> 50,
	PlotStyle -> {},
	PlotRange -> Automatic,
	ColorFunction -> Hue,
	Colors -> Automatic,
	DefaultColor -> RGBColor[0,0,0],
	ShowRoots -> False,
	RootPlotStyle -> {PointSize[0.04]}
};

Begin["`Private`"];

(*  Extra color conversion functions  *)
(*  Extension of color conversions in Graphics`Colors`  *)
HSBColor[h_] := Graphics`Colors`HSBColor[h, 1, 1];
CMYColor[c_, m_, y_, k_] := Graphics`Colors`CMYColor[(1 - k) c, (1 - k) m, (1 - k) y];

(*  My own color converters  *)
Gray2RGB[GrayLevel[g_]] := RGBColor[g, g, g];
AnyColor2RGB[c_] := Switch[Head[c],
	RGBColor, c, 
	Hue, HSBColor @@ c,
	GrayLevel, Gray2RGB[c], 
	CMYKColor, CMYColor@@c,
	_, c
];

(*  Color testing functions  *)
RGBColorQ[c_] := MatchQ[c, RGBColor[r_/;r >= 0 && r <= 1, g_/;g >= 0 && g <= 1, b_/;b >= 0 && b<=1]];
ColorQ[c_]:=RGBColorQ[AnyColor2RGB[c]];

(*  Convert Complex number to list  *)
ComplexToList := z_Complex -> {Re[z], Im[z]};

RealQ[r_] := NumericQ[r] && Im[r] == 0;

NewtonPoly::badlist = "`1` is not a list of complex roots.";
NewtonPoly::badvars = "Too many variables `1` in equation `2`.";
NewtonPoly::notpoly = "Equation must contain polynomials in `1`, instead of `2`.";
NewtonPoly::badargs = "Options expected (instead of `1`) beyond position 2 in NewtonPoly.";
NewtonPoly::badcov = "Coverage must be either Filled, Random or Uniform, instead of `1`.";
NewtonPoly::badmax = "MaxIterations must be a positive integer, instead of `1`.";
NewtonPoly::badran = "Plot range must be a list of the form {{xmin,xmax},{ymin,ymax}}, instead of `1`.";
NewtonPoly::badcf = "Color function `1` does not return a standard graphics color directive.";
NewtonPoly::badcl = "Color list `1` is not a list of color directives.  Using default color scheme.";
NewtonPoly::baddc = "Default color `1` is not a valid color directive.  Using black instead.";
NewtonPoly::badn = "Number of points `1` is not a positive integer.";

NewtonPoly[eqn_,n_,args___] := Module[{f, i, x, y, z, dx, dy, xmin, xmax, ymin, ymax, sols, var, cov, 
		max, pr, dc, cf, temp, len, findcolor, pointopts, plotopts, result, colors, showdots, rootopts},
	If[Head[eqn] === List,
		(*  See if eqn is a list of complex roots  *)
		If[!VectorQ[eqn, NumericQ],
			Message[NewtonPoly::badlist, eqn]; Return[],
			sols = Union[N[eqn + 0.0I]]; f[z_] = Expand[Times @@ (z - sols)]
		],
		
		(*  See if eqn is an equation  *)
		temp = If[Head[eqn] === Equal, eqn[[1]] - eqn[[2]], eqn];
		var = Variables[temp];
		If[Length[var] != 1 && And @@ (Head[#1] === Symbol& /@ var),
			Message[NewtonPoly::badvars,var,eqn]; Return[]];
		
		(*  See if eqn is a polynomial  *)
		var = var[[1]];
		f[z_] = N[temp /. var -> z];
		If[!PolynomialQ[f[z], z],
			Message[NewtonPoly::notpoly, var, eqn]; Return[]];
		
		(*  Find roots of the equation  *)
		sols = Union[(var /. #)& /@ NSolve[temp == 0, var] + 0.0I];
    ];
	len=Length[sols];
	
	If[!IntegerQ[n] || n < 1, Message[NewtonPoly::badn,n]; Return[]];
	
	(*  Check argument list  *)
	If[i=0;Scan[(i++;If[Head[#1]=!=Rule,Return[True]])&,{args}],
		Message[NewtonPoly::badargs,{args}[[i]]];Return[]];
	
	(*  Extract and verify various options  *)
	{cov, max, pr, pointopts, colors, cf, dc, showdots, rootopts} = {
		Coverage, MaxIterations, PlotRange, PlotStyle, Colors, ColorFunction, 
		DefaultColor, ShowRoots, RootPlotStyle
	} /. {args} /. Options[NewtonPoly];
	
	If[!MemberQ[{Filled,Random,Uniform}, cov],
		Message[NewtonPoly::badcov, cov]; Return[]];
	
	If[Head[max] =!= Integer || max < 1,
		Message[NewtonPoly::badmax, max]; Return[]];
	
	If[pr === Automatic,
		temp = Plus @@ sols / len;
		pr = (temp /. ComplexToList) + (Max[Abs[sols - temp]]){{-1, 1}, {-1, 1}},
		If[pr === All,
			pr = {Min[#1], Max[#1]}& /@ Transpose[sols /. ComplexToList],
			If[!MatrixQ[pr, RealQ] || Length[pr] != 2 || Length[pr[[1]]] != 2,
				Message[NewtonPoly::badran,pr];Return[]]
        ];
    ];
	{xmin, xmax, ymin, ymax} = N[Flatten[pr]];
	If[xmin == xmax, xmin -= 1.0; xmax += 1.0];
	If[ymin == ymax, ymin -= 1.0; ymax += 1.0];
	
	pr = AnyColor2RGB[cf[# / len]]& /@ Table[i, {i, 1, len}];
	If[Head[colors] === List,
		temp = AnyColor2RGB /@ colors;
		If[!VectorQ[temp, RGBColorQ],
			Message[NewtonPoly::badcl, colors]; colors = Automatic,
			If[Length[temp] < len,
				colors = Join[temp, Take[pr, Length[temp] - len]],
				colors = temp
			];
		];
	];
	If[colors === Automatic, colors = pr];
	
	(*  Look at default color  *)
	If[!RGBColorQ[temp = AnyColor2RGB[dc]],
		Message[NewtonPoly::baddc, dc];
		dc = DefaultColor /. Options[NewtonPoly],
		dc = temp
	];
	colors = Join[{dc}, colors];
	
	If[showdots == True, showdots = Point[#1 /. ComplexToList]& /@ sols, showdots = {}];
	
	plotopts = DeleteCases[{args}, (Coverage -> _) | (MaxIterations -> _) | (PlotRange -> _) | 
		(PlotStyle -> _) | (ColorFunction -> _) | (DefaultColor -> _) | (Colors->_) | 
		(ShowRoots -> _) | (RootPlotStyle -> _)
	];
	
	f = Compile[{{z, _Complex}}, Evaluate[z - f[z] / f'[z]]];
	findcolor = Compile[{{z, _Complex}}, Module[{i = 1},
		For[i = Floor[Length[sols]], i >= 1 && Abs[z - sols[[i]]] > 0.01, i--, Null];
		i], {{sols, _Complex, 1}}
	];
	
	(*  Need something like this to error-trap for cases where f'[z] is zero  *)		
	temp = Indeterminate -> ComplexInfinity;
	Off[CompiledFunction::cfcx, CompiledFunction::cfn, 
		CompiledFunction::cfsa,Power::infy,Infinity::indet];
	
	result = Show[Graphics[Switch[cov,
		Uniform, Flatten[{pointopts, Table[
			{colors[[findcolor[FixedPoint[f, p = x + I y, max] /. temp] + 1]], Point[{Re[p], Im[p]}]}, 
			{x, xmin, xmax, (xmax - xmin) / Floor[Sqrt[n] - 1]}, 
			{y, ymin, ymax, (ymax - ymin) / Floor[Sqrt[n] - 1]}
		]}, 1],
		Random, Flatten[{pointopts, Table[
			p = Random[Complex, {xmin + I ymin, xmax + I ymax}];
			{colors[[findcolor[FixedPoint[f, p, max] /. temp] + 1]], Point[{Re[p], Im[p]}]}, {n}
		]}, 1],
		Filled,
			dy = (ymax - ymin) / Floor[Sqrt[n]];
			dx = (xmax - xmin) / Floor[Sqrt[n]];
			RasterArray[Table[
				colors[[findcolor[FixedPoint[f, x + I y, max] /. temp] + 1]],
				{y, ymin + dy / 2, ymax, dy}, {x, xmin + dx / 2, xmax, dx}
			], {{xmin, ymin}, {xmax, ymax}}]
	]], Graphics[Flatten[{rootopts, showdots}, 1]], plotopts];
	
	On[CompiledFunction::cfcx, CompiledFunction::cfn, 
		CompiledFunction::cfsa, Power::infy,Infinity::indet];
	Return[result];
];
RandomFractal[deg_Integer, n_Integer, args___] := Module[{sols, z},
	sols = Table[Random[Complex, {-1 - I, 1 + I}], {deg}];
	Print[Expand[Times @@ (z - sols)] /. z -> "z"];
	Return[NewtonPoly[sols, n, args]];
];

End[];

SetAttributes[{NewtonPoly, RandomFractal}, ReadProtected];
Protect[NewtonPoly, RandomFractal];

EndPackage[];

Off[General::"spell1", General::"spell"];

Examples

NewtonPoly[z^3 - 1, 5000, Coverage -> Random, 
      PlotStyle -> PointSize[0.02]];

NewtonPoly[z^3 + 1, 1000, ShowRoots -> True];

NewtonPoly[z^4 + I z^2 + z - 10, 2000];