specifying LeastSquaresProblem from_tuples with mixed array-based and float goals

[smiet]

I encountered a bug() when trying to define a LeastSquaresProblem from tuples when the objective tuples consist of a mix of 1d-array objectives and scalar objectives.

Below is a MWE:

import numpy as np
from simsopt.objectives import LeastSquaresProblem
from simsopt.solve import least_squares_serial_solve
from simsopt import make_optimizable
from simsopt.objectives.functions import Identity

singleton = Identity()
obj_tuple1 = (singleton.f, 1, 1)


singleton2 = Identity()
def return_list(opt_obj): 
    return opt_obj.f()*np.ones(3)


list_opt = make_optimizable(return_list, singleton2)
obj_tuple2 = (list_opt.J, np.ones(3), 1) 

prob = LeastSquaresProblem.from_tuples([obj_tuple1, obj_tuple2])
least_squares_serial_solve(prob)

This throws the ValueError:

ValueError                                Traceback (most recent call last)
Cell In[1], line 19
     16 list_opt = make_optimizable(return_list, singleton2)
     17 obj_tuple2 = (list_opt.J, np.ones(3), 1) 
---> 19 prob = LeastSquaresProblem.from_tuples([obj_tuple1, obj_tuple2])
     21 least_squares_serial_solve(prob)

File ~/code/simsopt/src/simsopt/objectives/least_squares.py:136, in LeastSquaresProblem.from_tuples(cls, tuples, fail)
    127 """
    128 Initializes graph based LeastSquaresProblem from a sequence of tuples
    129 containing *f_in*, *goal*, and *weight*.
   (...)
    133         each tuple (the specified order matters).
    134 """
    135 funcs_in, goals, weights = zip(*tuples)
--> 136 return cls(goals, weights, funcs_in=funcs_in, fail=fail)

File ~/code/simsopt/src/simsopt/objectives/least_squares.py:68, in LeastSquaresProblem.__init__(self, goals, weights, funcs_in, depends_on, opt_return_fns, fail)
     66 if np.any(np.asarray(weights) < 0):
     67     raise ValueError('Weight cannot be negative')
---> 68 self.goals = np.asarray(goals)
     69 self.inp_weights = np.asarray(weights)
     70 self.fail = fail

This is because LeastSquaresProblem.from_tuples zips the inputs, and np.asarray receives a mix of integers and lists.

This could be solved by ravel-concatenating the inputs in from_tuples:

    @classmethod
    def from_tuples(cls,
                    tuples: Sequence[Tuple[Callable, Real, Real]],
                    fail: Union[None, float] = 1.0e12) -> LeastSquaresProblem:
        """
        Initializes graph based LeastSquaresProblem from a sequence of tuples
        containing *f_in*, *goal*, and *weight*.

        Args:
            tuples: A sequence of tuples containing (f_in, goal, weight) in
                each tuple (the specified order matters).
        """
        funcs_in, goals, weights = zip(*tuples)
        goal_array =  np.concatenate([np.ravel(x) for x in goals])
        weigth_array =  np.concatenate([np.ravel(x) for x in weights])
        return cls(goal_array, weigth_array, funcs_in=funcs_in, fail=fail)

The reason this is an issue, and not a pull-request, is that I don't know if the weight-arrays and goal-arrays are still correctly handled after this change. I don't see the code that expands the weights array to the same size of the goals array, but on the other hand, from_tuples only worked if weight was an integer, even if goals and obj_return.J() were arrays (a formatting function would fail in the argument to an objective_file.write call).

So: Should weight_array be ravel-concatenated as well? And should the user be able to specify weight as an array when the goal is an array? Are the weights still correctly assigned to the goals when mixed input is provided?

[also as an aside, Identity and some other Optimizable classes are defined with a .f() evaluation function, wheras the standard int he rest of simsopt seems to be .J(). Would it be good to re-factor this?]