diff --git a/pystencils/fd/derivation.py b/pystencils/fd/derivation.py index 5a9d79e8fe4038a090c8da84507747ed97633b86..af1f4e387b0bdc8bafccd03bad624d986c5c610d 100644 --- a/pystencils/fd/derivation.py +++ b/pystencils/fd/derivation.py @@ -303,11 +303,12 @@ class FiniteDifferenceStaggeredStencilDerivation: zero_counts[zero_count].append(weights) best = zero_counts[max(zero_counts.keys())] if len(best) > 1: # if there are multiple, pick the one that contains a nonzero center weight - center = [tuple(p + pos) for p in points].index((0, 0, 0)) + center = [tuple(p + pos) for p in points].index((0, 0, 0)[:dim]) best = [b for b in best if b[center] != 0] - if len(best) > 1: - raise NotImplementedError("more than one suitable set of weights found, don't know how to proceed") - weights = best[0] + if len(best) > 1: # if there are still multiple, they are equivalent, so we average + weights = sp.Add(*[sp.Matrix(b) for b in best]) / len(best) + else: + weights = best[0] assert weights points_tuple = tuple([tuple(p + pos) for p in points]) diff --git a/pystencils_tests/test_fd_derivation.ipynb b/pystencils_tests/test_fd_derivation.ipynb index f992c8ac5051dd83553cb7da3e6842c9be1aa8c3..1be42210e774a62513e0d73ce70c57ef51e0240f 100644 --- a/pystencils_tests/test_fd_derivation.ipynb +++ b/pystencils_tests/test_fd_derivation.ipynb @@ -332,6 +332,9 @@ "assert FiniteDifferenceStaggeredStencilDerivation(\"T\", 3, (2,)).apply(c3) == c3[0, 0, 1] - c3[0, 0, 0]\n", "assert FiniteDifferenceStaggeredStencilDerivation(\"B\", 3, (2,)).apply(c3) == c3[0, 0, 0] - c3[0, 0, -1]\n", "\n", + "assert FiniteDifferenceStaggeredStencilDerivation(\"S\", 2, (0,)).apply(c) == \\\n", + " (c[1, 0] + c[1, -1] - c[-1, 0] - c[-1, -1])/4\n", + "\n", "assert FiniteDifferenceStaggeredStencilDerivation(\"NE\", 2, (0,)).apply(c) + \\\n", " FiniteDifferenceStaggeredStencilDerivation(\"NE\", 2, (1,)).apply(c) == c[1, 1] - c[0, 0]\n", "assert FiniteDifferenceStaggeredStencilDerivation(\"NE\", 3, (0,)).apply(c3) + \\\n",