start work on fast simulation solvers

.gitignore

@@ -6,4 +6,4 @@ __pycache__/
 # Emacs stuff
 *~
 ~*~
-#*#
+#*#

AUTHORS.md

@@ -1,7 +1,7 @@
 
 ## Original Authors ##
 
-FEM4INAS follows the doctoral research of Alvaro Cea in the [PhD thesis](https://spiral.imperial.ac.uk/handle/10044/1/89976). The software has been subsequently implemented for maximum performance.
+FENIAX follows the doctoral research at Imperial College London of Alvaro Cea in the [PhD thesis](https://spiral.imperial.ac.uk/handle/10044/1/89976) under the supervision of Rafael Palacios. The software has been subsequently implemented for maximum performance in JAX.
 ```
 Alvaro Cea (Imperial College London)
 Rafael Palacios (Imperial College London)

docs/schemas/draft2.org

@@ -109,6 +109,36 @@ Systems with labels:
 [[file:~/projects/FENIAX/feniax/systems/intrinsicAD.py::label = f"main_{label_sys}_{label_ad}"][AD_system]]
 [[file:~/projects/FENIAX/feniax/systems/intrinsicShard.py::self.label = f"main_{label_sys}_{label_shard}"][shard_system]]
 
+*** Architecture for the solution of systems
+Chain of requirements between the various ways to run the code.
+Fast needs to use the functions in Flexible:
+Fast -> Flexible  and similarly:
+AD -> Fast
+Shard -> Flexible
+ShardF -> Shard
+
+
+- Flexible: workflow of simulations happening sequentially
+  major functions for the systems to be solved can be found in dq_...:
+  [[file:~/projects/FENIAX/feniax/intrinsic/dq_dynamic.py][dq_dynamic]]
+  [[file:~/projects/FENIAX/feniax/intrinsic/dq_static.py][dq_static]]
+
+
+- Fast: entire solution within one function such that memory copies to cuda devices are avoided
+  Computation of intrinsic modes, modal couplings, aerodynamic matrices happen within a single function, from within the solution of the system of equations is also called.
+  Importantly, the functions within the dq_[] modules are used for the solution, thereby avoiding code duplication and promoting modular design.
+
+
+- AD: the entire solution within one function as well, but needs inputs/ outputs for the differentiation to take place
+
+  The function to be differentiated will call the function in the   
+
+- Shard flexible: workflow as in flexible but with inputs over which to build solutions in parallel.
+
+- Shard fast: Similarly, everything happens within a function, 
+
+- Shard AD: shard the inputs, take a function for the output,
+
 
 ** XForces
 *** prescribed_follower
@@ -350,7 +380,6 @@ Critical for large design optimisation problems.
 Well in place for Nastran Models  except for the derivatives provided by Nastran using Sol 200.
 
 
-
 * System based solutions :noexport:
 TODO: make automatic label as the first
 | Type         | Target | Gravity    | BC1        | ModalAero | SteadyAero   | UnsteadyAero | Point loads | q0 approx | Rigid-body           | Nonlinearities         | residualised |

feniax/drivers/intrinsic_driver.py

@@ -49,7 +49,7 @@ def __init__(self, config: configuration.Config):
     def pre_simulation(self):
         # TODO: here the RFA for aerodynamics should be included
         # TODO: condensation methods of K and M should be included
-        if not self._config.driver.ad_on:
+        if not self._config.driver.ad_on and not self._config.driver.fast_on:
             if self._config.driver.compute_fem:
                 self._compute_modalshapes()
                 self._compute_modalcouplings()

feniax/intrinsic/staticAD.py

@@ -6,7 +6,7 @@
 import feniax.intrinsic.ad_common as adcommon
 import feniax.intrinsic.couplings as couplings
 import feniax.systems.sollibs.diffrax as libdiffrax
-
+import feniax.intrinsic.staticFast as staticFast
 
 newton = partial(jax.jit, static_argnames=["F", "sett"])(libdiffrax.newton)
 _solve = partial(jax.jit, static_argnames=["eqsolver", "dq", "sett"])(isys._staticSolve)
@@ -22,55 +22,42 @@ def main_10g11_1(
     *args,
     **kwargs,
 ):
+    """
+    Static follower force, pseudo time input
+    """
+
     t = inputs["t"]
     Ka = config.fem.Ka
     Ma = config.fem.Ma
 
     t_loads = jnp.hstack([config.system.t, t])
-    tn = len(t_loads)
     config.system.build_states(config.fem.num_modes, config.fem.num_nodes)
-    q2_index = config.system.states["q2"]
     eigenvals = jnp.load(config.fem.folder / config.fem.eig_names[0])
     eigenvecs = jnp.load(config.fem.folder / config.fem.eig_names[1])
     reduced_eigenvals = eigenvals[: config.fem.num_modes]
     reduced_eigenvecs = eigenvecs[:, : config.fem.num_modes]
-    # solver_args = config.system.solver_settings
-    X = config.fem.X
-    (
-        phi1,
-        psi1,
-        phi2,
-        phi1l,
-        phi1ml,
-        psi1l,
-        phi2l,
-        psi2l,
-        omega,
-        X_xdelta,
-        C0ab,
-        C06ab
-    ) = adcommon._compute_modes(X, Ka, Ma, reduced_eigenvals, reduced_eigenvecs, config)
-
-    # gamma1 = couplings.f_gamma1(phi1, psi1)
-    gamma2 = couplings.f_gamma2(phi1ml, phi2l, psi2l, X_xdelta)
-    eta0 = jnp.zeros(config.fem.num_modes)
-    config.system.xloads.build_point_follower(config.fem.num_nodes, C06ab)
-    x_forceinterpol = config.system.xloads.x
-    y_forceinterpol = config.system.xloads.force_follower
-    dq_args = (eta0, gamma2, omega, phi1l, x_forceinterpol, y_forceinterpol)
-
-    # q = _solve(dq_static.dq_10g11, t_loads, q0, dq_args, config.system.solver_settings)
-    # q = _solve(dq_static.dq_10g11, t_loads, q0, dq_args, config.system.solver_settings)
-    q = _solve(
-        newton, dq_static.dq_10g11, t_loads, q0, dq_args, config.system.solver_settings
-    )
-    q2 = q[:, q2_index]
-    # X2, X3, ra, Cab = recover_staticfields(q, tn, X, q2_index,
-    #                                        phi2l, psi2l, X_xdelta, C0ab, config.fem)
-    X2, X3, ra, Cab = isys.recover_staticfields(
-        q2, tn, X, phi2l, psi2l, X_xdelta, C0ab, config
-    )
-    X1 = jnp.zeros_like(X2)
+    (phi1,psi1,
+     phi2,
+     phi1l,
+     phi1ml,
+     psi1l,
+     phi2l,
+     psi2l,
+     omega,
+     X_xdelta,
+     C0ab,
+     C06ab,
+     gamma2,
+     q,
+     X1, X2, X3, ra, Cab
+     ) = staticFast.main_10g11(q0,
+                               config,
+                               Ka=Ka,
+                               Ma=Ma,
+                               eigenvals=reduced_eigenvals,
+                               eigenvecs=reduced_eigenvecs,
+                               t_loads=t_loads)
+
     return adcommon._objective_output(
         q=q,
         X1=X1,
@@ -95,60 +82,47 @@ def main_10g11_3(
     *args,
     **kwargs,
 ):
+    """
+    Static follower force, Ka, Ma, phi1
+    """
+
     Ka = inputs["Ka"]
     Ma = inputs["Ma"]
     eigenvals = inputs[
         "eigenvals"
-    ]  # jnp.load(config.fem.folder / config.fem.eig_names[0])
+    ]
     eigenvecs = inputs[
         "eigenvecs"
-    ]  # jnp.load(config.fem.folder / config.fem.eig_names[1])
+    ]
 
     t_loads = config.system.t
-    tn = len(t_loads)
     config.system.build_states(config.fem.num_modes, config.fem.num_nodes)
-    q2_index = config.system.states["q2"]
-    # eigenvals = jnp.load(config.fem.folder / config.fem.eig_names[0])
-    # eigenvecs = jnp.load(config.fem.folder / config.fem.eig_names[1])
     reduced_eigenvals = eigenvals[: config.fem.num_modes]
     reduced_eigenvecs = eigenvecs[:, : config.fem.num_modes]
     # solver_args = config.system.solver_settings
-    X = config.fem.X
-    (
-        phi1,
-        psi1,
-        phi2,
-        phi1l,
-        phi1ml,
-        psi1l,
-        phi2l,
-        psi2l,
-        omega,
-        X_xdelta,
-        C0ab,
-        C06ab,
-    ) = adcommon._compute_modes(X, Ka, Ma, reduced_eigenvals, reduced_eigenvecs, config)
-
-    # gamma1 = couplings.f_gamma1(phi1, psi1)
-    gamma2 = couplings.f_gamma2(phi1ml, phi2l, psi2l, X_xdelta)
-    eta0 = jnp.zeros(config.fem.num_modes)
-    config.system.xloads.build_point_follower(config.fem.num_nodes, C06ab)
-    x_forceinterpol = config.system.xloads.x
-    y_forceinterpol = config.system.xloads.force_follower
-    dq_args = (eta0, gamma2, omega, phi1l, x_forceinterpol, y_forceinterpol)
-
-    # q = _solve(dq_static.dq_10g11, t_loads, q0, dq_args, config.system.solver_settings)
-    # q = _solve(dq_static.dq_10g11, t_loads, q0, dq_args, config.system.solver_settings)
-    q = _solve(
-        newton, dq_static.dq_10g11, t_loads, q0, dq_args, config.system.solver_settings
-    )
-    q2 = q[:, q2_index]
-    # X2, X3, ra, Cab = recover_staticfields(q, tn, X, q2_index,
-    #                                        phi2l, psi2l, X_xdelta, C0ab, config.fem)
-    X2, X3, ra, Cab = isys.recover_staticfields(
-        q2, tn, X, phi2l, psi2l, X_xdelta, C0ab, config
-    )
-    X1 = jnp.zeros_like(X2)
+    (phi1,
+     psi1,
+     phi2,
+     phi1l,
+     phi1ml,
+     psi1l,
+     phi2l,
+     psi2l,
+     omega,
+     X_xdelta,
+     C0ab,
+     C06ab,
+     gamma2,
+     q,
+     X1, X2, X3, ra, Cab
+     ) = staticFast.main_10g11(q0,
+                               config,
+                               Ka=Ka,
+                               Ma=Ma,
+                               eigenvals=reduced_eigenvals,
+                               eigenvecs=reduced_eigenvecs,
+                               t_loads=t_loads)
+
     return adcommon._objective_output(
         q=q,
         X1=X1,
@@ -160,5 +134,4 @@ def main_10g11_3(
         nodes=jnp.array(obj_args.nodes),
         components=jnp.array(obj_args.components),
         t=jnp.array(obj_args.t),
-        axis=obj_args.axis,
     )

feniax/intrinsic/staticFast.py

@@ -0,0 +1,82 @@
+import jax.numpy as jnp
+import jax
+from functools import partial
+import feniax.systems.intrinsic_system as isys
+import feniax.intrinsic.dq_static as dq_static
+import feniax.intrinsic.ad_common as adcommon
+import feniax.intrinsic.couplings as couplings
+import feniax.systems.sollibs.diffrax as libdiffrax
+
+
+newton = partial(jax.jit, static_argnames=["F", "sett"])(libdiffrax.newton)
+_solve = partial(jax.jit, static_argnames=["eqsolver", "dq", "sett"])(isys._staticSolve)
+
+
+@partial(jax.jit, static_argnames=["config"])
+def main_10g11(
+    q0,
+    config,
+    *args,
+    **kwargs,
+):
+    """
+    Static computation with follower forces
+    """
+
+    Ka = kwargs["Ka"] #config.fem.Ka
+    Ma = kwargs["Ma"] #config.fem.Ma
+    reduced_eigenvals = kwargs["eigenvals"] 
+    reduced_eigenvecs = kwargs["eigenvecs"] 
+    t_loads = kwargs["t_loads"]
+    tn = len(t_loads)
+    config.system.build_states(config.fem.num_modes, config.fem.num_nodes)
+    q2_index = config.system.states["q2"]
+    # solver_args = config.system.solver_settings
+    X = config.fem.X
+    (
+        phi1,
+        psi1,
+        phi2,
+        phi1l,
+        phi1ml,
+        psi1l,
+        phi2l,
+        psi2l,
+        omega,
+        X_xdelta,
+        C0ab,
+        C06ab
+    ) = adcommon._compute_modes(X, Ka, Ma, reduced_eigenvals, reduced_eigenvecs, config)
+
+    gamma2 = couplings.f_gamma2(phi1ml, phi2l, psi2l, X_xdelta)
+    eta0 = jnp.zeros(config.fem.num_modes)
+    config.system.xloads.build_point_follower(config.fem.num_nodes, C06ab)
+    x_forceinterpol = config.system.xloads.x
+    y_forceinterpol = config.system.xloads.force_follower
+    dq_args = (eta0, gamma2, omega, phi1l, x_forceinterpol, y_forceinterpol)
+
+    q = _solve(
+        newton, dq_static.dq_10g11, t_loads, q0, dq_args, config.system.solver_settings
+    )
+    q2 = q[:, q2_index]
+    X2, X3, ra, Cab = isys.recover_staticfields(
+        q2, tn, X, phi2l, psi2l, X_xdelta, C0ab, config
+    )
+    X1 = jnp.zeros_like(X2)
+    return (phi1,psi1,
+            phi2,
+            phi1l,
+            phi1ml,
+            psi1l,
+            phi2l,
+            psi2l,
+            omega,
+            X_xdelta,
+            C0ab,
+            C06ab,
+            gamma2,
+            q,
+            X1, X2, X3, ra, Cab
+            )
+
+

feniax/preprocessor/containers/intrinsicmodal.py

@@ -144,6 +144,7 @@ class Ddriver(DataContainer):
     compute_fem: bool = dfield("", default=True)
     save_fem: bool = dfield("", default=True)
     ad_on: bool = dfield("", default=False)
+    fast_on: bool = dfield("", default=False)
 
     def __post_init__(self):
         if self.sol_path is not None: