arpastrana/jax_fdm: Release v0.9.0

Read full text
Authors: Rafael Pastrana, Deniz Oktay
Cited by: 0

Abstract

Added Implemented a vanilla GradientDescent optimizer. Implemented splu_cpu and splu_solve to cache sparse stiffness matrix factorization using scipy's sparse infrastructure. These functions specifically instantiate scipy.sparse.linalg.splu, a sparse LU factorization using the SuperLU package under the hood. We implemented both these functions by wrapping them into jax.pure_callback for compatibility with jax.jit. The fixed_point_bwd_adjoint of the fixed point solver solver_fixedpoint_implicit now caches a sparse stiffness matrix to avoid refactorizing it every time the a linear solve is invoked by lineax to resolve the adjoint system. Caching reduces computation cost considerably (about 30% in the British Museum example). Created GradientFreeOptimizer() class for other gradient-free and evolutionary algorithms to subclass it. This new parent class implements specific problem() and solve() methods that accounts for the different naming convention of the objective function in scipy.minimize (fun vs. func). Moreover, the new method does not calculate loss and gradients per iteration, only the loss, since the gradient is unnecessary for gradient-free optimization. Omitting the gradient calculation and focusing on the loss alone speeds up optimization wall time. Added attribute EquilibriumModel.iterload_fn to host a callable that modifies a load state before starting iterative FDM. This is our current alternative to the previous option that zeroed out the vertex/node loads without a warning. It also opens the door to other types of load changes between the first and second step of form-finding, giving users more flexibility. Modified the functions in fdm.py by exposing iterload_fn as an argument. These functions are model_from_sparsity(), fdm(), and constrained_fdm(). Added Viewer.save() to automatically save images of the current viewer scene. 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(). Pass implicit_diff argument to solver_fixedpoint. Print out statistics with ndigits of precision in FDDatastructure.print_stats(). Listed lineax and optimistix as dependencies. Added EquilibriumModel.residual_free_matrix() to compute the matrix with the residual force vectors on the free vertices of a structure. Set up a custom_vjp with the implicit function theorem for the nonlinear equilibrium solvers (least squares and root finding). Wrapped up 3 different optimistix optimizers to solve the nonlinear equilibrium problem with shape dependent loads. These solvers are least-squares and root-finding optimizers: Newton, Dogleg, and LevenbergMarquardt. These solvers minimize the residual function explicitly, which differs from the fixed-point iterators that solve the equilibrium problem by minimizing the difference between the XYZ coordinates of the free vertices of a structure over two consecutive iterations. These solvers are listed in the API as solver_newton, solver_dogleg, and solver_levenger_marquardt. Exposed the maxcor argument in scipy's LBFGSB() wrapper. This argument controls the number of approximation terms of the full Hessian. Added error_terms argument in LossPlotter.plot() to select what error and regularization terms are plotted. Implemented EquilibriumModel.load_xyz_matrix to calculate the load matrices for shape dependent loads. Implemented EquilibriumModel.load_xyz_matrix_from_r_fixed to calculate the load matrices for shape dependent loads. Added is_solver_fixedpoint and is_solver_leastsquares to check the type of an iterative solver. Added solver_gauss_newton to calculate equilibrium states in the presence of shape-dependent loads by minimizing a residual function explicitly. This solver is wrapped up from jaxopt. Exposed report_breakdown argument in LossPlotter.plot() to optionally plot the contributions of the error and regularization terms of a Loss function. Implemented Optimizer.options() to allow for method-specific setup in scipy.optimize.minimize. This new method assembles the options dictionary required by scipy in a way that can be customizable per optimizer. Changed Renamed network argument input to Optimizer.solve() to datastructure for correctness. This method can both ingest an FDNetwork or an FDMesh. The callback function is called once before optimization starts in Optimizer.solve() and in GradientFreeOptimizer.solve(). Fixed bug that forgot to fix the random seed of DualAnnealing() despite being passed in as an argument at initialization. Stop passing EquilibriumModel.linearsolve_fn() as solver_kwargs of fixed point solver. This function was used only for the custom_vjp operations. The custom_vjp now selects and appropriate linear solver based on whether EquilibriumModel.stiffness_matrix() is a sparse jax.CSC object or not. Now, goals.NetworkSmoothGoal() calculates the fairness energy on all the vertices. Previously, it only considered the free vertices. FDDatastructure.print_stats() doesn't reporting positive or negative force and force densities if the datastructure doesn't contain them. Changed the diagonal matrix generated by EquilibriumStructureSparse._get_sparse_diag_data() from jax.experimental.sparse.CSC to jax.experimental.sparse.BCSR for enabling Jacobian computations. Exposed arrow parameters (head width, head portion, body width and minimum width) in FDVectorPlotterArtist. Fixed bug in FDNetworkViewerArtist() that ignored custom colors when drawing node loads. LossPlotter.print_error_stats() reports the loss and error values with up 6 digits of precision. Renamed EquilibriumModel.force_fixed_matrix() to EquilibriumModel.residual_fixed_matrix(). Renamed EquilibriumModel.force_matrix() to EquilibriumModel.load_matrix(). EquilibriumParametersState.from_datastructure takes dtype as optional input. It defaults to jax.numpy.float64. Changed DTYPE_INT to int64 instead of int32. To calculate the local coordinate system of a mesh face, loads.face_load_lcs() no longer replaces vertex indices that were padded with a -1 with face[0]. Instead, it takes all the vertices in a face to get the XYZ coordinates of the face polygon. The previous behavior led to due excessive compilation time and XLA warnings due to "constant folding" problems because of the index replacement with a vmapped jnp.where(). ~~The faces generated by EquilibriumStructureMesh.from_mesh are now padded with first index face[0] instead of a -1.~~ EquilibriumModel automatically picks the iterative equilibrium function based on solver input as iterativesolve_fn. Functions EquilibriumModel.equilibrium() and EquilibriumModel.equilibrium_iterative() return xyz_free instead of xyz. The concatenation of xyz_free and xyz_fixed needed to build xyz is now handled by EquilibriumModel.__call__(). Set implicit=False and unroll=False in solver_anderson when performing implicit differentiation on iterative equilibrium calculation with implicit_diff=True. Set implicit=False and unroll=False in solver_fixedpoint when performing implicit differentiation on iterative equilibrium calculation with implicit_diff=True. jax_fdm.equilibrium.datastructure_validate now reports the number of edges with zero force densities. LBFGSBS became a subclass of LBFGSB instead of Optimizer. Disabled hard assertion test that ensured that every edge in a topology.Mesh() object was connected to at most 2 faces (manifoldness preservation). Now we print out a warning since we are all consenting adults over here. The implications of this change is that area load calculations might be incorrect, but this needs to be more thoroughly tested at a later time. Removed Removed attribute EquilibriumModel.ignore_nodes_load/EquilibriumModel.nodes_loads_iter because zeroing out node lodes behind the scenes can come bite you. Better to explicitly zero out them. Deleted jax_fdm.loads._faces_load_2 because it was not longer used.

Publication details

Journal / source
Zenodo (CERN European Organization for Nuclear Research)
Publication date
6 January 2026