%% @doc Network compiler for NIF-accelerated evaluation. %% %% This module compiles TWEANN genotypes from Mnesia records into %% a flat representation suitable for the Rust NIF evaluator. %% %% == Compilation Process == %% %% 1. Load genotype records from Mnesia (cortex, neurons, sensors, actuators) %% 2. Build a node graph with connections %% 3. Topologically sort nodes (inputs -> hidden -> outputs) %% 4. Convert to flat indexed representation %% 5. Call tweann_nif:compile_network/3 %% %% == Usage == %% %% Compile a genotype for fast evaluation: %% {ok, Network} = network_compiler:compile(AgentId) %% Outputs = tweann_nif:evaluate(Network, Inputs) %% %% Or compile from in-memory records: %% {ok, Network} = network_compiler:compile_from_records(Cortex, Neurons, Sensors, Actuators) %% %% @copyright 2025 Macula.io -module(network_compiler). -include("records.hrl"). -export([ compile/1, compile_from_records/4, compile_simple/3 ]). %% @doc Compile a genotype from Mnesia for NIF evaluation. %% %% Loads the agent's neural network from Mnesia and compiles it %% to a format suitable for the Rust NIF. %% %% @param AgentId The agent identifier %% @returns {ok, NetworkRef} | {error, Reason} -spec compile(AgentId :: term()) -> {ok, reference()} | {error, {mnesia_error, term()}} | {error, {compilation_failed, term(), [{atom(), atom(), non_neg_integer(), term()}]}}. compile(AgentId) -> case load_genotype(AgentId) of {ok, Cortex, Neurons, Sensors, Actuators} -> compile_from_records(Cortex, Neurons, Sensors, Actuators); {error, Reason} -> {error, Reason} end. %% @doc Compile from in-memory records. %% %% Use this when you already have the network records loaded. %% %% @param Cortex The cortex record %% @param Neurons List of neuron records %% @param Sensors List of sensor records %% @param Actuators List of actuator records %% @returns {ok, NetworkRef} | {error, Reason} -spec compile_from_records( Cortex :: #cortex{}, Neurons :: [#neuron{}], Sensors :: [#sensor{}], Actuators :: [#actuator{}] ) -> {ok, reference()} | {error, term()}. compile_from_records(_Cortex, Neurons, Sensors, Actuators) -> try %% Build ID to index mapping {IdToIndex, InputCount, OutputIndices} = build_index_mapping(Sensors, Neurons, Actuators), %% Build node list in topological order Nodes = build_node_list(Sensors, Neurons, Actuators, IdToIndex), %% Compile via NIF Network = tweann_nif:compile_network(Nodes, InputCount, OutputIndices), {ok, Network} catch error:Reason:Stack -> {error, {compilation_failed, Reason, Stack}} end. %% @doc Compile a simple feedforward network. %% %% Convenience function for creating simple networks without Mnesia. %% Useful for testing and examples. %% %% @param InputCount Number of input nodes %% @param HiddenLayers List of hidden layer sizes, e.g., [4, 3] for 2 layers %% @param OutputCount Number of output nodes %% @returns {ok, NetworkRef, Weights} where Weights can be used to set weights -spec compile_simple( InputCount :: pos_integer(), HiddenLayers :: [pos_integer()], OutputCount :: pos_integer() ) -> {ok, reference(), [{non_neg_integer(), non_neg_integer(), float()}]}. compile_simple(InputCount, HiddenLayers, OutputCount) -> %% Generate nodes and random weights {Nodes, Weights, TotalNodes} = generate_simple_network(InputCount, HiddenLayers, OutputCount), %% Output indices are the last OutputCount nodes OutputIndices = lists:seq(TotalNodes - OutputCount, TotalNodes - 1), %% Compile Network = tweann_nif:compile_network(Nodes, InputCount, OutputIndices), {ok, Network, Weights}. %%============================================================================== %% Internal Functions %%============================================================================== %% @private Load genotype from Mnesia load_genotype(AgentId) -> case mnesia:transaction(fun() -> [Agent] = mnesia:read({agent, AgentId}), CxId = Agent#agent.cx_id, [Cortex] = mnesia:read({cortex, CxId}), Neurons = [N || NId <- Cortex#cortex.neuron_ids, [N] <- [mnesia:read({neuron, NId})]], Sensors = [S || SId <- Cortex#cortex.sensor_ids, [S] <- [mnesia:read({sensor, SId})]], Actuators = [A || AId <- Cortex#cortex.actuator_ids, [A] <- [mnesia:read({actuator, AId})]], {Cortex, Neurons, Sensors, Actuators} end) of {atomic, {Cortex, Neurons, Sensors, Actuators}} -> {ok, Cortex, Neurons, Sensors, Actuators}; {aborted, Reason} -> {error, {mnesia_error, Reason}} end. %% @private Build mapping from record IDs to flat indices build_index_mapping(Sensors, Neurons, Actuators) -> %% Sensors become inputs (index 0..N-1) SensorIds = [S#sensor.id || S <- Sensors], InputCount = length(SensorIds), %% Sort neurons by layer coordinate for topological order SortedNeurons = lists:sort( fun(N1, N2) -> {{L1, _}, _} = N1#neuron.id, {{L2, _}, _} = N2#neuron.id, L1 =< L2 end, Neurons ), NeuronIds = [N#neuron.id || N <- SortedNeurons], %% Actuators become outputs (last indices) ActuatorIds = [A#actuator.id || A <- Actuators], %% Build ID -> Index map AllIds = SensorIds ++ NeuronIds ++ ActuatorIds, IdToIndex = maps:from_list( [{Id, Idx} || {Id, Idx} <- lists:zip(AllIds, lists:seq(0, length(AllIds) - 1))] ), %% Output indices OutputIndices = [maps:get(Id, IdToIndex) || Id <- ActuatorIds], {IdToIndex, InputCount, OutputIndices}. %% @private Build node list in format expected by NIF build_node_list(Sensors, Neurons, Actuators, IdToIndex) -> %% Sensors (inputs) - no connections, linear activation SensorNodes = [ {maps:get(S#sensor.id, IdToIndex), input, linear, 0.0, []} || S <- Sensors ], %% Sort neurons by layer for topological order SortedNeurons = lists:sort( fun(N1, N2) -> {{L1, _}, _} = N1#neuron.id, {{L2, _}, _} = N2#neuron.id, L1 =< L2 end, Neurons ), %% Neurons (hidden) NeuronNodes = [ begin Idx = maps:get(N#neuron.id, IdToIndex), Activation = N#neuron.af, %% Extract bias from weights (first weight with bias marker or 0.0) Bias = extract_bias(N#neuron.input_idps), %% Convert input connections Connections = convert_connections(N#neuron.input_idps, IdToIndex), {Idx, hidden, Activation, Bias, Connections} end || N <- SortedNeurons ], %% Actuators (outputs) - receive from neurons, linear activation ActuatorNodes = [ begin Idx = maps:get(A#actuator.id, IdToIndex), %% Actuators receive from their fanin neurons with weight 1.0 Connections = [ {maps:get(FromId, IdToIndex), 1.0} || FromId <- A#actuator.fanin_ids, maps:is_key(FromId, IdToIndex) ], {Idx, output, linear, 0.0, Connections} end || A <- Actuators ], SensorNodes ++ NeuronNodes ++ ActuatorNodes. %% @private Extract bias from weighted inputs extract_bias(InputIdps) -> %% Look for bias input (special marker) or return 0.0 case lists:keyfind(bias, 1, InputIdps) of {bias, [{W, _, _, _} | _]} -> W; {bias, W} when is_number(W) -> W; false -> 0.0 end. %% @private Convert input_idps to connection list convert_connections(InputIdps, IdToIndex) -> lists:flatmap( fun ({bias, _}) -> %% Skip bias - handled separately []; ({FromId, WeightList}) when is_list(WeightList) -> case maps:find(FromId, IdToIndex) of {ok, FromIdx} -> %% Sum weights if multiple (shouldn't happen normally) TotalWeight = lists:sum([W || {W, _, _, _} <- WeightList]), [{FromIdx, TotalWeight}]; error -> [] end; ({FromId, Weight}) when is_number(Weight) -> case maps:find(FromId, IdToIndex) of {ok, FromIdx} -> [{FromIdx, Weight}]; error -> [] end end, InputIdps ). %% @private Generate a simple feedforward network generate_simple_network(InputCount, HiddenLayers, OutputCount) -> %% Calculate total nodes HiddenTotal = lists:sum(HiddenLayers), TotalNodes = InputCount + HiddenTotal + OutputCount, %% Generate input nodes InputNodes = [ {I, input, linear, 0.0, []} || I <- lists:seq(0, InputCount - 1) ], %% Generate hidden layers with random weights {HiddenNodes, Weights1, NextIdx1, PrevLayerStart, PrevLayerEnd} = generate_hidden_layers(InputCount, HiddenLayers), %% Generate output layer - connects to last hidden layer, or inputs if no hidden {OutputNodes, Weights2} = generate_output_layer(NextIdx1, OutputCount, PrevLayerStart, PrevLayerEnd), AllNodes = InputNodes ++ HiddenNodes ++ OutputNodes, AllWeights = Weights1 ++ Weights2, {AllNodes, AllWeights, TotalNodes}. %% @private Generate hidden layers generate_hidden_layers(InputCount, []) -> %% No hidden layers - outputs will connect directly to inputs {[], [], InputCount, 0, InputCount - 1}; generate_hidden_layers(InputCount, HiddenLayers) -> generate_hidden_layers_acc(HiddenLayers, 0, InputCount - 1, InputCount, [], []). generate_hidden_layers_acc([], PrevStart, PrevEnd, NextIdx, NodesAcc, WeightsAcc) -> {lists:reverse(NodesAcc), lists:flatten(lists:reverse(WeightsAcc)), NextIdx, PrevStart, PrevEnd}; generate_hidden_layers_acc([LayerSize | Rest], PrevStart, PrevEnd, StartIdx, NodesAcc, WeightsAcc) -> %% Generate nodes for this layer {LayerNodes, LayerWeights} = lists:unzip([ begin Idx = StartIdx + I, %% Connect to all nodes in previous layer Connections = [ {PrevIdx, random_weight()} || PrevIdx <- lists:seq(PrevStart, PrevEnd) ], Bias = random_weight() * 0.1, Node = {Idx, hidden, tanh, Bias, Connections}, Weights = [{Idx, PrevIdx, W} || {PrevIdx, W} <- Connections], {Node, Weights} end || I <- lists:seq(0, LayerSize - 1) ]), NewPrevStart = StartIdx, NewPrevEnd = StartIdx + LayerSize - 1, generate_hidden_layers_acc(Rest, NewPrevStart, NewPrevEnd, StartIdx + LayerSize, LayerNodes ++ NodesAcc, [LayerWeights | WeightsAcc]). %% @private Generate output layer generate_output_layer(StartIdx, OutputCount, PrevStart, PrevEnd) -> {Nodes, NestedWeights} = lists:unzip([ begin Idx = StartIdx + I, Connections = [ {PrevIdx, random_weight()} || PrevIdx <- lists:seq(PrevStart, PrevEnd) ], Node = {Idx, output, tanh, 0.0, Connections}, Weights = [{Idx, PrevIdx, W} || {PrevIdx, W} <- Connections], {Node, Weights} end || I <- lists:seq(0, OutputCount - 1) ]), {Nodes, lists:flatten(NestedWeights)}. %% @private Generate random weight in [-1, 1] random_weight() -> rand:uniform() * 2 - 1.