MatrixFlow.h

#ifndef MATRIXFLOW_H
#define MATRIXFLOW_H

#include<vector>
#include<cmath>
#include"Matrix.h"

// Helper functions
double sigmoid(double x) {
    return 1.0 / (1.0 + std::exp(-x));
}

double anti_sigmoid(double x){
    double s = sigmoid(x); 
    return s * (1 - s);
}

Matrix sum_cols(Matrix& m) {
    Matrix res(m.num_of_cols, 1);
    for(int j=0; j<m.num_of_cols; j++){
        double sum = 0.0;
        for(int i=0; i<m.num_of_rows; i++){
            sum += m.at(i, j);
        }
        res.at(j, 0) = sum;
    }
    return res;
}


// Layer class
class Layer {
public:
    // each layer shold have a forward and backward pass
    virtual Matrix forward(Matrix& input) = 0;
    virtual Matrix backward(Matrix error_gradient, double lr) = 0;
    virtual ~Layer() {}
};


// Dense layer
class Dense : public Layer{
public:
    Matrix weights;
    Matrix baises;
    Matrix input;


    Dense(int n_prev, int n_neurons): weights(n_prev, n_neurons), baises(n_neurons, 1) {
        weights.randomize();
        baises.randomize();
    }

    Matrix forward(Matrix& input){ // Z = X * W + b
        this->input = input;
        Matrix Z = input.dot_product(weights);

        for(int i=0; i<Z.num_of_rows; i++){
            for(int j=0; j<Z.num_of_cols; j++){
                Z.at(i, j) += baises.at(j, 0);
            }
        }
        return Z;
    }

    Matrix backward(Matrix error_gradient, double lr){ // dW = X^T * dZ, dB = sum(dZ), dA = dZ * W^T
        Matrix dW = input.transpose().dot_product(error_gradient);
        Matrix dB = sum_cols(error_gradient);
        Matrix dA =  error_gradient.dot_product(weights.transpose()); // this will be passed back

        this->weights = this->weights - (dW * lr);
        this->baises = this->baises - (dB * lr);
        return dA;
    }
};

// Activation layer
class Activation : public Layer{
public:
    double (*activation)(double);
    double (*activation_prime)(double);
    Matrix input;

    Activation(double (*function)(double), double (*function_prime)(double)):
        activation(function), activation_prime(function_prime){};

    Matrix forward(Matrix& input){ // A = activation(Z)
        this->input = input;
        return input.apply(activation);
    }

    Matrix backward(Matrix error_gradeint, double lr){ // dZ = dZ * activation_prime(Z)
        Matrix Slope = input.apply(activation_prime);
        Matrix dZ3 = error_gradeint * Slope;
        return dZ3;
    }
};


// Sigmoid activation layer
class Sigmoid : public Activation {
public:
    Sigmoid(): Activation(
        [](double x) { return 1.0 / (1.0 + std::exp(-x)); },
        [](double x) { double s = 1.0 / (1.0 + std::exp(-x)); return s * (1 - s); }
    ){}
};

// ReLU activation layar
class ReLU : public Activation {
public:
    ReLU(): Activation(
        [](double x) { return x > 0 ? x : 0.0; },
        [](double x) { return x > 0 ? 1.0 : 0.0; }
    ){}
};

// Model class
class Sequential{
public:
    std::vector<Layer*> layers;
    Sequential(){}
    Sequential(std::vector<Layer*> layers): layers(layers){}

    ~Sequential(){
        for(Layer* layer: layers){
            delete layer;
        }
    }

    void add(Layer* layer){
        layers.push_back(layer);
    }

    Matrix forward(Matrix& input){ // forward pass 
        Matrix output = input; 
        for(Layer* layer: this->layers){
            output = layer->forward(output);
        }
        return output;
    }

    void backward(Matrix error_gradient, double lr){ // backward pass
        Matrix grad = error_gradient;
        for (auto it = layers.rbegin(); it != layers.rend(); ++it) {
            grad = (*it)->backward(grad, lr);
        }
    }
};


// trainer class
class Trainer{
public:
    Sequential& model;
    int epochs;
    int batch_size;
    double learning_rate;
    int loss_verbose_epochs;
    std::vector<Matrix> batches_X;
    std::vector<Matrix> batches_Y;

    Trainer(
        Sequential& model,
        Matrix X, Matrix Y,
        int epochs, int batch_size, double lr,
        int loss_verbose_epochs
    ): model(model), epochs(epochs), batch_size(batch_size), learning_rate(lr), loss_verbose_epochs(loss_verbose_epochs){
        this->batches_X = this->generate_batchs(X);
        this->batches_Y = this->generate_batchs(Y);
    }
    
    std::vector<Matrix> generate_batchs(Matrix data){ // generates batches of data
        std::vector<Matrix> batches;
        int minimum_num_batches = data.num_of_rows / batch_size;
        int last_batch_size = data.num_of_rows % batch_size;

        for (int i = 0; i < minimum_num_batches; i++) {
            Matrix batch(batch_size, data.num_of_cols);
            for (int r = 0; r < batch_size; r++) {
                for (int c = 0; c < data.num_of_cols; c++) {
                    batch.at(r, c) = data.at(i * batch_size + r, c);
                }
            }
            batches.push_back(batch);
        }

        if (last_batch_size > 0) {
            Matrix batch(last_batch_size, data.num_of_cols);
            for (int r = 0; r < last_batch_size; r++) {
                for (int c = 0; c < data.num_of_cols; c++) {
                    batch.at(r, c) = data.at(minimum_num_batches * batch_size + r, c);
                }
            }
            batches.push_back(batch);
        }
        return batches;
    }

    void train(){
        for (int e = 0; e < epochs; e++) {
            double epoch_loss = 0;
            for (size_t b = 0; b < batches_X.size(); b++) {
                // Forward pass
                Matrix output = model.forward(batches_X[b]);

                // loss
                Matrix error_gradient = output - batches_Y[b]; // mse

                // Accumulate loss
                for (int i = 0; i < error_gradient.data.size(); i++) {
                    epoch_loss += std::pow(error_gradient.data[i], 2);
                }

                // Backward Pass
                model.backward(error_gradient, learning_rate);
            }

            // loss verbose
            if (epochs > 0 && e % loss_verbose_epochs == 0) { 
                std::cout << "Epoch: " << e << " Loss: " << epoch_loss / loss_verbose_epochs << std::endl;
                epoch_loss=0;
            }
        }
    }
};

#endif