NNExperiments/src/nn/network.nim

# Copyright 2022 Mattia Giambirtone & All Contributors
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import util/matrix


import std/strformat
import std/tables
import std/random


randomize()


type
    NeuralNetwork* = ref object
        ## A generic feed-forward 
        ## neural network
        layers*: seq[Layer]
        activation: Activation                                 # The activation function along with its derivative
        loss: Loss                                             # The cost function along with its derivative
        # This parameter has a different meaning depending on
        # whether we're learning using backpropagation with gradient 
        # descent (in which case it is the amount by which we increase 
        # our input for the next epoch) or using a genetic approach 
        # (where it will be the rate of mutation for each layer)
        learnRate*: float
        # Extra parameters
        params*: TableRef[string, float]
    # Note: The derivatives of the loss and activation
    # function are only meaningful when performing gradient
    # descent!
    Loss* = ref object
        ## A loss function
        function: proc (params: TableRef[string, float]): float
        derivative: proc (x, y: float): float {.noSideEffect.}
    Activation* = ref object
        ## An activation function
        function: proc (input: float): float {.noSideEffect.}
        derivative: proc (x, y: float): float {.noSideEffect.}
    Layer* = ref object
        ## A generic neural network
        ## layer
        inputSize*: int                                     # The number of inputs we process
        outputSize*: int                                    # The number of outputs we produce
        weights*: Matrix[float]                              # The weights for each connection (2D)
        biases*: Matrix[float]                               # The biases for each neuron (1D)
        gradients: tuple[weights, biases: Matrix[float]]    # Gradient coefficients for weights and biases, if using gradient descent


proc `$`*(self: Layer): string =
    ## Returns a string representation
    ## of the layer
    result = &"Layer(inputs={self.inputSize}, outputs={self.outputSize})"


proc `$`*(self: NeuralNetwork): string =
    ## Returns a string representation
    ## of the network
    result = &"NeuralNetwork(learnRate={self.learnRate}, layers={self.layers})"


proc newLoss*(function: proc (params: TableRef[string, float]): float, derivative: proc (x, y: float): float {.noSideEffect.}): Loss =
    ## Creates a new Loss object
    new(result)
    result.function = function
    result.derivative = derivative


proc newActivation*(function: proc (input: float): float {.noSideEffect.}, derivative: proc (x, y: float): float {.noSideEffect.}): Activation =
    ## Creates a new Activation object
    new(result)
    result.function = function
    result.derivative = derivative


proc newLayer*(inputSize: int, outputSize: int, weightRange, biasRange: tuple[start, stop: float]): Layer =
    ## Creates a new layer with inputSize input 
    ## parameters and outputSize outgoing outputs.
    ## Weights are initialized with random values
    ## in the chosen range
    new(result)
    result.inputSize = inputSize
    result.outputSize = outputSize
    var biases = newSeqOfCap[float](outputSize)
    var biasGradients = newSeqOfCap[float](outputSize)
    for _ in 0..<outputSize:
        biases.add(rand(biasRange.start..biasRange.stop))
        biasGradients.add(0.0)
    var weights = newSeqOfCap[seq[float]](inputSize * outputSize)
    var weightGradients = newSeqOfCap[seq[float]](inputSize * outputSize)
    for _ in 0..<outputSize:
        weights.add(@[])
        weightGradients.add(@[])
        for _ in 0..<inputSize:
            weights[^1].add(rand(weightRange.start..weightRange.stop))
            weightGradients[^1].add(0)
    result.biases = newMatrix[float](biases)
    result.weights = newMatrix[float](weights)
    result.gradients = (weights: newMatrix[float](weightGradients), biases: newMatrix[float](biasGradients))

    

proc newNeuralNetwork*(layers: seq[int], activationFunc: Activation, lossFunc: Loss,
                       learnRate: float, weightRange, biasRange: tuple[start, stop: float]): NeuralNetwork =
    ## Initializes a new neural network
    ## with the given layer layout
    new(result)
    result.layers = newSeqOfCap[Layer](len(layers))
    for i in 0..<layers.high():
        result.layers.add(newLayer(layers[i], layers[i + 1], weightRange, biasRange))
    result.activation = activationFunc
    result.loss = lossFunc
    result.learnRate = learnRate
    result.params = newTable[string, float]()


proc compute*(self: NeuralNetwork, data: Matrix[float]): Matrix[float] =
    ## Performs a computation and returns a 1D array
    ## with the output
    when not defined(release):
        if data.shape.rows > 1:
            raise newException(ValueError, "input data must be one-dimensional")
        if data.shape.cols != self.layers[0].inputSize:
            raise newException(ValueError, &"input is of the wrong shape (expecting (1, {self.layers[0].inputSize}), got ({data.shape.rows}, {data.shape.cols}) instead)")
    result = data
    for layer in self.layers:
        result = (layer.weights.dot(result).sum() + layer.biases).apply(self.activation.function, axis= -1)