morphopt 3.1.1__py3-none-any.whl

morphopt/optcore/updaters/geometry/update_geometry.py

@@ -0,0 +1,347 @@
+
+import numpy as np
+
+import torch
+from .. import optimizer
+
+from ...modelparams.params import Params
+from ...modelparams import GeometryParams, FEAParams, Materials
+from tabulate import tabulate
+from ..base_updater import BaseUpdater
+
+class UpdaterGeometries(BaseUpdater):
+    """
+    The Updater class is responsible for updating the parameters of the optimization process.
+    It contains methods to update the parameters based on the optimization algorithm used.
+    """
+    from . import objectivefuncs
+
+    def __init__(self, params: Params, max_step_iter: int = 50, max_step_length: float = 0.5, reset_sensitivity_scaler_per_iter: int = 1) -> None:
+        """
+        Initialize the Updater class with the given parameters.
+        
+        Parameters:
+            surfaces (Surfaces): The surfaces object that contains the design variables.
+            max_step_iter (int): The maximum number of iterations for the sub-optimization process.
+            max_step_length_surf (list[float]): The maximum step length for each surface in the optimization process.
+        """
+
+        super().__init__(params)
+
+        self.max_step_iter = max_step_iter
+        """
+        The maximum number of iterations for the sub-optimization process.
+        """
+
+        self.constraints_funcs: dict[str, UpdaterGeometries.objectivefuncs.BaseConstraints] = {}
+        """
+        A list of penalty functions to be optimized. \n
+        L = sum_{i=1}^{n} w_i * f_i(x)
+        """
+
+        self.obj_funcs: dict[str, UpdaterGeometries.objectivefuncs.BaseObjective] = {}
+        """
+        A list of objective functions to be optimized. \n
+        L = sum_{i=1}^{n} w_i * f_i(x)
+        """
+
+        self._delta_control_points_previous: list[np.ndarray] = None
+        """
+        The previous change in control points for each surface.
+        """
+
+        self._weight_points: list[torch.Tensor] = []
+        """
+        The weights for the points in the optimization process.
+        """
+
+        self.params_update: GeometryParams = params.geometry
+        """
+        The surfaces object that contains the design variables.
+        """
+
+        self._reset_sensitivity_scaler_per_iter = reset_sensitivity_scaler_per_iter
+        """
+        The number of iterations after which the scaler is reset.
+        """
+
+        self._max_step_length_max: float = max_step_length
+        """
+        The maximum step length for each surface in the optimization process.
+        """
+
+        self._max_step_length: list[torch.Tensor] = []
+        """
+        The current maximum step length for each surface in the optimization process.
+        """
+
+        self._step_length_min_ratio: float = 0.1
+        """
+        The minimum ratio for the step length relative to the maximum step length.
+        """
+
+        self._step_length_decay: float = 0.5
+        """
+        The decay factor for the step length relative to the maximum step length.
+        """
+
+        self._step_length_increase: float = 1.5
+        """
+        The increase factor for the step length relative to the maximum step length.
+        """
+ 
+        self.if_update: list[float] = None
+        """
+        A list indicating whether each surface needs to be updated.
+        True means the surface needs to be updated, False means it does not.
+        """
+
+    def pathlog_required(self):
+        return ['geometryupdater']
+
+    def add_constraints(self,
+                               obj_func: objectivefuncs.basefuncs,
+                               name: str = None) -> None:
+        """
+        Add an objective function to the list of objective functions.
+
+        Parameters:
+            obj_func (ObjectiveFuncs.BaseObj): The objective function to be added.
+        """
+        if name is None:
+            name = obj_func.__class__.__name__
+
+        extra_num = 0
+        while name + '_%d' % extra_num in self.constraints_funcs.keys():
+            extra_num += 1
+
+        name = name + '_%d' % extra_num
+        self.constraints_funcs[name] = obj_func
+
+    def add_objective_function(self,
+                               obj_func: objectivefuncs.basefuncs,
+                               name: str = None) -> None:
+        """
+        Add an objective function to the list of objective functions.
+
+        Parameters:
+            obj_func (ObjectiveFuncs.BaseObj): The objective function to be added.
+        """
+        if name is None:
+            name = obj_func.__class__.__name__
+
+        extra_num = 0
+        while name + '_%d' % extra_num in self.obj_funcs.keys():
+            extra_num += 1
+
+        name = name + '_%d' % extra_num
+        self.obj_funcs[name] = obj_func
+
+    def reinitialize(self, *args, **kwargs) -> None:
+        """
+        Initialize the parameters of the optimization process.
+
+        Parameters:
+            iter_now (int): The current iteration number.
+        """
+
+        # get the weights for the points in the optimization process
+        self._weight_points = self.params_update.get_points_weight()
+
+        # initialize the objective function
+        r0, rdu0, rdu20 = self.params_update.get_geometry_values()
+
+        # initialize the shape derivative sensitivity
+        self._initialize_objectives(r0=r0, rdu0=rdu0, rdu20=rdu20)
+        
+        # get the total sensitivity
+        sensitivity = self._get_total_sensitivity()
+
+        # initialize the constraint functions
+        self._initialize_constraints(r0=r0, rdu0=rdu0, rdu20=rdu20, sensitivity=sensitivity)
+            
+        # initialize the optimizer
+        self._initialize_optimizer()
+
+        # initialize the step length
+        for i in range(len(self._max_step_length)):
+            if (self._max_step_length[i].shape[0] != self.params.geometry.surface_list[i].control_points.flatten().shape[0] // 3):
+                self._max_step_length[i] = self._max_step_length[i].mean().repeat(self.params.geometry.surface_list[i].num_variables // 3)
+
+        # save the sensitivity for the next iteration
+        self.sensitivity_previous = sensitivity
+          
+    def initialize(self):
+
+        if self.if_update is None:
+            self.if_update = [True for _ in range(self.params.geometry.num_surface)]
+
+        if len(self._max_step_length) == 0:
+            for i in range(self.params_update.num_surface):
+                self._max_step_length.append(torch.ones(self.params.geometry.surface_list[i].num_variables // 3) * self._max_step_length_max * 0.5)
+
+        for i in range(len(self._max_step_length)):
+            if not self.if_update[i]:
+                self._max_step_length[i] *= 0.0
+
+    def _initialize_objectives(self, r0: list[torch.Tensor], rdu0: list[torch.Tensor], rdu20: list[torch.Tensor]) -> None:
+        for obj_func in self.obj_funcs.values():
+            obj_func.initialize(r0=r0, rdu0=rdu0, rdu20=rdu20, weights=self._weight_points)
+
+    def _get_total_sensitivity(self) -> None:
+        # get the total sensitivity
+        sensitivity_all = []
+        for obj_func in self.obj_funcs.values():
+            sensitivity_all.append(obj_func.sensitivity)
+        sensitivity: list[torch.Tensor] = sensitivity_all[0]
+        for surf_ind in range(len(sensitivity)):
+            for obj_ind in range(1, len(sensitivity_all)):
+                sensitivity[surf_ind] += sensitivity_all[obj_ind][surf_ind]
+        return sensitivity
+    
+    def _initialize_constraints(self, r0: list[torch.Tensor], rdu0: list[torch.Tensor], rdu20: list[torch.Tensor], sensitivity: list[torch.Tensor]) -> None:
+        for constraints in self.constraints_funcs.values():
+            constraints.initialize(r0=r0, rdu0=rdu0, rdu20=rdu20, sensitivity=sensitivity, weights=self._weight_points)
+
+    def _initialize_optimizer(self) -> None:
+        self.optimizer = optimizer.LBFGS(closure=self.closure, num_limit=20, tol_error=1e-10)
+        self.iteration_total = 0
+
+    def _update_step_length(self, delta_control_points: list[np.ndarray]) -> None:
+
+        # update the step length based on the number of variables
+        if self._delta_control_points_previous is not None:
+            for i in range(len(self._max_step_length)):
+                # check if the mesh has improved
+                if (self._delta_control_points_previous[i].shape != self.params.geometry.surface_list[i].control_points.shape):
+                    self._max_step_length[i] = self._max_step_length[i].mean().repeat(self.params.geometry.surface_list[i].num_variables // 3)
+                else:
+                    delta_difference: np.ndarray = np.sum(self._delta_control_points_previous[i] * delta_control_points[i], axis=-1) / (np.linalg.norm(delta_control_points[i], axis=-1)) / np.linalg.norm(self._delta_control_points_previous[i], axis=-1)
+                    delta_difference[np.isnan(delta_difference)] = 0.0
+
+                    index_increase = (delta_difference > -0.5).flatten()
+                    index_decrease = (delta_difference <= -0.5).flatten()
+                    self._max_step_length[i][index_increase] = torch.clamp(self._max_step_length[i][index_increase] * self._step_length_increase,
+                                                                          max=self._max_step_length_max)
+                    self._max_step_length[i][index_decrease] = torch.clamp(self._max_step_length[i][index_decrease] * self._step_length_decay,
+                                                                          min=self._max_step_length_max * self._step_length_min_ratio)
+        for i in range(len(self._max_step_length)):
+            if not self.if_update[i]:
+                self._max_step_length[i] *= 0.0
+
+    def closure(self, x: torch.Tensor, return_list=False) -> float:
+        """
+        The closure function for the optimization process.
+
+        Parameters:
+            x (torch.Tensor): The current point in the optimization process.
+
+        Returns:
+            float: The objective function value at the current point.
+        """
+        # save the current point
+        x0 = self.params_update.get_parameters()
+
+        # Set the design variables to the current point
+        self.params_update.update_variables(x_change=x, max_step_length=self._max_step_length)
+
+        # Calculate the objective function value
+        r, rdu, rdu2 = self.params_update.get_geometry_values()
+        
+        constraints_value: list[torch.Tensor] = []
+        for constraints in self.constraints_funcs.values():
+            constraints_value.append(constraints(r=r, rdu=rdu, rdu2=rdu2))
+
+        obj_value: list[torch.Tensor] = []
+        for obj_func in self.obj_funcs.values():
+            obj_value.append(obj_func(r=r, rdu=rdu, rdu2=rdu2))
+
+        # enroll the design variables
+        self.params_update.set_parameters(xlist=x0)
+
+        if return_list:
+            return obj_value, constraints_value
+        else:
+            return sum(obj_value) + sum(constraints_value)
+
+    def update(self) -> torch.Tensor:
+        """
+        Update the parameters of the optimization process.
+        """
+
+        # initialize the optimizer
+        self.reinitialize()
+
+        # update the objective function
+        variables = self.params_update.get_variables().detach().clone()
+
+        # print the information
+        print("\n\n")
+        print("Start updating the surfaces...")
+
+
+        low_step_length_iter = 0
+        gk_new = None
+        for iteration in range(self.max_step_iter):
+            self.iteration_total += 1
+
+            # get the current variables of the surfaces
+            alpha, delta_var, gk_new = self.optimizer.step(x_now=variables, gk_now= gk_new)
+            variables.data += delta_var * alpha
+
+            # check if the step length is too small
+            if abs(alpha) < 1e-10:
+                low_step_length_iter += 1
+
+            if low_step_length_iter > 10:
+                print(
+                    f"Low step length detected ({low_step_length_iter} iterations), stopping optimization."
+                )
+                # break
+
+            # get current objective function value
+            if self.iteration_total % 10 == 0:
+                with torch.no_grad():
+                    obj_values, constraints_values = self.closure(x=variables, return_list=True)
+
+                # print the objective function value
+                # Print a pretty table showing objective values and iteration progress
+                # Clear previous output (move cursor up and clear lines)
+                if iteration > 0:
+                    print("\033[F\033[K" * 4, end="\r")
+
+                headers = ["Iteration"] + ["Total"] + list(self.obj_funcs.keys()) + list(self.constraints_funcs.keys())
+                data = [[f"{iteration+1}/{self.max_step_iter}"] +
+                        [f"{sum(obj_values).item():.6e}"] +
+                        [f"{val.item():.6e}" for val in obj_values] +
+                        [f"{val.item():.6e}" for val in constraints_values]]
+
+                string = tabulate(data, headers=headers, tablefmt="grid")
+                print(string, end="\r")
+
+        return variables.detach().clone()
+
+    def update_variables(self, dx: torch.Tensor) -> None:
+        """
+        Update the variables of the surfaces.
+        """
+        control_points0 = self.params_update.get_control_points_list()
+        control_points0 = [cp.detach().clone().cpu().numpy() for cp in control_points0]
+
+        self.params_update.update_variables(x_change=dx, max_step_length=self._max_step_length)
+
+        control_points_new = self.params_update.get_control_points_list()
+        control_points_new = [cp.detach().clone().cpu().numpy() for cp in control_points_new]
+        delta_control_points = [control_points_new[i] - control_points0[i] for i in range(len(control_points0))]
+        self._update_step_length(delta_control_points=delta_control_points)
+
+        self._delta_control_points_previous = [cp.copy() for cp in delta_control_points]
+
+    def save(self, foldpath, iteration):
+        step_length_numpy = [self._max_step_length[i].detach().cpu().numpy() for i in range(len(self._max_step_length))]
+        data = np.savez(foldpath + self.pathlog_required()[0] + f"/step_length_{iteration}.npz", *step_length_numpy)
+
+    def load(self, foldpath, iteration):
+        data = np.load(foldpath + self.pathlog_required()[0] + f"/step_length_{iteration}.npz")
+        for i in range(len(self._max_step_length)):
+            self._max_step_length[i] = torch.tensor(data['arr_%d' % i]).to(self.params_update.surface_list[i].model.control_points.device)

morphopt/optcore/updaters/optimizer.py

@@ -0,0 +1,177 @@
+import torch
+
+
+class BaseOpt():
+    def __init__(self, closure: callable):
+        self.closure = closure
+        """
+        the objective function to be optimized, which should return the objective value
+        """
+
+        self.grad = torch.func.jacrev(self.closure)
+        """        the gradient of the objective function, which should return the gradient vector
+        """
+
+        # params for line search
+        self._c1 = 1e-4
+        self._rou1 = 0.2
+
+    def set_line_search_params(self, c1: float = None, rou1: float = None):
+        """
+        Set the parameters for the line search method.
+
+        Parameters:
+            c1 (float): The parameter for the Wolfe condition.
+            rou1 (float): The parameter for the backtracking line search.
+        """
+        if c1 is not None:
+            self._c1 = c1
+        if rou1 is not None:
+            self._rou1 = rou1
+
+    def _LineSearchBacktracking(self, x0: torch.Tensor, dx: torch.Tensor, direction: torch.Tensor, alpha0: float = 1.0, obj0: float = None):
+        
+        """
+        Backtracking line search to find the step length for the optimization process.
+        
+        Parameters:
+            x0 (torch.Tensor): The current point in the optimization process.
+            dx (torch.Tensor): The gradient vector at the current point.
+            direction (torch.Tensor): The search direction.
+            alpha0 (float): The initial step length.
+            c1 (float): The parameter for the Wolfe condition.
+            rou1 (float): The parameter for the backtracking line search.
+            obj0 (float): The initial objective value.
+
+        Returns:
+            float: The step length that satisfies the Wolfe condition.
+        """
+
+        if obj0 is None:
+            obj0 = self.closure(x0)
+
+        alpha = alpha0
+        while True:
+            x_new = x0 + alpha * direction
+            with torch.no_grad():
+                obj_new = self.closure(x_new)
+            if ~torch.isnan(obj_new) and ~torch.isinf(
+                    obj_new
+            ) and obj_new < obj0 + self._c1 * alpha * dx.dot(direction):
+                return alpha, obj_new
+            alpha *= self._rou1
+            if (alpha * direction).abs().max() < 1e-14:
+                return 0., obj0
+
+    def step(x_now: torch.Tensor, gk_now: torch.Tensor=None) -> torch.Tensor:
+        """
+        Perform a single optimization step.
+        """
+        pass
+
+class LBFGS(BaseOpt):
+    def __init__(self,
+                 closure: callable,
+                 num_limit:int = 10,
+                 tol_error:float = 1e-10) -> None:
+        super(LBFGS, self).__init__(closure=closure)
+        self.SK = []
+        self.YK = []
+        self.rhok = []
+
+        self.num_limit = num_limit
+        """
+        the number of pairs of Hessian and gradient stored in the memory
+        """
+
+        self.tol_error = tol_error
+        """
+        the tolerance for the convergence of the optimization process
+        """
+
+    def Hg_loop(self, dv: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the product of the inverse Hessian matrix and the gradient vector.
+        This is done using the BFGS update formula.
+
+        Parameters:
+            dv (torch.Tensor): The gradient vector.
+
+        Returns:
+            torch.Tensor: The product of the inverse Hessian matrix and the gradient vector.
+        """
+        
+        q = dv.clone()
+        alpha = torch.zeros(len(self.SK))
+        for i in range(len(self.SK) - 1, -1, -1):
+            alpha[i] = self.rhok[i] * self.SK[i].dot(q)
+            q = q - alpha[i] * self.YK[i]
+        
+        y = q / (self.rhok[-1] * self.YK[-1].dot(self.YK[-1]))
+        
+        for i in range(len(self.SK)):
+            beta = self.rhok[i] * self.YK[i].dot(y)
+            y = y + (alpha[i] - beta) * self.SK[i]
+            
+        return y
+    
+    def step(self, x_now: torch.Tensor, gk_now: torch.Tensor=None):
+        
+        obj_now = self.closure(x_now)
+        if gk_now is None:
+            gk_now: torch.Tensor = self.grad(x_now).flatten()
+
+        gk_now.view(-1)[gk_now.view(-1).isnan()] = 0
+        
+        # if the first iteration, use steepest descent direction
+        if len(self.SK) == 0:
+            dk = -gk_now
+        else:
+            dk = self.Hg_loop(-gk_now)
+        if dk.view(-1).isnan().any():
+            dk = -gk_now
+            self.SK = []
+            self.YK = []
+            self.rhok = []
+
+        # if the gradient is not positive, use steepest descent direction
+        if (dk * gk_now).sum() > 0:
+            dk = -dk
+
+        # line search
+        alpha, obj_new = self._LineSearchBacktracking(x0=x_now, dx=gk_now, direction=dk, alpha0=1., obj0=obj_now)
+
+        # if the step length is too small, stop the iteration
+        if abs(alpha) <= self.tol_error:
+            self.SK = []
+            self.YK = []
+            self.rhok = []
+
+        x_new = x_now + alpha * dk
+        
+
+        if obj_new>obj_now:
+            self.SK = []
+            self.YK = []
+            self.rhok = []
+            return 0., torch.zeros_like(x_now), gk_now
+        
+        gk_new: torch.Tensor =self.grad(x_new).flatten()
+
+        yk = gk_new.flatten() - gk_now.flatten()
+        sk = alpha * dk.flatten()
+
+        # BFGS method
+        if yk.norm() > 1e-14 and sk.norm() > 1e-14:
+            self.SK.append(sk)
+            self.YK.append(yk)
+            self.rhok.append(1 / yk.dot(sk))
+
+            if len(self.SK) > self.num_limit:
+                self.SK = self.SK[1:]
+                self.YK = self.YK[1:]
+                self.rhok = self.rhok[1:]
+        
+        return alpha, dk, gk_new
+
+

morphopt/optcore/updaters/updaters.py

@@ -0,0 +1,82 @@
+
+import torch
+
+from ..baseobject import BaseObject
+
+from .geometry.update_geometry import UpdaterGeometries
+
+
+class Updaters(BaseObject):
+    """
+    This class is responsible for updating the morphologies of the neurons.
+    """
+
+    def __init__(self, surfaces:UpdaterGeometries=None, *args, **kwargs):
+        """
+        Initialize the Updaters class with a neuron object.
+
+        Args:
+            surfaces (UpdaterSurfaces, optional): An instance of the UpdaterSurfaces class for updating the surfaces.
+            loads (UpdaterLoads, optional): An instance of the UpdaterLoads class for updating the loads.
+            materials (UpdaterMaterials, optional): An instance of the UpdaterMaterials class for updating the materials.
+        """
+        self._surface: UpdaterGeometries = None
+        """
+        UpdaterSurfaces: An instance of the UpdaterSurfaces class for updating the surfaces.
+        """
+        self.if_update_surface = False
+        """
+        if_update_surface: A flag indicating whether the surface needs to be updated.
+        """
+        
+        if surfaces is not None:
+            self._surface = surfaces
+            self.if_update_surface = True
+        self._var_surface: torch.Tensor = None
+        """
+        var_surface: The updated surface variables.
+        """
+
+    def reinitialize(self, iteration: int) -> None:
+        """
+        reInitialize the Updaters class.
+        This method reinitializes the surfaces, loads, and materials if they are present.
+        """
+        pass
+
+    def initialize(self) -> None:
+        """
+        Initialize the Updaters class.
+        This method initializes the surfaces, loads, and materials if they are present.
+        """
+        if self.if_update_surface:
+            self._surface.initialize()
+
+    def update(self) -> torch.Tensor:
+        """
+        Update the morphology of the neuron.
+        """
+        
+        if self.if_update_surface:
+            self._var_surface = self._surface.update()
+    
+    def update_variables(self) -> None:
+        """
+        Update the variables of the surfaces and loads.
+        """
+        if self.if_update_surface:
+            self._surface.update_variables(dx=self._var_surface)
+
+    def save(self, foldpath: str, iteration: int) -> None:
+        """
+        Save the updater state to a file.
+        """
+        if self.if_update_surface:
+            self._surface.save(foldpath=foldpath, iteration=iteration)
+
+
+    def pathlog_required(self):
+        paths = []
+        if self.if_update_surface:
+            paths += self._surface.pathlog_required()
+        return paths