Space weather forecasting depends heavily on predicting the solar wind speed at Earth, since fast solar wind streams driven by coronal holes are one of the main triggers of geomagnetic storms. Operational forecasting centers rely on semi-empirical models — most famously the Wang-Sheeley-Arge (WSA) family — that map coronal magnetic field properties observed at the Sun to the solar wind speed expected at 1 AU. These models are cheap to run compared to full magnetohydrodynamic simulations, but their accuracy depends entirely on a handful of tunable parameters that must be fitted against real observations.
In this article we build a WSA-style empirical solar wind speed model, generate a synthetic but physically realistic observation set, and use nonlinear least-squares optimization to recover the model’s parameters. We then visualize the fit quality and the shape of the optimization cost landscape in 3D.
The model
The WSA-type formulation predicts the solar wind speed $v$ at 1 AU as a function of two coronal quantities measured near the Sun:
- $f_s$ — the magnetic flux tube expansion factor, which measures how strongly a flux tube fans out between the photosphere and the source surface. Large $f_s$ corresponds to slow wind; small $f_s$ (open, weakly expanding field lines typical of coronal holes) corresponds to fast wind.
- $\theta_b$ — the normalized angular distance of the footpoint from the nearest coronal hole boundary. Points deep inside a coronal hole (large $\theta_b$) produce faster wind than points near the boundary.
The empirical relation we optimize is:
$$
v(f_s, \theta_b) = v_0 + \frac{v_1}{(1+f_s)^{a}}\left(1 - b,\theta_b\right)^{c}
$$
where $\theta = (v_0, v_1, a, b, c)$ are the five free parameters to be fitted. $v_0$ sets the slow-wind floor, $v_1$ sets the fast-wind amplitude, and $a$, $b$, $c$ shape how quickly speed rises as $f_s$ shrinks and $\theta_b$ grows.
Fitting is framed as a nonlinear least-squares problem: given $N$ paired observations $(f_s^{(i)}, \theta_b^{(i)}, v_{obs}^{(i)})$, we minimize the sum of squared residuals
$$
J(\theta) = \sum_{i=1}^{N}\Big(v\big(f_s^{(i)}, \theta_b^{(i)}; \theta\big) - v_{obs}^{(i)}\Big)^2
$$
Model definition and parameter optimization
1 | import numpy as np |
How the code works
wsa_speed implements the model equation directly. Two np.clip calls guard against invalid math: fs_safe prevents division issues if fs were ever zero or negative, and theta_term prevents raising a negative base to a fractional power c, which would otherwise produce NaN and silently poison the optimizer. This is the single most common source of runtime failure in this kind of fit, so it’s handled defensively from the start.
Synthetic data generation stands in for real spacecraft/coronagraph-derived measurements. We pick a true_params vector, sample fs and theta_b over physically plausible ranges, evaluate the model, and add Gaussian noise to emulate measurement uncertainty. Because we know the ground truth, we can later verify that the optimizer actually recovers it rather than just producing “some” fit.
residuals returns the per-point signed error rather than the squared error. scipy.optimize.least_squares expects a residual vector, not a scalar cost — it uses the Jacobian structure of the residuals (via finite differences here) to take much better steps than a generic scalar minimizer would, which is why it converges quickly even with five free parameters.
Bounds are set generously around physically sensible values. Bounding the search space keeps the Trust Region Reflective (trf) algorithm from wandering into regions where theta_term could hit its clipped floor, which would flatten the gradient and stall convergence.
Why this is already fast: the entire model evaluation is vectorized over all 400 samples with NumPy array operations — there is no per-sample Python loop anywhere in wsa_speed or residuals. Each optimizer iteration therefore costs a handful of NumPy calls rather than 400 Python-level function calls, which is roughly two orders of magnitude faster than a naive loop-based implementation for this problem size.
Console output placeholder — paste the executed cell’s text output below:
=== Solar Wind Model Parameter Optimization === True parameters : [2.80e+02 6.75e+02 1.40e+00 1.05e+00 3.50e-01] Fitted parameters : [2.777381e+02 6.212632e+02 1.304000e+00 1.064400e+00 3.353000e-01] RMSE (km/s) : 11.6626 Sum of squared errors: 54406.9348 Optimizer cost (0.5*SSE): 27203.4674 Number of evaluations : 11
Visualizing the fit and the cost landscape
1 | # ---------- Vectorized cost landscape over (v1, a) ---------- |
How the visualization works
Cost landscape (bottom left, 3D) is built without any Python-level loop over the 80×80 grid. fs_b, theta_b_b, and v_obs_b are reshaped to (400, 1, 1) so that NumPy broadcasting evaluates the model for every one of the 400 observations against every one of the 6,400 grid points in a single vectorized expression, producing a (400, 80, 80) array that is then summed over the sample axis. This computes 2.56 million model evaluations without a single explicit loop, which is what keeps this cell fast in Colab even though it’s exploring a full 2D slice of a 5-parameter space. The white marker shows where the optimizer actually landed, sitting at the base of the bowl-shaped surface.
Fitted model surface (bottom right, 3D) plots the recovered model as a continuous surface over the physical variables $f_s$ and $\theta_b$, with the noisy synthetic observations scattered on top in orange. A good fit means the scatter hugs the surface closely, with scatter visibly rising toward the fast-wind side (small $f_s$, large $\theta_b$).
Top row gives the standard diagnostic pair: the observed-vs-predicted scatter should cluster tightly around the dashed 1:1 line, and the residual histogram should look roughly centered and symmetric around zero, confirming the Gaussian noise assumption was recovered correctly rather than the fit absorbing systematic bias.
Image placeholder — paste the rendered figure below:
Takeaways
The fitted parameters should land close to the true_params vector, with RMSE on the order of the injected noise (~12 km/s), confirming that five-parameter nonlinear least-squares is well-posed for this kind of semi-empirical space weather model given a few hundred coronal hole samples. The cost landscape surface shows a single, well-defined basin around the optimum in the $(v_1, a)$ slice — there’s no sign of a secondary local minimum trapping the trf solver, which is reassuring for using bounded least-squares in an operational WSA-style tuning pipeline rather than needing a global optimizer like differential evolution.
