Solver and penalties update

[JR-1991]

This pull request introduces several improvements and new features to the neural modeling codebase, focusing on solver flexibility, penalty extensions, and model configuration. The most significant changes are the refactoring of solver handling to allow dynamic selection and configuration, the addition of a null space penalty for stoichiometric matrices, and enhancements to model initialization and plotting.

Solver Flexibility and API Improvements

• Refactored all neural model classes (NeuralBase, NeuralODE, RateFlowODE, UniversalODE) to allow passing the solver type and its configuration (rtol, atol, dt0) dynamically, both in __call__ and predict methods. This enables users to select and configure ODE solvers at runtime, improving flexibility and usability. Helper methods _create_controller and _instantiate_solver were added for controller and solver instantiation. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11]

Penalty Extensions for Stoichiometric Matrices

• Added a new penalty function penalize_null_space to encourage the stoichiometric matrix to have a null space of zero, which can be used to enforce physical constraints (e.g., conservation laws). Integrated this penalty into the for_rateflow penalty collection, allowing it to be configured via null_space_alpha. [1] [2] [3] [4]

Model Initialization and Configuration

• Improved RateFlowODE initialization to accept and validate a mass_constraint matrix, ensuring it matches the species count and enforcing shape requirements. This supports more robust modeling of conservation laws. Also, model fields now use equinox field annotations for better static/dynamic parameter handling. [1] [2] [3]

Plotting Enhancements

• Enhanced rate plotting in plot_learned_rates by improving grid display (major and minor ticks) and updating the right panel to overlay model predictions with measurements for clearer visualization.

Miscellaneous

• Added a new test script test.py demonstrating model creation, neural ODE instantiation, and parameter flattening for debugging or analysis.

[Copilot]

Pull Request Overview

This pull request introduces significant improvements to the neural modeling framework by adding solver flexibility, penalty extensions, and model configuration enhancements. The main focus is on enabling dynamic solver selection and configuration at runtime, along with new penalty functions for stoichiometric matrices.

• Refactored all neural model classes to support dynamic solver configuration (type, rtol, atol, dt0) in both __call__ and predict methods
• Added null space penalty for stoichiometric matrices to enforce conservation laws
• Enhanced RateFlowODE initialization with mass constraint validation and equinox field annotations

Reviewed Changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
test.py	Simple test script demonstrating model creation and parameter flattening
catalax/neural/neuralbase.py	Core refactoring for solver flexibility with helper methods and abstract interface
catalax/neural/neuralode.py	Updated to use new dynamic solver configuration system
catalax/neural/rateflow.py	Enhanced with mass constraint validation, equinox fields, and solver flexibility
catalax/neural/universalode.py	Updated to support dynamic solver configuration
catalax/neural/penalties/stoich_mat.py	Added new null space penalty function
catalax/neural/penalties/penalties.py	Integrated null space penalty into penalty collection
catalax/neural/plots/rateflow.py	Improved grid display and model prediction overlay

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}