mirror of
https://github.com/triqs/dft_tools
synced 2024-12-26 06:14:14 +01:00
f1c32012a6
- for all functions, except when GIL option (not implemented) is here we use a triqs::py_stream, which redirect the stream to python sys.write - so that a simple C++ code with a std::cout print its output in the notebook... - NB : it only works for the code within the cell. If it calls another function on the lib which uses cout, print is not redirected. Designed for simple tests.... - cerr not (yet) redirected. (useful ?).
230 lines
7.5 KiB
Python
230 lines
7.5 KiB
Python
"""
|
|
=====================
|
|
Triqs magic
|
|
=====================
|
|
|
|
{TRIQS_DOC}
|
|
|
|
"""
|
|
import imp,os,sys,subprocess, hashlib,re
|
|
from IPython.core.error import UsageError
|
|
from IPython.core.magic import Magics, magics_class, line_magic, cell_magic
|
|
from IPython.core import display, magic_arguments
|
|
from IPython.utils import py3compat
|
|
from IPython.utils.io import capture_output
|
|
from IPython.utils.path import get_ipython_cache_dir
|
|
|
|
__version__ = '0.1.0'
|
|
|
|
triqs_path = "@CMAKE_INSTALL_PREFIX@"
|
|
cxx_compiler = "@CMAKE_CXX_COMPILER@"
|
|
converter_include_path = triqs_path + "/include/pytriqs/converters/"
|
|
generator_path = triqs_path + "/share/triqs/wrap_generator"
|
|
|
|
cmakelist = """
|
|
list(APPEND CMAKE_MODULE_PATH %s/share/triqs/cmake)
|
|
cmake_minimum_required(VERSION 2.8)
|
|
project(triqs_magic CXX)
|
|
set(CMAKE_BUILD_TYPE Release)
|
|
option(BUILD_SHARED_LIBS "Build shared libraries" ON)
|
|
find_package(TRIQS REQUIRED)
|
|
include_directories(${CMAKE_SOURCE_DIR} ${TRIQS_INCLUDE_ALL})
|
|
add_library(ext MODULE ext_wrap.cpp)
|
|
#add_library(ext MODULE ext.cpp ext_wrap.cpp)
|
|
set_target_properties(ext PROPERTIES PREFIX "") #eliminate the lib in front of the module name
|
|
target_link_libraries(ext ${TRIQS_LIBRARY_ALL})
|
|
triqs_set_rpath_for_target(ext)
|
|
"""%triqs_path
|
|
|
|
mod_dict = {
|
|
'gf' :
|
|
"""
|
|
module.use_module('gf')
|
|
module.add_using('namespace triqs::gfs')
|
|
""",
|
|
'parameters' :
|
|
"""
|
|
module.use_module('parameters')
|
|
module.add_using('namespace triqs::params')
|
|
""",}
|
|
|
|
def desc_file(c_decl, use) :
|
|
s = """
|
|
from wrap_generator import *
|
|
module = module_(full_name = "ext", doc = "")
|
|
module.add_include("<triqs/arrays.hpp>")
|
|
module.add_include("<triqs/python_tools/py_stream.hpp>")
|
|
module.add_using("namespace triqs::arrays")
|
|
module.add_preamble('#include "./ext.cpp"')
|
|
"""
|
|
|
|
mod_list = sum((x.split(',') for x in use), [])
|
|
for m in mod_list :
|
|
s+= mod_dict[m]
|
|
|
|
for f in c_decl :
|
|
s+= 'module.add_function("%s")\n'%f
|
|
|
|
return s + "module.generate_code()\n"
|
|
|
|
@magics_class
|
|
class TriqsMagics(Magics):
|
|
|
|
def __init__(self, shell):
|
|
super(TriqsMagics, self).__init__(shell=shell)
|
|
self._reloads = {}
|
|
self._code_cache = {}
|
|
self._lib_dir = os.path.join(get_ipython_cache_dir(), 'triqs')
|
|
if not os.path.exists(self._lib_dir):
|
|
os.makedirs(self._lib_dir)
|
|
|
|
def _import_all(self, module, verbosity=0):
|
|
imported = []
|
|
for k, v in module.__dict__.items():
|
|
if not k.startswith('__'):
|
|
self.shell.push({k: v})
|
|
imported.append(k)
|
|
if verbosity > 0 and imported:
|
|
print("\nOk. The following objects are ready to use: %s" % ", ".join(imported))
|
|
|
|
def extract_signature(self,code):
|
|
acc = []
|
|
for l in code.splitlines() :
|
|
ll = l.strip()
|
|
if ll.startswith("#") or ll.startswith("using") : continue
|
|
acc.append(l)
|
|
if l.count('{') : break
|
|
s = ''.join(acc).split('{',1)[0] # everything before {
|
|
#print "signature", s
|
|
return [s]
|
|
|
|
@magic_arguments.magic_arguments()
|
|
@magic_arguments.argument(
|
|
"-v", "--verbosity", action="count", default=0,
|
|
help="increase output verbosity"
|
|
)
|
|
@magic_arguments.argument(
|
|
'-u', "--use", action='append', default=[],
|
|
help="""Modules used"""
|
|
)
|
|
@cell_magic
|
|
def triqs(self, line, cell=None):
|
|
"""Compile and import everything from a Triqs code cell.
|
|
|
|
The content of the cell is written to a `.cpp` file in the
|
|
directory `IPYTHONDIR/triqs` using a filename with the hash of the
|
|
code. A python wrapper is then generated from the function signature.
|
|
Those files are then compiled.
|
|
The resulting module is imported and all of its symbols are injected into the user's namespace.
|
|
|
|
Usage
|
|
=====
|
|
Prepend ``%%triqs`` to your triqs code in a cell::
|
|
|
|
``%%triqs
|
|
|
|
! put your code here.
|
|
``
|
|
"""
|
|
try:
|
|
# custom saved arguments
|
|
saved_defaults = vars(
|
|
magic_arguments.parse_argstring(self.triqs,
|
|
self.shell.db['triqs']))
|
|
self.triqs.parser.set_defaults(**saved_defaults)
|
|
except KeyError:
|
|
saved_defaults = {'verbosity': 0}
|
|
|
|
if '-v' in line:
|
|
self.triqs.parser.set_defaults(verbosity=0)
|
|
|
|
## Add GIL argument ?
|
|
use_GIL = False
|
|
args = magic_arguments.parse_argstring(self.triqs, line)
|
|
|
|
code = cell if cell.endswith('\n') else cell + '\n'
|
|
|
|
#if not GIL, we replace std::cout by triqs::py_out for capture in the notebook
|
|
if not use_GIL :
|
|
code = re.sub("std::cout", "triqs::py_stream()", code)
|
|
|
|
key = code, line, sys.version_info, sys.executable
|
|
|
|
c_decl = self.extract_signature(code)
|
|
if args.verbosity>1 :
|
|
print "Found function of signature"
|
|
for f in c_decl :
|
|
print f
|
|
|
|
module_name = "ext"
|
|
module_dirname = os.path.join(self._lib_dir, "_triqs_magic_" + hashlib.md5(str(key).encode('utf-8')).hexdigest())
|
|
module_path = os.path.join(module_dirname, 'ext.so')
|
|
try :
|
|
os.mkdir(module_dirname)
|
|
except :
|
|
pass
|
|
#print "dir", module_dirname
|
|
|
|
old_cwd = os.getcwd()
|
|
try:
|
|
os.chdir(module_dirname)
|
|
|
|
with open('ext_desc.py', 'w') as f:
|
|
f.write(desc_file(c_decl, args.use))
|
|
|
|
with open('CMakeLists.txt', 'w') as f:
|
|
f.write(cmakelist)
|
|
|
|
with open('ext.cpp', 'w') as f:
|
|
f.write(code)
|
|
|
|
# Call the wrapper generator
|
|
command_generator = """
|
|
PYTHONPATH={generator_path} python ext_desc.py {generator_path}/wrapper.mako.cpp ext_wrap.cpp {generator_path}/py_converter_wrapper.mako.hpp conv.hpp {converter_include_path}
|
|
""".format(**globals())
|
|
|
|
try :
|
|
out_generator = subprocess.check_output(command_generator, stderr=subprocess.STDOUT, shell=True)
|
|
except subprocess.CalledProcessError as E :
|
|
print '---------- Wrapper generator error -------\n' + E.output
|
|
return
|
|
|
|
if args.verbosity>0 : print "---------- Wrapper generator ------", out_generator
|
|
|
|
# Call cmake
|
|
#command_cmake = "CXX =" + cxx_compiler+ " cmake . -DTRIQS_PATH=" + triqs_path
|
|
command_cmake = "cmake . -DCMAKE_CXX_COMPILER="+ cxx_compiler+ " -DTRIQS_PATH=" + triqs_path
|
|
|
|
try :
|
|
out_cmake = subprocess.check_output(command_cmake, stderr=subprocess.STDOUT, shell=True)
|
|
except subprocess.CalledProcessError as E :
|
|
print '---------- Cmake error -------\n' + E.output
|
|
return
|
|
|
|
if args.verbosity>0 : print "---------- Cmake ------", out_cmake
|
|
|
|
# Call make
|
|
command_make = "make -j2"
|
|
|
|
try :
|
|
out_make = subprocess.check_output(command_make, stderr=subprocess.STDOUT, shell=True)
|
|
except subprocess.CalledProcessError as E :
|
|
print '---------- Make error -------\n' + E.output
|
|
return
|
|
|
|
if args.verbosity>0 : print "---------- Make ------", out_make
|
|
|
|
finally:
|
|
os.chdir(old_cwd)
|
|
|
|
self._code_cache[key] = module_path
|
|
module = imp.load_dynamic(module_name, module_path)
|
|
self._import_all(module, verbosity=args.verbosity)
|
|
|
|
__doc__ = __doc__.format(TRIQS_DOC=' ' * 8 + TriqsMagics.triqs.__doc__)
|
|
|
|
def load_ipython_extension(ip):
|
|
"""Load the extension in IPython."""
|
|
ip.register_magics(TriqsMagics)
|
|
|