diff --git a/code/rocket_functions.py b/code/rocket_functions.py
index f4f84ae..0b0507e 100644
--- a/code/rocket_functions.py
+++ b/code/rocket_functions.py
@@ -1,88 +1,111 @@
 # Angus Dempster, Francois Petitjean, Geoff Webb
+#
+# @article{dempster_etal_2020,
+#   author  = {Dempster, Angus and Petitjean, Fran\c{c}ois and Webb, Geoffrey I},
+#   title   = {ROCKET: Exceptionally fast and accurate time classification using random convolutional kernels},
+#   year    = {2020},
+#   journal = {Data Mining and Knowledge Discovery},
+#   doi     = {https://doi.org/10.1007/s10618-020-00701-z}
+# }
+#
+# https://arxiv.org/abs/1910.13051 (preprint)
 
-# Dempster A, Petitjean F, Webb GI (2019) ROCKET: Exceptionally fast and
-# accurate time series classification using random convolutional kernels.
-# arXiv:1910.13051
-
-from numba import njit, prange
 import numpy as np
+from numba import njit, prange
 
-@njit
+@njit("Tuple((float64[:],int32[:],float64[:],int32[:],int32[:]))(int64,int64)")
 def generate_kernels(input_length, num_kernels):
 
-    candidate_lengths = np.array((7, 9, 11))
+    candidate_lengths = np.array((7, 9, 11), dtype = np.int32)
+    lengths = np.random.choice(candidate_lengths, num_kernels)
 
-    # initialise kernel parameters
-    weights = np.zeros((num_kernels, candidate_lengths.max())) # see note
-    lengths = np.zeros(num_kernels, dtype = np.int32) # see note
-    biases = np.zeros(num_kernels)
+    weights = np.zeros(lengths.sum(), dtype = np.float64)
+    biases = np.zeros(num_kernels, dtype = np.float64)
     dilations = np.zeros(num_kernels, dtype = np.int32)
     paddings = np.zeros(num_kernels, dtype = np.int32)
 
-    # note: only the first *lengths[i]* values of *weights[i]* are used
+    a1 = 0
 
     for i in range(num_kernels):
 
-        length = np.random.choice(candidate_lengths)
-        _weights = np.random.normal(0, 1, length)
-        bias = np.random.uniform(-1, 1)
-        dilation = 2 ** np.random.uniform(0, np.log2((input_length - 1) // (length - 1)))
-        padding = ((length - 1) * dilation) // 2 if np.random.randint(2) == 1 else 0
+        _length = lengths[i]
+
+        _weights = np.random.normal(0, 1, _length)
+
+        b1 = a1 + _length
+        weights[a1:b1] = _weights - _weights.mean()
 
-        weights[i, :length] = _weights - _weights.mean()
-        lengths[i], biases[i], dilations[i], paddings[i] = length, bias, dilation, padding
+        biases[i] = np.random.uniform(-1, 1)
+
+        dilation = 2 ** np.random.uniform(0, np.log2((input_length - 1) / (_length - 1)))
+        dilation = np.int32(dilation)
+        dilations[i] = dilation
+
+        padding = ((_length - 1) * dilation) // 2 if np.random.randint(2) == 1 else 0
+        paddings[i] = padding
+
+        a1 = b1
 
     return weights, lengths, biases, dilations, paddings
 
 @njit(fastmath = True)
 def apply_kernel(X, weights, length, bias, dilation, padding):
 
-    # zero padding
-    if padding > 0:
-        _input_length = len(X)
-        _X = np.zeros(_input_length + (2 * padding))
-        _X[padding:(padding + _input_length)] = X
-        X = _X
-
     input_length = len(X)
 
-    output_length = input_length - ((length - 1) * dilation)
+    output_length = (input_length + (2 * padding)) - ((length - 1) * dilation)
 
-    _ppv = 0 # "proportion of positive values"
+    _ppv = 0
     _max = np.NINF
 
-    for i in range(output_length):
+    end = (input_length + padding) - ((length - 1) * dilation)
+
+    for i in range(-padding, end):
 
         _sum = bias
 
+        index = i
+
         for j in range(length):
 
-            _sum += weights[j] * X[i + (j * dilation)]
+            if index > -1 and index < input_length:
 
-        if _sum > 0:
-            _ppv += 1
+                _sum = _sum + weights[j] * X[index]
+
+            index = index + dilation
 
         if _sum > _max:
             _max = _sum
 
+        if _sum > 0:
+            _ppv += 1
+
     return _ppv / output_length, _max
 
-@njit(parallel = True, fastmath = True)
+@njit("float64[:,:](float64[:,:],Tuple((float64[::1],int32[:],float64[:],int32[:],int32[:])))", parallel = True, fastmath = True)
 def apply_kernels(X, kernels):
 
     weights, lengths, biases, dilations, paddings = kernels
 
-    num_examples = len(X)
-    num_kernels = len(weights)
+    num_examples, _ = X.shape
+    num_kernels = len(lengths)
 
-    # initialise output
-    _X = np.zeros((num_examples, num_kernels * 2)) # 2 features per kernel
+    _X = np.zeros((num_examples, num_kernels * 2), dtype = np.float64) # 2 features per kernel
 
     for i in prange(num_examples):
 
+        a1 = 0 # for weights
+        a2 = 0 # for features
+
         for j in range(num_kernels):
 
-            _X[i, (j * 2):((j * 2) + 2)] = \
-            apply_kernel(X[i], weights[j][:lengths[j]], lengths[j], biases[j], dilations[j], paddings[j])
+            b1 = a1 + lengths[j]
+            b2 = a2 + 2
+
+            _X[i, a2:b2] = \
+            apply_kernel(X[i], weights[a1:b1], lengths[j], biases[j], dilations[j], paddings[j])
+
+            a1 = b1
+            a2 = b2
 
     return _X
checked out 4c2461070e59c10d467b94a5304ca2df70925969 2019-10-30T11:13:58+11:00 Initial commit
