Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[fitting] model fitting functionality with optimization methods #582

Draft
wants to merge 8 commits into
base: master
Choose a base branch
from
Draft

[fitting] model fitting functionality with optimization methods #582

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
6154f66
start model fitting, add optimization methods
chaoming0625 Jan 4, 2024
a303e99
fix
chaoming0625 Jan 4, 2024
3bf1fb7
Merge branch 'master' into modelfitting
chaoming0625 Mar 9, 2024
9ca2d6c
add `skopt` Optimization methods
chaoming0625 Mar 9, 2024
799c41f
upgrade optimizers
chaoming0625 Mar 10, 2024
bde2c1e
move optimization methods in `brainpy.analysis` into `brainpy.optimiz…
chaoming0625 Mar 10, 2024
b1c0997
standardize SGD optimizers
chaoming0625 Mar 10, 2024
78a5c0d
fix
chaoming0625 Mar 10, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
45 changes: 26 additions & 19 deletions brainpy/_src/analysis/lowdim/lowdim_analyzer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,17 +3,18 @@
import warnings
from functools import partial

import numpy as np
import jax
import numpy as np
from jax import numpy as jnp
from jax import vmap
from jax.scipy.optimize import minimize

from brainpy import errors, tools
import brainpy._src.math as bm
from brainpy._src.math.object_transform.base import Collector
from brainpy import errors, tools
from brainpy._src.analysis import constants as C, utils
from brainpy._src.analysis.base import DSAnalyzer
from brainpy._src.math.object_transform.base import Collector
from brainpy._src.optimizers.brentq import jax_brentq, ECONVERGED, brentq_candidates, brentq_roots

pyplot = None

Expand Down Expand Up @@ -316,7 +317,9 @@ def F_vmap_fp_aux(self):
def F_fixed_point_opt(self):
if C.F_fixed_point_opt not in self.analyzed_results:
def f(start_and_end, *args):
return utils.jax_brentq(self.F_fx)(start_and_end[0], start_and_end[1], args)
return jax_brentq(utils.f_without_jaxarray_return(self.F_fx))(
start_and_end[0], start_and_end[1], args
)

self.analyzed_results[C.F_fixed_point_opt] = f
return self.analyzed_results[C.F_fixed_point_opt]
Expand Down Expand Up @@ -387,7 +390,7 @@ def _get_fixed_points(self, candidates, *args, num_seg=None, tol_aux=1e-7, loss_
# optimize the fixed points
res = self.F_vmap_fp_opt(X, *args)
losses = self.F_vmap_fp_aux(res['root'], *args)
valid_or_not = jnp.logical_and(res['status'] == utils.ECONVERGED, losses <= tol_aux)
valid_or_not = jnp.logical_and(res['status'] == ECONVERGED, losses <= tol_aux)
ids = np.asarray(jnp.where(valid_or_not)[0])
fps = np.asarray(res['root'])[ids]
args = tuple(a[ids] for a in args)
Expand Down Expand Up @@ -569,10 +572,14 @@ def F_fixed_point_opt(self):
if self._can_convert_to_one_eq():
if self.convert_type() == C.x_by_y:
def f(start_and_end, *args):
return utils.jax_brentq(self.F_y_convert[1])(start_and_end[0], start_and_end[1], args)
return jax_brentq(utils.f_without_jaxarray_return(self.F_y_convert[1]))(
start_and_end[0], start_and_end[1], args
)
else:
def f(start_and_end, *args):
return utils.jax_brentq(self.F_x_convert[1])(start_and_end[0], start_and_end[1], args)
return jax_brentq(utils.f_without_jaxarray_return(self.F_x_convert[1]))(
start_and_end[0], start_and_end[1], args
)
self.analyzed_results[C.F_fixed_point_opt] = f

else:
Expand Down Expand Up @@ -718,23 +725,23 @@ def _get_fx_nullcline_points(self, coords=None, tol=1e-7, num_segments=1, fp_aux
# auxiliary functions
f2 = lambda y, x, *pars: self.F_fx(x, y, *pars)
vmap_f2 = jax.jit(vmap(f2), device=self.jit_device)
vmap_brentq_f2 = jax.jit(vmap(utils.jax_brentq(f2)), device=self.jit_device)
vmap_brentq_f1 = jax.jit(vmap(utils.jax_brentq(self.F_fx)), device=self.jit_device)
vmap_brentq_f2 = jax.jit(vmap(jax_brentq(utils.f_without_jaxarray_return(f2))), device=self.jit_device)
vmap_brentq_f1 = jax.jit(vmap(jax_brentq(utils.f_without_jaxarray_return(self.F_fx))), device=self.jit_device)

# num segments
for _j, Ps in enumerate(par_seg):
if len(par_seg.arg_id_segments[0]) > 1:
utils.output(f"{C.prefix}segment {_j} ...")
if coords == self.x_var + '-' + self.y_var:
x0s, x1s, vps = utils.brentq_candidates(self.F_vmap_fx, *((xs, ys) + Ps))
x_values_in_fx, out_args = utils.brentq_roots2(vmap_brentq_f1, x0s, x1s, *vps)
x0s, x1s, vps = brentq_candidates(self.F_vmap_fx, *((xs, ys) + Ps))
x_values_in_fx, out_args = brentq_roots(vmap_brentq_f1, x0s, x1s, *vps)
y_values_in_fx = out_args[0]
p_values_in_fx = out_args[1:]
x_values_in_fx, y_values_in_fx, p_values_in_fx = \
self._fp_filter(x_values_in_fx, y_values_in_fx, p_values_in_fx, fp_aux_filter)
elif coords == self.y_var + '-' + self.x_var:
x0s, x1s, vps = utils.brentq_candidates(vmap_f2, *((ys, xs) + Ps))
y_values_in_fx, out_args = utils.brentq_roots2(vmap_brentq_f2, x0s, x1s, *vps)
x0s, x1s, vps = brentq_candidates(vmap_f2, *((ys, xs) + Ps))
y_values_in_fx, out_args = brentq_roots(vmap_brentq_f2, x0s, x1s, *vps)
x_values_in_fx = out_args[0]
p_values_in_fx = out_args[1:]
x_values_in_fx, y_values_in_fx, p_values_in_fx = \
Expand Down Expand Up @@ -812,21 +819,21 @@ def _get_fy_nullcline_points(self, coords=None, tol=1e-7, num_segments=1, fp_aux
# auxiliary functions
f2 = lambda y, x, *pars: self.F_fy(x, y, *pars)
vmap_f2 = jax.jit(vmap(f2), device=self.jit_device)
vmap_brentq_f2 = jax.jit(vmap(utils.jax_brentq(f2)), device=self.jit_device)
vmap_brentq_f1 = jax.jit(vmap(utils.jax_brentq(self.F_fy)), device=self.jit_device)
vmap_brentq_f2 = jax.jit(vmap(jax_brentq(utils.f_without_jaxarray_return(f2))), device=self.jit_device)
vmap_brentq_f1 = jax.jit(vmap(jax_brentq(utils.f_without_jaxarray_return(self.F_fy))), device=self.jit_device)

for j, Ps in enumerate(par_seg):
if len(par_seg.arg_id_segments[0]) > 1: utils.output(f"{C.prefix}segment {j} ...")
if coords == self.x_var + '-' + self.y_var:
starts, ends, vps = utils.brentq_candidates(self.F_vmap_fy, *((xs, ys) + Ps))
x_values_in_fy, out_args = utils.brentq_roots2(vmap_brentq_f1, starts, ends, *vps)
starts, ends, vps = brentq_candidates(self.F_vmap_fy, *((xs, ys) + Ps))
x_values_in_fy, out_args = brentq_roots(vmap_brentq_f1, starts, ends, *vps)
y_values_in_fy = out_args[0]
p_values_in_fy = out_args[1:]
x_values_in_fy, y_values_in_fy, p_values_in_fy = \
self._fp_filter(x_values_in_fy, y_values_in_fy, p_values_in_fy, fp_aux_filter)
elif coords == self.y_var + '-' + self.x_var:
starts, ends, vps = utils.brentq_candidates(vmap_f2, *((ys, xs) + Ps))
y_values_in_fy, out_args = utils.brentq_roots2(vmap_brentq_f2, starts, ends, *vps)
starts, ends, vps = brentq_candidates(vmap_f2, *((ys, xs) + Ps))
y_values_in_fy, out_args = brentq_roots(vmap_brentq_f2, starts, ends, *vps)
x_values_in_fy = out_args[0]
p_values_in_fy = out_args[1:]
x_values_in_fy, y_values_in_fy, p_values_in_fy = \
Expand Down
1 change: 0 additions & 1 deletion brainpy/_src/analysis/utils/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,6 @@
from .function import *
from .measurement import *
from .model import *
from .optimization import *
from .others import *
from .outputs import *
from .visualization import *
Loading
Loading