Skip to content

Commit

Permalink
Minor changes.
Browse files Browse the repository at this point in the history
  • Loading branch information
Ig-dolci committed Apr 26, 2024
1 parent ddbccab commit 469cf3a
Showing 1 changed file with 4 additions and 6 deletions.
10 changes: 4 additions & 6 deletions pyop3/array/petsc.py
Original file line number Diff line number Diff line change
Expand Up @@ -380,8 +380,8 @@ def _map(self, axes):
loop_index = just_one(axes.outer_loops)
iterset = AxisTree(loop_index.iterset.node_map)
map_axes = iterset.add_subtree(axes, *iterset.leaf)
global_map = HierarchicalArray(map_axes, dtype=IntType)
global_map = global_map[loop_index.local_index]
mapping = HierarchicalArray(map_axes, dtype=IntType)
mapping = mapping[loop_index.local_index]
for idxs in my_product(axes.outer_loops):
# target_indices = {idx.index.id: idx.target_exprs for idx in idxs}
target_indices = merge_dicts([idx.replace_map for idx in idxs])
Expand All @@ -390,20 +390,18 @@ def _map(self, axes):
offset = axes.unindexed.offset(
p.target_exprs, p.target_path, loop_exprs=target_indices
)
global_map.set_value(
mapping.set_value(
p.source_exprs,
offset,
p.source_path,
loop_exprs=target_indices,
)
return global_map
return mapping


@cached_property
@PETSc.Log.EventDecorator()
def maps(self):
from pyop3.axtree.layout import my_product

# TODO: Don't think these need to be lists here.
# FIXME: This will only work for singly-nested matrices
if self.nested:
Expand Down

0 comments on commit 469cf3a

Please sign in to comment.