{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# ASM1: equations for the 0-level set of teh 2nd time derivative of SO\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# Copyright (C) 2025 Juan Pablo Carbajal\n# Copyright (C) 2025 Mariane Yvonne Schneider\n#\n# This program is free software; you can redistribute it and/or modify\n# it under the terms of the GNU General Public License as published by\n# the Free Software Foundation; either version 3 of the License, or\n# (at your option) any later version.\n#\n# This program is distributed in the hope that it will be useful,\n# but WITHOUT ANY WARRANTY; without even the implied warranty of\n# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the\n# GNU General Public License for more details.\n#\n# You should have received a copy of the GNU General Public License\n# along with this program. If not, see <http://www.gnu.org/licenses/>.\n\n# Author: Juan Pablo Carbajal <ajuanpi@gmail.com>\n\nimport matplotlib.pyplot as plt\nfrom sympy import *\n\n\ntry:\n    import dsafeatures\nexcept ModuleNotFoundError:\n    import sys\n    import os\n\n    sys.path.insert(0, os.path.abspath(\"../..\"))\n\nfrom dsafeatures.printing import *\nfrom dsafeatures.odemodel import ODEModel\n\ninit_printing()\n\nplt.style.use('figures.mplstyle')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Select model\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "model = ODEModel.from_csv(name='ASM1')\n\nX = model.states\nr_sym, r = model.rates, model.rates_expr\nM = model.coef_matrix\nO2, idxO2 = model.state_by_name(\"S_O2\")\n\n# Aeration\n# Define aeration\nO2_max = Symbol(r'\\text{SO}_\\text{sat}', real=True, positive=True)\nmax_O2_inflow = Symbol(r'\\dot{q}_{\\text{SO}\\text{max}}', real=True, positive=True)\naer = Function(r'\\dot{q}_\\text{SO}', real=True, postivive=True)(O2, O2_max, max_O2_inflow)\n\n# Extended system\nr_sym_ext = r_sym.col_join(Matrix([aer]))\none_O2 = eye(len(X))[:, [idxO2]]\nM = M.row_join(one_O2)\n\n# First derivative\ndotX = M @ r_sym_ext\n# 2nd derivative os SO\nJr = r_sym_ext.jacobian(X)\nddotO2 = (M[idxO2, :] @ (Jr @ dotX))[0]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We just use numbered variables for states and paramters\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "x = symbols(fr'x1:{len(X)+1}', real=True, positive=True)\nstate_subs = dict(zip(X, x))\np = symbols(fr'p1:{len(model.parameters())+1}', real=True, positive=True)\nparam_subs = dict(zip(model.parameters(), p))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We use static symbols for the rates and their derivatives\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "f = symbols(fr'r1:{len(r_sym_ext)+1}', real=True)\nrate_subs = dict(zip(r_sym_ext, f))\npartials = [d for d in Jr.flat() if not d.is_zero]\ndf = [Symbol(fr'\\,{{}}_{{{latex(state_subs[d.args[1][0]])}}}{latex(rate_subs[d.args[0]])}', real=True) for d in partials]\ndr_subs = dict(zip(partials, df))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "replace these on the 2nd derivative\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ddO2_s = ddotO2.subs(dr_subs).subs(rate_subs).subs(param_subs)\n\n# collect variables until nothing changes\nddO2_vars = [s for s in ddO2_s.free_symbols if s not in param_subs.values()]\nddO2_s = ddO2_s.expand()\nwhile (expr_ := ddO2_s.collect(ddO2_vars)) != ddO2_s:\n    ddO2_s = expr_"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Further, we replace constant expressions with new parameter names.\nTo do this we need to detect all expressions that contain only parameters\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def has_only(expr, syms):\n    \"\"\"Check if expressions has only syms.\"\"\"\n    return not bool(expr.free_symbols - set(syms))\n\ndef subs_expr_withall(expr, syms, x):\n    if has_only(expr, syms):\n        return x\n\n\npnew = numbered_symbols(prefix='k', start=1)\npbundle_subs = {}\nfor a in preorder_traversal(ddO2_s):\n    if a.is_Atom or not has_only(a, p):\n        continue\n    elif a not in pbundle_subs.keys():\n        p_ = next(pnew)\n        pbundle_subs[a] = p_"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We now replace the new symbols.\nHowever, this might leave us with repeated new constants.\nHence, we find the remaining constants and renumber them\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ddO2_s = ddO2_s.subs(pbundle_subs).collect(pbundle_subs.values())\nremain_ = ddO2_s.free_symbols - set(ddO2_vars)\nsubs_ = {r: s for r, s in zip(remain_, numbered_symbols(prefix='C', start=1))}\nddO2_s = ddO2_s.subs(subs_)\npbundle_subs = {k:subs_[v] for k,v in pbundle_subs.items() if v in subs_.keys()}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Final expression\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "ddO2_s.expand().collect(ddO2_vars).collect(pbundle_subs.values())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "plt.show()"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3",
      "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.13.2"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}