autowrap.py 40 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134
  1. """Module for compiling codegen output, and wrap the binary for use in
  2. python.
  3. .. note:: To use the autowrap module it must first be imported
  4. >>> from sympy.utilities.autowrap import autowrap
  5. This module provides a common interface for different external backends, such
  6. as f2py, fwrap, Cython, SWIG(?) etc. (Currently only f2py and Cython are
  7. implemented) The goal is to provide access to compiled binaries of acceptable
  8. performance with a one-button user interface, e.g.,
  9. >>> from sympy.abc import x,y
  10. >>> expr = (x - y)**25
  11. >>> flat = expr.expand()
  12. >>> binary_callable = autowrap(flat)
  13. >>> binary_callable(2, 3)
  14. -1.0
  15. Although a SymPy user might primarily be interested in working with
  16. mathematical expressions and not in the details of wrapping tools
  17. needed to evaluate such expressions efficiently in numerical form,
  18. the user cannot do so without some understanding of the
  19. limits in the target language. For example, the expanded expression
  20. contains large coefficients which result in loss of precision when
  21. computing the expression:
  22. >>> binary_callable(3, 2)
  23. 0.0
  24. >>> binary_callable(4, 5), binary_callable(5, 4)
  25. (-22925376.0, 25165824.0)
  26. Wrapping the unexpanded expression gives the expected behavior:
  27. >>> e = autowrap(expr)
  28. >>> e(4, 5), e(5, 4)
  29. (-1.0, 1.0)
  30. The callable returned from autowrap() is a binary Python function, not a
  31. SymPy object. If it is desired to use the compiled function in symbolic
  32. expressions, it is better to use binary_function() which returns a SymPy
  33. Function object. The binary callable is attached as the _imp_ attribute and
  34. invoked when a numerical evaluation is requested with evalf(), or with
  35. lambdify().
  36. >>> from sympy.utilities.autowrap import binary_function
  37. >>> f = binary_function('f', expr)
  38. >>> 2*f(x, y) + y
  39. y + 2*f(x, y)
  40. >>> (2*f(x, y) + y).evalf(2, subs={x: 1, y:2})
  41. 0.e-110
  42. When is this useful?
  43. 1) For computations on large arrays, Python iterations may be too slow,
  44. and depending on the mathematical expression, it may be difficult to
  45. exploit the advanced index operations provided by NumPy.
  46. 2) For *really* long expressions that will be called repeatedly, the
  47. compiled binary should be significantly faster than SymPy's .evalf()
  48. 3) If you are generating code with the codegen utility in order to use
  49. it in another project, the automatic Python wrappers let you test the
  50. binaries immediately from within SymPy.
  51. 4) To create customized ufuncs for use with numpy arrays.
  52. See *ufuncify*.
  53. When is this module NOT the best approach?
  54. 1) If you are really concerned about speed or memory optimizations,
  55. you will probably get better results by working directly with the
  56. wrapper tools and the low level code. However, the files generated
  57. by this utility may provide a useful starting point and reference
  58. code. Temporary files will be left intact if you supply the keyword
  59. tempdir="path/to/files/".
  60. 2) If the array computation can be handled easily by numpy, and you
  61. do not need the binaries for another project.
  62. """
  63. import sys
  64. import os
  65. import shutil
  66. import tempfile
  67. from subprocess import STDOUT, CalledProcessError, check_output
  68. from string import Template
  69. from warnings import warn
  70. from sympy.core.cache import cacheit
  71. from sympy.core.function import Lambda
  72. from sympy.core.relational import Eq
  73. from sympy.core.symbol import Dummy, Symbol
  74. from sympy.tensor.indexed import Idx, IndexedBase
  75. from sympy.utilities.codegen import (make_routine, get_code_generator,
  76. OutputArgument, InOutArgument,
  77. InputArgument, CodeGenArgumentListError,
  78. Result, ResultBase, C99CodeGen)
  79. from sympy.utilities.iterables import iterable
  80. from sympy.utilities.lambdify import implemented_function
  81. from sympy.utilities.decorator import doctest_depends_on
  82. _doctest_depends_on = {'exe': ('f2py', 'gfortran', 'gcc'),
  83. 'modules': ('numpy',)}
  84. class CodeWrapError(Exception):
  85. pass
  86. class CodeWrapper:
  87. """Base Class for code wrappers"""
  88. _filename = "wrapped_code"
  89. _module_basename = "wrapper_module"
  90. _module_counter = 0
  91. @property
  92. def filename(self):
  93. return "%s_%s" % (self._filename, CodeWrapper._module_counter)
  94. @property
  95. def module_name(self):
  96. return "%s_%s" % (self._module_basename, CodeWrapper._module_counter)
  97. def __init__(self, generator, filepath=None, flags=[], verbose=False):
  98. """
  99. generator -- the code generator to use
  100. """
  101. self.generator = generator
  102. self.filepath = filepath
  103. self.flags = flags
  104. self.quiet = not verbose
  105. @property
  106. def include_header(self):
  107. return bool(self.filepath)
  108. @property
  109. def include_empty(self):
  110. return bool(self.filepath)
  111. def _generate_code(self, main_routine, routines):
  112. routines.append(main_routine)
  113. self.generator.write(
  114. routines, self.filename, True, self.include_header,
  115. self.include_empty)
  116. def wrap_code(self, routine, helpers=None):
  117. helpers = helpers or []
  118. if self.filepath:
  119. workdir = os.path.abspath(self.filepath)
  120. else:
  121. workdir = tempfile.mkdtemp("_sympy_compile")
  122. if not os.access(workdir, os.F_OK):
  123. os.mkdir(workdir)
  124. oldwork = os.getcwd()
  125. os.chdir(workdir)
  126. try:
  127. sys.path.append(workdir)
  128. self._generate_code(routine, helpers)
  129. self._prepare_files(routine)
  130. self._process_files(routine)
  131. mod = __import__(self.module_name)
  132. finally:
  133. sys.path.remove(workdir)
  134. CodeWrapper._module_counter += 1
  135. os.chdir(oldwork)
  136. if not self.filepath:
  137. try:
  138. shutil.rmtree(workdir)
  139. except OSError:
  140. # Could be some issues on Windows
  141. pass
  142. return self._get_wrapped_function(mod, routine.name)
  143. def _process_files(self, routine):
  144. command = self.command
  145. command.extend(self.flags)
  146. try:
  147. retoutput = check_output(command, stderr=STDOUT)
  148. except CalledProcessError as e:
  149. raise CodeWrapError(
  150. "Error while executing command: %s. Command output is:\n%s" % (
  151. " ".join(command), e.output.decode('utf-8')))
  152. if not self.quiet:
  153. print(retoutput)
  154. class DummyWrapper(CodeWrapper):
  155. """Class used for testing independent of backends """
  156. template = """# dummy module for testing of SymPy
  157. def %(name)s():
  158. return "%(expr)s"
  159. %(name)s.args = "%(args)s"
  160. %(name)s.returns = "%(retvals)s"
  161. """
  162. def _prepare_files(self, routine):
  163. return
  164. def _generate_code(self, routine, helpers):
  165. with open('%s.py' % self.module_name, 'w') as f:
  166. printed = ", ".join(
  167. [str(res.expr) for res in routine.result_variables])
  168. # convert OutputArguments to return value like f2py
  169. args = filter(lambda x: not isinstance(
  170. x, OutputArgument), routine.arguments)
  171. retvals = []
  172. for val in routine.result_variables:
  173. if isinstance(val, Result):
  174. retvals.append('nameless')
  175. else:
  176. retvals.append(val.result_var)
  177. print(DummyWrapper.template % {
  178. 'name': routine.name,
  179. 'expr': printed,
  180. 'args': ", ".join([str(a.name) for a in args]),
  181. 'retvals': ", ".join([str(val) for val in retvals])
  182. }, end="", file=f)
  183. def _process_files(self, routine):
  184. return
  185. @classmethod
  186. def _get_wrapped_function(cls, mod, name):
  187. return getattr(mod, name)
  188. class CythonCodeWrapper(CodeWrapper):
  189. """Wrapper that uses Cython"""
  190. setup_template = """\
  191. try:
  192. from setuptools import setup
  193. from setuptools import Extension
  194. except ImportError:
  195. from distutils.core import setup
  196. from distutils.extension import Extension
  197. from Cython.Build import cythonize
  198. cy_opts = {cythonize_options}
  199. {np_import}
  200. ext_mods = [Extension(
  201. {ext_args},
  202. include_dirs={include_dirs},
  203. library_dirs={library_dirs},
  204. libraries={libraries},
  205. extra_compile_args={extra_compile_args},
  206. extra_link_args={extra_link_args}
  207. )]
  208. setup(ext_modules=cythonize(ext_mods, **cy_opts))
  209. """
  210. pyx_imports = (
  211. "import numpy as np\n"
  212. "cimport numpy as np\n\n")
  213. pyx_header = (
  214. "cdef extern from '{header_file}.h':\n"
  215. " {prototype}\n\n")
  216. pyx_func = (
  217. "def {name}_c({arg_string}):\n"
  218. "\n"
  219. "{declarations}"
  220. "{body}")
  221. std_compile_flag = '-std=c99'
  222. def __init__(self, *args, **kwargs):
  223. """Instantiates a Cython code wrapper.
  224. The following optional parameters get passed to ``distutils.Extension``
  225. for building the Python extension module. Read its documentation to
  226. learn more.
  227. Parameters
  228. ==========
  229. include_dirs : [list of strings]
  230. A list of directories to search for C/C++ header files (in Unix
  231. form for portability).
  232. library_dirs : [list of strings]
  233. A list of directories to search for C/C++ libraries at link time.
  234. libraries : [list of strings]
  235. A list of library names (not filenames or paths) to link against.
  236. extra_compile_args : [list of strings]
  237. Any extra platform- and compiler-specific information to use when
  238. compiling the source files in 'sources'. For platforms and
  239. compilers where "command line" makes sense, this is typically a
  240. list of command-line arguments, but for other platforms it could be
  241. anything. Note that the attribute ``std_compile_flag`` will be
  242. appended to this list.
  243. extra_link_args : [list of strings]
  244. Any extra platform- and compiler-specific information to use when
  245. linking object files together to create the extension (or to create
  246. a new static Python interpreter). Similar interpretation as for
  247. 'extra_compile_args'.
  248. cythonize_options : [dictionary]
  249. Keyword arguments passed on to cythonize.
  250. """
  251. self._include_dirs = kwargs.pop('include_dirs', [])
  252. self._library_dirs = kwargs.pop('library_dirs', [])
  253. self._libraries = kwargs.pop('libraries', [])
  254. self._extra_compile_args = kwargs.pop('extra_compile_args', [])
  255. self._extra_compile_args.append(self.std_compile_flag)
  256. self._extra_link_args = kwargs.pop('extra_link_args', [])
  257. self._cythonize_options = kwargs.pop('cythonize_options', {})
  258. self._need_numpy = False
  259. super().__init__(*args, **kwargs)
  260. @property
  261. def command(self):
  262. command = [sys.executable, "setup.py", "build_ext", "--inplace"]
  263. return command
  264. def _prepare_files(self, routine, build_dir=os.curdir):
  265. # NOTE : build_dir is used for testing purposes.
  266. pyxfilename = self.module_name + '.pyx'
  267. codefilename = "%s.%s" % (self.filename, self.generator.code_extension)
  268. # pyx
  269. with open(os.path.join(build_dir, pyxfilename), 'w') as f:
  270. self.dump_pyx([routine], f, self.filename)
  271. # setup.py
  272. ext_args = [repr(self.module_name), repr([pyxfilename, codefilename])]
  273. if self._need_numpy:
  274. np_import = 'import numpy as np\n'
  275. self._include_dirs.append('np.get_include()')
  276. else:
  277. np_import = ''
  278. with open(os.path.join(build_dir, 'setup.py'), 'w') as f:
  279. includes = str(self._include_dirs).replace("'np.get_include()'",
  280. 'np.get_include()')
  281. f.write(self.setup_template.format(
  282. ext_args=", ".join(ext_args),
  283. np_import=np_import,
  284. include_dirs=includes,
  285. library_dirs=self._library_dirs,
  286. libraries=self._libraries,
  287. extra_compile_args=self._extra_compile_args,
  288. extra_link_args=self._extra_link_args,
  289. cythonize_options=self._cythonize_options
  290. ))
  291. @classmethod
  292. def _get_wrapped_function(cls, mod, name):
  293. return getattr(mod, name + '_c')
  294. def dump_pyx(self, routines, f, prefix):
  295. """Write a Cython file with Python wrappers
  296. This file contains all the definitions of the routines in c code and
  297. refers to the header file.
  298. Arguments
  299. ---------
  300. routines
  301. List of Routine instances
  302. f
  303. File-like object to write the file to
  304. prefix
  305. The filename prefix, used to refer to the proper header file.
  306. Only the basename of the prefix is used.
  307. """
  308. headers = []
  309. functions = []
  310. for routine in routines:
  311. prototype = self.generator.get_prototype(routine)
  312. # C Function Header Import
  313. headers.append(self.pyx_header.format(header_file=prefix,
  314. prototype=prototype))
  315. # Partition the C function arguments into categories
  316. py_rets, py_args, py_loc, py_inf = self._partition_args(routine.arguments)
  317. # Function prototype
  318. name = routine.name
  319. arg_string = ", ".join(self._prototype_arg(arg) for arg in py_args)
  320. # Local Declarations
  321. local_decs = []
  322. for arg, val in py_inf.items():
  323. proto = self._prototype_arg(arg)
  324. mat, ind = [self._string_var(v) for v in val]
  325. local_decs.append(" cdef {} = {}.shape[{}]".format(proto, mat, ind))
  326. local_decs.extend([" cdef {}".format(self._declare_arg(a)) for a in py_loc])
  327. declarations = "\n".join(local_decs)
  328. if declarations:
  329. declarations = declarations + "\n"
  330. # Function Body
  331. args_c = ", ".join([self._call_arg(a) for a in routine.arguments])
  332. rets = ", ".join([self._string_var(r.name) for r in py_rets])
  333. if routine.results:
  334. body = ' return %s(%s)' % (routine.name, args_c)
  335. if rets:
  336. body = body + ', ' + rets
  337. else:
  338. body = ' %s(%s)\n' % (routine.name, args_c)
  339. body = body + ' return ' + rets
  340. functions.append(self.pyx_func.format(name=name, arg_string=arg_string,
  341. declarations=declarations, body=body))
  342. # Write text to file
  343. if self._need_numpy:
  344. # Only import numpy if required
  345. f.write(self.pyx_imports)
  346. f.write('\n'.join(headers))
  347. f.write('\n'.join(functions))
  348. def _partition_args(self, args):
  349. """Group function arguments into categories."""
  350. py_args = []
  351. py_returns = []
  352. py_locals = []
  353. py_inferred = {}
  354. for arg in args:
  355. if isinstance(arg, OutputArgument):
  356. py_returns.append(arg)
  357. py_locals.append(arg)
  358. elif isinstance(arg, InOutArgument):
  359. py_returns.append(arg)
  360. py_args.append(arg)
  361. else:
  362. py_args.append(arg)
  363. # Find arguments that are array dimensions. These can be inferred
  364. # locally in the Cython code.
  365. if isinstance(arg, (InputArgument, InOutArgument)) and arg.dimensions:
  366. dims = [d[1] + 1 for d in arg.dimensions]
  367. sym_dims = [(i, d) for (i, d) in enumerate(dims) if
  368. isinstance(d, Symbol)]
  369. for (i, d) in sym_dims:
  370. py_inferred[d] = (arg.name, i)
  371. for arg in args:
  372. if arg.name in py_inferred:
  373. py_inferred[arg] = py_inferred.pop(arg.name)
  374. # Filter inferred arguments from py_args
  375. py_args = [a for a in py_args if a not in py_inferred]
  376. return py_returns, py_args, py_locals, py_inferred
  377. def _prototype_arg(self, arg):
  378. mat_dec = "np.ndarray[{mtype}, ndim={ndim}] {name}"
  379. np_types = {'double': 'np.double_t',
  380. 'int': 'np.int_t'}
  381. t = arg.get_datatype('c')
  382. if arg.dimensions:
  383. self._need_numpy = True
  384. ndim = len(arg.dimensions)
  385. mtype = np_types[t]
  386. return mat_dec.format(mtype=mtype, ndim=ndim, name=self._string_var(arg.name))
  387. else:
  388. return "%s %s" % (t, self._string_var(arg.name))
  389. def _declare_arg(self, arg):
  390. proto = self._prototype_arg(arg)
  391. if arg.dimensions:
  392. shape = '(' + ','.join(self._string_var(i[1] + 1) for i in arg.dimensions) + ')'
  393. return proto + " = np.empty({shape})".format(shape=shape)
  394. else:
  395. return proto + " = 0"
  396. def _call_arg(self, arg):
  397. if arg.dimensions:
  398. t = arg.get_datatype('c')
  399. return "<{}*> {}.data".format(t, self._string_var(arg.name))
  400. elif isinstance(arg, ResultBase):
  401. return "&{}".format(self._string_var(arg.name))
  402. else:
  403. return self._string_var(arg.name)
  404. def _string_var(self, var):
  405. printer = self.generator.printer.doprint
  406. return printer(var)
  407. class F2PyCodeWrapper(CodeWrapper):
  408. """Wrapper that uses f2py"""
  409. def __init__(self, *args, **kwargs):
  410. ext_keys = ['include_dirs', 'library_dirs', 'libraries',
  411. 'extra_compile_args', 'extra_link_args']
  412. msg = ('The compilation option kwarg {} is not supported with the f2py '
  413. 'backend.')
  414. for k in ext_keys:
  415. if k in kwargs.keys():
  416. warn(msg.format(k))
  417. kwargs.pop(k, None)
  418. super().__init__(*args, **kwargs)
  419. @property
  420. def command(self):
  421. filename = self.filename + '.' + self.generator.code_extension
  422. args = ['-c', '-m', self.module_name, filename]
  423. command = [sys.executable, "-c", "import numpy.f2py as f2py2e;f2py2e.main()"]+args
  424. return command
  425. def _prepare_files(self, routine):
  426. pass
  427. @classmethod
  428. def _get_wrapped_function(cls, mod, name):
  429. return getattr(mod, name)
  430. # Here we define a lookup of backends -> tuples of languages. For now, each
  431. # tuple is of length 1, but if a backend supports more than one language,
  432. # the most preferable language is listed first.
  433. _lang_lookup = {'CYTHON': ('C99', 'C89', 'C'),
  434. 'F2PY': ('F95',),
  435. 'NUMPY': ('C99', 'C89', 'C'),
  436. 'DUMMY': ('F95',)} # Dummy here just for testing
  437. def _infer_language(backend):
  438. """For a given backend, return the top choice of language"""
  439. langs = _lang_lookup.get(backend.upper(), False)
  440. if not langs:
  441. raise ValueError("Unrecognized backend: " + backend)
  442. return langs[0]
  443. def _validate_backend_language(backend, language):
  444. """Throws error if backend and language are incompatible"""
  445. langs = _lang_lookup.get(backend.upper(), False)
  446. if not langs:
  447. raise ValueError("Unrecognized backend: " + backend)
  448. if language.upper() not in langs:
  449. raise ValueError(("Backend {} and language {} are "
  450. "incompatible").format(backend, language))
  451. @cacheit
  452. @doctest_depends_on(exe=('f2py', 'gfortran'), modules=('numpy',))
  453. def autowrap(expr, language=None, backend='f2py', tempdir=None, args=None,
  454. flags=None, verbose=False, helpers=None, code_gen=None, **kwargs):
  455. """Generates Python callable binaries based on the math expression.
  456. Parameters
  457. ==========
  458. expr
  459. The SymPy expression that should be wrapped as a binary routine.
  460. language : string, optional
  461. If supplied, (options: 'C' or 'F95'), specifies the language of the
  462. generated code. If ``None`` [default], the language is inferred based
  463. upon the specified backend.
  464. backend : string, optional
  465. Backend used to wrap the generated code. Either 'f2py' [default],
  466. or 'cython'.
  467. tempdir : string, optional
  468. Path to directory for temporary files. If this argument is supplied,
  469. the generated code and the wrapper input files are left intact in the
  470. specified path.
  471. args : iterable, optional
  472. An ordered iterable of symbols. Specifies the argument sequence for the
  473. function.
  474. flags : iterable, optional
  475. Additional option flags that will be passed to the backend.
  476. verbose : bool, optional
  477. If True, autowrap will not mute the command line backends. This can be
  478. helpful for debugging.
  479. helpers : 3-tuple or iterable of 3-tuples, optional
  480. Used to define auxiliary expressions needed for the main expr. If the
  481. main expression needs to call a specialized function it should be
  482. passed in via ``helpers``. Autowrap will then make sure that the
  483. compiled main expression can link to the helper routine. Items should
  484. be 3-tuples with (<function_name>, <sympy_expression>,
  485. <argument_tuple>). It is mandatory to supply an argument sequence to
  486. helper routines.
  487. code_gen : CodeGen instance
  488. An instance of a CodeGen subclass. Overrides ``language``.
  489. include_dirs : [string]
  490. A list of directories to search for C/C++ header files (in Unix form
  491. for portability).
  492. library_dirs : [string]
  493. A list of directories to search for C/C++ libraries at link time.
  494. libraries : [string]
  495. A list of library names (not filenames or paths) to link against.
  496. extra_compile_args : [string]
  497. Any extra platform- and compiler-specific information to use when
  498. compiling the source files in 'sources'. For platforms and compilers
  499. where "command line" makes sense, this is typically a list of
  500. command-line arguments, but for other platforms it could be anything.
  501. extra_link_args : [string]
  502. Any extra platform- and compiler-specific information to use when
  503. linking object files together to create the extension (or to create a
  504. new static Python interpreter). Similar interpretation as for
  505. 'extra_compile_args'.
  506. Examples
  507. ========
  508. >>> from sympy.abc import x, y, z
  509. >>> from sympy.utilities.autowrap import autowrap
  510. >>> expr = ((x - y + z)**(13)).expand()
  511. >>> binary_func = autowrap(expr)
  512. >>> binary_func(1, 4, 2)
  513. -1.0
  514. """
  515. if language:
  516. if not isinstance(language, type):
  517. _validate_backend_language(backend, language)
  518. else:
  519. language = _infer_language(backend)
  520. # two cases 1) helpers is an iterable of 3-tuples and 2) helpers is a
  521. # 3-tuple
  522. if iterable(helpers) and len(helpers) != 0 and iterable(helpers[0]):
  523. helpers = helpers if helpers else ()
  524. else:
  525. helpers = [helpers] if helpers else ()
  526. args = list(args) if iterable(args, exclude=set) else args
  527. if code_gen is None:
  528. code_gen = get_code_generator(language, "autowrap")
  529. CodeWrapperClass = {
  530. 'F2PY': F2PyCodeWrapper,
  531. 'CYTHON': CythonCodeWrapper,
  532. 'DUMMY': DummyWrapper
  533. }[backend.upper()]
  534. code_wrapper = CodeWrapperClass(code_gen, tempdir, flags if flags else (),
  535. verbose, **kwargs)
  536. helps = []
  537. for name_h, expr_h, args_h in helpers:
  538. helps.append(code_gen.routine(name_h, expr_h, args_h))
  539. for name_h, expr_h, args_h in helpers:
  540. if expr.has(expr_h):
  541. name_h = binary_function(name_h, expr_h, backend='dummy')
  542. expr = expr.subs(expr_h, name_h(*args_h))
  543. try:
  544. routine = code_gen.routine('autofunc', expr, args)
  545. except CodeGenArgumentListError as e:
  546. # if all missing arguments are for pure output, we simply attach them
  547. # at the end and try again, because the wrappers will silently convert
  548. # them to return values anyway.
  549. new_args = []
  550. for missing in e.missing_args:
  551. if not isinstance(missing, OutputArgument):
  552. raise
  553. new_args.append(missing.name)
  554. routine = code_gen.routine('autofunc', expr, args + new_args)
  555. return code_wrapper.wrap_code(routine, helpers=helps)
  556. @doctest_depends_on(exe=('f2py', 'gfortran'), modules=('numpy',))
  557. def binary_function(symfunc, expr, **kwargs):
  558. """Returns a SymPy function with expr as binary implementation
  559. This is a convenience function that automates the steps needed to
  560. autowrap the SymPy expression and attaching it to a Function object
  561. with implemented_function().
  562. Parameters
  563. ==========
  564. symfunc : SymPy Function
  565. The function to bind the callable to.
  566. expr : SymPy Expression
  567. The expression used to generate the function.
  568. kwargs : dict
  569. Any kwargs accepted by autowrap.
  570. Examples
  571. ========
  572. >>> from sympy.abc import x, y
  573. >>> from sympy.utilities.autowrap import binary_function
  574. >>> expr = ((x - y)**(25)).expand()
  575. >>> f = binary_function('f', expr)
  576. >>> type(f)
  577. <class 'sympy.core.function.UndefinedFunction'>
  578. >>> 2*f(x, y)
  579. 2*f(x, y)
  580. >>> f(x, y).evalf(2, subs={x: 1, y: 2})
  581. -1.0
  582. """
  583. binary = autowrap(expr, **kwargs)
  584. return implemented_function(symfunc, binary)
  585. #################################################################
  586. # UFUNCIFY #
  587. #################################################################
  588. _ufunc_top = Template("""\
  589. #include "Python.h"
  590. #include "math.h"
  591. #include "numpy/ndarraytypes.h"
  592. #include "numpy/ufuncobject.h"
  593. #include "numpy/halffloat.h"
  594. #include ${include_file}
  595. static PyMethodDef ${module}Methods[] = {
  596. {NULL, NULL, 0, NULL}
  597. };""")
  598. _ufunc_outcalls = Template("*((double *)out${outnum}) = ${funcname}(${call_args});")
  599. _ufunc_body = Template("""\
  600. static void ${funcname}_ufunc(char **args, npy_intp *dimensions, npy_intp* steps, void* data)
  601. {
  602. npy_intp i;
  603. npy_intp n = dimensions[0];
  604. ${declare_args}
  605. ${declare_steps}
  606. for (i = 0; i < n; i++) {
  607. ${outcalls}
  608. ${step_increments}
  609. }
  610. }
  611. PyUFuncGenericFunction ${funcname}_funcs[1] = {&${funcname}_ufunc};
  612. static char ${funcname}_types[${n_types}] = ${types}
  613. static void *${funcname}_data[1] = {NULL};""")
  614. _ufunc_bottom = Template("""\
  615. #if PY_VERSION_HEX >= 0x03000000
  616. static struct PyModuleDef moduledef = {
  617. PyModuleDef_HEAD_INIT,
  618. "${module}",
  619. NULL,
  620. -1,
  621. ${module}Methods,
  622. NULL,
  623. NULL,
  624. NULL,
  625. NULL
  626. };
  627. PyMODINIT_FUNC PyInit_${module}(void)
  628. {
  629. PyObject *m, *d;
  630. ${function_creation}
  631. m = PyModule_Create(&moduledef);
  632. if (!m) {
  633. return NULL;
  634. }
  635. import_array();
  636. import_umath();
  637. d = PyModule_GetDict(m);
  638. ${ufunc_init}
  639. return m;
  640. }
  641. #else
  642. PyMODINIT_FUNC init${module}(void)
  643. {
  644. PyObject *m, *d;
  645. ${function_creation}
  646. m = Py_InitModule("${module}", ${module}Methods);
  647. if (m == NULL) {
  648. return;
  649. }
  650. import_array();
  651. import_umath();
  652. d = PyModule_GetDict(m);
  653. ${ufunc_init}
  654. }
  655. #endif\
  656. """)
  657. _ufunc_init_form = Template("""\
  658. ufunc${ind} = PyUFunc_FromFuncAndData(${funcname}_funcs, ${funcname}_data, ${funcname}_types, 1, ${n_in}, ${n_out},
  659. PyUFunc_None, "${module}", ${docstring}, 0);
  660. PyDict_SetItemString(d, "${funcname}", ufunc${ind});
  661. Py_DECREF(ufunc${ind});""")
  662. _ufunc_setup = Template("""\
  663. def configuration(parent_package='', top_path=None):
  664. import numpy
  665. from numpy.distutils.misc_util import Configuration
  666. config = Configuration('',
  667. parent_package,
  668. top_path)
  669. config.add_extension('${module}', sources=['${module}.c', '${filename}.c'])
  670. return config
  671. if __name__ == "__main__":
  672. from numpy.distutils.core import setup
  673. setup(configuration=configuration)""")
  674. class UfuncifyCodeWrapper(CodeWrapper):
  675. """Wrapper for Ufuncify"""
  676. def __init__(self, *args, **kwargs):
  677. ext_keys = ['include_dirs', 'library_dirs', 'libraries',
  678. 'extra_compile_args', 'extra_link_args']
  679. msg = ('The compilation option kwarg {} is not supported with the numpy'
  680. ' backend.')
  681. for k in ext_keys:
  682. if k in kwargs.keys():
  683. warn(msg.format(k))
  684. kwargs.pop(k, None)
  685. super().__init__(*args, **kwargs)
  686. @property
  687. def command(self):
  688. command = [sys.executable, "setup.py", "build_ext", "--inplace"]
  689. return command
  690. def wrap_code(self, routines, helpers=None):
  691. # This routine overrides CodeWrapper because we can't assume funcname == routines[0].name
  692. # Therefore we have to break the CodeWrapper private API.
  693. # There isn't an obvious way to extend multi-expr support to
  694. # the other autowrap backends, so we limit this change to ufuncify.
  695. helpers = helpers if helpers is not None else []
  696. # We just need a consistent name
  697. funcname = 'wrapped_' + str(id(routines) + id(helpers))
  698. workdir = self.filepath or tempfile.mkdtemp("_sympy_compile")
  699. if not os.access(workdir, os.F_OK):
  700. os.mkdir(workdir)
  701. oldwork = os.getcwd()
  702. os.chdir(workdir)
  703. try:
  704. sys.path.append(workdir)
  705. self._generate_code(routines, helpers)
  706. self._prepare_files(routines, funcname)
  707. self._process_files(routines)
  708. mod = __import__(self.module_name)
  709. finally:
  710. sys.path.remove(workdir)
  711. CodeWrapper._module_counter += 1
  712. os.chdir(oldwork)
  713. if not self.filepath:
  714. try:
  715. shutil.rmtree(workdir)
  716. except OSError:
  717. # Could be some issues on Windows
  718. pass
  719. return self._get_wrapped_function(mod, funcname)
  720. def _generate_code(self, main_routines, helper_routines):
  721. all_routines = main_routines + helper_routines
  722. self.generator.write(
  723. all_routines, self.filename, True, self.include_header,
  724. self.include_empty)
  725. def _prepare_files(self, routines, funcname):
  726. # C
  727. codefilename = self.module_name + '.c'
  728. with open(codefilename, 'w') as f:
  729. self.dump_c(routines, f, self.filename, funcname=funcname)
  730. # setup.py
  731. with open('setup.py', 'w') as f:
  732. self.dump_setup(f)
  733. @classmethod
  734. def _get_wrapped_function(cls, mod, name):
  735. return getattr(mod, name)
  736. def dump_setup(self, f):
  737. setup = _ufunc_setup.substitute(module=self.module_name,
  738. filename=self.filename)
  739. f.write(setup)
  740. def dump_c(self, routines, f, prefix, funcname=None):
  741. """Write a C file with Python wrappers
  742. This file contains all the definitions of the routines in c code.
  743. Arguments
  744. ---------
  745. routines
  746. List of Routine instances
  747. f
  748. File-like object to write the file to
  749. prefix
  750. The filename prefix, used to name the imported module.
  751. funcname
  752. Name of the main function to be returned.
  753. """
  754. if funcname is None:
  755. if len(routines) == 1:
  756. funcname = routines[0].name
  757. else:
  758. msg = 'funcname must be specified for multiple output routines'
  759. raise ValueError(msg)
  760. functions = []
  761. function_creation = []
  762. ufunc_init = []
  763. module = self.module_name
  764. include_file = "\"{}.h\"".format(prefix)
  765. top = _ufunc_top.substitute(include_file=include_file, module=module)
  766. name = funcname
  767. # Partition the C function arguments into categories
  768. # Here we assume all routines accept the same arguments
  769. r_index = 0
  770. py_in, _ = self._partition_args(routines[0].arguments)
  771. n_in = len(py_in)
  772. n_out = len(routines)
  773. # Declare Args
  774. form = "char *{0}{1} = args[{2}];"
  775. arg_decs = [form.format('in', i, i) for i in range(n_in)]
  776. arg_decs.extend([form.format('out', i, i+n_in) for i in range(n_out)])
  777. declare_args = '\n '.join(arg_decs)
  778. # Declare Steps
  779. form = "npy_intp {0}{1}_step = steps[{2}];"
  780. step_decs = [form.format('in', i, i) for i in range(n_in)]
  781. step_decs.extend([form.format('out', i, i+n_in) for i in range(n_out)])
  782. declare_steps = '\n '.join(step_decs)
  783. # Call Args
  784. form = "*(double *)in{0}"
  785. call_args = ', '.join([form.format(a) for a in range(n_in)])
  786. # Step Increments
  787. form = "{0}{1} += {0}{1}_step;"
  788. step_incs = [form.format('in', i) for i in range(n_in)]
  789. step_incs.extend([form.format('out', i, i) for i in range(n_out)])
  790. step_increments = '\n '.join(step_incs)
  791. # Types
  792. n_types = n_in + n_out
  793. types = "{" + ', '.join(["NPY_DOUBLE"]*n_types) + "};"
  794. # Docstring
  795. docstring = '"Created in SymPy with Ufuncify"'
  796. # Function Creation
  797. function_creation.append("PyObject *ufunc{};".format(r_index))
  798. # Ufunc initialization
  799. init_form = _ufunc_init_form.substitute(module=module,
  800. funcname=name,
  801. docstring=docstring,
  802. n_in=n_in, n_out=n_out,
  803. ind=r_index)
  804. ufunc_init.append(init_form)
  805. outcalls = [_ufunc_outcalls.substitute(
  806. outnum=i, call_args=call_args, funcname=routines[i].name) for i in
  807. range(n_out)]
  808. body = _ufunc_body.substitute(module=module, funcname=name,
  809. declare_args=declare_args,
  810. declare_steps=declare_steps,
  811. call_args=call_args,
  812. step_increments=step_increments,
  813. n_types=n_types, types=types,
  814. outcalls='\n '.join(outcalls))
  815. functions.append(body)
  816. body = '\n\n'.join(functions)
  817. ufunc_init = '\n '.join(ufunc_init)
  818. function_creation = '\n '.join(function_creation)
  819. bottom = _ufunc_bottom.substitute(module=module,
  820. ufunc_init=ufunc_init,
  821. function_creation=function_creation)
  822. text = [top, body, bottom]
  823. f.write('\n\n'.join(text))
  824. def _partition_args(self, args):
  825. """Group function arguments into categories."""
  826. py_in = []
  827. py_out = []
  828. for arg in args:
  829. if isinstance(arg, OutputArgument):
  830. py_out.append(arg)
  831. elif isinstance(arg, InOutArgument):
  832. raise ValueError("Ufuncify doesn't support InOutArguments")
  833. else:
  834. py_in.append(arg)
  835. return py_in, py_out
  836. @cacheit
  837. @doctest_depends_on(exe=('f2py', 'gfortran', 'gcc'), modules=('numpy',))
  838. def ufuncify(args, expr, language=None, backend='numpy', tempdir=None,
  839. flags=None, verbose=False, helpers=None, **kwargs):
  840. """Generates a binary function that supports broadcasting on numpy arrays.
  841. Parameters
  842. ==========
  843. args : iterable
  844. Either a Symbol or an iterable of symbols. Specifies the argument
  845. sequence for the function.
  846. expr
  847. A SymPy expression that defines the element wise operation.
  848. language : string, optional
  849. If supplied, (options: 'C' or 'F95'), specifies the language of the
  850. generated code. If ``None`` [default], the language is inferred based
  851. upon the specified backend.
  852. backend : string, optional
  853. Backend used to wrap the generated code. Either 'numpy' [default],
  854. 'cython', or 'f2py'.
  855. tempdir : string, optional
  856. Path to directory for temporary files. If this argument is supplied,
  857. the generated code and the wrapper input files are left intact in
  858. the specified path.
  859. flags : iterable, optional
  860. Additional option flags that will be passed to the backend.
  861. verbose : bool, optional
  862. If True, autowrap will not mute the command line backends. This can
  863. be helpful for debugging.
  864. helpers : iterable, optional
  865. Used to define auxiliary expressions needed for the main expr. If
  866. the main expression needs to call a specialized function it should
  867. be put in the ``helpers`` iterable. Autowrap will then make sure
  868. that the compiled main expression can link to the helper routine.
  869. Items should be tuples with (<funtion_name>, <sympy_expression>,
  870. <arguments>). It is mandatory to supply an argument sequence to
  871. helper routines.
  872. kwargs : dict
  873. These kwargs will be passed to autowrap if the `f2py` or `cython`
  874. backend is used and ignored if the `numpy` backend is used.
  875. Notes
  876. =====
  877. The default backend ('numpy') will create actual instances of
  878. ``numpy.ufunc``. These support ndimensional broadcasting, and implicit type
  879. conversion. Use of the other backends will result in a "ufunc-like"
  880. function, which requires equal length 1-dimensional arrays for all
  881. arguments, and will not perform any type conversions.
  882. References
  883. ==========
  884. .. [1] http://docs.scipy.org/doc/numpy/reference/ufuncs.html
  885. Examples
  886. ========
  887. >>> from sympy.utilities.autowrap import ufuncify
  888. >>> from sympy.abc import x, y
  889. >>> import numpy as np
  890. >>> f = ufuncify((x, y), y + x**2)
  891. >>> type(f)
  892. <class 'numpy.ufunc'>
  893. >>> f([1, 2, 3], 2)
  894. array([ 3., 6., 11.])
  895. >>> f(np.arange(5), 3)
  896. array([ 3., 4., 7., 12., 19.])
  897. For the 'f2py' and 'cython' backends, inputs are required to be equal length
  898. 1-dimensional arrays. The 'f2py' backend will perform type conversion, but
  899. the Cython backend will error if the inputs are not of the expected type.
  900. >>> f_fortran = ufuncify((x, y), y + x**2, backend='f2py')
  901. >>> f_fortran(1, 2)
  902. array([ 3.])
  903. >>> f_fortran(np.array([1, 2, 3]), np.array([1.0, 2.0, 3.0]))
  904. array([ 2., 6., 12.])
  905. >>> f_cython = ufuncify((x, y), y + x**2, backend='Cython')
  906. >>> f_cython(1, 2) # doctest: +ELLIPSIS
  907. Traceback (most recent call last):
  908. ...
  909. TypeError: Argument '_x' has incorrect type (expected numpy.ndarray, got int)
  910. >>> f_cython(np.array([1.0]), np.array([2.0]))
  911. array([ 3.])
  912. """
  913. if isinstance(args, Symbol):
  914. args = (args,)
  915. else:
  916. args = tuple(args)
  917. if language:
  918. _validate_backend_language(backend, language)
  919. else:
  920. language = _infer_language(backend)
  921. helpers = helpers if helpers else ()
  922. flags = flags if flags else ()
  923. if backend.upper() == 'NUMPY':
  924. # maxargs is set by numpy compile-time constant NPY_MAXARGS
  925. # If a future version of numpy modifies or removes this restriction
  926. # this variable should be changed or removed
  927. maxargs = 32
  928. helps = []
  929. for name, expr, args in helpers:
  930. helps.append(make_routine(name, expr, args))
  931. code_wrapper = UfuncifyCodeWrapper(C99CodeGen("ufuncify"), tempdir,
  932. flags, verbose)
  933. if not isinstance(expr, (list, tuple)):
  934. expr = [expr]
  935. if len(expr) == 0:
  936. raise ValueError('Expression iterable has zero length')
  937. if len(expr) + len(args) > maxargs:
  938. msg = ('Cannot create ufunc with more than {0} total arguments: '
  939. 'got {1} in, {2} out')
  940. raise ValueError(msg.format(maxargs, len(args), len(expr)))
  941. routines = [make_routine('autofunc{}'.format(idx), exprx, args) for
  942. idx, exprx in enumerate(expr)]
  943. return code_wrapper.wrap_code(routines, helpers=helps)
  944. else:
  945. # Dummies are used for all added expressions to prevent name clashes
  946. # within the original expression.
  947. y = IndexedBase(Dummy('y'))
  948. m = Dummy('m', integer=True)
  949. i = Idx(Dummy('i', integer=True), m)
  950. f_dummy = Dummy('f')
  951. f = implemented_function('%s_%d' % (f_dummy.name, f_dummy.dummy_index), Lambda(args, expr))
  952. # For each of the args create an indexed version.
  953. indexed_args = [IndexedBase(Dummy(str(a))) for a in args]
  954. # Order the arguments (out, args, dim)
  955. args = [y] + indexed_args + [m]
  956. args_with_indices = [a[i] for a in indexed_args]
  957. return autowrap(Eq(y[i], f(*args_with_indices)), language, backend,
  958. tempdir, args, flags, verbose, helpers, **kwargs)