Dispersion: Symbolic Regression

# import doModels.RefractiveIndex as ri
# from doModels import fem, PhotonicStack
# from numpy.polynomial import Polynomial

# import matplotlib.pyplot as plt
# import numpy as np
# from tqdm.notebook import tqdm


# stack = PhotonicStack.create(
#     total_film_thickness=0.210,
#     etch_depth=0.210,
#     top_width=0.450,
#     side_wall_angle=0,
#     top_oxide_thickness=2.0,
#     bottom_oxide_thickness=2.0,
# )

# wavelengths = np.arange(1.4, 1.7, 2*0.005)
# top_widths = np.arange(0.300, 0.700, 2*0.02)

# num_modes = 1

# all_neffs = np.zeros((top_widths.shape[0], wavelengths.shape[0], num_modes), dtype=complex)

# for i, top_width in enumerate(tqdm(top_widths)):
#     stack.top_width.target = top_width
#     m = fem.mesh_waveguide(total_film_thickness=stack.total_film_thickness.target,
#             etch_depth=stack.etch_depth.target,
#             top_width=stack.top_width.target,
#             side_wall_angle=stack.side_wall_angle.target, 
#             show_plot=False)

#     for j, wavelength in enumerate(wavelengths):    
#         modes = fem.fem_waveguide(mesh=m, n_core=ri.silicon, n_clad=ri.silica, n_box=ri.silica, wavelength=wavelength, num_modes=num_modes)
#         all_neffs[i,j] = [mode.n_eff for mode in modes]


# import xarray as xr

# neff_te0 = np.real(all_neffs[:, :, 0])

# xarr = xr.DataArray(
#     data=neff_te0, 
#     coords={
#         "top_width": top_widths,
#         "wavelength": wavelengths,
#     },
#     dims=["top_width", "wavelength"],
#     name="neff_te0"
# )

# xarr.to_netcdf("data/neff_te0_db.nc")

import matplotlib.pyplot as plt
import numpy as np
import xarray as xr


ds = xr.load_dataset("data/neff_te0_db.nc")
print(ds)

wavelengths = ds["wavelength"].values
top_widths = ds["top_width"].values
neff = ds["neff_te0"].values

plt.figure()
plt.xlabel("Wavelength (µm)")
plt.ylabel("Effective refractive index")

for i in range(0, len(top_widths), 4):
    plt.plot(wavelengths, neff[i], label=f"Top width: {top_widths[i]:.3f} µm")

plt.legend()
plt.show()

Symbolic Regression

import numpy as np
from pysr import PySRRegressor
import xarray as xr
from itertools import product
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score, mean_squared_error
from matplotlib import pyplot as plt

ds = xr.load_dataset("data/neff_te0_db.nc")

coord_names = ["top_width", "wavelength"]
coords = [ds.coords[name].values for name in coord_names]

X = np.array(list(product(*coords)))
y = ds["neff_te0"].values.flatten()

# Normalize features
scaler_X = StandardScaler()
X_scaled = scaler_X.fit_transform(X)

scaler_y = StandardScaler()
y_scaled = scaler_y.fit_transform(y.reshape(-1, 1)).flatten()


model = PySRRegressor(
    niterations=400,
    populations=40,
    maxsize=40,
    verbosity=0
)
model.fit(X_scaled, y_scaled)

# Get all equations and their metrics
equations_df = model.equations_

# Plot complexity vs loss (Pareto front)
plt.figure(figsize=(10, 4))

plt.subplot(1,2,1)
plt.scatter(equations_df['complexity'], equations_df['loss'], alpha=0.7)
plt.xlabel('Complexity')
plt.ylabel('Loss')
plt.title('Pareto Front: Model Complexity vs Loss')
plt.yscale('log')  # Often helpful since loss can vary by orders of magnitude
plt.grid(True, alpha=0.3)

# Plot loss vs equation index (discovery order)
plt.subplot(1,2,2)
plt.plot(equations_df.index, equations_df['loss'], 'o-', alpha=0.7)
plt.xlabel('Equation Discovery Order')
plt.ylabel('Loss')
plt.title('Loss Evolution During Search')
plt.yscale('log')
plt.grid(True, alpha=0.3)

# For predictions, transform back
y_pred_scaled = model.predict(X_scaled)
y_pred = scaler_y.inverse_transform(y_pred_scaled.reshape(-1, 1)).flatten()


print(f"R² score: {r2_score(y, y_pred):.4f}")
print(f"RMSE: {np.sqrt(mean_squared_error(y, y_pred)):.4f}")

# Plot comparison
plt.figure(figsize=(10, 4))
plt.subplot(1, 2, 1)
plt.scatter(y, y_pred, alpha=0.5)
plt.plot([y.min(), y.max()], [y.min(), y.max()], "r--")
plt.xlabel("True values")
plt.ylabel("Predicted values")
plt.title("Prediction vs True")

plt.subplot(1, 2, 2)
residuals = y - y_pred
plt.scatter(y_pred, residuals, alpha=0.5)
plt.axhline(y=0, color="r", linestyle="--")
plt.xlabel("Predicted values")
plt.ylabel("Residuals")
plt.title("Residual Plot")

plt.tight_layout()

import numpy as np
import matplotlib.pyplot as plt

wavelengths = ds.wavelength.values
selected_top_widths = ds.top_width.values[::4]


plt.figure(figsize=(10, 4))

plt.subplot(1,2,1)
for tw in selected_top_widths:
    if tw in ds.top_width.values:  # Only plot if the exact value exists
        plt.plot(ds.wavelength, ds.neff_te0.sel(top_width=tw), 
                label=f"{tw:.2f} µm", linestyle="--", alpha=0.7)
plt.xlabel("Wavelength (µm)")
plt.ylabel("n_eff")
plt.title("Original Data")
plt.legend()

plt.subplot(1,2,2)
for tw in selected_top_widths:
    # Create input array for this top_width across all wavelengths
    X_pred = np.column_stack([
        np.full(len(wavelengths), tw),  # constant top_width
        wavelengths                     # varying wavelength
    ])

    X_pred_scaled = scaler_X.transform(X_pred)
    y_pred_scaled = model.predict(X_pred_scaled)
    y_pred = scaler_y.inverse_transform(y_pred_scaled.reshape(-1, 1)).flatten()

    plt.plot(wavelengths, y_pred, label=f"{tw:.2f} µm", linewidth=2)

plt.xlabel("Wavelength (µm)")
plt.ylabel("n_eff")
plt.title("PySR Model Predictions")
plt.legend()
plt.tight_layout()

# from sympy import symbols
# a = model.sympy()
# a.subs([(symbols("x0"), symbols("w")),
#         (symbols("x1"), symbols("λ"))])
# # type(a)


# 1. Get the scaled symbolic expression
expr_scaled = model.sympy()

# 3. Apply inverse y-scaling
y_mean = scaler_y.mean_[0]
y_std = scaler_y.scale_[0]
expr_rescaled = y_std * expr_scaled + y_mean

model.sympy()

x₁: 0.50