{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b1a31788",
   "metadata": {},
   "outputs": [],
   "source": [
    "def sample_k(a, probs):\n",
    "    n = a.shape[0] - 1\n",
    "    k = a.shape[1] - 1\n",
    "\n",
    "    j = k\n",
    "    sample = np.empty([n], dtype=int)\n",
    "    for i in np.arange(n, 0, -1):\n",
    "\n",
    "        # Unnormalized probabilities of Xi and -Xi\n",
    "        z = a[i, j]\n",
    "        p = (a[i-1, j-1] + probs[i-1]) - z\n",
    "\n",
    "        # Sample according to normalized dist.\n",
    "        X = np.random.binomial(1, np.exp(p))\n",
    "\n",
    "        # Pick next state based on value of sample\n",
    "        j = j-1 if X > 0 else j\n",
    "\n",
    "        # Concatenate to samples: preallocating and slicing\n",
    "        # should be faster than concatenating/stacking\n",
    "        sample[i-1] = X\n",
    "\n",
    "    # Our samples should always satisfy the constraint\n",
    "    assert(np.sum(sample) == k - 1)\n",
    "    return sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c66890c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "torch.set_default_dtype(torch.float64)\n",
    "\n",
    "from torch import Tensor\n",
    "def log1mexp(x):\n",
    "    return torch.where(x < 0.6931471805599453094, torch.log(-torch.expm1(-x)), torch.log1p(-torch.exp(-x)))\n",
    "\n",
    "def logaddexp(x: Tensor, y: Tensor) -> Tensor:\n",
    "    with torch.no_grad():\n",
    "        m = torch.maximum(x, y)\n",
    "        m = m.masked_fill_(torch.isneginf(m), 0.)\n",
    "\n",
    "    z = (x - m).exp_() + (y - m).exp_()\n",
    "    mask = z == 0\n",
    "    z = z.masked_fill_(mask, 1.).log_().add_(m)\n",
    "    z = z.masked_fill_(mask, -float('inf'))\n",
    "\n",
    "    return z\n",
    "\n",
    "# def pr_exactly_k(probs, k, log_space=True):\n",
    "#     \"\"\"\n",
    "#     probs: a tensor of shape: (batch_size, num_vars)\n",
    "#     where probs[:, i] corresponds to the batch probabilities\n",
    "#     of Bernoulli variable Xi\n",
    "\n",
    "#     if log_space=True, we expect log_probabilities as input\n",
    "#     \"\"\"\n",
    "#     probs = torch.from_numpy(probs)\n",
    "#     probs.requires_grad_()\n",
    "    \n",
    "#     n = probs.shape[0]\n",
    "\n",
    "#     # a[:, i, j] = Pr(Sum(X1, ... , Xi) = j)\n",
    "#     # Note: a[:, 0, 0] corresponds to the Pr\n",
    "#     # that an empty sequence summing up to -1\n",
    "#     # which is always 0\n",
    "#     a = torch.full((n+1, k+2), -torch.inf)\n",
    "\n",
    "#     # The probability of an empty sequence\n",
    "#     # summing to 0 is 1\n",
    "#     a[:, 0, 1] = 0\n",
    "    \n",
    "#     for i in range(1, n+1):\n",
    "#         # To get a sequence of length i with only\n",
    "#         # j true variables, I either take a seq.\n",
    "#         # with k true variables, and set variable\n",
    "#         # i to false, or take a sequence with k-1\n",
    "#         # true variables and set variable i to true\n",
    "#         a[:, i, 1:] = logaddexp(a[:, i-1, :-1] + probs[:, i-1:i],  a[:, i-1, 1:] + log1mexp(-probs[:, i-1:i].detach()))\n",
    "    \n",
    "#     # backward to compute marginals\n",
    "#     a[-1][-1].mean().backward()\n",
    "\n",
    "#     return a, probs.grad\n",
    "\n",
    "\n",
    "def pr_exactly_k(probs, k, log_space=True):\n",
    "    \"\"\"\n",
    "    probs: a tensor of shape: (batch_size, num_vars)\n",
    "    where probs[:, i] corresponds to the batch probabilities\n",
    "    of Bernoulli variable Xi\n",
    "\n",
    "    if log_space=True, we expect log_probabilities as input\n",
    "    \"\"\"\n",
    "    probs = torch.from_numpy(probs)\n",
    "    probs.requires_grad_()\n",
    "    \n",
    "    n = probs.shape[0]\n",
    "\n",
    "    # a[:, i, j] = Pr(Sum(X1, ... , Xi) = j)\n",
    "    # Note: a[:, 0, 0] corresponds to the Pr\n",
    "    # that an empty sequence summing up to -1\n",
    "    # which is always 0\n",
    "    if log_space:\n",
    "        a = torch.full((n+1, k+2), -torch.inf)\n",
    "    else:\n",
    "        a = torch.zeros((n+1, k+2))\n",
    "\n",
    "    # The probability of an empty sequence\n",
    "    # summing to 0 is 1\n",
    "    a[0, 1] = 0 if log_space else 1\n",
    "    for i in range(1, n+1):\n",
    "\n",
    "        # To get a sequence of length i with only\n",
    "        # j true variables, I either take a seq.\n",
    "        # with k true variables, and set variable\n",
    "        # i to false, or take a sequence with k-1\n",
    "        # true variables and set variable i to true\n",
    "        if log_space:\n",
    "            a[i, 1:] = logaddexp(a[i-1, :-1] + probs[i-1:i],  a[i-1, 1:] + log1mexp(-probs[i-1:i].detach()))\n",
    "        else:\n",
    "            a[i, 1:] =  a[i-1, :-1].clone() * probs[i-1] + a[i-1, 1:].clone() * (1 - probs[i-1]).detach()\n",
    "    a[-1][-1].mean().backward()\n",
    "\n",
    "    return a, probs.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "cb5378b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# def log1mexp(x):\n",
    "#     return np.where(x < 0.6931471805599453094, np.log(-np.expm1(-x)), np.log1p(-np.exp(-x)))\n",
    "\n",
    "def logsigmoid(x):\n",
    "    return -np.logaddexp(0, -x)\n",
    "\n",
    "def pr_exactly_k_np(probs, k, log_space=True):\n",
    "    \"\"\"\n",
    "    probs: a tensor of shape: (batch_size, num_vars)\n",
    "    where probs[:, i] corresponds to the batch probabilities\n",
    "    of Bernoulli variable Xi\n",
    "\n",
    "    if log_space=True, we expect log_probabilities as input\n",
    "    \"\"\"\n",
    "    n = probs.shape[0]\n",
    "\n",
    "    # a[:, i, j] = Pr(Sum(X1, ... , Xi) = j)\n",
    "    # Note: a[:, 0, 0] corresponds to the Pr\n",
    "    # that an empty sequence summing up to -1\n",
    "    # which is always 0\n",
    "    if log_space:\n",
    "        a = np.full((n+1, k+2), -np.inf)\n",
    "    else:\n",
    "        a = np.zeros((n+1, k+2))\n",
    "\n",
    "    # The probability of an empty sequence\n",
    "    # summing to 0 is 1\n",
    "    a[0, 1] = 0 if log_space else 1\n",
    "    for i in range(1, n+1):\n",
    "\n",
    "        # To get a sequence of length i with only\n",
    "        # j true variables, I either take a seq.\n",
    "        # with k true variables, and set variable\n",
    "        # i to false, or take a sequence with k-1\n",
    "        # true variables and set variable i to true\n",
    "        if log_space:\n",
    "            a[i, 1:] = np.logaddexp(a[i-1, :-1] + probs[i-1],  a[i-1, 1:] + log1mexp(-probs[i-1]))\n",
    "        else:\n",
    "            a[i, 1:] =  a[i-1, :-1] * probs[i-1] + a[i-1, 1:] * (1 - probs[i-1])\n",
    "\n",
    "    return a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "557595ff",
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "\n",
    "import numpy as np\n",
    "from sklearn.exceptions import ConvergenceWarning\n",
    "from sklearn.linear_model import ridge_regression\n",
    "from sklearn.utils.validation import check_is_fitted\n",
    "\n",
    "from pysindy.optimizers.base import BaseOptimizer\n",
    "\n",
    "\n",
    "class STLSQ_(BaseOptimizer):\n",
    "    \"\"\"Sequentially thresholded least squares algorithm. Defaults to doing\n",
    "    Sequentially thresholded Ridge regression.\n",
    "\n",
    "    Attempts to minimize the objective function\n",
    "    :math:`\\\\|y - Xw\\\\|^2_2 + \\\\alpha \\\\|w\\\\|^2_2`\n",
    "    by iteratively performing least squares and masking out\n",
    "    elements of the weight array w that are below a given threshold.\n",
    "\n",
    "    See the following reference for more details:\n",
    "\n",
    "        Brunton, Steven L., Joshua L. Proctor, and J. Nathan Kutz.\n",
    "        \"Discovering governing equations from data by sparse\n",
    "        identification of nonlinear dynamical systems.\"\n",
    "        Proceedings of the national academy of sciences\n",
    "        113.15 (2016): 3932-3937.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    threshold : float, optional (default 0.1)\n",
    "        Minimum magnitude for a coefficient in the weight vector.\n",
    "        Coefficients with magnitude below the threshold are set\n",
    "        to zero.\n",
    "\n",
    "    alpha : float, optional (default 0.05)\n",
    "        Optional L2 (ridge) regularization on the weight vector.\n",
    "\n",
    "    max_iter : int, optional (default 20)\n",
    "        Maximum iterations of the optimization algorithm.\n",
    "\n",
    "    ridge_kw : dict, optional (default None)\n",
    "        Optional keyword arguments to pass to the ridge regression.\n",
    "\n",
    "    fit_intercept : boolean, optional (default False)\n",
    "        Whether to calculate the intercept for this model. If set to false, no\n",
    "        intercept will be used in calculations.\n",
    "\n",
    "    normalize_columns : boolean, optional (default False)\n",
    "        Normalize the columns of x (the SINDy library terms) before regression\n",
    "        by dividing by the L2-norm. Note that the 'normalize' option in sklearn\n",
    "        is deprecated in sklearn versions >= 1.0 and will be removed.\n",
    "\n",
    "    copy_X : boolean, optional (default True)\n",
    "        If True, X will be copied; else, it may be overwritten.\n",
    "\n",
    "    initial_guess : np.ndarray, shape (n_features) or (n_targets, n_features),\n",
    "            optional (default None)\n",
    "        Initial guess for coefficients ``coef_``.\n",
    "        If None, least-squares is used to obtain an initial guess.\n",
    "\n",
    "    verbose : bool, optional (default False)\n",
    "        If True, prints out the different error terms every iteration.\n",
    "\n",
    "    Attributes\n",
    "    ----------\n",
    "    coef_ : array, shape (n_features,) or (n_targets, n_features)\n",
    "        Weight vector(s).\n",
    "\n",
    "    ind_ : array, shape (n_features,) or (n_targets, n_features)\n",
    "        Array of 0s and 1s indicating which coefficients of the\n",
    "        weight vector have not been masked out, i.e. the support of\n",
    "        ``self.coef_``.\n",
    "\n",
    "    history_ : list\n",
    "        History of ``coef_``. ``history_[k]`` contains the values of\n",
    "        ``coef_`` at iteration k of sequentially thresholded least-squares.\n",
    "\n",
    "    Examples\n",
    "    --------\n",
    "    >>> import numpy as np\n",
    "    >>> from scipy.integrate import odeint\n",
    "    >>> from pysindy import SINDy\n",
    "    >>> from pysindy.optimizers import STLSQ\n",
    "    >>> lorenz = lambda z,t : [10*(z[1] - z[0]),\n",
    "    >>>                        z[0]*(28 - z[2]) - z[1],\n",
    "    >>>                        z[0]*z[1] - 8/3*z[2]]\n",
    "    >>> t = np.arange(0,2,.002)\n",
    "    >>> x = odeint(lorenz, [-8,8,27], t)\n",
    "    >>> opt = STLSQ(threshold=.1, alpha=.5)\n",
    "    >>> model = SINDy(optimizer=opt)\n",
    "    >>> model.fit(x, t=t[1]-t[0])\n",
    "    >>> model.print()\n",
    "    x0' = -9.999 1 + 9.999 x0\n",
    "    x1' = 27.984 1 + -0.996 x0 + -1.000 1 x1\n",
    "    x2' = -2.666 x1 + 1.000 1 x0\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        threshold=0.1,\n",
    "        alpha=0.05,\n",
    "        max_iter=1000,\n",
    "        ridge_kw=None,\n",
    "        normalize_columns=False,\n",
    "        fit_intercept=False,\n",
    "        copy_X=True,\n",
    "        initial_guess=None,\n",
    "        verbose=False,\n",
    "        k=3,\n",
    "        lr=1.0,\n",
    "        _lambda=1.0\n",
    "    ):\n",
    "        super(STLSQ_, self).__init__(\n",
    "            max_iter=max_iter,\n",
    "            fit_intercept=fit_intercept,\n",
    "            copy_X=copy_X,\n",
    "            normalize_columns=normalize_columns,\n",
    "        )\n",
    "\n",
    "        if threshold < 0:\n",
    "            raise ValueError(\"threshold cannot be negative\")\n",
    "        if alpha < 0:\n",
    "            raise ValueError(\"alpha cannot be negative\")\n",
    "\n",
    "        self.threshold = threshold\n",
    "        self.alpha = alpha\n",
    "        self.ridge_kw = ridge_kw\n",
    "        self.initial_guess = initial_guess\n",
    "        self.verbose = verbose\n",
    "        self.k = k\n",
    "        self.lr = lr\n",
    "        self._lambda = _lambda\n",
    "\n",
    "#     def _sparse_coefficients(self, dim, ind, coef, threshold):\n",
    "#         \"\"\"Perform thresholding of the weight vector(s)\"\"\"\n",
    "#         print(\"Sparsifying Coefficients\")\n",
    "        \n",
    "# #         # Sampling exactly k\n",
    "# #         probs = logsigmoid(self.logits)\n",
    "# #         pr_k = pr_exactly_k(probs, self.k)\n",
    "# #         marginals = pr_k[-1][-1].mean().backward()\n",
    "# #         sample = sample_k(pr_k.detach().numpy(), probs) \n",
    "        \n",
    "#         c = np.zeros(dim)\n",
    "#         c[ind] = coef\n",
    "#         big_ind = sample != 0\n",
    "#         c[~big_ind] = 0\n",
    "#         return c, big_ind\n",
    "\n",
    "    def _sparse_coefficients(self, dim, ind, coef, threshold):\n",
    "        \"\"\"Perform thresholding of the weight vector(s)\"\"\"\n",
    "        c = np.zeros(dim)\n",
    "        c[ind] = coef\n",
    "        big_ind = np.abs(c) >= threshold\n",
    "        c[~big_ind] = 0\n",
    "        return c, big_ind\n",
    "\n",
    "    def _regress(self, x, y):\n",
    "        \"\"\"Perform the ridge regression\"\"\"\n",
    "        kw = self.ridge_kw or {}\n",
    "        coef = ridge_regression(x, y, self.alpha, **kw)\n",
    "        self.iters += 1\n",
    "        return coef\n",
    "    \n",
    "#     def _imle(self, x, y):\n",
    "        \n",
    "#         p = logsigmoid(self.logits)\n",
    "#         a_p = pr_exactly_k(p, self.k)\n",
    "#         marginals_p = a_p[-1][-1].backward()\n",
    "#         p_sample = sample_k(a_p.detach().numpy(), probs) \n",
    "        \n",
    "#         # perturb p to get q\n",
    "#         q = self.logits - _lambda*(p_sample*coef*x - y)*(coef*x)\n",
    "        \n",
    "#         # get the gradient\n",
    "#         a_q = pr_exactly_k(q, self.k)\n",
    "#         marginals_q = a_q[-1][-1].backward()\n",
    "#         grad = marginals_p - marginals_q\n",
    "        \n",
    "#         # update\n",
    "#         self.logits = self.logits - lr * grad\n",
    "        \n",
    "    def _imle_fwd(self):\n",
    "        p = logsigmoid(self.logits)\n",
    "#         p = np.broadcast_to(p, (257024, 14))\n",
    "        a_p, self.fwd_marginals = pr_exactly_k(p, self.k)\n",
    "\n",
    "#         self.fwd_marginals = p.grad\n",
    "\n",
    "        self.fwd_sample = sample_k(a_p.detach().numpy(), p)\n",
    "\n",
    "        \n",
    "        \n",
    "    def _imle_bwd(self, x, y, coef):\n",
    "        \n",
    "        # perturb p to get q\n",
    "        n_samples = y.shape[0]\n",
    "        coef_x = np.matmul(x, coef)\n",
    "        diff = coef_x - y.reshape(n_samples)\n",
    "        diff = diff.reshape(-1, 1)\n",
    "        tmp = diff * x\n",
    "        tmp_grad = tmp.mean(axis=0)\n",
    "        perturb = tmp_grad * coef\n",
    "        q = self.logits - self._lambda * 2 * perturb\n",
    "        q = logsigmoid(q)\n",
    "        \n",
    "        # get the gradient\n",
    "        a_q, marginals_q = pr_exactly_k(q, self.k)\n",
    "#         a_q[-1][-1].mean().backward()\n",
    "#         marginals_q = q.grad\n",
    "        grad = self.fwd_marginals - marginals_q\n",
    "        \n",
    "        # update\n",
    "        self.logits = self.logits - self.lr * grad.detach().numpy()\n",
    "        \n",
    "\n",
    "    def _no_change(self):\n",
    "        \"\"\"Check if the coefficient mask has changed after thresholding\"\"\"\n",
    "        this_coef = self.history_[-1].flatten()\n",
    "        if len(self.history_) > 1:\n",
    "            last_coef = self.history_[-2].flatten()\n",
    "        else:\n",
    "            last_coef = np.zeros_like(this_coef)\n",
    "        return all(bool(i) == bool(j) for i, j in zip(this_coef, last_coef))\n",
    "\n",
    "    def _reduce(self, x, y):\n",
    "        \"\"\"Performs at most ``self.max_iter`` iterations of the\n",
    "        sequentially-thresholded least squares algorithm.\n",
    "\n",
    "        Assumes an initial guess for coefficients and support are saved in\n",
    "        ``self.coef_`` and ``self.ind_``.\n",
    "        \"\"\"\n",
    "        if self.initial_guess is not None:\n",
    "            self.coef_ = self.initial_guess\n",
    "\n",
    "        ind = self.ind_\n",
    "        n_samples, n_features = x.shape\n",
    "        n_targets = y.shape[1]\n",
    "        n_features_selected = np.sum(ind)\n",
    "        \n",
    "#         self.logits =np.random.rand(n_features)# np.ones(n_features, dtype=float)\n",
    "        self.logits = np.ones(n_features, dtype=float)\n",
    "\n",
    "        # Print initial values for each term in the optimization\n",
    "        if self.verbose:\n",
    "            row = [\n",
    "                \"Iteration\",\n",
    "                \"|y - Xw|^2\",\n",
    "                \"a * |w|_2\",\n",
    "                \"|w|_0\",\n",
    "                \"Total error: |y - Xw|^2 + a * |w|_2\",\n",
    "            ]\n",
    "            print(\n",
    "                \"{: >10} ... {: >10} ... {: >10} ... {: >10}\"\n",
    "                \" ... {: >10}\".format(*row)\n",
    "            )\n",
    "        \n",
    "        for k in range(self.max_iter):\n",
    "            if np.count_nonzero(ind) == 0:\n",
    "                warnings.warn(\n",
    "                    \"Sparsity parameter is too big ({}) and eliminated all \"\n",
    "                    \"coefficients\".format(self.threshold)\n",
    "                )\n",
    "                coef = np.zeros((n_targets, n_features))\n",
    "                break\n",
    "\n",
    "            coef = np.zeros((n_targets, n_features))\n",
    "            for i in range(n_targets):\n",
    "                if np.count_nonzero(ind[i]) == 0:\n",
    "                    warnings.warn(\n",
    "                        \"Sparsity parameter is too big ({}) and eliminated all \"\n",
    "                        \"coefficients\".format(self.threshold)\n",
    "                    )\n",
    "                    continue\n",
    "                \n",
    "                \n",
    "                # Sample a mask of exactly-k features\n",
    "                self._imle_fwd()\n",
    "                \n",
    "                # Use the selected features to solve the ridge regression\n",
    "                mask = self.fwd_sample != 0\n",
    "                                    \n",
    "                mask = np.broadcast_to(mask, (257024, 14))\n",
    "\n",
    "                coef_i = self._regress(np.reshape(x[:][mask], (-1, self.k)), y[:, i])\n",
    "                \n",
    "                coef[i][mask[0]] = coef_i\n",
    "                ind[i] = mask[0]\n",
    "                \n",
    "                # Update the logits using imle\n",
    "                self._imle_bwd(x, y, coef[i,:]) # coef[i, :].shape should be num_features\n",
    "                                \n",
    "#                 coef_i, ind_i = self._sparse_coefficients(\n",
    "#                     n_features, ind[i], coef_i, self.threshold\n",
    "#                 )\n",
    "#                 coef[i][mask[0]] = coef_i\n",
    "#                 print(\"coef\", coef_i)\n",
    "\n",
    "# #                 coef[i] = coef_i\n",
    "#                 ind[i] = mask[0]\n",
    "\n",
    "            self.history_.append(coef)\n",
    "            if self.verbose:\n",
    "                R2 = np.sum((y - np.dot(x, coef.T)) ** 2)\n",
    "                L2 = self.alpha * np.sum(coef ** 2)\n",
    "                L0 = np.count_nonzero(coef)\n",
    "                row = [k, R2, L2, L0, R2 + L2]\n",
    "                print(\n",
    "                    \"{0:10d} ... {1:10.4e} ... {2:10.4e} ... {3:10d}\"\n",
    "                    \" ... {4:10.4e}\".format(*row)\n",
    "                )\n",
    "#             if np.sum(ind) == n_features_selected or self._no_change():\n",
    "                # could not (further) select important features\n",
    "#                 break\n",
    "        else:\n",
    "            warnings.warn(\n",
    "                \"STLSQ._reduce did not converge after {} iterations.\".format(\n",
    "                    self.max_iter\n",
    "                ),\n",
    "                ConvergenceWarning,\n",
    "            )\n",
    "            try:\n",
    "                coef\n",
    "            except NameError:\n",
    "                coef = self.coef_\n",
    "                warnings.warn(\n",
    "                    \"STLSQ._reduce has no iterations left to determine coef\",\n",
    "                    ConvergenceWarning,\n",
    "                )\n",
    "        self.ind_ = self.logits.argsort()[-self.k:]\n",
    "        mask = np.zeros(n_features, dtype=int)\n",
    "        mask[self.ind_] = 1\n",
    "        mask = np.broadcast_to(mask, (257024, 14))\n",
    "        self.coef_ = self._regress(np.reshape(x[:][mask!=0], (-1, self.k)), y[:, i])\n",
    "        \n",
    "        coef = np.zeros((n_features))\n",
    "        coef[mask[0]!=0] = self.coef_\n",
    "        self.coef_ = coef\n",
    "        \n",
    "#         self.ind_ = ind\n",
    "\n",
    "    @property\n",
    "    def complexity(self):\n",
    "        print(\"Complexity\")\n",
    "        check_is_fitted(self)\n",
    "\n",
    "        return np.count_nonzero(self.coef_) + np.count_nonzero(\n",
    "            [abs(self.intercept_) >= self.threshold]\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a8783ce0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot Kuramoto-Sivashinsky data and its derivative\n",
    "def plot_u_and_u_dot(t, x, u):\n",
    "    dt = t[1] - t[0]\n",
    "    plt.figure(figsize=(10, 4))\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.pcolormesh(t, x, u[:, :, 0])\n",
    "    plt.xlabel(\"t\", fontsize=16)\n",
    "    plt.ylabel(\"x\", fontsize=16)\n",
    "    plt.title(r\"$u(x, t)$\", fontsize=16)\n",
    "    u_dot = ps.FiniteDifference(axis=1)._differentiate(u, t=dt)\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.pcolormesh(t, x, u_dot[:, :, 0])\n",
    "    plt.xlabel(\"t\", fontsize=16)\n",
    "    plt.ylabel(\"x\", fontsize=16)\n",
    "    ax = plt.gca()\n",
    "    ax.set_yticklabels([])\n",
    "    plt.title(r\"$\\dot{u}(x, t)$\", fontsize=16)\n",
    "    plt.show()\n",
    "    return u_dot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "109a4903",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 143360 characters]\n",
      "text/plain": [
       "<Figure size 720x288 with 2 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "from scipy.io import loadmat\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pysindy as ps\n",
    "\n",
    "# Load data from .mat file\n",
    "data = loadmat('kuramoto_sivishinky.mat')\n",
    "t = np.ravel(data['tt'])\n",
    "dt = t[1] - t[0]\n",
    "x = np.ravel(data['x'])\n",
    "u = data['uu']\n",
    "u = u.reshape(len(x), len(t), 1)\n",
    "u_dot = plot_u_and_u_dot(t, x, u)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "635060fe",
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "import itertools\n",
    "import random\n",
    "def run():\n",
    "    lambdas = [1e4, 1e3, 1e2, 1e1, 1e-1, 1e-2, 1e-3, 1e-4]\n",
    "    lrs = [1e3, 1e2, 1e1, 1e-1, 1e-2, 1e-3, 1e-4]\n",
    "    \n",
    "#     lambdas = [1e1]\n",
    "#     lrs = [1e1]\n",
    "    for _lambda, lr in itertools.product(lambdas, lrs): \n",
    "        print(f\"_lambda  = {_lambda}, lr = {lr}\")\n",
    "        random.seed(1234)\n",
    "        torch.manual_seed(1234)\n",
    "        np.random.seed(1234)\n",
    "        \n",
    "        # Define PDE library that is quadratic in u,\n",
    "        # and fourth-order in spatial derivatives of u.\n",
    "        library_functions = [lambda x: x, lambda x: x * x]\n",
    "        library_function_names = [lambda x: x, lambda x: x + x]\n",
    "        pde_lib = ps.PDELibrary(\n",
    "            library_functions=library_functions,\n",
    "            function_names=library_function_names,\n",
    "            derivative_order=4,\n",
    "            spatial_grid=x,\n",
    "            is_uniform=True,\n",
    "        )\n",
    "\n",
    "        # Again, loop through all the optimizers\n",
    "        print('STLSQ model: ')\n",
    "        optimizer = STLSQ_(threshold=30, normalize_columns=True, k=3, lr=lr, _lambda=_lambda)\n",
    "        model = ps.SINDy(feature_library=pde_lib,\n",
    "                         feature_names=['u'],\n",
    "                         optimizer=optimizer)\n",
    "        model.fit(u, t=dt)\n",
    "        model.print(precision=8)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "a9e7e393",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "_lambda  = 10.0, lr = 10.0\n",
      "STLSQ model: \n",
      "(u)' = -0.99407255 u_11 + -0.99628600 u_1111 + -0.99233116 uu_1\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1217320/3535453974.py:316: ConvergenceWarning: STLSQ._reduce did not converge after 1000 iterations.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "run()# = -0.994 u_11 + -0.996 u_1111 + -0.992 uu_1 with idx 3, 5, 6"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b95bd5c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# -0.994 u_11 + -0.996 u_1111 + -0.992 uu_1"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
