Merge branch 'main' into add_basic_deterministic_advi

Author: Martin Ingram

notebooks/Exponential Trend Smoothing.ipynb

@@ -15,18 +15,19 @@
     "\n",
     "warnings.filterwarnings(action=\"ignore\", message=\"The RandomType SharedVariables\")\n",
     "\n",
-    "import matplotlib.pyplot as plt\n",
-    "from matplotlib.dates import DayLocator, MonthLocator, YearLocator, DateFormatter\n",
-    "\n",
-    "import pymc_extras.statespace as pmss\n",
-    "import pymc as pm\n",
+    "import os\n",
     "\n",
-    "import preliz as pz\n",
-    "import pandas as pd\n",
     "import arviz as az\n",
-    "import yfinance as yf\n",
+    "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
-    "import os\n",
+    "import pandas as pd\n",
+    "import preliz as pz\n",
+    "import pymc as pm\n",
+    "import yfinance as yf\n",
+    "\n",
+    "from matplotlib.dates import DateFormatter, DayLocator, MonthLocator, YearLocator\n",
+    "\n",
+    "import pymc_extras.statespace as pmss\n",
     "\n",
     "plt.rcParams.update(\n",
     "    {\n",
@@ -61,11 +62,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import ipywidgets as widgets\n",
-    "from IPython.display import display\n",
     "from functools import partial\n",
+    "\n",
+    "import ipywidgets as widgets\n",
     "import pytensor\n",
     "import pytensor.tensor as pt\n",
+    "\n",
+    "from IPython.display import display\n",
+    "\n",
     "from pymc_extras.statespace.filters.distributions import LinearGaussianStateSpace\n",
     "\n",
     "\n",
@@ -344,8 +348,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import datetime\n",
-    "\n",
     "FORCE_UPDATE = False\n",
     "\n",
     "if not os.path.isdir(\"data\"):\n",

notebooks/Making a Custom Statespace Model.ipynb

@@ -15,13 +15,13 @@
     }
    ],
    "source": [
-    "import numpy as np\n",
-    "import matplotlib.pyplot as plt\n",
     "import arviz as az\n",
-    "\n",
-    "from pymc_extras.statespace.core.statespace import PyMCStateSpace\n",
+    "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import pymc as pm\n",
     "import pytensor.tensor as pt\n",
-    "import pymc as pm"
+    "\n",
+    "from pymc_extras.statespace.core.statespace import PyMCStateSpace"
    ]
   },
   {
@@ -45,7 +45,7 @@
     "\n",
     "\n",
     "def print_model_ssm(mod, how=\"eval\"):\n",
-    "    nice_heading = f'{\"name\":<20}{\"__repr__\":<50}{\"shape\":<10}{\"value\":<20}'\n",
+    "    nice_heading = f\"{'name':<20}{'__repr__':<50}{'shape':<10}{'value':<20}\"\n",
     "    print(nice_heading)\n",
     "    print(\"=\" * len(nice_heading))\n",
     "    if how == \"eval\":\n",
@@ -1270,7 +1270,7 @@
    ],
    "source": [
     "az.plot_posterior(\n",
-    "    idata, var_names=[\"ar_params\", \"sigma_x\"], ref_val=true_ar.tolist() + [true_sigma_x]\n",
+    "    idata, var_names=[\"ar_params\", \"sigma_x\"], ref_val=[*true_ar.tolist(), true_sigma_x]\n",
     ");"
    ]
   },
@@ -1333,13 +1333,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from pymc_extras.statespace.models.utilities import make_default_coords\n",
     "from pymc_extras.statespace.utils.constants import (\n",
-    "    ALL_STATE_DIM,\n",
     "    ALL_STATE_AUX_DIM,\n",
-    "    OBS_STATE_DIM,\n",
+    "    ALL_STATE_DIM,\n",
     "    SHOCK_DIM,\n",
     ")\n",
-    "from pymc_extras.statespace.models.utilities import make_default_coords\n",
     "\n",
     "\n",
     "class AutoRegressiveThree(PyMCStateSpace):\n",

notebooks/SARMA Example.ipynb

@@ -31,20 +31,19 @@
     "\n",
     "jax.config.update(\"jax_platform_name\", \"cpu\")\n",
     "\n",
-    "import pymc as pm\n",
-    "from pytensor import tensor as pt\n",
+    "import warnings\n",
     "\n",
     "import arviz as az\n",
-    "import statsmodels.api as sm\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
     "import pandas as pd\n",
-    "from scipy import stats\n",
+    "import pymc as pm\n",
+    "import statsmodels.api as sm\n",
     "\n",
-    "import pymc_extras.statespace as pmss\n",
     "from pymc.model.transform.optimization import freeze_dims_and_data\n",
+    "from pytensor import tensor as pt\n",
     "\n",
-    "import warnings\n",
+    "import pymc_extras.statespace as pmss\n",
     "\n",
     "warnings.filterwarnings(action=\"ignore\", message=\"The RandomType SharedVariables\")\n",
     "\n",
@@ -2697,8 +2696,8 @@
    "source": [
     "fig, ax = plt.subplots()\n",
     "post = az.extract(post_pred).map(np.exp)\n",
-    "hdi = az.hdi(post_pred.map(np.exp))[f\"predicted_posterior_observed\"]\n",
-    "post[f\"predicted_posterior_observed\"].isel(observed_state=0).mean(dim=\"sample\").plot.line(\n",
+    "hdi = az.hdi(post_pred.map(np.exp))[\"predicted_posterior_observed\"]\n",
+    "post[\"predicted_posterior_observed\"].isel(observed_state=0).mean(dim=\"sample\").plot.line(\n",
     "    x=\"time\", ax=ax, add_legend=False, label=\"Posterior Mean, Predicted\"\n",
     ")\n",
     "ax.fill_between(\n",

notebooks/Structural Timeseries Modeling.ipynb

@@ -20,18 +20,20 @@
     }
    ],
    "source": [
-    "from pymc_extras.statespace import structural as st\n",
-    "from pymc_extras.statespace.utils.constants import SHORT_NAME_TO_LONG\n",
+    "import warnings\n",
+    "\n",
+    "import arviz as az\n",
     "import matplotlib.pyplot as plt\n",
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import preliz as pz\n",
     "import pymc as pm\n",
-    "import arviz as az\n",
     "import pytensor.tensor as pt\n",
-    "import preliz as pz\n",
     "\n",
-    "import numpy as np\n",
-    "import pandas as pd\n",
     "from patsy import dmatrix\n",
-    "import warnings\n",
+    "\n",
+    "from pymc_extras.statespace import structural as st\n",
+    "from pymc_extras.statespace.utils.constants import SHORT_NAME_TO_LONG\n",
     "\n",
     "warnings.filterwarnings(\"ignore\", category=UserWarning, message=\"The RandomType SharedVariables\")\n",
     "\n",
@@ -76,13 +78,14 @@
    },
    "outputs": [],
    "source": [
-    "from pymc_extras.statespace.filters.distributions import LinearGaussianStateSpace\n",
     "from pymc.pytensorf import inputvars\n",
     "\n",
+    "from pymc_extras.statespace.filters.distributions import LinearGaussianStateSpace\n",
+    "\n",
     "\n",
     "def make_numpy_function(mod):\n",
     "    mod = mod.build(verbose=False)\n",
-    "    data = pt.matrix(\"data\", shape=(None, 1))\n",
+    "    pt.matrix(\"data\", shape=(None, 1))\n",
     "    steps = pt.iscalar(\"steps\")\n",
     "    x0, _, c, d, T, Z, R, H, Q = mod._unpack_statespace_with_placeholders()\n",
     "    sequence_names = [x.name for x in [c, d] if x.ndim == 2]\n",
@@ -3113,7 +3116,6 @@
    },
    "outputs": [],
    "source": [
-    "from nutpie import transform_adapter\n",
     "import nutpie as ntp\n",
     "\n",
     "# Uncomment to see arguments for with_transform_adapt\n",

notebooks/VARMAX Example.ipynb

@@ -11,21 +11,21 @@
     "\n",
     "jax.config.update(\"jax_platform_name\", \"cpu\")\n",
     "\n",
+    "import sys\n",
+    "\n",
+    "import arviz as az\n",
+    "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
-    "import statsmodels.api as sm\n",
     "import pandas as pd\n",
-    "\n",
     "import pymc as pm\n",
     "import pytensor.tensor as pt\n",
-    "import arviz as az\n",
-    "\n",
-    "import matplotlib.pyplot as plt\n",
-    "import sys\n",
+    "import statsmodels.api as sm\n",
     "\n",
     "sys.path.append(\"..\")\n",
-    "import pymc_extras.statespace as pmss\n",
     "import re\n",
     "\n",
+    "import pymc_extras.statespace as pmss\n",
+    "\n",
     "config = {\n",
     "    \"figure.figsize\": [12.0, 4.0],\n",
     "    \"figure.dpi\": 72.0 * 2,\n",
@@ -841,7 +841,7 @@
     "    new_labels = []\n",
     "    for label in axis.yaxis.get_majorticklabels():\n",
     "        old_text = \"[\" + label.get_text().split(\"[\")[-1]\n",
-    "        labels = eval(re.sub(\"([\\d\\w]+)\", '\"\\g<1>\"', old_text))\n",
+    "        labels = eval(re.sub(r\"([\\d\\w]+)\", r'\"\\g<1>\"', old_text))\n",
     "        lag, other_var = labels\n",
     "        new_text = f\"L{lag}.{other_var}\"\n",
     "        new_labels.append(new_text)\n",

notebooks/discrete_markov_chain.ipynb

@@ -27,13 +27,13 @@
    "outputs": [],
    "source": [
     "import arviz as az\n",
+    "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
+    "import pandas as pd\n",
     "import pymc as pm\n",
     "import pytensor\n",
     "import pytensor.tensor as pt\n",
-    "import pandas as pd\n",
     "import statsmodels.api as sm\n",
-    "import matplotlib.pyplot as plt\n",
     "\n",
     "from matplotlib import ticker as mtick\n",
     "\n",
@@ -525,7 +525,7 @@
     "    \"dates\": dta_hamilton.index,\n",
     "    \"obs_dates\": dta_hamilton.index[order:],\n",
     "    \"states\": [\"State_1\", \"State_2\"],\n",
-    "    \"ar_params\": [f\"L{i+1}.phi\" for i in range(order)],\n",
+    "    \"ar_params\": [f\"L{i + 1}.phi\" for i in range(order)],\n",
     "}\n",
     "\n",
     "with pm.Model(coords=coords) as hmm:\n",

notebooks/latent_approx.ipynb

@@ -25,13 +25,9 @@
     }
    ],
    "source": [
-    "import arviz as az\n",
     "import matplotlib.pyplot as plt\n",
     "import numpy as np\n",
-    "import pymc as pm\n",
-    "from aesara import tensor as at\n",
-    "\n",
-    "import pymc_extras as pmx"
+    "import pymc as pm"
    ]
   },
   {
@@ -42,7 +38,7 @@
    "outputs": [],
    "source": [
     "from pymc_extras.gp import HSGP, KarhunenLoeveExpansion, ProjectedProcess\n",
-    "from pymc_extras.gp.latent_approx import ExpQuad, Matern12, Matern32, Matern52"
+    "from pymc_extras.gp.latent_approx import ExpQuad"
    ]
   },
   {

notebooks/marginalized_changepoint_model.ipynb

@@ -6,11 +6,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import pymc as pm\n",
-    "from pymc_extras.model.marginal.marginal_model import MarginalModel\n",
-    "import pandas as pd\n",
+    "import arviz as az\n",
     "import numpy as np\n",
-    "import arviz as az"
+    "import pandas as pd\n",
+    "import pymc as pm\n",
+    "\n",
+    "from pymc_extras.model.marginal.marginal_model import MarginalModel"
    ]
   },
   {

pymc_extras/distributions/__init__.py

@@ -29,14 +29,14 @@
 from pymc_extras.distributions.transforms import PartialOrder
 
 __all__ = [
+    "R2D2M2CP",
+    "BetaNegativeBinomial",
     "Chi",
-    "Maxwell",
     "DiscreteMarkovChain",
-    "GeneralizedPoisson",
-    "BetaNegativeBinomial",
     "GenExtreme",
-    "R2D2M2CP",
+    "GeneralizedPoisson",
+    "Maxwell",
+    "PartialOrder",
     "Skellam",
     "histogram_approximation",
-    "PartialOrder",
 ]

pymc_extras/distributions/histogram_utils.py

@@ -18,7 +18,7 @@
 
 from numpy.typing import ArrayLike
 
-__all__ = ["quantile_histogram", "discrete_histogram", "histogram_approximation"]
+__all__ = ["discrete_histogram", "histogram_approximation", "quantile_histogram"]
 
 
 def quantile_histogram(