diff --git a/src/mace/mace.py b/src/mace/mace.py index 616c9f4..b4c8825 100644 --- a/src/mace/mace.py +++ b/src/mace/mace.py @@ -13,7 +13,6 @@ import torchode as to # Lienen, M., & Günnemann, S. 2022, in The Symbiosis of Deep Learning and Differential Equations II, NeurIPS. https://openreview.net/forum?id=uiKVKTiUYB0 import src.mace.autoencoder as ae import src.mace.latentODE as lODE -from scipy.stats import gmean from time import time @@ -199,87 +198,90 @@ def forward(self, n_0, p, tstep): ## ---------- OLD VERSION OF THE SOLVER CLASS ---------- ## ## This class is compatible with an older version of the autoencoder -class Solver_old(nn.Module): - ''' - The Solver class presents the architecture of MACE. - Components: - 1) Encoder; neural network with adjustable amount of nodes and layers - 2) Neural ODE; ODE given by function g, with trainable elements - 3) Decoder; neural network with adjustable amount of nodes and layers - - ''' - def __init__(self, p_dim, z_dim, DEVICE, n_dim=466, g_nn = False, atol = 1e-5, rtol = 1e-2): - super(Solver_old, self).__init__() # type: ignore - - self.status_train = list() - self.status_test = list() - - self.z_dim = z_dim - self.n_dim = n_dim - self.DEVICE = DEVICE - self.g_nn = g_nn - - ## Setting the neural ODE - input_ae_dim = n_dim - if not self.g_nn: - self.g = lODE.G(z_dim) - input_ae_dim = input_ae_dim+p_dim - self.odeterm = to.ODETerm(self.g, with_args=False) - if self.g_nn: - self.g = lODE.Gnn(p_dim, z_dim) - self.odeterm = to.ODETerm(self.g, with_args=True) - - self.step_method = to.Dopri5(term=self.odeterm) - self.step_size_controller = to.IntegralController(atol=atol, rtol=rtol, term=self.odeterm) - self.adjoint = to.AutoDiffAdjoint(self.step_method, self.step_size_controller).to(self.DEVICE) # type: ignore - - self.jit_solver = torch.compile(self.adjoint) - - ## Setting the autoencoder (enocder + decoder) - hidden_ae_dim = int(gmean([input_ae_dim, z_dim])) - self.encoder = ae.Encoder_old(input_dim=input_ae_dim, hidden_dim=hidden_ae_dim, latent_dim=z_dim) - self.decoder = ae.Decoder_old(latent_dim=z_dim , hidden_dim=hidden_ae_dim, output_dim=n_dim) - def set_status(self, status, phase): - if phase == 'train': - self.status_train.append(status) - elif phase == 'test': - self.status_test.append(status) - - def get_status(self, phase): - if phase == 'train': - return np.array(self.status_train) - elif phase == 'test': - return np.array(self.status_test) - - - def forward(self, n_0, p, tstep): - ''' - Forward function giving the workflow of the MACE architecture. - ''' - - x_0 = n_0 ## use NN version of G - if not self.g_nn: ## DON'T use NN version of G - ## Ravel the abundances n_0 and physical input p to x_0 - x_0 = torch.cat((p, n_0), axis=-1) # type: ignore - - ## Encode x_0, returning the encoded z_0 in latent space - z_0 = self.encoder(x_0) +# from scipy.stats import gmean + +# class Solver_old(nn.Module): +# ''' +# The Solver class presents the architecture of MACE. +# Components: +# 1) Encoder; neural network with adjustable amount of nodes and layers +# 2) Neural ODE; ODE given by function g, with trainable elements +# 3) Decoder; neural network with adjustable amount of nodes and layers + +# ''' +# def __init__(self, p_dim, z_dim, DEVICE, n_dim=466, g_nn = False, atol = 1e-5, rtol = 1e-2): +# super(Solver_old, self).__init__() # type: ignore + +# self.status_train = list() +# self.status_test = list() + +# self.z_dim = z_dim +# self.n_dim = n_dim +# self.DEVICE = DEVICE +# self.g_nn = g_nn + +# ## Setting the neural ODE +# input_ae_dim = n_dim +# if not self.g_nn: +# self.g = lODE.G(z_dim) +# input_ae_dim = input_ae_dim+p_dim +# self.odeterm = to.ODETerm(self.g, with_args=False) +# if self.g_nn: +# self.g = lODE.Gnn(p_dim, z_dim) +# self.odeterm = to.ODETerm(self.g, with_args=True) + +# self.step_method = to.Dopri5(term=self.odeterm) +# self.step_size_controller = to.IntegralController(atol=atol, rtol=rtol, term=self.odeterm) +# self.adjoint = to.AutoDiffAdjoint(self.step_method, self.step_size_controller).to(self.DEVICE) # type: ignore + +# self.jit_solver = torch.compile(self.adjoint) + +# ## Setting the autoencoder (enocder + decoder) +# hidden_ae_dim = int(gmean([input_ae_dim, z_dim])) +# self.encoder = ae.Encoder_old(input_dim=input_ae_dim, hidden_dim=hidden_ae_dim, latent_dim=z_dim) +# self.decoder = ae.Decoder_old(latent_dim=z_dim , hidden_dim=hidden_ae_dim, output_dim=n_dim) + +# def set_status(self, status, phase): +# if phase == 'train': +# self.status_train.append(status) +# elif phase == 'test': +# self.status_test.append(status) + +# def get_status(self, phase): +# if phase == 'train': +# return np.array(self.status_train) +# elif phase == 'test': +# return np.array(self.status_test) + + +# def forward(self, n_0, p, tstep): +# ''' +# Forward function giving the workflow of the MACE architecture. +# ''' + +# x_0 = n_0 ## use NN version of G +# if not self.g_nn: ## DON'T use NN version of G +# ## Ravel the abundances n_0 and physical input p to x_0 +# x_0 = torch.cat((p, n_0), axis=-1) # type: ignore + +# ## Encode x_0, returning the encoded z_0 in latent space +# z_0 = self.encoder(x_0) - ## Create initial value problem - problem = to.InitialValueProblem( - y0 = z_0.to(self.DEVICE), ## "view" is om met de batches om te gaan - t_eval = tstep.view(z_0.shape[0],-1).to(self.DEVICE), - ) +# ## Create initial value problem +# problem = to.InitialValueProblem( +# y0 = z_0.to(self.DEVICE), ## "view" is om met de batches om te gaan +# t_eval = tstep.view(z_0.shape[0],-1).to(self.DEVICE), +# ) - ## Solve initial value problem. Details are set in the __init__() of this class. - solution = self.jit_solver.solve(problem, args=p) - z_s = solution.ys.view(-1, self.z_dim) ## want batches +# ## Solve initial value problem. Details are set in the __init__() of this class. +# solution = self.jit_solver.solve(problem, args=p) +# z_s = solution.ys.view(-1, self.z_dim) ## want batches - ## Decode the resulting values from latent space z_s back to physical space - n_s_ravel = self.decoder(z_s) +# ## Decode the resulting values from latent space z_s back to physical space +# n_s_ravel = self.decoder(z_s) - ## Reshape correctly - n_s = n_s_ravel.reshape(1,tstep.shape[-1], self.n_dim) +# ## Reshape correctly +# n_s = n_s_ravel.reshape(1,tstep.shape[-1], self.n_dim) - return n_s, z_s, solution.status \ No newline at end of file +# return n_s, z_s, solution.status