Pythonのscipy.optimize.least_squares()で複素数のデータを非線形の複素関数に最小二乗法でフィッティングする。

複素数のデータ、例えばSパラメータを入力にしたとき、あるパラメータを持つ出力関数（これも複素数、かつ普通は非線形）にフィッティングしたいことは良くある。が、scipyの最小二乗法でどうやるのかあまり書いてあるものをみたことがない。

とりあえずあまりスマートじゃないができるようになったので記録として残す。

ポイントは、複素数はoptimize.least_squaresで最小化できないので、いったん実数のラッパーを作ること。

こんな感じ。

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import least_squares

# z2 = a (z1 - b)^2 + c という複素関数とする。
def fit_func(a, b, c, z1):
    return a * (z1 - b)**2 + c
def res_func(a, b, c, z1, z2):
    return fit_func(a, b, c, z1) - z2
#複素数を使うには実数のラッパーが必要
def res_func_wrap(param, x1, x2):
    fwrap = res_func(param[0] + 1j*param[1], param[2] + 1j*param[3], param[4] + 1j*param[5],
                                 x1[:,0]+1j*x1[:,1], x2[:,0]+1j*x2[:,1])
    return fwrap.real ** 2 + fwrap.imag ** 2

#推定するa, b, c
a = 1 + 2j
b = -1 + 3j
c = -2 - 1j

#データはn点
n=1000

#z1 と z2を計算(乱数でばらつかせる)
z1 = np.zeros(n, dtype = complex)
z2 = np.zeros(n, dtype = complex)
for i in range(n):
    z1[i] = (2 * np.random.rand() - 1) + 1j*(2 * np.random.rand() - 1)
    z2[i] = fit_func(a, b, c, z1[i]) + 0.1*( (2 * np.random.rand() - 1) + 1j*(2 * np.random.rand() - 1) )

#実数のアレイに直す。
x1 = np.zeros((n,2))
x2 = np.zeros((n,2))
for i in range(n):
    x1[i, 0] = z1[i].real
    x1[i, 1] = z1[i].imag
    x2[i, 0] = z2[i].real
    x2[i, 1] = z2[i].imag
#初期値
initial_values = [0, 0, 0, 0, 0, 0]

res = least_squares(res_func_wrap, initial_values, args = (x1, x2))
print(res.x)

実行すると

[ 1.00113061 2.00065746 -1.00031483 2.9985901 -1.99946214 -1.00700744]

となった。実際は

[1,2,-1,3,-2,-1]

なのでちゃんとパラメータ推定できていそう。ただなんかもっさりしているので、もっといいやり方を誰か知っていたら教えてほしいなあ。