forked from jan/opencl_fdfd
		
	Use fdfd_tools.solvers.generic for csr solve
This commit is contained in:
		
							parent
							
								
									80592dff79
								
							
						
					
					
						commit
						848f86f6ee
					
				| @ -14,7 +14,7 @@ satisfy the constraints for the 'conjugate gradient' algorithm | |||||||
| (positive definite, symmetric) and some that don't. | (positive definite, symmetric) and some that don't. | ||||||
| """ | """ | ||||||
| 
 | 
 | ||||||
| from typing import List, Dict, Any | from typing import Dict, Any | ||||||
| import time | import time | ||||||
| 
 | 
 | ||||||
| import numpy | import numpy | ||||||
| @ -22,7 +22,7 @@ from numpy.linalg import norm | |||||||
| import pyopencl | import pyopencl | ||||||
| import pyopencl.array | import pyopencl.array | ||||||
| 
 | 
 | ||||||
| import fdfd_tools.operators | import fdfd_tools.solvers | ||||||
| 
 | 
 | ||||||
| from . import ops | from . import ops | ||||||
| 
 | 
 | ||||||
| @ -43,7 +43,7 @@ class CSRMatrix(object): | |||||||
|         self.data = pyopencl.array.to_device(queue, m.data.astype(numpy.complex128)) |         self.data = pyopencl.array.to_device(queue, m.data.astype(numpy.complex128)) | ||||||
| 
 | 
 | ||||||
| 
 | 
 | ||||||
| def cg(a: 'scipy.sparse.csr_matrix', | def cg(A: 'scipy.sparse.csr_matrix', | ||||||
|        b: numpy.ndarray, |        b: numpy.ndarray, | ||||||
|        max_iters: int = 10000, |        max_iters: int = 10000, | ||||||
|        err_threshold: float = 1e-6, |        err_threshold: float = 1e-6, | ||||||
| @ -54,7 +54,7 @@ def cg(a: 'scipy.sparse.csr_matrix', | |||||||
|     """ |     """ | ||||||
|     General conjugate-gradient solver for sparse matrices, where A @ x = b. |     General conjugate-gradient solver for sparse matrices, where A @ x = b. | ||||||
| 
 | 
 | ||||||
|     :param a: Matrix to solve (CSR format) |     :param A: Matrix to solve (CSR format) | ||||||
|     :param b: Right-hand side vector (dense ndarray) |     :param b: Right-hand side vector (dense ndarray) | ||||||
|     :param max_iters: Maximum number of iterations |     :param max_iters: Maximum number of iterations | ||||||
|     :param err_threshold: Error threshold for successful solve, relative to norm(b) |     :param err_threshold: Error threshold for successful solve, relative to norm(b) | ||||||
| @ -84,7 +84,7 @@ def cg(a: 'scipy.sparse.csr_matrix', | |||||||
|     rho = 1.0 + 0j |     rho = 1.0 + 0j | ||||||
|     errs = [] |     errs = [] | ||||||
| 
 | 
 | ||||||
|     m = CSRMatrix(queue, a) |     m = CSRMatrix(queue, A) | ||||||
| 
 | 
 | ||||||
|     ''' |     ''' | ||||||
|     Generate OpenCL kernels |     Generate OpenCL kernels | ||||||
| @ -102,7 +102,8 @@ def cg(a: 'scipy.sparse.csr_matrix', | |||||||
| 
 | 
 | ||||||
|     _, err2 = rhoerr_step(r, []) |     _, err2 = rhoerr_step(r, []) | ||||||
|     b_norm = numpy.sqrt(err2) |     b_norm = numpy.sqrt(err2) | ||||||
|     print('b_norm check: ', b_norm) |     if verbose: | ||||||
|  |         print('b_norm check: ', b_norm) | ||||||
| 
 | 
 | ||||||
|     success = False |     success = False | ||||||
|     for k in range(max_iters): |     for k in range(max_iters): | ||||||
| @ -126,7 +127,7 @@ def cg(a: 'scipy.sparse.csr_matrix', | |||||||
|         e = a_step(v, m, p, e) |         e = a_step(v, m, p, e) | ||||||
|         alpha = rho / dot(p, v, e) |         alpha = rho / dot(p, v, e) | ||||||
| 
 | 
 | ||||||
|         if k % 1000 == 0: |         if verbose and k % 1000 == 0: | ||||||
|             print(k) |             print(k) | ||||||
| 
 | 
 | ||||||
|     ''' |     ''' | ||||||
| @ -136,71 +137,43 @@ def cg(a: 'scipy.sparse.csr_matrix', | |||||||
| 
 | 
 | ||||||
|     x = x.get() |     x = x.get() | ||||||
| 
 | 
 | ||||||
|     if success: |     if verbose: | ||||||
|         print('Success', end='') |         if success: | ||||||
|     else: |             print('Success', end='') | ||||||
|         print('Failure', end=', ') |         else: | ||||||
|     print(', {} iterations in {} sec: {} iterations/sec \ |             print('Failure', end=', ') | ||||||
|                   '.format(k, time_elapsed, k / time_elapsed)) |         print(', {} iterations in {} sec: {} iterations/sec \ | ||||||
|     print('final error', errs[-1]) |                       '.format(k, time_elapsed, k / time_elapsed)) | ||||||
|     print('overhead {} sec'.format(start_time2 - start_time)) |         print('final error', errs[-1]) | ||||||
|  |         print('overhead {} sec'.format(start_time2 - start_time)) | ||||||
| 
 | 
 | ||||||
|     print('Final residual:', norm(a @ x - b) / norm(b)) |         print('Final residual:', norm(A @ x - b) / norm(b)) | ||||||
|     return x |     return x | ||||||
| 
 | 
 | ||||||
| 
 | 
 | ||||||
| def fdfd_cg_solver(omega: complex, | def fdfd_cg_solver(solver_opts: Dict[str, Any] = None, | ||||||
|                    dxes: List[List[numpy.ndarray]], |                    **fdfd_args | ||||||
|                    J: numpy.ndarray, |  | ||||||
|                    epsilon: numpy.ndarray, |  | ||||||
|                    mu: numpy.ndarray = None, |  | ||||||
|                    pec: numpy.ndarray = None, |  | ||||||
|                    pmc: numpy.ndarray = None, |  | ||||||
|                    adjoint: bool = False, |  | ||||||
|                    solver_opts: Dict[str, Any] = None, |  | ||||||
|                    ) -> numpy.ndarray: |                    ) -> numpy.ndarray: | ||||||
|     """ |     """ | ||||||
|     Conjugate gradient FDFD solver using CSR sparse matrices, mainly for |     Conjugate gradient FDFD solver using CSR sparse matrices, mainly for | ||||||
|      testing and development since it's much slower than the solver in main.py. |      testing and development since it's much slower than the solver in main.py. | ||||||
| 
 | 
 | ||||||
|     All ndarray arguments should be 1D arrays. To linearize a list of 3 3D ndarrays, |     Calls fdfd_tools.solvers.generic(**fdfd_args, | ||||||
|      either use fdfd_tools.vec() or numpy: |                                      matrix_solver=opencl_fdfd.csr.cg, | ||||||
|      f_1D = numpy.hstack(tuple((fi.flatten(order='F') for fi in [f_x, f_y, f_z]))) |                                      matrix_solver_opts=solver_opts) | ||||||
| 
 | 
 | ||||||
|     :param omega: Complex frequency to solve at. |     :param solver_opts: Passed as matrix_solver_opts to fdfd_tools.solver.generic(...). | ||||||
|     :param dxes: [[dx_e, dy_e, dz_e], [dx_h, dy_h, dz_h]] (complex cell sizes) |         Default {}. | ||||||
|     :param J: Electric current distribution (at E-field locations) |     :param fdfd_args: Passed as **fdfd_args to fdfd_tools.solver.generic(...). | ||||||
|     :param epsilon: Dielectric constant distribution (at E-field locations) |         Should include all of the arguments **except** matrix_solver and matrix_solver_opts | ||||||
|     :param mu: Magnetic permeability distribution (at H-field locations) |  | ||||||
|     :param pec: Perfect electric conductor distribution |  | ||||||
|         (at E-field locations; non-zero value indicates PEC is present) |  | ||||||
|     :param pmc: Perfect magnetic conductor distribution |  | ||||||
|         (at H-field locations; non-zero value indicates PMC is present) |  | ||||||
|     :param adjoint: If true, solves the adjoint problem. |  | ||||||
|     :param solver_opts: Passed as kwargs to opencl_fdfd.csr.cg(**solver_opts) |  | ||||||
|     :return: E-field which solves the system. |     :return: E-field which solves the system. | ||||||
|     """ |     """ | ||||||
| 
 | 
 | ||||||
|     if solver_opts is None: |     if solver_opts is None: | ||||||
|         solver_opts = dict() |         solver_opts = dict() | ||||||
| 
 | 
 | ||||||
|     b0 = -1j * omega * J |     x = fdfd_tools.solvers.generic(matrix_solver=cg, | ||||||
|     A0 = fdfd_tools.operators.e_full(omega, dxes, epsilon=epsilon, mu=mu, pec=pec, pmc=pmc) |                                    matrix_solver_opts=solver_opts, | ||||||
|  |                                    **fdfd_args) | ||||||
| 
 | 
 | ||||||
|     Pl, Pr = fdfd_tools.operators.e_full_preconditioners(dxes) |     return x | ||||||
| 
 |  | ||||||
|     if adjoint: |  | ||||||
|         A = (Pl @ A0 @ Pr).H |  | ||||||
|         b = Pr.H @ b0 |  | ||||||
|     else: |  | ||||||
|         A = Pl @ A0 @ Pr |  | ||||||
|         b = Pl @ b0 |  | ||||||
| 
 |  | ||||||
|     x = cg(A.tocsr(), b, **solver_opts) |  | ||||||
| 
 |  | ||||||
|     if adjoint: |  | ||||||
|         x0 = Pl.H @ x |  | ||||||
|     else: |  | ||||||
|         x0 = Pr @ x |  | ||||||
| 
 |  | ||||||
|     return x0 |  | ||||||
|  | |||||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user