From 8949920b7f0912a614ec6d47535cf04540e01cd5 Mon Sep 17 00:00:00 2001 From: Graham McIntire Date: Tue, 31 Mar 2026 16:26:34 -0500 Subject: [PATCH] Add Nx/Axon/EXLA ML model skeleton for propagation prediction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 13-feature feed-forward network (atmospheric + temporal + frequency). Includes build, init, predict, encode_features, save/load to disk. Model weights saved to priv/models/propagation_v1.nx (gitignored). Not yet trained — scaffolding only. --- .gitignore | 3 + config/config.exs | 3 + lib/microwaveprop/propagation/model.ex | 182 ++++++++++++++++++ mix.exs | 5 +- mix.lock | 6 + test/microwaveprop/propagation/model_test.exs | 132 +++++++++++++ 6 files changed, 330 insertions(+), 1 deletion(-) create mode 100644 lib/microwaveprop/propagation/model.ex create mode 100644 test/microwaveprop/propagation/model_test.exs diff --git a/.gitignore b/.gitignore index b96440f8..cb30a2a8 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,9 @@ microwaveprop-*.tar .env .env.* +# Trained ML model weights (binary, large) +/priv/models/*.nx + # GRIB2 test fixtures (large binary files, downloaded on-demand) /test/fixtures/grib2/*.grib2 diff --git a/config/config.exs b/config/config.exs index eea5e9b5..3886eb39 100644 --- a/config/config.exs +++ b/config/config.exs @@ -61,6 +61,9 @@ config :microwaveprop, ecto_repos: [Microwaveprop.Repo], generators: [timestamp_type: :utc_datetime, binary_id: true] +# Use EXLA as default Nx backend for accelerated tensor operations +config :nx, :default_backend, EXLA.Backend + # Use Jason for JSON parsing in Phoenix config :phoenix, :json_library, Jason diff --git a/lib/microwaveprop/propagation/model.ex b/lib/microwaveprop/propagation/model.ex new file mode 100644 index 00000000..a5e834fc --- /dev/null +++ b/lib/microwaveprop/propagation/model.ex @@ -0,0 +1,182 @@ +defmodule Microwaveprop.Propagation.Model do + @moduledoc """ + Neural network model for microwave propagation prediction. + + Uses Axon to define a model that predicts propagation conditions + (score 0-100 per band) from atmospheric and temporal features. + + ## Features (inputs) + + Atmospheric (from HRRR): + - surface_temp_c — surface temperature + - surface_dewpoint_c — surface dewpoint + - surface_pressure_mb — surface pressure + - abs_humidity — absolute humidity (g/m³), derived + - td_depression — temp minus dewpoint (°C) + - min_refractivity_gradient — minimum dN/dh from profile + - hpbl_m — planetary boundary layer height + - pwat_mm — precipitable water + + Temporal: + - utc_hour — hour of day (0-23), sin/cos encoded + - month — month of year (1-12), sin/cos encoded + + Band: + - freq_mhz — operating frequency, log-scaled + + ## Target (output) + + - score — propagation score (0-100), normalized to 0-1 for training + + ## Architecture + + Feed-forward neural network with: + - Input: 13 features (8 atmospheric + 4 cyclical temporal + 1 frequency) + - Hidden: 2 dense layers with ReLU activation + - Output: 1 sigmoid unit (score 0-1, scaled to 0-100) + + The model is intentionally simple to start. We can add complexity + (LSTM for temporal sequences, attention, ensemble) once we have + baseline performance metrics. + """ + + @feature_count 13 + @models_dir Path.join(:code.priv_dir(:microwaveprop), "models") + @default_path Path.join(@models_dir, "propagation_v1.nx") + + @doc """ + Builds the Axon model graph. Does not initialize parameters. + + Returns an `%Axon{}` struct ready for `Axon.build/2` or `Axon.Loop.trainer/3`. + """ + def build do + "features" + |> Axon.input(shape: {nil, @feature_count}) + |> Axon.dense(64, activation: :relu, name: "hidden_1") + |> Axon.dropout(rate: 0.2, name: "dropout_1") + |> Axon.dense(32, activation: :relu, name: "hidden_2") + |> Axon.dropout(rate: 0.1, name: "dropout_2") + |> Axon.dense(1, activation: :sigmoid, name: "output") + end + + @doc """ + Initializes random model parameters. + + Returns `%{...}` map of layer name → tensor params. + """ + def init do + model = build() + {init_fn, _predict_fn} = Axon.build(model) + template = Nx.template({1, @feature_count}, :f32) + init_fn.(template, Axon.ModelState.empty()) + end + + @doc """ + Saves model parameters to disk. Defaults to `priv/models/propagation_v1.nx`. + """ + def save(params, path \\ @default_path) do + File.mkdir_p!(Path.dirname(path)) + binary = Nx.serialize(params) + File.write!(path, binary) + :ok + end + + @doc """ + Loads model parameters from disk. Returns `{:ok, params}` or `:error`. + """ + def load(path \\ @default_path) do + case File.read(path) do + {:ok, binary} -> {:ok, Nx.deserialize(binary)} + {:error, _} -> :error + end + end + + @doc """ + Loads saved parameters if available, otherwise initializes random ones. + """ + def load_or_init(path \\ @default_path) do + case load(path) do + {:ok, params} -> params + :error -> init() + end + end + + @doc """ + Runs a forward pass with the given parameters and input features. + + `features` should be an `{batch_size, 13}` tensor of float32. + Returns an `{batch_size, 1}` tensor of scores in [0, 1]. + """ + def predict(params, features) do + model = build() + {_init_fn, predict_fn} = Axon.build(model) + predict_fn.(params, features) + end + + @doc """ + Returns the ordered list of feature names expected by the model. + """ + def feature_names do + [ + :surface_temp_c, + :surface_dewpoint_c, + :surface_pressure_mb, + :abs_humidity, + :td_depression, + :min_refractivity_gradient, + :hpbl_m, + :pwat_mm, + :utc_hour_sin, + :utc_hour_cos, + :month_sin, + :month_cos, + :log_freq_mhz + ] + end + + @doc """ + Encodes raw condition data into a feature tensor row. + + Takes a map of raw values and returns a flat list of 13 floats + ready to be stacked into a batch tensor. + """ + def encode_features(%{} = conditions) do + temp_c = conditions[:surface_temp_c] || 20.0 + dewpoint_c = conditions[:surface_dewpoint_c] || 10.0 + pressure_mb = conditions[:surface_pressure_mb] || 1013.0 + grad = conditions[:min_refractivity_gradient] || -70.0 + hpbl = conditions[:hpbl_m] || 500.0 + pwat = conditions[:pwat_mm] || 20.0 + utc_hour = conditions[:utc_hour] || 12 + month = conditions[:month] || 6 + freq_mhz = conditions[:freq_mhz] || 10_000 + + abs_humidity = abs_humidity(temp_c, dewpoint_c) + td_depression = temp_c - dewpoint_c + + # Cyclical encoding for time features + hour_rad = 2 * :math.pi() * utc_hour / 24 + month_rad = 2 * :math.pi() * (month - 1) / 12 + + [ + temp_c, + dewpoint_c, + pressure_mb, + abs_humidity, + td_depression, + grad, + hpbl, + pwat, + :math.sin(hour_rad), + :math.cos(hour_rad), + :math.sin(month_rad), + :math.cos(month_rad), + :math.log(freq_mhz) + ] + end + + defp abs_humidity(temp_c, dewpoint_c) do + e_sat = 6.112 * :math.exp(17.67 * dewpoint_c / (dewpoint_c + 243.5)) + 217.0 * e_sat / (temp_c + 273.15) + end +end diff --git a/mix.exs b/mix.exs index 38e31237..8ae881fe 100644 --- a/mix.exs +++ b/mix.exs @@ -64,7 +64,10 @@ defmodule Microwaveprop.MixProject do {:styler, "~> 1.11", only: [:dev, :test], runtime: false}, {:oban, "~> 2.19"}, {:credo, "~> 1.7", only: [:dev, :test], runtime: false}, - {:earmark, "~> 1.4"} + {:earmark, "~> 1.4"}, + {:nx, "~> 0.9"}, + {:axon, "~> 0.7"}, + {:exla, "~> 0.9"} ] end diff --git a/mix.lock b/mix.lock index e83afa72..985021bc 100644 --- a/mix.lock +++ b/mix.lock @@ -1,7 +1,9 @@ %{ + "axon": {:hex, :axon, "0.8.1", "c4a975e62a14ab6c374997b77367ec3a4c2740952ac474d3b0f202c91b7f75c4", [:mix], [{:kino, "~> 0.7", [hex: :kino, repo: "hexpm", optional: true]}, {:kino_vega_lite, "~> 0.1.7", [hex: :kino_vega_lite, repo: "hexpm", optional: true]}, {:nx, "~> 0.10", [hex: :nx, repo: "hexpm", optional: false]}, {:polaris, "~> 0.1", [hex: :polaris, repo: "hexpm", optional: false]}, {:table_rex, "~> 3.1 or ~> 4.1", [hex: :table_rex, repo: "hexpm", optional: true]}], "hexpm", "682a3517489300507ac9345f28341e7fa95bc5b4960d645816074ce551795d37"}, "bandit": {:hex, :bandit, "1.10.4", "02b9734c67c5916a008e7eb7e2ba68aaea6f8177094a5f8d95f1fb99069aac17", [:mix], [{:hpax, "~> 1.0", [hex: :hpax, repo: "hexpm", optional: false]}, {:plug, "~> 1.18", [hex: :plug, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}, {:thousand_island, "~> 1.0", [hex: :thousand_island, repo: "hexpm", optional: false]}, {:websock, "~> 0.5", [hex: :websock, repo: "hexpm", optional: false]}], "hexpm", "a5faf501042ac1f31d736d9d4a813b3db4ef812e634583b6a457b0928798a51d"}, "bunt": {:hex, :bunt, "1.0.0", "081c2c665f086849e6d57900292b3a161727ab40431219529f13c4ddcf3e7a44", [:mix], [], "hexpm", "dc5f86aa08a5f6fa6b8096f0735c4e76d54ae5c9fa2c143e5a1fc7c1cd9bb6b5"}, "cc_precompiler": {:hex, :cc_precompiler, "0.1.11", "8c844d0b9fb98a3edea067f94f616b3f6b29b959b6b3bf25fee94ffe34364768", [:mix], [{:elixir_make, "~> 0.7", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "3427232caf0835f94680e5bcf082408a70b48ad68a5f5c0b02a3bea9f3a075b9"}, + "complex": {:hex, :complex, "0.6.0", "b0130086a7a8c33574d293b2e0e250f4685580418eac52a5658a4bd148f3ccf1", [:mix], [], "hexpm", "0a5fa95580dcaf30fcd60fe1aaf24327c0fe401e98c24d892e172e79498269f9"}, "credo": {:hex, :credo, "1.7.17", "f92b6aa5b26301eaa5a35e4d48ebf5aa1e7094ac00ae38f87086c562caf8a22f", [:mix], [{:bunt, "~> 0.2.1 or ~> 1.0", [hex: :bunt, repo: "hexpm", optional: false]}, {:file_system, "~> 0.2 or ~> 1.0", [hex: :file_system, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "1eb5645c835f0b6c9b5410f94b5a185057bcf6d62a9c2b476da971cde8749645"}, "db_connection": {:hex, :db_connection, "2.9.0", "a6a97c5c958a2d7091a58a9be40caf41ab496b0701d21e1d1abff3fa27a7f371", [:mix], [{:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "17d502eacaf61829db98facf6f20808ed33da6ccf495354a41e64fe42f9c509c"}, "decimal": {:hex, :decimal, "2.3.0", "3ad6255aa77b4a3c4f818171b12d237500e63525c2fd056699967a3e7ea20f62", [:mix], [], "hexpm", "a4d66355cb29cb47c3cf30e71329e58361cfcb37c34235ef3bf1d7bf3773aeac"}, @@ -11,6 +13,7 @@ "ecto_sql": {:hex, :ecto_sql, "3.13.5", "2f8282b2ad97bf0f0d3217ea0a6fff320ead9e2f8770f810141189d182dc304e", [:mix], [{:db_connection, "~> 2.4.1 or ~> 2.5", [hex: :db_connection, repo: "hexpm", optional: false]}, {:ecto, "~> 3.13.0", [hex: :ecto, repo: "hexpm", optional: false]}, {:myxql, "~> 0.7", [hex: :myxql, repo: "hexpm", optional: true]}, {:postgrex, "~> 0.19 or ~> 1.0", [hex: :postgrex, repo: "hexpm", optional: true]}, {:tds, "~> 2.1.1 or ~> 2.2", [hex: :tds, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "aa36751f4e6a2b56ae79efb0e088042e010ff4935fc8684e74c23b1f49e25fdc"}, "elixir_make": {:hex, :elixir_make, "0.9.0", "6484b3cd8c0cee58f09f05ecaf1a140a8c97670671a6a0e7ab4dc326c3109726", [:mix], [], "hexpm", "db23d4fd8b757462ad02f8aa73431a426fe6671c80b200d9710caf3d1dd0ffdb"}, "esbuild": {:hex, :esbuild, "0.10.0", "b0aa3388a1c23e727c5a3e7427c932d89ee791746b0081bbe56103e9ef3d291f", [:mix], [{:jason, "~> 1.4", [hex: :jason, repo: "hexpm", optional: false]}], "hexpm", "468489cda427b974a7cc9f03ace55368a83e1a7be12fba7e30969af78e5f8c70"}, + "exla": {:hex, :exla, "0.11.0", "1428de9edcb297480a64611d3a72fcefe13c93c115bba6d38e910583c37e38c8", [:make, :mix], [{:elixir_make, "~> 0.6", [hex: :elixir_make, repo: "hexpm", optional: false]}, {:fine, "~> 0.1", [hex: :fine, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 1.0", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:nx, "~> 0.11.0", [hex: :nx, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}, {:xla, "~> 0.10.0", [hex: :xla, repo: "hexpm", optional: false]}], "hexpm", "1067207c802bd6f28cded6a2664979ee2e25dddce95cb84be3f0a3ebfbab2c74"}, "expo": {:hex, :expo, "1.1.1", "4202e1d2ca6e2b3b63e02f69cfe0a404f77702b041d02b58597c00992b601db5", [:mix], [], "hexpm", "5fb308b9cb359ae200b7e23d37c76978673aa1b06e2b3075d814ce12c5811640"}, "file_system": {:hex, :file_system, "1.1.1", "31864f4685b0148f25bd3fbef2b1228457c0c89024ad67f7a81a3ffbc0bbad3a", [:mix], [], "hexpm", "7a15ff97dfe526aeefb090a7a9d3d03aa907e100e262a0f8f7746b78f8f87a5d"}, "finch": {:hex, :finch, "0.21.0", "b1c3b2d48af02d0c66d2a9ebfb5622be5c5ecd62937cf79a88a7f98d48a8290c", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:mint, "~> 1.6.2 or ~> 1.7", [hex: :mint, repo: "hexpm", optional: false]}, {:nimble_options, "~> 0.4 or ~> 1.0", [hex: :nimble_options, repo: "hexpm", optional: false]}, {:nimble_pool, "~> 1.1", [hex: :nimble_pool, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "87dc6e169794cb2570f75841a19da99cfde834249568f2a5b121b809588a4377"}, @@ -25,6 +28,7 @@ "mint": {:hex, :mint, "1.7.1", "113fdb2b2f3b59e47c7955971854641c61f378549d73e829e1768de90fc1abf1", [:mix], [{:castore, "~> 0.1.0 or ~> 1.0", [hex: :castore, repo: "hexpm", optional: true]}, {:hpax, "~> 0.1.1 or ~> 0.2.0 or ~> 1.0", [hex: :hpax, repo: "hexpm", optional: false]}], "hexpm", "fceba0a4d0f24301ddee3024ae116df1c3f4bb7a563a731f45fdfeb9d39a231b"}, "nimble_options": {:hex, :nimble_options, "1.1.1", "e3a492d54d85fc3fd7c5baf411d9d2852922f66e69476317787a7b2bb000a61b", [:mix], [], "hexpm", "821b2470ca9442c4b6984882fe9bb0389371b8ddec4d45a9504f00a66f650b44"}, "nimble_pool": {:hex, :nimble_pool, "1.1.0", "bf9c29fbdcba3564a8b800d1eeb5a3c58f36e1e11d7b7fb2e084a643f645f06b", [:mix], [], "hexpm", "af2e4e6b34197db81f7aad230c1118eac993acc0dae6bc83bac0126d4ae0813a"}, + "nx": {:hex, :nx, "0.11.0", "d37723dbd6cfa274a5def6d6664f5680c32e2eb8a1ce25ec6d91751967fa0abf", [:mix], [{:complex, "~> 0.6", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "36157b21239aeb251d6cbac23eb0eb3495a5e1e0cbc2e6df16afd2ede1575205"}, "oban": {:hex, :oban, "2.21.1", "4b6af7b901ef9baca09e239b5a991ef2fa429cf5a13799bc429a131d610ff692", [:mix], [{:ecto_sql, "~> 3.10", [hex: :ecto_sql, repo: "hexpm", optional: false]}, {:ecto_sqlite3, "~> 0.9", [hex: :ecto_sqlite3, repo: "hexpm", optional: true]}, {:igniter, "~> 0.5", [hex: :igniter, repo: "hexpm", optional: true]}, {:jason, "~> 1.1", [hex: :jason, repo: "hexpm", optional: true]}, {:myxql, "~> 0.7", [hex: :myxql, repo: "hexpm", optional: true]}, {:postgrex, "~> 0.20", [hex: :postgrex, repo: "hexpm", optional: true]}, {:telemetry, "~> 1.3", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "8162a160924cf4a25905fed2a9242e7787d88e320e3b5b0dcf324eb17c51c4e6"}, "phoenix": {:hex, :phoenix, "1.8.5", "919db335247e6d4891764dc3063415b0d2457641c5f9b3751b5df03d8e20bbcf", [:mix], [{:bandit, "~> 1.0", [hex: :bandit, repo: "hexpm", optional: true]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:phoenix_pubsub, "~> 2.1", [hex: :phoenix_pubsub, repo: "hexpm", optional: false]}, {:phoenix_template, "~> 1.0", [hex: :phoenix_template, repo: "hexpm", optional: false]}, {:phoenix_view, "~> 2.0", [hex: :phoenix_view, repo: "hexpm", optional: true]}, {:plug, "~> 1.14", [hex: :plug, repo: "hexpm", optional: false]}, {:plug_cowboy, "~> 2.7", [hex: :plug_cowboy, repo: "hexpm", optional: true]}, {:plug_crypto, "~> 1.2 or ~> 2.0", [hex: :plug_crypto, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}, {:websock_adapter, "~> 0.5.3", [hex: :websock_adapter, repo: "hexpm", optional: false]}], "hexpm", "83b2bb125127e02e9f475c8e3e92736325b5b01b0b9b05407bcb4083b7a32485"}, "phoenix_ecto": {:hex, :phoenix_ecto, "4.7.0", "75c4b9dfb3efdc42aec2bd5f8bccd978aca0651dbcbc7a3f362ea5d9d43153c6", [:mix], [{:ecto, "~> 3.5", [hex: :ecto, repo: "hexpm", optional: false]}, {:phoenix_html, "~> 2.14.2 or ~> 3.0 or ~> 4.1", [hex: :phoenix_html, repo: "hexpm", optional: true]}, {:plug, "~> 1.9", [hex: :plug, repo: "hexpm", optional: false]}, {:postgrex, "~> 0.16 or ~> 1.0", [hex: :postgrex, repo: "hexpm", optional: true]}], "hexpm", "1d75011e4254cb4ddf823e81823a9629559a1be93b4321a6a5f11a5306fbf4cc"}, @@ -36,6 +40,7 @@ "phoenix_template": {:hex, :phoenix_template, "1.0.4", "e2092c132f3b5e5b2d49c96695342eb36d0ed514c5b252a77048d5969330d639", [:mix], [{:phoenix_html, "~> 2.14.2 or ~> 3.0 or ~> 4.0", [hex: :phoenix_html, repo: "hexpm", optional: true]}], "hexpm", "2c0c81f0e5c6753faf5cca2f229c9709919aba34fab866d3bc05060c9c444206"}, "plug": {:hex, :plug, "1.19.1", "09bac17ae7a001a68ae393658aa23c7e38782be5c5c00c80be82901262c394c0", [:mix], [{:mime, "~> 1.0 or ~> 2.0", [hex: :mime, repo: "hexpm", optional: false]}, {:plug_crypto, "~> 1.1.1 or ~> 1.2 or ~> 2.0", [hex: :plug_crypto, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.3 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "560a0017a8f6d5d30146916862aaf9300b7280063651dd7e532b8be168511e62"}, "plug_crypto": {:hex, :plug_crypto, "2.1.1", "19bda8184399cb24afa10be734f84a16ea0a2bc65054e23a62bb10f06bc89491", [:mix], [], "hexpm", "6470bce6ffe41c8bd497612ffde1a7e4af67f36a15eea5f921af71cf3e11247c"}, + "polaris": {:hex, :polaris, "0.1.0", "dca61b18e3e801ecdae6ac9f0eca5f19792b44a5cb4b8d63db50fc40fc038d22", [:mix], [{:nx, "~> 0.5", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "13ef2b166650e533cb24b10e2f3b8ab4f2f449ba4d63156e8c569527f206e2c2"}, "postgrex": {:hex, :postgrex, "0.22.0", "fb027b58b6eab1f6de5396a2abcdaaeb168f9ed4eccbb594e6ac393b02078cbd", [:mix], [{:db_connection, "~> 2.9", [hex: :db_connection, repo: "hexpm", optional: false]}, {:decimal, "~> 1.5 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:table, "~> 0.1.0", [hex: :table, repo: "hexpm", optional: true]}], "hexpm", "a68c4261e299597909e03e6f8ff5a13876f5caadaddd0d23af0d0a61afcc5d84"}, "req": {:hex, :req, "0.5.17", "0096ddd5b0ed6f576a03dde4b158a0c727215b15d2795e59e0916c6971066ede", [:mix], [{:brotli, "~> 0.3.1", [hex: :brotli, repo: "hexpm", optional: true]}, {:ezstd, "~> 1.0", [hex: :ezstd, repo: "hexpm", optional: true]}, {:finch, "~> 0.17", [hex: :finch, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: false]}, {:mime, "~> 2.0.6 or ~> 2.1", [hex: :mime, repo: "hexpm", optional: false]}, {:nimble_csv, "~> 1.0", [hex: :nimble_csv, repo: "hexpm", optional: true]}, {:plug, "~> 1.0", [hex: :plug, repo: "hexpm", optional: true]}], "hexpm", "0b8bc6ffdfebbc07968e59d3ff96d52f2202d0536f10fef4dc11dc02a2a43e39"}, "styler": {:hex, :styler, "1.11.0", "35010d970689a23c2bcc8e97bd8bf7d20e3561d60c49be84654df5c37d051a9c", [:mix], [], "hexpm", "70f36165d0cf238a32b7a456fdef6a9c72e77e657d7ac4a0ace33aeba3f2b8c0"}, @@ -48,4 +53,5 @@ "unicode_util_compat": {:hex, :unicode_util_compat, "0.7.1", "a48703a25c170eedadca83b11e88985af08d35f37c6f664d6dcfb106a97782fc", [:rebar3], [], "hexpm", "b3a917854ce3ae233619744ad1e0102e05673136776fb2fa76234f3e03b23642"}, "websock": {:hex, :websock, "0.5.3", "2f69a6ebe810328555b6fe5c831a851f485e303a7c8ce6c5f675abeb20ebdadc", [:mix], [], "hexpm", "6105453d7fac22c712ad66fab1d45abdf049868f253cf719b625151460b8b453"}, "websock_adapter": {:hex, :websock_adapter, "0.5.9", "43dc3ba6d89ef5dec5b1d0a39698436a1e856d000d84bf31a3149862b01a287f", [:mix], [{:bandit, ">= 0.6.0", [hex: :bandit, repo: "hexpm", optional: true]}, {:plug, "~> 1.14", [hex: :plug, repo: "hexpm", optional: false]}, {:plug_cowboy, "~> 2.6", [hex: :plug_cowboy, repo: "hexpm", optional: true]}, {:websock, "~> 0.5", [hex: :websock, repo: "hexpm", optional: false]}], "hexpm", "5534d5c9adad3c18a0f58a9371220d75a803bf0b9a3d87e6fe072faaeed76a08"}, + "xla": {:hex, :xla, "0.10.0", "41121e9f011456242d3a79b9289910ce43419be0b0e7ebe67cc1292c6b3f232f", [:make, :mix], [{:elixir_make, "~> 0.4", [hex: :elixir_make, repo: "hexpm", optional: false]}], "hexpm", "f57d91aea6e661b52bf12239316c598679e9170628122bbd941235f040122bc6"}, } diff --git a/test/microwaveprop/propagation/model_test.exs b/test/microwaveprop/propagation/model_test.exs new file mode 100644 index 00000000..5780cce5 --- /dev/null +++ b/test/microwaveprop/propagation/model_test.exs @@ -0,0 +1,132 @@ +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, "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, 13}) + 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, 13}) + 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, 13 features + input = Nx.broadcast(0.5, {1, 13}) + 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, 13}) + output = Model.predict(params, input) + + assert Nx.shape(output) == {4, 1} + end + end + + describe "encode_features/1" do + test "returns 13-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) == 13 + assert Enum.all?(features, &is_float/1) + end + + test "uses defaults for missing values" do + features = Model.encode_features(%{}) + assert length(features) == 13 + 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), 0.001 + end + end + + describe "feature_names/0" do + test "returns 13 feature names" do + assert length(Model.feature_names()) == 13 + end + end +end