Add pystencils.autodiff

README.md

@@ -55,6 +55,7 @@ All options:
 - `alltrafos`: pulls in additional dependencies for loop simplification e.g. libisl
 - `bench_db`: functionality to store benchmark result in object databases
 - `interactive`: installs dependencies to work in Jupyter including image I/O, plotting etc.
+- `autodiff`: enable derivation of adjoint kernels and generation of Torch/Tensorflow operations
 - `doc`: packages to build documentation
 
 Options can be combined e.g.

pystencils/__init__.py

@@ -13,6 +13,12 @@ from .simp import AssignmentCollection
 from .slicing import make_slice
 from .sympyextensions import SymbolCreator
 
+try:
+    import pystencils_autodiff
+    autodiff = pystencils_autodiff
+except ImportError:
+    pass
+
 __all__ = ['Field', 'FieldType', 'fields',
            'TypedSymbol',
            'make_slice',

pystencils/autodiff.py

+"""
+Provides tools for generation of auto-differentiable operations.
+
+See https://github.com/theHamsta/pystencils_autodiff
+
+Installation:
+
+.. code-block:: bash
+    pip install pystencils-autodiff
+"""
+import os
+
+if 'CI' not in os.environ:
+    raise NotImplementedError('pystencils-autodiff is not installed. Run `pip install pystencils-autodiff`')

setup.py

+import distutils
+import io
 import os
 import sys
-import io
-from setuptools import setup, find_packages
-import distutils
-from distutils.extension import Extension
 from contextlib import redirect_stdout
+from distutils.extension import Extension
 from importlib import import_module
 
+from setuptools import find_packages, setup
+
 if '--use-cython' in sys.argv:
     USE_CYTHON = True
     sys.argv.remove('--use-cython')
 else:
     USE_CYTHON = False
 
-
 quick_tests = [
     'test_datahandling.test_kernel',
     'test_blocking_staggered.test_blocking_staggered',
@@ -52,6 +52,7 @@ class SimpleTestRunner(distutils.cmd.Command):
         for test in quick_tests:
             self._run_tests_in_module(test)
 
+
 def readme():
     with open('README.md') as f:
         return f.read()
@@ -69,13 +70,13 @@ def cython_extensions(*extensions):
 try:
     sys.path.insert(0, os.path.abspath('doc'))
     from version_from_git import version_number_from_git
-    version=version_number_from_git()
+
+    version = version_number_from_git()
     with open("RELEASE-VERSION", "w") as f:
         f.write(version)
 except ImportError:
     version = open('RELEASE-VERSION', 'r').read()
 
-
 setup(name='pystencils',
       description='Speeding up stencil computations on CPUs and GPUs',
       version=version,
@@ -88,7 +89,7 @@ setup(name='pystencils',
       packages=['pystencils'] + ['pystencils.' + s for s in find_packages('pystencils')],
       install_requires=['sympy>=1.1', 'numpy', 'appdirs', 'joblib'],
       package_data={'pystencils': ['include/*.h', 'backends/cuda_known_functions.txt']},
-      ext_modules = cython_extensions("pystencils.boundaries.createindexlistcython"),
+      ext_modules=cython_extensions("pystencils.boundaries.createindexlistcython"),
       classifiers=[
           'Development Status :: 4 - Beta',
           'Framework :: Jupyter',
@@ -108,6 +109,7 @@ setup(name='pystencils',
           'alltrafos': ['islpy', 'py-cpuinfo'],
           'bench_db': ['blitzdb', 'pymongo', 'pandas'],
           'interactive': ['matplotlib', 'ipy_table', 'imageio', 'jupyter', 'pyevtk'],
+          'autodiff': ['pystencils-autodiff'],
           'doc': ['sphinx', 'sphinx_rtd_theme', 'nbsphinx',
                   'sphinxcontrib-bibtex', 'sphinx_autodoc_typehints', 'pandoc'],
       },