[Models] Stopped zeroing out node/vertex loads if tmax>1

Author: arpastrana

CHANGELOG.md

@@ -9,7 +9,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-- Added `jax.lax.custom_linear_solver()` to backward rule of implicit `custom_vjp` of fixed-point solver. This new function acts as a thin wrapper around the *sparse* linear solver of `EquilibriumModel.linearsolve_fn()` that defines a custom transpose rule to be compatible with `jax.linear_transpose()`. The transpose is needed by `lineax`, inside `FunctionLinearOperator`. Without the wrapper and the transpose, we cannot use implicit differentiation with a sparse linear solver and a fixed-point solver. Now we can.
+- Added automatic support for dense and sparse stiffness matrices in `custom_vjp` of `solver_fixed_point_implicit()`. For sparse matrices, we apply `jax.lax.custom_linear_solver()` as a thin wrapper around the sparse linear solve defined in `EquilibriumModel.linearsolve_fn()` to generate a transpose rule for it. The transpose rule is required by `lineax`, inside `FunctionLinearOperator`. Without the wrapper and the transpose, we cannot use implicit differentiation with a sparse linear solver and a fixed-point solver. Now we can.
 - Implemented `geometry.length_vector_sqrd()`.
 - Print out statistics with `ndigits` of precision in `FDDatastructure.print_stats()`.
 - Listed `lineax` and `optimistix` as dependencies.

src/jax_fdm/equilibrium/models.py

@@ -172,8 +172,6 @@ def __call__(self, params, structure):
         xyz_free = self.equilibrium(q, xyz_fixed, load_nodes, structure)
 
         if tmax > 1:
-            load_nodes = jnp.zeros_like(load_nodes)
-            load_state = LoadState(load_nodes, load_state.edges, load_state.faces)
 
             xyz_free = self.eq_iterative_fn(
                 q,