Merge pull request #10 from JuBiotech/sunode-fixes

Fix import route of `sunode`'s `as_aesara` wrapper

Author: lhelleckes

murefi/__init__.py

@@ -11,4 +11,4 @@
 )
 from .ode import BaseODEModel
 
-__version__ = "5.1.0"
+__version__ = "5.1.1"

murefi/datastructures.py

@@ -246,7 +246,12 @@ def make_template(tmin: float, tmax: float, independent_keys: list, rids: list,
 
     @staticmethod
     def make_template_like(
-        dataset, independent_keys: typing.Iterable[str], *, N: int = 200, tmin: typing.Optional[float] = None
+        dataset,
+        independent_keys: typing.Iterable[str],
+        *,
+        N: int = 200,
+        tmin: typing.Optional[float] = None,
+        tmax: typing.Optional[float] = None,
     ):
         """Create a dense template Dataset that has the same start and end times as another Dataset.
 
@@ -255,6 +260,7 @@ def make_template_like(
             independent_keys (list): list of independent variable keys to include in the template
             N (int): total number of timepoints (default: 200)
             tmin (float, optional): override for the start time (when tmin=None, the first timepoint of the template replicate is used)
+            tmax (float): override for the last timepoint
 
         Returns:
             dataset (Dataset): dataset object containing Replicates with dense timeseries of random y data
@@ -263,7 +269,7 @@ def make_template_like(
         for rid, rep in dataset.items():
             ds[rid] = Replicate.make_template(
                 tmin=rep.t_any[0] if tmin is None else tmin,
-                tmax=rep.t_max,
+                tmax=rep.t_max if tmax is None else tmax,
                 independent_keys=independent_keys,
                 rid=rid,
                 N=N,

murefi/symbolic.py

@@ -40,7 +40,7 @@
 
 try:
     import sunode
-    import sunode.wrappers.as_theano
+    from sunode.wrappers import as_aesara
 
     HAS_SUNODE = True
 except ModuleNotFoundError:
@@ -154,13 +154,13 @@ def dydt_dict(t, y, params):
     y0 = named_with_shapes_dict(y0, independent_keys)
     params = named_with_shapes_dict(ode_parameters, parameter_names)
     params["extra"] = numpy.zeros(1)
-    solution, *_ = sunode.wrappers.as_theano.solve_ivp(
+    solution, *_ = as_aesara.solve_ivp(
         y0=y0,
         params=params,
         rhs=dydt_dict,
         tvals=t,
         t0=t[0],
         derivatives="forward",
-        solver_kwargs=dict(sens_mode="simultaneous", compute_sens=True),
+        solver_kwargs=dict(sens_mode="simultaneous"),
     )
     return solution