diff --git a/flodym/__init__.py b/flodym/__init__.py index 4016652..f346c16 100644 --- a/flodym/__init__.py +++ b/flodym/__init__.py @@ -20,6 +20,7 @@ DynamicStockModel as DynamicStockModel, InflowDrivenDSM as InflowDrivenDSM, StockDrivenDSM as StockDrivenDSM, + FlexibleDSM as FlexibleDSM, ) from flodym.lifetime_models import ( LifetimeModel as LifetimeModel, diff --git a/flodym/flodym_arrays.py b/flodym/flodym_arrays.py index c324310..0cac1d6 100644 --- a/flodym/flodym_arrays.py +++ b/flodym/flodym_arrays.py @@ -8,8 +8,13 @@ from collections import defaultdict import numpy as np import pandas as pd -from pydantic import BaseModel as PydanticBaseModel, ConfigDict, model_validator -from typing import Optional, Union, Callable +from pydantic import ( + BaseModel as PydanticBaseModel, + ConfigDict, + model_validator, + ModelWrapValidatorHandler, +) +from typing import Optional, Union, Callable, Any, Self from copy import copy from .processes import Process @@ -61,6 +66,8 @@ class FlodymArray(PydanticBaseModel): """Values of the FlodymArray. Must have the same shape as the dimensions of the FlodymArray. If None, an array of zeros is created.""" name: Optional[str] = "unnamed" """Name of the FlodymArray.""" + _is_set: bool = False + """Flag indicating whether the flow has been set or not.""" @model_validator(mode="after") def validate_values(self): @@ -84,6 +91,14 @@ def _check_value_format(self): f"Values shape: {self.values.shape}\n" ) + @model_validator(mode="wrap") + @classmethod + def mark_set_or_unset(cls, data: Any, handler: ModelWrapValidatorHandler[Self]) -> Self: + obj = handler(data) + if isinstance(data, dict): + obj._is_set = "values" in data + return obj + @classmethod def from_dims_superset( cls, dims_superset: DimensionSet, dim_letters: tuple = None, **kwargs @@ -171,6 +186,7 @@ def set_values(self, values: np.ndarray): self._check_value_format() else: self.values[...] = values + self._is_set = True def sum_values(self): """Return the sum of all values in the FlodymArray.""" @@ -447,6 +463,7 @@ def __setitem__(self, keys, item): if isinstance(item, FlodymArray): slice_obj = self._sub_array_handler(keys) self.values[slice_obj.ids] = item.sum_values_to(slice_obj.dim_letters) + self._is_set = True else: self.set_values(copy(item)) return @@ -566,6 +583,19 @@ def items_where(self, condition: Callable) -> np.array: ] return np.array(items).transpose() + @property + def is_set(self) -> bool: + """A boolean to indicate whether the flow's values are known or not.""" + return self._is_set + + def mark_set(self): + """Mark the flow as having values.""" + self._is_set = True + + def mark_unset(self): + """Mark the flow as not having values""" + self._is_set = False + def __str__(self): base = f"{self.__class__.__name__} '{self.name}'" dims = f" with dims ({','.join(self.dims.letters)}) and shape {self.shape};" diff --git a/flodym/mfa_system.py b/flodym/mfa_system.py index 12fe0e4..3d4c2c5 100644 --- a/flodym/mfa_system.py +++ b/flodym/mfa_system.py @@ -302,6 +302,15 @@ def check_flows( if all_good: logging.info(f"Success - No negative flows or NaN values in {self.__class__.__name__}") + def mark_all_unset(self): + """Mark all parameters as unset.""" + for flow in self.flows.values(): + flow.mark_unset() + for stock in self.stocks.values(): + stock.inflow.mark_unset() + stock.outflow.mark_unset() + stock.stock.mark_unset() + @staticmethod def _error_or_warning(message: str, raise_error: bool) -> bool: if raise_error: diff --git a/flodym/stocks.py b/flodym/stocks.py index 290c3c3..a1a06b3 100644 --- a/flodym/stocks.py +++ b/flodym/stocks.py @@ -4,17 +4,29 @@ from abc import abstractmethod import numpy as np -from scipy.linalg import solve_triangular -from pydantic import BaseModel as PydanticBaseModel, ConfigDict, model_validator +from pydantic import BaseModel as PydanticBaseModel, ConfigDict, model_validator, Field from typing import Optional, Union import logging from .processes import Process from .flodym_arrays import StockArray, FlodymArray -from .dimensions import DimensionSet +from .dimensions import Dimension, DimensionSet from .lifetime_models import LifetimeModel, UnevenTimeDim +def stock_compute_decorator(func): + """Adds checks before and after every stock compute routine""" + + def wrapper(self: "Stock", *args, **kwargs): + self._check_needed_arrays() + func(self, *args, **kwargs) + self.mark_computed() + + wrapper.is_decorated = True + + return wrapper + + class Stock(PydanticBaseModel): """Stock objects are components of an MFASystem, where materials can accumulate over time. They consist of three :py:class:`flodym.FlodymArray` objects: @@ -73,20 +85,34 @@ def validate_time_first_dim(self): ) return self + @model_validator(mode="after") + def check_compute_decorator(self): + if not getattr(self.compute, "is_decorated", False): + raise RuntimeError( + "Stock.compute method must have the stock_compute_decorator applied to it." + ) + return self + @model_validator(mode="after") def init_t(self): self._t = UnevenTimeDim(dim=self.dims[self.time_letter]) return self @abstractmethod + @stock_compute_decorator def compute(self): - # always add this check first - self._check_needed_arrays() + # Add stock_compute_decorator to all subclasses! + pass @abstractmethod def _check_needed_arrays(self): pass + def mark_computed(self): + self.inflow.mark_set() + self.outflow.mark_set() + self.stock.mark_set() + @property def shape(self) -> tuple: """Shape of the stock, inflow, outflow arrays, defined by the dimensions.""" @@ -121,6 +147,13 @@ def get_stock_balance(self) -> np.ndarray: ) # stock_change(t) = stock(t) - stock(t-1) return self.inflow.values - self.outflow.values - dsdt + @property + def is_computed(self) -> bool: + """Check whether the stock has been computed, i.e. whether the stock, inflow, and outflow arrays + are all set. + """ + return self.inflow.is_set and self.outflow.is_set and self.stock.is_set + def _to_whole_period(self, annual_flow: np.ndarray) -> np.ndarray: """multiply annual flow by interval length to get flow over whole period.""" return np.einsum("t...,t->t...", annual_flow, self._t.interval_lengths) @@ -139,14 +172,13 @@ class SimpleFlowDrivenStock(Stock): """Given inflows and outflows, the stock can be calculated without a lifetime model or cohorts.""" def _check_needed_arrays(self): - if ( - np.max(np.abs(self.inflow.values)) < 1e-10 - and np.max(np.abs(self.outflow.values)) < 1e-10 - ): - logging.warning("Inflow and Outflow are zero. This will lead to a zero stock.") + if not self.inflow.is_set and not self.outflow.is_set: + logging.warning( + "Neither inflow and outflow are set (is_set=False). If this is intended, perform mark_set() on one of them." + ) + @stock_compute_decorator def compute(self): - self._check_needed_arrays() annual_net_inflow = self.inflow.values - self.outflow.values net_inflow_whole_period = self._to_whole_period(annual_net_inflow) self.stock.values[...] = np.cumsum(net_inflow_whole_period, axis=0) @@ -162,8 +194,31 @@ class DynamicStockModel(Stock): Can be input either as a LifetimeModel subclass, or as an instance of a LifetimeModel subclass. For available subclasses, see `flodym.lifetime_models`. """ + cohort_dim: Optional[Dimension] = None + _outflow_by_cohort: np.ndarray = None _stock_by_cohort: np.ndarray = None + _dims_cohort: DimensionSet = None + _initial_stock_dsm: "InflowByCohortDrivenDSM" = None + _initial_stock_year: int = None + + @model_validator(mode="after") + def validate_cohort_dim(self): + if self.cohort_dim is not None: + t = self.dims[self.time_letter] + c = self.cohort_dim + if c.letter == t.letter or c.name == t.name: + raise ValueError( + "Cohort dimension letter and name must be different from time dimension letter and name." + ) + if c.items != t.items: + raise ValueError("Cohort dimension size must be the same as time dimension size.") + if c.letter in self.dims.letters or c.name in self.dims.names: + raise ValueError("Cohort dimension must not be part of the stock dimensions.") + self._dims_cohort = DimensionSet( + dim_list=[t, c] + list(self.dims.drop(self.time_letter, inplace=False).dim_list) + ) + return self @model_validator(mode="after") def init_cohort_arrays(self): @@ -184,6 +239,12 @@ def init_lifetime_model(self): def _check_needed_arrays(self): self.lifetime_model._check_prms_set() + def _check_cohort_dim(self, application: str): + if self.cohort_dim is None: + raise ValueError( + f"Cohort dimension must be provided at DSM initialization for {application} to work." + ) + @property def _n_t(self) -> int: return list(self.shape)[0] @@ -200,13 +261,22 @@ def _shape_no_t(self) -> tuple: def _t_diag_indices(self) -> tuple: return np.diag_indices(self._n_t) + (slice(None),) * len(self._shape_no_t) - def get_outflow_by_cohort(self) -> np.ndarray: + def get_outflow_by_cohort(self) -> StockArray: """Outflow by cohort, i.e. the outflow of each production year at each time step.""" - return self._outflow_by_cohort + return self._get_by_cohort_array(self._outflow_by_cohort, "outflow_by_cohort") - def get_stock_by_cohort(self) -> np.ndarray: + def get_stock_by_cohort(self) -> StockArray: """Stock by cohort, i.e. the stock of each production year at each time step.""" - return self._stock_by_cohort + return self._get_by_cohort_array(self._stock_by_cohort, "stock_by_cohort") + + def _get_by_cohort_array(self, values: np.ndarray, name: str) -> StockArray: + if self.cohort_dim is None: # leave in for backwards compatibility + logging.warning( + f"Cohort dimension is not defined; {name} cannot be retrieved as FlodymArray. Returning raw ndarray instead." + ) + return values + else: + return StockArray(dims=self._dims_cohort, values=values, name=f"{self.name}_{name}") def _compute_outflow(self): self._outflow_by_cohort = np.einsum( @@ -214,6 +284,83 @@ def _compute_outflow(self): ) self.outflow.values[...] = self._outflow_by_cohort.sum(axis=1) + def set_initial_stock(self, initial_stock: FlodymArray, initial_year: int): + """Set an initial stock at a given year. + The initial stock is added to the stock by cohort internally. + + Args: + initial_stock (FlodymArray): Initial stock to be set. Must have same dimensions as the stock, except for time dimension. + initial_year (int): Year in which the initial stock is given. Must be an item of the time dimension. + """ + self._check_cohort_dim("set_initial_stock") + if initial_stock.dims.letters != self._dims_cohort.drop(self.time_letter).letters: + raise ValueError( + f"Initial stock dimensions {initial_stock.dims.letters} do not match expected dims {self._dims_cohort.drop(self.time_letter).letters}." + ) + if initial_year not in self.dims[self.time_letter].items: + raise ValueError( + f"Initial year {initial_year} is not in time dimension items {self.dims[self.time_letter].items}." + ) + + inflow_by_cohort = StockArray(dims=self._dims_cohort, name=f"{self.name}_inflow_by_cohort") + # Manually set the slice for the initial_year time index + initial_year_idx = self.dims[self.time_letter].items.index(initial_year) + inflow_by_cohort.values[initial_year_idx, :, ...] = self._to_annual(initial_stock.values) + if np.any((inflow_by_cohort.values > 0) & (self.lifetime_model.sf == 0)): + raise ValueError( + "Initial stock is inconsistent with the lifetime model: survival function is zero for some cohort/inflow year combinations." + ) + # Back-calculate the inflow from the initial stock by dividing by survival function at initial_year + # Then multiply by survival function at all times to get inflow_by_cohort[t, c] + inflow_reconstructed = inflow_by_cohort.values[initial_year_idx, :, ...] / np.where( + self.lifetime_model.sf[initial_year_idx, :, ...] == 0, + 1, + self.lifetime_model.sf[initial_year_idx, :, ...], + ) + # Set inflow_by_cohort for all time steps + for t_idx in range(self._n_t): + inflow_by_cohort.values[t_idx, :, ...] = ( + inflow_reconstructed * self.lifetime_model.sf[t_idx, :, ...] + ) + inflow_by_cohort.mark_set() # Mark as set before passing to InflowByCohortDrivenDSM + self._initial_stock_dsm = InflowByCohortDrivenDSM( + dims=self.dims, + cohort_dim=self.cohort_dim, + lifetime_model=self.lifetime_model, + inflow_by_cohort=inflow_by_cohort, + name=f"{self.name}_ISD", + ) + self._initial_stock_year = initial_year + + @property + def _initial_year_index(self): + return self.dims[self.time_letter].items.index(self._initial_stock_year) + + def _add_initial_stock_to_inflow(self): + if self._initial_stock_dsm is not None: + self._initial_stock_dsm.compute(stop_after="inflow") + if np.any(self.inflow.values[: self._initial_year_index + 1, ...] > 0): + raise ValueError( + f"Prescribed inflow before or in the initial stock year {self._initial_stock_year} is non-zero." + ) + self.inflow.values[...] += self._initial_stock_dsm.inflow.values + + def _add_initial_stock_to_stock(self) -> np.ndarray: + if self._initial_stock_dsm is not None: + if np.any(self.stock.values[: self._initial_year_index + 1, ...] > 0): + raise ValueError( + f"Prescribed stock before or in the initial stock year {self._initial_stock_year} is non-zero." + ) + # Initial stock only gives the remaining(!) stock by cohort in one year; + # From this, we first need to reconstruct what the total stock was over historic time + # This needs two steps: + # 1) Construct what inflow in a year would lead to the given initial stock cohort + # surviving + # 2) Sum this up over time in an inflow-driven way to get total stock over historic time + self._initial_stock_dsm.compute(stop_after="stock") + iiy = self._initial_year_index + self.stock.values[: iiy + 1, ...] = self._initial_stock_dsm.stock.values[: iiy + 1, ...] + def __str__(self): base = super().__str__() lifetime_model = self.lifetime_model.__class__.__name__ @@ -226,13 +373,16 @@ class InflowDrivenDSM(DynamicStockModel): """ def _check_needed_arrays(self): - super()._check_needed_arrays() - if np.allclose(self.inflow.values, np.zeros(self.shape)): - logging.warning("Inflow is zero. This will lead to a zero stock and outflow.") + DynamicStockModel._check_needed_arrays(self) + if not self.inflow.is_set: + logging.warning( + "Inflow is not set (is_set=False). If this is intended, perform mark_set() on it." + ) + @stock_compute_decorator def compute(self): """Determine stocks and outflows and store values in the class instance.""" - self._check_needed_arrays() + self._add_initial_stock_to_inflow() self._compute_stock() self._compute_outflow() @@ -252,27 +402,16 @@ class StockDrivenDSM(DynamicStockModel): where A is the survival function matrix, x is the inflow vector, and b is the stock vector. """ - solver: str = "manual" - """Algorithm to use for solving the equation system. Options are: "manual" (default), which uses - an own python implementation, and "lapack", which calls the lapack trtrs routine via scipy. - The lapack implementation may be more precise. Speed depends on the dimensionality, - but the manual implementation is usually faster. - """ - - @model_validator(mode="after") - def init_solver(self): - if self.solver not in ["manual", "lapack"]: - raise ValueError("Solver must be either 'manual' or 'lapack'.") - return self - def _check_needed_arrays(self): - super()._check_needed_arrays() - if np.allclose(self.stock.values, np.zeros(self.shape)): - logging.warning("Stock is zero. This will lead to a zero inflow and outflow.") + DynamicStockModel._check_needed_arrays(self) + if not self.stock.is_set: + logging.warning( + "Stock is not set (is_set=False). If this is intended, perform mark_set() on it." + ) + @stock_compute_decorator def compute(self): """Determine inflows and outflows and store values in the class instance.""" - self._check_needed_arrays() self._compute_cohorts_and_inflow() self._compute_outflow() @@ -282,22 +421,6 @@ def _compute_cohorts_and_inflow(self): This involves solving the lower triangular equation system A*x=b, where A is the survival function matrix, x is the inflow vector, and b is the stock vector. """ - if self.solver == "manual": - self._compute_inflow_manual() - elif self.solver == "lapack": - self._compute_inflow_lapack() - else: - raise ValueError(f"Unknown engine: {self.solver}") - - self._stock_by_cohort = np.einsum( - "c...,tc...->tc...", self.inflow.values, self.lifetime_model.sf - ) - - def _compute_inflow_manual(self) -> tuple[np.ndarray]: - """With given total stock and lifetime distribution, - the method builds the stock by cohort and the inflow, - using a manual algorithm for solving of the equation system (see "solver" doc for details). - """ # Maths behind implementation: # Solve square linear equation system # sf * inflow = stock @@ -307,6 +430,9 @@ def _compute_inflow_manual(self) -> tuple[np.ndarray]: # solve for inflow_i: # inflow_i = ( stock_i - sum_{j=1...i-1}(sf_i,j * inflow_j) ) / sf_ii inflow_whole_period = np.zeros_like(self.inflow.values) + + self._add_initial_stock_to_stock() + for i in range(self._n_t): stock_i = self.stock.values[i, ...] sf_ij = self.lifetime_model.sf[i, :i, ...] @@ -315,17 +441,74 @@ def _compute_inflow_manual(self) -> tuple[np.ndarray]: inflow_whole_period[i, ...] = (stock_i - (sf_ij * inflow_j).sum(axis=0)) / sf_ii self.inflow.values[...] = self._to_annual(inflow_whole_period) + self._stock_by_cohort = np.einsum( + "c...,tc...->tc...", inflow_whole_period, self.lifetime_model.sf + ) - def _compute_inflow_lapack(self) -> tuple[np.ndarray]: - """With given total stock and lifetime distribution, - the method builds the stock by cohort and the inflow, - using lapack for solving of the equation system (see "engine" doc for details). - """ - sf = self.lifetime_model.sf - slt = (slice(None),) - inflow_whole_period = np.zeros_like(self.inflow.values) - for i in np.ndindex(self._shape_no_t): - inflow_whole_period[slt + i] = solve_triangular( - sf[2 * slt + i], self.stock.values[slt + i], lower=True + +class InflowByCohortDrivenDSM(InflowDrivenDSM): + + inflow_by_cohort: StockArray = None + cohort_dim: Dimension # require for this subclass + + @model_validator(mode="after") + def init_cohort_arrays(self): + if self.inflow_by_cohort is None: + self.inflow_by_cohort = StockArray( + dims=self._dims_cohort, name=f"{self.name}_inflow_by_cohort" ) - self.inflow.values[...] = self._to_annual(inflow_whole_period) + else: + if self.inflow_by_cohort.dims.letters != self._dims_cohort.letters: + raise ValueError( + f"Inflow by cohort dimensions {self.inflow_by_cohort.dims.letters} do not match expected dims {self._dims_cohort.letters}." + ) + return self + + def _check_needed_arrays(self): + DynamicStockModel._check_needed_arrays(self) + if not self.inflow_by_cohort.is_set: + logging.warning( + "Inflow_by_cohort is not set (is_set=False). If this is intended, perform mark_set() on it." + ) + if self.inflow.is_set: + logging.warning( + "Inflow is not set (is_set=True). It will be overwritten. Use either inflow or inflow_by_cohort." + ) + + def _compute_inflow(self): + # Extract diagonal where t==c (inflow occurs in cohort year) + self.inflow.values[...] = self.inflow_by_cohort.values[self._t_diag_indices] + + @stock_compute_decorator + def compute(self, stop_after: str = None): + """Determine stocks and outflows and store values in the class instance.""" + self._compute_inflow() + if stop_after == "inflow": + return + self._compute_stock() + if stop_after == "stock": + return + self._compute_outflow() + + +class FlexibleDSM(DynamicStockModel): + """Computes either stock-driven or inflow-driven dynamic stock model, depending on which of the + stock or inflow is set. + """ + + compute_stock_driven = StockDrivenDSM.compute + compute_inflow_driven = InflowDrivenDSM.compute + + _compute_cohorts_and_inflow = StockDrivenDSM._compute_cohorts_and_inflow + _compute_stock = InflowDrivenDSM._compute_stock + + def compute(self): + if self.stock.is_set: + self.compute_stock_driven() + elif self.inflow.is_set: + self.compute_inflow_driven() + else: + raise ValueError("Either stock or inflow must be set for FlexibleDSM.compute().") + + # replaces decorator, since inner functions are already decorated + compute.is_decorated = True diff --git a/tests/test_stocks.py b/tests/test_stocks.py index 997a7cb..385f756 100644 --- a/tests/test_stocks.py +++ b/tests/test_stocks.py @@ -4,7 +4,7 @@ from flodym.dimensions import Dimension, DimensionSet from flodym.flodym_arrays import StockArray -from flodym.stocks import InflowDrivenDSM, StockDrivenDSM +from flodym.stocks import InflowDrivenDSM, StockDrivenDSM, InflowByCohortDrivenDSM from flodym.lifetime_models import LogNormalLifetime dim_list = [ @@ -166,3 +166,292 @@ def test_unequal_time_steps(): stocks_select = stocks[1] values_select = stocks_select.values[-10:] assert np.max(np.abs(values_all - values_select)) < 0.01 + + +def test_inflow_by_cohort_driven_dsm(): + """Test InflowByCohortDrivenDSM basic functionality.""" + # Create simple dimensions for testing + time_items = list(range(2000, 2011)) + dims_test = DimensionSet( + dim_list=[ + Dimension(name="time", letter="t", items=time_items, dtype=int), + Dimension(name="product", letter="p", items=["A", "B"], dtype=str), + ] + ) + + # Create cohort dimension + cohort_dim = Dimension(name="cohort", letter="c", items=time_items, dtype=int) + + # Create lifetime model + lifetime_model = LogNormalLifetime(dims=dims_test, time_letter="t", mean=5, std=2) + + # Create InflowByCohortDrivenDSM + dsm = InflowByCohortDrivenDSM( + dims=dims_test, + cohort_dim=cohort_dim, + lifetime_model=lifetime_model, + time_letter="t", + name="test_dsm", + ) + + # Set inflow by cohort - diagonal pattern (inflow only occurs in cohort year) + dsm.inflow_by_cohort.values[...] = 0.0 + for i in range(len(time_items)): + dsm.inflow_by_cohort.values[i, i, :] = 1.0 # 1 unit per product per year + + # Compute the DSM + dsm.compute() + + # Check that inflow was computed correctly (diagonal extraction) + assert dsm.inflow.is_set + assert dsm.stock.is_set + assert dsm.outflow.is_set + + # Verify inflow values are approximately 1.0 at each time step + assert np.allclose(dsm.inflow.values, 1.0) + + # Verify stock balance + dsm.check_stock_balance() + + +def test_inflow_driven_dsm_with_initial_stock(): + """Test InflowDrivenDSM with initial stock set via set_initial_stock.""" + # Create dimensions + time_items = list(range(2000, 2021)) + dims_test = DimensionSet( + dim_list=[ + Dimension(name="time", letter="t", items=time_items, dtype=int), + Dimension(name="product", letter="p", items=["A", "B"], dtype=str), + ] + ) + + # Create cohort dimension (required for initial stock functionality) + cohort_dim = Dimension(name="cohort", letter="c", items=time_items, dtype=int) + + # Create lifetime model + lifetime_model = LogNormalLifetime(dims=dims_test, time_letter="t", mean=8, std=3) + + # Create InflowDrivenDSM with cohort_dim + dsm = InflowDrivenDSM( + dims=dims_test, + cohort_dim=cohort_dim, + lifetime_model=lifetime_model, + time_letter="t", + name="test_dsm_initial", + ) + + # Set initial stock at year 2005 + initial_year = 2005 + # Initial stock dims should be (cohort, product) = _dims_cohort.drop(time_letter) + initial_stock_dims = DimensionSet( + dim_list=[cohort_dim] + list(dims_test.drop("t", inplace=False).dim_list) + ) + initial_stock = StockArray(dims=initial_stock_dims, name="initial_stock") + initial_stock.values[...] = 0.0 + # Set stock from various cohorts + for i in range(5): # 5 historical cohorts + initial_stock.values[i, :] = 10.0 - i # Decreasing stock from older cohorts + + dsm.set_initial_stock(initial_stock, initial_year) + + # Set inflow for years AFTER initial year only + # Inflow before and at initial year will be computed from initial stock + dsm.inflow.values[...] = 0.0 + initial_year_idx = time_items.index(initial_year) + dsm.inflow.values[initial_year_idx + 1 :, :] = 2.0 # 2 units per year after initial + dsm.inflow.mark_set() # Mark inflow as set + + # Compute + dsm.compute() + + # Verify that stock is set + assert dsm.stock.is_set + assert dsm.inflow.is_set + assert dsm.outflow.is_set + + # Stock at initial year should be non-zero + assert dsm.stock.values[initial_year_idx, :].sum() > 0 + + # Stock before initial year should also be non-zero (reconstructed from initial stock) + assert dsm.stock.values[initial_year_idx - 1, :].sum() > 0 + + # Verify stock balance + dsm.check_stock_balance() + + +def test_stock_driven_dsm_with_initial_stock(): + """Test StockDrivenDSM with initial stock set via set_initial_stock.""" + # Create dimensions + time_items = list(range(2000, 2021)) + dims_test = DimensionSet( + dim_list=[ + Dimension(name="time", letter="t", items=time_items, dtype=int), + Dimension(name="product", letter="p", items=["A", "B"], dtype=str), + ] + ) + + # Create cohort dimension (required for initial stock functionality) + cohort_dim = Dimension(name="cohort", letter="c", items=time_items, dtype=int) + + # Create lifetime model + lifetime_model = LogNormalLifetime(dims=dims_test, time_letter="t", mean=8, std=3) + + # Create StockDrivenDSM with cohort_dim + dsm = StockDrivenDSM( + dims=dims_test, + cohort_dim=cohort_dim, + lifetime_model=lifetime_model, + time_letter="t", + name="test_stock_driven_initial", + ) + + # Set initial stock at year 2005 + initial_year = 2005 + # Initial stock dims should be (cohort, product) = _dims_cohort.drop(time_letter) + initial_stock_dims = DimensionSet( + dim_list=[cohort_dim] + list(dims_test.drop("t", inplace=False).dim_list) + ) + initial_stock = StockArray(dims=initial_stock_dims, name="initial_stock") + initial_stock.values[...] = 0.0 + # Set stock from various cohorts + for i in range(5): # 5 historical cohorts + initial_stock.values[i, :] = 10.0 - i # Decreasing stock from older cohorts + + dsm.set_initial_stock(initial_stock, initial_year) + + # Set total stock for all years AFTER initial year only + dsm.stock.values[...] = 0.0 + initial_year_idx = time_items.index(initial_year) + # Stock before and at initial year will be computed from initial stock + # We only set stock after the initial year + for i in range(initial_year_idx + 1, len(time_items)): + dsm.stock.values[i, :] = ( + 50.0 + (i - initial_year_idx) * 2 + ) # Growing stock after initial year + dsm.stock.mark_set() # Mark stock as set + + # Compute + dsm.compute() + + # Verify that inflow is set + assert dsm.inflow.is_set + assert dsm.stock.is_set + assert dsm.outflow.is_set + + # Inflow should be non-zero + assert dsm.inflow.values.sum() > 0 + + # Verify stock balance + dsm.check_stock_balance() + + +def test_get_stock_by_cohort(): + """Test get_stock_by_cohort method.""" + # Create dimensions + time_items = list(range(2000, 2011)) + dims_test = DimensionSet( + dim_list=[ + Dimension(name="time", letter="t", items=time_items, dtype=int), + Dimension(name="product", letter="p", items=["A", "B"], dtype=str), + ] + ) + + # Create cohort dimension + cohort_dim = Dimension(name="cohort", letter="c", items=time_items, dtype=int) + + # Create lifetime model + lifetime_model = LogNormalLifetime(dims=dims_test, time_letter="t", mean=5, std=2) + + # Create InflowDrivenDSM with cohort dimension + dsm = InflowDrivenDSM( + dims=dims_test, + cohort_dim=cohort_dim, + lifetime_model=lifetime_model, + time_letter="t", + name="test_cohort", + ) + + # Set inflow + dsm.inflow.values[...] = 1.0 + + # Compute + dsm.compute() + + # Get stock by cohort + stock_by_cohort = dsm.get_stock_by_cohort() + + # Verify it's a StockArray + assert isinstance(stock_by_cohort, StockArray) + + # Verify dimensions include both time and cohort + assert "t" in stock_by_cohort.dims.letters + assert "c" in stock_by_cohort.dims.letters + + # Verify shape is correct (n_t, n_t, ...) + assert stock_by_cohort.values.shape[0] == len(time_items) + assert stock_by_cohort.values.shape[1] == len(time_items) + + # Verify that stock by cohort sums to total stock along cohort dimension + total_stock_from_cohort = stock_by_cohort.values.sum(axis=1) + assert np.allclose(total_stock_from_cohort, dsm.stock.values) + + # Verify that future cohorts (t < c) have zero stock + for t_idx in range(len(time_items)): + for c_idx in range(t_idx + 1, len(time_items)): + assert np.allclose(stock_by_cohort.values[t_idx, c_idx, :], 0.0) + + +def test_get_outflow_by_cohort(): + """Test get_outflow_by_cohort method.""" + # Create dimensions + time_items = list(range(2000, 2011)) + dims_test = DimensionSet( + dim_list=[ + Dimension(name="time", letter="t", items=time_items, dtype=int), + Dimension(name="product", letter="p", items=["A", "B"], dtype=str), + ] + ) + + # Create cohort dimension + cohort_dim = Dimension(name="cohort", letter="c", items=time_items, dtype=int) + + # Create lifetime model + lifetime_model = LogNormalLifetime(dims=dims_test, time_letter="t", mean=5, std=2) + + # Create InflowDrivenDSM with cohort dimension + dsm = InflowDrivenDSM( + dims=dims_test, + cohort_dim=cohort_dim, + lifetime_model=lifetime_model, + time_letter="t", + name="test_cohort", + ) + + # Set inflow + dsm.inflow.values[...] = 1.0 + + # Compute + dsm.compute() + + # Get outflow by cohort + outflow_by_cohort = dsm.get_outflow_by_cohort() + + # Verify it's a StockArray + assert isinstance(outflow_by_cohort, StockArray) + + # Verify dimensions include both time and cohort + assert "t" in outflow_by_cohort.dims.letters + assert "c" in outflow_by_cohort.dims.letters + + # Verify shape is correct (n_t, n_t, ...) + assert outflow_by_cohort.values.shape[0] == len(time_items) + assert outflow_by_cohort.values.shape[1] == len(time_items) + + # Verify that outflow by cohort sums to total outflow along cohort dimension + total_outflow_from_cohort = outflow_by_cohort.values.sum(axis=1) + assert np.allclose(total_outflow_from_cohort, dsm.outflow.values) + + # Verify that future cohorts (t < c) have zero outflow + for t_idx in range(len(time_items)): + for c_idx in range(t_idx + 1, len(time_items)): + assert np.allclose(outflow_by_cohort.values[t_idx, c_idx, :], 0.0)