Two-phase training: pretrain on algorithm scores, fine-tune on QSOs

- 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

lib/microwaveprop/propagation/model.ex

@@ -31,7 +31,7 @@ defmodule Microwaveprop.Propagation.Model do
   ## Architecture
 
   Feed-forward neural network with:
-    - Input: 13 features (8 atmospheric + 4 cyclical temporal + 1 frequency)
+    - Input: 15 features (10 atmospheric + 4 cyclical temporal + 1 frequency)
     - Hidden: 2 dense layers with ReLU activation
     - Output: 1 sigmoid unit (score 0-1, scaled to 0-100)
 
@@ -40,7 +40,7 @@ defmodule Microwaveprop.Propagation.Model do
   baseline performance metrics.
   """
 
-  @feature_count 13
+  @feature_count 15
   @models_dir Path.join(:code.priv_dir(:microwaveprop), "models")
   @default_path Path.join(@models_dir, "propagation_v1.nx")
 
@@ -106,7 +106,7 @@ defmodule Microwaveprop.Propagation.Model do
   @doc """
   Trains the model on the given features and targets.
 
-  Features should be an `{n, 13}` tensor and targets an `{n, 1}` tensor
+  Features should be an `{n, #{@feature_count}}` tensor and targets an `{n, 1}` tensor
   with values in [0, 1].
 
   Returns `{trained_state, %{final_loss: float}}`.
@@ -116,11 +116,13 @@ defmodule Microwaveprop.Propagation.Model do
     * `:epochs` - number of training epochs (default: 50)
     * `:batch_size` - batch size (default: 256)
     * `:learning_rate` - Adam learning rate (default: 0.001)
+    * `:initial_state` - pre-trained model state to resume from (default: empty)
   """
   def train(features, targets, opts \\ []) do
     epochs = Keyword.get(opts, :epochs, 50)
     batch_size = Keyword.get(opts, :batch_size, 256)
     lr = Keyword.get(opts, :learning_rate, 0.001)
+    initial_state = Keyword.get(opts, :initial_state)
 
     model = build()
 
@@ -132,8 +134,10 @@ defmodule Microwaveprop.Propagation.Model do
       |> Stream.zip(Nx.to_batched(targets, batch_size))
       |> Stream.map(fn {x, y} -> {%{"features" => x}, y} end)
 
+    init_state = initial_state || Axon.ModelState.empty()
+
     trained_state =
-      Axon.Loop.run(loop, data, Axon.ModelState.empty(),
+      Axon.Loop.run(loop, data, init_state,
         epochs: epochs,
         compiler: EXLA
       )
@@ -219,6 +223,8 @@ defmodule Microwaveprop.Propagation.Model do
       :min_refractivity_gradient,
       :hpbl_m,
       :pwat_mm,
+      :surface_refractivity,
+      :latitude,
       :solar_hour_sin,
       :solar_hour_cos,
       :month_sin,
@@ -240,6 +246,8 @@ defmodule Microwaveprop.Propagation.Model do
     grad = conditions[:min_refractivity_gradient] || -70.0
     hpbl = conditions[:hpbl_m] || 500.0
     pwat = conditions[:pwat_mm] || 20.0
+    refractivity = conditions[:surface_refractivity] || 320.0
+    latitude = conditions[:latitude] || 37.0
     utc_hour = conditions[:utc_hour] || 12
     longitude = conditions[:longitude] || -97.0
     month = conditions[:month] || 6
@@ -264,6 +272,8 @@ defmodule Microwaveprop.Propagation.Model do
       (grad + 300) / 300,
       hpbl / 3000.0,
       pwat / 60.0,
+      (refractivity - 250) / 120,
+      (latitude - 25) / 25,
       :math.sin(hour_rad),
       :math.cos(hour_rad),
       :math.sin(month_rad),

lib/mix/tasks/propagation_train.ex

@@ -1,30 +1,32 @@
 defmodule Mix.Tasks.PropagationTrain do
-  @shortdoc "Train ML model from QSO-HRRR matched data"
+  @shortdoc "Train ML model from QSO-HRRR data with algorithm pre-training"
 
   @moduledoc """
-  Trains the propagation ML model on real QSO outcomes matched to HRRR
+  Two-phase ML training for propagation prediction:
-  atmospheric conditions. The target is QSO distance normalized per-band
-  (0-1), so the model learns what atmospheric conditions actually produce
-  good propagation from historical contacts — not from the hand-tuned algorithm.
 
-  Features: HRRR conditions averaged across both QSO endpoints + solar time + frequency.
+  Phase 1 (pre-train): Learn broad atmospheric patterns from 500K+ algorithm
-  Target: distance_km / band_p99_distance (capped at 1.0).
+  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 --epochs 100 --batch-size 512
+      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: p99 distances from the QSO dataset.
+  # Per-band normalization for QSO distance percentile computation
-  # Distance is divided by these values and capped at 1.0.
   @band_max_km %{
     10_000 => 800.0,
     24_000 => 350.0,
@@ -39,7 +41,12 @@ defmodule Mix.Tasks.PropagationTrain do
   def run(args) do
     {opts, _, _} =
       OptionParser.parse(args,
-        strict: [epochs: :integer, batch_size: :integer]
+        strict: [
+          pretrain_epochs: :integer,
+          finetune_epochs: :integer,
+          pretrain_sample: :integer,
+          batch_size: :integer
+        ]
       )
 
     Application.put_env(
@@ -54,79 +61,68 @@ defmodule Mix.Tasks.PropagationTrain do
 
     Mix.Task.run("app.start")
 
-    epochs = Keyword.get(opts, :epochs, 100)
+    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 (QSO-HRRR)")
+    IO.puts("PROPAGATION MODEL TRAINING (Two-Phase)")
     "=" |> String.duplicate(70) |> IO.puts()
-    IO.puts("Config: epochs=#{epochs}, batch_size=#{batch_size}")
+    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")
 
-    IO.puts("Loading QSO-HRRR training data...")
+    # ── Phase 1: Pre-train on algorithm scores ──────────────────────
-    {features, targets, band_counts} = load_training_data()
+    IO.puts("=" <> String.duplicate("─", 50))
-    n = elem(Nx.shape(features), 0)
+    IO.puts("PHASE 1: Pre-training on algorithm scores")
+    IO.puts("=" <> String.duplicate("─", 50))
 
-    IO.puts("\nLoaded #{n} training samples from real QSOs:")
+    {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)
 
-    band_counts
+    {pt_features, pt_targets} = shuffle(pt_features, pt_targets, pt_n)
-    |> Enum.sort_by(fn {band, _} -> band end)
+    {pt_train_x, pt_train_y, pt_val_x, pt_val_y, _, _} = split(pt_features, pt_targets, pt_n)
-    |> Enum.each(fn {band, count} ->
-      max_km = Map.get(@band_max_km, band)
-      suffix = if max_km, do: " (norm: #{trunc(max_km)} km)", else: " (excluded)"
-      IO.puts("  #{div(band, 1000)} GHz: #{count} QSOs#{suffix}")
-    end)
 
-    IO.puts("")
+    IO.puts("\nPre-training for #{pretrain_epochs} epochs (lr=0.001)...")
 
-    # Shuffle the dataset
+    {pretrained_state, pt_metrics} =
-    IO.puts("Shuffling and splitting dataset (80/10/10)...")
+      Model.train(pt_train_x, pt_train_y, epochs: pretrain_epochs, batch_size: batch_size)
-    key = Nx.Random.key(System.os_time())
-    {indices, _} = Nx.Random.shuffle(key, Nx.iota({n}))
-    indices = Nx.as_type(indices, :s64)
-    features = Nx.take(features, indices)
-    targets = Nx.take(targets, indices)
 
-    # 80/10/10 split
+    print_loss(pt_metrics.final_loss)
-    train_end = trunc(n * 0.8)
+    print_eval("Pre-train val", Model.evaluate(pretrained_state, pt_val_x, pt_val_y))
-    val_end = trunc(n * 0.9)
 
-    train_features = Nx.slice(features, [0, 0], [train_end, 13])
+    # ── Phase 2: Fine-tune on real QSO data ─────────────────────────
-    train_targets = Nx.slice(targets, [0, 0], [train_end, 1])
+    IO.puts("\n" <> "=" <> String.duplicate("─", 50))
-    val_features = Nx.slice(features, [train_end, 0], [val_end - train_end, 13])
+    IO.puts("PHASE 2: Fine-tuning on real QSO-HRRR data")
-    val_targets = Nx.slice(targets, [train_end, 0], [val_end - train_end, 1])
+    IO.puts("=" <> String.duplicate("─", 50))
-    test_features = Nx.slice(features, [val_end, 0], [n - val_end, 13])
-    test_targets = Nx.slice(targets, [val_end, 0], [n - val_end, 1])
 
-    IO.puts("  Train: #{train_end}, Val: #{val_end - train_end}, Test: #{n - val_end}\n")
+    {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)
 
-    # Train
+    {ft_features, ft_targets} = shuffle(ft_features, ft_targets, ft_n)
-    IO.puts("Training for #{epochs} epochs...")
+    {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)
-    {trained_state, train_metrics} = Model.train(train_features, train_targets, epochs: epochs, batch_size: batch_size)
-    final_loss = train_metrics.final_loss
 
-    if is_float(final_loss) and not (final_loss != final_loss) do
+    IO.puts("\nFine-tuning for #{finetune_epochs} epochs (lr=0.0003)...")
-      IO.puts("Training complete. Final loss: #{Float.round(final_loss, 6)}\n")
-    else
-      IO.puts("Training complete. Final loss: NaN (training diverged)\n")
-      IO.puts("Try: --sample 100000 or reduce epochs")
-      System.halt(1)
-    end
 
-    # Evaluate on validation set
+    {trained_state, ft_metrics} =
-    val_metrics = Model.evaluate(trained_state, val_features, val_targets)
+      Model.train(ft_train_x, ft_train_y,
-    IO.puts("Validation metrics:")
+        epochs: finetune_epochs,
-    IO.puts("  RMSE:      #{Float.round(val_metrics.rmse * 100, 2)} points (on 0-100 scale)")
+        batch_size: batch_size,
-    IO.puts("  R-squared: #{Float.round(val_metrics.r_squared, 4)}\n")
+        learning_rate: 0.0003,
+        initial_state: pretrained_state
+      )
 
-    # Evaluate on test set
+    print_loss(ft_metrics.final_loss)
-    test_metrics = Model.evaluate(trained_state, test_features, test_targets)
+    print_eval("Fine-tune val", Model.evaluate(trained_state, ft_val_x, ft_val_y))
-    IO.puts("Test metrics:")
+    print_eval("Fine-tune test", Model.evaluate(trained_state, ft_test_x, ft_test_y))
-    IO.puts("  RMSE:      #{Float.round(test_metrics.rmse * 100, 2)} points (on 0-100 scale)")
-    IO.puts("  R-squared: #{Float.round(test_metrics.r_squared, 4)}\n")
 
-    # Save model
+    # ── Save ────────────────────────────────────────────────────────
-    IO.puts("Saving model to priv/models/propagation_v1.nx...")
+    IO.puts("\nSaving model to priv/models/propagation_v1.nx...")
     Model.save(trained_state)
     IO.puts("Done!")
 
@@ -134,9 +130,86 @@ defmodule Mix.Tasks.PropagationTrain do
     IO.puts("Finished: #{DateTime.to_string(DateTime.utc_now())}")
   end
 
-  defp load_training_data do
+  # ── Phase 1 data: algorithm scores from propagation_scores ────────
-    # Join QSOs to HRRR at both endpoints, average conditions.
+
-    # Same query as the analysis task but returns raw values for encoding.
+  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,
@@ -144,6 +217,7 @@ defmodule Mix.Tasks.PropagationTrain do
       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) +
@@ -155,7 +229,9 @@ defmodule Mix.Tasks.PropagationTrain do
       (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(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
@@ -172,37 +248,45 @@ defmodule Mix.Tasks.PropagationTrain do
       AND q.pos2->>'lng' IS NOT NULL
     """
 
-    IO.puts("  Running QSO-HRRR join query...")
     %{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)
 
-    # Group by band and compute within-band distance percentile as target.
+    # Compute within-band distance percentile
-    # This converts raw distance to "how good were conditions relative to
-    # what's possible on this band" — reduces noise from operator/equipment.
     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_lookup =
+    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 this distance within its band (0.0 to 1.0)
           rank = Enum.count(distances, &(&1 <= distance_km))
-          pct = rank / n
+          {row, rank / n}
-          {row, pct}
         end)
       end)
 
     {feature_rows, target_rows} =
-      percentile_lookup
+      percentile_rows
-      |> Enum.map(fn {[band_mhz, _distance_km, utc_hour, month, lon, temp, dewpoint, pressure, grad, hpbl, pwat],
+      |> Enum.map(fn {[
-                      percentile} ->
+                        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),
@@ -211,6 +295,8 @@ defmodule Mix.Tasks.PropagationTrain do
             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),
@@ -223,10 +309,55 @@ defmodule Mix.Tasks.PropagationTrain do
 
     IO.puts("  After filtering: #{length(feature_rows)} training samples")
 
-    features_tensor = Nx.tensor(feature_rows, type: :f32)
+    {Nx.tensor(feature_rows, type: :f32), Nx.tensor(target_rows, type: :f32), band_counts}
-    targets_tensor = Nx.tensor(target_rows, type: :f32)
+  end
 
-    {features_tensor, targets_tensor, band_counts}
+  # ── 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

test/microwaveprop/propagation/model_test.exs

@@ -33,7 +33,7 @@ defmodule Microwaveprop.Propagation.ModelTest do
       assert {:ok, loaded} = Model.load(path)
 
       # Predict with both and compare
-      input = Nx.broadcast(0.5, {1, 13})
+      input = Nx.broadcast(0.5, {1, 15})
       original = Model.predict(params, input)
       reloaded = Model.predict(loaded, input)
 
@@ -53,7 +53,7 @@ defmodule Microwaveprop.Propagation.ModelTest do
       Model.save(params, path)
 
       loaded = Model.load_or_init(path)
-      input = Nx.broadcast(0.5, {1, 13})
+      input = Nx.broadcast(0.5, {1, 15})
       assert Nx.to_flat_list(Model.predict(params, input)) == Nx.to_flat_list(Model.predict(loaded, input))
     end
 
@@ -66,8 +66,8 @@ defmodule Microwaveprop.Propagation.ModelTest do
   describe "predict/2" do
     test "produces output in [0, 1] range" do
       params = Model.init()
-      # Single sample, 13 features
+      # Single sample, 15 features
-      input = Nx.broadcast(0.5, {1, 13})
+      input = Nx.broadcast(0.5, {1, 15})
       output = Model.predict(params, input)
 
       assert Nx.shape(output) == {1, 1}
@@ -78,7 +78,7 @@ defmodule Microwaveprop.Propagation.ModelTest do
 
     test "handles batch of multiple samples" do
       params = Model.init()
-      input = Nx.broadcast(0.5, {4, 13})
+      input = Nx.broadcast(0.5, {4, 15})
       output = Model.predict(params, input)
 
       assert Nx.shape(output) == {4, 1}
@@ -86,7 +86,7 @@ defmodule Microwaveprop.Propagation.ModelTest do
   end
 
   describe "encode_features/1" do
-    test "returns 13-element list" do
+    test "returns 15-element list" do
       features =
         Model.encode_features(%{
           surface_temp_c: 25.0,
@@ -100,13 +100,13 @@ defmodule Microwaveprop.Propagation.ModelTest do
           freq_mhz: 10_000
         })
 
-      assert length(features) == 13
+      assert length(features) == 15
       assert Enum.all?(features, &is_float/1)
     end
 
     test "uses defaults for missing values" do
       features = Model.encode_features(%{})
-      assert length(features) == 13
+      assert length(features) == 15
       assert Enum.all?(features, &is_float/1)
     end
 
@@ -129,7 +129,7 @@ defmodule Microwaveprop.Propagation.ModelTest do
     test "trains on synthetic data and returns state with metrics" do
       # Generate small synthetic dataset: 64 samples, 13 features
       key = Nx.Random.key(42)
-      {features, key} = Nx.Random.uniform(key, shape: {64, 13}, type: :f32)
+      {features, key} = Nx.Random.uniform(key, shape: {64, 15}, type: :f32)
       {targets, _key} = Nx.Random.uniform(key, shape: {64, 1}, type: :f32)
 
       {trained_state, metrics} = Model.train(features, targets, epochs: 2, batch_size: 32)
@@ -144,7 +144,7 @@ defmodule Microwaveprop.Propagation.ModelTest do
     test "returns rmse and r_squared" do
       params = Model.init()
       key = Nx.Random.key(42)
-      {features, key} = Nx.Random.uniform(key, shape: {32, 13}, type: :f32)
+      {features, key} = Nx.Random.uniform(key, shape: {32, 15}, type: :f32)
       {targets, _key} = Nx.Random.uniform(key, shape: {32, 1}, type: :f32)
 
       metrics = Model.evaluate(params, features, targets)
@@ -213,8 +213,8 @@ defmodule Microwaveprop.Propagation.ModelTest do
   end
 
   describe "feature_names/0" do
-    test "returns 13 feature names" do
+    test "returns 15 feature names" do
-      assert length(Model.feature_names()) == 13
+      assert length(Model.feature_names()) == 15
     end
   end
 end