Skip to content

Commit

Permalink
Switch from pyrsistent to immutabledict
Browse files Browse the repository at this point in the history
  • Loading branch information
inducer committed Jul 18, 2024
1 parent 4698624 commit 7f3cd5e
Show file tree
Hide file tree
Showing 4 changed files with 14 additions and 16 deletions.
1 change: 0 additions & 1 deletion .test-conda-env-py3.yml
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,6 @@ dependencies:
- python=3
- python-symengine
- pyfmmlib
- pyrsistent
- pyvkfft
- mpi4py

Expand Down
2 changes: 1 addition & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
numpy
sympy
pyrsistent
immutabledict
pyvkfft

# used in mpi-based tests
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ def write_git_revision(package_name):
"loopy>=2021.1",
"boxtree>=2018.1",
"arraycontext",
"pyrsistent>=0.16.0",
"immutabledict",
"sympy>=0.7.2",
"pymbolic>=2021.1",
"pyvkfft>=2022.1",
Expand Down
25 changes: 12 additions & 13 deletions sumpy/expansion/diff_op.py
Original file line number Diff line number Diff line change
Expand Up @@ -33,7 +33,7 @@
import numpy as np
import sympy as sp
import sympy.polys.agca.modules as sp_modules
from pyrsistent import pmap
from immutabledict import immutabledict

from pytools import memoize

Expand Down Expand Up @@ -110,12 +110,12 @@ def order(self) -> int:
return deg

def __mul__(self, param: Number_ish) -> LinearPDESystemOperator:
eqs = []
eqs: list[Mapping[DerivativeIdentifier, sp.Expr]] = []
for eq in self.eqs:
deriv_ident_to_coeff = {}
for k, v in eq.items():
deriv_ident_to_coeff[k] = v * param
eqs.append(pmap(deriv_ident_to_coeff))
eqs.append(immutabledict(deriv_ident_to_coeff))
return LinearPDESystemOperator(self.dim, tuple(eqs))

__rmul__ = __mul__
Expand All @@ -125,15 +125,15 @@ def __add__(
) -> LinearPDESystemOperator:
assert self.dim == other_diff_op.dim
assert len(self.eqs) == len(other_diff_op.eqs)
eqs = []
eqs: list[Mapping[DerivativeIdentifier, sp.Expr]] = []
for eq, other_eq in zip(self.eqs, other_diff_op.eqs):
res = dict(eq)
for k, v in other_eq.items():
if k in res:
res[k] += v
else:
res[k] = v
eqs.append(pmap(res))
eqs.append(immutabledict(res))
return LinearPDESystemOperator(self.dim, tuple(eqs))

__radd__ = __add__
Expand Down Expand Up @@ -272,7 +272,7 @@ def intersect(
DerivativeIdentifier(mi, 0): sym.sympify(coeff.as_expr().simplify()) for
(mi, coeff) in zip(scalar_pde.monoms(), scalar_pde.coeffs())
}
results.append(LinearPDESystemOperator(pde.dim, (pmap(pde_dict),)))
results.append(LinearPDESystemOperator(pde.dim, (immutabledict(pde_dict),)))

return results

Expand Down Expand Up @@ -351,7 +351,7 @@ def as_scalar_pde(pde: LinearPDESystemOperator, comp_idx: int) \
def laplacian(diff_op: LinearPDESystemOperator) -> LinearPDESystemOperator:
dim = diff_op.dim
empty: tuple[Mapping[DerivativeIdentifier, sp.Expr], ...] = \
(pmap(),) * len(diff_op.eqs)
(immutabledict(),) * len(diff_op.eqs)
res = LinearPDESystemOperator(dim, empty)
for j in range(dim):
mi = [0]*diff_op.total_dims
Expand All @@ -363,19 +363,19 @@ def laplacian(diff_op: LinearPDESystemOperator) -> LinearPDESystemOperator:
def diff(
diff_op: LinearPDESystemOperator, mi: tuple[int, ...]
) -> LinearPDESystemOperator:
eqs = []
eqs: list[Mapping[DerivativeIdentifier, sp.Expr]] = []
for eq in diff_op.eqs:
res = {}
for deriv_ident, v in eq.items():
new_mi = add_mi(deriv_ident.mi, mi)
res[DerivativeIdentifier(new_mi, deriv_ident.vec_idx)] = v
eqs.append(pmap(res))
eqs.append(immutabledict(res))
return LinearPDESystemOperator(diff_op.dim, tuple(eqs))


def divergence(diff_op: LinearPDESystemOperator) -> LinearPDESystemOperator:
assert len(diff_op.eqs) == diff_op.dim
res = LinearPDESystemOperator(diff_op.dim, (pmap(),))
res = LinearPDESystemOperator(diff_op.dim, (immutabledict(),))
for i in range(diff_op.dim):
mi = [0]*diff_op.total_dims
mi[i] = 1
Expand Down Expand Up @@ -439,7 +439,6 @@ def make_identity_diff_op(
mi = tuple([0]*(ninput + 1))
else:
mi = tuple([0]*ninput)
eqs = tuple(pmap(
return LinearPDESystemOperator(ninput, tuple(immutabledict(
{DerivativeIdentifier(mi, i): sp.sympify(1)})
for i in range(noutput))
return LinearPDESystemOperator(ninput, eqs)
for i in range(noutput)))

0 comments on commit 7f3cd5e

Please sign in to comment.