import numpy as np
import pyomo.core as pyo


def portfolio_optimization(
    covariances: np.ndarray, returns: np.ndarray, budget: int
) -> pyo.ConcreteModel:
    model = pyo.ConcreteModel()
    num_assets = len(returns)
    model.x = pyo.Var(range(num_assets), domain=pyo.Binary)
    x_array = model.x.values()

    model.budget = pyo.Constraint(expr=(sum(x_array) == budget))

    risk: float = x_array @ covariances @ x_array
    profit: float = returns @ x_array
    model.risk, model.profit = risk, profit

    model.cost = pyo.Objective(expr=model.risk - model.profit, sense=pyo.minimize)

    return model

This function generates a PYOMO model which formulates the Portfolio Optimization. It is the process of selecting the best asset distribution according such that it minimizes the difference between the risk and the expected return.

The model consists of:

• Binary variable declaration for each asset (model.x) indicating whether the asset is chosen to the portfolio
• Constraint – ensure that the portfolio consists of specific amount of assets.
• Objective rule – minimizes the difference between the risk and the expected returns.