Update code#27
Merged
chaoming0625 merged 6 commits intomainfrom Feb 19, 2025
Merged
Conversation
Contributor
Reviewer's Guide by SourceryThis pull request refactors the Updated class diagram for HHTypedNeuronclassDiagram
class HHTypedNeuron {
-size: brainstate.typing.Size
-name: Optional[str]
-ion_channels: Dict[str, IonChannel]
+__init__(size: brainstate.typing.Size, name: Optional[str], **ion_channels)
+pop_size: Tuple[int, ...]
+n_compartment: int
+current(*args, **kwargs)
+pre_integral(*args, **kwargs)
+compute_derivative(*args, **kwargs)
+post_integral(*args, **kwargs)
+init_state(batch_size=None)
+reset_state(batch_size=None)
+add(**elements)
}
HHTypedNeuron ..> IonChannel : contains
HHTypedNeuron --|> brainstate.nn.Dynamics
HHTypedNeuron --|> Container
HHTypedNeuron --|> DiffEqModule
note for HHTypedNeuron "Base class for Hodgkin-Huxley type neuron models"
Updated class diagram for IonChannelclassDiagram
class IonChannel {
-size: brainstate.typing.Size
-name: Optional[str]
+__init__(size: brainstate.typing.Size, name: Optional[str])
+varshape
+current(*args, **kwargs)
+pre_integral(*args, **kwargs)
+compute_derivative(*args, **kwargs)
+post_integral(*args, **kwargs)
+reset_state(*args, **kwargs)
+init_state(*args, **kwargs)
}
IonChannel --|> brainstate.graph.Node
IonChannel --|> TreeNode
IonChannel --|> DiffEqModule
note for IonChannel "Base class for ion channel modeling"
Updated class diagram for IonclassDiagram
class Ion {
-channels: Dict[str, Channel]
-_external_currents: Dict[str, Callable]
+__init__(size: brainstate.typing.Size, name: Optional[str], **channels)
+external_currents
+pre_integral(V)
+compute_derivative(V)
+post_integral(V)
+current(V, include_external: bool = False)
+init_state(V, batch_size: int = None)
+reset_state(V, batch_size: int = None)
+register_external_current(key: str, fun: Callable)
+pack_info() IonInfo
+add(**elements)
}
Ion --|> IonChannel
Ion --|> Container
Ion ..> Channel : contains
note for Ion "Base class for modeling the Ion dynamics"
Updated class diagram for MixIonsclassDiagram
class MixIons {
-ions: Tuple[Ion]
-_ion_types: Tuple[Type[Ion]]
-channels: Dict[str, Channel]
+__init__(*ions, name: Optional[str] = None, **channels)
+ion_types
+pre_integral(V)
+compute_derivative(V)
+post_integral(V)
+current(V)
+init_state(V, batch_size: int = None)
+reset_state(V, batch_size=None)
+add(**elements)
-_check_hierarchy(ions, leaf)
-_get_ion(ion)
}
MixIons --|> IonChannel
MixIons --|> Container
MixIons ..> Channel : contains
note for MixIons "Mixing Ions"
Updated class diagram for ChannelclassDiagram
class Channel {
}
Channel --|> IonChannel
note for Channel "Base class for modeling channel dynamics."
Updated class diagram for PotassiumChannelclassDiagram
class PotassiumChannel {
}
PotassiumChannel --|> Channel
note for PotassiumChannel "Base class for potassium channel dynamics."
Updated class diagram for SodiumChannelclassDiagram
class SodiumChannel {
}
SodiumChannel --|> Channel
note for SodiumChannel "Base class for sodium channel dynamics."
File-Level Changes
Tips and commandsInteracting with Sourcery
Customizing Your ExperienceAccess your dashboard to:
Getting Help
|
![sourcery-ai[bot]](https://avatars.githubusercontent.com/in/48477?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Hey @chaoming0625 - I've reviewed your changes - here's some feedback:
Overall Comments:
- Consider adding a brief migration guide or deprecation warnings for users who are upgrading from
bsttobrainstate. - The documentation strings are very detailed, but consider adding some simple examples to the docstrings to improve usability.
Here's what I looked at during the review
- 🟡 General issues: 1 issue found
- 🟢 Security: all looks good
- 🟢 Testing: all looks good
- 🟡 Complexity: 2 issues found
- 🟢 Documentation: all looks good
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
| channel.reset_state(self.V.value, batch_size=batch_size) | ||
|
|
||
| def add_elem(self, **elements): | ||
| def add(self, **elements): |
Contributor
There was a problem hiding this comment.
suggestion: Renaming 'add_elem' to 'add' improves clarity.
Ensure that all external references to this method are updated accordingly and that its behavior remains consistent.
Suggested implementation:
result = self.add(**new_elements)# Usage example: creating a new ion channel via add
channel = neuron_group.add(channel_param=value)Review other parts of your codebase for references to the old add_elem method and update them to use add.
| self.phi = brainstate.init.param(phi, self.varshape, allow_none=False) | ||
| self.g_max = brainstate.init.param(g_max, self.varshape, allow_none=False) | ||
|
|
||
| def init_state(self, V, Na: IonInfo, batch_size=None): |
Contributor
There was a problem hiding this comment.
issue (complexity): Consider refactoring the repetitive DiffEqState initialization and derivative calculations in the INa_Rsg subclass into helper functions or loops to reduce code duplication.
To reduce the complexity and duplication in the `INa_Rsg` subclass (without changing behavior), consider refactoring the repeated DiffEqState initialization and derivative calls into helper functions or loops.
For example, in the `init_state` method you can use a dictionary of state names and their initializers:
```python
def init_state(self, V, Na: IonInfo, batch_size=None):
init_funcs = {
"C1": u.math.ones,
"C2": u.math.ones,
"C3": u.math.ones,
"C4": u.math.ones,
"C5": u.math.ones,
"I1": u.math.ones,
"I2": u.math.ones,
"I3": u.math.ones,
"I4": u.math.ones,
"I5": u.math.ones,
"I6": u.math.ones,
"O": u.math.zeros,
"B": u.math.zeros,
}
for name, init_func in init_funcs.items():
setattr(self, name, DiffEqState(bst.init.param(init_func, self.varshape, batch_size)))
self.normalize_states([getattr(self, key) for key in init_funcs])Likewise, you can capture common derivative patterns into helper functions. For instance, if a derivative has the form
(A.value * f_A(V) + B.value * f_B(V) − state.value * (f_C(V) + f_D(V))) / u.ms
you could factor it out (you’ll need to adapt this for each unique state):
def calc_derivative(self, state, contributions, losses, dt=u.ms):
num = sum(func(V) * val for val, func in contributions) - state.value * sum(func(V) for func in losses)
state.derivative = num / dt
def compute_derivative(self, V, Na: IonInfo):
# Example for C1: (I1 * bi1(V) + C2 * b01(V) - C1 * (fi1(V) + f01(V))) / u.ms
self.calc_derivative(
self.C1,
contributions=[(self.I1.value, self.bi1), (self.C2.value, self.b01)],
losses=[self.fi1, self.f01]
)
# Similar refactoring can be applied for other states
# ...These small refactoring steps abstract the repetitive patterns while keeping functionality intact.
| # k1: first derivative step | ||
| assert len(tableau.A[0]) == 0, f'The first row of A must be empty. Got {tableau.A[0]}' | ||
| with bst.environ.context(t=t + tableau.C[0] * dt), bst.StateTraceStack() as trace: | ||
| with brainstate.environ.context(t=t + tableau.C[0] * dt), brainstate.StateTraceStack() as trace: |
Contributor
There was a problem hiding this comment.
issue (complexity): Consider creating a context manager to encapsulate the combined contexts and reduce nesting.
You can reduce the nesting by extracting the repeated context‐combinations into a helper/context manager. For example, you might define
from contextlib import contextmanager
@contextmanager
def env_context_trace(t, *, ms=False):
time_val = t * u.ms if ms else t
with brainstate.environ.context(t=time_val), brainstate.StateTraceStack() as trace:
yield traceThen, in the integration routines (e.g. in _general_rk_step and _diffrax_solve), replace
with brainstate.environ.context(t=t + tableau.C[0] * dt), brainstate.StateTraceStack() as trace:
...with
with env_context_trace(t + tableau.C[0] * dt, ms=False) as trace:
...This refactoring keeps the functionality unchanged while reducing cognitive load by flattening the nesting.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This pull request includes a significant refactor of the
braincellpackage, mainly focusing on replacing the usage of thebstmodule with thebrainstatemodule across various files. Additionally, it includes improvements to theContainerandTreeNodeclasses inbraincell/_misc.py.Module replacement:
braincell/_integrators.py: Replaced all instances ofbstwithbrainstatefor context management, state tracing, and typing. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] [12] [13] [14] [15]braincell/_protocol.py: UpdatedDiffEqStateandDiffEqModuleclasses to usebrainstateinstead ofbst. [1] [2] [3] [4]braincell/channel/calcium.py: Replacedbstwithbrainstate.Class improvements:
braincell/_misc.py: EnhancedContainerandTreeNodeclasses with detailed docstrings and additional methods for better type checking and element management. [1] [2] [3] [4] [5] [6] [7]Documentation:
README.md: Removed outdated comments about thebraincellpackage.Summary by Sourcery
Refactors the
braincellpackage to use thebrainstatemodule instead ofbst, enhances theContainerandTreeNodeclasses, and updates documentation.Enhancements:
ContainerandTreeNodeclasses inbraincell/_misc.pywith detailed docstrings and additional methods for better type checking and element management.braincell/_base.pyto provide a more comprehensive overview of the module's purpose, class hierarchy, and usage.Documentation:
braincellpackage inREADME.md.