Use pycompilation if installed.

opty/direct_collocation.py

@@ -15,7 +15,7 @@
  import_kwargs={'fromlist': ['']},
  catch=(RuntimeError,))
 
-from .utils import ufuncify_matrix, parse_free, _optional_plt_dep
+from .utils import ufuncify_matrix, parse_free, _optional_plt_dep, forward_jacobian
 
 __all__ = ['Problem', 'ConstraintCollocator']
 
@@ -115,7 +115,9 @@ def __init__(self, obj, obj_grad, *args, **kwargs):
  self.obj = obj
  self.obj_grad = obj_grad
  self.con = self.collocator.generate_constraint_function()
+ print('Constraint function generated.')
  self.con_jac = self.collocator.generate_jacobian_function()
+ print('Jacobian function generated.')
 
  self.con_jac_rows, self.con_jac_cols = \
  self.collocator.jacobian_indices()
@@ -1261,19 +1263,31 @@ def _gen_multi_arg_con_jac_func(self):
 
  # This creates a matrix with all of the symbolic partial derivatives
  # necessary to compute the full Jacobian.
- symbolic_partials = self.discrete_eom.jacobian(wrt)
+ print('Starting the Jacobian differentiation.')
+ # symbolic_partials = self.discrete_eom.jacobian(wrt)
+ discrete_eom_matrix = sm.ImmutableDenseMatrix(self.discrete_eom)
+ wrt_matrix = sm.ImmutableDenseMatrix([list(wrt)])
+ symbolic_partials = forward_jacobian(discrete_eom_matrix, wrt_matrix.T)
+ print('Jacobian differentiation finished.')
 
  # This generates a numerical function that evaluates the matrix of
  # partial derivatives. This function returns the non-zero elements
  # needed to build the sparse constraint Jacobian.
+ print('Starting the Jacobian compilation.')
  eval_partials = ufuncify_matrix(args, symbolic_partials,
  const=constant_syms + (h_sym,),
  tmp_dir=self.tmp_dir,
  parallel=self.parallel)
+ print('Jacobian compilation finished.')
 
+ if isinstance(symbolic_partials, tuple) and len(symbolic_partials) == 2:
+ num_rows = symbolic_partials[1][0].shape[0]
+ num_cols = symbolic_partials[1][0].shape[1]
+ else:
+ num_rows = symbolic_partials.shape[0]
+ num_cols = symbolic_partials.shape[1]
  result = np.empty((self.num_collocation_nodes - 1,
- symbolic_partials.shape[0] *
- symbolic_partials.shape[1]))
+ num_rows * num_cols))
 
  def constraints_jacobian(state_values, specified_values,
  parameter_values, interval_value):
@@ -1443,11 +1457,13 @@ def constraints(free):
  def generate_constraint_function(self):
  """Returns a function which evaluates the constraints given the
  array of free optimization variables."""
+ print('Generating constraint function.')
  self._gen_multi_arg_con_func()
  return self._wrap_constraint_funcs(self._multi_arg_con_func, 'con')
 
  def generate_jacobian_function(self):
  """Returns a function which evaluates the Jacobian of the
  constraints given the array of free optimization variables."""
+ print('Generating jacobian function.')
  self._gen_multi_arg_con_jac_func()
  return self._wrap_constraint_funcs(self._multi_arg_con_jac_func, 'jac')

opty/utils.py

@@ -11,9 +11,12 @@
 from distutils.ccompiler import new_compiler
 from distutils.errors import CompileError
 from distutils.sysconfig import customize_compiler
+from collections import Counter
+from timeit import default_timer as timer
 
 import numpy as np
 import sympy as sm
+from sympy.utilities.iterables import numbered_symbols
 try:
  plt = sm.external.import_module('matplotlib.pyplot',
  __import__kwargs={'fromlist': ['']},
@@ -23,6 +26,132 @@
  import_kwargs={'fromlist': ['']},
  catch=(RuntimeError,))
 
+pycompilation = sm.external.import_module('pycompilation')
+
+
+def forward_jacobian(expr, wrt):
+
+ def add_to_cache(node):
+ if node in expr_to_replacement_cache:
+ replacement_symbol = expr_to_replacement_cache[node]
+ return replacement_symbol, replacement_to_reduced_expr_cache[replacement_symbol]
+ elif node in replacement_to_reduced_expr_cache:
+ return node, replacement_to_reduced_expr_cache[node]
+ elif isinstance(node, sm.Tuple):
+ return None, None
+ elif not node.free_symbols:
+ return node, node
+
+ replacement_symbol = replacement_symbols.__next__()
+ replaced_subexpr = node.xreplace(expr_to_replacement_cache)
+ replacement_to_reduced_expr_cache[replacement_symbol] = replaced_subexpr
+ expr_to_replacement_cache[node] = replacement_symbol
+ return replacement_symbol, replaced_subexpr
+
+ if not isinstance(expr, sm.ImmutableDenseMatrix):
+ msg = (
+ 'The forward Jacobian differentiation algorithm can only be used '
+ 'to differentiate a single matrix expression at a time.'
+ )
+ raise NotImplementedError(msg)
+ elif expr.shape[1] != 1:
+ msg = 'Can only compute the Jacobian for column matrices.'
+ raise NotImplementedError(msg)
+ elif not isinstance(wrt, sm.ImmutableDenseMatrix) or wrt.shape[1] != 1:
+ msg = (
+ 'The forward Jacobian differentiation algorithm can compute '
+ 'Jacobians with respect to column matrices.'
+ )
+ raise NotImplementedError
+
+ replacement_symbols = numbered_symbols(
+ prefix='_z',
+ cls=sm.Symbol,
+ exclude=expr.free_symbols,
+ )
+
+ expr_to_replacement_cache = {}
+ replacement_to_reduced_expr_cache = {}
+
+ print('Adding expression nodes to cache...')
+ start = timer()
+ replacements, reduced_exprs = sm.cse(expr.args[2], replacement_symbols)
+ for replacement_symbol, reduced_subexpr in replacements:
+ replaced_subexpr = reduced_subexpr.xreplace(expr_to_replacement_cache)
+ replacement_to_reduced_expr_cache[replacement_symbol] = replaced_subexpr
+ expr_to_replacement_cache[reduced_subexpr] = replacement_symbol
+ for node in sm.postorder_traversal(reduced_subexpr):
+ _ = add_to_cache(node)
+ for reduced_expr in reduced_exprs:
+ for node in reduced_expr:
+ _ = add_to_cache(node)
+ finish = timer()
+ print(f'Completed in {finish - start:.2f}s')
+
+ reduced_matrix = sm.ImmutableDenseMatrix(reduced_exprs).xreplace(expr_to_replacement_cache)
+ replacements = list(replacement_to_reduced_expr_cache.items())
+
+ partial_derivative_mapping = {}
+ absolute_derivative_mapping = {}
+ for i, wrt_symbol in enumerate(wrt.args[2]):
+ absolute_derivative = [sm.S.Zero] * len(wrt)
+ absolute_derivative[i] = sm.S.One
+ absolute_derivative_mapping[wrt_symbol] = sm.ImmutableDenseMatrix([absolute_derivative])
+
+ print('Differentiating expression nodes...')
+ start = timer()
+ zeros = sm.ImmutableDenseMatrix.zeros(1, len(wrt))
+ for symbol, subexpr in replacements:
+ free_symbols = subexpr.free_symbols
+ absolute_derivative = zeros
+ for free_symbol in free_symbols:
+ replacement_symbol, partial_derivative = add_to_cache(subexpr.diff(free_symbol))
+ absolute_derivative += partial_derivative * absolute_derivative_mapping.get(free_symbol, zeros)
+ absolute_derivative_mapping[symbol] = sm.ImmutableDenseMatrix([[add_to_cache(a)[0] for a in absolute_derivative]])
+
+ replaced_jacobian = sm.ImmutableDenseMatrix.vstack(*[absolute_derivative_mapping[e] for e in reduced_matrix])
+ finish = timer()
+ print(f'Completed in {finish - start:.2f}s')
+
+ print('Determining required replacements...')
+ start = timer()
+ required_replacement_symbols = set()
+ stack = [entry for entry in replaced_jacobian if entry.free_symbols]
+ while stack:
+ entry = stack.pop()
+ if entry in required_replacement_symbols or entry in wrt:
+ continue
+ children = list(replacement_to_reduced_expr_cache.get(entry, entry).free_symbols)
+ for child in children:
+ if child not in required_replacement_symbols and child not in wrt:
+ stack.append(child)
+ required_replacement_symbols.add(entry)
+ finish = timer()
+ print(f'Completed in {finish - start:.2f}s')
+
+ required_replacements_dense = {
+ replacement_symbol: replaced_subexpr
+ for replacement_symbol, replaced_subexpr in replacement_to_reduced_expr_cache.items()
+ if replacement_symbol in required_replacement_symbols
+ }
+
+ counter = Counter(replaced_jacobian.free_symbols)
+ for replaced_subexpr in required_replacements_dense.values():
+ counter.update(replaced_subexpr.free_symbols)
+
+ print('Substituting required replacements...')
+ required_replacements = {}
+ unrequired_replacements = {}
+ for replacement_symbol, replaced_subexpr in required_replacements_dense.items():
+ if isinstance(replaced_subexpr, sm.Symbol) or counter[replacement_symbol] == 1:
+ unrequired_replacements[replacement_symbol] = replaced_subexpr.xreplace(unrequired_replacements)
+ else:
+ required_replacements[replacement_symbol] = replaced_subexpr.xreplace(unrequired_replacements)
+ finish = timer()
+ print(f'Completed in {finish - start:.2f}s')
+
+ return (list(required_replacements.items()), [replaced_jacobian.xreplace(unrequired_replacements)])
+
 
 def building_docs():
  if 'READTHEDOCS' in os.environ:
@@ -103,26 +232,20 @@ def parse_free(free, n, q, N):
 
 _c_template = """\
 #include <math.h>
-#include "{file_prefix}_h.h"
 
 void {routine_name}(double matrix[{matrix_output_size}], {input_args})
 {{
 {eval_code}
 }}
 """
 
-_h_template = """\
-void {routine_name}(double matrix[{matrix_output_size}], {input_args});
-"""
-
 _cython_template = """\
 import numpy as np
 from cython.parallel import prange
 cimport numpy as np
 cimport cython
 
-cdef extern from "{file_prefix}_h.h"{head_gil}:
- void {routine_name}(double matrix[{matrix_output_size}], {input_args})
+cdef extern void c_{routine_name} "{routine_name}" (double matrix[{matrix_output_size}], {input_args}){head_gil}
 
 @cython.boundscheck(False)
 @cython.wraparound(False)
@@ -133,7 +256,7 @@ def {routine_name}_loop(np.ndarray[np.double_t, ndim=2] matrix, {numpy_typed_inp
  cdef int i
 
  for i in {loop_sig}:
- {routine_name}(&matrix[i, 0], {indexed_input_args})
+ c_{routine_name}(&matrix[i, 0], {indexed_input_args})
 
  return matrix.reshape(n, {num_rows}, {num_cols})
 """
@@ -230,7 +353,14 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False):
  # not sure if this current version counts sequentially.
  global module_counter
 
- matrix_size = expr.shape[0] * expr.shape[1]
+ if hasattr(expr, 'shape'):
+ num_rows = expr.shape[0]
+ num_cols = expr.shape[1]
+ else:
+ num_rows = expr[1][0].shape[0]
+ num_cols = expr[1][0].shape[1]
+
+ matrix_size = num_rows * num_cols
 
  file_prefix_base = 'ufuncify_matrix'
  file_prefix = '{}_{}'.format(file_prefix_base, module_counter)
@@ -257,8 +387,8 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False):
  d = {'routine_name': 'eval_matrix',
  'file_prefix': file_prefix,
  'matrix_output_size': matrix_size,
- 'num_rows': expr.shape[0],
- 'num_cols': expr.shape[1]}
+ 'num_rows': num_rows,
+ 'num_cols': num_cols}
 
  if parallel:
  if openmp_installed():
@@ -280,9 +410,12 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False):
  d['compile_args'] = ""
  d['link_args'] = ""
 
- matrix_sym = sm.MatrixSymbol('matrix', expr.shape[0], expr.shape[1])
+ matrix_sym = sm.MatrixSymbol('matrix', num_rows, num_cols)
 
- sub_exprs, simple_mat = sm.cse(expr, sm.numbered_symbols('z_'))
+ if isinstance(expr, tuple) and len(expr) == 2:
+ sub_exprs, simple_mat = expr
+ else:
+ sub_exprs, simple_mat = sm.cse(expr, sm.numbered_symbols('z_'))
 
  sub_expr_code = '\n'.join(['double ' + sm.ccode(sub_expr[1], sub_expr[0])
  for sub_expr in sub_exprs])
@@ -320,31 +453,45 @@ def ufuncify_matrix(args, expr, const=None, tmp_dir=None, parallel=False):
 
  files = {}
  files[d['file_prefix'] + '_c.c'] = _c_template.format(**d)
- files[d['file_prefix'] + '_h.h'] = _h_template.format(**d)
  files[d['file_prefix'] + '.pyx'] = _cython_template.format(**d)
  files[d['file_prefix'] + '_setup.py'] = _setup_template.format(**d)
 
- workingdir = os.getcwd()
- os.chdir(codedir)
+ if pycompilation:
+ sources = [
+ (d['file_prefix'] + '_c.c', files[d['file_prefix'] + '_c.c']),
+ (d['file_prefix'] + '.pyx', files[d['file_prefix'] + '.pyx']),
+ ]
 
- try:
- sys.path.append(codedir)
- for filename, code in files.items():
- with open(filename, 'w') as f:
- f.write(code)
- cmd = [sys.executable, d['file_prefix'] + '_setup.py', 'build_ext',
- '--inplace']
- subprocess.call(cmd, stderr=subprocess.STDOUT, stdout=subprocess.PIPE)
- cython_module = importlib.import_module(d['file_prefix'])
- finally:
- module_counter += 1
- sys.path.remove(codedir)
- os.chdir(workingdir)
- if tmp_dir is None:
- # NOTE : I can't figure out how to get rmtree to work on Windows,
- # so I don't delete the directory on Windows.
- if sys.platform != "win32":
- shutil.rmtree(codedir)
+ options = []
+ if parallel:
+ options += ['openmp']
+
+ cython_module = pycompilation.compile_link_import_strings(
+ sources, options=options, std='c99',
+ include_dirs=[np.get_include()], build_dir=codedir)
+ else:
+ workingdir = os.getcwd()
+ os.chdir(codedir)
+
+ try:
+ sys.path.append(codedir)
+ for filename, code in files.items():
+ with open(filename, 'w') as f:
+ f.write(code)
+ cmd = [sys.executable, d['file_prefix'] + '_setup.py', 'build_ext',
+ '--inplace']
+ subprocess.call(cmd, stderr=subprocess.STDOUT,
+ stdout=subprocess.PIPE)
+ cython_module = importlib.import_module(d['file_prefix'])
+ finally:
+ module_counter += 1
+ sys.path.remove(codedir)
+ os.chdir(workingdir)
+ if tmp_dir is None:
+ # NOTE : I can't figure out how to get rmtree to work on Windows,
+ # so I don't delete the directory on Windows.
+ if sys.platform != "win32":
+ shutil.rmtree(codedir)
 
  return getattr(cython_module, d['routine_name'] + '_loop')