feat: use constantdict instead of immutabledict

Author: alexfikl

pyproject.toml

@@ -31,7 +31,7 @@ classifiers = [
 dependencies = [
     "arraycontext>=2021.1",
     "boxtree>=2023.1",
-    "immutabledict",
+    "constantdict>=2024.4",
     "loopy>=2024.1",
     "numpy",
     "pyopencl>=2022.1",

requirements.txt

@@ -1,6 +1,6 @@
 numpy
 sympy
-immutabledict
+constantdict
 pyvkfft
 
 # used in mpi-based tests

sumpy/codegen.py

@@ -28,7 +28,7 @@
 import re
 
 import numpy as np
-from immutabledict import immutabledict
+from constantdict import constantdict
 
 import loopy as lp
 import pymbolic.primitives as prim
@@ -142,7 +142,7 @@ def with_types(self, arg_id_to_dtype, clbl_inf_ctx):
 
         if z_dtype.numpy_dtype.kind == "c":
             return (self.copy(name_in_target="bessel_jv_two_complex",
-                              arg_id_to_dtype=immutabledict({
+                              arg_id_to_dtype=constantdict({
                                   -2: NumpyType(np.complex128),
                                   -1: NumpyType(np.complex128),
                                   0: NumpyType(np.int32),
@@ -151,7 +151,7 @@ def with_types(self, arg_id_to_dtype, clbl_inf_ctx):
                     clbl_inf_ctx)
         else:
             return (self.copy(name_in_target="bessel_jv_two",
-                              arg_id_to_dtype=immutabledict({
+                              arg_id_to_dtype=constantdict({
                                   -2: NumpyType(np.float64),
                                   -1: NumpyType(np.float64),
                                   0: NumpyType(np.int32),
@@ -184,15 +184,15 @@ def with_types(self, arg_id_to_dtype, clbl_inf_ctx):
 
         if z_dtype.numpy_dtype.kind == "c":
             return (self.copy(name_in_target="hank1_01_complex",
-                              arg_id_to_dtype=immutabledict({
+                              arg_id_to_dtype=constantdict({
                                   -2: NumpyType(np.complex128),
                                   -1: NumpyType(np.complex128),
                                   0: NumpyType(np.complex128),
                                   })),
                     clbl_inf_ctx)
         else:
             return (self.copy(name_in_target="hank1_01",
-                              arg_id_to_dtype=immutabledict({
+                              arg_id_to_dtype=constantdict({
                                   -2: NumpyType(np.complex128),
                                   -1: NumpyType(np.complex128),
                                   0: NumpyType(np.float64),

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 immutabledict import immutabledict
+from constantdict import constantdict
 
 from pytools import memoize
 
@@ -115,7 +115,8 @@ def __mul__(self, param: Number_ish) -> LinearPDESystemOperator:
             deriv_ident_to_coeff = {}
             for k, v in eq.items():
                 deriv_ident_to_coeff[k] = v * param
-            eqs.append(immutabledict(deriv_ident_to_coeff))
+            eqs.append(constantdict(deriv_ident_to_coeff))
+
         return LinearPDESystemOperator(self.dim, tuple(eqs))
 
     __rmul__ = __mul__
@@ -125,6 +126,7 @@ def __add__(
             ) -> LinearPDESystemOperator:
         assert self.dim == other_diff_op.dim
         assert len(self.eqs) == len(other_diff_op.eqs)
+
         eqs: list[Mapping[DerivativeIdentifier, sp.Expr]] = []
         for eq, other_eq in zip(self.eqs, other_diff_op.eqs, strict=True):
             res = dict(eq)
@@ -133,7 +135,8 @@ def __add__(
                     res[k] += v
                 else:
                     res[k] = v
-            eqs.append(immutabledict(res))
+            eqs.append(constantdict(res))
+
         return LinearPDESystemOperator(self.dim, tuple(eqs))
 
     __radd__ = __add__
@@ -269,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(), strict=True)
         }
-        results.append(LinearPDESystemOperator(pde.dim, (immutabledict(pde_dict),)))
+        results.append(LinearPDESystemOperator(pde.dim, (constantdict(pde_dict),)))
 
     return results
 
@@ -347,8 +350,8 @@ 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], ...] = \
-        (immutabledict(),) * len(diff_op.eqs)
+    empty: tuple[Mapping[DerivativeIdentifier, sp.Expr], ...] = (
+        (constantdict(),) * len(diff_op.eqs))
     res = LinearPDESystemOperator(dim, empty)
     for j in range(dim):
         mi = [0]*diff_op.total_dims
@@ -366,17 +369,20 @@ def diff(
         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(immutabledict(res))
+        eqs.append(constantdict(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, (immutabledict(),))
+
+    res = LinearPDESystemOperator(diff_op.dim, (constantdict(),))
     for i in range(diff_op.dim):
         mi = [0]*diff_op.total_dims
         mi[i] = 1
         res += diff(diff_op[i], tuple(mi))
+
     return res
 
 
@@ -436,6 +442,6 @@ def make_identity_diff_op(
         mi = tuple([0]*(ninput + 1))
     else:
         mi = tuple([0]*ninput)
-    return LinearPDESystemOperator(ninput, tuple(immutabledict(
-                    {DerivativeIdentifier(mi, i): sp.sympify(1)})
-                    for i in range(noutput)))
+    return LinearPDESystemOperator(ninput, tuple(
+            constantdict({DerivativeIdentifier(mi, i): sp.sympify(1)})
+            for i in range(noutput)))