code: """Complex Step derivative approximations."""
  from __future__ import division, print_function
  
  from itertools import groupby
  from six.moves import range
  
  import numpy as np
  
  from openmdao.approximation_schemes.approximation_scheme import ApproximationScheme
  from openmdao.utils.name_maps import abs_key2rel_key
  
  
  DEFAULT_CS_OPTIONS = {
      'step': 1e-15,
      'form': 'forward',
  }
  
  
  class ComplexStep(ApproximationScheme):
      r"""
      Approximation scheme using complex step to calculate derivatives.
  
      For example, using  a step size of 'h' will approximate the derivative in
      the following way:
  
      .. math::
  
          f'(x) = \Im{\frac{f(x+ih)}{h}}.
  
      Attributes
      ----------
      _exec_list : list
          A list of which derivatives (in execution order) to compute.
          The entries are of the form (of, wrt, options), where of and wrt are absolute names
          and options is a dictionary.
      """
  
      def __init__(self):
          """
          Initialize the ApproximationScheme.
          """
          super(ComplexStep, self).__init__()
          self._exec_list = []
  
      def add_approximation(self, abs_key, kwargs):
          """
          Use this approximation scheme to approximate the derivative d(of)/d(wrt).
  
          Parameters
          ----------
          abs_key : tuple(str,str)
              Absolute name pairing of (of, wrt) for the derivative.
          kwargs : dict
              Additional keyword arguments, to be interpreted by sub-classes.
          """
          of, wrt = abs_key
          options = DEFAULT_CS_OPTIONS.copy()
          options.update(kwargs)
          self._exec_list.append((of, wrt, options))
  
      @staticmethod
      def _key_fun(approx_tuple):
          """
          Compute the sorting key for an approximation tuple.
  
          Parameters
          ----------
          approx_tuple : tuple(str, str, dict)
              A given approximated derivative (of, wrt, options)
  
          Returns
          -------
          tuple(str, str, float)
              Sorting key (wrt, form, step_size)
  
          """
          options = approx_tuple[2]
          return (approx_tuple[1], options['form'], options['step'])
  
      def _init_approximations(self):
          """
          Prepare for later approximations.
          """
          # itertools.groupby works like `uniq` rather than the SQL query, meaning that it will only
          # group adjacent items with identical keys.
          self._exec_list.sort(key=self._key_fun)
  
          # TODO: Automatic sparse FD by constructing a graph of variable dependence?
  
      def compute_approximations(self, system, jac=None, deriv_type='partial'):
          """
          Execute the system to compute the approximate sub-Jacobians.
  
          Parameters
          ----------
          system : System
              System on which the execution is run.
          jac : None or dict-like
              If None, update system with the approximated sub-Jacobians. Otherwise, store the
              approximations in the given dict-like object.
          deriv_type : str
              One of 'total' or 'partial', indicating if total or partial derivatives are
              being approximated.
          """
          if jac is None:
              jac = system._jacobian
  
          if deriv_type == 'total':
              current_vec = system._outputs
          elif deriv_type == 'partial':
              current_vec = system._residuals
          else:
              raise ValueError('deriv_type must be one of "total" or "partial"')
  
          # Turn on complex step.
          system._inputs._vector_info._under_complex_step = True
  
          # create a scratch array
          out_tmp = system._outputs.get_data()
          results_clone = current_vec._clone(True)
  
          # To support driver src_indices, we need to override some checks in Jacobian, but do it
          # selectively.
          uses_src_indices = (system._owns_approx_of_idx or system._owns_approx_wrt_idx) and \
              not isinstance(jac, dict)
  
          for key, approximations in groupby(self._exec_list, self._key_fun):
              # groupby (along with this key function) will group all 'of's that have the same wrt and
              # step size.
              wrt, form, delta = key
              if form == 'reverse':
                  delta *= -1.0
              fact = 1.0 / delta
  
              if wrt in system._owns_approx_wrt_idx:
                  in_idx = system._owns_approx_wrt_idx[wrt]
                  in_size = len(in_idx)
              else:
                  if wrt in system._var_abs2meta:
                      in_size = system._var_abs2meta[wrt]['size']
  
                  in_idx = range(in_size)
  
              outputs = []
  
              # Note: If access to `approximations` is required again in the future, we will need to
              # throw it in a list first. The groupby iterator only works once.
              for approx_tuple in approximations:
                  of = approx_tuple[0]
                  # TODO: Sparse derivatives
                  if of in system._owns_approx_of_idx:
                      out_idx = system._owns_approx_of_idx[of]
                      out_size = len(out_idx)
                  else:
                      out_size = system._var_abs2meta[of]['size']
  
                  outputs.append((of, np.zeros((out_size, in_size))))
  
              for i_count, idx in enumerate(in_idx):
                  # Run the Finite Difference
                  input_delta = [(wrt, idx, delta)]
                  result = self._run_point_complex(system, input_delta, out_tmp, results_clone,
                                                   deriv_type)
  
                  for of, subjac in outputs:
                      if of in system._owns_approx_of_idx:
                          out_idx = system._owns_approx_of_idx[of]
                          subjac[:, i_count] = result._imag_views_flat[of][out_idx] * fact
                      else:
                          subjac[:, i_count] = result._imag_views_flat[of] * fact
  
              for of, subjac in outputs:
                  rel_key = abs_key2rel_key(system, (of, wrt))
                  if uses_src_indices:
                      jac._override_checks = True
                  jac[rel_key] = subjac
                  if uses_src_indices:
                      jac._override_checks = False
  
          # Turn off complex step.
          system._inputs._vector_info._under_complex_step = False
  
      def _run_point_complex(self, system, input_deltas, out_tmp, result_clone, deriv_type='partial'):
          """
          Perturb the system inputs with a complex step, runs, and returns the results.
  
          Parameters
          ----------
          system : System
              The system having its derivs approximated.
          input_deltas : list
              List of (input name, indices, delta) tuples, where input name is an absolute name.
          out_tmp : ndarray
              An array the same size as the system outputs that is used for temporary storage.
          result_clone : Vector
              A vector cloned from the outputs vector. Used to store the results.
          deriv_type : str
              One of 'total' or 'partial', indicating if total or partial derivatives are being
              approximated.
  
          Returns
          -------
          Vector
              Copy of the results from running the perturbed system.
          """
          # TODO: MPI
  
          inputs = system._inputs
          outputs = system._outputs
  
          if deriv_type == 'total':
              run_model = system.run_solve_nonlinear
              results_vec = outputs
          elif deriv_type == 'partial':
              run_model = system.run_apply_nonlinear
              results_vec = system._residuals
          else:
              raise ValueError('deriv_type must be one of "total" or "partial"')
  
          for in_name, idxs, delta in input_deltas:
              if in_name in outputs._imag_views_flat:
                  outputs._imag_views_flat[in_name][idxs] += delta
              else:
                  inputs._imag_views_flat[in_name][idxs] += delta
  
          results_vec.get_data(out_tmp)
          run_model()
  
          # TODO: Grab only results of interest
          result_clone.set_vec(results_vec)
          results_vec.set_data(out_tmp)
  
          for in_name, idxs, delta in input_deltas:
              if in_name in outputs._imag_views_flat:
                  outputs._imag_views_flat[in_name][idxs] -= delta
              else:
                  inputs._imag_views_flat[in_name][idxs] -= delta
  
          return result_clone
  
apis: numpy.zeros