Mercurial > repos > pointsource_pde / file comparison

comparison: src/laser_sampling.py

src/laser_sampling.py

changeset 1
a4137aedcb3a
child 3
c3a4f4bb87f7
equal deleted inserted replaced
0:7ec1cfe19a24 1:a4137aedcb3a
1 import numpy as np
2 from dolfinx import fem, geometry
3 from mpi4py import MPI
4 from petsc4py import PETSc
5 from pointsource_pde.dolfinx_extras import get_mass_matrix
6 from scipy.sparse.linalg import svds
7 from slepc4py import SLEPc
8
9 from dolfinx_access import cell_FunctionSpace_f64
10
11
12 class LaserSampling:
13 """
14 A linear operator class for laser sampling of a concentration (fem.Function)
15 """
16
17 def __init__(
18 self,
19 V,
20 domain,
21 nx,
22 ny,
23 beams,
24 num_segments=100,
25 domain_size=[-2, 2, -2, 2],
26 ):
27 self.V = V
28 self.domain = domain
29 self.nx, self.ny = nx, ny
30 self.M = len(beams)
31 self.domain_size = domain_size
32 self.num_segments = num_segments
33
34 # Build R matrix (single function does everything)
35 self.R = self.matrix(beams)
36
37 A = get_mass_matrix(V)
38 solver = PETSc.KSP().create(A.comm)
39 solver.setOperators(A)
40 solver.setType(PETSc.KSP.Type.PREONLY)
41 solver.getPC().setType(PETSc.PC.Type.LU)
42 self.solver = solver
43
44 # ‖x‖^2 = ‖Ax.array‖^2
45 # ‖Lx‖^2 = ‖Rx.array‖^2 = ‖RA^{-1}Ax.array‖^2 ≤ ‖RA^{-1}‖‖x‖^2 ≤ ‖R‖‖A^{-1}‖‖x‖^2,
46 # so we need the inverse of the minimal eigenvalue.
47 E = SLEPc.EPS().create(A.comm)
48 E.setOperators(A)
49 E.setProblemType(SLEPc.EPS.ProblemType.NHEP)
50 E.setWhichEigenpairs(SLEPc.EPS.Which.SMALLEST_MAGNITUDE)
51 # E.setWhichEigenpairs(SLEPc.EPS.Which.LARGEST_MAGNITUDE)
52 E.setFromOptions()
53 E.solve()
54
55 sigma = abs(E.getEigenvalue(0))
56 # self._opnorm = 1.0 / sigma
57 # print("Full sampling opnorm = %f" % self._opnorm)
58
59 # SVD operator norm
60 _u, s, _v = svds(self.R, k=1)
61 self._opnorm = s[0] # / sigma
62 print(f"Laser sampling opnorm = {self._opnorm:.6f}")
63
64 def matrix(self, beams):
65 """Combined: generate beams + segment + cell collision + dof weighting"""
66 dofmap = self.V.dofmap
67 tdim = self.domain.topology.dim
68
69 # Geometry trees
70
71 bb_tree = geometry.bb_tree(self.domain, tdim)
72
73 n_dofs = self.V.dofmap.index_map.size_global
74 R = np.zeros((self.M, n_dofs))
75
76 xmin, xmax, ymin, ymax = self.domain_size
77
78 for beam_idx, (start_pt, end_pt) in enumerate(beams):
79 # Segment into midpoints
80 vec = end_pt - start_pt
81 total_len = np.linalg.norm(vec)
82 seg_len = total_len / self.num_segments
83
84 midpoints = []
85 for i in range(self.num_segments):
86 t = (i + 0.5) / self.num_segments
87 midpoint = start_pt + t * vec
88 midpoints.append(midpoint)
89 # midpoints = np.array(midpoints, dtype=np.float64) # Force float64
90
91 # Ensure shape (N, 3), C-contiguous, read-only for DOLFINx
92 # assert midpoints.shape[1] == 2, (
93 # "Expected 2D points, got shape[1]=2 but need padding to 3D"
94 # )
95 # midpoints = np.pad(
96 # midpoints, ((0, 0), (0, 1)), mode="constant"
97 # ) # (N, 3) with z=0
98 # midpoints = np.ascontiguousarray(midpoints) # C-order
99 # midpoints.setflags(write=False) # Read-only requirement
100
101 # Find colliding cells
102 # cell_candidates = geometry.compute_collisions_points(bb_tree, midpoints)
103 # colliding_cells = geometry.compute_colliding_cells(
104 # self.domain, cell_candidates, midpoints
105 # )
106
107 # Cell→DOFs→weight accumulation
108 # cell_lists = list(colliding_cells.array)
109 # for seg_i, cell_idx_raw in enumerate(cell_lists):
110 # cell_idx = int(cell_idx_raw) # Convert np.int32 → Python int
111 for point in midpoints:
112 pad_point = np.array([point[0], point[1], 0.0])
113 cell_idx = cell_FunctionSpace_f64(self.V._cpp_object, pad_point)
114 if cell_idx >= 0: # Valid cell index
115 # Get global DOFs
116 cell_dofs_local = dofmap.cell_dofs(cell_idx)
117 # Convert local dofs to numpy int32 array (DOLFINx requirement)
118 # local_dofs_array = np.asarray(cell_dofs_local, dtype=np.int32)
119 # local_dofs_array.setflags(write=False) # Read-only requirement
120
121 # Batch convert all local indices to global in one call
122 cell_dofs_global = dofmap.index_map.local_to_global(cell_dofs_local)
123 cell_dofs = np.unique(cell_dofs_global)
124
125 weight = seg_len / len(cell_dofs)
126 R[beam_idx, cell_dofs] += weight # ADD for overlaps!
127
128 return R
129
130 def apply(self, x):
131 """
132 This does not work with MPI
133 """
134 x.x.scatter_forward()
135 return self.R @ x.x.array
136
137 def diff_adjdir(self, j, _x):
138 """
139 This does not work with MPI
140 """
141 # We need ⟨Rx,v⟩_{ℝ^n} = ⟨x,R^*v⟩_V
142 # We have ⟨x,R^*v⟩_V = ⟨Ax.array,[R^*v].array⟩_{ℝ^m}
143 # But Rx = R₀ x.array, so we need
144 # ⟨x.array,R₀^* v⟩_{ℝ^n} = ⟨Ax.array,[R^*v].array⟩_{ℝ^m}
145 # Taking [R^*v].array=A^{-1}R₀^* v, the conjugate works correctly.
146 tmp = fem.Function(self.V)
147 np.matmul(self.R.T, j, out=tmp.x.array)
148 tmp.x.scatter_forward()
149 return tmp
150 # TODO: is convection_diffusion actualy based on norms in ℝ^n?
151 # res = fem.Function(self.V)
152 # res.x.array[:] = 0.0
153 # self.solver.solve(tmp.x.petsc_vec, res.x.petsc_vec)
154 # res.x.scatter_forward()
155 # return res
156
157 def opnorm(self, *args):
158 # raise NotImplementedError("Need mesh weighting?")
159 return self._opnorm
160
161 def lipschitz_factor(self, *args):
162 return self.opnorm(*args)
163
164 def diff_adj_lipschitz_factor(self, *args):
165 return self.opnorm(*args)
166
167 def diff_bound(self, *args):
168 return self.opnorm(*args)
169
170 def codomain_bound(self, xbound=None):
171 if xbound is None:
172 raise Exception("Linear operators have unbounded range")
173 else:
174 return self.opnorm() * xbound
175
176 # Construct observation noise
177 def noise(self, noise_level):
178 M = self.M
179 # Generate noise with shape (M, 1) using scalar standard deviation
180 noise = np.random.normal(0, noise_level, size=(M,))
181 return noise

Mercurial Repository: pointsource_pde

mercurial