Updated PINN-based solver & fixed grid issues. ODE - tested, PDE - to be done

Author: Maslyaev

epde/integrate/interface.py

@@ -30,32 +30,32 @@ def _term_solver_form(term, grids, default_domain, variables: List[str] = ['u',]
 
         try:
             coeff_tensor = torch.ones_like(grids[0]).to(device)
-            
         except KeyError:
             raise NotImplementedError('No cache implemented')
+        
         for factor in term.structure:
             if factor.is_deriv:
                 for param_idx, param_descr in factor.params_description.items():
                     if param_descr['name'] == 'power':
                         power_param_idx = param_idx
                 deriv_orders.append(factor.deriv_code)
+
                 if factor.evaluator._evaluator != simple_function_evaluator:
                     if factor.evaluator._evaluator._single_function_token:
                         eval_func = factor.evaluator._evaluator._evaluation_functions_torch 
                     else:
                         eval_func = factor.evaluator._evaluator._evaluation_functions_torch[factor.label]
                     if not isinstance(eval_func, torch.nn.Sequential):
-                        # print(f'for term {factor.name} eval func {eval_func} is non')
                         eval_func_kwargs = dict()
                         for key in factor.evaluator._evaluator.eval_fun_params_labels:
                             for param_idx, param_descr in factor.params_description.items():
                                 if param_descr['name'] == key:
                                     eval_func_kwargs[key] = factor.params[param_idx]
-                        # print(f'eval_func_kwargs for {factor.name} are {eval_func_kwargs}')
                         lbd_eval_func = make_eval_func(eval_func, eval_func_kwargs)
                     deriv_powers.append(lbd_eval_func)
                 else:
                     deriv_powers.append(factor.params[power_param_idx])
+
                 try:
                     if isinstance(factor.variable, str):
                         cur_deriv_var = variables.index(factor.variable)
@@ -83,8 +83,6 @@ def _term_solver_form(term, grids, default_domain, variables: List[str] = ['u',]
         if deriv_vars == []:
             if isinstance(deriv_powers, int) and deriv_powers != 0:
                 raise Exception('Something went wrong with parsing an equation for solver')
-            # elif isinstance(deriv_powers, list) and all([spec_power != 0 for spec_power in deriv_powers]):
-            #     raise Exception('Something went wrong with parsing an equation for solver')
             else:
                 deriv_vars = [0,]
 
@@ -96,7 +94,6 @@ def _term_solver_form(term, grids, default_domain, variables: List[str] = ['u',]
                'pow': deriv_powers,
                'var': deriv_vars}
 
-        # print(f'Translated {term.name} to "term" {deriv_orders}, "pow" {deriv_powers}, "var" {deriv_vars} ')
         return res
 
     @singledispatchmethod
@@ -155,18 +152,24 @@ def adjust_shape(tensor, mode = 'NN'):
         return _solver_form
 
     def use_grids(self, grids=None): # 
+        print('=' * 10 + ' USE GRIDS ' + '=' * 10)
+        print('before: ', grids)
+
         if grids is None and self.grids is None:
             _, self.grids = global_var.grid_cache.get_all(mode = 'torch')
+            self.grids = [grid[global_var.grid_cache.g_func != 0] for grid in self.grids]
         elif grids is not None:
             if len(grids) != len(global_var.grid_cache.get_all(mode = 'torch')[1]):
                 raise ValueError(
                     'Number of passed grids does not match the problem')
             if isinstance(grids[0], np.ndarray):
                 grids = [torch.from_numpy(subgrid).to(self._device) for subgrid in grids]
             self.grids = grids
+        print(self.grids[0].shape)
+
 
 
-    def form(self, grids=None, mode = 'NN'):
+    def form(self, grids=None, mode = 'NN'): # -> List[str, ]:
         self.use_grids(grids=grids)
         equation_forms = []
 

epde/integrate/pinn_integration.py

@@ -9,7 +9,7 @@
 import numpy as np
 import torch
 
-from typing import Callable, Union, Dict, List
+from typing import Callable, Union, Dict, List, Tuple
 from functools import singledispatchmethod, singledispatch
 
 from torch.nn import Sequential
@@ -327,24 +327,24 @@ def create_domain(variables: List[str], grids : List[Union[np.ndarray, torch.Ten
     def solve_epde_system(self, system: Union[SoEq, dict], grids: list=None, boundary_conditions=None,
                           mode='NN', data=None, use_cache: bool = False, use_fourier: bool = False,
                           fourier_params: dict = None, use_adaptive_lambdas: bool = False,
-                          to_numpy: bool = False, grid_var_keys = None, *args, **kwargs):
+                          to_numpy: bool = False, grid_var_keys = None,
+                          *args, **kwargs) -> Tuple[float, Union[torch.Tensor, np.ndarray]]:
         solver_device(device = self._device)
 
         if isinstance(system, SoEq):
             system_interface = SystemSolverInterface(system_to_adapt=system)
             system_solver_forms = system_interface.form(grids = grids, mode = mode)
         elif isinstance(system, dict):
-            system_solver_forms = list(system.values())
+            system_solver_forms = list(system.values()) # TODO: refactor instead of quickfixes
         elif isinstance(system, list):
             system_solver_forms = system
         else:
             raise TypeError(f'Incorrect type of the equations passed into solver. Expected dict or SoEq, got {type(system)}.')
-        
+
         if boundary_conditions is None:
-            raise NotImplementedError('TBD')
             op_gen = PregenBOperator(system=system,
                                      system_of_equation_solver_form=[sf_labeled[1] for sf_labeled
-                                                                     in system.values()])
+                                                                     in system_solver_forms]) # system.values .vals()
             op_gen.generate_default_bc(vals = data, grids = grids)
             boundary_conditions = op_gen.conditions
 
@@ -355,6 +355,7 @@ def solve_epde_system(self, system: Union[SoEq, dict], grids: list=None, boundar
 
         if grids is None:
             grid_var_keys, grids = global_var.grid_cache.get_all(mode = 'torch')
+            grids = [grid[global_var.grid_cache.g_func != 0] for grid in grids]
         elif grid_var_keys is None:
             grid_var_keys, _ = global_var.grid_cache.get_all(mode = 'torch')
 
@@ -367,15 +368,23 @@ def solve_epde_system(self, system: Union[SoEq, dict], grids: list=None, boundar
 
     def solve(self, equations: Union[List, SoEq, SolverEquation], domain: Domain,
               boundary_conditions = None, mode = 'NN', use_cache: bool = False, 
-              use_fourier: bool = False, fourier_params: dict = None, #  epochs = 1e3, 
-              use_adaptive_lambdas: bool = False, to_numpy = False, *args, **kwargs):
+              use_fourier: bool = False, fourier_params: dict = None,
+              use_adaptive_lambdas: bool = False, to_numpy = False,
+              *args, **kwargs) -> Tuple[float, Union[torch.Tensor, np.ndarray]]:
 
         if isinstance(equations, SolverEquation):
             equations_prepared = equations
         else:
             equations_prepared = SolverEquation()
             for form in equations:
-                equations_prepared.add(form)
+                print(f'form is solve has a type of {type(form)}: {form}')
+                if isinstance(form, dict):
+                    equations_prepared.add(form)
+                elif (isinstance(form, list) or isinstance(form, tuple)) and len(form) == 2:
+                    equations_prepared.add(form[1])
+                else:
+                    raise ValueError()
+                
         if self.net is None:
             self.net = self.get_net(equations_prepared, mode, domain, use_fourier,
                                     fourier_params, device=self._device)

epde/interface/interface.py

@@ -699,7 +699,7 @@ 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 = 1e4
+                epochs_max = 1e4 # 1e4
                 global_var.reset_data_repr_nn(data = data, derivs = base_derivs, epochs_max=epochs_max,
                                               grids = grid, predefined_ann = None, device = self._device, 
                                               use_fourier = fourier_layers, fourier_params = fourier_params)

epde/operators/common/fitness.py

@@ -240,8 +240,6 @@ def apply(self, objective : SoEq, arguments : dict, force_out_of_place: bool = F
             else:
                 referential_data = global_var.tensor_cache.get((eq.main_var_to_explain, (1.0,)))
 
-                print(f'solution shape {solution.shape}')
-                print(f'solution[..., eq_idx] {solution[..., eq_idx].shape}, eq_idx {eq_idx}')
                 discr = (solution[..., eq_idx] - referential_data.reshape(solution[..., eq_idx].shape))
                 discr = np.multiply(discr, self.g_fun_vals.reshape(discr.shape))
                 rl_error = np.linalg.norm(discr, ord = 2) 

epde/solver/callbacks/save_model.py

@@ -0,0 +1,29 @@
+import torch
+from epde.solver.callbacks.callback import Callback
+import os
+
+
+
+class SaveModel(Callback):
+    """Class for saving model during train
+    """
+    def __init__(self,
+                 path_to_folder: str,
+                 every_step : int = 1):
+        """
+        Args:
+            path (str): path_to_folder to save model 
+            every_step (int): save model every n steps. Defaults 1.
+        """
+        super().__init__()
+        self.path_to_folder = path_to_folder
+        self.every_step = every_step
+    
+    def save_model(self):
+        model_name = "model-{}.pt".format(self.model.t-1)
+        save_path = os.path.join(self.path_to_folder, model_name)
+        torch.save(self.model.net, save_path)
+
+    def on_epoch_end(self, logs=None):
+        if (self.model.t-1) % self.every_step == 0:
+            self.save_model()

epde/solver/data_CSG.py

@@ -0,0 +1,98 @@
+""" Module for construct a domain with complex or irregular geometry"""
+import torch
+from abc import ABC, abstractmethod
+
+class Shape(ABC):
+    @abstractmethod
+    def contains(self, pts: torch.Tensor) -> torch.BoolTensor:
+        """
+        Returns a mask (N,) indicating which pts are inside the shape.
+        """
+        ...
+
+    @abstractmethod
+    def boundary(self, pts: torch.Tensor, rtol: float = 1e-4, atol: float = 0.0) -> torch.BoolTensor:
+        """
+        Returns a mask (N,) indicating which pts lie on the boundary of the shape.
+        """
+        ...
+
+class Rectangle(Shape):
+    def __init__(self, lower, upper, dims=None):
+        """
+        lower: sequence of lower bounds for each spatial dimension
+        upper: sequence of upper bounds for each spatial dimension
+        dims: indices of dimensions in the grid to apply the rectangle (default first len(lower))
+        """
+        self.lower = torch.as_tensor(lower, dtype=torch.float32)
+        self.upper = torch.as_tensor(upper, dtype=torch.float32)
+        self.dims = dims if dims is not None else list(range(self.lower.numel()))
+
+    def _select(self, pts: torch.Tensor) -> torch.Tensor:
+        return pts[:, self.dims]
+
+    def contains(self, pts: torch.Tensor) -> torch.BoolTensor:
+        sub = self._select(pts)
+        return ((sub >= self.lower) & (sub <= self.upper)).all(dim=1)
+
+    def boundary(self, pts: torch.Tensor, rtol: float = 1e-4, atol: float = 0.0) -> torch.BoolTensor:
+        sub = self._select(pts)
+        inside_or_bound = self.contains(pts)
+        on_lower = torch.isclose(sub, self.lower.unsqueeze(0), rtol=rtol, atol=atol).any(dim=1)
+        on_upper = torch.isclose(sub, self.upper.unsqueeze(0), rtol=rtol, atol=atol).any(dim=1)
+        return inside_or_bound & (on_lower | on_upper)
+
+class Circle(Shape):
+    def __init__(self, center, radius, dims=None):
+        """
+        center: sequence of center coordinates for each spatial dimension
+        radius: scalar radius
+        dims: indices of dimensions in the grid to apply the circle (default first len(center))
+        """
+        self.center = torch.as_tensor(center, dtype=torch.float32)
+        self.radius_sq = float(radius) ** 2
+        self.dims = dims if dims is not None else list(range(self.center.numel()))
+
+    def _select(self, pts: torch.Tensor) -> torch.Tensor:
+        return pts[:, self.dims]
+
+    def contains(self, pts: torch.Tensor) -> torch.BoolTensor:
+        sub = self._select(pts)
+        sqd = ((sub - self.center) ** 2).sum(dim=1)
+        return sqd <= self.radius_sq
+
+    def boundary(self, pts: torch.Tensor, rtol: float = 1e-4, atol: float = 0.0) -> torch.BoolTensor:
+        sub = self._select(pts)
+        sqd = ((sub - self.center) ** 2).sum(dim=1)
+        return torch.isclose(sqd, torch.tensor(self.radius_sq, dtype=pts.dtype), rtol=rtol, atol=atol)
+
+# CSG operations
+
+def csg_difference(grid: torch.Tensor, shape: Shape) -> torch.Tensor:
+    """
+    Returns points of `grid` outside the given `shape`.
+
+    Args:
+        grid: (N, D) tensor of coordinates.
+        shape: a Shape instance.
+    """
+    mask_outside = ~shape.contains(grid)
+    return grid[mask_outside]
+
+
+def csg_boundary(grid: torch.Tensor, shape: Shape, rtol: float = 1e-4, atol: float = 0.0) -> torch.Tensor:
+    """
+    Returns points of `grid` that lie on the boundary of the given `shape`.
+
+    Args:
+        grid: (N, D) tensor of coordinates.
+        shape: a Shape instance.
+        rtol: relative tolerance for boundary detection.
+        atol: absolute tolerance for boundary detection.
+    """
+    mask_bnd = shape.boundary(grid, rtol=rtol, atol=atol)
+    return grid[mask_bnd]
+
+
+
+

epde/solver/derivative.py

@@ -133,6 +133,7 @@ def take_derivative(self, term: dict, grid_points:  torch.Tensor) -> torch.Tenso
         Returns:
             der_term (torch.Tensor): resulting field, computed on a grid.
         """
+        # print(f'In take_derivative {grid_points.shape}')
         dif_dir = list(term.keys())[1]
         # it is may be int, function of grid or torch.Tensor
         if callable(term['coeff']):
@@ -156,16 +157,20 @@ def take_derivative(self, term: dict, grid_points:  torch.Tensor) -> torch.Tenso
                     factor_val = term['pow'][j](der_args)
                 else:             
                     factor_val = term['pow'][j](*der_args)
+                # print(f'factor_val.shape is {factor_val.shape}')
                 der_term = der_term * factor_val
             else:
                 if derivative == [None] or derivative is None:
                     der = self.model(grid_points)[:, term['var'][j]].reshape(-1, 1)
                 else:
                     der = self._nn_autograd(self.model, grid_points, term['var'][j], axis=derivative)
+                # print(f'der.shape is {der.shape} from grid_points of shape {grid_points.shape}')
+
                 if isinstance(term['pow'][j],(int,float)):
                     der_term = der_term * der ** term['pow'][j]
                 elif isinstance(term['pow'][j], Callable):
                     der_term = der_term * term['pow'][j](der)
+        # print(f'coeff.shape {coeff.shape} & der_term.shape {der_term.shape}')
         der_term = coeff * der_term
         return der_term