def get_interpolation_coeffs(directory, data, times, use_noskip, reduced, interpolation_method='cubic'):
# Create new folder for storing coefficients as datasets (interpolation is expensive)
if not os.path.exists(directory):
os.mkdir(directory)
# Dataset name
noskip = 'noskip' if use_noskip else ''
red = 'red' if reduced else ''
timing = 'eqspaced' if times is None or interpolation_method == 'rectilinear' else 'irrspaced' # quick naming fix. TODO: if another equally spaced time series timestamps is in times it wouldn't name it properly...
coeffs_filename = f'{interpolation_method}_coeffs{noskip}{red}_{timing}.hdf5'
absolute_coeffs_filename_path = os.path.join(directory, coeffs_filename)
# Interpolate and save, or load it for use if it exists
coefficients = {}
if not os.path.exists(absolute_coeffs_filename_path):
if interpolation_method == 'cubic':
coefficients['train_coeffs'] = torchcde.natural_cubic_spline_coeffs(data['train_data'], t=times)
coefficients['val_coeffs'] = torchcde.natural_cubic_spline_coeffs(data['val_data'], t=times)
coefficients['test_coeffs'] = torchcde.natural_cubic_spline_coeffs(data['test_data'], t=times)
elif interpolation_method == 'linear':
coefficients['train_coeffs'] = torchcde.linear_interpolation_coeffs(data['train_data'], t=times)
coefficients['val_coeffs'] = torchcde.linear_interpolation_coeffs(data['val_data'], t=times)
coefficients['test_coeffs'] = torchcde.linear_interpolation_coeffs(data['test_data'], t=times)
elif interpolation_method == 'rectilinear': # rectifilinear doesn't work when passing time argument
if timing == 'irrspaced': print('Warning: will do default equally spaced time array instead, rectifilinear interpolation currently works with it only.')
coefficients['train_coeffs'] = torchcde.linear_interpolation_coeffs(data['train_data'], rectilinear=0)
coefficients['val_coeffs'] = torchcde.linear_interpolation_coeffs(data['val_data'], rectilinear=0)
coefficients['test_coeffs'] = torchcde.linear_interpolation_coeffs(data['test_data'], rectilinear=0)
# Save coefficients in the new directory
print('Saving interpolation coefficients ...')
train_nobs, train_ntimes, train_nfeatures = coefficients['train_coeffs'].shape
val_nobs, val_ntimes, val_nfeatures = coefficients['val_coeffs'].shape
test_nobs, test_ntimes, test_nfeatures = coefficients['test_coeffs'].shape
hdf5_coeffs = h5py.File(absolute_coeffs_filename_path , mode='w')
hdf5_coeffs.create_dataset('train', (train_nobs, train_ntimes, train_nfeatures), np.float, data=coefficients['train_coeffs'])
hdf5_coeffs.create_dataset('val', (val_nobs, val_ntimes, val_nfeatures), np.float, data=coefficients['val_coeffs'])
hdf5_coeffs.create_dataset('test', (test_nobs, test_ntimes, test_nfeatures), np.float, data=coefficients['test_coeffs'])
else:
print('Loading corresponding interpolation coefficients ...')
coeffs_dataset = h5py.File(absolute_coeffs_filename_path, mode='r')
coefficients['train_coeffs'] = torch.Tensor(coeffs_dataset['train'][:])
coefficients['val_coeffs'] = torch.Tensor(coeffs_dataset['val'][:])
coefficients['test_coeffs'] = torch.Tensor(coeffs_dataset['test'][:])
train_coeffs = coefficients['train_coeffs']
val_coeffs = coefficients['val_coeffs']
test_coeffs = coefficients['test_coeffs']
print(f'Train data interpolation coefficients shape: {train_coeffs.shape}')
print(f'Validation data interpolation coefficients shape: {val_coeffs.shape}')
print(f'Test data interpolation coefficients shape: {test_coeffs.shape}')
return coefficients
def test_grad_paths():
for method in ('rk4', 'dopri5'):
for adjoint in (True, False):
t = torch.linspace(0, 9, 10, requires_grad=True)
path = torch.rand(1, 10, 3, requires_grad=True)
coeffs = torchcde.natural_cubic_spline_coeffs(path, t)
cubic_spline = torchcde.NaturalCubicSpline(coeffs, t)
z0 = torch.rand(1, 3, requires_grad=True)
func = _Func(input_size=3, hidden_size=3)
t_ = torch.tensor([0., 9.], requires_grad=True)
z = torchcde.cdeint(X=cubic_spline,
func=func,
z0=z0,
t=t_,
adjoint=adjoint,
method=method,
rtol=1e-4,
atol=1e-6)
assert z.shape == (1, 2, 3)
assert t.grad is None
assert path.grad is None
assert z0.grad is None
assert func.variable.grad is None
assert t_.grad is None
z[:, 1].sum().backward()
assert isinstance(t.grad, torch.Tensor)
assert isinstance(path.grad, torch.Tensor)
assert isinstance(z0.grad, torch.Tensor)
assert isinstance(func.variable.grad, torch.Tensor)
assert isinstance(t_.grad, torch.Tensor)
def _solve_cde(x):
# x should be of shape (batch, length, channels)
batch_size = x.size(0)
input_channels = x.size(2)
hidden_channels = 4 # hyperparameter, we can pick whatever we want for this
coeffs = torchcde.natural_cubic_spline_coeffs(x)
X = torchcde.NaturalCubicSpline(coeffs)
z0 = torch.rand(batch_size, hidden_channels)
class F(torch.nn.Module):
def __init__(self):
super(F, self).__init__()
self.linear = torch.nn.Linear(hidden_channels,
hidden_channels * input_channels)
def forward(self, t, z):
return self.linear(z).view(batch_size, hidden_channels, input_channels)
func = F()
zt = torchcde.cdeint(X=X, func=func, z0=z0, t=X.interval)
zT = zt[:, -1] # get the terminal value of the CDE
return zT
def test_specification_and_derivative():
for _ in range(10):
for use_t in (False, True):
for num_batch_dims in (0, 1, 2, 3):
batch_dims = []
for _ in range(num_batch_dims):
batch_dims.append(torch.randint(low=1, high=3, size=(1,)).item())
length = torch.randint(low=5, high=10, size=(1,)).item()
channels = torch.randint(low=1, high=5, size=(1,)).item()
if use_t:
t = torch.linspace(0, 1, length, dtype=torch.float64)
else:
t = torch.linspace(0, length - 1, length, dtype=torch.float64)
x = torch.rand(*batch_dims, length, channels, dtype=torch.float64)
coeffs = torchcde.natural_cubic_spline_coeffs(x, t)
spline = torchcde.NaturalCubicSpline(coeffs, t)
# Test specification
for i, point in enumerate(t):
evaluate = spline.evaluate(point)
xi = x[..., i, :]
assert evaluate.allclose(xi, atol=1e-5, rtol=1e-5)
# Test derivative
for point in torch.rand(100, dtype=torch.float64):
point_with_grad = point.detach().requires_grad_(True)
evaluate = spline.evaluate(point_with_grad)
derivative = spline.derivative(point)
autoderivative = []
for elem in evaluate.view(-1):
elem.backward(retain_graph=True)
with torch.no_grad():
autoderivative.append(point_with_grad.grad.clone())
point_with_grad.grad.zero_()
autoderivative = torch.stack(autoderivative).view(*evaluate.shape)
assert derivative.shape == autoderivative.shape
assert derivative.allclose(autoderivative, atol=1e-5, rtol=1e-5)
def test_short():
for use_t in (False, True):
if use_t:
t = torch.tensor([0., 1.])
else:
t = None
values = torch.rand(2, 1)
coeffs = torchcde.natural_cubic_spline_coeffs(values, t)
spline = torchcde.NaturalCubicSpline(coeffs, t)
coeffs2 = torchcde.linear_interpolation_coeffs(values, t)
linear = torchcde.LinearInterpolation(coeffs2, t)
batch_dims = []
num_channels = 1
_test_equal(batch_dims, num_channels, linear, spline, torch.float32, -1.5, 1.5, 1e-4)
def main(num_epochs=30):
train_X, train_y = get_data()
######################
# input_channels=3 because we have both the horizontal and vertical position of a point in the spiral, and time.
# hidden_channels=8 is the number of hidden channels for the evolving z_t, which we get to choose.
# output_channels=1 because we're doing binary classification.
######################
model = NeuralCDE(input_channels=3, hidden_channels=8, output_channels=1)
optimizer = torch.optim.Adam(model.parameters())
######################
# Now we turn our dataset into a continuous path. We do this here via natural cubic spline interpolation.
# The resulting `train_coeffs` is a tensor describing the path.
# For most problems, it's probably easiest to save this tensor and treat it as the dataset.
######################
train_coeffs = torchcde.natural_cubic_spline_coeffs(train_X)
train_dataset = torch.utils.data.TensorDataset(train_coeffs, train_y)
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=32)
for epoch in range(num_epochs):
for batch in train_dataloader:
batch_coeffs, batch_y = batch
pred_y = model(batch_coeffs).squeeze(-1)
loss = torch.nn.functional.binary_cross_entropy_with_logits(pred_y, batch_y)
loss.backward()
optimizer.step()
optimizer.zero_grad()
print('Epoch: {} Training loss: {}'.format(epoch, loss.item()))
test_X, test_y = get_data()
test_coeffs = torchcde.natural_cubic_spline_coeffs(test_X)
pred_y = model(test_coeffs).squeeze(-1)
binary_prediction = (torch.sigmoid(pred_y) > 0.5).to(test_y.dtype)
prediction_matches = (binary_prediction == test_y).to(test_y.dtype)
proportion_correct = prediction_matches.sum() / test_y.size(0)
print('Test Accuracy: {}'.format(proportion_correct))
def test_interp():
for _ in range(3):
for use_t in (True, False):
for drop in (False, True):
num_points = torch.randint(low=5, high=100, size=(1,)).item()
half_num_points = num_points // 2
num_points = 2 * half_num_points + 1
if use_t:
times1 = torch.rand(half_num_points, dtype=torch.float64) - 1
times2 = torch.rand(half_num_points, dtype=torch.float64)
t = torch.cat([times1, times2, torch.tensor([0.], dtype=torch.float64)]).sort().values
t_ = t
start, end = -1.5, 1.5
del times1, times2
else:
t = torch.linspace(0, num_points - 1, num_points, dtype=torch.float64)
t_ = None
start = 0
end = num_points - 0.5
num_channels = torch.randint(low=1, high=3, size=(1,)).item()
num_batch_dims = torch.randint(low=0, high=3, size=(1,)).item()
batch_dims = []
for _ in range(num_batch_dims):
batch_dims.append(torch.randint(low=1, high=3, size=(1,)).item())
if use_t:
cubic = _Cubic(batch_dims, num_channels, start=t[0], end=t[-1])
knot = 0
else:
cubic = _Offset(batch_dims, num_channels, start=t[0], end=t[-1], offset=t[1] - t[0])
knot = t[1] - t[0]
values = cubic.evaluate(t)
if drop:
for values_slice in values.unbind(dim=-1):
num_drop = int(num_points * torch.randint(low=1, high=4, size=(1,)).item() / 10)
num_drop = min(num_drop, num_points - 4)
# don't drop first or last
to_drop = torch.randperm(num_points - 2)[:num_drop] + 1
to_drop = [x for x in to_drop if x != knot]
values_slice[..., to_drop] = float('nan')
del num_drop, to_drop, values_slice
coeffs = torchcde.natural_cubic_spline_coeffs(values, t_)
spline = torchcde.NaturalCubicSpline(coeffs, t_)
_test_equal(batch_dims, num_channels, cubic, spline, torch.float64, start, end, 1e-3)
def test_linear():
for use_t in (False, True):
start = torch.rand(1).item() * 5 - 2.5
end = torch.rand(1).item() * 5 - 2.5
start, end = min(start, end), max(start, end)
num_points = torch.randint(low=2, high=10, size=(1,)).item()
num_channels = torch.randint(low=1, high=4, size=(1,)).item()
m = torch.rand(num_channels) * 5 - 2.5
c = torch.rand(num_channels) * 5 - 2.5
if use_t:
t = torch.linspace(start, end, num_points)
t_ = t
else:
t = torch.linspace(0, num_points - 1, num_points)
t_ = None
values = m * t.unsqueeze(-1) + c
coeffs = torchcde.natural_cubic_spline_coeffs(values, t_)
spline = torchcde.NaturalCubicSpline(coeffs, t_)
coeffs2 = torchcde.linear_interpolation_coeffs(values, t_)
linear = torchcde.LinearInterpolation(coeffs2, t_)
batch_dims = []
_test_equal(batch_dims, num_channels, linear, spline, torch.float32, -1.5, 1.5, 1e-4)
def test_shape():
for method in ('rk4', 'dopri5'):
for _ in range(10):
num_points = torch.randint(low=5, high=100, size=(1,)).item()
num_channels = torch.randint(low=1, high=3, size=(1,)).item()
num_hidden_channels = torch.randint(low=1, high=5, size=(1,)).item()
num_batch_dims = torch.randint(low=0, high=3, size=(1,)).item()
batch_dims = []
for _ in range(num_batch_dims):
batch_dims.append(torch.randint(low=1, high=3, size=(1,)).item())
t = torch.rand(num_points).sort().values
values = torch.rand(*batch_dims, num_points, num_channels)
coeffs = torchcde.natural_cubic_spline_coeffs(values, t)
spline = torchcde.NaturalCubicSpline(coeffs, t)
class _Func(torch.nn.Module):
def __init__(self):
super(_Func, self).__init__()
self.variable = torch.nn.Parameter(torch.rand(*[1 for _ in range(num_batch_dims)], 1, num_channels))
def forward(self, t, z):
return z.sigmoid().unsqueeze(-1) + self.variable
f = _Func()
z0 = torch.rand(*batch_dims, num_hidden_channels)
num_out_times = torch.randint(low=2, high=10, size=(1,)).item()
out_times = torch.rand(num_out_times, dtype=torch.float64).sort().values * (t[-1] - t[0]) + t[0]
options = {}
if method == 'rk4':
options['step_size'] = 1. / num_points
out = torchcde.cdeint(spline, f, z0, out_times, method=method, options=options, rtol=1e-4, atol=1e-6)
assert out.shape == (*batch_dims, num_out_times, num_hidden_channels)
def interp_():
coeffs = torchcde.natural_cubic_spline_coeffs(path)
yield torchcde.NaturalCubicSpline(coeffs)
coeffs = torchcde.linear_interpolation_coeffs(path)
yield torchcde.LinearInterpolation(coeffs, reparameterise='bump')
