Forecast the Future

Mix.install([
  {:explorer, "~> 0.8.1"},
  {:nx, "~> 0.7.1"},
  {:exla, "~> 0.7.1"},
  {:axon, "~> 0.6.1"},
  {:vega_lite, "~> 0.1.8"},
  {:kino, "~> 0.12.3"},
  {:kino_vega_lite, "~> 0.1.11"}
])

alias VegaLite, as: Vl

Application.put_env(
  :exla,
  :clients,
  cuda: [
    platforms: :cuda,
    lazy_transformers: :never
  ]
)

Nx.global_default_backend(EXLA.Backend)
Nx.Defn.default_options(compiler: EXLA)

Exploring the Data

csv_file = "all_stocks_2006-01-01_to_2018-01-01.csv"

df =
  csv_file
  |> Explorer.DataFrame.from_csv!(parse_dates: true)
  |> Explorer.DataFrame.select(["Date", "Close", "Name"])

Vl.new(title: "DIJA Stock Prices", width: 640, height: 480)
|> Vl.data_from_values(Explorer.DataFrame.to_columns(df))
|> Vl.mark(:line)
|> Vl.encode_field(:x, "Date", type: :temporal)
|> Vl.encode_field(:y, "Close", type: :quantitative)
|> Vl.encode_field(:color, "Name", type: :nominal)
|> Kino.VegaLite.new()

aapl_df =
  Explorer.DataFrame.filter_with(df, fn df ->
    Explorer.Series.equal(df["Name"], "AAPL")
  end)

Vl.new(title: "AAPL Stock Prices", width: 640, height: 480)
|> Vl.data_from_values(Explorer.DataFrame.to_columns(aapl_df))
|> Vl.mark(:line)
|> Vl.encode_field(:x, "Date", type: :temporal)
|> Vl.encode_field(:y, "Close", type: :quantitative)
|> Vl.encode_field(:color, "Name", type: :nominal)
|> Kino.VegaLite.new()

Normalize the Data

normalized_aapl_df =
  Explorer.DataFrame.mutate_with(aapl_df, fn df ->
    var = Explorer.Series.variance(df["Close"])
    mean = Explorer.Series.mean(df["Close"])
    centered = Explorer.Series.subtract(df["Close"], mean)
    norm = Explorer.Series.divide(centered, var)
    [Close: norm]
  end)

Vl.new(title: "AAPL Stock Prices", width: 640, height: 480)
|> Vl.data_from_values(Explorer.DataFrame.to_columns(normalized_aapl_df))
|> Vl.mark(:line)
|> Vl.encode_field(:x, "Date", type: :temporal)
|> Vl.encode_field(:y, "Close", type: :quantitative)
|> Vl.encode_field(:color, "Name", type: :nominal)
|> Kino.VegaLite.new()

defmodule Data do
  def window(inputs, window_size, target_window_size) do
    inputs
    |> Stream.chunk_every(window_size + target_window_size, 1, :discard)
    |> Stream.map(fn window ->
      features =
        window
        |> Enum.take(window_size)
        |> Nx.tensor()
        |> Nx.new_axis(1)

      targets =
        window
        |> Enum.drop(window_size)
        |> Nx.tensor()
        |> Nx.new_axis(1)

      {features, targets}
    end)
  end

  def batch(inputs, batch_size) do
    inputs
    |> Stream.chunk_every(batch_size, batch_size, :discard)
    |> Stream.map(fn windows ->
      {features, targets} = Enum.unzip(windows)
      {Nx.stack(features), Nx.stack(targets)}
    end)
  end
end

train_df =
  Explorer.DataFrame.filter_with(normalized_aapl_df, fn df ->
    Explorer.Series.less(df["Date"], Date.new!(2016, 1, 1))
  end)

test_df =
  Explorer.DataFrame.filter_with(normalized_aapl_df, fn df ->
    Explorer.Series.greater_equal(df["Date"], Date.new!(2016, 1, 1))
  end)

window_size = 5
batch_size = 32

train_prices = Explorer.Series.to_list(train_df["Close"])
test_prices = Explorer.Series.to_list(test_df["Close"])

single_step_train_data =
  train_prices
  |> Data.window(window_size, 1)
  |> Data.batch(batch_size)

single_step_test_data =
  test_prices
  |> Data.window(window_size, 1)
  |> Data.batch(batch_size)

Enum.take(single_step_train_data, 1)

CNN

cnn_model =
  Axon.input("stock_price")
  |> Axon.nx(&Nx.new_axis(&1, -1))
  |> Axon.conv(32, kernel_size: {window_size, 1}, activation: :relu)
  |> Axon.dense(32, activation: :relu)
  |> Axon.dense(1)

template = Nx.template({32, 10, 1}, :f32)
Axon.Display.as_graph(cnn_model, template)

cnn_trained_model_state =
  cnn_model
  |> Axon.Loop.trainer(:mean_squared_error, :adam)
  |> Axon.Loop.metric(:mean_absolute_error)
  |> Axon.Loop.run(single_step_train_data, %{}, epochs: 50, compiler: EXLA)

cnn_model
|> Axon.Loop.evaluator()
|> Axon.Loop.metric(:mean_absolute_error)
|> Axon.Loop.run(single_step_test_data, cnn_trained_model_state, compiler: EXLA)

defmodule Analysis do
  def visualize_predictions(
        model,
        model_state,
        prices,
        window_size,
        target_window_size,
        batch_size
      ) do
    {_, predict_fn} = Axon.build(model, compiler: EXLA)

    windows =
      prices
      |> Data.window(window_size, target_window_size)
      |> Data.batch(batch_size)
      |> Stream.map(&elem(&1, 0))

    predicted =
      Enum.flat_map(windows, fn window ->
        predict_fn.(model_state, window) |> Nx.to_flat_list()
      end)

    predicted = List.duplicate(nil, 10) ++ predicted

    types =
      List.duplicate("AAPL", length(prices)) ++
        List.duplicate("Predicted", length(prices))

    days =
      Enum.to_list(0..(length(prices) - 1)) ++
        Enum.to_list(0..(length(prices) - 1))

    prices = prices ++ predicted

    plot(
      %{
        "days" => days,
        "prices" => prices,
        "types" => types
      },
      "AAPL Stock Price vs. Predicted, CNN Single-Shot"
    )
  end

  defp plot(values, title) do
    Vl.new(title: title, width: 640, height: 480)
    |> Vl.data_from_values(values)
    |> Vl.mark(:line)
    |> Vl.encode_field(:x, "days", type: :temporal)
    |> Vl.encode_field(:y, "prices", type: :quantitative)
    |> Vl.encode_field(:color, "types", type: :nominal)
    |> Kino.VegaLite.new()
  end
end

Analysis.visualize_predictions(
  cnn_model,
  cnn_trained_model_state,
  Explorer.Series.to_list(aapl_df["Close"]),
  window_size,
  1,
  batch_size
)