Graham McIntire
a6b89961f6
Retrained on 60k month-balanced HRRR profiles through 2026-04-05 (test RMSE 1.7 score points, R² 0.97). Added explain_prediction/4 returning ranked feature contributions via batched finite-difference attribution, so the UI can show users which weather factors drove each prediction.
306 lines
9 KiB
Elixir
306 lines
9 KiB
Elixir
defmodule Microwaveprop.Propagation.ModelTest do
|
|
use ExUnit.Case, async: true
|
|
|
|
alias Microwaveprop.Propagation.Model
|
|
|
|
describe "build/0" do
|
|
test "returns an Axon model" do
|
|
model = Model.build()
|
|
assert %Axon{} = model
|
|
end
|
|
end
|
|
|
|
describe "init/0" do
|
|
test "returns initialized model state with parameters" do
|
|
params = Model.init()
|
|
assert %Axon.ModelState{} = params
|
|
assert Map.has_key?(params.data, "hidden_1")
|
|
assert Map.has_key?(params.data, "hidden_2")
|
|
assert Map.has_key?(params.data, "hidden_3")
|
|
assert Map.has_key?(params.data, "output")
|
|
end
|
|
end
|
|
|
|
describe "save/2 and load/1" do
|
|
@tag :tmp_dir
|
|
test "round-trips parameters to disk", %{tmp_dir: tmp_dir} do
|
|
path = Path.join(tmp_dir, "test_model.nx")
|
|
params = Model.init()
|
|
|
|
assert :ok = Model.save(params, path)
|
|
assert File.exists?(path)
|
|
|
|
assert {:ok, loaded} = Model.load(path)
|
|
|
|
# Predict with both and compare
|
|
input = Nx.broadcast(0.5, {1, 20})
|
|
original = Model.predict(params, input)
|
|
reloaded = Model.predict(loaded, input)
|
|
|
|
assert Nx.to_flat_list(original) == Nx.to_flat_list(reloaded)
|
|
end
|
|
|
|
test "load returns :error for missing file" do
|
|
assert :error = Model.load("/tmp/nonexistent_model.nx")
|
|
end
|
|
end
|
|
|
|
describe "load_or_init/1" do
|
|
@tag :tmp_dir
|
|
test "loads from file when it exists", %{tmp_dir: tmp_dir} do
|
|
path = Path.join(tmp_dir, "test_model.nx")
|
|
params = Model.init()
|
|
Model.save(params, path)
|
|
|
|
loaded = Model.load_or_init(path)
|
|
input = Nx.broadcast(0.5, {1, 20})
|
|
assert Nx.to_flat_list(Model.predict(params, input)) == Nx.to_flat_list(Model.predict(loaded, input))
|
|
end
|
|
|
|
test "initializes when file missing" do
|
|
params = Model.load_or_init("/tmp/nonexistent_model.nx")
|
|
assert %Axon.ModelState{} = params
|
|
end
|
|
end
|
|
|
|
describe "predict/2" do
|
|
test "produces output in [0, 1] range" do
|
|
params = Model.init()
|
|
# Single sample, 20 features
|
|
input = Nx.broadcast(0.5, {1, 20})
|
|
output = Model.predict(params, input)
|
|
|
|
assert Nx.shape(output) == {1, 1}
|
|
|
|
value = output |> Nx.to_flat_list() |> hd()
|
|
assert value >= 0.0 and value <= 1.0
|
|
end
|
|
|
|
test "handles batch of multiple samples" do
|
|
params = Model.init()
|
|
input = Nx.broadcast(0.5, {4, 20})
|
|
output = Model.predict(params, input)
|
|
|
|
assert Nx.shape(output) == {4, 1}
|
|
end
|
|
end
|
|
|
|
describe "encode_features/1" do
|
|
test "returns 20-element list" do
|
|
features =
|
|
Model.encode_features(%{
|
|
surface_temp_c: 25.0,
|
|
surface_dewpoint_c: 15.0,
|
|
surface_pressure_mb: 1015.0,
|
|
min_refractivity_gradient: -80.0,
|
|
hpbl_m: 600.0,
|
|
pwat_mm: 25.0,
|
|
utc_hour: 18,
|
|
month: 7,
|
|
freq_mhz: 10_000
|
|
})
|
|
|
|
assert length(features) == 20
|
|
assert Enum.all?(features, &is_float/1)
|
|
end
|
|
|
|
test "uses defaults for missing values" do
|
|
features = Model.encode_features(%{})
|
|
assert length(features) == 20
|
|
assert Enum.all?(features, &is_float/1)
|
|
end
|
|
|
|
test "cyclical encoding: hour 0 and hour 24 produce same sin/cos" do
|
|
f0 = Model.encode_features(%{utc_hour: 0})
|
|
f24 = Model.encode_features(%{utc_hour: 24})
|
|
|
|
# sin/cos at indices 8, 9
|
|
assert_in_delta Enum.at(f0, 8), Enum.at(f24, 8), 0.001
|
|
assert_in_delta Enum.at(f0, 9), Enum.at(f24, 9), 0.001
|
|
end
|
|
|
|
test "log_freq_mhz is last feature" do
|
|
features = Model.encode_features(%{freq_mhz: 10_000})
|
|
assert_in_delta List.last(features), :math.log(10_000) / 13.0, 0.001
|
|
end
|
|
end
|
|
|
|
describe "train/3" 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, 20}, 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)
|
|
|
|
assert %Axon.ModelState{} = trained_state
|
|
assert is_float(metrics.final_loss)
|
|
assert metrics.final_loss >= 0.0
|
|
end
|
|
end
|
|
|
|
describe "evaluate/3" do
|
|
test "returns rmse and r_squared" do
|
|
params = Model.init()
|
|
key = Nx.Random.key(42)
|
|
{features, key} = Nx.Random.uniform(key, shape: {32, 20}, type: :f32)
|
|
{targets, _key} = Nx.Random.uniform(key, shape: {32, 1}, type: :f32)
|
|
|
|
metrics = Model.evaluate(params, features, targets)
|
|
|
|
assert is_float(metrics.rmse)
|
|
assert metrics.rmse >= 0.0
|
|
assert is_float(metrics.r_squared)
|
|
end
|
|
end
|
|
|
|
describe "predict_score/3" do
|
|
test "returns integer between 0 and 100" do
|
|
params = Model.init()
|
|
predict_fn = Model.compile_predict()
|
|
|
|
score =
|
|
Model.predict_score(predict_fn, params, %{
|
|
surface_temp_c: 25.0,
|
|
surface_dewpoint_c: 15.0,
|
|
surface_pressure_mb: 1015.0,
|
|
min_refractivity_gradient: -80.0,
|
|
hpbl_m: 600.0,
|
|
pwat_mm: 25.0,
|
|
utc_hour: 18,
|
|
month: 7,
|
|
freq_mhz: 10_000
|
|
})
|
|
|
|
assert is_integer(score)
|
|
assert score >= 0
|
|
assert score <= 100
|
|
end
|
|
|
|
test "returns different scores for different conditions" do
|
|
params = Model.init()
|
|
predict_fn = Model.compile_predict()
|
|
|
|
score_hot =
|
|
Model.predict_score(predict_fn, params, %{
|
|
surface_temp_c: 35.0,
|
|
surface_dewpoint_c: 25.0,
|
|
surface_pressure_mb: 1010.0,
|
|
min_refractivity_gradient: -300.0,
|
|
hpbl_m: 200.0,
|
|
pwat_mm: 40.0,
|
|
utc_hour: 3,
|
|
month: 7,
|
|
freq_mhz: 10_000
|
|
})
|
|
|
|
score_cold =
|
|
Model.predict_score(predict_fn, params, %{
|
|
surface_temp_c: -10.0,
|
|
surface_dewpoint_c: -20.0,
|
|
surface_pressure_mb: 1030.0,
|
|
min_refractivity_gradient: -40.0,
|
|
hpbl_m: 1500.0,
|
|
pwat_mm: 5.0,
|
|
utc_hour: 14,
|
|
month: 1,
|
|
freq_mhz: 10_000
|
|
})
|
|
|
|
# With random init, scores likely differ (not guaranteed but very probable)
|
|
assert is_integer(score_hot)
|
|
assert is_integer(score_cold)
|
|
end
|
|
end
|
|
|
|
describe "feature_names/0" do
|
|
test "returns 20 feature names" do
|
|
assert length(Model.feature_names()) == 20
|
|
end
|
|
end
|
|
|
|
describe "explain_prediction/3" do
|
|
test "returns score and one contribution entry per feature" do
|
|
params = Model.init()
|
|
predict_fn = Model.compile_predict()
|
|
|
|
result =
|
|
Model.explain_prediction(predict_fn, params, %{
|
|
surface_temp_c: 25.0,
|
|
surface_dewpoint_c: 18.0,
|
|
surface_pressure_mb: 1013.0,
|
|
min_refractivity_gradient: -250.0,
|
|
hpbl_m: 800.0,
|
|
pwat_mm: 30.0,
|
|
surface_refractivity: 340.0,
|
|
latitude: 32.5,
|
|
longitude: -97.0,
|
|
sfi: 130.0,
|
|
kp_max: 2.0,
|
|
ducting_detected: true,
|
|
k_index: 25.0,
|
|
lifted_index: -2.0,
|
|
utc_hour: 23,
|
|
month: 7,
|
|
freq_mhz: 10_000
|
|
})
|
|
|
|
assert is_integer(result.score)
|
|
assert result.score >= 0 and result.score <= 100
|
|
assert length(result.contributions) == 20
|
|
|
|
Enum.each(result.contributions, fn c ->
|
|
assert is_atom(c.feature)
|
|
assert is_float(c.normalized) or is_integer(c.normalized)
|
|
assert is_float(c.sensitivity)
|
|
assert is_float(c.contribution)
|
|
end)
|
|
end
|
|
|
|
test "contributions are sorted by absolute contribution descending" do
|
|
params = Model.init()
|
|
predict_fn = Model.compile_predict()
|
|
|
|
result =
|
|
Model.explain_prediction(predict_fn, params, %{
|
|
surface_temp_c: 15.0,
|
|
surface_dewpoint_c: 5.0,
|
|
freq_mhz: 24_000
|
|
})
|
|
|
|
magnitudes = Enum.map(result.contributions, &abs(&1.contribution))
|
|
assert magnitudes == Enum.sort(magnitudes, :desc)
|
|
end
|
|
|
|
test "every emitted feature name matches feature_names/0" do
|
|
params = Model.init()
|
|
predict_fn = Model.compile_predict()
|
|
|
|
result = Model.explain_prediction(predict_fn, params, %{freq_mhz: 47_000})
|
|
|
|
emitted = result.contributions |> Enum.map(& &1.feature) |> Enum.sort()
|
|
expected = Enum.sort(Model.feature_names())
|
|
assert emitted == expected
|
|
end
|
|
|
|
test "perturbing a feature changes the predicted score" do
|
|
# Verify the sensitivity output is non-trivial: the model must
|
|
# actually move when at least one feature moves. With random init
|
|
# this is essentially guaranteed unless every weight collapses.
|
|
params = Model.init()
|
|
predict_fn = Model.compile_predict()
|
|
|
|
result =
|
|
Model.explain_prediction(predict_fn, params, %{
|
|
surface_temp_c: 20.0,
|
|
surface_dewpoint_c: 10.0,
|
|
surface_pressure_mb: 1013.0,
|
|
freq_mhz: 10_000
|
|
})
|
|
|
|
max_sensitivity = result.contributions |> Enum.map(&abs(&1.sensitivity)) |> Enum.max()
|
|
assert max_sensitivity > 0.0
|
|
end
|
|
end
|
|
end
|