Graham McIntire
f85fadecec
Replaces the two-phase (algorithm-scores pretrain + QSO-distance finetune) pipeline with a single unbiased phase that samples hrrr_profiles stratified uniformly by calendar month and trains on the physical algorithm score from Scorer.composite_score/2. Previous Phase 2 target was 'within-band distance percentile' computed from the contacts table. Because non-contest months have almost no QSOs, the model learned 'these conditions → QSO happened → good' and implicitly 'no QSO → bad', so it systematically under-predicted propagation during quiet months (Feb, Mar, Nov). The new pipeline decouples the target from QSO activity entirely. Stratification uses ROW_NUMBER() OVER (PARTITION BY month ORDER BY random()) so every month contributes the same number of profile rows regardless of contest-driven density. Each profile is exploded into one training row per band. Trained model: val RMSE=1.66 pts, R²=0.9744 on 840k rows (10k profiles/month × 7 bands, 50 epochs). Monthly sanity check at fixed weather conditions (10 GHz, 15°C/10°C) now shows Jan=Feb=Mar=68 — the model no longer penalizes quiet months.
394 lines
13 KiB
Elixir
394 lines
13 KiB
Elixir
defmodule Mix.Tasks.PropagationTrain do
|
||
@shortdoc "Train propagation ML model from month-balanced HRRR + algorithm scores"
|
||
|
||
@moduledoc """
|
||
Trains the propagation Nx/Axon model to predict the physical algorithm
|
||
score from weather conditions alone.
|
||
|
||
## Why not use QSO outcomes as the target?
|
||
|
||
An earlier version of this task fine-tuned on "within-band distance
|
||
percentile" from the `contacts` table. That target is structurally
|
||
biased: non-contest months (Feb, Mar, Nov) have almost no QSOs, so the
|
||
model learned "these conditions → QSO happened → good" and implicitly
|
||
"no QSO happened → bad." The result was a model that systematically
|
||
under-predicted propagation during months when nobody is on the air,
|
||
which is exactly the wrong thing for a forecast.
|
||
|
||
This task instead:
|
||
|
||
1. Samples `hrrr_profiles` stratified *uniformly by calendar month*
|
||
— every month contributes the same number of rows regardless of
|
||
contest-driven density. Contest months still exist in the data,
|
||
they just don't drown out the quiet months.
|
||
2. Joins each sample to `solar_indices` (same date) and the nearest
|
||
`soundings` row (±6 h) for K-index / lifted-index features.
|
||
3. Computes the target via `Scorer.composite_score/2` — the *same*
|
||
physical scoring algorithm that `PropagationGridWorker` uses in
|
||
production. Sky / wind / rain fall back to neutral defaults
|
||
because `hrrr_profiles` doesn't store them; the model learns the
|
||
weather-only residual and the missing factors remain constant.
|
||
4. Trains on every band for every sampled profile, producing one
|
||
training row per (profile, band) pair.
|
||
|
||
## Usage
|
||
|
||
mix propagation_train
|
||
mix propagation_train --samples-per-month 20000 --epochs 100
|
||
"""
|
||
|
||
use Mix.Task
|
||
|
||
alias Microwaveprop.Propagation.BandConfig
|
||
alias Microwaveprop.Propagation.Model
|
||
alias Microwaveprop.Propagation.Scorer
|
||
alias Microwaveprop.Repo
|
||
|
||
@query_timeout 600_000
|
||
|
||
@impl Mix.Task
|
||
def run(args) do
|
||
{opts, _, _} =
|
||
OptionParser.parse(args,
|
||
strict: [
|
||
samples_per_month: :integer,
|
||
epochs: :integer,
|
||
batch_size: :integer,
|
||
learning_rate: :float
|
||
]
|
||
)
|
||
|
||
# Boot the app without starting Oban queues so training doesn't have
|
||
# the HRRR / ERA5 / weather cron firing in the background against the
|
||
# local DB we're reading from.
|
||
Application.put_env(
|
||
:microwaveprop,
|
||
Oban,
|
||
Keyword.put(Application.get_env(:microwaveprop, Oban, []), :queues, false)
|
||
)
|
||
|
||
Mix.Task.run("app.start")
|
||
|
||
samples_per_month = Keyword.get(opts, :samples_per_month, 20_000)
|
||
epochs = Keyword.get(opts, :epochs, 100)
|
||
batch_size = Keyword.get(opts, :batch_size, 256)
|
||
learning_rate = Keyword.get(opts, :learning_rate, 0.001)
|
||
|
||
print_header(samples_per_month, epochs, batch_size, learning_rate)
|
||
|
||
{features, targets, month_counts} = load_training_data(samples_per_month)
|
||
|
||
n = elem(Nx.shape(features), 0)
|
||
IO.puts("Loaded #{n} training rows (one per profile × band).")
|
||
print_month_counts(month_counts)
|
||
|
||
{features, targets} = shuffle(features, targets, n)
|
||
{train_x, train_y, val_x, val_y, test_x, test_y} = split(features, targets, n)
|
||
|
||
IO.puts("\nTraining for #{epochs} epochs (lr=#{learning_rate}, batch=#{batch_size})...")
|
||
|
||
{trained_state, metrics} =
|
||
Model.train(train_x, train_y,
|
||
epochs: epochs,
|
||
batch_size: batch_size,
|
||
learning_rate: learning_rate
|
||
)
|
||
|
||
print_loss(metrics.final_loss)
|
||
print_eval("Validation", Model.evaluate(trained_state, val_x, val_y))
|
||
print_eval("Test", Model.evaluate(trained_state, test_x, test_y))
|
||
|
||
IO.puts("\nSaving model to priv/models/propagation_v1.nx...")
|
||
Model.save(trained_state)
|
||
IO.puts("Done.")
|
||
|
||
check_monthly_bias(trained_state)
|
||
|
||
"-" |> String.duplicate(70) |> IO.puts()
|
||
IO.puts("Finished: #{DateTime.to_string(DateTime.utc_now())}")
|
||
end
|
||
|
||
# ── Data loading ────────────────────────────────────────────────────
|
||
|
||
# Returns {features_tensor, targets_tensor, %{month => count}}.
|
||
# Stratifies by calendar month so every month contributes `per_month`
|
||
# profile rows, then explodes each profile into one row per band.
|
||
defp load_training_data(per_month) do
|
||
IO.puts("Sampling hrrr_profiles stratified by month (#{per_month}/month)...")
|
||
|
||
# The stratification is the whole point: random_row_in_month ranks each
|
||
# profile within its (year, month) bucket by a random key, then we keep
|
||
# the first `per_month` rows per month across all years. This is much
|
||
# cheaper than ORDER BY random() over the whole 43M-row table because
|
||
# the window function can stream.
|
||
sql = """
|
||
WITH ranked AS (
|
||
SELECT
|
||
h.valid_time,
|
||
h.lat,
|
||
h.lon,
|
||
h.surface_temp_c,
|
||
h.surface_dewpoint_c,
|
||
h.surface_pressure_mb,
|
||
h.min_refractivity_gradient,
|
||
h.hpbl_m,
|
||
h.pwat_mm,
|
||
h.surface_refractivity,
|
||
h.ducting_detected,
|
||
extract(month from h.valid_time)::int AS month,
|
||
extract(hour from h.valid_time)::int AS utc_hour,
|
||
ROW_NUMBER() OVER (
|
||
PARTITION BY extract(month from h.valid_time)::int
|
||
ORDER BY random()
|
||
) AS rn
|
||
FROM hrrr_profiles h
|
||
WHERE h.surface_temp_c IS NOT NULL
|
||
AND h.surface_dewpoint_c IS NOT NULL
|
||
AND h.surface_pressure_mb IS NOT NULL
|
||
)
|
||
SELECT
|
||
r.valid_time, r.lat, r.lon,
|
||
r.surface_temp_c, r.surface_dewpoint_c, r.surface_pressure_mb,
|
||
r.min_refractivity_gradient, r.hpbl_m, r.pwat_mm,
|
||
r.surface_refractivity, r.ducting_detected,
|
||
r.month, r.utc_hour,
|
||
si.sfi,
|
||
COALESCE((SELECT MAX(v) FROM unnest(si.kp_values) AS v), 2.0) AS kp_max,
|
||
snd.k_index,
|
||
snd.lifted_index
|
||
FROM ranked r
|
||
LEFT JOIN solar_indices si ON si.date = r.valid_time::date
|
||
LEFT JOIN LATERAL (
|
||
SELECT s.k_index, s.lifted_index
|
||
FROM soundings s
|
||
WHERE s.observed_at BETWEEN r.valid_time - interval '6 hours'
|
||
AND r.valid_time + interval '6 hours'
|
||
ORDER BY ABS(EXTRACT(EPOCH FROM s.observed_at - r.valid_time))
|
||
LIMIT 1
|
||
) snd ON true
|
||
WHERE r.rn <= $1
|
||
"""
|
||
|
||
%{rows: rows} = Repo.query!(sql, [per_month], timeout: @query_timeout)
|
||
IO.puts(" #{length(rows)} profile rows after stratified sample")
|
||
|
||
month_counts = Enum.frequencies_by(rows, fn row -> Enum.at(row, 11) end)
|
||
|
||
bands = BandConfig.all_freqs()
|
||
|
||
# Each profile row turns into one training row per band.
|
||
{feature_rows, target_rows} =
|
||
rows
|
||
|> Enum.flat_map(fn row -> encode_profile_rows(row, bands) end)
|
||
|> Enum.unzip()
|
||
|
||
{Nx.tensor(feature_rows, type: :f32), Nx.tensor(target_rows, type: :f32), month_counts}
|
||
end
|
||
|
||
# Builds one training row per (profile, band). Target is the algorithm
|
||
# composite score normalized to [0, 1]; features are the ML model's
|
||
# 20-element encoding of the same conditions.
|
||
defp encode_profile_rows(row, bands) do
|
||
[
|
||
_valid_time,
|
||
lat,
|
||
lon,
|
||
temp_c,
|
||
dewpoint_c,
|
||
pressure_mb,
|
||
grad,
|
||
hpbl_m,
|
||
pwat_mm,
|
||
refractivity,
|
||
ducting,
|
||
month,
|
||
utc_hour,
|
||
sfi,
|
||
kp_max,
|
||
k_index,
|
||
lifted_index
|
||
] = row
|
||
|
||
temp_c = to_float(temp_c)
|
||
dewpoint_c = to_float(dewpoint_c)
|
||
pressure_mb = to_float(pressure_mb)
|
||
grad = to_float(grad)
|
||
hpbl_m = to_float(hpbl_m)
|
||
pwat_mm = to_float(pwat_mm)
|
||
refractivity = to_float(refractivity)
|
||
lat_f = to_float(lat)
|
||
lon_f = to_float(lon)
|
||
sfi_f = to_float(sfi)
|
||
kp_f = to_float(kp_max)
|
||
k_idx_f = to_float(k_index)
|
||
li_f = to_float(lifted_index)
|
||
|
||
scorer_conditions =
|
||
build_scorer_conditions(
|
||
temp_c,
|
||
dewpoint_c,
|
||
pressure_mb,
|
||
grad,
|
||
hpbl_m,
|
||
pwat_mm,
|
||
lat_f,
|
||
lon_f,
|
||
trunc(to_float(month)),
|
||
trunc(to_float(utc_hour))
|
||
)
|
||
|
||
Enum.map(bands, fn freq_mhz ->
|
||
band_config = BandConfig.get(freq_mhz)
|
||
%{score: score} = Scorer.composite_score(scorer_conditions, band_config)
|
||
|
||
features =
|
||
Model.encode_features(%{
|
||
surface_temp_c: temp_c,
|
||
surface_dewpoint_c: dewpoint_c,
|
||
surface_pressure_mb: pressure_mb,
|
||
min_refractivity_gradient: grad,
|
||
hpbl_m: hpbl_m,
|
||
pwat_mm: pwat_mm,
|
||
surface_refractivity: refractivity,
|
||
latitude: lat_f,
|
||
sfi: sfi_f,
|
||
kp_max: kp_f,
|
||
ducting_detected: ducting,
|
||
k_index: k_idx_f,
|
||
lifted_index: li_f,
|
||
utc_hour: to_float(utc_hour),
|
||
longitude: lon_f,
|
||
month: trunc(to_float(month)),
|
||
freq_mhz: freq_mhz
|
||
})
|
||
|
||
{features, [score / 100.0]}
|
||
end)
|
||
end
|
||
|
||
# Builds a conditions map matching the shape Scorer.composite_score/2
|
||
# expects. Sky / wind / rain default to nil → Scorer uses neutral
|
||
# values (50, 50, 100). The ML model's feature set also omits those
|
||
# channels, so the missing-factor contribution is constant across all
|
||
# samples and the model learns the weather-derived residual cleanly.
|
||
defp build_scorer_conditions(temp_c, dewpoint_c, pressure_mb, grad, hpbl_m, pwat_mm, lat, lon, month, utc_hour) do
|
||
%{
|
||
abs_humidity: Scorer.absolute_humidity(temp_c, dewpoint_c),
|
||
temp_f: Scorer.c_to_f(temp_c),
|
||
dewpoint_f: Scorer.c_to_f(dewpoint_c),
|
||
pressure_mb: pressure_mb,
|
||
prev_pressure_mb: nil,
|
||
min_refractivity_gradient: grad,
|
||
bl_depth_m: hpbl_m,
|
||
pwat_mm: pwat_mm,
|
||
sky_cover_pct: nil,
|
||
wind_speed_kts: nil,
|
||
rain_rate_mmhr: 0.0,
|
||
utc_hour: utc_hour,
|
||
utc_minute: 0,
|
||
month: month,
|
||
longitude: lon,
|
||
latitude: lat,
|
||
best_duct_band_ghz: nil
|
||
}
|
||
end
|
||
|
||
# ── Post-training sanity check ──────────────────────────────────────
|
||
|
||
# After training, evaluate the model on a synthetic set of
|
||
# (temp, month) combinations at otherwise-identical conditions. The
|
||
# goal is to confirm the model doesn't prefer / penalize specific
|
||
# months independently of the weather they carry — i.e. that a
|
||
# generic March day doesn't rank below a generic June day just
|
||
# because the training data had fewer March contacts.
|
||
defp check_monthly_bias(params) do
|
||
IO.puts("\nMonthly sanity check at fixed conditions (10 GHz, 15°C, 10°C dewpoint)...")
|
||
predict_fn = Model.compile_predict()
|
||
|
||
for month <- 1..12 do
|
||
conditions = %{
|
||
surface_temp_c: 15.0,
|
||
surface_dewpoint_c: 10.0,
|
||
surface_pressure_mb: 1013.0,
|
||
min_refractivity_gradient: -70.0,
|
||
hpbl_m: 800.0,
|
||
pwat_mm: 20.0,
|
||
surface_refractivity: 320.0,
|
||
latitude: 32.5,
|
||
longitude: -97.0,
|
||
sfi: 120.0,
|
||
kp_max: 2.0,
|
||
ducting_detected: false,
|
||
k_index: 25.0,
|
||
lifted_index: 0.0,
|
||
utc_hour: 12,
|
||
month: month,
|
||
freq_mhz: 10_000
|
||
}
|
||
|
||
score = Model.predict_score(predict_fn, params, conditions)
|
||
IO.puts(" month=#{month |> Integer.to_string() |> String.pad_leading(2)}: #{score}")
|
||
end
|
||
end
|
||
|
||
# ── Helpers ──────────────────────────────────────────────────────────
|
||
|
||
defp print_header(per_month, epochs, batch_size, lr) do
|
||
"=" |> String.duplicate(70) |> IO.puts()
|
||
IO.puts("PROPAGATION MODEL TRAINING — month-balanced HRRR → algorithm score")
|
||
"=" |> String.duplicate(70) |> IO.puts()
|
||
IO.puts("Samples per month: #{per_month} epochs: #{epochs} batch: #{batch_size} lr: #{lr}")
|
||
IO.puts("Started: #{DateTime.to_string(DateTime.utc_now())}\n")
|
||
end
|
||
|
||
defp print_month_counts(counts) do
|
||
counts
|
||
|> Enum.sort_by(fn {m, _} -> m end)
|
||
|> Enum.each(fn {m, count} ->
|
||
IO.puts(" month #{m |> Integer.to_string() |> String.pad_leading(2)}: #{count}")
|
||
end)
|
||
end
|
||
|
||
defp shuffle(features, targets, n) do
|
||
key = Nx.Random.key(System.os_time())
|
||
{indices, _} = Nx.Random.shuffle(key, Nx.iota({n}))
|
||
indices = Nx.as_type(indices, :s64)
|
||
{Nx.take(features, indices), Nx.take(targets, indices)}
|
||
end
|
||
|
||
defp split(features, targets, n) do
|
||
nc = elem(Nx.shape(features), 1)
|
||
train_end = trunc(n * 0.8)
|
||
val_end = trunc(n * 0.9)
|
||
test_size = n - val_end
|
||
|
||
IO.puts(" Split: train=#{train_end}, val=#{val_end - train_end}, test=#{test_size}")
|
||
|
||
{
|
||
Nx.slice(features, [0, 0], [train_end, nc]),
|
||
Nx.slice(targets, [0, 0], [train_end, 1]),
|
||
Nx.slice(features, [train_end, 0], [val_end - train_end, nc]),
|
||
Nx.slice(targets, [train_end, 0], [val_end - train_end, 1]),
|
||
Nx.slice(features, [val_end, 0], [test_size, nc]),
|
||
Nx.slice(targets, [val_end, 0], [test_size, 1])
|
||
}
|
||
end
|
||
|
||
defp print_loss(loss) do
|
||
if is_float(loss) and loss == loss do
|
||
IO.puts(" Final loss: #{Float.round(loss, 6)}")
|
||
else
|
||
IO.puts(" Final loss: NaN (training diverged)")
|
||
System.halt(1)
|
||
end
|
||
end
|
||
|
||
defp print_eval(label, metrics) do
|
||
IO.puts(" #{label}: RMSE=#{Float.round(metrics.rmse * 100, 2)} pts, R²=#{Float.round(metrics.r_squared, 4)}")
|
||
end
|
||
|
||
defp to_float(nil), do: 0.0
|
||
defp to_float(%Decimal{} = d), do: Decimal.to_float(d)
|
||
defp to_float(v) when is_float(v), do: v
|
||
defp to_float(v) when is_integer(v), do: v / 1
|
||
end
|