Graham McIntire
07558d17eb
- 15 features: add surface_refractivity and latitude - Bigger network: 128→64→32 (3 hidden layers) - Phase 1: pretrain on 500K stratified algorithm scores (all seasons/locations) - Phase 2: fine-tune on 57K real QSO-HRRR matched data (percentile target) - Lower LR (0.0003) for fine-tuning to preserve pretrained knowledge - Model.train accepts :initial_state option for transfer learning
367 lines
13 KiB
Elixir
367 lines
13 KiB
Elixir
defmodule Mix.Tasks.PropagationTrain do
|
|
@shortdoc "Train ML model from QSO-HRRR data with algorithm pre-training"
|
|
|
|
@moduledoc """
|
|
Two-phase ML training for propagation prediction:
|
|
|
|
Phase 1 (pre-train): Learn broad atmospheric patterns from 500K+ algorithm
|
|
scores covering all seasons, times, and locations across CONUS.
|
|
|
|
Phase 2 (fine-tune): Calibrate against real QSO outcomes (57K+ contacts
|
|
matched to HRRR conditions). Target is within-band distance percentile.
|
|
|
|
Features: 15 inputs (10 atmospheric + latitude + 4 cyclical temporal + frequency).
|
|
|
|
## Usage
|
|
|
|
mix propagation_train
|
|
mix propagation_train --pretrain-epochs 50 --finetune-epochs 200
|
|
"""
|
|
|
|
use Mix.Task
|
|
|
|
alias Microwaveprop.Propagation.BandConfig
|
|
alias Microwaveprop.Propagation.Model
|
|
alias Microwaveprop.Repo
|
|
|
|
@query_timeout 600_000
|
|
|
|
# Per-band normalization for QSO distance percentile computation
|
|
@band_max_km %{
|
|
10_000 => 800.0,
|
|
24_000 => 350.0,
|
|
47_000 => 200.0,
|
|
75_000 => 200.0,
|
|
122_000 => 140.0,
|
|
134_000 => 160.0,
|
|
241_000 => 115.0
|
|
}
|
|
|
|
@impl Mix.Task
|
|
def run(args) do
|
|
{opts, _, _} =
|
|
OptionParser.parse(args,
|
|
strict: [
|
|
pretrain_epochs: :integer,
|
|
finetune_epochs: :integer,
|
|
pretrain_sample: :integer,
|
|
batch_size: :integer
|
|
]
|
|
)
|
|
|
|
Application.put_env(
|
|
:microwaveprop,
|
|
Oban,
|
|
Keyword.put(
|
|
Application.get_env(:microwaveprop, Oban, []),
|
|
:queues,
|
|
false
|
|
)
|
|
)
|
|
|
|
Mix.Task.run("app.start")
|
|
|
|
pretrain_epochs = Keyword.get(opts, :pretrain_epochs, 50)
|
|
finetune_epochs = Keyword.get(opts, :finetune_epochs, 200)
|
|
pretrain_sample = Keyword.get(opts, :pretrain_sample, 500_000)
|
|
batch_size = Keyword.get(opts, :batch_size, 256)
|
|
|
|
"=" |> String.duplicate(70) |> IO.puts()
|
|
IO.puts("PROPAGATION MODEL TRAINING (Two-Phase)")
|
|
"=" |> String.duplicate(70) |> IO.puts()
|
|
IO.puts("Phase 1: #{pretrain_epochs} epochs on #{pretrain_sample} algorithm scores")
|
|
IO.puts("Phase 2: #{finetune_epochs} epochs on real QSO data")
|
|
IO.puts("Started: #{DateTime.to_string(DateTime.utc_now())}\n")
|
|
|
|
# ── Phase 1: Pre-train on algorithm scores ──────────────────────
|
|
IO.puts("=" <> String.duplicate("─", 50))
|
|
IO.puts("PHASE 1: Pre-training on algorithm scores")
|
|
IO.puts("=" <> String.duplicate("─", 50))
|
|
|
|
{pt_features, pt_targets, pt_counts} = load_pretrain_data(pretrain_sample)
|
|
pt_n = elem(Nx.shape(pt_features), 0)
|
|
IO.puts("Loaded #{pt_n} algorithm score samples:")
|
|
print_band_counts(pt_counts)
|
|
|
|
{pt_features, pt_targets} = shuffle(pt_features, pt_targets, pt_n)
|
|
{pt_train_x, pt_train_y, pt_val_x, pt_val_y, _, _} = split(pt_features, pt_targets, pt_n)
|
|
|
|
IO.puts("\nPre-training for #{pretrain_epochs} epochs (lr=0.001)...")
|
|
|
|
{pretrained_state, pt_metrics} =
|
|
Model.train(pt_train_x, pt_train_y, epochs: pretrain_epochs, batch_size: batch_size)
|
|
|
|
print_loss(pt_metrics.final_loss)
|
|
print_eval("Pre-train val", Model.evaluate(pretrained_state, pt_val_x, pt_val_y))
|
|
|
|
# ── Phase 2: Fine-tune on real QSO data ─────────────────────────
|
|
IO.puts("\n" <> "=" <> String.duplicate("─", 50))
|
|
IO.puts("PHASE 2: Fine-tuning on real QSO-HRRR data")
|
|
IO.puts("=" <> String.duplicate("─", 50))
|
|
|
|
{ft_features, ft_targets, ft_counts} = load_qso_data()
|
|
ft_n = elem(Nx.shape(ft_features), 0)
|
|
IO.puts("Loaded #{ft_n} QSO training samples:")
|
|
print_band_counts(ft_counts)
|
|
|
|
{ft_features, ft_targets} = shuffle(ft_features, ft_targets, ft_n)
|
|
{ft_train_x, ft_train_y, ft_val_x, ft_val_y, ft_test_x, ft_test_y} = split(ft_features, ft_targets, ft_n)
|
|
|
|
IO.puts("\nFine-tuning for #{finetune_epochs} epochs (lr=0.0003)...")
|
|
|
|
{trained_state, ft_metrics} =
|
|
Model.train(ft_train_x, ft_train_y,
|
|
epochs: finetune_epochs,
|
|
batch_size: batch_size,
|
|
learning_rate: 0.0003,
|
|
initial_state: pretrained_state
|
|
)
|
|
|
|
print_loss(ft_metrics.final_loss)
|
|
print_eval("Fine-tune val", Model.evaluate(trained_state, ft_val_x, ft_val_y))
|
|
print_eval("Fine-tune test", Model.evaluate(trained_state, ft_test_x, ft_test_y))
|
|
|
|
# ── Save ────────────────────────────────────────────────────────
|
|
IO.puts("\nSaving model to priv/models/propagation_v1.nx...")
|
|
Model.save(trained_state)
|
|
IO.puts("Done!")
|
|
|
|
"-" |> String.duplicate(70) |> IO.puts()
|
|
IO.puts("Finished: #{DateTime.to_string(DateTime.utc_now())}")
|
|
end
|
|
|
|
# ── Phase 1 data: algorithm scores from propagation_scores ────────
|
|
|
|
defp load_pretrain_data(sample_size) do
|
|
bands = BandConfig.all_freqs()
|
|
per_band = div(sample_size, length(bands))
|
|
|
|
IO.puts(" Loading algorithm scores (stratified #{per_band}/band)...")
|
|
|
|
band_queries =
|
|
Enum.map_join(bands, "\nUNION ALL\n", fn band_mhz ->
|
|
"""
|
|
(SELECT
|
|
ps.score, ps.band_mhz,
|
|
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,
|
|
ps.lat,
|
|
EXTRACT(HOUR FROM ps.valid_time) AS utc_hour,
|
|
EXTRACT(MONTH FROM ps.valid_time) AS month,
|
|
ps.lon
|
|
FROM propagation_scores ps
|
|
JOIN hrrr_profiles h
|
|
ON h.lat = ps.lat AND h.lon = ps.lon AND h.valid_time = ps.valid_time
|
|
WHERE ps.band_mhz = #{band_mhz}
|
|
AND h.surface_temp_c IS NOT NULL
|
|
AND h.surface_dewpoint_c IS NOT NULL
|
|
AND h.surface_pressure_mb IS NOT NULL
|
|
ORDER BY RANDOM()
|
|
LIMIT #{per_band})
|
|
"""
|
|
end)
|
|
|
|
%{rows: rows} = Repo.query!(band_queries, [], timeout: @query_timeout)
|
|
band_counts = Enum.frequencies_by(rows, fn row -> Enum.at(row, 1) end)
|
|
|
|
{feature_rows, target_rows} =
|
|
rows
|
|
|> Enum.map(fn [
|
|
score,
|
|
band_mhz,
|
|
temp_c,
|
|
dewpoint_c,
|
|
pressure_mb,
|
|
grad,
|
|
hpbl,
|
|
pwat,
|
|
refractivity,
|
|
lat,
|
|
utc_hour,
|
|
month,
|
|
lon
|
|
] ->
|
|
features =
|
|
Model.encode_features(%{
|
|
surface_temp_c: to_float(temp_c),
|
|
surface_dewpoint_c: to_float(dewpoint_c),
|
|
surface_pressure_mb: to_float(pressure_mb),
|
|
min_refractivity_gradient: to_float(grad),
|
|
hpbl_m: to_float(hpbl),
|
|
pwat_mm: to_float(pwat),
|
|
surface_refractivity: to_float(refractivity),
|
|
latitude: to_float(lat),
|
|
utc_hour: to_float(utc_hour),
|
|
month: trunc(to_float(month)),
|
|
longitude: to_float(lon),
|
|
freq_mhz: band_mhz
|
|
})
|
|
|
|
{features, [score / 100.0]}
|
|
end)
|
|
|> Enum.unzip()
|
|
|
|
{Nx.tensor(feature_rows, type: :f32), Nx.tensor(target_rows, type: :f32), band_counts}
|
|
end
|
|
|
|
# ── Phase 2 data: real QSO-HRRR matches ──────────────────────────
|
|
|
|
defp load_qso_data do
|
|
IO.puts(" Running QSO-HRRR join query...")
|
|
|
|
sql = """
|
|
SELECT
|
|
q.band::integer AS band_mhz,
|
|
q.distance_km::float AS distance_km,
|
|
EXTRACT(HOUR FROM q.qso_timestamp)::int AS utc_hour,
|
|
EXTRACT(MONTH FROM q.qso_timestamp)::int AS month,
|
|
((q.pos1->>'lng')::float + (q.pos2->>'lng')::float) / 2.0 AS avg_lon,
|
|
((q.pos1->>'lat')::float + (q.pos2->>'lat')::float) / 2.0 AS avg_lat,
|
|
(COALESCE(h1.surface_temp_c, h2.surface_temp_c) +
|
|
COALESCE(h2.surface_temp_c, h1.surface_temp_c)) / 2.0 AS avg_temp,
|
|
(COALESCE(h1.surface_dewpoint_c, h2.surface_dewpoint_c) +
|
|
COALESCE(h2.surface_dewpoint_c, h1.surface_dewpoint_c)) / 2.0 AS avg_dewpoint,
|
|
(COALESCE(h1.surface_pressure_mb, h2.surface_pressure_mb) +
|
|
COALESCE(h2.surface_pressure_mb, h1.surface_pressure_mb)) / 2.0 AS avg_pressure,
|
|
(COALESCE(h1.min_refractivity_gradient, h2.min_refractivity_gradient) +
|
|
COALESCE(h2.min_refractivity_gradient, h1.min_refractivity_gradient)) / 2.0 AS avg_gradient,
|
|
(COALESCE(h1.hpbl_m, h2.hpbl_m) +
|
|
COALESCE(h2.hpbl_m, h1.hpbl_m)) / 2.0 AS avg_hpbl,
|
|
(COALESCE(h1.pwat_mm, h2.pwat_mm) +
|
|
COALESCE(h2.pwat_mm, h1.pwat_mm)) / 2.0 AS avg_pwat,
|
|
(COALESCE(h1.surface_refractivity, h2.surface_refractivity) +
|
|
COALESCE(h2.surface_refractivity, h1.surface_refractivity)) / 2.0 AS avg_refractivity
|
|
FROM qsos q
|
|
INNER JOIN hrrr_profiles h1
|
|
ON h1.lat = ROUND((q.pos1->>'lat')::numeric * 8) / 8
|
|
AND h1.lon = ROUND((q.pos1->>'lng')::numeric * 8) / 8
|
|
AND h1.valid_time = date_trunc('hour', q.qso_timestamp)
|
|
INNER JOIN hrrr_profiles h2
|
|
ON h2.lat = ROUND((q.pos2->>'lat')::numeric * 8) / 8
|
|
AND h2.lon = ROUND((q.pos2->>'lng')::numeric * 8) / 8
|
|
AND h2.valid_time = date_trunc('hour', q.qso_timestamp)
|
|
WHERE q.distance_km > 0
|
|
AND q.distance_km < 3000
|
|
AND q.qso_timestamp >= '2016-06-30'
|
|
AND q.pos1->>'lng' IS NOT NULL
|
|
AND q.pos2->>'lng' IS NOT NULL
|
|
"""
|
|
|
|
%{rows: rows} = Repo.query!(sql, [], timeout: @query_timeout)
|
|
IO.puts(" Matched #{length(rows)} QSOs to HRRR conditions")
|
|
|
|
band_counts = Enum.frequencies_by(rows, fn row -> Enum.at(row, 0) end)
|
|
|
|
# Compute within-band distance percentile
|
|
rows_by_band =
|
|
rows
|
|
|> Enum.filter(fn [band_mhz | _] -> Map.has_key?(@band_max_km, band_mhz) end)
|
|
|> Enum.group_by(fn [band_mhz | _] -> band_mhz end)
|
|
|
|
percentile_rows =
|
|
Enum.flat_map(rows_by_band, fn {_band, band_rows} ->
|
|
distances = band_rows |> Enum.map(fn [_, d | _] -> d end) |> Enum.sort()
|
|
n = length(distances)
|
|
|
|
Enum.map(band_rows, fn [_, distance_km | _] = row ->
|
|
rank = Enum.count(distances, &(&1 <= distance_km))
|
|
{row, rank / n}
|
|
end)
|
|
end)
|
|
|
|
{feature_rows, target_rows} =
|
|
percentile_rows
|
|
|> Enum.map(fn {[
|
|
band_mhz,
|
|
_dist,
|
|
utc_hour,
|
|
month,
|
|
lon,
|
|
lat,
|
|
temp,
|
|
dewpoint,
|
|
pressure,
|
|
grad,
|
|
hpbl,
|
|
pwat,
|
|
refractivity
|
|
], percentile} ->
|
|
features =
|
|
Model.encode_features(%{
|
|
surface_temp_c: to_float(temp),
|
|
surface_dewpoint_c: to_float(dewpoint),
|
|
surface_pressure_mb: to_float(pressure),
|
|
min_refractivity_gradient: to_float(grad),
|
|
hpbl_m: to_float(hpbl),
|
|
pwat_mm: to_float(pwat),
|
|
surface_refractivity: to_float(refractivity),
|
|
latitude: to_float(lat),
|
|
utc_hour: to_float(utc_hour),
|
|
month: trunc(to_float(month)),
|
|
longitude: to_float(lon),
|
|
freq_mhz: band_mhz
|
|
})
|
|
|
|
{features, [percentile]}
|
|
end)
|
|
|> Enum.unzip()
|
|
|
|
IO.puts(" After filtering: #{length(feature_rows)} training samples")
|
|
|
|
{Nx.tensor(feature_rows, type: :f32), Nx.tensor(target_rows, type: :f32), band_counts}
|
|
end
|
|
|
|
# ── Helpers ───────────────────────────────────────────────────────
|
|
|
|
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_band_counts(counts) do
|
|
counts
|
|
|> Enum.sort_by(fn {band, _} -> band end)
|
|
|> Enum.each(fn {band, count} ->
|
|
IO.puts(" #{div(band, 1000)} GHz: #{count}")
|
|
end)
|
|
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
|