Rewriten finite difference logging

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:
@@ -46,7 +118,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,56 +137,94 @@ 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
+        if x0 is not None:
+            self.x0 = np.asarray(x0)
+        else:
+            self.x0 = x0
 
         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
+        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")
+            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":
             # Centered differences:
-            for j in range(len(x0)):
-                x = np.copy(x0)
+            self.f0 = np.zeros(self.jac_size[0])
 
-                x[j] = x0[j] + steps[j]
-                self.opt.x = x
+            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]
                 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[:, 2 * j + 1]
                 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])
+                self.f0 = self.f0 + (fplus + fminus)/2
+            self.f0 = self.f0 / self.nparams
 
         elif self.diff_method == "forward":
             # 1-sided differences
-            self.opt.x = x0
-            f0 = np.asarray(self.fn())
-            for j in range(len(x0)):
-                x = np.copy(x0)
-                x[j] = x0[j] + steps[j]
-                self.opt.x = x
+            self.xs[:, 0] = x0[:]
+            self.opt.x = self.xs[:, 0]
+            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]
+                self.opt.x = self.xs[:, j + 1]
                 fplus = np.asarray(self.fn())
 
-                jac[:, j] = (fplus - f0) / steps[j]
+                if self.flatten_out:
+                    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
-
         return jac
 
 
@@ -162,6 +274,14 @@ 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
+        self.xs = np.zeros((self.nparams, self.nevals_jac))
 
     def __enter__(self):
         self.mpi_apart()
@@ -202,10 +322,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,46 +333,42 @@ 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):
-                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]
+            for j in range(self.nparams):
+                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
-            nevals_jac = nparams + 1
-            xs = np.zeros((nparams, nevals_jac))
-            xs[:, 0] = x0[:]
-            for j in range(nparams):
-                xs[:, j + 1] = x0[:]
-                xs[j, j + 1] = x0[j] + steps[j]
+            self.xs[:, 0] = x0[:]
+            for j in range(self.nparams):
+                self.xs[:, j + 1] = x0[:]
+                self.xs[j, j + 1] = x0[j] + steps[j]
 
         evals = 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)
-                x = xs[:, j]
+                x = self.xs[:, j]
                 mpi.comm_groups.bcast(x, root=0)
                 opt.x = x
                 out = np.asarray(self.fn())
 
                 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,18 +385,21 @@ 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])
+            # 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(nparams):
+            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
-        return jac, xs, evals
+        return jac, evals
 
     def mpi_leaders_task(self, *args):
         """
@@ -329,7 +447,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
         """
@@ -354,28 +472,11 @@ def jac(self, x: RealArray = None, *args, **kwargs):
         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}')
+        # 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

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