{
  "nbformat": 4,
  "nbformat_minor": 2,
  "metadata": {
    "orig_nbformat": 4,
    "language_info": {
      "name": "python",
      "version": "3.8.10",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3.8.10 64-bit ('tf': conda)"
    },
    "interpreter": {
      "hash": "859cadc33cf9c93d9354f622d0bfdc75f7d75fc23a9b22a18dfc646072acf83e"
    },
    "colab": {
      "name": "508_Assignment_1.ipynb",
      "provenance": [],
      "collapsed_sections": []
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# AME508 Assignment 1\n",
        "**Zhilong Li**\n",
        "\n",
        "**USC ID:** 3884-8441-02"
      ],
      "metadata": {
        "id": "MjkK5icEhhqZ"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "source": [
        "from tensorflow import keras\n",
        "import tensorflow as tf\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "import os\n",
        "# from icecream import ic\n",
        "%matplotlib inline\n",
        "os.environ['CUDA_VISIBLE_DEVICES'] = '-1'\n",
        "os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'\n",
        "\n",
        "global_random_seed = np.random.randint(0, 99999) # 3884\n",
        "np.random.seed(global_random_seed)"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Init Plugin\n",
            "Init Graph Optimizer\n",
            "Init Kernel\n"
          ]
        }
      ],
      "metadata": {
        "id": "LOBC5ynYhhqc"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "source": [
        "def get_network(input_dim, output_dim, width, depth,\n",
        "                act_func='tanh', output_act_func=True):\n",
        "    \"\"\"Construct a sequential model based on the input configuration\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    input_dim : int\n",
        "    output_dim : int\n",
        "    width : int\n",
        "    depth : int\n",
        "        Number of hidden layers + output layer\n",
        "    act_func : str, optional\n",
        "        Which activation function this model will be using, by default 'tanh'\n",
        "    output_act_func : bool, optional\n",
        "        Should the output layer use an activation function, by default True\n",
        "\n",
        "    Returns\n",
        "    -------\n",
        "    keras.Sequential\n",
        "        The constructed model based on the input arguments\n",
        "    \"\"\"\n",
        "    # initializer = keras.initializers.RandomUniform(\n",
        "    #     minval=-1., maxval=1., seed=global_random_seed)\n",
        "    initializer = keras.initializers.RandomUniform(\n",
        "        minval=-1., maxval=1.)\n",
        "\n",
        "    if act_func == 'sin':\n",
        "        act_func = tf.math.sin\n",
        "\n",
        "    hidden_layers = [keras.layers.InputLayer(input_shape=[input_dim])]\n",
        "    hidden_layers += [keras.layers.Dense(width, activation=act_func,\n",
        "                                         kernel_initializer=initializer,\n",
        "                                         bias_initializer=initializer) for _ in range(depth-1)]\n",
        "    if output_act_func:\n",
        "        hidden_layers.append(keras.layers.Dense(\n",
        "            output_dim, activation=act_func,\n",
        "            kernel_initializer=initializer,\n",
        "            bias_initializer=initializer))\n",
        "    else:\n",
        "        hidden_layers.append(keras.layers.Dense(\n",
        "            output_dim, kernel_initializer=initializer,\n",
        "            bias_initializer=initializer))\n",
        "    model = keras.Sequential(hidden_layers)\n",
        "    model.build()\n",
        "    return model\n"
      ],
      "outputs": [],
      "metadata": {
        "id": "9ycHIJVTw2Z6",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 134
        },
        "outputId": "d2e8c419-48a1-4311-88ac-255f27b14e7f"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "source": [
        "def get_models_prediction(model_list, input: list):\n",
        "    \"\"\"get the result from all models in the model list\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    model_list : list\n",
        "        List that contains all models need to test\n",
        "    input : list\n",
        "        All input data\n",
        "    \"\"\"\n",
        "    result = []\n",
        "    for model in model_list:\n",
        "        result.append(model.predict(input))\n",
        "    return result\n",
        "\n",
        "def gene_test_data(num_points: int, low_val=-1, high_val=1):\n",
        "    return np.linspace(low_val, high_val, num_points)"
      ],
      "outputs": [],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "source": [
        "def get_result(act_func: str, enable_output_act_func: bool):\n",
        "    all_models = []\n",
        "    width_list = [5, 10, 20, 40]\n",
        "    params_counter = []\n",
        "    for width in width_list:\n",
        "        new_model = get_network(1, 1, width, 15,\n",
        "                                      act_func, enable_output_act_func)\n",
        "        all_models.append(new_model)\n",
        "        params_counter.append(new_model.count_params())\n",
        "    result = get_models_prediction(all_models, test_data)\n",
        "    # get FFT result\n",
        "    result_fft = []\n",
        "    for i in range(len(result)):\n",
        "        result_fft.append(list(np.fft.fft(result[i], axis=0)))\n",
        "    # Plotting\n",
        "    plt.style.use(\"seaborn\")\n",
        "    fig, ((p1, p2, p3, p4), (p5, p6, p7, p8)) = plt.subplots(\n",
        "        2, 4, figsize=(11, 6), dpi=150)\n",
        "    pics = [p1, p2, p3, p4, p5, p6, p7, p8]\n",
        "    for i in range(len(result)):\n",
        "        # Output from the network\n",
        "        pics[i].plot(test_data, list(result[i]), c='orange', linewidth=1)\n",
        "        pics[i].set(adjustable='box', aspect=1.0/pics[i].get_data_ratio())\n",
        "        pics[i].set_title(f\"Width={width_list[i]}\\nParams: {params_counter[i]}\")\n",
        "        pics[i].annotate(f'subPic ({i+1})', xy=(0, 0.95),\n",
        "                         xycoords='axes fraction', c='red')\n",
        "        # Plot the FFT result\n",
        "        pics[i+4].plot(np.fft.fftfreq(data_num), np.abs(np.real(list(result_fft[i]))),\n",
        "                       c='g', linewidth=0.7)\n",
        "        pics[i+4].set(adjustable='box', aspect=1.0/pics[i+4].get_data_ratio())\n",
        "        pics[i+4].set_title(f\"FFT Width={width_list[i]}\")\n",
        "        pics[i+4].annotate(f'subPic ({i+5})', xy=(0, 0.95),\n",
        "                           xycoords='axes fraction', c='red')\n",
        "    fig.suptitle(f\"Result of {act_func}, Output actv function {'ENABLED' if enable_output_act_func else 'DISABLED'}\"\n",
        "                 + f\"               Pic({pic_counter})\")\n",
        "    plt.show()\n"
      ],
      "outputs": [],
      "metadata": {
        "id": "hBEFy4bUqy6G"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "source": [
        "data_num = 10000\n",
        "test_data = gene_test_data(data_num)\n",
        "pic_counter = 0\n",
        "for enable_output_act_func in [True, False]:\n",
        "    for act_func in ['sin', 'tanh']:\n",
        "        pic_counter+=1\n",
        "        get_result(act_func, enable_output_act_func)\n"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Metal device set to: Apple M1\n",
            "\n",
            "systemMemory: 16.00 GB\n",
            "maxCacheSize: 5.33 GB\n",
            "\n"
          ]
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 149668 characters]",
            "text/plain": [
              "<Figure size 1650x900 with 8 Axes>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 145964 characters]",
            "text/plain": [
              "<Figure size 1650x900 with 8 Axes>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 171296 characters]",
            "text/plain": [
              "<Figure size 1650x900 with 8 Axes>"
            ]
          },
          "metadata": {}
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 156244 characters]",
            "text/plain": [
              "<Figure size 1650x900 with 8 Axes>"
            ]
          },
          "metadata": {}
        }
      ],
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 1000
        },
        "id": "CF0P__53qy6H",
        "outputId": "c3d98fa6-6e71-46db-c247-3ee535c2666c"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2-(a) Calculate and print the total number of network parameters\n",
        "\n",
        "The cell below is calculating the number of params by formula, and those results is the **same** as *model.count_params()*."
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "source": [
        "def calc_params_num(input_dim, output_dim, width, depth):\n",
        "    L = depth-1\n",
        "    H = width\n",
        "    d = input_dim\n",
        "    D = output_dim\n",
        "    N = (L-1)*H**2+d*H+D*H+L*H+D\n",
        "    print(\n",
        "        f\"For network that: input dim={input_dim}, output dim={output_dim}, width={width}, depth={depth}\\nthe total params is: {N}\\n\")\n",
        "    return N\n",
        "\n",
        "\n",
        "for w in [5, 10, 20, 40]:\n",
        "    calc_params_num(1, 1, w, 15)\n"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "For network that: input dim=1, output dim=1, width=5, depth=15\n",
            "the total params is: 406\n",
            "\n",
            "For network that: input dim=1, output dim=1, width=10, depth=15\n",
            "the total params is: 1461\n",
            "\n",
            "For network that: input dim=1, output dim=1, width=20, depth=15\n",
            "the total params is: 5521\n",
            "\n",
            "For network that: input dim=1, output dim=1, width=40, depth=15\n",
            "the total params is: 21441\n",
            "\n"
          ]
        }
      ],
      "metadata": {}
    },
    {
      "cell_type": "markdown",
      "source": [
        "## 2-(d) What trends do you observe as a function of width?\n",
        "\n",
        "With a larger width value in the network, looks we can have a more tortuous output curve. \n",
        "\n",
        "And in the FFT plot we can see the frequency decays more slowly, also a wider frequency distribution.\n",
        "\n",
        "## 2-(e) What effect does the choice of activation function appear to have.\n",
        "\n",
        "If the activation function is turned **on** in the output layer.According to Pic(1) & Pic(2), For the same width, the network using *sin* function will have more zigzags than the network using *tanh*. This will be more obvious when comparing *subPic(2)* of *Pic(1)* and *Pic(2)*. But in *subPic(4)* of *Pic(1)* and *Pic(2)*, *sin* function will still be able to remain a wide frequency range while the output of *Tanh* is cutoff.\n",
        "\n",
        "If the activation function is turned **off** in the output layer.According to Pic(3) & Pic(4), for the same width, *sin* function still can have curver output. But compare to *Tanh* it's more periodic. And since there is no activation function at the output layer, *Tanh* network's output is having more zigzags in a wider range (see *subPic(4)* of *Pic(4)* and *Pic(2)*).\n",
        "\n",
        "## 2-(f) How much do these trends vary if you re-run the script? What is the cause of this variation, if any?\n",
        "\n",
        "This trend doesn't change much with multiple runs, a larger width will always has a more tortuous output curve. \n",
        "\n",
        "## 2-(g) What is your conclusion about the expressivity of neural networks as a function of width and/or activation function?\n",
        "\n",
        "I think the expressivity of the neural network is positively correlated with the width value, a large value of width will improve the expressivity of the neural network.\n",
        "\n",
        "The model with *sin* function can have better expressivity than *tanh* function in my experiments. Although the impact of activate function is smaller than the width value (at least in my results above). But I think the activation function is also a very important factor of a model. For different jobs we do need to choose the activation function wisely to achieve a good accuracy.\n",
        "\n",
        "\n",
        "# Some extra work\n",
        "\n",
        "I am running the network for multiple times and get the average value of the fourier transform. It becomes smoother and we can see the frequncy distribution more clearly."
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": 38,
      "source": [
        "repeat_time = 4\n",
        "all_fft_result = []\n",
        "test_data = gene_test_data(100000)\n",
        "for t in range(repeat_time):\n",
        "    print(t)\n",
        "    model = get_network(1,1,20,15,act_func='sin')\n",
        "    result = model.predict(test_data)\n",
        "    fft_result = np.real(np.abs(np.fft.fft(result, axis=0)))\n",
        "    all_fft_result.append(fft_result)\n"
      ],
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "0\n"
          ]
        }
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": 39,
      "source": [
        "fft_avg = np.sum(all_fft_result, axis=0)/repeat_time\n",
        "plt.figure()\n",
        "plt.plot(np.fft.fftfreq(100000), fft_avg)"
      ],
      "outputs": [
        {
          "output_type": "execute_result",
          "data": {
            "text/plain": [
              "[<matplotlib.lines.Line2D at 0x178bf0dc0>]"
            ]
          },
          "metadata": {},
          "execution_count": 39
        },
        {
          "output_type": "display_data",
          "data": {
            "image/png": "[encoded data, 12596 characters]",
            "text/plain": [
              "<Figure size 576x396 with 1 Axes>"
            ]
          },
          "metadata": {}
        }
      ],
      "metadata": {}
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "source": [],
      "outputs": [],
      "metadata": {}
    }
  ]
}