Neptunya mst

.gitignore

@@ -139,6 +139,9 @@ ENV/
 # Rope project settings
 .ropeproject
 
+# meson build
+.meson*
+
 # mkdocs documentation
 /site
 

capytaine/bem/engines.py

@@ -60,6 +60,8 @@ class BasicMatrixEngine(MatrixEngine):
                                 'lu_decomposition': linear_solvers.LUSolverWithCache().solve,
                                 'gmres': linear_solvers.solve_gmres,
                                 }
+    if linear_solvers.gpu_direct:
+        available_linear_solvers['gpu_direct'] = linear_solvers.gpu_direct
 
     def __init__(self, *, linear_solver='lu_decomposition', matrix_cache_size=1):
 

capytaine/bem/solver.py

@@ -23,10 +23,12 @@
 from capytaine.io.xarray import problems_from_dataset, assemble_dataset, kochin_data_array
 from capytaine.tools.optional_imports import silently_import_optional_dependency
 from capytaine.tools.lists_of_points import _normalize_points, _normalize_free_surface_points
+
 from capytaine.tools.symbolic_multiplication import supporting_symbolic_multiplication
 
 LOG = logging.getLogger(__name__)
 
+
 class BEMSolver:
     """
     Solver for linear potential flow problems.
@@ -160,7 +162,9 @@ def solve_all(self, problems, *, method='indirect', n_jobs=1, progress_bar=True,
             joblib = silently_import_optional_dependency("joblib")
             if joblib is None:
                 raise ImportError(f"Setting the `n_jobs` argument to {n_jobs} requires the missing optional dependency 'joblib'.")
+
             groups_of_problems = LinearPotentialFlowProblem._group_for_parallel_resolution(problems)
+
             parallel = joblib.Parallel(return_as="generator", n_jobs=n_jobs)
             groups_of_results = parallel(joblib.delayed(self.solve_all)(grp, method=method, n_jobs=1, progress_bar=False, _check_wavelength=False, **kwargs) for grp in groups_of_problems)
             if progress_bar:
@@ -169,6 +173,8 @@ def solve_all(self, problems, *, method='indirect', n_jobs=1, progress_bar=True,
                                           description=f"Solving BEM problems with {n_jobs} threads:")
             results = [res for grp in groups_of_results for res in grp]  # flatten the nested list
             return results
+
+
 
     @staticmethod
     def _check_wavelength(problems):

capytaine/matrices/linear_solvers.py

@@ -7,6 +7,8 @@
 
 import logging
 
+import os
+
 import numpy as np
 from scipy import linalg as sl
 from scipy.sparse import linalg as ssl
@@ -149,3 +151,71 @@ def gmres_no_fft(A, b):
         LOG.warning(f"No convergence of the GMRES. Error code: {info}")
 
     return x
+
+##### GPU Solver #####
+class IgnoreGPUException(Exception):
+    pass
+
+try:
+    if not os.environ.get('CAPYTAINE_GPU',None) == 'True':
+        raise IgnoreGPUException()
+
+    import torch
+    def solve_gpu(a:np.ndarray,b:np.ndarray)->np.ndarray:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+        a = torch.from_numpy(a).cfloat().to(device)
+        b = torch.from_numpy(b).cfloat().to(device)
+
+        res = torch.linalg.solve(a, b)
+
+        return res.cpu().numpy()
+
+    def gpu_direct(A, b):
+        assert isinstance(b, np.ndarray) and A.ndim == b.ndim+1 and A.shape[-2] == b.shape[-1]
+        if isinstance(A, BlockCirculantMatrix):
+            LOG.debug("\tSolve linear circulant system %s", A)
+            blocks_of_diagonalization = A.block_diagonalize()
+            fft_of_rhs = np.fft.fft(np.reshape(b, (A.nb_blocks[0], A.block_shape[0])), axis=0)
+            try:  # Try to run it as vectorized numpy arrays.
+                fft_of_result = solve_gpu(blocks_of_diagonalization, fft_of_rhs)
+            except np.linalg.LinAlgError:  # Or do the same thing with list comprehension.
+                fft_of_result = np.array([gpu_direct(block, vec) for block, vec in zip(blocks_of_diagonalization, fft_of_rhs)])
+            result = np.fft.ifft(fft_of_result, axis=0).reshape((A.shape[1],))
+            return result
+
+        elif isinstance(A, BlockSymmetricToeplitzMatrix):
+            if A.nb_blocks == (2, 2):
+                LOG.debug("\tSolve linear system %s", A)
+                A1, A2 = A._stored_blocks[0, :]
+                b1, b2 = b[:len(b)//2], b[len(b)//2:]
+                x_plus = gpu_direct(A1 + A2, b1 + b2)
+                x_minus = gpu_direct(A1 - A2, b1 - b2)
+                return np.concatenate([x_plus + x_minus, x_plus - x_minus])/2
+            else:
+                # Not implemented
+                LOG.debug("\tSolve linear system %s", A)
+                return gpu_direct(A.full_matrix(), b)
+
+        elif isinstance(A, BlockMatrix):
+            LOG.debug("\tSolve linear system %s", A)
+            return gpu_direct(A.full_matrix(), b)
+
+        elif isinstance(A, np.ndarray):
+            LOG.debug(f"\tSolve linear system (size: {A.shape}) with numpy direct solver.")
+            return solve_gpu(A, b)
+
+        else:
+            raise ValueError(f"Unrecognized type of matrix to solve: {A}")
+
+    print(f'GPU available: {torch.cuda.get_device_name()}')
+except IgnoreGPUException:
+    pass
+except ImportError as e:
+    LOG.debug(f'GPU solver not available. Install pytorch and cuda to enable. Error: {e}')
+    gpu_direct = None
+
+except Exception as e:
+    LOG.debug(f'Unknown GPU Import Error: {e}')
+    gpu_direct = None    
+