Dispersion: Symbolic Regression
```python
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")
```
```python import matplotlib.pyplot as plt import numpy as np import xarray as xr
ds = xr.load_dataset("data/neff_te0_db.nc") print(ds) ```
```python 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
```python 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)
```
```python
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) ```
```python
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() ```
```python 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() ```
```python
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() ```