33
33
import numpy as np
34
34
import sympy as sp
35
35
import sympy .polys .agca .modules as sp_modules
36
- from pyrsistent import pmap
36
+ from immutabledict import immutabledict
37
37
38
38
from pytools import memoize
39
39
@@ -110,12 +110,12 @@ def order(self) -> int:
110
110
return deg
111
111
112
112
def __mul__ (self , param : Number_ish ) -> LinearPDESystemOperator :
113
- eqs = []
113
+ eqs : list [ Mapping [ DerivativeIdentifier , sp . Expr ]] = []
114
114
for eq in self .eqs :
115
115
deriv_ident_to_coeff = {}
116
116
for k , v in eq .items ():
117
117
deriv_ident_to_coeff [k ] = v * param
118
- eqs .append (pmap (deriv_ident_to_coeff ))
118
+ eqs .append (immutabledict (deriv_ident_to_coeff ))
119
119
return LinearPDESystemOperator (self .dim , tuple (eqs ))
120
120
121
121
__rmul__ = __mul__
@@ -125,15 +125,15 @@ def __add__(
125
125
) -> LinearPDESystemOperator :
126
126
assert self .dim == other_diff_op .dim
127
127
assert len (self .eqs ) == len (other_diff_op .eqs )
128
- eqs = []
128
+ eqs : list [ Mapping [ DerivativeIdentifier , sp . Expr ]] = []
129
129
for eq , other_eq in zip (self .eqs , other_diff_op .eqs ):
130
130
res = dict (eq )
131
131
for k , v in other_eq .items ():
132
132
if k in res :
133
133
res [k ] += v
134
134
else :
135
135
res [k ] = v
136
- eqs .append (pmap (res ))
136
+ eqs .append (immutabledict (res ))
137
137
return LinearPDESystemOperator (self .dim , tuple (eqs ))
138
138
139
139
__radd__ = __add__
@@ -272,7 +272,7 @@ def intersect(
272
272
DerivativeIdentifier (mi , 0 ): sym .sympify (coeff .as_expr ().simplify ()) for
273
273
(mi , coeff ) in zip (scalar_pde .monoms (), scalar_pde .coeffs ())
274
274
}
275
- results .append (LinearPDESystemOperator (pde .dim , (pmap (pde_dict ),)))
275
+ results .append (LinearPDESystemOperator (pde .dim , (immutabledict (pde_dict ),)))
276
276
277
277
return results
278
278
@@ -351,7 +351,7 @@ def as_scalar_pde(pde: LinearPDESystemOperator, comp_idx: int) \
351
351
def laplacian (diff_op : LinearPDESystemOperator ) -> LinearPDESystemOperator :
352
352
dim = diff_op .dim
353
353
empty : tuple [Mapping [DerivativeIdentifier , sp .Expr ], ...] = \
354
- (pmap (),) * len (diff_op .eqs )
354
+ (immutabledict (),) * len (diff_op .eqs )
355
355
res = LinearPDESystemOperator (dim , empty )
356
356
for j in range (dim ):
357
357
mi = [0 ]* diff_op .total_dims
@@ -363,19 +363,19 @@ def laplacian(diff_op: LinearPDESystemOperator) -> LinearPDESystemOperator:
363
363
def diff (
364
364
diff_op : LinearPDESystemOperator , mi : tuple [int , ...]
365
365
) -> LinearPDESystemOperator :
366
- eqs = []
366
+ eqs : list [ Mapping [ DerivativeIdentifier , sp . Expr ]] = []
367
367
for eq in diff_op .eqs :
368
368
res = {}
369
369
for deriv_ident , v in eq .items ():
370
370
new_mi = add_mi (deriv_ident .mi , mi )
371
371
res [DerivativeIdentifier (new_mi , deriv_ident .vec_idx )] = v
372
- eqs .append (pmap (res ))
372
+ eqs .append (immutabledict (res ))
373
373
return LinearPDESystemOperator (diff_op .dim , tuple (eqs ))
374
374
375
375
376
376
def divergence (diff_op : LinearPDESystemOperator ) -> LinearPDESystemOperator :
377
377
assert len (diff_op .eqs ) == diff_op .dim
378
- res = LinearPDESystemOperator (diff_op .dim , (pmap (),))
378
+ res = LinearPDESystemOperator (diff_op .dim , (immutabledict (),))
379
379
for i in range (diff_op .dim ):
380
380
mi = [0 ]* diff_op .total_dims
381
381
mi [i ] = 1
@@ -439,7 +439,6 @@ def make_identity_diff_op(
439
439
mi = tuple ([0 ]* (ninput + 1 ))
440
440
else :
441
441
mi = tuple ([0 ]* ninput )
442
- eqs = tuple (pmap (
442
+ return LinearPDESystemOperator ( ninput , tuple (immutabledict (
443
443
{DerivativeIdentifier (mi , i ): sp .sympify (1 )})
444
- for i in range (noutput ))
445
- return LinearPDESystemOperator (ninput , eqs )
444
+ for i in range (noutput )))
0 commit comments