照度分布計算その6 - 複素2次元配列の演算の実装

回折計算によるレーザ光学系の照度分布の計算。やるぞ、実装、実装。まず、計算の中心となる複素数2次元配列のクラスから。

まえもちょっと書いたけどSAmplitudePlaneクラスは

@interface SAmplitudePlane : NSObject  {
    complex *cPlane;
    int     sSize;	//	square shape
}

みたいなインスタンス変数を持ってる。cPlaneは複素数の2次元配列へのポインタで、sSizeはその1次元分の長さ。cPlaneは縦横とも同じ長さの正方配列だけにするのでサイズはひとつでいいことにする。

このクラスのインスタンスの初期化の本体は

- (id)initWithAmplitudeFunction:(SAmplitudeFunction *)aFunc
            coordinateConverter:(SCoordinateConverter *)converter
                         ofSize:(int)squareSize
{
    self= [super init];
    sSize = squareSize;
    cPlane = (complex *)malloc(sSize * sSize * sizeof(complex));
    if (cPlane == NULL) {
        [self release];
        return nil;
    }
    if ((aFunc != nil) && (converter != nil))
        [self fillBy:aFunc converter:converter];
    return self;
}

こんなの。振幅関数SAmplitudeFunctionと、配列のインデクス値を振幅関数の実数値に変換するSCoordinateConverterを受け取って複素数の配列を-fillBy:converter:メソッドで初期化する。どっちかがnilなら配列は初期化しないで放っておく。もちろんdealloc（あるいはガベージコレクション環境ならfinalize）を

- (void)dealloc
{
    if (cPlane != NULL)
        free(cPlane);
    [super dealloc];
}

みたいに書くのを忘れないようにする。

SAmplitudePlaneは他に

- (int)squareSize;
- (complex *)cPlane;
- (id)addByPlane:(SAmplitudePlane *)otherplane;
- (id)multiplyByPlane:(SAmplitudePlane *)otherplane;
- (id)scaleBy:(double)factor;
- (id)rotateHalf;
- (id)fourier;
- (id)inverseFourier;
- (id)intensityMap;

のようなメソッドを持つ。始めのふたつはいいとしてその後は、例えば-addByPlane:メソッドはふたつのSAmplitudePlaneのインスタンスの配列同士を足し算した新しいインスタンスを返すメソッドで

- (id)addByPlane:(SAmplitudePlane *)otherPlane
{
    if (sSize != [otherPlane squareSize])
        return nil;
    //	else
    SAmplitudePlane	*newPlane = [[[SAmplitudePlane alloc]
                 initWithAmplitudeFunction:nil
                       coordinateConverter:nil
                                    ofSize:sSize] autorelease];
    complex *p = cPlane;
    complex *op = [otherPlane cPlane];
    complex *np = [newPlane cPlane];
    int     i;
    int     s2 = sSize * sSize;
	
    for (i = 0 ; i < s2 ; i ++)
        *(np ++) = cAdd(*(p ++), *(op ++));
    return newPlane;
}

というようなもの。関数cAdd()は複素数の足し算をする関数。こういうのはいろんなところにライブラリがあるけど探しているうちに書けてしまうのでいちいち書いてる。上手くやればベクトルユニットが使えたりするはずだけど、今回は面倒なのでやらない。

ほかのかけ算や、定数倍（-scaleBy:メソッド）も全く同じ。Fourier変換をやる-fourierメソッドは内部でfftwライブラリを呼ぶだけ。

- (id)fourier
{
    return [self switchFourier:FFTW_FORWARD];
}

- (id)inverseFourier
{
    return [self switchFourier:FFTW_BACKWARD];
}

- (id)switchFourier:(int)direction
{
    int     i;
    int     s2 = sSize * sSize;
    double  r = 1.0 / sSize;
    SAmplitudePlane	*newPlane = [[[SAmplitudePlane alloc]
                            initWithAmplitudeFunction:nil
                                  coordinateConverter:nil
                                               ofSize:sSize] autorelease];
    complex *np = [newPlane cPlane];
    fftw_plan plan = fftw_plan_dft_2d(sSize, sSize,
                                      (fftw_complex *)cPlane,
                                      (fftw_complex *)np,
                                      direction, FFTW_ESTIMATE);
    fftw_execute(plan);
    fftw_destroy_plan(plan);
    for (i = 0 ; i < s2 ; i ++) {
        *np = cScale(*np, r);
        np ++;
    }
    return newPlane;
}

説明するまでもない簡単なもの。後半で前回書いた総エネルギーを一定にするためにスケーリングしている。ほかも同じ発想で、上手く動いたやつをコピペしてループの中で呼んでる関数だけを変える、という方法で書いてある。ダサイって？いや、この単純さがバグを呼び込まない方法のひとつ、と僕は思っている。

この他のメソッドとして

- (complex)bilinearedValueAtIndexX:(double)xi andY:(double)yi;	//	take fractional index
- (double)bilinearedRealValueAtIndexX:(double)xi andY:(double)yi;
- (id)shrinkedPlaneWithSqareSize:(int)newSize;

を書いておく。bilineared~というのは前回説明した整数インデクス値の途中の値をバイリニア近似で返す物で

- (complex)bilinearedValueAtIndexX:(double)xi andY:(double)yi
{
    int     x = (int)xi;
    int     xp1 = (x + 1) % sSize;
    int     y = (int)yi;
    int     yp1 = (y + 1) % sSize;
    double  dx = xi - (double)x;
    double  dy = yi - (double)y;
    complex i11 = cPlane[sSize * y + x];
    complex i21 = cPlane[sSize * yp1 + x];
    complex i12 = cPlane[sSize * y + xp1];
    complex i22 = cPlane[sSize * yp1 + xp1];
	
    double  c1 = 1.0 - dx - dy + dx * dy;
    double  c2 = dx * (1.0 - dy);
    double  c3 = dy * (1.0 - dx);
    return complexBy(c1 * i11.re + c3 * i21.re + c2 * i12.re + dx * dy * i22.re,
                     c1 * i11.im + c3 * i21.im + c2 * i12.im + dx * dy * i22.im);
}

みたいなもの。これは前回書いた式をそのままコーディングしている。cPlaneがFFT対象なので、トーラスになっていることだけ注意すれば問題ない。

もうひとつの-shrinkedPlaneWithSqareSize:メソッドは表示とかに使うために小さい配列を返す。このとき配列の中心が[0,0]の位置で、負の添字はトーラスになった逆の端の方から取り出す必要があることだけ注意すればいい。ソースを示すまでもない。

ああっ、たったここまで書いてもうジジイのおねむの時間になってしまった。ホーガンのほんの感想なんか書いてたせいもあるけど、コードを書き出すと時間の経つのは早い。残りは明日にしよ。忘れないように寝ながら頭の中でコーディングする。そうするとすぐ寝てしまう。それがいいのやらわるいのやらわからんけど。