Missing eq_param_index when loading BundleSolution

[ptflores1]

When calling BundleSolution1D.load() the new instance is initialized without an eq_param_index argument. Since the _diff_eqs_wrapper function uses the local variable eq_param_index there is no way of recovering the original value and the functionality is lost.

[ptflores1]

The code in

neurodiffeq/neurodiffeq/solvers_utils.py

Lines 519 to 533 in 1b6d8c5

	elif load_dict["type_name"] == "BundleSolver1D" or load_dict["parent_type_name"] == "BundleSolver1D":
	t_min = load_dict['solver'].r_min[0]
	t_max = load_dict['solver'].r_max[0]
	solver = cls(ode_system=de_system,
	conditions=cond,
	metrics=load_dict['metrics'],
	nets=nets,
	optimizer=optimizer,
	train_generator=train_generator,
	valid_generator=valid_generator,
	t_min=t_min,
	t_max=t_max,
	theta_min=tuple(load_dict['solver'].r_min[1:]),
	theta_max=tuple(load_dict['solver'].r_max[1:]))

can be modified to

elif load_dict["type_name"] == "BundleSolver1D" or load_dict["parent_type_name"] == "BundleSolver1D": 
     t_min = load_dict['solver'].r_min[0] 
     t_max = load_dict['solver'].r_max[0] 
  
     solver = cls(ode_system=de_system, 
                  conditions=cond, 
                  metrics=load_dict['metrics'], 
                  nets=nets, 
                  optimizer=optimizer, 
                  train_generator=train_generator, 
                  valid_generator=valid_generator, 
                  t_min=t_min, 
                  t_max=t_max, 
                  theta_min=tuple(load_dict['solver'].r_min[1:]), 
                  theta_max=tuple(load_dict['solver'].r_max[1:]),
                  eq_param_index=(index - len(cond) - 1 for index in load_dict['solver'].eq_param_index)) # new line
) 

[ptflores1]

@shuheng-liu would that be okay for a PR?

[sathvikbhagavan]

Thanks for the report! I don't think we need to recompute as it is stored in https://github.com/NeuroDiffGym/neurodiffeq/blob/master/neurodiffeq/solvers.py#L1354

It should be fixed in #212

[ptflores1]

Great! Thanks