Mercurial > repos > pointsource_pde / file comparison
comparison: src/dolfinx_extras.py
src/dolfinx_extras.py
- changeset 1
- a4137aedcb3a
equal
deleted
inserted
replaced
| 0:7ec1cfe19a24 | 1:a4137aedcb3a |
|---|---|
| 1 # import weakref | |
| 2 | |
| 3 import numpy as np | |
| 4 import ufl | |
| 5 from dolfinx import fem | |
| 6 from dolfinx.fem.function import Function, FunctionSpace | |
| 7 | |
| 8 # _mass_matrices = weakref.WeakKeyDictionary() | |
| 9 _mass_matrices = {} | |
| 10 | |
| 11 | |
| 12 def get_mass_matrix(space: FunctionSpace): | |
| 13 # idex = space # Does not work due to fenicx being garbage | |
| 14 idx = id(space) # FIXME: This leaks memory 🤯 | |
| 15 if idx not in _mass_matrices: | |
| 16 u = ufl.TrialFunction(space) | |
| 17 v = ufl.TestFunction(space) | |
| 18 a = fem.form(ufl.inner(u, v) * ufl.dx) | |
| 19 m = fem.petsc.assemble_matrix(a) | |
| 20 m.assemble() | |
| 21 _mass_matrices[idx] = m | |
| 22 return _mass_matrices[idx] | |
| 23 | |
| 24 | |
| 25 def dot(f: Function, g: Function): | |
| 26 if g.function_space != f.function_space: | |
| 27 raise ValueError("Function spaces need to agree") | |
| 28 m = get_mass_matrix(f.function_space) | |
| 29 f_vec = f.x.petsc_vec | |
| 30 g_vec = g.x.petsc_vec | |
| 31 return g_vec.dot(m @ f_vec) | |
| 32 | |
| 33 | |
| 34 def norm2_squared(f: Function): | |
| 35 return dot(f, f) | |
| 36 | |
| 37 | |
| 38 def norm2(f: Function): | |
| 39 return norm2_squared(f).sqrt() | |
| 40 | |
| 41 | |
| 42 def dist2_squared(f: Function, g: Function): | |
| 43 if g.function_space != f.function_space: | |
| 44 raise ValueError("Function spaces need to agree") | |
| 45 m = get_mass_matrix(f.function_space) | |
| 46 f_vec = f.x.petsc_vec | |
| 47 g_vec = g.x.petsc_vec | |
| 48 mf = m @ f_vec | |
| 49 mg = m @ g_vec | |
| 50 return f_vec.dot(mf) + 2 * g_vec.dot(mg) + g_vec.dot(mf) | |
| 51 | |
| 52 | |
| 53 def dist2(f: Function, g: Function): | |
| 54 return np.sqrt(dist2_squared(f, g)) |
