Mutlidimensional and Simultaneous Curve Fitting in Python using 'lmfit'

import numpy as np
grid = 101
x = y = np.linspace(0, 1, grid)
z = np.linspace(-0.1, 0.1, grid)
xx,yy,zz = np.meshgrid(x,y,z)

np.random.seed(1)
M = 5 # we arbitrarily choose 'M' to be 5 and 'N' to be 4
N = 4
E_x = np.zeros_like(xx)
E_y = np.zeros_like(xx)
E_z = np.zeros_like(xx)
As_true = {} # We'll store the 'true' values of A_m,n here.
for m in range(1, M+1):
    for n in range(1, N+1):
        A = np.random.uniform(0,3) # Setting A's to random 'true' values.
        As_true[f'A_{m}_{n}'] = A
        k = np.sqrt((m*np.pi)**2+(n*np.pi)**2)
        # Vector components are the sum of all the terms we defined 
        E_x += A*m*np.pi*np.cos(m*np.pi*xx)*np.sin(n*np.pi*yy)*np.sinh(k*zz) 
        E_y += A*n*np.pi*np.sin(m*np.pi*xx)*np.cos(n*np.pi*yy)*np.sinh(k*zz)
        E_z += A*k*np.sin(m*np.pi*xx)*np.sin(n*np.pi*yy)*np.cosh(k*zz)

E_x.shape

import pandas as pd 
df_true = pd.DataFrame({'x':xx.flatten(), 'y':yy.flatten(), 'z':zz.flatten(),
                        'E_x':E_x.flatten(), 'E_y':E_y.flatten(), 'E_z':E_z.flatten()})

df_slice = df_true.query('x==0.25') # Select all data in the y-z plane for x=0.25

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D # Needed for 3D plots
%matplotlib inline
plt.rcParams['figure.figsize'] = (8,6)

# We'll plot the 3 individual field components
for component in ['E_x', 'E_y', 'E_z']:
    fig = plt.figure()
    ax = plt.axes(projection='3d')
    # We reshape the 1D pandas columns into 2D arrays,
    # which is the format needed for the `plot_surface` arguments
    ax.plot_surface(df_slice.y.values.reshape((grid,-1)),    
                    df_slice.z.values.reshape((grid,-1)),
                    df_slice[component].values.reshape((grid,-1)),
                    cmap='viridis')
    ax.xaxis.labelpad = ax.yaxis.labelpad = 15
    plt.xlabel('Y')
    plt.ylabel('Z')
    ax.set_zlabel('$'+component+'$')
    plt.title('$'+component+'$ vs Y and Z for X=0.25')
plt.show()

df_test = df_true.sample(n=10000) # Take 10,000 randomly sampled points; this will be our input data.

# Lets make a scatter plot of the input data for a single 2D plane.
df_test_slice = df_test.query('x==0.25')
fig=plt.figure()
ax = plt.axes(projection='3d')
ax.scatter(df_test_slice.y, df_test_slice.z, df_test_slice.E_x,
           c=df_test_slice.E_x, s=15, cmap='viridis')
ax.xaxis.labelpad = ax.yaxis.labelpad = 15
plt.xlabel('X')
plt.ylabel('Y')
ax.set_zlabel(r'$E_x$')
plt.title('Input Data Sample')
plt.show()

# x, y and z are our input independent parameters.
# 'ms' and 'ns' are the total number of 'm' and 'n' integers that dictate our sum.
# **As will be our free parameters, and they'll be named like A_1_1, A_1_2...
def electric_field_3d(x, y, z, M, N, **As):
    E_x = np.zeros_like(x)
    E_y = np.zeros_like(x)
    E_z = np.zeros_like(x)
    for m in range(1, M+1):
        for n in range(1, N+1):
            A = As[f'A_{m}_{n}']
            k = np.sqrt((m*np.pi)**2+(n*np.pi)**2)
            E_x += A*m*np.pi*np.cos(m*np.pi*x)*np.sin(n*np.pi*y)*np.sinh(k*z)
            E_y += A*n*np.pi*np.sin(m*np.pi*x)*np.cos(n*np.pi*y)*np.sinh(k*z)
            E_z += A*k*np.sin(m*np.pi*x)*np.sin(n*np.pi*y)*np.cosh(k*z)
    # We return our estimates for E_x, E_y, E_z
    # as a single long 1D array using 'concatenate'.
    # We will have to remember this order, and mimic it
    # when we input data (but that shouldn't be to hard).
    return np.concatenate([E_x, E_y, E_z])
    

import lmfit as lm
# We define our model and tell it which inputs are independent variables
model = lm.Model(electric_field_3d, independent_vars=['x','y','z'])
# Now we can define our free parameters and initialize them.
params = lm.Parameters()
params.add('M', value=M, vary=False) # ms and ns shouldn't vary, they just
params.add('N', value=N, vary=False) # let us decide when to truncate our sum
for m in range(1, M+1):
    for n in range(1, N+1):
        params.add(f'A_{m}_{n}', value=0) # We define our free params, for our 'N'*'M' number of terms

# Our inputs need to be concatenated into a 1D array, the same order as we defined in our model
samples = np.concatenate([df_test.E_x.values, df_test.E_y.values, df_test.E_z.values])
# Now we input the dependent and independent variables, the free parameters, and make a call to Model.fit
result = model.fit(samples, x=df_test.x.values, y=df_test.y.values,
                   z=df_test.z.values, params=params)

# We can print the output of the fit, and we can compare it with the 'true' A_(m,n) that were used to generated.
print(result.fit_report())
print('True A values:')
for i in As_true:
    print(i, As_true[i])

# We evaluate the ModelResult using the high-granularity grid of x,y,z coordinates:
results_3_full = result.eval(x=df_true.x, y=df_true.y, z=df_true.z).reshape((3,-1))
init_3_full = result.eval(x=df_true.x, y=df_true.y, z=df_true.z, params=params).reshape((3,-1))
df_true['E_x_fit'] = results_3_full[0]
df_true['E_y_fit'] = results_3_full[1]
df_true['E_z_fit'] = results_3_full[2]
df_true['E_x_init'] = init_3_full[0]
df_true['E_y_init'] = init_3_full[1]
df_true['E_z_init'] = init_3_full[2]

df_test_slice = df_test.query('x==0.25') # Looking at the same 2D slice
df_slice = df_true.query('x==0.25')      # from our initial plots

fig=plt.figure()
ax = plt.axes(projection='3d')
ax.plot_wireframe(df_slice.y.values.reshape((grid,-1)),
                  df_slice.z.values.reshape((grid,-1)),
                  df_slice.E_x_fit.values.reshape((grid,-1)),
                  color='green', label='Best Fit', alpha=0.5)
ax.plot_wireframe(df_slice.y.values.reshape((grid,-1)),
                  df_slice.z.values.reshape((grid,-1)),
                  df_slice.E_x_init.values.reshape((grid,-1)),
                  color='k', label='Init Fit', alpha=0.3,
                  linestyle='--', rcount=10, ccount=10)
ax.scatter(df_test_slice.y, df_test_slice.z, df_test_slice.E_x,
           c='k', s=15, label='Data',
           depthshade=False, alpha=1, cmap='copper')
ax.xaxis.labelpad = ax.yaxis.labelpad = 15
ax.grid(False)
plt.xlabel('Y')
plt.ylabel('Z')
ax.set_zlabel(r'$E_x$')
plt.title('Fit Result Comparisons', pad=20)
plt.legend()
plt.show()