diff --git a/evoxels/solvers.py b/evoxels/solvers.py index 11b6890..69112c3 100644 --- a/evoxels/solvers.py +++ b/evoxels/solvers.py @@ -1,13 +1,14 @@ from IPython.display import clear_output from dataclasses import dataclass from typing import Callable, Any, Type +from abc import ABC, abstractmethod from timeit import default_timer as timer import sys from .problem_definition import ODE from .timesteppers import TimeStepper @dataclass -class TimeDependentSolver: +class BaseSolver(ABC): """Generic wrapper for solving one or more fields with a time stepper.""" vf: Any # VoxelFields object fieldnames: str | list[str] @@ -33,46 +34,24 @@ def __post_init__(self): self.profiler = JAXMemoryProfiler() else: raise ValueError(f"Unsupported backend: {self.backend}") - - def solve( - self, - time_increment=0.1, - frames=10, - max_iters=100, - problem_kwargs=None, - jit=True, - verbose=True, - vtk_out=False, - plot_bounds=None, - colormap='viridis' - ): - """Run the time integration loop. - - Args: - time_increment (float): Size of a single time step. - frames (int): Number of output frames (for plotting, vtk, checks). - max_iters (int): Number of time steps to compute. - problem_kwargs (dict | None): Problem-specific input arguments. - jit (bool): Create just-in-time compiled kernel if ``True`` - verbose (bool | str): If ``True`` prints memory stats, ``'plot'`` - updates an interactive plot. - vtk_out (bool): Write VTK files for each frame if ``True``. - plot_bounds (tuple | None): Optional value range for plots. - """ - - problem_kwargs = problem_kwargs or {} + if isinstance(self.fieldnames, str): self.fieldnames = [self.fieldnames] else: self.fieldnames = list(self.fieldnames) + def _init_fields(self): + """Initialize fields in the voxel grid.""" u_list = [self.vg.init_scalar_field(self.vf.fields[name]) for name in self.fieldnames] u = self.vg.concatenate(u_list, 0) u = self.vg.bc.trim_boundary_nodes(u) - + return u + + def _init_stepper(self, time_increment, problem_kwargs, jit): + problem_kwargs = problem_kwargs or {} if self.step_fn is not None: - step = self.step_fn self.problem = None + step = self.step_fn else: if self.problem_cls is None or self.timestepper_cls is None: raise ValueError("Either provide step_fn or both problem_cls and timestepper_cls") @@ -88,23 +67,47 @@ def solve( import torch step = torch.compile(step) - n_out = max_iters // frames - frame = 0 - slice_idx = self.vf.Nz // 2 + return step + + @abstractmethod + def _run_loop(self, u, step, time_increment, frames, max_iters, + vtk_out, verbose, plot_bounds, colormap): + """Abstract method for running the time integration loop.""" + raise NotImplementedError("Subclasses must implement _run_loop method.") - start = timer() - for i in range(max_iters): - time = i * time_increment - if i % n_out == 0: - self._handle_outputs(u, frame, time, slice_idx, vtk_out, verbose, plot_bounds, colormap) - frame += 1 + def solve( + self, + time_increment=0.1, + frames=10, + max_iters=100, + problem_kwargs=None, + jit=True, + verbose=True, + vtk_out=False, + plot_bounds=None, + colormap='viridis' + ): + """Run the time integration loop. - u = step(time, u) + Args: + time_increment (float): Size of a single time step. + frames (int): Number of output frames (for plotting, vtk, checks). + max_iters (int): Number of time steps to compute. + problem_kwargs (dict | None): Problem-specific input arguments. + jit (bool): Create just-in-time compiled kernel if ``True`` + verbose (bool | str): If ``True`` prints memory stats, ``'plot'`` + updates an interactive plot. + vtk_out (bool): Write VTK files for each frame if ``True``. + plot_bounds (tuple | None): Optional value range for plots. + """ + u = self._init_fields() + step = self._init_stepper(time_increment, problem_kwargs, jit) + start = timer() + u = self._run_loop(u, step, time_increment, frames, max_iters, + vtk_out, verbose, plot_bounds, colormap) end = timer() - time = max_iters * time_increment - self._handle_outputs(u, frame, time, slice_idx, vtk_out, verbose, plot_bounds, colormap) - + self.computation_time = end - start if verbose: self.profiler.print_memory_stats(start, end, max_iters) @@ -113,10 +116,10 @@ def _handle_outputs(self, u, frame, time, slice_idx, vtk_out, verbose, plot_boun if getattr(self, 'problem', None) is not None: u_out = self.vg.bc.trim_ghost_nodes(self.problem.pad_bc(u)) else: - u_out = u + u_out = self.vg.bc.trim_ghost_nodes(self.vg.pad_zeros(u)) for i, name in enumerate(self.fieldnames): - self.vf.fields[name] = self.vg.export_scalar_field_to_numpy(u_out[i:i+1]) + self.vf.set_field(name, self.vg.export_scalar_field_to_numpy(u_out[i:i+1])) if verbose: self.profiler.update_memory_stats() @@ -133,3 +136,66 @@ def _handle_outputs(self, u, frame, time, slice_idx, vtk_out, verbose, plot_boun if verbose == 'plot': clear_output(wait=True) self.vf.plot_slice(self.fieldnames[0], slice_idx, time=time, colormap=colormap, value_bounds=plot_bounds) + +@dataclass +class TimeDependentSolver(BaseSolver): + """Solver for time-dependent problems.""" + def _run_loop(self, u, step, time_increment, frames, max_iters, + vtk_out, verbose, plot_bounds, colormap): + n_out = max_iters // frames + frame = 0 + slice_idx = self.vf.Nz // 2 + + for i in range(max_iters): + time = i * time_increment + if i % n_out == 0: + self._handle_outputs(u, frame, time, slice_idx, vtk_out, + verbose, plot_bounds, colormap) + frame += 1 + + u = step(time, u) + time = max_iters * time_increment + self._handle_outputs(u, frame, time, slice_idx, vtk_out, + verbose, plot_bounds, colormap) + return u + +@dataclass +class SteadyStatePseudoTimeSolver(BaseSolver): + """Solver for steady-state problems.""" + conv_crit: float = 1e-6 + check_freq: int = 10 + + def _run_loop(self, u, step, time_increment, frames, max_iters, + vtk_out, verbose, plot_bounds, colormap): + slice_idx = self.vf.Nz // 2 + self.converged = False + self.iter = 0 + + while not self.converged and self.iter < max_iters: + time = self.iter * time_increment + diff = u - step(time, u) + u = step(time, u) + + if self.iter % self.check_freq == 0: + self.converged = self.check_convergence(diff, verbose) + + self._handle_outputs(u, 0, time, slice_idx, vtk_out, + verbose, plot_bounds, colormap) + return u + + def check_convergence(self, diff, verbose): + """Check for convergence based on relative change in fields.""" + converged = True + for i, name in enumerate(self.fieldnames): + # Check if Frobenius norm of change is below threshold + rel_change = self.vg.lib.linalg.norm(diff[i]) / \ + self.vg.lib.sqrt(self.vf.Nx * self.vf.Ny * self.vf.Nz) + if rel_change > self.conv_crit: + converged = False + if verbose: + print(f"Iter {self.iter}: Field '{name}' relative change: {rel_change:.2e}") + + if converged and verbose: + print(f"Converged after {self.iter} iterations.") + + return converged diff --git a/evoxels/voxelfields.py b/evoxels/voxelfields.py index 048af16..5bf689f 100644 --- a/evoxels/voxelfields.py +++ b/evoxels/voxelfields.py @@ -120,30 +120,40 @@ def grid_info(self): grid = Grid(self.shape, self.origin, self.spacing, self.convention) return grid - def add_field(self, name: str, array=None): + def set_field(self, name: str, array: np.ndarray): """ - Adds a field to the voxel grid. + Set field values for an existing field in the voxel grid. Args: name (str): Name of the field. - array (numpy.ndarray, optional): 3D array to initialize the field. If None, initializes with zeros. + array (numpy.ndarray, optional): 3D array. Raises: ValueError: If the provided array does not match the voxel grid dimensions. TypeError: If the provided array is not a numpy array. """ - if array is not None: - if isinstance(array, np.ndarray): - if array.shape == self.shape: - self.fields[name] = array - else: - raise ValueError( - f"The provided array must have the shape {self.shape}." - ) + if isinstance(array, np.ndarray): + if array.shape == self.shape: + self.fields[name] = array else: - raise TypeError("The provided array must be a numpy array.") + raise ValueError( + f"The provided array must have the shape {self.shape}." + ) + else: + raise TypeError("The provided array must be a numpy array.") + + def add_field(self, name: str, array=None): + """ + Adds a field to the voxel grid. + + Args: + name (str): Name of the field. + array (numpy.ndarray, optional): 3D array to initialize the field. If None, initializes with zeros. + """ + if array is not None: + self.set_field(name, array) else: - self.fields[name] = np.zeros(self.shape) + self.set_field(name, np.zeros(self.shape)) def set_voxel_sphere(self, name: str, center, radius, label: int | float = 1): """Create a voxelized representation of a sphere in 3D