Imperative rework+

.gitignore

@@ -1,3 +1,4 @@
 **/__pycache__
 .DS_Store
+.python-version
 dist/

README.md

@@ -33,61 +33,63 @@ poetry install
 poetry shell
 ```
 
-<!-- TODO 
+<!-- TODO
 Figuring out how to install this this more permanently is a hanging question.
 Poetry will let you build a wheel and tarball with `poetry build`, but getting this into PyPI would be better.
 -->
 
-## Examples
-
-Note that `python` commands in this section should be run inside the virtual environment.
-
-### Default
-
-```bash
-python dcolorExample.py
-```
+## Usage
 
-It generates a series of examples.
-When each image appears clicking the close button starts the next example.
-Lambda expressions are used to define and pass functions to the plot() function.
-The last example is:
+Heres's a short script which plots a function with poles and zeros.
+A lambda expression is used to define the function, which is then passed to `dcolor.plot`.
 
 ```python
-dc.plot(lambda z : ((z**2-1)*(z-2- 1j)**2)/(z**2 +2+ 2j),
-    title='((z**2-1)*(z-2- 1j)**2)/(z**2 +2+ 2j)')
+import dcolor
+import matplotlib.pyplot as plt
+
+plt.title("$f(z) = \\frac{(z^2 - 1)(z - 2 - 1j)^2}{z^2 +2+ 2j}$")
+dcolor.dcolor(
+  lambda z : ((z**2 - 1)*(z - 2 - 1j)**2)/(z**2 + 2 + 2j),
+  xlim=(-8, 8),
+  ylim=(-8, 8),
+)
+# plt.show() called implicitly
 ```
 
-Which results in the following plot:
-
-![dcolor example](images/dcolor.png)
+## Examples
 
-### hsvcolor
+A selection of example functions can be plotted from `examples/example.py`.
+It contains several command-line options for both spawning interactive plots and
+generating static images.
 
-This is like dcolor except that it does not convert the HSV image to RGB
+The example images below can be plotted with the following command, which plots the
+same function from the usage section:
 
 ```bash
-python hsvcolorExample.py
+python examples/example.py -s algebraic.3 -c {METHOD}
 ```
 
-Will create the images.
-The last image produces:
+Note that this script may not function outside of the Poetry virtual environment.
 
-![hsvcolor example](images/hsvcolor.png)
+### Default method (magnitude_oscillating)
 
-### rgbcolor
+The default method assigns the complex argument to hue.
+Increasing magnitudes produce an effect resembling logarithmically-spaced contours.
 
-This is designed to show  the magnitude of abs(z).
-White means big, shades of green means intermediate, and black means small.
+![magnitude_oscillating example](images/magnitude_oscillating_algebraic_3.png)
 
-```bash
-python rgbcolorExample.py
-```
+### raw_magnitude_oscillating
+
+This is like the default, except that it does not convert the HSV image to RGB.
 
-Will create the images.
-The last image produces:
+![raw_magnitude_oscillating example](images/raw_magnitude_oscillating_algebraic_3.png)
+
+### green_magnitude
+
+This is designed to show  the magnitude of z.
+White means big, shades of green means intermediate, and black means small.
 
-![rgbcolor example](images/rgbcolor.png)
+![green_magnitude_example](images/green_magnitude_algebraic_3.png)
 
 
 ## Website and Documentation

dcolor/__init__.py

@@ -1,3 +1,31 @@
-from .dcolor import DColor
-from .hsvcolor import DColor as DColorHSV
-from .rgbcolor import DColor as DColorRGB
+"""
+dcolor
+
+Package for generating complex plots using [domain coloring](https://en.wikipedia.org/wiki/Domain_coloring).
+
+The main plotting function is `dcolor.dcolor`, with arguments as:
+
+    dcolor.dcolor(
+      lambda z : ...,
+      xlim=(-8, 8),
+      ylim=(-8, 8),
+    )
+
+By default, this will automatically show the plot in a matplotlib figure.
+For more information, see the documentation for `dcolor.dcolor`
+"""
+from .dcolor import dcolor, DColor, ComplexFunction
+from .color_maps import (
+    magnitude_oscillating,
+    raw_magnitude_oscillating,
+    green_magnitude,
+)
+
+__all__ = [
+    "dcolor",
+    "DColor",
+    "ComplexFunction",
+    "magnitude_oscillating",
+    "raw_magnitude_oscillating",
+    "green_magnitude"
+]

dcolor/color_maps.py

@@ -0,0 +1,125 @@
+"""
+dcolor.color_maps
+
+Module providing complex color maps used in domain coloring.
+These functions take a 2D array of complex values and convert them to RGB triples.
+"""
+from typing import Callable, Literal
+from typing_extensions import TypeAlias, TypeAliasType
+
+from matplotlib.colors import hsv_to_rgb
+import numpy as np
+import numpy.typing as npt
+
+__all__ = ["magnitude_oscillating", "raw_magnitude_oscillating", "green_magnitude"]
+
+ComplexPlane = TypeAliasType(
+    "ComplexPlane", np.ndarray[Literal[2], np.dtype[np.complexfloating]]
+)  # 2D representation of the complex plane
+Image = TypeAliasType(
+    "Image", np.ndarray[Literal[3], np.dtype[np.floating]]
+)  # RGB image
+# TODO: if going the Matplotlib-style norm, cmap kwargs, this may need its own class
+ComplexColorMap: TypeAlias = Callable[[ComplexPlane], Image]  # C -> [0, 1]^3
+
+
+def normalize(arr: npt.NDArray) -> npt.NDArray:
+    """Used for normalizing data in array based on min/max values"""
+    arrMin = np.min(arr)
+    arrMax = np.max(arr)
+    arr = arr - arrMin
+    return arr / (arrMax - arrMin)
+
+
+def magnitude_oscillating(zz: ComplexPlane) -> Image:
+    """
+    Converts a complex 2D array `zz` to an RGB image with normal domain coloring.
+
+    Hue is taken from the complex argument of `zz`.
+    Saturation and value vary with the logarithm of the magnitude of `zz`,
+    giving the appearance of logarithmic countours.
+    """
+    H = (np.angle(zz) % (2.0 * np.pi)) / (2.0 * np.pi)  # Hue determined by arg(z)
+    r = np.log2(1.0 + np.abs(zz))
+    S = (1.0 + np.abs(np.sin(2.0 * np.pi * r))) / 2.0
+    V = (1.0 + np.abs(np.cos(2.0 * np.pi * r))) / 2.0
+
+    return hsv_to_rgb(np.dstack((H, S, V)))
+
+
+def raw_magnitude_oscillating(zz: ComplexPlane) -> Image:
+    """
+    Converts a complex 2D array `zz` to an RGB image.
+
+    Same as `magnitude_oscillating`, but with a "rawer" conversion to RGB,
+    rather than going through matplotlib.colors
+    """
+    h = normalize(np.angle(zz) % (2.0 * np.pi))  # Hue determined by arg(z)
+    r = np.log2(1.0 + np.abs(zz))
+    s = (1.0 + np.abs(np.sin(2.0 * np.pi * r))) / 2.0
+    v = (1.0 + np.abs(np.cos(2.0 * np.pi * r))) / 2.0
+
+    r = np.empty_like(h)
+    g = np.empty_like(h)
+    b = np.empty_like(h)
+
+    i = (h * 6.0).astype(int)
+    f = (h * 6.0) - i
+    p = v * (1.0 - s)
+    q = v * (1.0 - s * f)
+    t = v * (1.0 - s * (1.0 - f))
+
+    idx = i % 6 == 0
+    r[idx] = v[idx]
+    g[idx] = t[idx]
+    b[idx] = p[idx]
+
+    idx = i == 1
+    r[idx] = q[idx]
+    g[idx] = v[idx]
+    b[idx] = p[idx]
+
+    idx = i == 2
+    r[idx] = p[idx]
+    g[idx] = v[idx]
+    b[idx] = t[idx]
+
+    idx = i == 3
+    r[idx] = p[idx]
+    g[idx] = q[idx]
+    b[idx] = v[idx]
+
+    idx = i == 4
+    r[idx] = t[idx]
+    g[idx] = p[idx]
+    b[idx] = v[idx]
+
+    idx = i == 5
+    r[idx] = v[idx]
+    g[idx] = p[idx]
+    b[idx] = q[idx]
+
+    idx = s == 0
+    r[idx] = v[idx]
+    g[idx] = v[idx]
+    b[idx] = v[idx]
+
+    return np.stack([r, g, b], axis=-1)
+
+
+def green_magnitude(zz: ComplexPlane, expand=1.0) -> Image:
+    """
+    Converts a complex 2D array `zz` to an RGB image.
+
+    Small magnitues are colored black, large ones are white, and values
+    between appear green.
+    The rate at which this occurs is controlled by `expand`.
+    """
+    absz = np.abs(zz)
+    r = absz * 0.5 / expand
+    g = absz * 1.00 / expand
+    b = absz * 0.5 / expand
+    r = np.clip(r, 0, 1)
+    g = np.clip(g, 0, 1)
+    b = np.clip(b, 0, 1)
+    return np.dstack((r, g, b))