Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
Stephan Seitz
pystencils
Commits
b1e9b1e4
Commit
b1e9b1e4
authored
Jan 30, 2020
by
Michael Kuron
Browse files
Merge branch 'master' of i10git.cs.fau.de:pycodegen/pystencils into abs
parents
659319bd
a5ab1070
Changes
55
Hide whitespace changes
Inline
Side-by-side
.gitlab-ci.yml
View file @
b1e9b1e4
...
...
@@ -65,16 +65,21 @@ minimal-windows:
-
python -c "import numpy"
-
python setup.py quicktest
minimal-
ubuntu
:
ubuntu
:
stage
:
test
except
:
variables
:
-
$ENABLE_NIGHTLY_BUILDS
image
:
i10git.cs.fau.de:5005/pycodegen/pycodegen/
minimal_
ubuntu
image
:
i10git.cs.fau.de:5005/pycodegen/pycodegen/ubuntu
script
:
-
python3 setup.py quicktest
-
mkdir -p ~/.config/matplotlib
-
echo "backend:template" > ~/.config/matplotlib/matplotlibrc
-
sed -i 's/--doctest-modules //g' pytest.ini
-
pytest-3 -v -m "not longrun"
tags
:
-
docker
-
cuda
-
AVX
minimal-conda
:
stage
:
test
...
...
conftest.py
View file @
b1e9b1e4
import
os
import
runpy
import
sys
import
warnings
import
tempfile
import
warnings
import
nbformat
import
pytest
...
...
@@ -34,7 +34,7 @@ def add_path_to_ignore(path):
collect_ignore
=
[
os
.
path
.
join
(
SCRIPT_FOLDER
,
"doc"
,
"conf.py"
),
os
.
path
.
join
(
SCRIPT_FOLDER
,
"pystencils"
,
"opencl"
,
"opencl.autoinit"
)]
os
.
path
.
join
(
SCRIPT_FOLDER
,
"pystencils"
,
"opencl"
,
"opencl.autoinit"
)]
add_path_to_ignore
(
'pystencils_tests/benchmark'
)
add_path_to_ignore
(
'_local_tmp'
)
...
...
@@ -44,7 +44,7 @@ collect_ignore += [os.path.join(SCRIPT_FOLDER, "pystencils/autodiff.py")]
try
:
import
pycuda
except
ImportError
:
collect_ignore
+=
[
os
.
path
.
join
(
SCRIPT_FOLDER
,
"pystencils
/pystencils
_tests/test_cudagpu.py"
)]
collect_ignore
+=
[
os
.
path
.
join
(
SCRIPT_FOLDER
,
"pystencils_tests/test_cudagpu.py"
)]
add_path_to_ignore
(
'pystencils/gpucuda'
)
try
:
...
...
@@ -73,7 +73,22 @@ try:
import
blitzdb
except
ImportError
:
add_path_to_ignore
(
'pystencils/runhelper'
)
collect_ignore
+=
[
os
.
path
.
join
(
SCRIPT_FOLDER
,
"pystencils_tests/test_parameterstudy.py"
)]
try
:
import
islpy
except
ImportError
:
collect_ignore
+=
[
os
.
path
.
join
(
SCRIPT_FOLDER
,
"pystencils/integer_set_analysis.py"
)]
try
:
import
graphviz
except
ImportError
:
collect_ignore
+=
[
os
.
path
.
join
(
SCRIPT_FOLDER
,
"pystencils/backends/dot.py"
)]
try
:
import
pyevtk
except
ImportError
:
collect_ignore
+=
[
os
.
path
.
join
(
SCRIPT_FOLDER
,
"pystencils/datahandling/vtk.py"
)]
collect_ignore
+=
[
os
.
path
.
join
(
SCRIPT_FOLDER
,
'setup.py'
)]
...
...
@@ -83,8 +98,6 @@ for root, sub_dirs, files in os.walk('.'):
collect_ignore
.
append
(
f
)
class
IPythonMockup
:
def
run_line_magic
(
self
,
*
args
,
**
kwargs
):
pass
...
...
@@ -131,7 +144,7 @@ class IPyNbFile(pytest.File):
exporter
.
exclude_markdown
=
True
exporter
.
exclude_input_prompt
=
True
notebook_contents
=
self
.
fspath
.
open
()
notebook_contents
=
self
.
fspath
.
open
(
encoding
=
'utf-8'
)
with
warnings
.
catch_warnings
():
warnings
.
filterwarnings
(
"ignore"
,
"IPython.core.inputsplitter is deprecated"
)
...
...
doc/notebooks/01_tutorial_getting_started.ipynb
View file @
b1e9b1e4
...
...
@@ -441,18 +441,22 @@
Now lets grab an image to apply this filter to:
%% Cell type:code id: tags:
```
python
import
requests
import
imageio
from
io
import
BytesIO
try
:
import
requests
import
imageio
from
io
import
BytesIO
response
=
requests
.
get
(
"https://www.python.org/static/img/python-logo.png"
)
img
=
imageio
.
imread
(
BytesIO
(
response
.
content
)).
astype
(
np
.
double
)
img
/=
img
.
max
()
plt
.
imshow
(
img
);
response
=
requests
.
get
(
"https://www.python.org/static/img/python-logo.png"
)
img
=
imageio
.
imread
(
BytesIO
(
response
.
content
)).
astype
(
np
.
double
)
img
/=
img
.
max
()
plt
.
imshow
(
img
);
except
ImportError
:
print
(
"No requests installed"
)
img
=
np
.
random
.
random
((
82
,
290
,
4
))
```
%%%% Output: display_data

...
...
doc/notebooks/demo_plotting_and_animation.ipynb
View file @
b1e9b1e4
%% Cell type:code id: tags:
``` python
from pystencils.session import *
import shutil
```
%% Cell type:markdown id: tags:
# Plotting and Animation
...
...
@@ -211,13 +213,16 @@
```
%% Cell type:code id: tags:
``` python
plt.figure()
animation = plt.scalar_field_animation(run_func, frames=60)
ps.jupyter.display_as_html_video(animation)
if shutil.which("ffmpeg") is not None:
plt.figure()
animation = plt.scalar_field_animation(run_func, frames=60)
ps.jupyter.display_as_html_video(animation)
else:
print("No ffmpeg installed")
```
%%%% Output: execute_result
<IPython.core.display.HTML object>
...
...
@@ -243,12 +248,15 @@
For surface plots there is also an animated version:
%% Cell type:code id: tags:
``` python
animation = plt.surface_plot_animation(run_func, frames=60)
ps.jupyter.display_as_html_video(animation)
if shutil.which("ffmpeg") is not None:
animation = plt.surface_plot_animation(run_func, frames=60)
ps.jupyter.display_as_html_video(animation)
else:
print("No ffmpeg installed")
```
%%%% Output: execute_result
<IPython.core.display.HTML object>
...
...
@@ -332,13 +340,16 @@
```
%% Cell type:code id: tags:
``` python
plt.figure()
animation = plt.vector_field_animation(run_func, frames=60)
ps.jupyter.display_as_html_video(animation)
if shutil.which("ffmpeg") is not None:
plt.figure()
animation = plt.vector_field_animation(run_func, frames=60)
ps.jupyter.display_as_html_video(animation)
else:
print("No ffmpeg installed")
```
%%%% Output: execute_result
<IPython.core.display.HTML object>
...
...
@@ -348,12 +359,15 @@
...and magnitude plots
%% Cell type:code id: tags:
``` python
animation = plt.vector_field_magnitude_animation(run_func, frames=60)
ps.jupyter.display_as_html_video(animation)
if shutil.which("ffmpeg") is not None:
animation = plt.vector_field_magnitude_animation(run_func, frames=60)
ps.jupyter.display_as_html_video(animation)
else:
print("No ffmpeg installed")
```
%%%% Output: execute_result
<IPython.core.display.HTML object>
doc/notebooks/demo_wave_equation.ipynb
View file @
b1e9b1e4
%% Cell type:code id: tags:
```
python
from
pystencils.session
import
*
import
shutil
```
%% Cell type:markdown id: tags:
# Demo: Finite differences - 2D wave equation
...
...
@@ -183,12 +185,15 @@
Lets create an animation of the solution:
%% Cell type:code id: tags:
```
python
ani
=
plt
.
surface_plot_animation
(
run
,
zlim
=
(
-
1
,
1
))
ps
.
jupyter
.
display_as_html_video
(
ani
)
if
shutil
.
which
(
"ffmpeg"
)
is
not
None
:
ani
=
plt
.
surface_plot_animation
(
run
,
zlim
=
(
-
1
,
1
))
ps
.
jupyter
.
display_as_html_video
(
ani
)
else
:
print
(
"No ffmpeg installed"
)
```
%%%% Output: execute_result
<IPython.core.display.HTML object>
...
...
@@ -230,13 +235,16 @@
for
t
in
range
(
timesteps
):
kernel
(
u0
=
u_arrays
[
0
],
u1
=
u_arrays
[
1
],
u2
=
u_arrays
[
2
])
u_arrays
[
0
],
u_arrays
[
1
],
u_arrays
[
2
]
=
u_arrays
[
1
],
u_arrays
[
2
],
u_arrays
[
0
]
return
u_arrays
[
2
]
ani
=
plt
.
surface_plot_animation
(
run_LLVM
,
zlim
=
(
-
1
,
1
))
assert
np
.
isfinite
(
np
.
max
(
u_arrays
[
2
]))
ps
.
jupyter
.
display_as_html_video
(
ani
)
if
shutil
.
which
(
"ffmpeg"
)
is
not
None
:
ani
=
plt
.
surface_plot_animation
(
run_LLVM
,
zlim
=
(
-
1
,
1
))
ps
.
jupyter
.
display_as_html_video
(
ani
)
else
:
print
(
"No ffmpeg installed"
)
```
%%%% Output: execute_result
<IPython.core.display.HTML object>
...
...
pystencils/__init__.py
View file @
b1e9b1e4
...
...
@@ -5,7 +5,7 @@ from . import stencil as stencil
from
.assignment
import
Assignment
,
assignment_from_stencil
from
.data_types
import
TypedSymbol
from
.datahandling
import
create_data_handling
from
.display_utils
import
show_code
,
to_dot
from
.display_utils
import
show_code
,
get_code_obj
,
get_code_str
,
to_dot
from
.field
import
Field
,
FieldType
,
fields
from
.kernel_decorator
import
kernel
from
.kernelcreation
import
create_indexed_kernel
,
create_kernel
,
create_staggered_kernel
...
...
@@ -25,7 +25,7 @@ __all__ = ['Field', 'FieldType', 'fields',
'TypedSymbol'
,
'make_slice'
,
'create_kernel'
,
'create_indexed_kernel'
,
'create_staggered_kernel'
,
'show_code'
,
'to_dot'
,
'show_code'
,
'to_dot'
,
'get_code_obj'
,
'get_code_str'
,
'AssignmentCollection'
,
'Assignment'
,
'assignment_from_stencil'
,
...
...
pystencils/assignment.py
View file @
b1e9b1e4
# -*- coding: utf-8 -*-
import
numpy
as
np
import
sympy
as
sp
from
sympy.printing.latex
import
LatexPrinter
...
...
@@ -24,9 +23,20 @@ def assignment_str(assignment):
if
Assignment
:
_old_new
=
sp
.
codegen
.
ast
.
Assignment
.
__new__
def
_Assignment__new__
(
cls
,
lhs
,
rhs
,
*
args
,
**
kwargs
):
if
isinstance
(
lhs
,
(
list
,
tuple
,
sp
.
Matrix
))
and
isinstance
(
rhs
,
(
list
,
tuple
,
sp
.
Matrix
)):
assert
len
(
lhs
)
==
len
(
rhs
),
f
'
{
lhs
}
and
{
rhs
}
must have same length when performing vector assignment!'
return
tuple
(
_old_new
(
cls
,
a
,
b
,
*
args
,
**
kwargs
)
for
a
,
b
in
zip
(
lhs
,
rhs
))
return
_old_new
(
cls
,
lhs
,
rhs
,
*
args
,
**
kwargs
)
Assignment
.
__str__
=
assignment_str
Assignment
.
__new__
=
_Assignment__new__
LatexPrinter
.
_print_Assignment
=
print_assignment_latex
sp
.
MutableDenseMatrix
.
__hash__
=
lambda
self
:
hash
(
tuple
(
self
))
else
:
# back port for older sympy versions that don't have Assignment yet
...
...
pystencils/astnodes.py
View file @
b1e9b1e4
...
...
@@ -110,10 +110,17 @@ class Conditional(Node):
return
result
def
__str__
(
self
):
return
'if:({!s}) '
.
format
(
self
.
condition_expr
)
return
self
.
__repr__
(
)
def
__repr__
(
self
):
return
'if:({!r}) '
.
format
(
self
.
condition_expr
)
repr
=
'if:({!r}) '
.
format
(
self
.
condition_expr
)
if
self
.
true_block
:
repr
+=
'
\n\t
{}) '
.
format
(
self
.
true_block
)
if
self
.
false_block
:
repr
=
'else: '
.
format
(
self
.
false_block
)
repr
+=
'
\n\t
{} '
.
format
(
self
.
false_block
)
return
repr
def
replace_by_true_block
(
self
):
"""Replaces the conditional by its True block"""
...
...
@@ -284,7 +291,10 @@ class Block(Node):
self
.
_nodes
=
nodes
self
.
parent
=
None
for
n
in
self
.
_nodes
:
n
.
parent
=
self
try
:
n
.
parent
=
self
except
AttributeError
:
pass
@
property
def
args
(
self
):
...
...
pystencils/backends/cbackend.py
View file @
b1e9b1e4
import
re
from
collections
import
namedtuple
from
typing
import
Set
...
...
@@ -10,11 +11,12 @@ from sympy.printing.ccode import C89CodePrinter
from
pystencils.astnodes
import
KernelFunction
,
Node
from
pystencils.cpu.vectorization
import
vec_all
,
vec_any
from
pystencils.data_types
import
(
PointerType
,
VectorType
,
address_of
,
cast_func
,
create_type
,
get_type_of_expression
,
reinterpret_cast_func
,
vector_memory_access
)
PointerType
,
VectorType
,
address_of
,
cast_func
,
create_type
,
get_type_of_expression
,
reinterpret_cast_func
,
vector_memory_access
)
from
pystencils.fast_approximation
import
fast_division
,
fast_inv_sqrt
,
fast_sqrt
from
pystencils.integer_functions
import
(
bit_shift_left
,
bit_shift_right
,
bitwise_and
,
bitwise_or
,
bitwise_xor
,
int_div
,
int_power_of_2
,
modulo_ceil
)
bit_shift_left
,
bit_shift_right
,
bitwise_and
,
bitwise_or
,
bitwise_xor
,
int_div
,
int_power_of_2
,
modulo_ceil
)
try
:
from
sympy.printing.ccode
import
C99CodePrinter
as
CCodePrinter
...
...
@@ -23,6 +25,9 @@ except ImportError:
__all__
=
[
'generate_c'
,
'CustomCodeNode'
,
'PrintNode'
,
'get_headers'
,
'CustomSympyPrinter'
]
HEADER_REGEX
=
re
.
compile
(
r
'^[<"].*[">]$'
)
KERNCRAFT_NO_TERNARY_MODE
=
False
...
...
@@ -91,7 +96,7 @@ def get_global_declarations(ast):
visit_node
(
ast
)
return
sorted
(
set
(
global_declarations
),
key
=
lambda
x
:
str
(
x
)
)
return
sorted
(
set
(
global_declarations
),
key
=
str
)
def
get_headers
(
ast_node
:
Node
)
->
Set
[
str
]:
...
...
@@ -111,6 +116,9 @@ def get_headers(ast_node: Node) -> Set[str]:
if
isinstance
(
g
,
Node
):
headers
.
update
(
get_headers
(
g
))
for
h
in
headers
:
assert
HEADER_REGEX
.
match
(
h
),
f
'header /
{
h
}
/ does not follow the pattern /"..."/ or /<...>/'
return
sorted
(
headers
)
...
...
@@ -392,6 +400,13 @@ class CustomSympyPrinter(CCodePrinter):
return
self
.
_print
(
expr
.
args
[
0
])
elif
isinstance
(
expr
,
fast_inv_sqrt
):
return
"({})"
.
format
(
self
.
_print
(
1
/
sp
.
sqrt
(
expr
.
args
[
0
])))
elif
isinstance
(
expr
,
sp
.
Abs
):
return
"abs({})"
.
format
(
self
.
_print
(
expr
.
args
[
0
]))
elif
isinstance
(
expr
,
sp
.
Mod
):
if
expr
.
args
[
0
].
is_integer
and
expr
.
args
[
1
].
is_integer
:
return
"({} % {})"
.
format
(
self
.
_print
(
expr
.
args
[
0
]),
self
.
_print
(
expr
.
args
[
1
]))
else
:
return
"fmod({}, {})"
.
format
(
self
.
_print
(
expr
.
args
[
0
]),
self
.
_print
(
expr
.
args
[
1
]))
elif
expr
.
func
in
infix_functions
:
return
"(%s %s %s)"
%
(
self
.
_print
(
expr
.
args
[
0
]),
infix_functions
[
expr
.
func
],
self
.
_print
(
expr
.
args
[
1
]))
elif
expr
.
func
==
int_power_of_2
:
...
...
pystencils/backends/cuda_backend.py
View file @
b1e9b1e4
...
...
@@ -3,7 +3,7 @@ from os.path import dirname, join
from
pystencils.astnodes
import
Node
from
pystencils.backends.cbackend
import
CBackend
,
CustomSympyPrinter
,
generate_c
from
pystencils.fast_approximation
import
fast_division
,
fast_inv_sqrt
,
fast_sqrt
from
pystencils.interpolation_astnodes
import
InterpolationMode
from
pystencils.interpolation_astnodes
import
DiffInterpolatorAccess
,
InterpolationMode
with
open
(
join
(
dirname
(
__file__
),
'cuda_known_functions.txt'
))
as
f
:
lines
=
f
.
readlines
()
...
...
@@ -70,10 +70,13 @@ class CudaSympyPrinter(CustomSympyPrinter):
super
(
CudaSympyPrinter
,
self
).
__init__
()
self
.
known_functions
.
update
(
CUDA_KNOWN_FUNCTIONS
)
def
_print_
Texture
Access
(
self
,
node
):
dtype
=
node
.
texture
.
field
.
dtype
.
numpy_dtype
def
_print_
Interpolator
Access
(
self
,
node
):
dtype
=
node
.
interpolator
.
field
.
dtype
.
numpy_dtype
if
node
.
texture
.
interpolation_mode
==
InterpolationMode
.
CUBIC_SPLINE
:
if
type
(
node
)
==
DiffInterpolatorAccess
:
# cubicTex3D_1st_derivative_x(texture tex, float3 coord)
template
=
f
"cubicTex%iD_1st_derivative_
{
list
(
reversed
(
'xyz'
[:
node
.
ndim
]))[
node
.
diff_coordinate_idx
]
}
(%s, %s)"
# noqa
elif
node
.
interpolator
.
interpolation_mode
==
InterpolationMode
.
CUBIC_SPLINE
:
template
=
"cubicTex%iDSimple(%s, %s)"
else
:
if
dtype
.
itemsize
>
4
:
...
...
@@ -84,8 +87,8 @@ class CudaSympyPrinter(CustomSympyPrinter):
template
=
"tex%iD(%s, %s)"
code
=
template
%
(
node
.
texture
.
field
.
spatial_dimensions
,
str
(
node
.
texture
),
node
.
interpolator
.
field
.
spatial_dimensions
,
str
(
node
.
interpolator
),
# + 0.5 comes from Nvidia's staggered indexing
', '
.
join
(
self
.
_print
(
o
+
0.5
)
for
o
in
reversed
(
node
.
offsets
))
)
...
...
pystencils/boundaries/boundaryhandling.py
View file @
b1e9b1e4
...
...
@@ -8,6 +8,8 @@ from pystencils.boundaries.createindexlist import (
create_boundary_index_array
,
numpy_data_type_for_boundary_object
)
from
pystencils.cache
import
memorycache
from
pystencils.data_types
import
TypedSymbol
,
create_type
from
pystencils.datahandling
import
ParallelDataHandling
from
pystencils.datahandling.pycuda
import
PyCudaArrayHandler
from
pystencils.field
import
Field
from
pystencils.kernelparameters
import
FieldPointerSymbol
...
...
@@ -96,16 +98,23 @@ class BoundaryHandling:
def
to_gpu
(
gpu_version
,
cpu_version
):
gpu_version
=
gpu_version
.
boundary_object_to_index_list
cpu_version
=
cpu_version
.
boundary_object_to_index_list
if
ParallelDataHandling
and
isinstance
(
self
.
data_handling
,
ParallelDataHandling
):
array_handler
=
PyCudaArrayHandler
()
else
:
array_handler
=
self
.
data_handling
.
array_handler
for
obj
,
cpu_arr
in
cpu_version
.
items
():
if
obj
not
in
gpu_version
or
gpu_version
[
obj
].
shape
!=
cpu_arr
.
shape
:
gpu_version
[
obj
]
=
self
.
data_handling
.
array_handler
.
to_gpu
(
cpu_arr
)
gpu_version
[
obj
]
=
array_handler
.
to_gpu
(
cpu_arr
)
else
:
self
.
data_handling
.
array_handler
.
upload
(
gpu_version
[
obj
],
cpu_arr
)
array_handler
.
upload
(
gpu_version
[
obj
],
cpu_arr
)
class_
=
self
.
IndexFieldBlockData
class_
.
to_cpu
=
to_cpu
class_
.
to_gpu
=
to_gpu
data_handling
.
add_custom_class
(
self
.
_index_array_name
,
class_
)
gpu
=
self
.
_target
in
data_handling
.
_GPU_LIKE_TARGETS
data_handling
.
add_custom_class
(
self
.
_index_array_name
,
class_
,
cpu
=
True
,
gpu
=
gpu
)
@
property
def
data_handling
(
self
):
...
...
@@ -253,11 +262,13 @@ class BoundaryHandling:
"""
Writes a VTK field where each cell with the given boundary is marked with 1, other cells are 0
This can be used to display the simulation geometry in Paraview
:param file_name: vtk filename
:param boundaries: boundary object, or special string 'domain' for domain cells or special string 'all' for all
boundary conditions.
can also be a sequence, to write multiple boundaries to VTK file
:param ghost_layers: number of ghost layers to write, or True for all, False for none
Params:
file_name: vtk filename
boundaries: boundary object, or special string 'domain' for domain cells or special string 'all' for all
boundary conditions.
can also be a sequence, to write multiple boundaries to VTK file
ghost_layers: number of ghost layers to write, or True for all, False for none
"""
if
boundaries
==
'all'
:
boundaries
=
list
(
self
.
_boundary_object_to_boundary_info
.
keys
())
+
[
'domain'
]
...
...
@@ -329,7 +340,7 @@ class BoundaryHandling:
self
.
kernel
=
kernel
class
IndexFieldBlockData
:
def
__init__
(
self
):
def
__init__
(
self
,
*
args
,
**
kwargs
):
self
.
boundary_object_to_index_list
=
{}
self
.
boundary_object_to_data_setter
=
{}
...
...
pystencils/cpu/cpujit.py
View file @
b1e9b1e4
...
...
@@ -270,7 +270,8 @@ if( PyErr_Occurred() ) {{ return NULL; }}
template_extract_complex
=
"""
PyObject * obj_{name} = PyDict_GetItemString(kwargs, "{name}");
if( obj_{name} == NULL) {{ PyErr_SetString(PyExc_TypeError, "Keyword argument '{name}' missing"); return NULL; }};
{target_type} {name}{{ {extract_function_real}( obj_{name} ), {extract_function_imag}( obj_{name} ) }};
{target_type} {name}{{ ({real_type}) {extract_function_real}( obj_{name} ),
({real_type}) {extract_function_imag}( obj_{name} ) }};
if( PyErr_Occurred() ) {{ return NULL; }}
"""
...
...
@@ -409,6 +410,8 @@ def create_function_boilerplate_code(parameter_info, name, insert_checks=True):
pre_call_code
+=
template_extract_complex
.
format
(
extract_function_real
=
extract_function
[
0
],
extract_function_imag
=
extract_function
[
1
],
target_type
=
target_type
,
real_type
=
"float"
if
target_type
==
"ComplexFloat"
else
"double"
,
name
=
param
.
symbol
.
name
)
else
:
pre_call_code
+=
template_extract_scalar
.
format
(
extract_function
=
extract_function
,
...
...
pystencils/data_types.py
View file @
b1e9b1e4
...
...
@@ -575,9 +575,7 @@ def get_type_of_expression(expr,
raise
NotImplementedError
(
"Could not determine type for"
,
expr
,
type
(
expr
))
class
Type
(
sp
.
Basic
):
is_Atom
=
True
class
Type
(
sp
.
Atom
):
def
__new__
(
cls
,
*
args
,
**
kwargs
):
return
sp
.
Basic
.
__new__
(
cls
)
...
...
pystencils/display_utils.py
View file @
b1e9b1e4
from
typing
import
Any
,
Dict
,
Optional
from
typing
import
Any
,
Dict
,
Optional
,
Union
import
sympy
as
sp
...
...
@@ -35,10 +35,10 @@ def highlight_cpp(code: str):
return
HTML
(
highlight
(
code
,
CppLexer
(),
HtmlFormatter
()))
def
show
_code
(
ast
:
KernelFunction
,
custom_backend
=
None
):
def
get
_code
_obj
(
ast
:
Union
[
KernelFunction
,
KernelWrapper
],
custom_backend
=
None
):
"""Returns an object to display generated code (C/C++ or CUDA)
Can either
be displayed as HTML in Jupyter notebooks or printed as normal string.
Can either be displayed as HTML in Jupyter notebooks or printed as normal string.
"""
from
pystencils.backends.cbackend
import
generate_c
...
...
@@ -65,3 +65,37 @@ def show_code(ast: KernelFunction, custom_backend=None):
def
__repr__
(
self
):
return
generate_c
(
self
.
ast
,
dialect
=
dialect
,
custom_backend
=
custom_backend
)
return
CodeDisplay
(
ast
)
def
get_code_str
(
ast
,
custom_backend
=
None
):
return
str
(
get_code_obj
(
ast
,
custom_backend
))
def
_isnotebook
():
try
:
shell
=
get_ipython
().
__class__
.
__name__
if
shell
==
'ZMQInteractiveShell'
:
return
True
# Jupyter notebook or qtconsole
elif
shell
==
'TerminalInteractiveShell'
:
return
False
# Terminal running IPython
else
:
return
False
# Other type (?)
except
NameError
:
return
False
def
show_code
(
ast
:
Union
[
KernelFunction
,
KernelWrapper
],
custom_backend
=
None
):
code
=
get_code_obj
(
ast
,
custom_backend
)
if
_isnotebook
():
from
IPython.display
import
display
display
(
code
)
else
:
try
:
import
rich.syntax
import
rich.console
syntax
=
rich
.
syntax
.
Syntax
(
str
(
code
),
"c++"
,
theme
=
"monokai"
,
line_numbers
=
True
)
console
=
rich
.
console
.
Console
()
console
.
print
(
syntax
)
except
ImportError
:
print
(
code
)
pystencils/fd/derivation.py
View file @
b1e9b1e4
import
itertools
import
warnings
from
collections
import
defaultdict
import
itertools
import
numpy
as
np
import
sympy
as
sp
...
...
@@ -184,6 +184,9 @@ class FiniteDifferenceStencilDerivation:
result
[
max_offset
-
direction
[
1
],
max_offset
+
direction
[
0
]]
=
weight
return
result
def
__array__
(
self
):
return
np
.
array
(
self
.
as_array
().
tolist
())
def
as_array
(
self
):
dim
=
len
(
self
.
stencil
[
0
])
assert
(
dim
==
2
or
dim
==
3
),
"Only 2D or 3D matrix representations are available"
...
...
@@ -205,12 +208,12 @@ class FiniteDifferenceStencilDerivation: