Graham McIntire
c12f8cf5ed
Score time-of-day per grid point using longitude/15 solar offset instead of hardcoded CST/CDT. Add PWAT as 10th scoring factor. Refine pressure thresholds. Update ML model and training pipeline to use local solar time.
275 lines
7.9 KiB
Elixir
275 lines
7.9 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: 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 """
|
|
Trains the model on the given features and targets.
|
|
|
|
Features should be an `{n, 13}` 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)
|
|
"""
|
|
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)
|
|
|
|
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)
|
|
|
|
trained_state =
|
|
Axon.Loop.run(loop, data, Axon.ModelState.empty(),
|
|
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 trained model parameters and a conditions map (same keys as `encode_features/1`).
|
|
Returns an integer score clamped to [0, 100].
|
|
"""
|
|
def predict_score(params, conditions) do
|
|
features =
|
|
conditions
|
|
|> encode_features()
|
|
|> Nx.tensor(type: :f32)
|
|
|> Nx.reshape({1, @feature_count})
|
|
|
|
model = build()
|
|
{_init_fn, predict_fn} = Axon.build(model, compiler: EXLA)
|
|
|
|
params
|
|
|> predict_fn.(%{"features" => features})
|
|
|> Nx.squeeze()
|
|
|> Nx.multiply(100)
|
|
|> Nx.round()
|
|
|> Nx.to_number()
|
|
|> trunc()
|
|
|> max(0)
|
|
|> min(100)
|
|
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
|
|
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
|
|
|
|
[
|
|
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
|