From d2cc95825f6acb92cf2cc0b263723d34b33c024f Mon Sep 17 00:00:00 2001 From: Stefan Buller Date: Thu, 3 Nov 2022 17:42:47 -0400 Subject: [PATCH 01/14] Added an objective to take precalculated values of another objective. This new objective wraps the other objective, but uses the precalculated values to evaluate points close to the precalculated points. The intended usage is to quickly redo parts of an optimization, that got interrupted due to for example time limits. --- src/simsopt/objectives/utilities.py | 68 +++++++++++++++++++++++++++++ tests/objectives/test_utilities.py | 56 +++++++++++++++++++++++- 2 files changed, 123 insertions(+), 1 deletion(-) diff --git a/src/simsopt/objectives/utilities.py b/src/simsopt/objectives/utilities.py index 6aea07942..89cd7f462 100644 --- a/src/simsopt/objectives/utilities.py +++ b/src/simsopt/objectives/utilities.py @@ -116,3 +116,71 @@ def __float__(self): def __imul__(self, alpha): self.value *= alpha return self + +class PrecalculatedObjective(Optimizable): + def __init__(self, obj, xs, Js, dJs, radius=1e-8): + r""" + Takes an objective and a set of points, the value of the objective at those points, and it's derivatives at those points. + When ``J(x)`` and ``dJ(x)`` are called for points ``x`` close to the set of points used to initialize this object, it returns the precalculated values and derivatives. Otherwise, the provided objective is called. + Specifically, if :math:`\exists x_s \ni \all |x_i - x_i^s|<\epsilon_i, i \in \{0,...,N-1\}` where ``N`` is ``ndofs``, the precalculated values are used. + This is useful when the provided objective function is expensive to evaluate, but is already known at specific points. The intended usage is to restart an optimization without recalculating the objective at known points. + + Args: + obj: the underlying objective. It should provide a ``.J()`` and ``.dJ()`` function. + xs: the points at which the objective ``obj`` is known. An ``m`` x ``ndofs`` array. + Js: the values at the ``xs`` points. An array of length ``m``. + dJs: the derivatives at the ``xs`` points. An ``m`` x ``ndofs`` array. + radius: A scalar or ``ndofs`` sized array with the radius within which the precalculated points will be used. + """ + Optimizable.__init__(self, x0=np.asarray([]), depends_on=[obj]) + self.obj = obj + self.precomputed_xs = xs + self.precomputed_Js = Js + self.precomputed_dJs = dJs + self.radius = radius + + def J(self): + for precomputed_x,precomputed_J in zip(self.precomputed_xs,self.precomputed_Js): + if np.all(np.fabs(self.x - precomputed_x) < self.radius): + # TODO: this returns the first matching value, not the best matching value + val = precomputed_J + break + else: + val = self.obj.J() + return val + + @derivative_dec + def dJ(self): + for precomputed_x, precomputed_dJ in zip(self.precomputed_xs, self.precomputed_dJs): + if np.all(np.fabs(self.x - precomputed_x) < self.radius): + # TODO: this returns the first matching value, not the best matching value + dval = precomputed_dJ + break + else: + dval = self.obj.dJ(partials=True) + return dval + + + def as_dict(self) -> dict: + d = {} + d["@class"] = self.__class__.__name__ + d["@module"] = self.__class__.__module__ + d["obj"] = self.obj + d["radius"] = np.array(self.radius) + d["xs"] = np.array(self.precomputed_xs) + d["Js"] = np.array(self.precomputed_Js) + d["dJs"] = np.array(self.precomputed_dJs) + return d + + @classmethod + def from_dict(cls, d): + decoder = MontyDecoder() + obj = decoder.process_decoded(d["obj"]) + xs = decoder.process_decoded(d["xs"]) + Js = decoder.process_decoded(d["Js"]) + dJs = decoder.process_decoded(d["dJs"]) + radius = decoder.process_decoded(d["radius"]) + + return cls(obj, xs, Js, dJs, radius) + + return_fn_map = {'J': J, 'dJ': dJ} diff --git a/tests/objectives/test_utilities.py b/tests/objectives/test_utilities.py index 8593bc654..5dbd503d5 100644 --- a/tests/objectives/test_utilities.py +++ b/tests/objectives/test_utilities.py @@ -4,10 +4,14 @@ import numpy as np from monty.json import MontyDecoder, MontyEncoder +from simsopt._core import Optimizable +from simsopt._core.derivative import derivative_dec from simsopt.geo.curvexyzfourier import CurveXYZFourier from simsopt.geo.curveobjectives import CurveLength, LpCurveCurvature, LpCurveTorsion -from simsopt.objectives.utilities import MPIObjective, QuadraticPenalty +from simsopt.objectives.utilities import MPIObjective, QuadraticPenalty, PrecalculatedObjective from simsopt.geo import parameters + + parameters['jit'] = False try: from mpi4py import MPI @@ -58,6 +62,53 @@ def test_quadratic_penalty(self): self.subtest_quadratic_penalty(curve, J.J()+0.1) self.subtest_quadratic_penalty(curve, J.J()-0.1) + def test_precalculated_objective(self): + class Myopt(Optimizable): + def __init__(self): + Optimizable.__init__(self,x0=[0.0,0.0], ames = ["x0","x1"], fixed=[False,False]) + + def J(self): + print("Performing objective calculation") + return self.x[0]**2 + np.cos(self.x[1]) + + @derivative_dec + def dJ(self): + print("Performing derivative calculation") + return np.array([2*self.x[0], - np.sin(self.x[1])]) + + wrapped_obj = Myopt() + precalc_x = [[0.0,0.0],[1.25,9.0],[0.6,0.5]] + precalc_J = [1.0, 0.6513697381153231, 1.2375825618903726] + precalc_dJ = [[0.0, 0.0],[2.5, -0.4121184852417566],[1.2, -0.479425538604203]] + obj = PrecalculatedObjective(wrapped_obj,precalc_x,precalc_J,precalc_dJ,radius = 1e-6) + + # objective function is evaluated away from precalculated points + obj.x = [7.5,2.3] + self.assertAlmostEqual(obj.J(), 55.58372397872017) + + # precalculated value is used when x is closer than 'radius' to a precalculated x + # where radius is applied to all coordinates. + obj.x = [0.99e-6, 0.99e-6] + # Equal to 1 within delta=1e-14 because precalculate is used + self.assertAlmostEqual(obj.J(), 1.0, delta = 1e-14) + # Not equal to the actual value of wrapped_obj at obj.x because precalculated value is used + self.assertNotAlmostEqual(obj.J(), 1.00000000000049, delta = 1e-14) # + + obj.x = [1.01e-6, 1.01e-6] + # Precalculated is not used because we are away from 0 + self.assertAlmostEqual(obj.J(), 1.00000000000051, delta = 1e-14) + # Not equal to 1 within delta=1e-14 because precalculate is not used + self.assertNotAlmostEqual(obj.J(), 1.0, delta = 1e-14) + + # test vector radii + obj2 = PrecalculatedObjective(wrapped_obj,precalc_x,precalc_J,precalc_dJ,radius = [1e-6, 1e-7]) + # not precalculated + obj2.x = [1.25 + 1.1e-7, 9.0 - 1.1e-7] + self.assertAlmostEqual(obj2.dJ(partials=True)[0], 2.50000022, delta=1-7) + # precalculated + obj2.x = [1.25 + 1.1e-7, 9.0 - 0.99e-7] + self.assertAlmostEqual(obj2.dJ(partials=True)[0], 2.5, delta=1-7) + def test_mpi_objective(self): if MPI is None: print("skip test_mpi_objective") @@ -81,3 +132,6 @@ def test_mpi_objective(self): Jmpi1 = MPIObjective(Js1subset, comm, needs_splitting=False) assert abs(Jmpi1.J() - sum(J.J() for J in Js)/n) < 1e-14 assert np.sum(np.abs(Jmpi1.dJ() - sum(J.dJ() for J in Js)/n)) < 1e-14 + +if __name__ == "__main__": + unittest.main() From c77720dec7c451726b9d5fd8dbb3778d129875a5 Mon Sep 17 00:00:00 2001 From: Stefan Buller Date: Thu, 22 Dec 2022 11:58:44 -0500 Subject: [PATCH 02/14] Rewrote jacobian logger to be more agnostic to whether the problem is a least squares problem or generic. Updated the "general" solver in serial. --- src/simsopt/_core/finite_difference.py | 114 ++++++++++++++----------- src/simsopt/objectives/utilities.py | 68 --------------- src/simsopt/solve/mpi.py | 6 +- src/simsopt/solve/serial.py | 71 +++++++++++---- tests/objectives/test_utilities.py | 49 +---------- 5 files changed, 125 insertions(+), 183 deletions(-) diff --git a/src/simsopt/_core/finite_difference.py b/src/simsopt/_core/finite_difference.py index 6ce780d8a..e39079dbf 100644 --- a/src/simsopt/_core/finite_difference.py +++ b/src/simsopt/_core/finite_difference.py @@ -46,7 +46,9 @@ def __init__(self, func: Callable, x0: RealArray = None, abs_step: Real = 1.0e-7, rel_step: Real = 0.0, - diff_method: str = "forward") -> None: + diff_method: str = "forward", + log_file: Union[str, typing.IO] = "jac_log", + flatten_out: bool = False) -> None: try: if not isinstance(func.__self__, Optimizable): @@ -63,29 +65,51 @@ def __init__(self, func: Callable, raise ValueError(f"Finite difference method {diff_method} not implemented. " "Supported methods are 'centered' and 'forward'.") self.diff_method = diff_method + self.log_file = log_file + self.new_log_file = False + self.log_header_written = False + self.flatten_out = flatten_out - self.x0 = np.asarray(x0) if x0 is not None else x0 + x0 = np.asarray(x0) if x0 is not None else x0 + self.x0 = x0 if x0 else self.opt.x self.jac_size = None + self.eval_cnt = 1 + self.nparams = self.opt.dof_size + + def init_log(self): + if isinstance(self.log_file, str): + datestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S") + log_file = self.log_file + "_" + datestr + ".dat" + self.log_file = open(log_file, 'w') + self.new_log_file = True + self.start_time = time() - def jac(self, x: RealArray = None) -> RealArray: + + def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: if x is not None: self.x0 = np.asarray(x) x0 = self.x0 if self.x0 is not None else self.opt.x opt_x0 = self.opt.x - + + self.init_log() + if self.jac_size is None: out = self.fn() if not isinstance(out, (np.ndarray, collections.abc.Sequence)): out = [out] - self.jac_size = (len(out), self.opt.dof_size) - + if self.flatten_out: + self.jac_size = (self.nparams) + else: + self.jac_size = (len(out), self.nparams) + jac = np.zeros(self.jac_size) steps = finite_difference_steps(x0, abs_step=self.abs_step, rel_step=self.rel_step) if self.diff_method == "centered": + self.nevals_jac = 2 * self.nparams # Centered differences: - for j in range(len(x0)): + for j in range(self.nparams): # len(x0) x = np.copy(x0) x[j] = x0[j] + steps[j] @@ -96,20 +120,28 @@ def jac(self, x: RealArray = None) -> RealArray: self.opt.x = x fminus = np.asarray(self.fn()) - jac[:, j] = (fplus - fminus) / (2 * steps[j]) - + if self.flatten_out: + jac[j] = (fplus - fminus) / (2 * steps[j]) + else: + jac[:, j] = (fplus - fminus) / (2 * steps[j]) + elif self.diff_method == "forward": # 1-sided differences + self.nevals_jac = self.nparams + 1 self.opt.x = x0 f0 = np.asarray(self.fn()) - for j in range(len(x0)): + for j in range(self.nparams): # len(x0) x = np.copy(x0) x[j] = x0[j] + steps[j] self.opt.x = x fplus = np.asarray(self.fn()) - jac[:, j] = (fplus - f0) / steps[j] - + if self.flatten_out: + jac[j] = (fplus - f0) / steps[j] + else: + jac[:, j] = (fplus - f0) / steps[j] + + self.eval_cnt += self.nevals_jac # Set the opt.x to the original x self.opt.x = opt_x0 @@ -162,7 +194,8 @@ def __init__(self, func: Callable, self.jac_size = None self.eval_cnt = 1 - + self.nparams = self.opt.dof_size + def __enter__(self): self.mpi_apart() self.init_log() @@ -202,10 +235,9 @@ def _jac(self, x: RealArray = None): logger.info('Beginning parallel finite difference gradient calculation') x0 = np.copy(opt.x) - nparams = opt.dof_size # Make sure all leaders have the same x0. mpi.comm_leaders.Bcast(x0) - logger.info(f'nparams: {nparams}') + logger.info(f'nparams: {self.nparams}') logger.info(f'x0: {x0}') # Set up the list of parameter values to try @@ -214,19 +246,19 @@ def _jac(self, x: RealArray = None): mpi.comm_leaders.Bcast(steps) diff_method = mpi.comm_leaders.bcast(self.diff_method) if diff_method == "centered": - nevals_jac = 2 * nparams - xs = np.zeros((nparams, nevals_jac)) - for j in range(nparams): + self.nevals_jac = 2 * self.nparams + xs = np.zeros((self.nparams, self.nevals_jac)) + for j in range(self.nparams): xs[:, 2 * j] = x0[:] # I don't think I need np.copy(), but not 100% sure. xs[j, 2 * j] = x0[j] + steps[j] xs[:, 2 * j + 1] = x0[:] xs[j, 2 * j + 1] = x0[j] - steps[j] else: # diff_method == "forward": # 1-sided differences - nevals_jac = nparams + 1 - xs = np.zeros((nparams, nevals_jac)) + self.nevals_jac = self.nparams + 1 + xs = np.zeros((self.nparams, self.nevals_jac)) xs[:, 0] = x0[:] - for j in range(nparams): + for j in range(self.nparams): xs[:, j + 1] = x0[:] xs[j, j + 1] = x0[j] + steps[j] @@ -234,15 +266,15 @@ def _jac(self, x: RealArray = None): # nvals = None # Work on this later if not mpi.proc0_world: # All procs other than proc0_world should initialize evals before - # the nevals_jac loop, since they may not have any evals. + # the self.nevals_jac loop, since they may not have any evals. self.jac_size = np.zeros(2, dtype=np.int32) self.jac_size = mpi.comm_leaders.bcast(self.jac_size) - evals = np.zeros((self.jac_size[0], nevals_jac)) + evals = np.zeros((self.jac_size[0], self.nevals_jac)) # Do the hard work of evaluating the functions. - logger.info(f'size of evals is ({self.jac_size[0]}, {nevals_jac})') + logger.info(f'size of evals is ({self.jac_size[0]}, {self.nevals_jac})') ARB_VAL = 100 - for j in range(nevals_jac): + for j in range(self.nevals_jac): # Handle only this group's share of the work: if np.mod(j, mpi.ngroups) == mpi.rank_leaders: mpi.mobilize_workers(ARB_VAL) @@ -253,7 +285,7 @@ def _jac(self, x: RealArray = None): if evals is None and mpi.proc0_world: self.jac_size = mpi.comm_leaders.bcast(self.jac_size) - evals = np.zeros((self.jac_size[0], nevals_jac)) + evals = np.zeros((self.jac_size[0], self.nevals_jac)) evals[:, j] = out # evals[:, j] = np.array([f() for f in dofs.funcs]) @@ -270,12 +302,12 @@ def _jac(self, x: RealArray = None): # Use the evals to form the Jacobian jac = np.zeros(self.jac_size) if diff_method == "centered": - for j in range(nparams): + for j in range(self.nparams): jac[:, j] = (evals[:, 2 * j] - evals[:, 2 * j + 1]) / ( 2 * steps[j]) else: # diff_method == "forward": # 1-sided differences: - for j in range(nparams): + for j in range(self.nparams): jac[:, j] = (evals[:, j + 1] - evals[:, 0]) / steps[j] # Weird things may happen if we do not reset the state vector @@ -329,7 +361,7 @@ def mpi_workers_task(self, *args): traceback.print_exc() # Print traceback # Call to jac function is made in proc0 - def jac(self, x: RealArray = None, *args, **kwargs): + def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: """ Called by proc0 """ @@ -356,26 +388,10 @@ def jac(self, x: RealArray = None, *args, **kwargs): jac, xs, evals = self._jac(x) logger.debug(f'jac is {jac}') + + # Log file is now written externally + # by a wrapper in the serial or mpi solver. - # Write to the log file: - logfile = self.log_file - if not self.log_header_written: - logfile.write(f'Problem type:\nleast_squares\nnparams:\n{len(x)}\n') - logfile.write('function_evaluation,seconds') - for j in range(len(x)): - logfile.write(f',x({j})') - logfile.write('\n') - self.log_header_written = True - nevals = evals.shape[1] - for j in range(nevals): - del_t = time() - self.start_time - j_eval = j + self.eval_cnt - 1 - logfile.write(f'{j_eval:6d},{del_t:12.4e}') - for xj in xs[:, j]: - logfile.write(f',{xj:24.16e}') - logfile.write('\n') - logfile.flush() - - self.eval_cnt += nevals + self.eval_cnt += self.nevals_jac return jac diff --git a/src/simsopt/objectives/utilities.py b/src/simsopt/objectives/utilities.py index d6e3f2b2f..af8ddeb79 100644 --- a/src/simsopt/objectives/utilities.py +++ b/src/simsopt/objectives/utilities.py @@ -101,71 +101,3 @@ def __float__(self): def __imul__(self, alpha): self.value *= alpha return self - -class PrecalculatedObjective(Optimizable): - def __init__(self, obj, xs, Js, dJs, radius=1e-8): - r""" - Takes an objective and a set of points, the value of the objective at those points, and it's derivatives at those points. - When ``J(x)`` and ``dJ(x)`` are called for points ``x`` close to the set of points used to initialize this object, it returns the precalculated values and derivatives. Otherwise, the provided objective is called. - Specifically, if :math:`\exists x_s \ni \all |x_i - x_i^s|<\epsilon_i, i \in \{0,...,N-1\}` where ``N`` is ``ndofs``, the precalculated values are used. - This is useful when the provided objective function is expensive to evaluate, but is already known at specific points. The intended usage is to restart an optimization without recalculating the objective at known points. - - Args: - obj: the underlying objective. It should provide a ``.J()`` and ``.dJ()`` function. - xs: the points at which the objective ``obj`` is known. An ``m`` x ``ndofs`` array. - Js: the values at the ``xs`` points. An array of length ``m``. - dJs: the derivatives at the ``xs`` points. An ``m`` x ``ndofs`` array. - radius: A scalar or ``ndofs`` sized array with the radius within which the precalculated points will be used. - """ - Optimizable.__init__(self, x0=np.asarray([]), depends_on=[obj]) - self.obj = obj - self.precomputed_xs = xs - self.precomputed_Js = Js - self.precomputed_dJs = dJs - self.radius = radius - - def J(self): - for precomputed_x,precomputed_J in zip(self.precomputed_xs,self.precomputed_Js): - if np.all(np.fabs(self.x - precomputed_x) < self.radius): - # TODO: this returns the first matching value, not the best matching value - val = precomputed_J - break - else: - val = self.obj.J() - return val - - @derivative_dec - def dJ(self): - for precomputed_x, precomputed_dJ in zip(self.precomputed_xs, self.precomputed_dJs): - if np.all(np.fabs(self.x - precomputed_x) < self.radius): - # TODO: this returns the first matching value, not the best matching value - dval = precomputed_dJ - break - else: - dval = self.obj.dJ(partials=True) - return dval - - - def as_dict(self) -> dict: - d = {} - d["@class"] = self.__class__.__name__ - d["@module"] = self.__class__.__module__ - d["obj"] = self.obj - d["radius"] = np.array(self.radius) - d["xs"] = np.array(self.precomputed_xs) - d["Js"] = np.array(self.precomputed_Js) - d["dJs"] = np.array(self.precomputed_dJs) - return d - - @classmethod - def from_dict(cls, d): - decoder = MontyDecoder() - obj = decoder.process_decoded(d["obj"]) - xs = decoder.process_decoded(d["xs"]) - Js = decoder.process_decoded(d["Js"]) - dJs = decoder.process_decoded(d["dJs"]) - radius = decoder.process_decoded(d["radius"]) - - return cls(obj, xs, Js, dJs, radius) - - return_fn_map = {'J': J, 'dJ': dJ} diff --git a/src/simsopt/solve/mpi.py b/src/simsopt/solve/mpi.py index de7b1759e..77427de5b 100644 --- a/src/simsopt/solve/mpi.py +++ b/src/simsopt/solve/mpi.py @@ -21,6 +21,9 @@ except ImportError as err: MPI = None +from .serial import finite_difference_jac_wrapper + +from .._core.types import RealArray from .._core.optimizable import Optimizable from ..util.mpi import MpiPartition from .._core.finite_difference import MPIFiniteDifference @@ -200,7 +203,8 @@ def _f_proc0(x): x0 = np.copy(prob.x) logger.info("Using finite difference method implemented in " "SIMSOPT for evaluating gradient") - result = least_squares(_f_proc0, x0, jac=fd.jac, verbose=2, + jac = finite_difference_jac_wrapper(fd) + result = least_squares(_f_proc0, x0, jac=jac, verbose=2, **kwargs) else: diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index cdf0665b7..a71e8c46e 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -16,6 +16,7 @@ import numpy as np from scipy.optimize import least_squares, minimize +from .._core.types import RealArray from ..objectives.least_squares import LeastSquaresProblem from .._core.optimizable import Optimizable from .._core.finite_difference import FiniteDifference @@ -26,6 +27,43 @@ __all__ = ['least_squares_serial_solve', 'serial_solve'] +def finite_difference_jac_wrapper(fd, problem_type = 'least_squares'): + """Wrapper for the `jac` method of the `MPIFiniteDifference` and `FiniteDifference` classes. + For logging the jacobian calculation when used with `scipy.optimize.least_squares`. + Also handles scipy.optimize.minimize, for now.""" + def jac(x: RealArray = None, *args, **kwargs): + ret = fd.jac(x, *args, **kwargs) + log_file = fd.log_file + nparams = fd.nparams + if not fd.log_header_written: + log_file.write(f'Problem type:\n{problem_type}\nnparams:\n{nparams}\n') + log_file.write('function_evaluation, seconds') + if problem_type == 'least_squares': + log_file.write(', d(residual_j)/d(x_i)') + else: + log_file.write(', d(f)/d(x_i)') + log_file.write('\n') + fd.log_header_written = True + del_t = time() - fd.start_time + j_eval = fd.eval_cnt//fd.nevals_jac + log_file.write(f'{j_eval:6d},{del_t:12.4e}') + + # Compute the jacobian of the square-summed residuals + # for scipy least-squares problem + #f = fd.fn() # function value at x0 + #total_jac = np.sum(2 * _jac * f[:, None],axis=0) + #for total_jacj in total_jac: + # log_file.write(f',{total_jacj:24.16e}') + + # Compute the derivative of each residual, as required by scipy.optimize.least_squares + log_file.write(", " + np.array_str(ret, max_line_width = np.inf, precision = None).replace('\n',',')) + log_file.write('\n') + log_file.flush() + return ret + + return jac + + def least_squares_serial_solve(prob: LeastSquaresProblem, grad: bool = None, abs_step: float = 1.0e-7, @@ -59,7 +97,7 @@ def least_squares_serial_solve(prob: LeastSquaresProblem, """ datestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S") - objective_file = open(f"simsopt_{datestr}.dat", 'w') + objective_file = open(f"objective_{datestr}.dat", 'w') residuals_file = open(f"residuals_{datestr}.dat", 'w') nevals = 0 @@ -145,7 +183,8 @@ def objective(x): fd = FiniteDifference(prob.residuals, abs_step=abs_step, rel_step=rel_step, diff_method=diff_method) logger.info("Using derivatives") - result = least_squares(objective, x0, verbose=2, jac=fd.jac, **kwargs) + jac = finite_difference_jac_wrapper(fd) + result = least_squares(objective, x0, verbose=2, jac=jac, **kwargs) else: logger.info("Using derivative-free method") result = least_squares(objective, x0, verbose=2, **kwargs) @@ -186,14 +225,12 @@ def serial_solve(prob: Union[Optimizable, Callable], be used. If ``"forward"``, one-sided finite differences will be used. Else, error is raised. kwargs: Any arguments to pass to - `scipy.optimize.least_squares `_. - For instance, you can supply ``max_nfev=100`` to set - the maximum number of function evaluations (not counting - finite-difference gradient evaluations) to 100. Or, you - can supply ``method`` to choose the optimization algorithm. + `scipy.optimize.minimize `_. + For instance, you can supply ``method`` + to choose the optimization algorithm. """ - filename = "simsopt_" + datetime.now().strftime("%Y-%m-%d-%H-%M-%S") \ + filename = "objective_" + datetime.now().strftime("%Y-%m-%d-%H-%M-%S") \ + ".dat" with open(filename, 'w') as objective_file: datalogging_started = False @@ -202,11 +239,12 @@ def serial_solve(prob: Union[Optimizable, Callable], def objective(x): nonlocal datalogging_started, objective_file, nevals + prob.x = x try: - result = prob(x) + result = prob.J() except: - result = 1e+12 - + logger.info("Exception caught during function evaluation") + result = 1.0e12 # Since the number of terms is not known until the first # evaluation of the objective function, we cannot write the # header of the output file until this first evaluation is @@ -215,7 +253,7 @@ def objective(x): # Initialize log file datalogging_started = True objective_file.write( - f"Problem type:\ngeneral\nnparams:\n{prob.dof_size}\n") + f"Problem type:\nminimize\nnparams:\n{prob.dof_size}\n") objective_file.write("function_evaluation,seconds") for j in range(prob.dof_size): objective_file.write(f",x({j})") @@ -243,12 +281,11 @@ def objective(x): logger.info("Beginning solve.") x0 = np.copy(prob.x) if grad: - raise RuntimeError("Need to convert least-squares Jacobian to " - "gradient of the scalar objective function") logger.info("Using derivatives") - fd = FiniteDifference(prob, abs_step=abs_step, - rel_step=rel_step, diff_method=diff_method) - result = least_squares(objective, x0, verbose=2, jac=fd.jac, + fd = FiniteDifference(prob.J, abs_step=abs_step, + rel_step=rel_step, diff_method=diff_method, flatten_out = True) + jac = finite_difference_jac_wrapper(fd, problem_type='minimize') + result = minimize(objective, x0, options={'disp': True}, jac=jac, **kwargs) else: logger.info("Using derivative-free method") diff --git a/tests/objectives/test_utilities.py b/tests/objectives/test_utilities.py index a41d3504d..b3c0fe407 100644 --- a/tests/objectives/test_utilities.py +++ b/tests/objectives/test_utilities.py @@ -7,7 +7,7 @@ from simsopt._core.derivative import derivative_dec from simsopt.geo.curvexyzfourier import CurveXYZFourier from simsopt.geo.curveobjectives import CurveLength, LpCurveCurvature, LpCurveTorsion -from simsopt.objectives.utilities import MPIObjective, QuadraticPenalty, PrecalculatedObjective +from simsopt.objectives.utilities import MPIObjective, QuadraticPenalty from simsopt.geo import parameters from simsopt._core.json import GSONDecoder, GSONEncoder, SIMSON @@ -60,53 +60,6 @@ def test_quadratic_penalty(self): J = CurveLength(curve) self.subtest_quadratic_penalty(curve, J.J()+0.1) self.subtest_quadratic_penalty(curve, J.J()-0.1) - - def test_precalculated_objective(self): - class Myopt(Optimizable): - def __init__(self): - Optimizable.__init__(self,x0=[0.0,0.0], ames = ["x0","x1"], fixed=[False,False]) - - def J(self): - print("Performing objective calculation") - return self.x[0]**2 + np.cos(self.x[1]) - - @derivative_dec - def dJ(self): - print("Performing derivative calculation") - return np.array([2*self.x[0], - np.sin(self.x[1])]) - - wrapped_obj = Myopt() - precalc_x = [[0.0,0.0],[1.25,9.0],[0.6,0.5]] - precalc_J = [1.0, 0.6513697381153231, 1.2375825618903726] - precalc_dJ = [[0.0, 0.0],[2.5, -0.4121184852417566],[1.2, -0.479425538604203]] - obj = PrecalculatedObjective(wrapped_obj,precalc_x,precalc_J,precalc_dJ,radius = 1e-6) - - # objective function is evaluated away from precalculated points - obj.x = [7.5,2.3] - self.assertAlmostEqual(obj.J(), 55.58372397872017) - - # precalculated value is used when x is closer than 'radius' to a precalculated x - # where radius is applied to all coordinates. - obj.x = [0.99e-6, 0.99e-6] - # Equal to 1 within delta=1e-14 because precalculate is used - self.assertAlmostEqual(obj.J(), 1.0, delta = 1e-14) - # Not equal to the actual value of wrapped_obj at obj.x because precalculated value is used - self.assertNotAlmostEqual(obj.J(), 1.00000000000049, delta = 1e-14) # - - obj.x = [1.01e-6, 1.01e-6] - # Precalculated is not used because we are away from 0 - self.assertAlmostEqual(obj.J(), 1.00000000000051, delta = 1e-14) - # Not equal to 1 within delta=1e-14 because precalculate is not used - self.assertNotAlmostEqual(obj.J(), 1.0, delta = 1e-14) - - # test vector radii - obj2 = PrecalculatedObjective(wrapped_obj,precalc_x,precalc_J,precalc_dJ,radius = [1e-6, 1e-7]) - # not precalculated - obj2.x = [1.25 + 1.1e-7, 9.0 - 1.1e-7] - self.assertAlmostEqual(obj2.dJ(partials=True)[0], 2.50000022, delta=1-7) - # precalculated - obj2.x = [1.25 + 1.1e-7, 9.0 - 0.99e-7] - self.assertAlmostEqual(obj2.dJ(partials=True)[0], 2.5, delta=1-7) def test_mpi_objective(self): if MPI is None: From e799bcb224279cc03bb5687cf77955d900fed55e Mon Sep 17 00:00:00 2001 From: Stefan Buller Date: Thu, 22 Dec 2022 14:48:51 -0500 Subject: [PATCH 03/14] Some small tweaks to work better with MPI (hopefully) --- src/simsopt/_core/finite_difference.py | 21 ++++++++++++++------- src/simsopt/solve/serial.py | 2 +- 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/src/simsopt/_core/finite_difference.py b/src/simsopt/_core/finite_difference.py index e39079dbf..e7c261867 100644 --- a/src/simsopt/_core/finite_difference.py +++ b/src/simsopt/_core/finite_difference.py @@ -76,7 +76,12 @@ def __init__(self, func: Callable, self.jac_size = None self.eval_cnt = 1 self.nparams = self.opt.dof_size - + if self.diff_method == "centered": + self.nevals_jac = 2 * self.nparams + else: + # 1-sided differences + self.nevals_jac = self.nparams + 1 + def init_log(self): if isinstance(self.log_file, str): datestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S") @@ -107,7 +112,6 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: steps = finite_difference_steps(x0, abs_step=self.abs_step, rel_step=self.rel_step) if self.diff_method == "centered": - self.nevals_jac = 2 * self.nparams # Centered differences: for j in range(self.nparams): # len(x0) x = np.copy(x0) @@ -127,7 +131,6 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: elif self.diff_method == "forward": # 1-sided differences - self.nevals_jac = self.nparams + 1 self.opt.x = x0 f0 = np.asarray(self.fn()) for j in range(self.nparams): # len(x0) @@ -195,7 +198,13 @@ def __init__(self, func: Callable, self.jac_size = None self.eval_cnt = 1 self.nparams = self.opt.dof_size - + # set self.nevals_jac here to set it for every process + if self.diff_method == "centered": + self.nevals_jac = 2 * self.nparams + else: + # 1-sided differences + self.nevals_jac = self.nparams + 1 + def __enter__(self): self.mpi_apart() self.init_log() @@ -218,7 +227,7 @@ def __exit__(self, exc_type, exc_value, tb): self.mpi.together() if self.mpi.proc0_world and self.new_log_file: self.log_file.close() - + # Called by MPI leaders def _jac(self, x: RealArray = None): # Use shortcuts for class variables @@ -246,7 +255,6 @@ def _jac(self, x: RealArray = None): mpi.comm_leaders.Bcast(steps) diff_method = mpi.comm_leaders.bcast(self.diff_method) if diff_method == "centered": - self.nevals_jac = 2 * self.nparams xs = np.zeros((self.nparams, self.nevals_jac)) for j in range(self.nparams): xs[:, 2 * j] = x0[:] # I don't think I need np.copy(), but not 100% sure. @@ -255,7 +263,6 @@ def _jac(self, x: RealArray = None): xs[j, 2 * j + 1] = x0[j] - steps[j] else: # diff_method == "forward": # 1-sided differences - self.nevals_jac = self.nparams + 1 xs = np.zeros((self.nparams, self.nevals_jac)) xs[:, 0] = x0[:] for j in range(self.nparams): diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index a71e8c46e..a3a51443e 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -54,7 +54,7 @@ def jac(x: RealArray = None, *args, **kwargs): #total_jac = np.sum(2 * _jac * f[:, None],axis=0) #for total_jacj in total_jac: # log_file.write(f',{total_jacj:24.16e}') - + # Compute the derivative of each residual, as required by scipy.optimize.least_squares log_file.write(", " + np.array_str(ret, max_line_width = np.inf, precision = None).replace('\n',',')) log_file.write('\n') From e2b708363a6219624f3c16261625820194419003 Mon Sep 17 00:00:00 2001 From: daringli Date: Thu, 22 Dec 2022 21:19:19 +0100 Subject: [PATCH 04/14] Fixed printing for big arrays --- src/simsopt/mhd/vmec_diagnostics.py | 2 +- src/simsopt/solve/serial.py | 3 ++- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/simsopt/mhd/vmec_diagnostics.py b/src/simsopt/mhd/vmec_diagnostics.py index 7a1d69d7a..389be5093 100644 --- a/src/simsopt/mhd/vmec_diagnostics.py +++ b/src/simsopt/mhd/vmec_diagnostics.py @@ -1143,7 +1143,7 @@ def vmec_compute_geometry(vs, s, theta, phi, phi_center=0): # temporary fix. Please see issue #238 and the discussion therein gbdrift = -1 * 2 * B_reference * L_reference * L_reference * sqrt_s[:, None, None] * B_cross_grad_B_dot_grad_alpha / (modB * modB * modB) * toroidal_flux_sign - gbdrift0 = B_cross_grad_B_dot_grad_psi * 2 * shat[:, None, None] / (modB * modB * modB * sqrt_s[:, None, None]) * toroidal_flux_sign + gbdrift0 = -B_cross_grad_B_dot_grad_psi * 2 * shat[:, None, None] / (modB * modB * modB * sqrt_s[:, None, None]) * toroidal_flux_sign # temporary fix. Please see issue #238 and the discussion therein cvdrift = gbdrift - 2 * B_reference * L_reference * L_reference * sqrt_s[:, None, None] * mu_0 * d_pressure_d_s[:, None, None] * toroidal_flux_sign / (edge_toroidal_flux_over_2pi * modB * modB) diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index a3a51443e..873b118cb 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -56,7 +56,8 @@ def jac(x: RealArray = None, *args, **kwargs): # log_file.write(f',{total_jacj:24.16e}') # Compute the derivative of each residual, as required by scipy.optimize.least_squares - log_file.write(", " + np.array_str(ret, max_line_width = np.inf, precision = None).replace('\n',',')) + with numpy.printoptions(threshold=numpy.inf): + log_file.write(", " + np.array_str(ret, max_line_width = np.inf, precision = None).replace('\n',',')) log_file.write('\n') log_file.flush() return ret From df3af0e5264b86000a4091f5bea9296617faaeaa Mon Sep 17 00:00:00 2001 From: daringli Date: Thu, 22 Dec 2022 21:22:48 +0100 Subject: [PATCH 05/14] Bugfix --- src/simsopt/solve/serial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index 873b118cb..2cd60d251 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -56,7 +56,7 @@ def jac(x: RealArray = None, *args, **kwargs): # log_file.write(f',{total_jacj:24.16e}') # Compute the derivative of each residual, as required by scipy.optimize.least_squares - with numpy.printoptions(threshold=numpy.inf): + with np.printoptions(threshold=np.inf): log_file.write(", " + np.array_str(ret, max_line_width = np.inf, precision = None).replace('\n',',')) log_file.write('\n') log_file.flush() From 1d11405e10314c0215a52b7355b1201ce0153224 Mon Sep 17 00:00:00 2001 From: daringli Date: Thu, 22 Dec 2022 23:24:34 +0100 Subject: [PATCH 06/14] Reverted jacobian to only print the derivative of the total objective. No longer useful for reseting the least squares optimizer, but that wasn't a good default due to the 1000s of residuals in the QS objective. --- src/simsopt/solve/serial.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index 2cd60d251..14e841ad6 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -50,16 +50,18 @@ def jac(x: RealArray = None, *args, **kwargs): # Compute the jacobian of the square-summed residuals # for scipy least-squares problem - #f = fd.fn() # function value at x0 - #total_jac = np.sum(2 * _jac * f[:, None],axis=0) - #for total_jacj in total_jac: - # log_file.write(f',{total_jacj:24.16e}') + f = fd.fn() # function value at x0 + total_jac = np.sum(2 * ret * f[:, None],axis=0) + for total_jacj in total_jac: + log_file.write(f',{total_jacj:24.16e}') # Compute the derivative of each residual, as required by scipy.optimize.least_squares - with np.printoptions(threshold=np.inf): - log_file.write(", " + np.array_str(ret, max_line_width = np.inf, precision = None).replace('\n',',')) - log_file.write('\n') - log_file.flush() + # We don't do this because the QS objective can have thousands of residuals + # TODO: maybe do it? + # with np.printoptions(threshold=np.inf): + # log_file.write(", " + np.array_str(ret, max_line_width = np.inf, precision = None).replace('\n',',')) + # log_file.write('\n') + # log_file.flush() return ret return jac From b795436c08ca808eb35ba7879ce9f46411d21464 Mon Sep 17 00:00:00 2001 From: daringli Date: Thu, 22 Dec 2022 23:29:23 +0100 Subject: [PATCH 07/14] Added missing newline --- src/simsopt/solve/serial.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index 14e841ad6..580c9241a 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -54,7 +54,8 @@ def jac(x: RealArray = None, *args, **kwargs): total_jac = np.sum(2 * ret * f[:, None],axis=0) for total_jacj in total_jac: log_file.write(f',{total_jacj:24.16e}') - + log_file.write('\n') + log_file.flush() # Compute the derivative of each residual, as required by scipy.optimize.least_squares # We don't do this because the QS objective can have thousands of residuals # TODO: maybe do it? From 4c2586e32c3334bdf67bde88056881dfa2f2c9e4 Mon Sep 17 00:00:00 2001 From: daringli Date: Thu, 22 Dec 2022 23:37:09 +0100 Subject: [PATCH 08/14] Added non-least-squares case --- src/simsopt/solve/serial.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index 580c9241a..731f915a2 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -51,7 +51,10 @@ def jac(x: RealArray = None, *args, **kwargs): # Compute the jacobian of the square-summed residuals # for scipy least-squares problem f = fd.fn() # function value at x0 - total_jac = np.sum(2 * ret * f[:, None],axis=0) + if problem_type == 'least_squares': + total_jac = np.sum(2 * ret * f[:, None],axis=0) + else: + total_jac = ret for total_jacj in total_jac: log_file.write(f',{total_jacj:24.16e}') log_file.write('\n') From 2f062d0ef9fbd19cd868d45f831b10d027202793 Mon Sep 17 00:00:00 2001 From: Stefan Buller Date: Tue, 10 Jan 2023 11:55:13 -0500 Subject: [PATCH 09/14] Updated jacobian logging to also support the legacy option (which is the default). --- src/simsopt/_core/finite_difference.py | 62 ++++++++++-------- src/simsopt/solve/mpi.py | 4 +- src/simsopt/solve/serial.py | 90 ++++++++++++++++++-------- 3 files changed, 99 insertions(+), 57 deletions(-) diff --git a/src/simsopt/_core/finite_difference.py b/src/simsopt/_core/finite_difference.py index e7c261867..5d72d6963 100644 --- a/src/simsopt/_core/finite_difference.py +++ b/src/simsopt/_core/finite_difference.py @@ -70,8 +70,10 @@ def __init__(self, func: Callable, self.log_header_written = False self.flatten_out = flatten_out - x0 = np.asarray(x0) if x0 is not None else x0 - self.x0 = x0 if x0 else self.opt.x + if x0 is not None: + self.x0 = np.asarray(x0) + else: + self.x0 = self.opt.x self.jac_size = None self.eval_cnt = 1 @@ -81,7 +83,8 @@ def __init__(self, func: Callable, else: # 1-sided differences self.nevals_jac = self.nparams + 1 - + self.xs = np.zeros((self.nparams, self.nevals_jac)) + def init_log(self): if isinstance(self.log_file, str): datestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S") @@ -114,14 +117,13 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: if self.diff_method == "centered": # Centered differences: for j in range(self.nparams): # len(x0) - x = np.copy(x0) - - x[j] = x0[j] + steps[j] - self.opt.x = x + self.xs[:, 2 * j] = x0[:] + self.xs[j, 2 * j] = x0[j] + steps[j] + self.opt.x = self.xs[:, 2 * j] fplus = np.asarray(self.fn()) - - x[j] = x0[j] - steps[j] - self.opt.x = x + self.xs[:, 2 * j + 1] = x0[:] + self.xs[j, 2 * j + 1] = x0[j] - steps[j] + self.opt.x = self.xs[j, 2 * j + 1] fminus = np.asarray(self.fn()) if self.flatten_out: @@ -131,12 +133,13 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: elif self.diff_method == "forward": # 1-sided differences - self.opt.x = x0 + self.xs[:, 0] = x0[:] + self.opt.x = self.xs[:, 0] f0 = np.asarray(self.fn()) for j in range(self.nparams): # len(x0) - x = np.copy(x0) - x[j] = x0[j] + steps[j] - self.opt.x = x + self.xs[:, j + 1] = x0[:] + self.xs[j, j + 1] = x0[j] + steps[j] + self.opt.x = self.xs[:, j + 1] fplus = np.asarray(self.fn()) if self.flatten_out: @@ -204,7 +207,8 @@ def __init__(self, func: Callable, else: # 1-sided differences self.nevals_jac = self.nparams + 1 - + self.xs = np.zeros((self.nparams, self.nevals_jac)) + def __enter__(self): self.mpi_apart() self.init_log() @@ -255,19 +259,19 @@ def _jac(self, x: RealArray = None): mpi.comm_leaders.Bcast(steps) diff_method = mpi.comm_leaders.bcast(self.diff_method) if diff_method == "centered": - xs = np.zeros((self.nparams, self.nevals_jac)) for j in range(self.nparams): - xs[:, 2 * j] = x0[:] # I don't think I need np.copy(), but not 100% sure. - xs[j, 2 * j] = x0[j] + steps[j] - xs[:, 2 * j + 1] = x0[:] - xs[j, 2 * j + 1] = x0[j] - steps[j] + self.xs[:, 2 * j] = x0[:] # I don't think I need np.copy(), but not 100% sure. + self.xs[j, 2 * j] = x0[j] + steps[j] + self.xs[:, 2 * j + 1] = x0[:] + self.xs[j, 2 * j + 1] = x0[j] - steps[j] else: # diff_method == "forward": # 1-sided differences - xs = np.zeros((self.nparams, self.nevals_jac)) - xs[:, 0] = x0[:] + self.xs[:, 0] = x0[:] for j in range(self.nparams): - xs[:, j + 1] = x0[:] - xs[j, j + 1] = x0[j] + steps[j] + self.xs[:, j + 1] = x0[:] + self.xs[j, j + 1] = x0[j] + steps[j] + print(self.xs) + print(self.xs.shape) evals = None # nvals = None # Work on this later @@ -285,7 +289,7 @@ def _jac(self, x: RealArray = None): # Handle only this group's share of the work: if np.mod(j, mpi.ngroups) == mpi.rank_leaders: mpi.mobilize_workers(ARB_VAL) - x = xs[:, j] + x = self.xs[:, j] mpi.comm_groups.bcast(x, root=0) opt.x = x out = np.asarray(self.fn()) @@ -320,7 +324,7 @@ def _jac(self, x: RealArray = None): # Weird things may happen if we do not reset the state vector # to x0: opt.x = x0 - return jac, xs, evals + return jac, evals def mpi_leaders_task(self, *args): """ @@ -393,9 +397,11 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: self.mpi.mobilize_leaders(ARB_VAL) # Any value not equal to STOP self.mpi.comm_leaders.bcast(x, root=0) - jac, xs, evals = self._jac(x) + jac, evals = self._jac(x) logger.debug(f'jac is {jac}') - + print("on master") + print(self.xs) + print(self.xs.shape) # Log file is now written externally # by a wrapper in the serial or mpi solver. diff --git a/src/simsopt/solve/mpi.py b/src/simsopt/solve/mpi.py index 77427de5b..36ef971e0 100644 --- a/src/simsopt/solve/mpi.py +++ b/src/simsopt/solve/mpi.py @@ -79,6 +79,7 @@ def least_squares_mpi_solve(prob: LeastSquaresProblem, abs_step: float = 1.0e-7, rel_step: float = 0.0, diff_method: str = "forward", + jac_verbose: str = "legacy", **kwargs): """ Solve a nonlinear-least-squares minimization problem using @@ -101,6 +102,7 @@ def least_squares_mpi_solve(prob: LeastSquaresProblem, "forward". If ``centered``, centered finite differences will be used. If ``forward``, one-sided finite differences will be used. Else, error is raised. + jac_verbose: Flag controlling how verbose the jacobian log is. kwargs: Any arguments to pass to `scipy.optimize.least_squares `_. For instance, you can supply ``max_nfev=100`` to set @@ -203,7 +205,7 @@ def _f_proc0(x): x0 = np.copy(prob.x) logger.info("Using finite difference method implemented in " "SIMSOPT for evaluating gradient") - jac = finite_difference_jac_wrapper(fd) + jac = finite_difference_jac_wrapper(fd, verbose=jac_verbose, comment="mpi") result = least_squares(_f_proc0, x0, jac=jac, verbose=2, **kwargs) diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index 731f915a2..cb054b4a3 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -27,45 +27,77 @@ __all__ = ['least_squares_serial_solve', 'serial_solve'] -def finite_difference_jac_wrapper(fd, problem_type = 'least_squares'): +def finite_difference_jac_wrapper(fd, problem_type = 'least_squares', verbose = "legacy", comment = ""): """Wrapper for the `jac` method of the `MPIFiniteDifference` and `FiniteDifference` classes. For logging the jacobian calculation when used with `scipy.optimize.least_squares`. - Also handles scipy.optimize.minimize, for now.""" + Also handles scipy.optimize.minimize. + verbose - controls the amount of information logged. + 'legacy': Only the points used by the jacobian calculation is logged. The default. + 'dfdx': The derivative of the scalar objective is logged. + 'dRdx': For least_squares problems. Logs the matrix dR_i/dx_j where R_i are the components of the residual vector. WARNING: can result in huge jacobian logs.""" + + if comment != "": + comment = " " + comment + + if verbose not in ["legacy", 'dfdx', 'dRdx']: + raise ValueError("Unrecognized verbose flag: '" + verbose + "' Recognized values are 'legacy', dfdx', 'dRdx'.") + if (verbose == 'dRdx') and (problem_type != 'least_squares'): + raise ValueError("Verbose flag 'dRdx' is only supported for problem_type 'lest_squares'.") + def jac(x: RealArray = None, *args, **kwargs): ret = fd.jac(x, *args, **kwargs) log_file = fd.log_file nparams = fd.nparams + # WRITE HEADER if not fd.log_header_written: - log_file.write(f'Problem type:\n{problem_type}\nnparams:\n{nparams}\n') + log_file.write(f'Problem type:\n{problem_type}{comment}\nnparams:\n{nparams}\n') log_file.write('function_evaluation, seconds') - if problem_type == 'least_squares': + if verbose == "dRdx": log_file.write(', d(residual_j)/d(x_i)') - else: + elif verbose == "dfdx": log_file.write(', d(f)/d(x_i)') + elif verbose == "legacy": + for j in range(nparams): + log_file.write(f', x({j})') log_file.write('\n') fd.log_header_written = True - del_t = time() - fd.start_time - j_eval = fd.eval_cnt//fd.nevals_jac - log_file.write(f'{j_eval:6d},{del_t:12.4e}') + # WRITE DATA + if verbose == "dfdx": + del_t = time() - fd.start_time + j_eval = fd.eval_cnt//fd.nevals_jac + log_file.write(f'{j_eval:6d},{del_t:12.4e}') - # Compute the jacobian of the square-summed residuals - # for scipy least-squares problem - f = fd.fn() # function value at x0 - if problem_type == 'least_squares': - total_jac = np.sum(2 * ret * f[:, None],axis=0) - else: - total_jac = ret - for total_jacj in total_jac: - log_file.write(f',{total_jacj:24.16e}') - log_file.write('\n') - log_file.flush() - # Compute the derivative of each residual, as required by scipy.optimize.least_squares - # We don't do this because the QS objective can have thousands of residuals - # TODO: maybe do it? - # with np.printoptions(threshold=np.inf): - # log_file.write(", " + np.array_str(ret, max_line_width = np.inf, precision = None).replace('\n',',')) - # log_file.write('\n') - # log_file.flush() + if problem_type == 'least_squares': + f = fd.fn() # function value at x0 + total_jac = np.sum(2 * ret * f[:, None],axis=0) + else: + total_jac = ret + for total_jacj in total_jac: + log_file.write(f',{total_jacj:24.16e}') + log_file.write('\n') + log_file.flush() + elif verbose == "dRdx": + # only least squares can use verbose = 'dRdx' + del_t = time() - fd.start_time + j_eval = fd.eval_cnt//fd.nevals_jac + log_file.write(f'{j_eval:6d},{del_t:12.4e}') + with np.printoptions(threshold=np.inf): + log_file.write(", " + np.array_str(ret, max_line_width = np.inf, precision = None).replace('\n',',')) + log_file.write('\n') + log_file.flush() + elif verbose == "legacy": + for j in range(fd.nevals_jac): + del_t = time() - fd.start_time + j_eval = j + fd.eval_cnt - fd.nevals_jac - 1 + print(fd.eval_cnt) + print("^^^") + log_file.write(f'{j_eval:6d},{del_t:12.4e}') + print(fd.xs[:, j]) + for xj in fd.xs[:, j]: + print(j) + log_file.write(f',{xj:24.16e}') + log_file.write('\n') + log_file.flush() return ret return jac @@ -76,6 +108,7 @@ def least_squares_serial_solve(prob: LeastSquaresProblem, abs_step: float = 1.0e-7, rel_step: float = 0.0, diff_method: str = "forward", + jac_verbose: str = "legacy", **kwargs): """ Solve a nonlinear-least-squares minimization problem using @@ -190,7 +223,7 @@ def objective(x): fd = FiniteDifference(prob.residuals, abs_step=abs_step, rel_step=rel_step, diff_method=diff_method) logger.info("Using derivatives") - jac = finite_difference_jac_wrapper(fd) + jac = finite_difference_jac_wrapper(fd, verbose=jac_verbose, comment="serial") result = least_squares(objective, x0, verbose=2, jac=jac, **kwargs) else: logger.info("Using derivative-free method") @@ -209,6 +242,7 @@ def serial_solve(prob: Union[Optimizable, Callable], abs_step: float = 1.0e-7, rel_step: float = 0.0, diff_method: str = "centered", + jac_verbose: str = "legacy", **kwargs): """ Solve a general minimization problem (i.e. one that need not be of @@ -291,7 +325,7 @@ def objective(x): logger.info("Using derivatives") fd = FiniteDifference(prob.J, abs_step=abs_step, rel_step=rel_step, diff_method=diff_method, flatten_out = True) - jac = finite_difference_jac_wrapper(fd, problem_type='minimize') + jac = finite_difference_jac_wrapper(fd, problem_type='minimize', verbose = jac_verbose, comment="serial") result = minimize(objective, x0, options={'disp': True}, jac=jac, **kwargs) else: From e84a0e185463945c263c073b44bdd73680e4d0c9 Mon Sep 17 00:00:00 2001 From: Stefan Buller Date: Wed, 11 Jan 2023 14:05:14 -0500 Subject: [PATCH 10/14] f0 is now calculated inside the finite difference object to avoid duplicate work and because calling fd.fn() in the logger resulted in the execution freezing on cobra. --- src/simsopt/_core/finite_difference.py | 20 +++++++++++--------- src/simsopt/solve/serial.py | 6 +----- 2 files changed, 12 insertions(+), 14 deletions(-) diff --git a/src/simsopt/_core/finite_difference.py b/src/simsopt/_core/finite_difference.py index 5d72d6963..9320bd1d7 100644 --- a/src/simsopt/_core/finite_difference.py +++ b/src/simsopt/_core/finite_difference.py @@ -116,6 +116,8 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: rel_step=self.rel_step) if self.diff_method == "centered": # Centered differences: + self.f0 = np.zeros(self.jac_size[0]) + for j in range(self.nparams): # len(x0) self.xs[:, 2 * j] = x0[:] self.xs[j, 2 * j] = x0[j] + steps[j] @@ -130,12 +132,14 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: jac[j] = (fplus - fminus) / (2 * steps[j]) else: jac[:, j] = (fplus - fminus) / (2 * steps[j]) + self.f0 = self.f0 + (fplus + fminus)/2 + self.f0 = self.f0 / self.nparams elif self.diff_method == "forward": # 1-sided differences self.xs[:, 0] = x0[:] self.opt.x = self.xs[:, 0] - f0 = np.asarray(self.fn()) + self.f0 = np.asarray(self.fn()) for j in range(self.nparams): # len(x0) self.xs[:, j + 1] = x0[:] self.xs[j, j + 1] = x0[j] + steps[j] @@ -143,9 +147,9 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: fplus = np.asarray(self.fn()) if self.flatten_out: - jac[j] = (fplus - f0) / steps[j] + jac[j] = (fplus - self.f0) / steps[j] else: - jac[:, j] = (fplus - f0) / steps[j] + jac[:, j] = (fplus - self.f0) / steps[j] self.eval_cnt += self.nevals_jac # Set the opt.x to the original x @@ -270,8 +274,6 @@ def _jac(self, x: RealArray = None): for j in range(self.nparams): self.xs[:, j + 1] = x0[:] self.xs[j, j + 1] = x0[j] + steps[j] - print(self.xs) - print(self.xs.shape) evals = None # nvals = None # Work on this later @@ -316,11 +318,14 @@ def _jac(self, x: RealArray = None): for j in range(self.nparams): jac[:, j] = (evals[:, 2 * j] - evals[:, 2 * j + 1]) / ( 2 * steps[j]) + # approximate f0 as the average of self.fn at stencil points + self.f0 = np.sum(evals,axis=1)/self.nevals_jac else: # diff_method == "forward": # 1-sided differences: for j in range(self.nparams): jac[:, j] = (evals[:, j + 1] - evals[:, 0]) / steps[j] - + self.f0 = evals[:, 0] + # Weird things may happen if we do not reset the state vector # to x0: opt.x = x0 @@ -399,9 +404,6 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: jac, evals = self._jac(x) logger.debug(f'jac is {jac}') - print("on master") - print(self.xs) - print(self.xs.shape) # Log file is now written externally # by a wrapper in the serial or mpi solver. diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index cb054b4a3..5f580b364 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -68,7 +68,7 @@ def jac(x: RealArray = None, *args, **kwargs): log_file.write(f'{j_eval:6d},{del_t:12.4e}') if problem_type == 'least_squares': - f = fd.fn() # function value at x0 + f = fd.f0 # function value at x0 total_jac = np.sum(2 * ret * f[:, None],axis=0) else: total_jac = ret @@ -89,12 +89,8 @@ def jac(x: RealArray = None, *args, **kwargs): for j in range(fd.nevals_jac): del_t = time() - fd.start_time j_eval = j + fd.eval_cnt - fd.nevals_jac - 1 - print(fd.eval_cnt) - print("^^^") log_file.write(f'{j_eval:6d},{del_t:12.4e}') - print(fd.xs[:, j]) for xj in fd.xs[:, j]: - print(j) log_file.write(f',{xj:24.16e}') log_file.write('\n') log_file.flush() From f023bc9df31816e43ebe71943c821c383e60d087 Mon Sep 17 00:00:00 2001 From: Stefan Buller Date: Wed, 11 Jan 2023 14:57:04 -0500 Subject: [PATCH 11/14] Fixed linting errors. --- src/simsopt/_core/finite_difference.py | 27 +++++++++++------------ src/simsopt/field/magneticfieldclasses.py | 2 +- src/simsopt/solve/mpi.py | 2 +- src/simsopt/solve/serial.py | 22 +++++++++--------- tests/field/test_magneticfields.py | 4 ++-- tests/objectives/test_utilities.py | 3 ++- 6 files changed, 30 insertions(+), 30 deletions(-) diff --git a/src/simsopt/_core/finite_difference.py b/src/simsopt/_core/finite_difference.py index 9320bd1d7..b5012ed81 100644 --- a/src/simsopt/_core/finite_difference.py +++ b/src/simsopt/_core/finite_difference.py @@ -84,7 +84,7 @@ def __init__(self, func: Callable, # 1-sided differences self.nevals_jac = self.nparams + 1 self.xs = np.zeros((self.nparams, self.nevals_jac)) - + def init_log(self): if isinstance(self.log_file, str): datestr = datetime.now().strftime("%Y-%m-%d-%H-%M-%S") @@ -93,15 +93,14 @@ def init_log(self): self.new_log_file = True self.start_time = time() - def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: if x is not None: self.x0 = np.asarray(x) x0 = self.x0 if self.x0 is not None else self.opt.x opt_x0 = self.opt.x - + self.init_log() - + if self.jac_size is None: out = self.fn() if not isinstance(out, (np.ndarray, collections.abc.Sequence)): @@ -110,15 +109,15 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: self.jac_size = (self.nparams) else: self.jac_size = (len(out), self.nparams) - + jac = np.zeros(self.jac_size) steps = finite_difference_steps(x0, abs_step=self.abs_step, rel_step=self.rel_step) if self.diff_method == "centered": # Centered differences: self.f0 = np.zeros(self.jac_size[0]) - - for j in range(self.nparams): # len(x0) + + for j in range(self.nparams): # len(x0) self.xs[:, 2 * j] = x0[:] self.xs[j, 2 * j] = x0[j] + steps[j] self.opt.x = self.xs[:, 2 * j] @@ -134,13 +133,13 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: jac[:, j] = (fplus - fminus) / (2 * steps[j]) self.f0 = self.f0 + (fplus + fminus)/2 self.f0 = self.f0 / self.nparams - + elif self.diff_method == "forward": # 1-sided differences self.xs[:, 0] = x0[:] self.opt.x = self.xs[:, 0] self.f0 = np.asarray(self.fn()) - for j in range(self.nparams): # len(x0) + for j in range(self.nparams): # len(x0) self.xs[:, j + 1] = x0[:] self.xs[j, j + 1] = x0[j] + steps[j] self.opt.x = self.xs[:, j + 1] @@ -150,7 +149,7 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: jac[j] = (fplus - self.f0) / steps[j] else: jac[:, j] = (fplus - self.f0) / steps[j] - + self.eval_cnt += self.nevals_jac # Set the opt.x to the original x self.opt.x = opt_x0 @@ -212,7 +211,7 @@ def __init__(self, func: Callable, # 1-sided differences self.nevals_jac = self.nparams + 1 self.xs = np.zeros((self.nparams, self.nevals_jac)) - + def __enter__(self): self.mpi_apart() self.init_log() @@ -235,7 +234,7 @@ def __exit__(self, exc_type, exc_value, tb): self.mpi.together() if self.mpi.proc0_world and self.new_log_file: self.log_file.close() - + # Called by MPI leaders def _jac(self, x: RealArray = None): # Use shortcuts for class variables @@ -319,13 +318,13 @@ def _jac(self, x: RealArray = None): jac[:, j] = (evals[:, 2 * j] - evals[:, 2 * j + 1]) / ( 2 * steps[j]) # approximate f0 as the average of self.fn at stencil points - self.f0 = np.sum(evals,axis=1)/self.nevals_jac + self.f0 = np.sum(evals, axis=1)/self.nevals_jac else: # diff_method == "forward": # 1-sided differences: for j in range(self.nparams): jac[:, j] = (evals[:, j + 1] - evals[:, 0]) / steps[j] self.f0 = evals[:, 0] - + # Weird things may happen if we do not reset the state vector # to x0: opt.x = x0 diff --git a/src/simsopt/field/magneticfieldclasses.py b/src/simsopt/field/magneticfieldclasses.py index 64d71e704..9f34fa75a 100644 --- a/src/simsopt/field/magneticfieldclasses.py +++ b/src/simsopt/field/magneticfieldclasses.py @@ -481,7 +481,7 @@ def _A_impl(self, A): k = np.sqrt(1-np.divide(np.square(alpha), np.square(beta))) ellipek2 = ellipe(k**2) ellipkk2 = ellipk(k**2) - + num = (2*self.r0+np.sqrt(points[:, 0]**2+points[:, 1]**2)*ellipek2+(self.r0**2+points[:, 0]**2+points[:, 1]**2+points[:, 2]**2)*(ellipe(k**2)-ellipkk2)) denom = ((points[:, 0]**2+points[:, 1]**2+1e-31)*np.sqrt(self.r0**2+points[:, 0]**2+points[:, 1]**2+2*self.r0*np.sqrt(points[:, 0]**2+points[:, 1]**2)+points[:, 2]**2+1e-31)) fak = num/denom diff --git a/src/simsopt/solve/mpi.py b/src/simsopt/solve/mpi.py index 36ef971e0..dd78b0aa5 100644 --- a/src/simsopt/solve/mpi.py +++ b/src/simsopt/solve/mpi.py @@ -22,7 +22,7 @@ MPI = None from .serial import finite_difference_jac_wrapper - + from .._core.types import RealArray from .._core.optimizable import Optimizable from ..util.mpi import MpiPartition diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index 5f580b364..9e3aa5040 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -27,7 +27,7 @@ __all__ = ['least_squares_serial_solve', 'serial_solve'] -def finite_difference_jac_wrapper(fd, problem_type = 'least_squares', verbose = "legacy", comment = ""): +def finite_difference_jac_wrapper(fd, problem_type='least_squares', verbose="legacy", comment=""): """Wrapper for the `jac` method of the `MPIFiniteDifference` and `FiniteDifference` classes. For logging the jacobian calculation when used with `scipy.optimize.least_squares`. Also handles scipy.optimize.minimize. @@ -38,12 +38,12 @@ def finite_difference_jac_wrapper(fd, problem_type = 'least_squares', verbose = if comment != "": comment = " " + comment - + if verbose not in ["legacy", 'dfdx', 'dRdx']: raise ValueError("Unrecognized verbose flag: '" + verbose + "' Recognized values are 'legacy', dfdx', 'dRdx'.") if (verbose == 'dRdx') and (problem_type != 'least_squares'): raise ValueError("Verbose flag 'dRdx' is only supported for problem_type 'lest_squares'.") - + def jac(x: RealArray = None, *args, **kwargs): ret = fd.jac(x, *args, **kwargs) log_file = fd.log_file @@ -66,10 +66,10 @@ def jac(x: RealArray = None, *args, **kwargs): del_t = time() - fd.start_time j_eval = fd.eval_cnt//fd.nevals_jac log_file.write(f'{j_eval:6d},{del_t:12.4e}') - + if problem_type == 'least_squares': - f = fd.f0 # function value at x0 - total_jac = np.sum(2 * ret * f[:, None],axis=0) + f = fd.f0 # function value at x0 + total_jac = np.sum(2 * ret * f[:, None], axis=0) else: total_jac = ret for total_jacj in total_jac: @@ -82,7 +82,7 @@ def jac(x: RealArray = None, *args, **kwargs): j_eval = fd.eval_cnt//fd.nevals_jac log_file.write(f'{j_eval:6d},{del_t:12.4e}') with np.printoptions(threshold=np.inf): - log_file.write(", " + np.array_str(ret, max_line_width = np.inf, precision = None).replace('\n',',')) + log_file.write(", " + np.array_str(ret, max_line_width=np.inf, precision=None).replace('\n', ',')) log_file.write('\n') log_file.flush() elif verbose == "legacy": @@ -95,7 +95,7 @@ def jac(x: RealArray = None, *args, **kwargs): log_file.write('\n') log_file.flush() return ret - + return jac @@ -320,10 +320,10 @@ def objective(x): if grad: logger.info("Using derivatives") fd = FiniteDifference(prob.J, abs_step=abs_step, - rel_step=rel_step, diff_method=diff_method, flatten_out = True) - jac = finite_difference_jac_wrapper(fd, problem_type='minimize', verbose = jac_verbose, comment="serial") + rel_step=rel_step, diff_method=diff_method, flatten_out=True) + jac = finite_difference_jac_wrapper(fd, problem_type='minimize', verbose=jac_verbose, comment="serial") result = minimize(objective, x0, options={'disp': True}, jac=jac, - **kwargs) + **kwargs) else: logger.info("Using derivative-free method") result = minimize(objective, x0, options={'disp': True}, **kwargs) diff --git a/tests/field/test_magneticfields.py b/tests/field/test_magneticfields.py index 919f535e5..e4bac9ed5 100644 --- a/tests/field/test_magneticfields.py +++ b/tests/field/test_magneticfields.py @@ -265,12 +265,12 @@ def test_circularcoil_Bfield(self): assert np.allclose(Bfield.dB_by_dX(), Bcircular.dB_by_dX()) assert np.allclose(dB1_by_dX[:, 0, 0]+dB1_by_dX[:, 1, 1]+dB1_by_dX[:, 2, 2], np.zeros((npoints))) # divergence assert np.allclose(dB1_by_dX, transpGradB1) # symmetry of the gradient - + # one points Bfield.set_points(np.asarray([[0.1, 0.2, 0.3]])) Afield = Bfield.A() assert np.allclose(Afield, [[0, 5.15785, -2.643056]]) - + # two points Bfield.set_points(np.asarray([[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]])) Afield = Bfield.A() diff --git a/tests/objectives/test_utilities.py b/tests/objectives/test_utilities.py index b3c0fe407..83c29ea06 100644 --- a/tests/objectives/test_utilities.py +++ b/tests/objectives/test_utilities.py @@ -60,7 +60,7 @@ def test_quadratic_penalty(self): J = CurveLength(curve) self.subtest_quadratic_penalty(curve, J.J()+0.1) self.subtest_quadratic_penalty(curve, J.J()-0.1) - + def test_mpi_objective(self): if MPI is None: print("skip test_mpi_objective") @@ -85,5 +85,6 @@ def test_mpi_objective(self): assert abs(Jmpi1.J() - sum(J.J() for J in Js)/n) < 1e-14 assert np.sum(np.abs(Jmpi1.dJ() - sum(J.dJ() for J in Js)/n)) < 1e-14 + if __name__ == "__main__": unittest.main() From 64a0d2be8e3c8a22635f45b74c8a5ca5fcbca3ad Mon Sep 17 00:00:00 2001 From: Stefan Buller Date: Wed, 11 Jan 2023 15:29:19 -0500 Subject: [PATCH 12/14] Moved some code around. Renamed 'wrapper' to decorator to conform to python nomenclature. --- src/simsopt/_core/finite_difference.py | 74 ++++++++++++++++++++++++- src/simsopt/solve/mpi.py | 5 +- src/simsopt/solve/serial.py | 77 +------------------------- 3 files changed, 78 insertions(+), 78 deletions(-) diff --git a/src/simsopt/_core/finite_difference.py b/src/simsopt/_core/finite_difference.py index b5012ed81..3fef1c12c 100644 --- a/src/simsopt/_core/finite_difference.py +++ b/src/simsopt/_core/finite_difference.py @@ -32,7 +32,79 @@ logger = logging.getLogger(__name__) -__all__ = ['FiniteDifference'] +__all__ = ['FiniteDifference', 'finite_difference_jac_decorator'] + + +def finite_difference_jac_decorator(fd, problem_type='least_squares', verbose="legacy", comment=""): + """Wrapper for the `jac` method of the `MPIFiniteDifference` and `FiniteDifference` classes. + For logging the jacobian calculation when used with `scipy.optimize.least_squares`. + Also handles scipy.optimize.minimize. + verbose - controls the amount of information logged. + 'legacy': Only the points used by the jacobian calculation is logged. The default. + 'dfdx': The derivative of the scalar objective is logged. + 'dRdx': For least_squares problems. Logs the matrix dR_i/dx_j where R_i are the components of the residual vector. WARNING: can result in huge jacobian logs.""" + + if comment != "": + comment = " " + comment + + if verbose not in ["legacy", 'dfdx', 'dRdx']: + raise ValueError("Unrecognized verbose flag: '" + verbose + "' Recognized values are 'legacy', dfdx', 'dRdx'.") + if (verbose == 'dRdx') and (problem_type != 'least_squares'): + raise ValueError("Verbose flag 'dRdx' is only supported for problem_type 'lest_squares'.") + + def wrapper(x: RealArray = None, *args, **kwargs): + ret = fd.jac(x, *args, **kwargs) + log_file = fd.log_file + nparams = fd.nparams + # WRITE HEADER + if not fd.log_header_written: + log_file.write(f'Problem type:\n{problem_type}{comment}\nnparams:\n{nparams}\n') + log_file.write('function_evaluation, seconds') + if verbose == "dRdx": + log_file.write(', d(residual_j)/d(x_i)') + elif verbose == "dfdx": + log_file.write(', d(f)/d(x_i)') + elif verbose == "legacy": + for j in range(nparams): + log_file.write(f', x({j})') + log_file.write('\n') + fd.log_header_written = True + # WRITE DATA + if verbose == "dfdx": + del_t = time() - fd.start_time + j_eval = fd.eval_cnt//fd.nevals_jac + log_file.write(f'{j_eval:6d},{del_t:12.4e}') + + if problem_type == 'least_squares': + f = fd.f0 # function value at x0 + total_jac = np.sum(2 * ret * f[:, None], axis=0) + else: + total_jac = ret + for total_jacj in total_jac: + log_file.write(f',{total_jacj:24.16e}') + log_file.write('\n') + log_file.flush() + elif verbose == "dRdx": + # only least squares can use verbose = 'dRdx' + del_t = time() - fd.start_time + j_eval = fd.eval_cnt//fd.nevals_jac + log_file.write(f'{j_eval:6d},{del_t:12.4e}') + with np.printoptions(threshold=np.inf): + log_file.write(", " + np.array_str(ret, max_line_width=np.inf, precision=None).replace('\n', ',')) + log_file.write('\n') + log_file.flush() + elif verbose == "legacy": + for j in range(fd.nevals_jac): + del_t = time() - fd.start_time + j_eval = j + fd.eval_cnt - fd.nevals_jac - 1 + log_file.write(f'{j_eval:6d},{del_t:12.4e}') + for xj in fd.xs[:, j]: + log_file.write(f',{xj:24.16e}') + log_file.write('\n') + log_file.flush() + return ret + + return wrapper class FiniteDifference: diff --git a/src/simsopt/solve/mpi.py b/src/simsopt/solve/mpi.py index dd78b0aa5..5cbc829df 100644 --- a/src/simsopt/solve/mpi.py +++ b/src/simsopt/solve/mpi.py @@ -21,12 +21,11 @@ except ImportError as err: MPI = None -from .serial import finite_difference_jac_wrapper - from .._core.types import RealArray from .._core.optimizable import Optimizable from ..util.mpi import MpiPartition from .._core.finite_difference import MPIFiniteDifference +from .._core.finite_difference import finite_difference_jac_decorator from ..objectives.least_squares import LeastSquaresProblem logger = logging.getLogger(__name__) @@ -205,7 +204,7 @@ def _f_proc0(x): x0 = np.copy(prob.x) logger.info("Using finite difference method implemented in " "SIMSOPT for evaluating gradient") - jac = finite_difference_jac_wrapper(fd, verbose=jac_verbose, comment="mpi") + jac = finite_difference_jac_decorator(fd, verbose=jac_verbose, comment="mpi") result = least_squares(_f_proc0, x0, jac=jac, verbose=2, **kwargs) diff --git a/src/simsopt/solve/serial.py b/src/simsopt/solve/serial.py index 9e3aa5040..1499158e4 100644 --- a/src/simsopt/solve/serial.py +++ b/src/simsopt/solve/serial.py @@ -20,6 +20,7 @@ from ..objectives.least_squares import LeastSquaresProblem from .._core.optimizable import Optimizable from .._core.finite_difference import FiniteDifference +from .._core.finite_difference import finite_difference_jac_decorator logger = logging.getLogger(__name__) @@ -27,78 +28,6 @@ __all__ = ['least_squares_serial_solve', 'serial_solve'] -def finite_difference_jac_wrapper(fd, problem_type='least_squares', verbose="legacy", comment=""): - """Wrapper for the `jac` method of the `MPIFiniteDifference` and `FiniteDifference` classes. - For logging the jacobian calculation when used with `scipy.optimize.least_squares`. - Also handles scipy.optimize.minimize. - verbose - controls the amount of information logged. - 'legacy': Only the points used by the jacobian calculation is logged. The default. - 'dfdx': The derivative of the scalar objective is logged. - 'dRdx': For least_squares problems. Logs the matrix dR_i/dx_j where R_i are the components of the residual vector. WARNING: can result in huge jacobian logs.""" - - if comment != "": - comment = " " + comment - - if verbose not in ["legacy", 'dfdx', 'dRdx']: - raise ValueError("Unrecognized verbose flag: '" + verbose + "' Recognized values are 'legacy', dfdx', 'dRdx'.") - if (verbose == 'dRdx') and (problem_type != 'least_squares'): - raise ValueError("Verbose flag 'dRdx' is only supported for problem_type 'lest_squares'.") - - def jac(x: RealArray = None, *args, **kwargs): - ret = fd.jac(x, *args, **kwargs) - log_file = fd.log_file - nparams = fd.nparams - # WRITE HEADER - if not fd.log_header_written: - log_file.write(f'Problem type:\n{problem_type}{comment}\nnparams:\n{nparams}\n') - log_file.write('function_evaluation, seconds') - if verbose == "dRdx": - log_file.write(', d(residual_j)/d(x_i)') - elif verbose == "dfdx": - log_file.write(', d(f)/d(x_i)') - elif verbose == "legacy": - for j in range(nparams): - log_file.write(f', x({j})') - log_file.write('\n') - fd.log_header_written = True - # WRITE DATA - if verbose == "dfdx": - del_t = time() - fd.start_time - j_eval = fd.eval_cnt//fd.nevals_jac - log_file.write(f'{j_eval:6d},{del_t:12.4e}') - - if problem_type == 'least_squares': - f = fd.f0 # function value at x0 - total_jac = np.sum(2 * ret * f[:, None], axis=0) - else: - total_jac = ret - for total_jacj in total_jac: - log_file.write(f',{total_jacj:24.16e}') - log_file.write('\n') - log_file.flush() - elif verbose == "dRdx": - # only least squares can use verbose = 'dRdx' - del_t = time() - fd.start_time - j_eval = fd.eval_cnt//fd.nevals_jac - log_file.write(f'{j_eval:6d},{del_t:12.4e}') - with np.printoptions(threshold=np.inf): - log_file.write(", " + np.array_str(ret, max_line_width=np.inf, precision=None).replace('\n', ',')) - log_file.write('\n') - log_file.flush() - elif verbose == "legacy": - for j in range(fd.nevals_jac): - del_t = time() - fd.start_time - j_eval = j + fd.eval_cnt - fd.nevals_jac - 1 - log_file.write(f'{j_eval:6d},{del_t:12.4e}') - for xj in fd.xs[:, j]: - log_file.write(f',{xj:24.16e}') - log_file.write('\n') - log_file.flush() - return ret - - return jac - - def least_squares_serial_solve(prob: LeastSquaresProblem, grad: bool = None, abs_step: float = 1.0e-7, @@ -219,7 +148,7 @@ def objective(x): fd = FiniteDifference(prob.residuals, abs_step=abs_step, rel_step=rel_step, diff_method=diff_method) logger.info("Using derivatives") - jac = finite_difference_jac_wrapper(fd, verbose=jac_verbose, comment="serial") + jac = finite_difference_jac_decorator(fd, verbose=jac_verbose, comment="serial") result = least_squares(objective, x0, verbose=2, jac=jac, **kwargs) else: logger.info("Using derivative-free method") @@ -321,7 +250,7 @@ def objective(x): logger.info("Using derivatives") fd = FiniteDifference(prob.J, abs_step=abs_step, rel_step=rel_step, diff_method=diff_method, flatten_out=True) - jac = finite_difference_jac_wrapper(fd, problem_type='minimize', verbose=jac_verbose, comment="serial") + jac = finite_difference_jac_decorator(fd, problem_type='minimize', verbose=jac_verbose, comment="serial") result = minimize(objective, x0, options={'disp': True}, jac=jac, **kwargs) else: From 43ccddafa7e338ae732a8c111e87137514928a3d Mon Sep 17 00:00:00 2001 From: Stefan Buller Date: Wed, 11 Jan 2023 17:32:03 -0500 Subject: [PATCH 13/14] Fixed an error where the FiniteDifference object would not update its DOFs if the underlying optimizable got updated. --- src/simsopt/_core/finite_difference.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/simsopt/_core/finite_difference.py b/src/simsopt/_core/finite_difference.py index 3fef1c12c..b3cee4145 100644 --- a/src/simsopt/_core/finite_difference.py +++ b/src/simsopt/_core/finite_difference.py @@ -145,7 +145,7 @@ def __init__(self, func: Callable, if x0 is not None: self.x0 = np.asarray(x0) else: - self.x0 = self.opt.x + self.x0 = x0 self.jac_size = None self.eval_cnt = 1 @@ -196,7 +196,7 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: fplus = np.asarray(self.fn()) self.xs[:, 2 * j + 1] = x0[:] self.xs[j, 2 * j + 1] = x0[j] - steps[j] - self.opt.x = self.xs[j, 2 * j + 1] + self.opt.x = self.xs[:, 2 * j + 1] fminus = np.asarray(self.fn()) if self.flatten_out: @@ -225,7 +225,6 @@ def jac(self, x: RealArray = None, *args, **kwargs) -> RealArray: self.eval_cnt += self.nevals_jac # Set the opt.x to the original x self.opt.x = opt_x0 - return jac From f52f8b84687850d13fbfbcf38a5fa91b41471c13 Mon Sep 17 00:00:00 2001 From: Stefan Buller Date: Wed, 11 Jan 2023 18:42:44 -0500 Subject: [PATCH 14/14] Reverted some changes not related to the feature of this branch. --- src/simsopt/mhd/vmec_diagnostics.py | 2 +- tests/core/test_finite_difference.py | 4 ++++ tests/objectives/test_utilities.py | 2 -- 3 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/simsopt/mhd/vmec_diagnostics.py b/src/simsopt/mhd/vmec_diagnostics.py index 389be5093..7a1d69d7a 100644 --- a/src/simsopt/mhd/vmec_diagnostics.py +++ b/src/simsopt/mhd/vmec_diagnostics.py @@ -1143,7 +1143,7 @@ def vmec_compute_geometry(vs, s, theta, phi, phi_center=0): # temporary fix. Please see issue #238 and the discussion therein gbdrift = -1 * 2 * B_reference * L_reference * L_reference * sqrt_s[:, None, None] * B_cross_grad_B_dot_grad_alpha / (modB * modB * modB) * toroidal_flux_sign - gbdrift0 = -B_cross_grad_B_dot_grad_psi * 2 * shat[:, None, None] / (modB * modB * modB * sqrt_s[:, None, None]) * toroidal_flux_sign + gbdrift0 = B_cross_grad_B_dot_grad_psi * 2 * shat[:, None, None] / (modB * modB * modB * sqrt_s[:, None, None]) * toroidal_flux_sign # temporary fix. Please see issue #238 and the discussion therein cvdrift = gbdrift - 2 * B_reference * L_reference * L_reference * sqrt_s[:, None, None] * mu_0 * d_pressure_d_s[:, None, None] * toroidal_flux_sign / (edge_toroidal_flux_over_2pi * modB * modB) diff --git a/tests/core/test_finite_difference.py b/tests/core/test_finite_difference.py index d7c284b81..cd73e2e22 100755 --- a/tests/core/test_finite_difference.py +++ b/tests/core/test_finite_difference.py @@ -253,3 +253,7 @@ def test_jac_mpi(self): mpi.together() if mpi.proc0_world: fd.log_file.close() + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/objectives/test_utilities.py b/tests/objectives/test_utilities.py index 83c29ea06..b824f4468 100644 --- a/tests/objectives/test_utilities.py +++ b/tests/objectives/test_utilities.py @@ -3,8 +3,6 @@ import numpy as np -from simsopt._core import Optimizable -from simsopt._core.derivative import derivative_dec from simsopt.geo.curvexyzfourier import CurveXYZFourier from simsopt.geo.curveobjectives import CurveLength, LpCurveCurvature, LpCurveTorsion from simsopt.objectives.utilities import MPIObjective, QuadraticPenalty