diff --git a/model_solver/model_solver.py b/model_solver/model_solver.py index 3263879..27a0fcf 100644 --- a/model_solver/model_solver.py +++ b/model_solver/model_solver.py @@ -15,36 +15,53 @@ class ModelSolver: """ - EXAMPLE OF USE USE: + ModelSolver is designed to handle and solve mathematical models represented by a system of equations. + It supports various mathematical functions such as min, max, log, and exp. + This class allows you to initialize a model with a list of equations and endogenous variables. + It subsequently solves the model using input data stored in a Pandas DataFrame. - Let "equations" and "endogenous" be lists containing equations and endogenous variables, respectively, stored as strings, e.g. + Usage Example: + + Let `equations` and `endogenous` be lists containing equations and endogenous variables, respectively, stored as strings, e.g., equations = [ - 'x+y=A', - 'x/y=B' + 'x + y = A', + 'x / y = B' ] endogenous = [ 'x', 'y' ] - where 'A' and 'B' are exogenous variables - The solver supports the mathematical functions min, max, log and exp + where 'A' and 'B' are exogenous variables. - A class instance called "model" is initialized by + To initialize a ModelSolver instance, use: model = ModelSolver(equations, endogenous) - This reads in the equations and endogenous variables and perform block analysis and ordering and generates simulation code - Upon completion, the model is ready to be solved subject to data (exogenous and initial values of endogenous variables) in a Pandas DataFrame - Let "input_df" be a dataframe containing data on A and B and initial values for x and y. Then the model can be solved by invoking + This reads in the equations and endogenous variables, performs block analysis and ordering, and generates simulation code. + + To solve the model using input data in a Pandas DataFrame, let's assume you have a DataFrame named "input_df" containing data on 'A' and 'B' as well as initial values for 'x' and 'y'. You can solve the model by invoking: solution_df = model.solve_model(input_df) - Now "solution_df" is a Pandas DataFrame with exactly the same dimensions as "input_df", but where the endogenous variables are replaced by the solutions to the model - The last solution is also stored in "model.last_solution" + Now, "solution_df" is a Pandas DataFrame with the same dimensions as "input_df," but with the endogenous variables replaced by the solutions to the model. The last solution is also stored in "model.last_solution." + + Attributes + ---------- + last_solution : pandas DataFrame + The last solved solution. + + Methods + ------- + solve_model(input_df) + Solves the model based on input data in a Pandas DataFrame. + Returns a DataFrame with the same dimensions as input_df. + + Analysis Methods + --------------- + (TBA - To Be Added) - ModelSolver also has a number of methods for analysis (TBA) """ # Reads in equations and endogenous variables and does a number of operations, e.g. analyzing block structure using graph theory. @@ -429,9 +446,7 @@ def _gen_def_or_obj_fun_and_jac(eqns: tuple[str], def switch_endo_var(self, old_endo_vars: list[str], new_endo_vars: list[str]): - """ - Sets old_endo_vars as exogenous and new_endo_vars as endogenous and performs block analysis - """ + if all([x in self.endo_vars for x in old_endo_vars]) is False: print('All variables in old_endo_vars are not endogenous') @@ -464,7 +479,34 @@ def find_endo_var(self, endo_var: str): def describe(self): """ - Describes model, that is number of equations, number of simultaneous blocks and how many equations are in each block + Sets old_endo_vars as exogenous and new_endo_vars as endogenous and performs block analysis. + + Parameters: + ---------- + old_endo_vars : list of str + List of old endogenous variables to be switched to exogenous. + + new_endo_vars : list of str + List of new endogenous variables to be switched from exogenous. + + Returns: + ------- + None + + Notes: + ------ + This function switches the endogenous and exogenous status of variables and performs block analysis on the model. + + Raises: + ------ + ValueError: + If any variable in `old_endo_vars` is not in the current list of endogenous variables. + If any variable in `new_endo_vars` is already in the list of endogenous variables. + + Example: + -------- + >>> model = YourModelClass() + >>> model.switch_endo_var(['var1', 'var2'], ['var3', 'var4']) """ print('-'*100) @@ -479,7 +521,65 @@ def describe(self): def show_blocks(self): """ - Prints endogenous and exogenous variables and equations for every block in the model + Prints endogenous and exogenous variables and equations for every block in the model. + + Iterates through all blocks in the model and calls the `show_block` function to display their details. + + Returns: + ------- + None + + Example: + -------- + >>> model = YourModelClass() + >>> model.show_blocks() + + -------------------------------------------------- + Block 1 + -------------------------------------------------- + Endogenous Variables: + - var1 + - var2 + + Exogenous Variables: + - exog_var1 + - exog_var2 + + Equations: + - eqn1: var1 = exog_var1 + exog_var2 + - eqn2: var2 = var1 + exog_var2 + + -------------------------------------------------- + Block 2 + -------------------------------------------------- + Endogenous Variables: + - var3 + - var4 + + Exogenous Variables: + - exog_var3 + - exog_var4 + + Equations: + - eqn3: var3 = exog_var3 + exog_var4 + - eqn4: var4 = var3 + exog_var4 + + ... + + -------------------------------------------------- + Block n + -------------------------------------------------- + Endogenous Variables: + - var_n1 + - var_n2 + + Exogenous Variables: + - exog_var_n1 + - exog_var_n2 + + Equations: + - eqn_n1: var_n1 = exog_var_n1 + exog_var_n2 + - eqn_n2: var_n2 = var_n1 + exog_var_n2 """ for key, _ in self._blocks.items(): @@ -489,7 +589,41 @@ def show_blocks(self): def show_block(self, i: int): """ - Prints endogenous and exogenous variables and equations for a given block + Prints endogenous and exogenous variables and equations for a given block. + + Parameters: + ----------- + i : int + The index of the block to display. + + Returns: + -------- + None + + Example: + -------- + >>> model = YourModelClass() + >>> model.show_block(1) + + Block consists of an equation or a system of equations + + 5 endogenous variables: + - var1 + - var2 + - var3 + - var4 + - var5 + + 3 predetermined variables: + - pred_var1 + - pred_var2 + - pred_var3 + + 4 equations: + - eqn1: var1 = pred_var1 + pred_var2 + - eqn2: var2 = var1 + pred_var2 + - eqn3: var3 = pred_var2 + pred_var3 + - eqn4: var4 = var3 + pred_var1 """ block = self._blocks.get(i) @@ -507,7 +641,32 @@ def show_block(self, i: int): def solve_model(self, input_df: pd.DataFrame, jit=True) -> pd.DataFrame: """ - Solves the model for a given DataFrame + Solves the model for a given DataFrame. + + Parameters: + ----------- + input_df : pd.DataFrame + A DataFrame containing input data for the model. + + jit : bool, optional + Flag indicating whether to use just-in-time (JIT) compilation for solving equations. + Default is True. + + Returns: + -------- + pd.DataFrame + A DataFrame containing the model's output data. + + Raises: + ------- + TypeError: + If any column in `input_df` is not of numeric data type. + + Example: + -------- + >>> model = YourModelClass() + >>> input_data = pd.DataFrame({'var1': [1.0, 2.0, 3.0], 'var2': [0.5, 1.0, 1.5]}) + >>> output_data = model.solve_model(input_data) """ if self._some_error: @@ -678,7 +837,35 @@ def draw_blockwise_graph( figsize=(7.5, 7.5) ): """ - Draws a directed graph of block in which variable is along with max number of ancestors and descendants. + Draws a directed graph of a block containing the given variable with a limited number of ancestors and descendants. + + Parameters: + ----------- + var : str + The variable for which the blockwise graph will be drawn. + + max_ancs_gens : int, optional + Maximum number of generations of ancestors to include in the graph. Default is 5. + + max_desc_gens : int, optional + Maximum number of generations of descendants to include in the graph. Default is 5. + + max_nodes : int, optional + Maximum number of nodes to include in the graph. If the graph has more nodes, it won't be plotted. Default is 50. + + figsize : tuple, optional + A tuple specifying the width and height of the figure for the graph. Default is (7.5, 7.5). + + Returns: + -------- + None + + Example: + -------- + >>> model = YourModelClass() + >>> model.draw_blockwise_graph('var1', max_ancs_gens=3, max_desc_gens=2, max_nodes=30, figsize=(10, 10)) + + Draws a directed graph of the block containing 'var1' with up to 3 generations of ancestors and 2 generations of descendants. """ if self._some_error: @@ -765,7 +952,26 @@ def _find_var_node(self, var): def trace_to_exog_vars(self, block: str): """ - Prints all exogenous variables that are ancestors to block + Prints all exogenous variables that are ancestors to the given block. + + Parameters: + ----------- + block : str + The block for which exogenous variables will be traced. + + Returns: + -------- + None + + Example: + -------- + >>> model = YourModelClass() + >>> model.trace_to_exog_vars('Block1') + + exog_var1 + exog_var2 + exog_var3 + ... """ if self._some_error: @@ -789,8 +995,32 @@ def _trace_to_exog_vars(self, block): def trace_to_exog_vals(self, block: int, period_index: int): """ - Traces block back to exogenous values + Traces the given block back to exogenous values and prints those values. + + Parameters: + ----------- + block : int + The block to be traced back to exogenous values. + + period_index : int + The index of the period for which exogenous values will be traced. + + Returns: + -------- + None + + Example: + -------- + >>> model = YourModelClass() + >>> model.trace_to_exog_vals(1, 3) + + Block 1 traces back to the following exogenous variable values in 2023-01-04: + exog_var1=12.5 + exog_var2=8.2 + exog_var3=10.0 + ... """ + try: output_array = self._last_solution.to_numpy(dtype=np.float64, copy=True) var_col_index = {var: i for i, var in enumerate(self._last_solution.columns.str.lower().to_list())} @@ -810,8 +1040,36 @@ def trace_to_exog_vals(self, block: int, period_index: int): def show_block_vals(self, i: int, period_index: int): """ - TBA + Prints the values of endogenous and predetermined variables in a given block for a specific period. + + Parameters: + ----------- + i : int + The index of the block for which variable values will be displayed. + + period_index : int + The index of the period for which variable values will be shown. + + Returns: + -------- + None + + Example: + -------- + >>> model = YourModelClass() + >>> model.show_block_vals(1, 3) + + Block 1 has endogenous variables in 2023-01-04 that evaluate to: + var1=10.5 + var2=15.2 + ... + + Block 1 has predetermined variables in 2023-01-04 that evaluate to: + pred_var1=8.1 + pred_var2=9.7 + ... """ + try: output_array = self._last_solution.to_numpy(dtype=np.float64, copy=True) var_col_index = {var: i for i, var in enumerate(self._last_solution.columns.str.lower().to_list())}