def test_dynamic_matrix_location_dependent():
try:
from pystencils.data_types import TypedMatrixSymbol
except ImportError:
import pytest
pytest.skip()
x, y = pystencils.fields('x, y: float32[3d]')
A = TypedMatrixSymbol('A', 3, 1, create_type('double'),
CustomCppType('Vector3<double>'))
my_fun_call = DynamicFunction(
TypedSymbol('my_fun', 'std::function<Vector3<double>(int, int, int)>'),
A.dtype, *pystencils.x_vector(3))
assignments = pystencils.AssignmentCollection({
A: my_fun_call,
y.center: A[0] + A[1] + A[2]
})
ast = pystencils.create_kernel(assignments)
pystencils.show_code(ast, custom_backend=FrameworkIntegrationPrinter())
my_fun_call = DynamicFunction(
TypedSymbol('my_fun', TemplateType('Functor_T')), A.dtype,
*pystencils.x_vector(3))
assignments = pystencils.AssignmentCollection({
A: my_fun_call,
y.center: A[0] + A[1] + A[2]
})
ast = pystencils.create_kernel(assignments)
pystencils.show_code(ast, custom_backend=FrameworkIntegrationPrinter())
def add_fixed_constant_boundary_handling(assignments, with_cse=True):
field_accesses = set().union(
itertools.chain.from_iterable(
[a.atoms(Field.Access) for a in assignments]))
if all(all(o == 0 for o in a.offsets) for a in field_accesses):
return assignments
common_shape = next(iter(field_accesses)).field.spatial_shape
ndim = len(common_shape)
def is_out_of_bound(access, shape):
return sp.Or(*[sp.Or(a < 0, a >= s) for a, s in zip(access, shape)])
safe_assignments = [
pystencils.Assignment(
assignment.lhs,
assignment.rhs.subs({
a: ConditionalFieldAccess(
a,
is_out_of_bound(
sp.Matrix(a.offsets) + x_vector(ndim), common_shape))
for a in assignment.rhs.atoms(Field.Access)
if not a.is_absolute_access
})) for assignment in assignments.all_assignments
]
if with_cse:
safe_assignments = sympy_cse(
pystencils.AssignmentCollection(safe_assignments))
return safe_assignments
else:
return pystencils.AssignmentCollection(safe_assignments)
def downsample(input: {'field_type': pystencils.field.FieldType.CUSTOM},
result, factor):
assert input.spatial_dimensions == result.spatial_dimensions
assert input.index_shape == result.index_shape
ndim = input.spatial_dimensions
assignments = AssignmentCollection({
result.center:
input.absolute_access(factor * pystencils.x_vector(ndim), ())
})
def create_autodiff(self, constant_fields=None, **kwargs):
backward_assignments = upsample(AdjointField(result),
AdjointField(input), factor)
self._autodiff = pystencils.autodiff.AutoDiffOp(
assignments,
"",
backward_assignments=backward_assignments,
**kwargs)
assignments._create_autodiff = types.MethodType(create_autodiff,
assignments)
return assignments
def upsample(input: {'field_type': pystencils.field.FieldType.CUSTOM}, result,
factor):
ndim = input.spatial_dimensions
here = pystencils.x_vector(ndim)
assignments = AssignmentCollection({
result.center:
pystencils.astnodes.ConditionalFieldAccess(
input.absolute_access(
tuple(
cast_func(sympy.S(1) / factor * h, create_type('int64'))
for h in here), ()),
sympy.Or(*[s % cast_func(factor, 'int64') > 0 for s in here]))
})
def create_autodiff(self, constant_fields=None, **kwargs):
backward_assignments = downsample(AdjointField(result),
AdjointField(input), factor)
self._autodiff = pystencils.autodiff.AutoDiffOp(
assignments,
"",
backward_assignments=backward_assignments,
**kwargs)
assignments._create_autodiff = types.MethodType(create_autodiff,
assignments)
return assignments
def max_pooling(input: {'field_type': FieldType.CUSTOM},
result):
assert input.spatial_dimensions == result.spatial_dimensions
assert input.index_shape == result.index_shape
assignments = []
ndim = input.spatial_dimensions
offsets = itertools.product((0, 1), repeat=ndim)
assignments.append(
pystencils.Assignment(result.center,
sympy.Max(*[input.absolute_access(2 * pystencils.x_vector(ndim) + sympy.Matrix(offset), ()) # noqa
for offset in offsets])
)
)
return assignments
def forward_projection(volume: pystencils.Field,
projection: pystencils.Field,
projection_matrix,
step_size=1,
cubic_bspline_interpolation=False,
add_to_projector=False,
central_ray_point=None):
# is_projection_stack = projection.spatial_dimensions == volume.spatial_dimensions
interpolation_mode = 'cubic_spline' if cubic_bspline_interpolation else 'linear'
volume_texture = pystencils.interpolation_astnodes.Interpolator(
volume, interpolation_mode)
ndim = volume.spatial_dimensions
projection_matrix = pystencils_reco.ProjectiveMatrix(projection_matrix)
t = pystencils_reco.typed_symbols('_parametrization', 'float32')
texture_coordinates = sympy.Matrix(
pystencils_reco.typed_symbols(f'_t:{ndim}', 'float32'))
u = projection.physical_coordinates_staggered
x = volume.index_to_physical(texture_coordinates)
is_perspective = projection_matrix.matrix.cols == ndim + 1
if is_perspective:
eqn = projection_matrix @ sympy.Matrix([*x, 1]) - sympy.Matrix(
[*(t * u), t])
else:
# this also works for perspective/cone beam projection (but may lead to instable parametrization)
eqn = projection_matrix @ x - u
ray_equations = sympy.solve(eqn, texture_coordinates, rational=False)
if not is_perspective:
t = [t for t in texture_coordinates
if t not in ray_equations.keys()][0]
assert len(
ray_equations.keys()
) == ndim - 1, "projection_matrix does not appear to define a projection"
ray_equations = sympy.Matrix(
[ray_equations[s] if s != t else t for s in texture_coordinates])
projection_vector = sympy.diff(ray_equations, t)
projection_vector_norm = projection_vector.norm()
projection_vector /= projection_vector_norm
conditions = pystencils_reco._geometry.coordinate_in_field_conditions(
volume, ray_equations)
if not central_ray_point:
central_ray_point = [0] * projection.spatial_dimensions
central_ray = projection_vector.subs({
i: j
for i, j in zip(pystencils.x_vector(projection.spatial_dimensions),
central_ray_point)
})
intersection_candidates = []
for i in range(ndim):
solution_min = sympy.solve(ray_equations[i], t, rational=False)
solution_max = sympy.solve(ray_equations[i] - volume.spatial_shape[i],
t,
rational=False)
intersection_candidates.extend(solution_min + solution_max)
intersection_point1 = sympy.Piecewise(
*[(f, sympy.And(*conditions).subs({t: f}))
for f in intersection_candidates], (-0, True))
intersection_point2 = sympy.Piecewise(
*[(f, sympy.And(*conditions).subs({t: f}))
for f in reversed(intersection_candidates)], (-0, True))
assert intersection_point1 != intersection_point2, \
"The intersections are unconditionally equal, reconstruction volume is not in detector FOV!"
# perform a integer set analysis here?
# space = isl.Space.create_from_names(isl.DEFAULT_CONTEXT, set=[str(t) for t in texture_coordinates])
# ray_set = isl.BasicSet.universe(space)
# for i, t in enumerate(texture_coordinates):
# # dafaq?
# ray_set.add_constraint(isl.Constraint.ineq_from_names(space, {str(texture_coordinates): 1}))
# ray_set.add_constraint(isl.Constraint.ineq_from_names(space,
# # {1: -volume.shape[i],
# str(texture_coordinates): -1}))
# ray_set.add_constraint(isl.Constraint.eq_from_name(space, ray_equations[i].subs({ #TODO
min_t = sympy.Min(intersection_point1, intersection_point2)
max_t = sympy.Max(intersection_point1, intersection_point2)
# parametrization_dim = list(ray_equations).index(t)
# min_t = 0
# max_t = volume.spatial_shape[parametrization_dim]
line_integral, num_steps, min_t_tmp, max_t_tmp, intensity_weighting, step = pystencils.data_types.typed_symbols(
'line_integral, num_steps, min_t_tmp, max_t_tmp, intensity_weighting, step',
'float32')
i = pystencils.data_types.TypedSymbol('i', 'int32')
num_steps = pystencils.data_types.TypedSymbol('num_steps', 'int32')
# step = step_size / projection_vector_norm
# tex_coord = ray_equations.subs({t: min_t_tmp + i * step})
tex_coord = ray_equations.subs({t: min_t_tmp}) + projection_vector * i
if callable(volume.coordinate_transform):
intensity_weighting_sym = projection_vector.dot(central_ray)**2
else:
intensity_weighting_sym = projection_vector.dot(central_ray)**2
assignments = {
min_t_tmp:
min_t,
max_t_tmp:
max_t,
num_steps:
sympy.ceiling(
(max_t_tmp - min_t_tmp) / (step_size / projection_vector_norm)),
line_integral:
sympy.Sum(volume_texture.at(tex_coord), (i, 0, num_steps)),
intensity_weighting:
intensity_weighting_sym,
projection.center():
(line_integral * step_size * intensity_weighting) +
(projection.center() if add_to_projector else 0)
# projection.center(): (max_t_tmp - min_t_tmp) / step # Uncomment to get path length
}
# def create_autodiff(self, constant_fields=None):
# backward_assignments = backward_projection(AdjointField(projections),
# AdjointField(volume),
# projection_matrix,
# 1)
# self._autodiff = pystencils.autodiff.AutoDiffOp(
# assignments, "op", constant_fields=constant_fields, backward_assignments=backward_assignments)
# assignments._create_autodiff = types.MethodType(create_autodiff, assignments)
return assignments
def upsample(input: {'field_type': FieldType.CUSTOM},
result,
sampling_factor=2):
assert input.spatial_dimensions == result.spatial_dimensions
assert input.field_type == FieldType.CUSTOM \
or result.spatial_shape == tuple([2 * x for x in input.spatial_shape])
assert input.index_shape == result.index_shape
assignments = []
ndim = input.spatial_dimensions
for i in range(result.index_shape[0]):
assignments.append(
pystencils.Assignment(result.center(i),
input.absolute_access(
sympy.Matrix(tuple([cast_func(x // sampling_factor, create_type("int")) for x in pystencils.x_vector(ndim)])), (i,)))
)
return assignments