More fine control for simple bounds in Opti stack

[JeyRunner]

More fine control for simple bounds in Opti stack

[JeyRunner]

In case someone wants similar fine control over the variable bounds, I have written a small workaround wrapper class for Opti. For solving with variable bounds, it will create a new nlpsol instance and pass the variable bounds instead of using the nlpsol instance inside the Opti object (its not possible to use the nlpsol instance inside the Opti object from python).

Note that this is just a workaround (using solve_with_simple_var_bounds(...) has some drawbacks, e.g. opti.callback does not work). This should be implemented in c++ in the OptiNode class.

Usage with the workaround wrapper:

# create wrapper for Opti
opti = OptiWithSimpleBounds.create()

var_x = opti.variable(6)
opti.subject_to(var_x[0:3] == var_x[3:]**2)

# variable bounds
opti.subject_to_var_bounds(-casadi.inf, var_x[2:], 1)  # lower bound non-restricted
opti.subject_to_var_bounds(0, var_x[0:2], 2)

# solve with new wrapper function
opti.solve_with_simple_var_bounds('ipopt', opts, ipopt_opts)

# get values
sol_x = opti.value(var_x)

This yields the ipopt output:

Total number of variables............................:        6
                     variables with only lower bounds:        0
                variables with lower and upper bounds:        2
                     variables with only upper bounds:        4

Implementation:

class OptiWithSimpleBounds(OptiAdvanced):

	solved_with_simple_bounds: bool
	solved_with_simple_bounds_solver_result: dict

	# key is the index of the variable
	class VarBounds:
		sym: MX
		lbx: DM
		ubx: DM
	simple_var_bounds_values: dict[VarBounds]


	@staticmethod
	def create_from_opti(opti: Opti) -> 'OptiWithSimpleBounds':
		"""
		Create OptiWithSimpleBounds from an existing opti object.
		Not that this will create a copy of the opti object, thus don't use the original opti object afterwards.
		:return: the new OptiWithSimpleBounds object.
		"""
		advanced = opti.advanced
		# patch object
		advanced.__class__ = OptiWithSimpleBounds
		# add attrs
		setattr(advanced, 'solved_with_simple_bounds', False)
		setattr(advanced, 'solved_with_simple_bounds_solver_result', None)
		setattr(advanced, 'simple_var_bounds_values', {})
		return advanced

	@staticmethod
	def create():
		"""
		Create OptiWithSimpleBounds object to be used inplace of Opti.
		"""
		return OptiWithSimpleBounds.create_from_opti(Opti())

	def __init__(self, *args):
		assert False, ("Never create this object directly, allways use: \n"
					   "OptiWithSimpleBounds.create()")


	def subject_to_var_bounds(self, min: DM, variable: MX, max: DM):
		"""
		Add simple bounds on optimization variables.
		If you do not want to restrict either min or max, just use (-)DM.inf() as value.

		Example usage:
		opti.subject_to_var_bounds(-DM.inf(), var_x[0:2], 42)
		opti.subject_to_var_bounds(0, var_x, 42)
		"""
		# helper
		def get_sym_slice_elements():
			# get symbols
			f = casadi.Function("f", [], [variable], {"allow_free": True})
			assert len(f.free_mx()) == 1, "currently just support using one var at a time"
			sym: MX = f.free_mx()[0]
			assert sym.is_symbolic(), "you need to provided a variable to constrain"

			# Evaluate jacobian of expr wrt symbols (taken from optistack_internal.cpp (set_value_internal))
			Jf = casadi.Function("f", [], [jacobian(variable, sym)], [], ['J'])#, {"compact": True})
			J: DM = Jf()['J']
			sp_JT: Sparsity = J.T.sparsity()
			sliced_elements_of_var = sp_JT.row()
			return sym, sliced_elements_of_var


		var_sym = None
		var_elements_to_constrain = []

		# if is just a var
		if variable.is_symbolic():
			var_sym = variable
			var_elements_to_constrain = np.arange(0, var_sym.rows())
		else:
			# slice of var
			sym, sliced_elements_of_var = get_sym_slice_elements()
			print(f'variable bounds for sym ({variable}) {sym} elements {sliced_elements_of_var}')
			var_sym = sym
			var_elements_to_constrain = sliced_elements_of_var

		var_i = self.__get_var_sym_index(var_sym)
		# print('variable index', var_i)

		# add var bounds
		var_bounds = None
		if var_i not in self.simple_var_bounds_values:
			var_nx = var_sym.rows()
			var_bounds = OptiWithSimpleBounds.VarBounds()
			var_bounds.sym = var_sym
			var_bounds.lbx = -DM.inf(var_nx)
			var_bounds.ubx = DM.inf(var_nx)
			self.simple_var_bounds_values[var_i] = var_bounds
		else:
			var_bounds = self.simple_var_bounds_values[var_i]

		# check that bounds are not constraint already
		if not (casadi.logic_all(var_bounds.lbx[var_elements_to_constrain] == -DM.inf()) and
			casadi.logic_all(var_bounds.ubx[var_elements_to_constrain] == DM.inf())):
			assert False, \
				(f"Variable {var_sym} has already defined bounds (from previous call to subject_to_var_bounds) for range {var_elements_to_constrain}.\n"
				  f"	lbx {var_bounds.lbx}\n"
				  f"	ubx {var_bounds.ubx}\n")

		# set bounds
		var_bounds.lbx[var_elements_to_constrain] = min
		var_bounds.ubx[var_elements_to_constrain] = max
		# print('lbx ', var_bounds.lbx)
		# print('ubx ', var_bounds.ubx)
		# print()

	def __get_var_sym_index(self, var_sym) -> int:
		var_meta: MetaVar = self.get_meta(var_sym)
		var_i = var_meta.i
		return var_i

	def __get_var_lower_and_upper_bounds(self) -> (DM, DM):
		"""
		Get upper and lower bounds for variables x.
		opti.bake() has to be called before.
		:return: lbx, ubx
		"""
		lbx = -DM.inf(self.nx)
		ubx = DM.inf(self.nx)
		print("__get_var_lower_and_upper_bounds:")
		for var in self.simple_var_bounds_values.keys():
			bounds_dict: OptiWithSimpleBounds.VarBounds = self.simple_var_bounds_values[var]
			var_meta: MetaVar = self.get_meta(bounds_dict.sym)
			start = var_meta.start
			stop = var_meta.stop
			lbx[start:stop] = bounds_dict.lbx
			ubx[start:stop] = bounds_dict.ubx
			print(f"> var '{var}' index range  {start} : {stop},  bounds = {bounds_dict}")
		print('lbx', lbx)
		print('ubx', ubx)
		print()
		return lbx, ubx


	def solve_with_simple_var_bounds(self, solver: str, p_opts, solver_opts) -> OptiAdvanced:
		"""
		Solve with previously defined simple_var_bounds.
		Note that this will create a new internal nlpsolver and copy back the results into this opti object.
		Thus, some methods of this opti object, e.g. opti.callback(...), will not work.
		But the resulting values can be normally obtained by using opti.value(...).
		"""
		self.solver(solver, p_opts, solver_opts)

		# get the args how the internal solver would be called:
		advanced: OptiAdvanced = self
		advanced.bake()
		advanced.solve_prepare()
		solver_args = advanced.arg()
		print('orig solver args', solver_args)

		# put the solver options into the options
		p_opts[solver] = solver_opts

		# create new solver
		solver = nlpsol('solver', 'ipopt', {
			'x': self.x,
			'p': self.p,
			'f': self.f,
			'g': self.g,
		}, p_opts)

		# get the lower and upper variable bounds
		lbx, ubx = self.__get_var_lower_and_upper_bounds()

		solved_with_simple_bounds_solver_result = solver(
			**solver_args,
			lbx=lbx,
			ubx=ubx
		)

		# copy back results form solver to x values of this object
		advanced.res(solved_with_simple_bounds_solver_result)

		self.solved_with_simple_bounds = True
		return advanced
``