Merge pull request #65 from Gromwud/main

Cumulitive update

Author: Gromwud

epde/cache/cache.py

@@ -150,6 +150,8 @@ class Cache(object):
     def __init__(self, device = 'cpu'):
         self._device = device
         self.max_allowed_tensors = None
+        self._g_func_flat_cache = None
+        self._g_func_mask_cache = None
 
         self.memory_default = {'torch' : dict(), 'numpy' : dict()} # TODO: add separate cache for torch tensors and numpy
         self.memory_normalized = {'torch' : dict(), 'numpy' : dict()}
@@ -226,6 +228,20 @@ def g_func(self):  # , g_func: Union[Callable, type(None)] = None
     @g_func.setter
     def g_func(self, function: Union[Callable, np.ndarray, list]):
         self._g_func = function
+        self._g_func_flat_cache = None
+        self._g_func_mask_cache = None
+
+    @property
+    def g_func_flat(self):
+        if self._g_func_flat_cache is None:
+            self._g_func_flat_cache = self.g_func.reshape(-1)
+        return self._g_func_flat_cache
+
+    @property
+    def g_func_mask(self):
+        if self._g_func_mask_cache is None:
+            self._g_func_mask_cache = self.g_func != 0
+        return self._g_func_mask_cache
 
     def add_base_matrix(self, label):
         assert label in self.memory_default['numpy'].keys()

epde/globals.py

@@ -97,12 +97,14 @@ class VerboseManager:
     show_iter_idx : bool
     iter_fitness : bool
     iter_stats : bool
-    show_ann_loss : bool    
+    show_ann_loss : bool
     show_warnings : bool
+    candidate_objectives : bool
 
-def init_verbose(plot_DE_solutions : bool = False, show_iter_idx : bool = True, 
-                 show_iter_fitness : bool = False, show_iter_stats : bool = False, 
-                 show_ann_loss : bool = False, show_warnings : bool = False):
+def init_verbose(plot_DE_solutions : bool = False, show_iter_idx : bool = True,
+                 show_iter_fitness : bool = False, show_iter_stats : bool = False,
+                 show_ann_loss : bool = False, show_warnings : bool = False,
+                 candidate_objectives : bool = False):
     """
     Method for initialized of manager for output in text form
 
@@ -121,8 +123,9 @@ def init_verbose(plot_DE_solutions : bool = False, show_iter_idx : bool = True,
     global verbose
     if not show_warnings:
         warnings.filterwarnings("ignore")
-    verbose = VerboseManager(plot_DE_solutions, show_iter_idx, show_iter_fitness, 
-                             show_iter_stats, show_ann_loss, show_warnings)
+    verbose = VerboseManager(plot_DE_solutions, show_iter_idx, show_iter_fitness,
+                             show_iter_stats, show_ann_loss, show_warnings,
+                             candidate_objectives)
 
 def reset_control_nn(n_control: int = 1, ann: torch.nn.Sequential = None, 
                      ctrl_args: list = [(0, [None,]),], device: str = 'cpu'):

epde/operators/common/fitness.py

@@ -78,7 +78,7 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
             discr_feats = np.dot(features, objective.weights_final[:-1][objective.weights_internal != 0])
 
         discr = (discr_feats + np.full(target.shape, objective.weights_final[-1]) - target)
-        self.g_fun_vals = global_var.grid_cache.g_func.reshape(-1)
+        self.g_fun_vals = global_var.grid_cache.g_func_flat
         discr = np.multiply(discr, self.g_fun_vals)
         rl_error = np.linalg.norm(discr, ord = 2)
 
@@ -153,7 +153,10 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
         objective.aic_calculated = True
 
         data_shape = global_var.grid_cache.inner_shape
-        weights = calculate_weights(features, target, self.g_fun_vals, data_shape)
+        if hasattr(objective, '_cached_sw_weights') and objective._cached_sw_weights is not None:
+            weights = objective._cached_sw_weights
+        else:
+            weights = calculate_weights(features, target, self.g_fun_vals, data_shape)
         weights_arr = np.array(weights)
         std = weights_arr.std(axis=0, ddof=1)
         mu = weights_arr.mean(axis=0)
@@ -164,7 +167,6 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
             cv[mu == 0] = 0.0  # Handle zero mean
 
         total_lr = sum(cv[:-1]) / len(data_shape)
-        # total_lr = sum(dim_results) / target_vals.ndim
 
         objective.fitness_calculated = True
         objective.fitness_value = fitness_value
@@ -173,7 +175,7 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
 
     def get_g_fun_vals(self):
         try:
-            self.g_fun_vals = global_var.grid_cache.g_func[global_var.grid_cache.g_func != 0].reshape(-1)
+            self.g_fun_vals = global_var.grid_cache.g_func[global_var.grid_cache.g_func_mask].reshape(-1)
         except AttributeError:
             self.g_fun_vals = None
 
@@ -230,7 +232,7 @@ def apply(self, objective : SoEq, arguments : dict, force_out_of_place: bool = F
         grids = torch.stack([grid.reshape(-1) for grid in grids], dim = 1).float()
         solution = solution_nn(grids).detach().cpu().numpy()
         self.g_fun_vals = global_var.grid_cache.g_func
-        
+
         if force_out_of_place:
             sum_err = 0
 
@@ -306,10 +308,11 @@ def apply(self, objective: SoEq, arguments: dict, force_out_of_place: bool = Fal
                                                                boundary_conditions=None, use_fourier=True)
 
         _, grids = global_var.grid_cache.get_all(mode='torch')
-        grids = [grid[global_var.grid_cache.g_func != 0] for grid in grids]
+        g_mask = global_var.grid_cache.g_func_mask
+        grids = [grid[g_mask] for grid in grids]
         grids = torch.stack([grid.reshape(-1) for grid in grids], dim=1).float()
         solution = solution_nn(grids).detach().cpu().numpy()
-        self.g_fun_vals = global_var.grid_cache.g_func[global_var.grid_cache.g_func != 0]
+        self.g_fun_vals = global_var.grid_cache.g_func[g_mask]
 
         if force_out_of_place:
             sum_err = 0
@@ -339,7 +342,10 @@ def apply(self, objective: SoEq, arguments: dict, force_out_of_place: bool = Fal
             # Calculate r-loss
             data_shape = global_var.grid_cache.inner_shape
             _, target, features = eq.evaluate(normalize=True, return_val=False)
-            weights = calculate_weights(features, target, self.g_fun_vals, data_shape)
+            if hasattr(eq, '_cached_sw_weights') and eq._cached_sw_weights is not None:
+                weights = eq._cached_sw_weights
+            else:
+                weights = calculate_weights(features, target, self.g_fun_vals, data_shape)
             weights_arr = np.array(weights)
             std = weights_arr.std(axis=0, ddof=1)
             mu = weights_arr.mean(axis=0)
@@ -369,7 +375,7 @@ def feature_reshape(self, features_vals):
 
     def get_g_fun_vals(self):
         try:
-            self.g_fun_vals = global_var.grid_cache.g_func.reshape(-1)
+            self.g_fun_vals = global_var.grid_cache.g_func_flat
         except AttributeError:
             self.g_fun_vals = None
 

epde/operators/common/right_part_selection.py

@@ -68,6 +68,7 @@ def apply(self, objective : Equation, arguments : dict):
                     min_idx = target_idx
                     weights_internal = objective.weights_internal
                     weights_final = [weight for weight in objective.weights_final if weight != 0]
+                    sw_weights = objective._cached_sw_weights
 
                 objective.weights_internal_evald = False
                 objective.weights_final_evald = False
@@ -78,6 +79,7 @@ def apply(self, objective : Equation, arguments : dict):
 
             objective.weights_internal = weights_internal
             objective.weights_final = weights_final
+            objective._cached_sw_weights = sw_weights
             objective.weights_internal_evald = True
             objective.weights_final_evald = True
             objective.target_idx = min_idx

epde/operators/common/sparsity.py

@@ -41,9 +41,11 @@ def get_cv(self, weights):
 
     def fit(self, X, y, sample_weights):
         self.n_samples, self.n_features = X.shape
+        self.cached_weights_ = None
 
         # 1. Initial Weights
         weights = calculate_weights(X, y, sample_weights=sample_weights, grid_shape=self.grid_shape)
+        self.cached_weights_ = weights
         cv = self.get_cv(weights[:, :-1])
 
         self.coef_ = weights.mean(axis=0)[:-1]
@@ -80,6 +82,7 @@ def fit(self, X, y, sample_weights):
 
                 if new_coef == 0:
                     weights = calculate_weights(X[:, self.coef_ != 0], y, sample_weights=sample_weights, grid_shape=self.grid_shape)
+                    self.cached_weights_ = weights
                     new_cv = iter(self.get_cv(weights[:, :-1]))
                     cv = np.array([next(new_cv) if _ else 0 for _ in self.coef_ != 0])
 
@@ -151,14 +154,16 @@ def apply(self, objective : Equation, arguments : dict):
 
         _, target, features = objective.evaluate(normalize = True, return_val = False)
 
-        self.g_fun_vals = global_var.grid_cache.g_func[global_var.grid_cache.g_func != 0]
+        self.g_fun_vals = global_var.grid_cache.g_func[global_var.grid_cache.g_func_mask]
 
         estimator.fit(features, target, self.g_fun_vals)
         objective.weights_internal = estimator.coef_
         objective.weights_internal_evald = True
         objective.weights_final = np.append(objective.weights_internal, estimator.intercept_)
         objective.weights_final_evald = True
         objective.weights_final = [weight for weight in objective.weights_final if weight != 0]
+        objective._cached_sw_weights = estimator.cached_weights_
+        objective._eval_cache = {}
 
 
     def use_default_tags(self):