Economics

def get_constraints(m: AbstractModel) -> Sequence[GeneralConstraint]:
    """
    # Economic module and production function

    ## Gross and net GDP

    The core of the model is the economic module, detailing how GDP, investments and
    consumptions vary over time. We use a traditional Cobb-Douglas production function. This means
    that the gross GDP is calculated by:

    :::mimosa.common.economics.calc_GDP

    The net GDP is then calculated by subtracting the damages and
    mitigation costs from the gross GDP. *(Note that in MIMOSA, the damages are expressed as a fraction of the gross GDP,
    whereas the mitigation costs are expressed in absolute terms.)*

    $$ \\text{GDP}_{\\text{net},t,r} = \\text{GDP}_{\\text{gross},t,r} \\cdot (1 - \\text{damage costs}_{t,r}) - \\text{mitigation costs}_{t,r}$$

    ## Investments and consumption

    This net GDP is then split in a part of investments ($I_t$) and a part of consumption ($C_t$), according to a fixed savings rate ($\\text{sr}$):

    $$ I_{t,r} = \\text{sr} \\cdot \\text{GDP}_{\\text{net},t,r}, $$

    $$ C_{t,r} = (1 - \\text{sr}) \\cdot \\text{GDP}_{\\text{net},t,r}. $$

    ## Capital stock

    The capital stock $K_t$ grows over time according to the investments and the depreciation of the capital stock:

    $$ K_{t,r} = K_{t-1,r} + \\Delta t \\cdot \\frac{\\partial K_{t,r}}{\\partial t}, $$

    with the change in capital stock calculated by:

    :::mimosa.common.economics.calc_dKdt

    Since this only gives the change in capital stock, we need to add the initial capital stock to get the actual capital stock.
    This is calculated as a region-dependent multiple of the initial GDP:

    $$ K_{t=0,r} = \\text{init_capitalstock_factor}_r \\cdot \\text{GDP}_{t=0,r}. $$

    The initial capital stock factor is a calibration factor to obtain the initial capital stock. TODO: Source (IMF)
    ``` plotly
    {"file_path": "./assets/plots/init_capital_factor.json"}
    ```

    ## Parameters defined in this module
    - init_capitalstock_factor
    - alpha: Output elasticity of capital. Type: float. Default: 0.3. Min: 0. Max: 1.
    - dk: Yearly depreciation rate of capital stock. Type: float. Default: 0.05. Min: 0. Max: inf.
    - sr: Fraction of GDP used for investments. Type: float. Default: 0.21. Min: 0. Max: 1.
    - ignore_damages: Flag to not take into account the damages in the GDP (but damages are calculated). Type: bool. Default: False.



    """
    constraints = []

    m.init_capitalstock_factor = Param(
        m.regions,
        units=quant.unit("currency_unit"),
        doc="regional::economics.init_capital_factor",
    )
    m.capital_stock = Var(
        m.t,
        m.regions,
        initialize=lambda m, t, r: m.init_capitalstock_factor[r] * m.GDP(m.year(t), r),
        units=quant.unit("currency_unit"),
    )

    # Parameters
    m.alpha = Param(doc="::economics.GDP.alpha")
    m.dk = Param(doc="::economics.GDP.depreciation of capital")
    m.sr = Param(doc="::economics.GDP.savings rate")

    m.GDP_gross = Var(
        m.t,
        m.regions,
        initialize=lambda m, t, r: m.GDP(m.year(0), r),
        units=quant.unit("currency_unit"),
    )
    m.GDP_net = Var(
        m.t,
        m.regions,
        units=quant.unit("currency_unit"),
        initialize=lambda m, t, r: m.GDP(m.year(0), r),
    )
    m.investments = Var(m.t, m.regions, units=quant.unit("currency_unit"))
    m.consumption = Var(m.t, m.regions, units=quant.unit("currency_unit"))

    m.ignore_damages = Param(doc="::economics.damages.ignore damages")

    # Cobb-Douglas, GDP, investments, capital and consumption
    constraints.extend(
        [
            RegionalConstraint(
                lambda m, t, r: (
                    m.GDP_gross[t, r]
                    == economics.calc_GDP(
                        m.TFP(m.year(t), r),
                        m.L(m.year(t), r),
                        soft_min(m.capital_stock[t, r], scale=10),
                        m.alpha,
                    )
                    if t > 0
                    else Constraint.Skip
                ),
                "GDP_gross",
            ),
            RegionalInitConstraint(
                lambda m, r: m.GDP_gross[0, r] == m.GDP(m.year(0), r), "GDP_gross_init"
            ),
            RegionalConstraint(
                lambda m, t, r: m.GDP_net[t, r]
                == m.GDP_gross[t, r]
                * (1 - (m.damage_costs[t, r] if not value(m.ignore_damages) else 0))
                - m.mitigation_costs[t, r]
                - m.financial_transfer[t, r],
                "GDP_net",
            ),
            RegionalConstraint(
                lambda m, t, r: m.investments[t, r] == m.sr * m.GDP_net[t, r],
                "investments",
            ),
            RegionalConstraint(
                lambda m, t, r: m.consumption[t, r] == (1 - m.sr) * m.GDP_net[t, r],
                "consumption",
            ),
            RegionalConstraint(
                lambda m, t, r: (
                    (
                        m.capital_stock[t, r]
                        == m.capital_stock[t - 1, r]
                        + m.dt
                        * economics.calc_dKdt(
                            m.capital_stock[t, r], m.dk, m.investments[t, r], m.dt
                        )
                    )
                    if t > 0
                    else Constraint.Skip
                ),
                "capital_stock",
            ),
            RegionalInitConstraint(
                lambda m, r: m.capital_stock[0, r]
                == m.init_capitalstock_factor[r] * m.GDP(m.year(0), r)
            ),
        ]
    )

    return constraints