Predicting Stock Prices with Elixir & Axon

Mix.install(
  [
    {:nx, "~> 0.6.0"},
    {:axon, "~> 0.6.0"},
    {:vega_lite, "~> 0.1.8"},
    {:explorer, "~> 0.7.2"},
    {:kino, "~> 0.12.2"},
    {:kino_vega_lite, "~> 0.1.11"},
    {:scholar, "~> 0.2.1"},
    {:jason, "~> 1.4.0"},
    {:httpoison, "~> 1.8"},
    {:exla, "~> 0.6.4"}
  ],
  config: [
    nx: [default_backend: EXLA.Backend],
    exla: [
      clients: [
        cuda: [platform: :cuda],
        rocm: [platform: :rocm],
        tpu: [platform: :tpu],
        host: [platform: :host]
      ]
    ]
  ]
)

Explanation

The following LiveBook is an Elixir and Axon adaptation of this YouTube video. The only difference is:

• I am looking at all the data rather than just from 1/1/2012 to 1/1/2020

Get the Alpha Vantage API Key

alpha_vantage_api_key_text = Kino.Input.text("Alpha Vantage API Key")

Get the Data

company = "META"
start = ~D[2012-01-01]
end_ = ~D[2020-01-01]

url =
  "https://www.alphavantage.co/query?function=TIME_SERIES_DAILY&symbol=#{company}&outputsize=full&apikey=#{Kino.Input.read(alpha_vantage_api_key_text)}"

%HTTPoison.Response{status_code: 200, body: response_body} = HTTPoison.get!(url)

json = Jason.decode!(response_body)["Time Series (Daily)"]
meta_prices = for tuple <- json, do: [elem(tuple, 0), elem(tuple, 1)["4. close"]]
meta_prices = Enum.sort_by(meta_prices, fn [date, _] -> date end)
meta_prices = Enum.map(meta_prices, fn [date, price] -> [date, String.to_float(price)] end)

Shape the Data

prediction_days_text = Kino.Input.text("Num Prediction Days")

prediction_days = String.to_integer(Kino.Input.read(prediction_days_text))
x = []
y = []

x =
  for i <- prediction_days..(length(meta_prices) - 1) do
    [
      meta_prices |> Enum.at(i) |> Enum.at(0),
      meta_prices
      |> Enum.map(fn lst -> Enum.at(lst, 1) end)
      |> Enum.slice((i - prediction_days)..(i - 1))
    ]
  end

y =
  for i <- prediction_days..(length(meta_prices) - 1) do
    Enum.at(meta_prices, i)
  end

dates = Enum.map(y, fn [date, _price] -> date end)

Set Up the Tensors

mapped_x = Enum.map(x, fn [_date, prices] -> prices end)
mapped_y = Enum.map(y, fn [_date, price] -> price end)

x_tensor = Nx.tensor(mapped_x)
y_tensor = Nx.tensor(mapped_y)

price_min = Explorer.Series.min(Explorer.Series.from_tensor(y_tensor))
price_max = Explorer.Series.max(Explorer.Series.from_tensor(y_tensor))

x_tensor = Scholar.Preprocessing.min_max_scale(x_tensor)
y_tensor = Scholar.Preprocessing.min_max_scale(y_tensor)

shape_1 = elem(Nx.shape(x_tensor), 0)
shape_2 = elem(Nx.shape(x_tensor), 1)
shape_3 = 1

shape = {shape_1, shape_2, shape_3}

x_tensor = Nx.reshape(x_tensor, shape)

Train/Test Split

test_num_text = Kino.Input.text("Number of Test Cases")

test_num = String.to_integer(Kino.Input.read(test_num_text))
[x_train, x_test] = [x_tensor[0..-test_num//1], x_tensor[-test_num..-1]]
[y_train, y_test] = [y_tensor[0..-test_num//1], y_tensor[-test_num..-1]]

Create the Model

input = Axon.input("input", shape: {nil, prediction_days, 1})

model =
  input
  |> Axon.lstm(50, activation: :sigmoid)
  |> then(fn {seq, _} -> seq end)
  |> Axon.dropout(rate: 0.2)
  |> Axon.lstm(50, activation: :sigmoid)
  |> then(fn {seq, _} -> seq end)
  |> Axon.dropout(rate: 0.2)
  |> Axon.lstm(50, activation: :sigmoid)
  |> then(fn {seq, _} -> seq end)
  |> Axon.dropout(rate: 0.2)
  |> Axon.nx(fn x ->
    Nx.gather(
      x,
      Nx.tensor(
        for i <- 0..(elem(Nx.shape(x), 0) - 1) do
          [
            for k <- 0..49 do
              [i, -1, k]
            end
          ]
        end
      )
    )
  end)
  |> Axon.dense(1, activation: :sigmoid)
  |> Axon.nx(fn x -> Nx.squeeze(x) end)

Batch the Data

batched_x_train =
  x_train
  |> Nx.to_batched(32, leftover: :discard)

batched_y_train =
  y_train
  |> Nx.to_batched(32, leftover: :discard)

Train the Model and Run It

params =
  model
  |> Axon.Loop.trainer(:mean_squared_error, :adam)
  |> Axon.Loop.metric(:mean_absolute_error, "MAE")
  |> Axon.Loop.run(Stream.zip(batched_x_train, batched_y_train), %{}, epochs: 25, compiler: EXLA)

Predict the Test Set

y_pred = Axon.predict(model, params, x_test)

Unscale the Predictions & Values

scale = price_max - price_min
y_pred_unscaled = Nx.add(price_min, Nx.multiply(scale, y_pred))
y_test_unscaled = Nx.add(price_min, Nx.multiply(scale, y_test))

data =
  Explorer.DataFrame.new(
    dates: Enum.take(dates, -test_num),
    y_pred: Nx.to_flat_list(y_pred_unscaled),
    y_test: Nx.to_flat_list(y_test_unscaled)
  )

VegaLite.new(width: 400, height: 400, title: "y_pred vs. y_test")
|> VegaLite.data_from_values(data, only: ["dates", "y_test", "y_pred"])
|> VegaLite.layers([
  VegaLite.new()
  |> VegaLite.mark(:line, color: "red")
  |> VegaLite.encode_field(:x, "dates", type: :temporal)
  |> VegaLite.encode_field(:y, "y_test", type: :quantitative),
  VegaLite.new()
  |> VegaLite.mark(:line, color: "green")
  |> VegaLite.encode_field(:x, "dates", type: :temporal)
  |> VegaLite.encode_field(:y, "y_pred", type: :quantitative)
])