cumulative refactor

epde/control/control.py

@@ -313,14 +313,15 @@ def modify_bc(operator: dict, scale: Union[float, torch.Tensor]) -> dict:
             loss_hist.append(loss)
 
             if fig_folder is not None and LV_exp:
-                plt.figure(figsize=(11, 6))
+                fig = plt.figure(figsize=(11, 6))
                 plt.plot(grids_merged.cpu().detach().numpy(), control_inputs.cpu().detach().numpy()[:, 0], color = 'k')
                 plt.plot(grids_merged.cpu().detach().numpy(), control_inputs.cpu().detach().numpy()[:, 1], color = 'r')
                 plt.plot(grids_merged.cpu().detach().numpy(), global_var.control_nn.net(control_inputs).cpu().detach().numpy(),
                          color = 'tab:orange')
                 plt.grid()
                 frame_name = f'Exp_{time.month}_{time.day}_at_{time.hour}_{time.minute}_{t}.png'
                 plt.savefig(os.path.join(fig_folder, frame_name))
+                plt.close(fig)
 
             if fig_folder is not None:
                 exp_res = {'state'   : control_inputs.cpu().detach().numpy(),

epde/globals.py

@@ -7,7 +7,6 @@
 """
 
 from dataclasses import dataclass
-import copy
 import warnings
 from typing import List, Union
 
@@ -246,7 +245,7 @@ def reset_data_repr_nn(data: List[np.ndarray], grids: List[np.ndarray], train: b
             scheduler.step(val_loss)
 
             if val_loss < min_val_loss:
-                best_state = copy.deepcopy(model.state_dict())
+                best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
                 min_val_loss = val_loss
                 val_no_improve = 0
             else:

epde/integrate/deepxde_integration.py

@@ -379,8 +379,8 @@ def solve(self, equation_or_system, grids: list, data):
             else:
                 data_list = data
         elif isinstance(equation_or_system, SoEq):
-            eq_list = list(equation_or_system.vals.values())
             var_names = equation_or_system.vars_to_describe
+            eq_list = [equation_or_system.vals[var] for var in equation_or_system.vars_to_describe]
             if isinstance(data, np.ndarray):
                 raise ValueError("For SoEq, data must be a list of arrays (one per variable).")
             data_list = data

epde/interface/interface.py

@@ -699,10 +699,10 @@ def create_pool(self, data: Union[np.ndarray, list, tuple], variable_names=['u',
                 global_var.reset_data_repr_nn(data = data, derivs = base_derivs, train = False, 
                                               grids = grid, predefined_ann = data_nn, device = self._device)
             else:
-                # epochs_max = 1e5 # 1e4
-                global_var.reset_data_repr_nn(data = data, derivs = base_derivs, epochs_max=ann_epochs_max,
-                                              grids = grid, predefined_ann = None, device = self._device, 
-                                              use_fourier = fourier_layers, fourier_params = fourier_params)
+                epochs_max = 1e5 # 1e4
+                # global_var.reset_data_repr_nn(data = data, derivs = base_derivs, epochs_max=ann_epochs_max,
+                #                               grids = grid, predefined_ann = None, device = self._device,
+                #                               use_fourier = fourier_layers, fourier_params = fourier_params)
 
         if isinstance(additional_tokens, list):
             if not all([isinstance(tf, (TokenFamily, PreparedTokens)) for tf in additional_tokens]):

epde/operators/common/fitness.py

@@ -117,7 +117,7 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
         None
         """
         self_args, subop_args = self.parse_suboperator_args(arguments=arguments)
-        # self.suboperators['sparsity'].apply(objective, subop_args['sparsity'])
+
         if force_out_of_place:
             self.suboperators['sparsity'].apply(objective, subop_args['sparsity'])
             if all(objective.weights_internal == 0):
@@ -128,6 +128,8 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
             _, target, features = objective.evaluate(normalize=False, return_val=False)
         else:
             _, target, features = objective.evaluate(normalize=True, return_val=False)
+
+        # self.suboperators['sparsity'].apply(objective, subop_args['sparsity'])
         # _, target, features = objective.evaluate(normalize=False, return_val=False)
 
         self.get_g_fun_vals()
@@ -146,8 +148,8 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
 
         fitness_value = rl_error
 
-        if force_out_of_place:
-            return fitness_value
+        # if force_out_of_place:
+        #     return fitness_value
 
         objective.aic = None
         objective.aic_calculated = True
@@ -156,17 +158,19 @@ def apply(self, objective: Equation, arguments: dict, force_out_of_place: bool =
         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 = calculate_weights(features, target, self.g_fun_vals, data_shape, objective.weights_final[-1] != 0)
         weights_arr = np.array(weights)
         std = weights_arr.std(axis=0, ddof=1)
         mu = weights_arr.mean(axis=0)
 
         # Safe division
         with np.errstate(divide='ignore', invalid='ignore'):
             cv = (std ** 2) / (mu ** 2)
-            cv[mu == 0] = 0.0  # Handle zero mean
 
-        total_lr = sum(cv[:-1]) / len(data_shape)
+        total_lr = sum(cv) / len(data_shape)
+
+        if force_out_of_place:
+            return fitness_value * total_lr
 
         objective.fitness_calculated = True
         objective.fitness_value = fitness_value
@@ -428,7 +432,7 @@ def apply(self, objective, arguments: dict, force_out_of_place: bool = False):
                 raise ValueError("NaN loss")
 
             if isinstance(objective, SoEq):
-                for idx, (var_name, eq) in enumerate(objective.vals.items()):
+                for idx, (var_name, eq) in enumerate({val: objective.vals[val] for val in objective.vars_to_describe}.items()):
                     err = self._compute_error(solution_list[idx], data_list[idx], eq)
                     if force_out_of_place:
                         pass
@@ -510,11 +514,13 @@ def plot_data_vs_solution(grid, data, solution):
         ax.set_xlabel("x1")
         ax.set_ylabel("x2")
         plt.show()
+        plt.close(fig)
     if grid.shape[1]==1:
         fig = plt.figure()
         plt.scatter(grid.reshape(-1), solution.reshape(-1), color = 'r')
         plt.scatter(grid.reshape(-1), data.reshape(-1), color = 'k')
         plt.show()
+        plt.close(fig)
     else:
         raise Exception('Infeasible dimensionality of the input dataset.')
 

epde/operators/common/right_part_selection.py

@@ -67,7 +67,7 @@ def apply(self, objective : Equation, arguments : dict):
                     min_fitness = fitness
                     min_idx = target_idx
                     weights_internal = objective.weights_internal
-                    weights_final = [weight for weight in objective.weights_final if weight != 0]
+                    weights_final = objective.weights_final
                     sw_weights = objective._cached_sw_weights
 
                 objective.weights_internal_evald = False