# Source code for gluonts.nursery.auto_ode.auto_ode

# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
#
# or in the "license" file accompanying this file. This file 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 numpy as np
import mxnet as mx
from mxnet import gluon, nd, autograd
import matplotlib.pyplot as plt
from mxnet.gluon import nn, rnn
from tqdm import tqdm_notebook as tqdm
import math

[docs]class NeuralLV:
def __init__(
self,
num_time_series,
num_time_steps,
low_rank_param,
is_full_matrix,
p0,
r,
k,
A,
is_sym,
):
# Define number of time series
self.num_time_series = num_time_series
# Define number of discrete time steps will assume the same for each time series so p_i(t) in R^(dxN)
self.num_time_steps = num_time_steps
self.low_rank_param = low_rank_param
self.is_full_matrix = is_full_matrix
self.p0 = p0
self.r = r
self.k = k
self.A = A
self.is_sym = is_sym
self.dtype = r.dtype
self.ctx = r.context

# mat1 = A is the default case.  If low_rank is on, we have the symmetric case where A=B^TB
# and mat1 = B and in the non-symmetric case we need to pass in mat2 = C so A = B^TC
[docs]    def solve_discrete_lv(self, mat1, mat2=None, is_full_matrix=True):
p = (
[]
)  # need to store as list for autograd won't let you append indices in same matrix
p.append(self.p0)
for n in range(
self.num_time_steps - 1
):  # element-wise vector division and multiplication
# Compute Ap to generate synthetic data for the full rank matrix A
if is_full_matrix:
mat_vec_prod = nd.dot(mat1, p[n])
else:
mat_vec_prod = compute_mat_vec_prod(mat1, mat2, p[n])
p.append((1 + self.r * (1 - mat_vec_prod / self.k)) * p[n])
# concat puts in nd array of size num_ts*N
# need to take size (N, num_ts) and transpose otherwise default is doing row major
# and we need column major storing of the linear list p
return (
nd.concat(*p, dim=0)
.reshape(self.num_time_steps, self.num_time_series)
.T
)

# p0: initial condition shape (num_ts, )
# r: growth rate shape (num_ts, ) TODO: Learn
# k: carrying capacity shape (num_ts, ) TODO: Learn
# A: interaction matrix shape (num_ts, num_ts)
[docs]    def run(self, num_epochs=1000, model=None):
# Compute exact solution with full rank matrix
p = self.solve_discrete_lv(self.A)
# To initialize model otherwise can feed in model and rerun
if model is None:
model = LowRankVectorEmbedding(self)
model.collect_params().initialize(
mx.init.Xavier(magnitude=2.24), force_reinit=True, ctx=self.ctx
)
p_approx = train(p, trainer, model, num_epochs)
B = model.embedding_low_rank_mat_B(model.feat_static_cat).T
C = (
None
if self.is_sym
else model.embedding_low_rank_mat_C(model.feat_static_cat).T
)
A_approx = compute_low_rank_product(B, C)
return p_approx, p, A_approx, model

[docs]class LowRankVectorEmbedding(gluon.HybridBlock):
def __init__(self, neural_lv, **kwargs):
super().__init__(**kwargs)
self.num_time_series = neural_lv.num_time_series  # num_ts
self.low_rank_param = neural_lv.low_rank_param  # low rank parameter k
self.is_full_matrix = neural_lv.is_full_matrix
self.dtype = neural_lv.dtype
self.ctx = neural_lv.ctx
self.feat_static_cat = nd.arange(
self.num_time_series, ctx=self.ctx, dtype=self.dtype
)
self.is_sym = neural_lv.is_sym
self.neural_lv = neural_lv

with self.name_scope():
self.embedding_low_rank_mat_B = gluon.nn.Embedding(
input_dim=self.num_time_series,
output_dim=self.low_rank_param,
dtype=self.dtype,
)
if not self.is_sym:
self.embedding_low_rank_mat_C = gluon.nn.Embedding(
input_dim=self.num_time_series,
output_dim=self.low_rank_param,
dtype=self.dtype,
)

[docs]    def forward(self):
# find low rank vector computed per time series
# feat_static_cat consists of the time series indices 0, ..., cardinality - 1
# embedding returns (num_ts, low_rank_param) need to transpose it
B = self.embedding_low_rank_mat_B(self.feat_static_cat).T
C = (
None
if self.is_sym
else self.embedding_low_rank_mat_C(self.feat_static_cat).T
)
# Explicitly form matrix matrix product A = B^T* CO(kd^2) expensive
if self.is_full_matrix:
return self.neural_lv.solve_discrete_lv(
compute_low_rank_product(B, C)
)
# Compute matrix vector product B^T * (C p) O(kd)
else:
return self.neural_lv.solve_discrete_lv(B, C, self.is_full_matrix)

[docs]def train(p, trainer, model, num_epochs=1000):
loss = gluon.loss.L2Loss()
tqdm_epochs = tqdm(range(num_epochs))
for e in tqdm_epochs:
# forward pass
p_approx = model.forward()
Loss = loss(p, p_approx)
Loss.backward()
trainer.step(model.num_time_series)
tqdm_epochs.set_postfix({"loss": nd.sum(Loss).asscalar()})
return p_approx

[docs]def lv_plot_ts(
p, p_approx, max_num_plots=10, num_rows=2, fig_size_width=10
):  # plots all time series at corresponding time point
plt.rcParams["figure.figsize"] = (fig_size_width, 5)
num_ts = p.shape[0]
N = p.shape[1]
t = np.arange(N)
num_plots = min(num_ts, max_num_plots)
num_cols = int(num_plots / num_rows)
fig, axs = plt.subplots(num_rows, num_cols)
for ts_idx in range(num_plots):
plt.subplot(num_rows, num_cols, ts_idx + 1)
plt.plot(
t, p[ts_idx, :].asnumpy(), t, p_approx[ts_idx, :].asnumpy(), "r--"
)
plt.ylabel(f"$p_{ts_idx}(t)$")
plt.xlabel("t")
plt.legend(("Exact", "Approx"))
plt.xlabel("time: $t$")
plt.ylabel(f"$p_{ts_idx}(t)$")

# Compute B^Tz, where z = C*p, C = B in the symmetic case
[docs]def compute_mat_vec_prod(B, C, p):
# Can replace dot with nd.linalg.gemm2 for the matrix vector multiplication
z = nd.dot(C, p) if C is not None else nd.dot(B, p)
return nd.dot(B, z, transpose_a=True)

# A = B^TC, where C = B in the symmetric case
[docs]def compute_low_rank_product(B, C):
return (
nd.dot(B, C, transpose_a=True)
if C is not None
else nd.dot(B, B, transpose_a=True)
)

# Returns random samples fromuniform distribution [0,1] can change to randn for normally distributed random values
[docs]def generate_data(
num_ts, ctx=mx.gpu(), dtype="float64", seed=100
):  # num_ts = d
np.random.seed(seed)
# vector of shape (num_ts, )
r = np.random.rand(num_ts)
# vector of shape (num_ts, )
k = np.random.rand(num_ts)
# matrix of shape (num_ts, num_ts)
A = np.random.rand(num_ts, num_ts)
# diagonal entries are 1 representing intraspecies competition
np.fill_diagonal(A, 1)
# initial condition vector of shape (num_ts, )
p0 = np.random.rand(num_ts)
return (
nd.array(r, ctx=ctx, dtype=dtype),
nd.array(k, ctx=ctx, dtype=dtype),
nd.array(p0, ctx=ctx, dtype=dtype),
nd.array(A, ctx=ctx, dtype=dtype),
)