import numpy as np
from alfi.plot import plot_spatiotemporal_data, tight_kwargs
from alfi.datasets import (
    ReactionDiffusion, ReactionDiffusionGenerator,
    HomogeneousReactionDiffusion, DrosophilaSpatialTranscriptomics,
    ToyTranscriptomicGenerator, ToyTranscriptomics, P53Data
)
from alfi.utilities.data import generate_neural_dataset_1d
import torch
from matplotlib import pyplot as plt

times = np.load('../../../experiments/reaction-diffusion/lfo-2d/0savedmodel_times.npy')
print(times)

[[1.62124280e+09 2.88591735e+00]]

maes = np.load('../../../experiments/maes.npy')
print(maes.shape)
print(maes)
# maes = np.concatenate([maes[0:1], maes[2:]])
print(maes.mean(0))
mus = maes.mean(0)
stds = maes.std(0)
print([f"{mu:.03f} \pm {std:.03f}" for mu, std in zip(mus, stds)])
print(maes.std(0))
print(maes)

(8, 6)
[[0.11592306 0.01499886 0.14960724 0.08717901 0.59995526 0.60820451]
 [0.23967709 0.17669429 0.27794567 0.5331775  0.96472377 0.95857757]
 [0.13342252 0.1116971  0.47449568 0.49677879 0.89472753 1.29174372]
 [0.11744934 0.01429605 0.10625371 0.07622159 0.20530783 0.20755504]
 [0.0447192  0.0193229  0.0799474  0.10034436 0.49640873 0.49233919]
 [0.18396279 0.01190908 0.22114587 0.33653194 0.29246566 0.29409794]
 [0.12110703 0.04262565 0.12458419 0.24493977 0.48721123 0.76658674]
 [0.23486331 0.04683529 0.09372982 0.64144421 0.85250157 0.97170943]]
[0.14889054 0.0547974  0.1909637  0.31457714 0.5991627  0.69885177]
['0.149 \\pm 0.062', '0.055 \\pm 0.055', '0.191 \\pm 0.124', '0.315 \\pm 0.208', '0.599 \\pm 0.264', '0.699 \\pm 0.345']
[0.06205021 0.0554866  0.12446819 0.20843943 0.26397228 0.34534045]
[[0.11592306 0.01499886 0.14960724 0.08717901 0.59995526 0.60820451]
 [0.23967709 0.17669429 0.27794567 0.5331775  0.96472377 0.95857757]
 [0.13342252 0.1116971  0.47449568 0.49677879 0.89472753 1.29174372]
 [0.11744934 0.01429605 0.10625371 0.07622159 0.20530783 0.20755504]
 [0.0447192  0.0193229  0.0799474  0.10034436 0.49640873 0.49233919]
 [0.18396279 0.01190908 0.22114587 0.33653194 0.29246566 0.29409794]
 [0.12110703 0.04262565 0.12458419 0.24493977 0.48721123 0.76658674]
 [0.23486331 0.04683529 0.09372982 0.64144421 0.85250157 0.97170943]]

maes = np.load('../../../experiments/smses.npy')
print(maes.shape)
print(maes)
# maes = np.concatenate([maes[0:1], maes[2:]])
maes = maes[:-1]
print(maes.mean(0))
mus = maes.mean(0)
stds = maes.std(0)
print(' & '.join([f"{mu:.03f} \pm {std:.03f}" for mu, std in zip(mus, stds)]))
print(maes.std(0))
print(maes)

maes = np.load('../../../experiments/q2s.npy')
print(maes.shape)
print(maes)
# maes = np.concatenate([maes[0:1], maes[2:]])
maes = maes[:-1]
print(maes.mean(0))
mus = maes.mean(0)
stds = maes.std(0)
print(' & '.join([f"{100*mu:.03f} \pm {100*std:.03f}" for mu, std in zip(mus, stds)]))
print(maes.std(0))
print(maes)

(8, 4)
[[0.90470374 0.9895197  0.84071165 0.9816863 ]
 [0.9166009  0.989699   0.9031095  0.7736031 ]
 [0.93124306 0.9781323  0.9372911  0.73299205]
 [0.98804474 0.999749   0.9425496  0.9812811 ]
 [0.99818945 0.99959695 0.995339   0.98071426]
 [0.9698974  0.99811107 0.9500379  0.82037437]
 [0.93465286 0.9769651  0.91638124 0.7176237 ]
 [0.81598    0.97432077 0.903434   0.3442139 ]]
[0.94904745 0.9902533  0.9264886  0.8554678 ]
94.905 \pm 3.362 & 99.025 \pm 0.897 & 92.649 \pm 4.420 & 85.547 \pm 11.300
[0.03362352 0.00897263 0.04420087 0.11299903]
[[0.90470374 0.9895197  0.84071165 0.9816863 ]
 [0.9166009  0.989699   0.9031095  0.7736031 ]
 [0.93124306 0.9781323  0.9372911  0.73299205]
 [0.98804474 0.999749   0.9425496  0.9812811 ]
 [0.99818945 0.99959695 0.995339   0.98071426]
 [0.9698974  0.99811107 0.9500379  0.82037437]
 [0.93465286 0.9769651  0.91638124 0.7176237 ]]

from lafomo.models import NeuralOperator
from lafomo.datasets import DrosophilaSpatialTranscriptomics
from torch.utils.data import DataLoader


block_dim = 2
learning_rate = 1e-3

modes = 8
width = 32
in_channels = 3

state_dict = torch.load('../nn/saved_model1205.pt')
model = NeuralOperator(block_dim, in_channels, 2, modes, width, num_layers=4)
print(model.count_params())
model.load_state_dict(state_dict)
# optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4)
# trainer = NeuralOperatorTrainer(model, [optimizer], train_loader, test_loader)

def show_result(model, loader):
    x, y, params = next(iter(loader))
    out, params_out = model(x)

    num_t = x.shape[1]
    num_x = x.shape[2]

    t0 = x[..., 1].min().numpy()
    t1 = x[..., 1].max().numpy()
    print(t0, t1)
    x0 = x[..., 2].min().numpy()
    x1 = x[..., 2].max().numpy()

    extent = [t0, t1, x0, x1]

    print(num_t, num_x)
    plot_spatiotemporal_data(
        [
            out[0, ..., 0].detach().view(num_t, num_x).t(),
            y[0, ..., 0].view(num_t, num_x).t(),
            x[0, ..., 0].view(num_t, num_x).t()
        ],
        extent, nrows=1, ncols=3,
        titles=['Latent (Prediction)', 'Latent (Target)', 'Test input']
    )
    out = out.squeeze()
    y_target = y.squeeze()
    print(params_out[0].detach(), params[0])
    # print(F.mse_loss(out[...,0], y_target))

dataset = DrosophilaSpatialTranscriptomics(gene='kr', data_dir='../../../data', scale=True, scale_tx=True, nn_format=True)

print(dataset.train_data[0][0].shape)
dros_loader = DataLoader(dataset.train_data)

show_result(model, dros_loader)

13641527
params torch.Size([1, 5])
torch.Size([8, 64, 3])
0.0 1.0
8 64
tensor([ 1.8658, -1.5222,  1.2546,  2.1713,  0.3376]) tensor([-1., -1., -1., -1., -1.])

C:\Users\Jacob\Documents\proj\lafomo\lafomo\utilities\data.py:44: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  grid = torch.tensor(grid.reshape(1, s1, s2, 2), dtype=torch.float)
C:\Users\Jacob\Documents\proj\lafomo\lafomo\utilities\data.py:48: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  data = torch.tensor(data.reshape(data.shape[0], s1, s2, 4), dtype=torch.float)

gen = ToyTranscriptomicGenerator(
    num_outputs=5,
    num_latents=1,
    softplus=False,
    num_times=12)
gen.generate(2000, 10, './')

  0%|          | 0/2010 [00:00<?, ?it/s]C:\Users\Jacob\Documents\proj\lafomo\lafomo\datasets\toy.py:112: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  basal = torch.tensor(ref.basal, dtype=self.dtype)
C:\Users\Jacob\Documents\proj\lafomo\lafomo\datasets\toy.py:113: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  decay = torch.tensor(ref.decay, dtype=self.dtype)
C:\Users\Jacob\Documents\proj\lafomo\lafomo\datasets\toy.py:114: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  sensitivity = torch.tensor(ref.sensitivity, dtype=self.dtype)
100%|██████████| 2010/2010 [20:13<00:00,  1.66it/s]

torch.Size([2010, 3, 5, 1])

dataset = ToyTranscriptomics(data_dir='./')
i = torch.randint(0, 100, torch.Size([1]))[0]
print(i)
train = dataset.train_data[i]
print(len(dataset.train_data))

for i in range(gen.num_outputs):
    plt.plot(train[0][..., 1+i])
plt.plot(train[1][..., 0], label='latent')
plt.legend()

tensor(25)
2000

<matplotlib.legend.Legend at 0x1f8084f4088>

dataset = HomogeneousReactionDiffusion('../../../data', one_fixed_sample=False, nn_format=True, ntest=0)
# dataset = HomogeneousReactionDiffusion('../../../data', nn_format=False)


plt.figure()
d = dataset.train_data[0][0][:, :, 0]
i = 2
tx = dataset.data[0][0]
ts = tx[:, :, 1].unique().sort()[0].numpy()
xs = tx[:, :, 2].unique().sort()[0].numpy()
extent = [ts[0], ts[-1], xs[0], xs[-1]]

print('x')
print(dataset.train_data[0][0].shape)
x = dataset.train_data[i][0]
y = dataset.train_data[i][1]
plt.imshow(x[..., 0].t(), origin='lower', extent=extent)
plt.colorbar()
plt.figure()
plt.imshow(y[..., 0].t(), origin='lower', extent=extent)
plt.colorbar()
# plt.imshow(d.t(), origin='lower', extent=extent)

params torch.Size([1050, 4])
x
torch.Size([41, 41, 3])

<matplotlib.colorbar.Colorbar at 0x2b61d8fd748>

import numpy as np
data = ReactionDiffusion(data_dir='../../../data', nn_format=True)
i = np.random.randint(0, 100, 1)[0]
# i = 70

print('shape', data[i][0].shape)
print(data[i][2])
plt.imshow(data[i][0][..., 0].view(41, 41).t(), origin='lower', extent=extent)
plt.figure()
plt.imshow(data[i][1][..., 0].view(41, 41).t(), origin='lower', extent=extent)

torch.Size([1681, 2])
tensor([   0,    1,    2,  ..., 1678, 1679, 1680])
params torch.Size([2000, 4])
shape torch.Size([41, 41, 3])
tensor([0.3000, 0.4000, 0.3500, 0.0010])

<matplotlib.image.AxesImage at 0x2b60dd1ce88>

13641326
params torch.Size([1, 4])
torch.Size([8, 64, 3])
0.0 1.0
8 64
tensor([  2.0080,   0.9432,  30.6544, -46.3498]) tensor([-1., -1., -1., -1.])

C:\Users\Jacob\Documents\proj\lafomo\lafomo\utilities\data.py:44: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  grid = torch.tensor(grid.reshape(1, s1, s2, 2), dtype=torch.float)
C:\Users\Jacob\Documents\proj\lafomo\lafomo\utilities\data.py:48: UserWarning: To copy construct from a tensor, it is recommended to use sourceTensor.clone().detach() or sourceTensor.clone().detach().requires_grad_(True), rather than torch.tensor(sourceTensor).
  data = torch.tensor(data.reshape(data.shape[0], s1, s2, 4), dtype=torch.float)

toydata1 = torch.load('../../../data/intermediate_toydata1_tomerge.pt')
toydata0 = torch.load('../../../data/intermediate_toydata0_tomerge.pt')

print(len(toydata0))
print(len(toydata1))
print(toydata1[0])

toydata = list()
for data in toydata0:
    toydata.append({**data, 'sensitivity': 1.0})


toydata.extend(toydata1)
print(len(toydata))
torch.save(toydata, '../../../data/intermediate_toydata.pt')

294
1050
{'samples': tensor([[ 0.0000,  0.0000,  0.0000,  ...,  0.7967,  0.6844,  0.4588],
        [ 0.0000,  0.0000,  0.0000,  ..., -0.0570, -0.3565, -0.5160],
        [ 0.0000,  0.0000,  0.0000,  ...,  1.7302,  1.2690,  0.4195],
        ...,
        [ 0.0000,  0.0000,  0.0000,  ..., -1.1586, -0.7418, -0.4727],
        [ 0.0000,  0.0000,  0.0000,  ...,  0.7132,  1.1986,  1.5983],
        [ 0.0000,  0.0000,  0.0000,  ..., -1.2674, -0.3059,  0.6457]]), 'l1': 0.1, 'l2': 0.1, 'sensitivity': 0.5, 'diffusion': 0.001, 'decay': 0.01}
1344

ReactionDiffusionGenerator().save_dataset(data_dir='../../../data')

gen = ReactionDiffusionGenerator()
dataset = HomogeneousReactionDiffusion('../../../data', one_fixed_sample=False, nn_format=False, ntest=0)
# dataset = HomogeneousReactionDiffusion('../../../data', nn_format=False)
tx = dataset[0][0]

print(tx.shape)
gen.kyy(tx.t(), tx.t()).shape

torch.Size([2, 1681])

torch.Size([1681, 1681])