Scalaで複素数とマンデルブロート集合

よくある題材です。
四則演算を作ろうとしたとき、引数が Complex のやつと Double のやつの２種類を作ろうと思ったけど、それだとメソッドが４つから８つになってしまう。
なので、初めて implicit を使ってみました。
あっさりできました。

で、四則演算ができれば、マンデルブロート集合も計算できるので、複素数版を作ってみました。
ちゃんと動きました。
計算式は、マンデルブロート集合の公式そのものです。シンプル♪

package asanmath

/**
 * 複素数を作ってみる。immutableで。
 */
case class Complex(val re:Double, val im:Double) {
    def this(re:Double) = this(re, 0)

    lazy val norm = re * re + im * im
    lazy val abs = math.sqrt(norm)
    /** 複素共役 */
    lazy val conjugate = Complex(re, -im)

    // 四則演算
    def +(that:Complex) = Complex(this.re + that.re, this.im + that.im)
    def -(that:Complex) = Complex(this.re - that.re, this.im - that.im)
    def *(that:Complex) = Complex(
            this.re * that.re - this.im * that.im,
            this.re * that.im + this.im * that.re)
    def /(that:Complex) = Complex(
            (this.re * that.re + this.im * that.im) / that.norm,
            (this.im * that.re - this.re * that.im) / that.norm)

    override def toString() = "%f%+fi".format(re, im)
}

object Complex {
    lazy val Zero = Complex(0, 0)
    implicit def Real2Complex(x: Double) = Complex(x, 0)
}

/** 複素数でマンデルブロート集合を作ってみる */
object ComplexMandelbrotTest {
    import java.awt.Color
    import java.awt.image.BufferedImage
    import java.io.File
    import javax.imageio.ImageIO
    /**
     * マンデルブロートの各点において、z^2 > 4 になるまでの回数を計算する。
     * maxCount回を超えると-1を返す。
     */
    def calcMandelbrot(c:Complex, maxCount:Int):Int = {
        var z = Complex.Zero
        for (i <- 0 until maxCount) {
            if (z.norm >= 4.0) return i
            z = z * z + c
        }
        -1
    }
    /**
     * 画像を作成する。
     */
    def createImage(m:Complex, zoom:Double, maxCount:Int):BufferedImage = {
        var image = new BufferedImage(500, 500, BufferedImage.TYPE_3BYTE_BGR)
        for (y <- 0 until image.getHeight; x <- 0 until image.getWidth) {
            val c = Complex(
                    m.re + (x-image.getWidth  / 2) / zoom,
                    m.im - (y-image.getHeight / 2) / zoom)
            val count = calcMandelbrot(c, maxCount)
            image.setRGB(x, y, 
                if (count != -1) {
                    Color.HSBtoRGB(count/100.0f, 1.0f, 1.0f)
                } else {
                    0   // BLACK
                })
        }
        image
    }
    def main(args:Array[String]) {
//      val image = createImage(Complex(-0.745428, 0.113009), 100, 100)
//      val image = createImage(Complex(-0.745428, 0.113009), 1000, 100)
//      val image = createImage(Complex(-0.745428, 0.113009), 10000, 500)
//      val image = createImage(Complex(-0.745428, 0.113009), 100000, 1000)
//      val image = createImage(Complex(-0.745428, 0.113009), 1000000, 1000)
//      val image = createImage(Complex(-0.745428, 0.113009), 10000000, 1500)
        val image = createImage(Complex(-0.745428, 0.113009), 100000000, 2000)  // 13秒ぐらい
        ImageIO.write(image, "png", new File("ComplexMandelbrotTest.png"))
    }       
}

/** テスト */
object ComplexTest {
    def main(args: Array[String]) {
        val a = new Complex(3, -4)
        println("a = " + a)             // "a = 3.000000-4.000000i"
        println("a.im = " + a.im)       // -4.0
        println("a.norm = " + a.norm)   // 25.0
        println("a.abs = " + a.abs)     // 5.0 ピタゴラスの三角形
        val b = Complex(1,1)
        println("b = " + b)             
        println("b.abs = "+ b.abs)      // 1.414213562373095i   √2
        println("a + b = "+ (a + b))    // 4.000000-3.000000i
        println("a - b = "+ (a - b))    // 2.000000-5.000000i
        println("a - 1 = "+ (a - 1))    // 2.000000-4.000000i
        println("a * 2 = "+ (a * 2))    // 6.000000-8.000000i
        println("a * b = "+ (a * b))    // 7.000000-1.000000i
        println("a / b = "+ (a / b))
        ComplexMandelbrotTest.main(args)
    }
}