diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 855635b..6ba49ed 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -28,7 +28,7 @@ jobs: python-version: ${{ matrix.python-version }} - name: Install Dependences run: | - conda install --yes conda-build conda-verify pytest pytest-cov hypothesis statsmodels aesara c-compiler + conda install --yes conda-build conda-verify pytest pytest-cov hypothesis statsmodels pytensor c-compiler - name: Build package run: | conda build --variants "{python: [${{ matrix.python-version }}]}" ./sunode/conda diff --git a/README.md b/README.md index c5ac1d3..9e1202b 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ functions using symbolic differentiation and common subexpression elimination. In either case the functions are compiled using numba and the resulting C-function is passed to sunode, so that there is no python overhead. -`sunode` comes with an Aesara wrapper so that parameters of an ODE can be estimated +`sunode` comes with an PyTensor wrapper so that parameters of an ODE can be estimated using PyMC. ### Installation @@ -134,7 +134,7 @@ We'll use some time artificial data: ```python import numpy as np import sunode -import sunode.wrappers.as_aesara +import sunode.wrappers.as_pytensor import pymc as pm times = np.arange(1900,1921,1) @@ -186,17 +186,17 @@ with pm.Model() as model: gamma = pm.Deterministic('gamma', freq / speed_ratio / ratio) delta = pm.Deterministic('delta', freq / speed_ratio / fixed_hares / ratio) - y_hat, _, problem, solver, _, _ = sunode.wrappers.as_aesara.solve_ivp( + y_hat, _, problem, solver, _, _ = sunode.wrappers.as_pytensor.solve_ivp( y0={ # The initial conditions of the ode. Each variable - # needs to specify a theano or numpy variable and a shape. + # needs to specify a PyTensor or numpy variable and a shape. # This dict can be nested. 'hares': (hares_start, ()), 'lynx': (lynx_start, ()), }, params={ # Each parameter of the ode. sunode will only compute derivatives - # with respect to theano variables. The shape needs to be specified + # with respect to PyTensor variables. The shape needs to be specified # as well. It it infered automatically for numpy variables. # This dict can be nested. 'alpha': (alpha, ()), diff --git a/conda/meta.yaml b/conda/meta.yaml index cd3cfd5..16fcfa3 100644 --- a/conda/meta.yaml +++ b/conda/meta.yaml @@ -42,7 +42,7 @@ test: - pytest - hypothesis - statsmodels - - aesara + - pytensor commands: - pytest --pyargs sunode diff --git a/doc/source/quickstart_pymc.rst b/doc/source/quickstart_pymc.rst index fce6a1e..ff6d271 100644 --- a/doc/source/quickstart_pymc.rst +++ b/doc/source/quickstart_pymc.rst @@ -42,7 +42,7 @@ We'll use some time artificial data::: import numpy as np import sunode - import sunode.wrappers.as_aesara + import sunode.wrappers.as_pytensor import pymc as pm times = np.arange(1900,1921,1) @@ -124,7 +124,7 @@ Now, we define the names, (symbolic) values and shapes of the parameters and ini We solve the ODE using the ``solve_ivp`` function from sunode:: with model: - from sunode.wrappers.as_aesara import solve_ivp + from sunode.wrappers.as_pytensor import solve_ivp solution, *_ = solve_ivp( y0=y0, params=params, diff --git a/doc/source/without_pymc.rst b/doc/source/without_pymc.rst index 8d0868d..4392cf1 100644 --- a/doc/source/without_pymc.rst +++ b/doc/source/without_pymc.rst @@ -12,7 +12,7 @@ current number, and die when eaten by a lynx. We get: .. math:: \frac{dH}{dt} = \alpha H - \beta LH \\ \frac{dL}{dt} = \delta LH - \gamma L -If we want to solve this ODE without the support of Aesara or PyMC, we need to +If we want to solve this ODE without the support of PyTensor or PyMC, we need to first declare the parameters and states we are using. We have four parameters and two states, and each one is a scalar values, so it has shape ():: @@ -71,7 +71,7 @@ ODE might look like this:: This right-hand-side function is usually only called once to collect the sympy expressions of the derivatives. Control flow within this function might behave in unexpected ways if you are new to this concept. It is the - same thing as with theano, pytorch or tensorflow in graph mode. This means + same thing as with PyTensor, pytorch or tensorflow in graph mode. This means that something like this will **not** work as expected:: value = 1 diff --git a/mypy.ini b/mypy.ini index 6900e7b..d8c2584 100644 --- a/mypy.ini +++ b/mypy.ini @@ -20,10 +20,10 @@ ignore_missing_imports = True [mypy-pandas] ignore_missing_imports = True -[mypy-theano] +[mypy-pytensor] ignore_missing_imports = True -[mypy-theano.tensor] +[mypy-pytensor.tensor] ignore_missing_imports = True [mypy-sympy.printing.pycode] diff --git a/notebooks/from_sympy.ipynb b/notebooks/from_sympy.ipynb index 2921a48..833faa5 100644 --- a/notebooks/from_sympy.ipynb +++ b/notebooks/from_sympy.ipynb @@ -22,17 +22,18 @@ "\n", "import sunode.symode.paramset\n", "import sunode.symode.problem\n", - "import sunode.wrappers.as_aesara\n", + "import sunode.wrappers.as_pytensor\n", "\n", - "import aesara\n", - "import aesara.tensor as aet" + "import pytensor\n", + "import pytensor.tensor as pt" ] }, { + "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ - "### Define ode using sympy and aesara" + "### Define ODE using sympy and Pytensor" ] }, { @@ -41,8 +42,8 @@ "metadata": {}, "outputs": [], "source": [ - "b = aet.dvector('b')\n", - "d = aet.dvector('d')\n", + "b = pt.dvector('b')\n", + "d = pt.dvector('d')\n", "\n", "def rhs(t, y, params):\n", " return {\n", @@ -53,7 +54,7 @@ " }\n", "\n", "\n", - "solution, problem, solver = sunode.wrappers.as_aesara.solve_ivp(\n", + "solution, problem, solver = sunode.wrappers.as_pytensor.solve_ivp(\n", " t0=0,\n", " y0={\n", " 'a': (np.arange(3, dtype=float) + d[0] ** 2, 3),\n", @@ -185,7 +186,7 @@ "outputs": [], "source": [ "val = (solution ** 2).sum()\n", - "grads = aet.grad(val, [b, d])" + "grads = pt.grad(val, [b, d])" ] }, { @@ -194,8 +195,8 @@ "metadata": {}, "outputs": [], "source": [ - "func = aesara.function([b, d], [val] + grads)\n", - "func2 = aesara.function([b, d], [val])" + "func = pytensor.function([b, d], [val] + grads)\n", + "func2 = pytensor.function([b, d], [val])" ] }, { @@ -661,10 +662,11 @@ ] }, { + "attachments": {}, "cell_type": "markdown", "metadata": {}, "source": [ - "### Integrate into aesara and check gradients" + "### Integrate into PyTensor and check gradients" ] }, { @@ -682,23 +684,23 @@ "metadata": {}, "outputs": [], "source": [ - "params = aet.dvector('params')\n", - "y0 = aet.dvector('y0')\n", - "solve_ode = sunode.wrappers.as_aesara.SolveODEAdjoint(solver, 0, tvals)\n", + "params = pt.dvector('params')\n", + "y0 = pt.dvector('y0')\n", + "solve_ode = sunode.wrappers.as_pytensor.SolveODEAdjoint(solver, 0, tvals)\n", "solution = solve_ode(y0, params)\n", "\n", "loss = (solution ** 2).sum()\n", - "grad_p, grad_y0 = aet.grad(loss, [params, y0])\n", - "func = aesara.function([y0, params], [loss, grad_p, grad_y0])\n", + "grad_p, grad_y0 = pt.grad(loss, [params, y0])\n", + "func = pytensor.function([y0, params], [loss, grad_p, grad_y0])\n", "\n", "# Explicit solution\n", "loss = ((\n", " ((0.5 * tvals ** 2 * params[1] + tvals * y0[1]) + y0[0]) ** 2\n", " + (tvals * params[1] + y0[1]) ** 2\n", ")).sum()\n", - "grad_p, grad_y0 = aet.grad(loss, [params, y0])\n", + "grad_p, grad_y0 = pt.grad(loss, [params, y0])\n", "\n", - "func2 = aesara.function([y0, params], [loss, grad_p, grad_y0])" + "func2 = pytensor.function([y0, params], [loss, grad_p, grad_y0])" ] }, { @@ -842,7 +844,7 @@ " params = pm.Normal('params', sigma=10, shape=ode.n_params)\n", " y0 = pm.Normal('y0', shape=ode.n_states)\n", " \n", - " solve_ode = sunode.wrappers.as_aesara.SolveODEAdjoint(solver, 0, tvals)\n", + " solve_ode = sunode.wrappers.as_pytensor.SolveODEAdjoint(solver, 0, tvals)\n", " mu = solve_ode(y0, params)\n", " error = 0.8 * np.random.randn(len(tvals))\n", " pm.Normal('y', mu=mu[:, 0], sigma=0.8, observed=tvals ** 2 + tvals + 5 + error)\n", diff --git a/notebooks/pymc_model.ipynb b/notebooks/pymc_model.ipynb index 3816b20..9b9191b 100644 --- a/notebooks/pymc_model.ipynb +++ b/notebooks/pymc_model.ipynb @@ -25,7 +25,7 @@ "outputs": [], "source": [ "import os\n", - "os.environ[\"AESARA_FLAGS\"] = \"floatX=float64\"" + "os.environ[\"pytensor_FLAGS\"] = \"floatX=float64\"" ] }, { @@ -48,8 +48,8 @@ "from scipy.integrate import odeint\n", "\n", "import pymc as pm\n", - "import aesara\n", - "import aesara.tensor as aet\n", + "import pytensor\n", + "import pytensor.tensor as pt\n", "\n", "# this notebook show DEBUG log messages\n", "# logging.getLogger('pymc').setLevel(logging.DEBUG)\n", @@ -150,9 +150,9 @@ "outputs": [], "source": [ "# To demonstrate that test-value computation works, but also for debugging\n", - "aesara.config.compute_test_value = 'raise'\n", - "aesara.config.exception_verbosity = 'high'\n", - "aesara.config.traceback.limit = 5" + "pytensor.config.compute_test_value = 'raise'\n", + "pytensor.config.exception_verbosity = 'high'\n", + "pytensor.config.traceback.limit = 5" ] }, { @@ -176,7 +176,7 @@ "metadata": {}, "outputs": [], "source": [ - "import sunode.wrappers.as_aesara\n", + "import sunode.wrappers.as_pytensor\n", "\n", "\n", "def get_model_sunode():\n", @@ -187,7 +187,7 @@ " s0 = pm.Normal('red_0', mu=10, sigma=2)\n", " extra = pm.Normal('extra', shape=n_extra)\n", "\n", - " y_hat, _, _ = sunode.wrappers.as_aesara.solve_ivp(\n", + " y_hat, _, _ = sunode.wrappers.as_pytensor.solve_ivp(\n", " y0={\n", " 'S': (s0, ()), # TODO Infer shape from model?\n", " 'P': np.array(y0_true[1], dtype='d'),\n", @@ -196,7 +196,7 @@ " params={\n", " 'K_S': (K_S, ()),\n", " 'vmax': (vmax, ()),\n", - " 'tmp': np.zeros(1), # TODO aesara wants at least one fixed param\n", + " 'tmp': np.zeros(1), # TODO pytensor wants at least one fixed param\n", " 'extra_p': (extra, (n_extra,))\n", " },\n", " rhs=reaction_sympy,\n", @@ -235,7 +235,7 @@ " \n", " extra = pm.Normal('extra', shape=n_extra)\n", "\n", - " y_hat, problem, _ = sunode.wrappers.as_aesara.solve_ivp(\n", + " y_hat, problem, _ = sunode.wrappers.as_pytensor.solve_ivp(\n", " y0={\n", " 'S': (s0, ()), # TODO Infer shape from model?\n", " 'P': np.array(y0_true[1], dtype='d'),\n", @@ -244,7 +244,7 @@ " params={\n", " 'K_S': (K_S, ()),\n", " 'vmax': (vmax, ()),\n", - " 'tmp': np.zeros(1), # TODO aesara wants at least one fixed param\n", + " 'tmp': np.zeros(1), # TODO pytensor wants at least one fixed param\n", " 'extra_p': (extra, (n_extra,))\n", " },\n", " rhs=reaction_sympy,\n", @@ -305,22 +305,22 @@ " \n", " # create a test function for evaluating the logp value\n", " print('Compiling f_logpt')\n", - " f_logpt = aesara.function(\n", + " f_logpt = pytensor.function(\n", " inputs=t_inputs,\n", " outputs=[pmodel.logpt],\n", " # with float32, allow downcast because the forward integration is always float64\n", - " allow_input_downcast=(aesara.config.floatX == 'float32')\n", + " allow_input_downcast=(pytensor.config.floatX == 'float32')\n", " )\n", " print(f'Test logpt:')\n", " print(f_logpt(*test_inputs))\n", " \n", " # and another test function for evaluating the gradient\n", " print('Compiling f_logpt')\n", - " f_grad = aesara.function(\n", + " f_grad = pytensor.function(\n", " inputs=t_inputs,\n", - " outputs=aet.grad(pmodel.logpt, t_inputs),\n", + " outputs=pt.grad(pmodel.logpt, t_inputs),\n", " # with float32, allow downcast because the forward integration is always float64\n", - " allow_input_downcast=(aesara.config.floatX == 'float32')\n", + " allow_input_downcast=(pytensor.config.floatX == 'float32')\n", " )\n", " print(f'Test gradient:')\n", " print(f_grad(*test_inputs))\n", @@ -984,7 +984,7 @@ "metadata": {}, "outputs": [], "source": [ - "aesara.printing.pydotprint(aet.grad(model_sunode.logpt, model_sunode.vmax), 'ODE_API_shapes_and_benchmarking.png')\n", + "pytensor.printing.pydotprint(pt.grad(model_sunode.logpt, model_sunode.vmax), 'ODE_API_shapes_and_benchmarking.png')\n", "IPython.display.Image('ODE_API_shapes_and_benchmarking.png')" ] }, @@ -1007,7 +1007,7 @@ "metadata": {}, "outputs": [], "source": [ - "from aesara import d3viz\n", + "from pytensor import d3viz\n", "d3viz.d3viz(model.logpt, 'ODE_API_shapes_and_benchmarking.html')" ] }, @@ -1041,4 +1041,3 @@ "nbformat": 4, "nbformat_minor": 4 } - diff --git a/sunode/test_aesara.py b/sunode/test_pytensor.py similarity index 78% rename from sunode/test_aesara.py rename to sunode/test_pytensor.py index 966d799..8ea46a4 100644 --- a/sunode/test_aesara.py +++ b/sunode/test_pytensor.py @@ -1,6 +1,6 @@ import numpy as np -import aesara -import aesara.tensor as aet +import pytensor +import pytensor.tensor as pt import sunode.wrappers @@ -12,7 +12,7 @@ def dydt_dict(t, y, p): 'B': y.B, 'C': y.C, } - A = aet.dscalar("A") + A = pt.dscalar("A") A.tag.test_value = np.array(0.9) @@ -30,7 +30,7 @@ def dydt_dict(t, y, p): 'extra': np.array([0.]) } - solution, *_ = sunode.wrappers.as_aesara.solve_ivp( + solution, *_ = sunode.wrappers.as_pytensor.solve_ivp( y0=y0, params=params, rhs=dydt_dict, @@ -40,5 +40,5 @@ def dydt_dict(t, y, p): solver_kwargs=dict(sens_mode="simultaneous") ) - func = aesara.function([A], [solution["A"], solution["B"]]) + func = pytensor.function([A], [solution["A"], solution["B"]]) assert func(0.2)[0].shape == time.shape \ No newline at end of file diff --git a/sunode/wrappers/__init__.py b/sunode/wrappers/__init__.py index 064c77d..9fb2e5e 100644 --- a/sunode/wrappers/__init__.py +++ b/sunode/wrappers/__init__.py @@ -1,4 +1,4 @@ -from . import as_aesara -from . import as_aesara as as_theano +from . import as_pytensor +from . import as_pytensor as as_aesara -__all__ = ['as_theano', 'as_aesara'] \ No newline at end of file +__all__ = ['as_aesara', 'as_pytensor'] \ No newline at end of file diff --git a/sunode/wrappers/as_aesara.py b/sunode/wrappers/as_pytensor.py similarity index 86% rename from sunode/wrappers/as_aesara.py rename to sunode/wrappers/as_pytensor.py index 841b23b..32d7d1b 100644 --- a/sunode/wrappers/as_aesara.py +++ b/sunode/wrappers/as_pytensor.py @@ -1,22 +1,15 @@ try: - import aesara - import aesara.tensor as aet + import pytensor.tensor as pt + from pytensor.graph.basic import Constant, Variable + from pytensor.graph.fg import MissingInputError + from pytensor.graph.op import Op + from pytensor.gradient import grad_not_implemented +except ModuleNotFoundError: + import aesara.tensor as pt from aesara.graph.basic import Constant, Variable from aesara.graph.fg import MissingInputError from aesara.graph.op import Op from aesara.gradient import grad_not_implemented -except ModuleNotFoundError: - import theano - import theano.tensor as aet - from theano.gradient import grad_not_implemented - if hasattr(theano, "gof"): - from theano.gof.fg import MissingInputError - from theano.gof.var import Constant, Variable - from theano.gof.op import Op - else: - from theano.graph.basic import Constant, Variable - from theano.graph.fg import MissingInputError - from theano.graph.op import Op import copy from typing import Dict, Optional, Any, Callable @@ -60,14 +53,14 @@ def read_dict(vals, name=None): tensor, dim_names = vals else: try: - tensor, dim_names = vals, aet.as_tensor_variable(vals).shape.eval() + tensor, dim_names = vals, pt.as_tensor_variable(vals).shape.eval() except MissingInputError as e: raise ValueError( 'Shapes of tensors need to be statically ' 'known or given explicitly.') from e if isinstance(dim_names, (str, int)): dim_names = (dim_names,) - tensor = aet.as_tensor_variable(tensor) + tensor = pt.as_tensor_variable(tensor) if tensor.ndim != len(dim_names): raise ValueError( f"Dimension mismatch for {name}: Value has rank {tensor.ndim}, " @@ -104,22 +97,22 @@ def read_dict(vals, name=None): tensor = flat_tensors[path] if isinstance(tensor, tuple): tensor, _ = tensor - vars.append(aet.as_tensor_variable(tensor).reshape((-1,))) + vars.append(pt.as_tensor_variable(tensor).reshape((-1,))) if vars: - params_subs_flat = aet.concatenate(vars) + params_subs_flat = pt.concatenate(vars) else: - params_subs_flat = aet.as_tensor_variable(np.zeros(0)) + params_subs_flat = pt.as_tensor_variable(np.zeros(0)) vars = [] for path in problem.params_subset.remainder.subset_paths: tensor = flat_tensors[path] if isinstance(tensor, tuple): tensor, _ = tensor - vars.append(aet.as_tensor_variable(tensor).reshape((-1,))) + vars.append(pt.as_tensor_variable(tensor).reshape((-1,))) if vars: - params_remaining_flat = aet.concatenate(vars) + params_remaining_flat = pt.concatenate(vars) else: - params_remaining_flat = aet.as_tensor_variable(np.zeros(0)) + params_remaining_flat = pt.as_tensor_variable(np.zeros(0)) flat_tensors = as_flattened(y0) vars = [] @@ -127,8 +120,8 @@ def read_dict(vals, name=None): tensor = flat_tensors[path] if isinstance(tensor, tuple): tensor, _ = tensor - vars.append(aet.as_tensor_variable(tensor).reshape((-1,))) - y0_flat = aet.concatenate(vars) + vars.append(pt.as_tensor_variable(tensor).reshape((-1,))) + y0_flat = pt.concatenate(vars) if derivatives == 'adjoint': sol = solver.AdjointSolver(problem, **solver_kwargs) @@ -150,8 +143,8 @@ def read_dict(vals, name=None): class SolveODE(Op): - itypes = [aet.dvector, aet.dvector, aet.dvector] - otypes = [aet.dmatrix, aet.dtensor3] + itypes = [pt.dvector, pt.dvector, pt.dvector] + otypes = [pt.dmatrix, pt.dtensor3] __props__ = ('_solver_id', '_t0', '_tvals_id') @@ -218,15 +211,15 @@ def grad(self, inputs, g): assert str(g_grad) == '' solution, sens = self(*inputs) return [ - aet.zeros_like(inputs[0]), - aet.sum(g[:, None, :] * sens, (0, -1)), + pt.zeros_like(inputs[0]), + pt.sum(g[:, None, :] * sens, (0, -1)), grad_not_implemented(self, 2, inputs[-1]) ] class SolveODEAdjoint(Op): - itypes = [aet.dvector, aet.dvector, aet.dvector] - otypes = [aet.dmatrix] + itypes = [pt.dvector, pt.dvector, pt.dvector] + otypes = [pt.dmatrix] __props__ = ('_solver_id', '_t0', '_tvals_id') @@ -263,8 +256,8 @@ def grad(self, inputs, g): class SolveODEAdjointBackward(Op): - itypes = [aet.dvector, aet.dvector, aet.dvector, aet.dmatrix] - otypes = [aet.dvector, aet.dvector] + itypes = [pt.dvector, pt.dvector, pt.dvector, pt.dmatrix] + otypes = [pt.dvector, pt.dvector] __props__ = ('_solver_id', '_t0', '_tvals_id')