Standardize integration methods

[chaoming0625]

This pull request includes major changes to the dendritex library, particularly focusing on refactoring and improving the differential equation integration modules. The most important changes include the removal of the State4Integral and DendriteDynamics classes, the introduction of the DiffEqState and DiffEqModule classes, and the renaming of several methods for consistency.

Refactoring and improvements:

• Removed State4Integral and DendriteDynamics classes, and introduced DiffEqState and DiffEqModule classes in dendritex/_integrators.py. This change simplifies the integration of differential equations and ensures better modularity. [1] [2]
• Replaced all instances of DendriteDynamics with DiffEqModule in dendritex/_base.py and dendritex/_integrators.py to reflect the new class structure. [1] [2] [3]
• Renamed methods before_integral and post_derivative to pre_integral and post_integral, respectively, for better clarity and consistency in dendritex/_base.py. [1] [2] [3]
• Added a check for the presence of the diffrax package and conditionally imported it in dendritex/_integrators.py to ensure compatibility and prevent import errors. [1] [2]
• Updated the Runge-Kutta integration methods to use the new DiffEqModule class and refactored the _rk_update function for better readability and maintainability in dendritex/_integrators.py. [1] [2]

Summary by Sourcery

Refactor the differential equation integration modules by replacing the State4Integral and DendriteDynamics classes with DiffEqState and DiffEqModule, respectively. Rename before_integral and post_derivative to pre_integral and post_integral for consistency. Add a check for the diffrax package and update the Runge-Kutta integration methods.

New Features:

• Introduce the DiffEqState class to manage the state of differential equations, including derivatives and diffusion.
• Introduce the DiffEqModule class to encapsulate the logic for computing derivatives and handling pre- and post-integration steps.

Enhancements:

• Standardize integration methods for differential equations, simplifying the process and improving modularity.

Tests:

• Add a conditional import for the diffrax package to ensure compatibility and prevent import errors.

[sourcery-ai]

Reviewer's Guide by Sourcery

This pull request refactors the differential equation integration modules in the dendritex library by replacing the State4Integral and DendriteDynamics classes with DiffEqState and DiffEqModule, respectively. This change simplifies integration and improves modularity. It also renames methods for consistency, adds a diffrax dependency check, and updates Runge-Kutta methods.

Sequence diagram showing the updated integration process

sequenceDiagram
    participant M as DiffEqModule
    participant S as DiffEqState

    Note over M,S: Integration process
    M->>M: pre_integral()
    M->>S: get current state
    M->>M: compute_derivative()
    M->>S: set derivative
    Note over M,S: Integration step happens here
    M->>M: post_integral()
    Note over M,S: Process repeats for next step

Loading

Class diagram showing the refactored integration classes

classDiagram
    class DiffEqState {
        +derivative: PyTree
        +diffusion: PyTree
        +value: PyTree
        +get_derivative()
        +set_derivative(value)
        +get_diffusion()
        +set_diffusion(value)
    }
    class DiffEqModule {
        +pre_integral()
        +compute_derivative()
        +post_integral()
    }

    DiffEqState --|> ShortTermState: inherits
    HHTypedNeuron --|> DiffEqModule: implements
    IonChannel --|> DiffEqModule: implements

    note for DiffEqState "Replaces State4Integral"
    note for DiffEqModule "Replaces DendriteDynamics"

Loading

State diagram for the integration lifecycle

stateDiagram-v2
    [*] --> Initialize
    Initialize --> PreIntegral: start integration
    PreIntegral --> ComputeDerivative: pre_integral()
    ComputeDerivative --> Integration: compute_derivative()
    Integration --> PostIntegral: integrate
    PostIntegral --> PreIntegral: post_integral()
    PostIntegral --> [*]: end simulation

Loading

File-Level Changes

Change	Details	Files
Replaced the State4Integral and DendriteDynamics classes with DiffEqState and DiffEqModule.	Introduced the DiffEqState class to manage differential equation state variables. Introduced the DiffEqModule class to encapsulate differential equation operations. Removed the State4Integral class. Removed the DendriteDynamics class. Moved Container and TreeNode classes from dendritex/_base.py to dendritex/_misc.py.	dendritex/_integrators.py dendritex/_base.py dendritex/_misc.py dendritex/channels/calcium.py dendritex/channels/potassium_calcium.py dendritex/channels/sodium.py dendritex/channels/potassium.py dendritex/channels/hyperpolarization_activated.py dendritex/neurons/multi_compartment.py dendritex/neurons/single_compartment.py examples/fitting_a_hh_neuron/main.py examples/simple_dendrite_model.py
Renamed before_integral and post_derivative to pre_integral and post_integral.	Renamed before_integral to pre_integral. Renamed post_derivative to post_integral.	dendritex/_base.py dendritex/_integrators.py dendritex/channels/calcium.py dendritex/channels/potassium_calcium.py dendritex/channels/sodium.py dendritex/channels/potassium.py dendritex/channels/hyperpolarization_activated.py dendritex/neurons/multi_compartment.py dendritex/neurons/single_compartment.py
Added a check for the diffrax package and updated Runge-Kutta methods.	Added conditional import for diffrax based on installation status. Updated Runge-Kutta integration methods to use the DiffEqModule class. Refactored the _rk_update function.	dendritex/_integrators.py
Updated import paths and method signatures.	Updated import paths to reflect the new class structure. Modified method signatures to accommodate the new classes.	dendritex/neurons/multi_compartment.py dendritex/neurons/single_compartment.py dendritex/ions/calcium.py examples/fitting_a_hh_neuron/main.py examples/simple_dendrite_model.py
Updated setup files and dependencies.	Removed diffrax from dependencies. Updated project URL.	setup.py pyproject.toml
Added an example script for calcium activation functions.	Created a new script to visualize calcium activation functions.	examples/calcium_activation_functions.py
Updated example scripts.	Modified example scripts to use the new classes and methods. Removed unnecessary comments and updated plotting code.	examples/hh_neuron.py examples/fitting_a_hh_neuron/main.py
Updated leaky channel implementation.	Renamed methods for consistency with other channel implementations.	dendritex/channels/leaky.py

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time. You can also use
  this command to specify where the summary should be inserted.

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey @chaoming0625 - I've reviewed your changes and they look great!

Here's what I looked at during the review

• 🟡 General issues: 1 issue found
• 🟢 Security: all looks good
• 🟢 Testing: all looks good
• 🟢 Complexity: all looks good
• 🟢 Documentation: all looks good

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[sourcery-ai]

suggestion (performance): Consider caching dt to avoid repeated environment lookups

Getting dt from the environment in each tree map operation could be inefficient. Consider getting it once at the start of the function.

Suggested implementation:

def _rk_update(
    dt: float,
    coeff: Sequence,
    st: bst.State,

The complete implementation will also require:

1. Updating all calls to _rk_update to pass the dt parameter
2. Modifying the internal logic of _rk_update to use the passed dt parameter instead of looking it up from the environment
3. Since I can't see the full implementation details of where dt is currently being accessed, you'll need to remove those environment lookups and use the passed dt parameter instead