Merge pull request #31 from Gromwud/main

Cumulative update

Author: Maslyaev

epde/cache/cache.py

@@ -361,15 +361,15 @@ def add(self, label, tensor, normalized: bool = False, structural: bool = False,
 
     def delete_entry(self, entry_label):
         print(f'Deleting {entry_label} from cache!')
-        if entry_label not in self.memory_default.keys():
+        if entry_label not in self.memory_default["numpy"].keys():
             raise ValueError('deleted element already not in memory')
-        del self.memory_default[entry_label]
+        del self.memory_default["numpy"][entry_label]
         try:
-            del self.memory_structural[entry_label]
+            del self.memory_structural["numpy"][entry_label]
         except KeyError:
             pass
         try:
-            del self.memory_normalized[entry_label]
+            del self.memory_normalized["numpy"][entry_label]
         except KeyError:
             pass
 

epde/interface/interface.py

@@ -708,6 +708,8 @@ def create_pool(self, data: Union[np.ndarray, list, tuple], variable_names=['u',
                                           for tf in additional_tokens])
         print(f'The cardinality of defined token pool is {self.pool.families_cardinality()}')
         print(f'Among them, the pool contains {self.pool.families_cardinality(meaningful_only=True)}')
+        for family in self.pool.families:
+            family.chech_constancy()
 
     def save_derivatives(self, variable:str, deriv:np.ndarray):
         '''

epde/interface/token_family.py

@@ -257,7 +257,7 @@ def test_evaluator(self, deriv=False):
             self.test_evaluation = self._evaluator.apply(self.test_token)
         print('Test evaluation performed correctly')
 
-    def chech_constancy(self, test_function, **tfkwargs):
+    def chech_constancy(self, **tfkwargs):
         '''
         Method to check, if any single simple token in the studied domain is constant, or close to it. The constant token is to be displayed and deleted from tokens and cache.
 
@@ -270,16 +270,22 @@ def chech_constancy(self, test_function, **tfkwargs):
         assert self.params_set
         constant_tokens_labels = []
         for label in self.tokens:
-            print(type(global_var.tensor_cache.memory[label + ' power 1']))
-            constancy = test_function(global_var.tensor_cache.memory[label + ' power 1'], **tfkwargs)
+            data_label = (label, (1.0,))
+            data = global_var.tensor_cache.memory_default["numpy"].get(data_label)
+            try:
+                constancy = np.isclose(np.min(data), np.max(data))
+            except TypeError:
+                print(f"No {label} data in cache!")
+                continue
             if constancy:
                 constant_tokens_labels.append(label)
 
         for label in constant_tokens_labels:
             print(f'Function {label} is assumed to be constant in the studied domain. \
-                  Removed from the equaton search.')
+                          Removed from the equaton search.')
+            data_label = (label, (1.0,))
             self.tokens.remove(label)
-            global_var.tensor_cache.delete_entry(label + ' power 1')
+            global_var.tensor_cache.delete_entry(data_label)
 
     def evaluate(self, token):
         """

epde/operators/common/fitness.py

@@ -201,7 +201,7 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
                     feature_window = features[start_idx:end_idx, :]
                     estimator = LinearRegression(fit_intercept=False)
                     estimator.fit(feature_window, target_window, sample_weight=self.g_fun_vals[start_idx:end_idx])
-                    valuable_weights = estimator.coef_[:-1]
+                    valuable_weights = estimator.coef_
                     eq_window_weights.append(valuable_weights)
                 eq_cv = np.array([np.abs(np.std(_) / np.mean(_)) for _ in zip(*eq_window_weights)])
                 lr = eq_cv.mean()
@@ -235,7 +235,7 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
                             target_window = target_vals[:, start_idx:end_idx].reshape(-1)
                             feature_window = features.reshape(*data_shape, -1)[:, start_idx:end_idx].reshape(-1, features.shape[-1])
                             estimator.fit(feature_window, target_window, sample_weight=self.g_fun_vals.reshape(*data_shape, -1)[:, start_idx:end_idx].reshape(-1))
-                        valuable_weights = estimator.coef_[:-1]
+                        valuable_weights = estimator.coef_
                         eq_window_weights.append(valuable_weights)
                     eq_cv = np.array([np.abs(np.std(_) / np.mean(_)) for _ in zip(*eq_window_weights)])
                     lr += eq_cv.mean()
@@ -275,7 +275,7 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
                             target_window = target_vals[:, :, start_idx:end_idx].reshape(-1)
                             feature_window = features.reshape(*data_shape, -1)[:, :, start_idx:end_idx].reshape(-1, features.shape[-1])
                             estimator.fit(feature_window, target_window, sample_weight=self.g_fun_vals.reshape(*data_shape, -1)[:, :, start_idx:end_idx].reshape(-1))
-                        valuable_weights = estimator.coef_[:-1]
+                        valuable_weights = estimator.coef_
                         eq_window_weights.append(valuable_weights)
                     eq_cv = np.array([np.abs(np.std(_) / np.mean(_)) for _ in zip(*eq_window_weights)])
                     lr += eq_cv.mean()

epde/preprocessing/smoothers.py

@@ -101,6 +101,7 @@ def __init__(self):
     def forward(self, x):
         x = torch.sin(x)
         return x
+
 class NN(torch.nn.Module):
     def __init__(self,
                  Num_Hidden_Layers   : int          = 3,
@@ -232,63 +233,105 @@ def forward(self, X: torch.Tensor) -> torch.Tensor:
         # Pass through the last layer (with no activation function) and return.
         return self.Layers[self.Num_Hidden_Layers](X);
 
-
 class ANNSmoother(AbstractSmoother):
     def __init__(self):
-        pass
+        super().__init__()  # Optional depending on AbstractSmoother
+        self.model = None
+
+    def __call__(self, data, grid, epochs_max=1000, loss_mean=1000, loss_threshold=1e-8,
+                 batch_frac=0.5, val_frac=0.1, learning_rate=1e-3, return_ann=False, device='cpu'):
+        if torch.cuda.is_available():
+            device = "cuda"
+        # Convert to int if passed as float
+        epochs_max = int(epochs_max)
 
-    def __call__(self, data, grid, epochs_max=1e3, loss_mean=1000, batch_frac=0.5,
-                 learining_rate=1e-4, return_ann: bool = False, device = 'cpu'):
+        # Infer input dimension
         dim = 1 if np.any([s == 1 for s in data.shape]) and data.ndim == 2 else data.ndim
-        model = baseline_ann(dim)
-        # model = NN(Num_Hidden_Layers=5, Neurons_Per_Layer=50, Input_Dim=dim, Activation_Function='Tanh')
-        grid_flattened = torch.from_numpy(np.array([subgrid.reshape(-1) for subgrid in grid])).float().T
 
+        # Initialize model
+        # model = baseline_ann(dim).to(device)
+        model = NN(Num_Hidden_Layers=5, Neurons_Per_Layer=50, Input_Dim=dim, Activation_Function='Sin').to(device)
+        self.model = model
+
+        # Flatten grid and reshape field
+        grid_flattened = torch.from_numpy(np.array([subgrid.reshape(-1) for subgrid in grid])).float().T.to(device)
+        field_ = torch.from_numpy(data.reshape(-1, 1)).float().to(device)
         original_shape = data.shape
 
-        field_ = torch.from_numpy(data.reshape(-1, 1)).float()
+        # Train/val split
+        N = grid_flattened.size(0)
+        val_size = int(N * val_frac)
+        train_size = N - val_size
+        indices = torch.randperm(N)
+        train_idx, val_idx = indices[:train_size], indices[train_size:]
+
+        train_x, train_y = grid_flattened[train_idx], field_[train_idx]
+        val_x, val_y = grid_flattened[val_idx], field_[val_idx]
 
-        # device = torch.device(device)        
-        grid_flattened.to(device)
-        field_.to(device)
-        optimizer = torch.optim.Adam(model.parameters(), lr=learining_rate)
+        # Optimizer and scheduler
+        optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+        scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=epochs_max // 10, gamma=0.5)
+        # scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min', factor=0.5)
+        # scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs_max // 10)
 
-        batch_size = int(data.size * batch_frac)
+        # Loss function
+        loss_fn = torch.nn.MSELoss()
 
-        t = 0
+        # Batch size
+        batch_size = max(1, int(data.size * batch_frac))
 
-        min_loss = np.inf
-        while loss_mean > 1e-5 and t < epochs_max:
+        # Training loop
+        min_val_loss = np.inf
+        best_model_state = None
 
-            permutation = torch.randperm(grid_flattened.size()[0])
+        model.train()
+        for epoch in range(epochs_max):
+            permutation = torch.randperm(train_x.size(0))
+            train_loss_list = []
 
-            loss_list = []
+            for i in range(0, train_x.size(0)-1, batch_size):
+                indices = permutation[i:i + batch_size]
+                batch_x = train_x[indices]
+                batch_y = train_y[indices]
 
-            for i in range(0, grid_flattened.size()[0], batch_size):
                 optimizer.zero_grad()
+                pred = model(batch_x)
+                loss = loss_fn(pred, batch_y)
+                # loss = torch.mean(torch.abs(batch_y - pred))
+                loss.backward()
+                optimizer.step()
+                train_loss_list.append(loss.item())
 
-                indices = permutation[i:i+batch_size]
-                batch_x, batch_y = grid_flattened[indices], field_[indices]
+            train_loss = np.mean(train_loss_list)
+            scheduler.step(train_loss)
 
-                loss = torch.mean(torch.abs(batch_y-model(batch_x)))
+            with torch.no_grad():
+                val_pred = model(val_x)
+                val_loss = loss_fn(val_pred, val_y).item()
+
+            if epoch % 100 == 0:
+                print(f"Epoch {epoch:4d} | Loss: {val_loss:.6e}")
+
+            if val_loss < min_val_loss:
+                min_val_loss = val_loss
+                best_model_state = model.state_dict()
+
+            if val_loss <= loss_threshold:
+                print(f"Early stopping at epoch {epoch}, loss = {val_loss:.4e}")
+                break
+
+        # Load best model and evaluate
+        model.load_state_dict(best_model_state)
+        model.eval()
+
+        with torch.no_grad():
+            prediction = model(grid_flattened).cpu().numpy().reshape(original_shape)
 
-                loss.backward()
-                optimizer.step()
-                loss_list.append(loss.item())
-            loss_mean = np.mean(loss_list)
-            if loss_mean < min_loss:
-                best_model = model
-                min_loss = loss_mean
-            # if global_var.verbose.show_ann_loss:
-            print('Surface training t={}, loss={}'.format(t, loss_mean))
-            t += 1
-
-        data_approx = best_model(grid_flattened).detach().numpy().reshape(original_shape)
         if return_ann:
-            warn('Returning ANN from smoother. This should not occur anywhere, except selected experiments.')
-            return data_approx, best_model
+            warn('Returning ANN from smoother. This should only happen in selected experiments.')
+            return prediction, model
         else:
-            return data_approx
+            return prediction
 
 
 class GaussianSmoother(AbstractSmoother):