Graham McIntire
2c16fd979d
- Nx, Axon, EXLA, Polaris deps restricted to only: [:dev, :test] - model.ex and training mix tasks moved to lib_ml/ (compiled via elixirc_paths in dev/test only) - load_ml_model uses Code.ensure_loaded? + apply/3 to avoid compile-time references to ML modules in production - Verified: MIX_ENV=prod compiles clean with no ML warnings
336 lines
10 KiB
Elixir
336 lines
10 KiB
Elixir
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: 20 features (12 atmospheric + 2 solar + 1 geographic + 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 20
|
|
@models_dir Path.join(:code.priv_dir(:microwaveprop), "models")
|
|
@default_path Path.join(@models_dir, "propagation_v1.nx")
|
|
|
|
def default_path, do: @default_path
|
|
|
|
@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(128, activation: :relu, name: "hidden_1")
|
|
|> Axon.dropout(rate: 0.2, name: "dropout_1")
|
|
|> Axon.dense(64, activation: :relu, name: "hidden_2")
|
|
|> Axon.dropout(rate: 0.15, name: "dropout_2")
|
|
|> Axon.dense(32, activation: :relu, name: "hidden_3")
|
|
|> Axon.dropout(rate: 0.1, name: "dropout_3")
|
|
|> 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 """
|
|
Trains the model on the given features and targets.
|
|
|
|
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}}`.
|
|
|
|
## Options
|
|
|
|
* `: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()
|
|
|
|
loop = Axon.Loop.trainer(model, :mean_squared_error, Polaris.Optimizers.adam(learning_rate: lr))
|
|
|
|
data =
|
|
features
|
|
|> Nx.to_batched(batch_size)
|
|
|> 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, init_state,
|
|
epochs: epochs,
|
|
compiler: EXLA
|
|
)
|
|
|
|
# Compute final loss on full dataset
|
|
{_init_fn, predict_fn} = Axon.build(model, compiler: EXLA)
|
|
preds = predict_fn.(trained_state, %{"features" => features})
|
|
final_loss = preds |> Nx.subtract(targets) |> Nx.pow(2) |> Nx.mean() |> Nx.to_number()
|
|
|
|
{trained_state, %{final_loss: final_loss}}
|
|
end
|
|
|
|
@doc """
|
|
Evaluates trained parameters against features and targets.
|
|
|
|
Returns `%{rmse: float, r_squared: float}`.
|
|
"""
|
|
def evaluate(params, features, targets) do
|
|
model = build()
|
|
{_init_fn, predict_fn} = Axon.build(model, compiler: EXLA)
|
|
preds = predict_fn.(params, %{"features" => features})
|
|
|
|
mse = targets |> Nx.subtract(preds) |> Nx.pow(2) |> Nx.mean() |> Nx.to_number()
|
|
rmse = :math.sqrt(mse)
|
|
|
|
ss_res = targets |> Nx.subtract(preds) |> Nx.pow(2) |> Nx.sum() |> Nx.to_number()
|
|
mean_y = targets |> Nx.mean() |> Nx.to_number()
|
|
ss_tot = targets |> Nx.subtract(mean_y) |> Nx.pow(2) |> Nx.sum() |> Nx.to_number()
|
|
r_squared = if ss_tot > 0, do: 1.0 - ss_res / ss_tot, else: 0.0
|
|
|
|
%{rmse: rmse, r_squared: r_squared}
|
|
end
|
|
|
|
@doc """
|
|
Predicts a propagation score (0-100) from raw conditions.
|
|
|
|
Takes a compiled predict function, model parameters, and a conditions map.
|
|
Returns an integer score clamped to [0, 100].
|
|
"""
|
|
def predict_score(predict_fn, params, conditions) do
|
|
features =
|
|
conditions
|
|
|> encode_features()
|
|
|> Nx.tensor(type: :f32)
|
|
|> Nx.reshape({1, @feature_count})
|
|
|
|
params
|
|
|> predict_fn.(%{"features" => features})
|
|
|> Nx.squeeze()
|
|
|> Nx.multiply(100)
|
|
|> Nx.round()
|
|
|> Nx.to_number()
|
|
|> trunc()
|
|
|> max(0)
|
|
|> min(100)
|
|
end
|
|
|
|
@doc """
|
|
Predicts scores for multiple condition maps in a single batched forward pass.
|
|
|
|
Takes a compiled predict function, model parameters, and a list of conditions maps.
|
|
Returns a list of integer scores (0-100).
|
|
"""
|
|
@batch_chunk_size 10_000
|
|
|
|
def predict_scores_batch(predict_fn, params, conditions_list) do
|
|
conditions_list
|
|
|> Enum.chunk_every(@batch_chunk_size)
|
|
|> Enum.flat_map(fn chunk ->
|
|
feature_rows = Enum.map(chunk, &encode_features/1)
|
|
batch = Nx.tensor(feature_rows, type: :f32)
|
|
|
|
params
|
|
|> predict_fn.(%{"features" => batch})
|
|
|> Nx.multiply(100)
|
|
|> Nx.round()
|
|
|> Nx.squeeze(axes: [1])
|
|
|> Nx.to_flat_list()
|
|
|> Enum.map(fn score -> score |> trunc() |> max(0) |> min(100) end)
|
|
end)
|
|
end
|
|
|
|
@doc """
|
|
Compiles the predict function for reuse. Call once at load time.
|
|
"""
|
|
def compile_predict do
|
|
model = build()
|
|
{_init_fn, predict_fn} = Axon.build(model, compiler: EXLA)
|
|
predict_fn
|
|
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,
|
|
:surface_refractivity,
|
|
:latitude,
|
|
:sfi,
|
|
:kp_max,
|
|
:ducting_detected,
|
|
:k_index,
|
|
:lifted_index,
|
|
:solar_hour_sin,
|
|
:solar_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
|
|
refractivity = conditions[:surface_refractivity] || 320.0
|
|
latitude = conditions[:latitude] || 37.0
|
|
sfi = conditions[:sfi] || 120.0
|
|
kp_max = conditions[:kp_max] || 2.0
|
|
ducting = if conditions[:ducting_detected], do: 1.0, else: 0.0
|
|
k_index = conditions[:k_index] || 20.0
|
|
lifted_index = conditions[:lifted_index] || 0.0
|
|
utc_hour = conditions[:utc_hour] || 12
|
|
longitude = conditions[:longitude] || -97.0
|
|
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 (local solar time)
|
|
local_hour = :math.fmod(utc_hour + longitude / 15 + 24, 24)
|
|
hour_rad = 2 * :math.pi() * local_hour / 24
|
|
month_rad = 2 * :math.pi() * (month - 1) / 12
|
|
|
|
# Normalize all features to ~[0, 1] using physical bounds.
|
|
[
|
|
(temp_c + 40) / 80,
|
|
(dewpoint_c + 40) / 80,
|
|
(pressure_mb - 960) / 80,
|
|
abs_humidity / 25.0,
|
|
td_depression / 40.0,
|
|
(grad + 300) / 300,
|
|
hpbl / 3000.0,
|
|
pwat / 60.0,
|
|
(refractivity - 250) / 120,
|
|
(latitude - 25) / 25,
|
|
(sfi - 60) / 200,
|
|
kp_max / 9.0,
|
|
ducting,
|
|
(k_index + 10) / 50,
|
|
(lifted_index + 10) / 20,
|
|
:math.sin(hour_rad),
|
|
:math.cos(hour_rad),
|
|
:math.sin(month_rad),
|
|
:math.cos(month_rad),
|
|
:math.log(freq_mhz) / 13.0
|
|
]
|
|
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
|