/
separate_complex.py
152 lines (128 loc) · 5.38 KB
/
separate_complex.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
from sympy.printing.python import PythonPrinter, StrPrinter
from sympy import Symbol, Function, symbols, re, im, S
class ComplexFormPrinter(PythonPrinter):
def __init__(self, skip_multiplication_by_one=False):
self.skip_multiplication_by_one = skip_multiplication_by_one
PythonPrinter.__init__(self)
def _print_Float(self,expr):
return str(expr)
def _print_Function(self, expr):
if expr.func == re:
arg = expr.args[0]
if hasattr(arg, 'name'):
return arg.name + '_re'
return str(arg.func) + '(' + ','.join([self._print(x) for x in arg.args]) + ')'
if expr.func == im:
arg = expr.args[0]
if hasattr(arg, 'name'):
return arg.name + '_im'
return str(arg.func) + '(' + ','.join([self._print(x) for x in arg.args]) + ')'
return PythonPrinter._print_Function(self, expr)
def _print_Symbol(self, expr):
return expr.name
def _print_Rational(self,expr):
return StrPrinter._print_Rational(self,expr)
def _print_Mul(self,expr):
def rearrange_multiplication(expr):
sign = ''
if expr.startswith('-'):
expr = expr[1:]
sign = '-'
measures = []
functions = []
numbers = []
symbols = []
args = expr.split('*')
for arg in args:
if arg.startswith('d'):
measures.append(arg)
elif arg.find('('):
functions.append(arg)
elif arg.find('.'):
numbers.append(arg)
else:
symbols.append(arg)
return sign + '*'.join(numbers + symbols + functions + measures)
return rearrange_multiplication(StrPrinter._print_Mul(self, expr)).replace('.00000000000000','')
class ComplexFunction(Function):
@classmethod
def eval(cls, *args):
for arg in args:
if arg.is_Number and arg is S.Zero:
return S.Zero
def _eval_is_real(self):
for arg in self.args:
if not arg.is_real:
return False
return True
class grad(ComplexFunction):
nargs = 1
def as_real_imag(self, deep=True, **hints):
u_re, u_im = self.args[0].as_real_imag()
return (grad(u_re), grad(u_im))
class div(grad):
nargs = 1
def as_real_imag(self, deep=True, **hints):
u_re, u_im = self.args[0].as_real_imag()
return (div(u_re), div(u_im))
class curl(grad):
nargs = 1
def as_real_imag(self, deep=True, **hints):
u_re, u_im = self.args[0].as_real_imag()
return (curl(u_re), curl(u_im))
class nabla_grad(grad):
nargs = 1
def as_real_imag(self, deep=True, **hints):
u_re, u_im = self.args[0].as_real_imag()
return (nabla_grad(u_re), nabla_grad(u_im))
class rot(curl):
pass
class inner(ComplexFunction):
nargs = 2
treat_arguments_as_vectors = True
@classmethod
def eval(cls, u, v):
if not cls.treat_arguments_as_vectors and u.is_real and v.is_real:
return u*v
return ComplexFunction.eval(cls, u, v)
def as_real_imag(self, deep=True, **hints):
u_re,u_im = self.args[0].as_real_imag()
v_re,v_im = self.args[1].as_real_imag()
return (inner(u_re,v_re) - inner(u_im,v_im), inner(u_re,v_im) + inner(v_re,u_im))
class cross(ComplexFunction):
nargs = 2
def as_real_imag(self, deep=True, **hints):
u_re,u_im = self.args[0].as_real_imag()
v_re,v_im = self.args[1].as_real_imag()
return (cross(u_re,v_re) - cross(u_im,v_im), cross(u_re,v_im) + cross(v_re,u_im))
def separate(form_as_str, **kwargs):
for var_name, var in kwargs.items():
if type(var) is tuple: # as complex value: (real, imaginary)
locals()[var_name] = symbols(var_name)
else: # as real value: value
locals()[var_name] = symbols(var_name, real=True)
form_printer = ComplexFormPrinter()
form_as_sympy = eval(form_as_str)
expanded_expression = form_as_sympy.expand(complex=True)
real_part_as_sympy, imag_part_as_sympy = expanded_expression.as_real_imag()
real_part_as_str = form_printer.doprint(real_part_as_sympy)
imag_part_as_str = form_printer.doprint(imag_part_as_sympy)
from dolfin import cross, curl, grad, inner, Function
for var_name, var in kwargs.items():
if type(var) is tuple: # as complex value: (real, imaginary)
locals()[var_name + '_re'] = var[0]
locals()[var_name + '_im'] = var[1]
elif type(var) is dict:
locals()[var_name + '_re'] = var['re']
locals()[var_name + '_im'] = var['im']
real_part_as_str = real_part_as_str.replace(var_name, var_name + '_re')
imag_part_as_str = imag_part_as_str.replace(var_name, var_name + '_im')
else: # as real value: value
locals()[var_name] = var
#print("Debug: real part of the expression after separation:\n\t'%s'" % real_part_as_str)
#print("Debug: imaginary part of the expression after separation:\n\t'%s'" % imag_part_as_str)
real_part_as_ufl = eval(real_part_as_str)
imaginary_part_as_ufl = eval(imag_part_as_str)
return real_part_as_ufl, imaginary_part_as_ufl
if __name__ == '__main__':
a_re, a_im = separate("T*E", T={'re':'T_re','im':'T_im'}, E=('E_re','E_im'))