From 782d1192def8cdf487930b5fb1d356789e339176 Mon Sep 17 00:00:00 2001 From: Andreas Kloeckner Date: Thu, 18 Jul 2024 11:00:06 -0500 Subject: [PATCH] Switch from pyrsistent to immutabledict --- .test-conda-env-py3.yml | 1 - requirements.txt | 2 +- setup.py | 2 +- sumpy/expansion/diff_op.py | 25 ++++++++++++------------- 4 files changed, 14 insertions(+), 16 deletions(-) diff --git a/.test-conda-env-py3.yml b/.test-conda-env-py3.yml index caa303f5..f60204d3 100644 --- a/.test-conda-env-py3.yml +++ b/.test-conda-env-py3.yml @@ -14,7 +14,6 @@ dependencies: - python=3 - python-symengine - pyfmmlib -- pyrsistent - pyvkfft - mpi4py diff --git a/requirements.txt b/requirements.txt index da6a8fd2..9a3889d8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ numpy sympy -pyrsistent +immutabledict pyvkfft # used in mpi-based tests diff --git a/setup.py b/setup.py index 59c7f9b8..a1216d6e 100644 --- a/setup.py +++ b/setup.py @@ -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", diff --git a/sumpy/expansion/diff_op.py b/sumpy/expansion/diff_op.py index 63c234a0..d185ebbf 100644 --- a/sumpy/expansion/diff_op.py +++ b/sumpy/expansion/diff_op.py @@ -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 @@ -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__ @@ -125,7 +125,7 @@ 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(): @@ -133,7 +133,7 @@ def __add__( res[k] += v else: res[k] = v - eqs.append(pmap(res)) + eqs.append(immutabledict(res)) return LinearPDESystemOperator(self.dim, tuple(eqs)) __radd__ = __add__ @@ -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 @@ -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 @@ -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 @@ -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)))