GPU Solver Enabled Via Env Variable

PyNite/Analysis.py

@@ -1,7 +1,7 @@
 from math import isclose
 from PyNite.LoadCombo import LoadCombo
 from numpy import array, atleast_2d, zeros, subtract, matmul, divide, seterr, nanmax
-from numpy.linalg import solve
+from PyNite.Solvers import solve
 
 def _prepare_model(model):
     """Prepares a model for analysis by ensuring at least one load combination is defined, generating all meshes that have not already been generated, activating all non-linear members, and internally numbering all nodes and elements.
@@ -147,7 +147,7 @@ def _PDelta_step(model, combo_name, P1, FER1, D1_indices, D2_indices, D2, log=Tr
 
     # Import `scipy` features if the sparse solver is being used
     if sparse == True:
-        from scipy.sparse.linalg import spsolve
+        from PyNite.Solvers import spsolve
 
     iter_count_TC = 1    # Tracks tension/compression-only iterations
     iter_count_PD = 1    # Tracks P-Delta iterations

PyNite/FEModel3D.py

@@ -3,8 +3,9 @@
 import warnings
 from math import isclose
 
+
+from PyNite.Solvers import solve
 from numpy import array, zeros, matmul, divide, subtract, atleast_2d, all
-from numpy.linalg import solve
 
 from PyNite.Node3D import Node3D
 from PyNite.Material import Material
@@ -1938,7 +1939,7 @@ def analyze(self, log=False, check_stability=True, check_statics=False, max_iter
 
         # Import `scipy` features if the sparse solver is being used
         if sparse == True:
-            from scipy.sparse.linalg import spsolve
+            from PyNite.Solvers import spsolve
 
         # Prepare the model for analysis
         Analysis._prepare_model(self)

PyNite/Solvers.py

@@ -0,0 +1,42 @@
+"""Module to import various solvers for PyNite
+The GPU solver method is preferred, but requires PyTorch and a CUDA capable GPU, and is not available on all platforms.
+
+The CPU solver is the next best option, and is available on all platforms via numpy and scipy
+"""
+import os
+try:
+    if os.environ.get('PYNITE_GPU',None) != 'True':
+        raise ImportError(f'PYNITE_GPU environment variable not set to `True`')
+
+    import torch
+    import numpy
+
+    if not torch.cuda.is_available():
+        raise ImportError(f'CUDA not available')
+
+    def solve(a:numpy.ndarray,b:numpy.ndarray)->numpy.ndarray:
+        device = torch.device("cuda")
+
+        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 spsolve(a:numpy.ndarray,b:numpy.ndarray)->numpy.ndarray:
+        device = torch.device("cuda")
+        a = a.todense() #FIXME: sparse solver is currently not in torch standard library, investigate other options
+        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()  
+
+    print(f'PyNite Running with GPU solver: {torch.cuda.get_device_name()}')
+
+except Exception as e:
+    print(f'GPU solver not available: {e}')
+    from numpy.linalg import solve
+    from scipy.sparse.linalg import spsolve