Add docs and examples (#18)

* Add v0 docs

* Add v0 docs

* Add scripts to generate notebook html

* Working towards passing tests

* All tests pass

* Reformat docs

* Never divide by zero when computing iterations per second

.github/workflows/release.yml

@@ -0,0 +1,47 @@
+name: release-pipeline
+
+on:
+  release:
+    types:
+      - created
+
+jobs:
+  release-job:
+    runs-on: ubuntu-latest
+    env:
+      PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
+    steps:
+    - uses: actions/checkout@692973e3d937129bcbf40652eb9f2f61becf3332
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.11
+    - name: Install release tooling
+      run: |
+        pip install twine wheel
+    - name: Build package
+      run: |
+        python setup.py sdist bdist_wheel
+    - name: Check version number match
+      run: |
+        echo "GITHUB_REF: ${GITHUB_REF}"
+        # The GITHUB_REF should be something like "refs/tags/v1.2.3"
+        # Make sure the package version is the same as the tag
+        grep -Rq "^Version: ${GITHUB_REF:11}$" cge_modeling.egg-info/PKG-INFO
+    - name: Publish to PyPI
+      run: |
+        twine check dist/*
+        twine upload --repository pypi --username __token__ --password ${PYPI_TOKEN} dist/*
+  test-install-job:
+    needs: release-job
+    runs-on: ubuntu-latest
+    steps:
+    - name: Set up Python
+      uses: actions/setup-python@v5
+      with:
+        python-version: 3.7
+    - name: Give PyPI a chance to update the index
+      run: sleep 240
+    - name: Install from PyPI
+      run: |
+        pip install cge_modeling==${GITHUB_REF:11}

.github/workflows/rtd-link-preview.yml

@@ -0,0 +1,16 @@
+name: Read the Docs Pull Request Preview
+on:
+  pull_request_target:
+    types:
+      - opened
+
+permissions:
+  pull-requests: write
+
+jobs:
+  documentation-links:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: readthedocs/actions/preview@v1
+        with:
+          project-slug: "cge_modeling"

.gitignore

@@ -170,3 +170,7 @@ cge_modeling/base/__pycache__/
 cge_modeling/pytensorf/__pycache__/
 cge_modeling/tools/__pycache__/
 tests/utilities/__pycache__/
+
+# Locally built docs
+docs/build/
+docs/jupyter_execute/

.readthedocs.yaml

@@ -0,0 +1,17 @@
+version: 2
+
+sphinx:
+    configuration: docs/source/conf.py
+
+python:
+   install:
+   - method: pip
+     path: .
+
+conda:
+  environment: "conda_envs/environment_docs.yml"
+
+build:
+  os: "ubuntu-22.04"
+  tools:
+    python: "mambaforge-4.10"

REQUIREMENTS.txt

@@ -2,10 +2,9 @@ pytest
 setuptools
 numba
 numpy
-sympy
+sympy<1.13
 scipy
 pandas
-ipython
 joblib
 pytensor
 latextable

cge_modeling/__init__.py

@@ -2,6 +2,10 @@
 from cge_modeling.base.primitives import Equation, Parameter, Variable
 from cge_modeling.plotting import plot_kateplot, plot_lines
 from cge_modeling.tools.output_tools import display_info_as_table, latex_print_equations
+from cge_modeling.pytensorf.rewrites import prod_to_no_zero_prod  # noqa: F401
+
+
+__version__ = "0.0.1"
 
 __all__ = [
     "CGEModel",

cge_modeling/base/cge.py

@@ -267,7 +267,6 @@ def __init__(
             self._compile(
                 backend=backend,
                 mode=mode,
-                inverse_method=inverse_method,
                 functions_to_compile=compile,
                 use_scan_euler=use_scan_euler,
             )
@@ -738,15 +737,12 @@ def _compile_pytensor(
         self,
         mode,
         functions_to_compile: list[CompiledFunctions],
-        inverse_method: Literal["solve", "pinv", "svd"] = "solve",
         sparse: bool = False,
         use_scan_euler: bool = True,
     ):
         _log.info("Compiling model equations into pytensor graph")
-        (variables, parameters), (system, jac, jac_inv, B) = (
-            compile_cge_model_to_pytensor(
-                self, inverse_method=inverse_method, sparse=sparse
-            )
+        (variables, parameters), (system, jac, B) = compile_cge_model_to_pytensor(
+            self, sparse=sparse
         )
         inputs = variables + parameters
         text_mode = "C" if mode is None else mode
@@ -807,6 +803,7 @@ def f_jac(*args, **kwargs):
                     pt.concatenate([pt.atleast_1d(eq).ravel() for eq in grad]),
                     self.n_variables,
                 )
+
                 # View grad as the function, compute jvp instead  (hessp)
                 # _log.info("Computing SSE Jacobian")
                 # hessp, p = make_jacobian(grad, variables, return_jvp=True)
@@ -863,7 +860,7 @@ def __pytensor_euler_helper(*args, n_steps=100, **kwargs):
                         )
                         self.__last_n_steps = n_steps
                         self.__compiled_f_euler = compile_euler_approximation_function(
-                            jac_inv,
+                            jac,
                             B,
                             variables,
                             parameters,
@@ -881,10 +878,9 @@ def __pytensor_euler_helper(*args, n_steps=100, **kwargs):
                     f_step = jax_euler_step(system, variables, parameters)
                 else:
                     inputs, outputs = pytensor_euler_step(
-                        system, jac_inv, B, variables, parameters
+                        system, jac, B, variables, parameters
                     )
                     # inputs = get_required_inputs(outputs)
-                    print(inputs, outputs)
                     f_step = pytensor.function(inputs, outputs, mode=mode)
                     f_step.trust_inputs = True
                     self.f_step = f_step
@@ -928,7 +924,10 @@ def f_euler(
                             n_steps=n_steps,
                         )
 
-                        current_variable_vals, current_parameter_vals = current_step
+                        current_variable_vals = current_step[: len(self.variable_names)]
+                        current_parameter_vals = current_step[
+                            len(self.variable_names) :
+                        ]
 
                         # current_variable_vals = flat_current_step[
                         #     : len(self.variable_names)
@@ -953,7 +952,8 @@ def f_euler(
                         elapsed = end - start
                         unit = "s/it"
                         if elapsed < 1:
-                            elapsed = 1 / elapsed
+                            # if elapsed is exactly zero just put a big number
+                            elapsed = 1 / elapsed if elapsed != 0 else 1000000
                             unit = "it/s"
 
                         progress.comment = desc.format(elapsed, unit, rmse)
@@ -980,7 +980,6 @@ def _compile(
         self,
         backend: Literal["pytensor", "numba"] | None = "pytensor",
         mode=None,
-        inverse_method: Literal["solve", "pinv", "svd"] = "solve",
         functions_to_compile: list[CompiledFunctions] = "all",
         use_scan_euler: bool = True,
     ):
@@ -993,9 +992,6 @@ def _compile(
             The backend to compile to. One of 'pytensor' or 'numba'.
         mode: str
             Pytensor compile mode. Ignored if mode is not 'pytensor'.
-        inverse_method: str
-            The method to use to compute the inverse of the Jacobian. One of "solve", "pinv", or "svd".
-            Defaults to "solve". Ignored if mode is not 'pytensor'
         functions_to_compile: list of str
             A list of functions to compile. Valid choices are 'root', 'minimum', 'euler', or 'all'. Defaults to 'all'.
         use_scan_euler: bool
@@ -1035,7 +1031,6 @@ def _compile(
         elif backend == "pytensor":
             self._compile_pytensor(
                 mode=mode,
-                inverse_method=inverse_method,
                 functions_to_compile=functions_to_compile,
                 use_scan_euler=use_scan_euler,
                 sparse=self.sparse,

cge_modeling/plotting.py

@@ -3,7 +3,6 @@
 import arviz as az
 import matplotlib.pyplot as plt
 import numpy as np
-import pandas as pd
 from matplotlib.colors import Colormap
 from matplotlib.gridspec import GridSpec
 from matplotlib.ticker import PercentFormatter
@@ -80,7 +79,7 @@ def plot_lines(
         The number of columns in the grid of plots.
     var_names: list of str, optional
         Name of the variables to plot. If None, all variables will be plotted.
-    initial_values: dict[str, np.ndarray], optional
+    initial_values: dict[str, np.array], optional
         The initial values of the variables in the model; those passed to the simulate method. If None, the initial
         values will be taken from the InferenceData object.
     plot_euler: bool, default True
@@ -169,12 +168,13 @@ def f_cmap(*args):
 
             t0 = np.zeros_like(initial_value)
             T = np.ones_like(final_value)
+
             if plot_euler:
                 T = T * idata["euler"].variables.step.max().item()
 
             axis.scatter(
                 t0,
-                final_value,
+                initial_value,
                 marker="*",
                 color="tab:green",
                 zorder=10,
@@ -198,7 +198,7 @@ def f_cmap(*args):
 
 def plot_kateplot(
     idata: az.InferenceData,
-    initial_values: dict[str, np.ndarray],
+    initial_values: dict[str, np.array],
     mod: CGEModel,
     var_names: str | list[str],
     shock_name: str | None = None,
@@ -213,7 +213,7 @@ def plot_kateplot(
     ----------
     idata: az.InferenceData
         The InferenceData object returned by the model's simulate method.
-    initial_values: dict[str, np.ndarray]
+    initial_values: dict[str, np.array]
         The initial values of the variables in the model; those passed to the simulate method.
     mod: CGEModel
         The model object.
@@ -347,17 +347,21 @@ def _plot_one_bar(
     ] = "pct_change",
     legend=True,
     threshhold=1e-6,
+    labelsize=12,
+    xlabel_rotation=0,
 ):
     try:
         data = data.to_dataframe().droplevel(level=drop_vars, axis=0)
         initial_data = initial_data.to_dataframe().droplevel(level=drop_vars, axis=0)
     except ValueError:
-        data = pd.DataFrame([data.to_dict()]).loc[:, ["data", "name"]].set_index("name")
+        data = data.to_dataarray().to_dataframe(name="data")
+        data.index.name = "name"
         initial_data = (
-            pd.DataFrame([initial_data.to_dict()])
-            .loc[:, ["data", "name"]]
-            .set_index("name")
+            initial_data.to_dataarray()
+            .to_dataframe(name="data")
+            .droplevel(level=drop_vars, axis=0)
         )
+
     if "step" in data.columns:
         data = data.drop(columns=["step"])
     if "step" in initial_data.columns:
@@ -374,11 +378,9 @@ def _plot_one_bar(
     to_plot = _compute_bar_data(data, initial_data, metric, threshhold)
 
     if orientation == "v":
-        # to_plot.plot.bar(ax=ax, legend='', facecolor='none', edgecolor="black")
         to_plot.plot.bar(ax=ax, legend=legend)
         ax.axhline(0, color="black", lw=2.0)
     else:
-        # to_plot.plot.barh(ax=ax, legend='', facecolor='none', edgecolor="black")
         to_plot.plot.barh(ax=ax, legend=legend)
         ax.axvline(0, color="black", lw=2.0)
 
@@ -388,7 +390,8 @@ def _plot_one_bar(
             ax.yaxis.set_major_formatter(ticker)
         elif orientation == "h":
             ax.xaxis.set_major_formatter(ticker)
-    ax.tick_params(which="both", labelsize=6)
+    ax.tick_params(which="both", labelsize=labelsize)
+    ax.tick_params(which="major", axis="x", rotation=xlabel_rotation)
     return ax
 
 
@@ -406,6 +409,9 @@ def plot_bar(
         "pct_change", "change", "abs_change", "final", "initial", "both"
     ] = "pct_change",
     threshhold=1e-6,
+    labelsize=12,
+    xlabel_rotation=0,
+    legend=True,
     **figure_kwargs,
 ):
     data = None
@@ -471,7 +477,10 @@ def plot_bar(
                 drop_vars=drop_vars,
                 orientation=orientation,
                 metric=metric,
+                legend=legend,
                 threshhold=threshhold,
+                labelsize=labelsize,
+                xlabel_rotation=xlabel_rotation,
             )
             axis.set(title=var)
             if metric == "final_initial":
@@ -497,7 +506,6 @@ def plot_bar(
             raise ValueError(msg)
         fig, ax = plt.subplots(dpi=dpi, figsize=figsize, **figure_kwargs)
 
-        # labels = [label for label in [make_labels(var_name, mod) for var_name in var_names]]
         _ = _plot_one_bar(
             data,
             initial_values,
@@ -506,6 +514,9 @@ def plot_bar(
             drop_vars=drop_vars,
             orientation=orientation,
             metric=metric,
+            legend=legend,
+            labelsize=labelsize,
+            xlabel_rotation=xlabel_rotation,
         )
 
     return fig