[Documentation] Update NumPy notebooks and add PyTorch examples

.devcontainer/devcontainer.json

@@ -1,7 +1,7 @@
 {
 	// For quick reference: https://containers.dev/implementors/json_reference/
-	"name": "Miniforge3",
-	"image": "condaforge/miniforge3:24.3.0-0",
+	"name": "Python 3",
+	"image": "mcr.microsoft.com/devcontainers/python:3.12-bullseye",
 	//"build": {
 	//	"context": "..",
 	//	"dockerfile": "Dockerfile"
@@ -22,7 +22,7 @@
 
 	// 4. Features to add to the Dev Container. More info: https://containers.dev/implementors/features.
 	"features": {
-		// ZSH without OMZ
+		// ZSH with OMZ
 		"ghcr.io/devcontainers/features/common-utils:2": {
 			"installZsh": true,
 			"configureZshAsDefaultShell": true,
@@ -41,15 +41,16 @@
 		"ghcr.io/schlich/devcontainer-features/powerlevel10k:1": {},
 		"ghcr.io/nils-geistmann/devcontainers-features/zsh:0": {
 			"setLocale": true,
-			"theme": "robbyrussell",
+			"theme": "agnoster",
 			"plugins": "git docker",
 			"desiredLocale": "en_US.UTF-8 UTF-8"
 		},
 		"ghcr.io/devcontainers-extra/features/zsh-plugins:0": {
 			"plugins": "ssh-agent npm",
 			"omzPlugins": "https://github.com/zsh-users/zsh-autosuggestions",
 			"username": "vscode"
-		}
+		},
+		"ghcr.io/devcontainers-extra/features/hatch:2": {}
 	},
 
 	// 5. Configure tool-specific properties.
@@ -78,7 +79,7 @@
 	},
 
 	// 6. Set `remoteUser` to `root` to connect as root instead. More info: https://aka.ms/vscode-remote/containers/non-root.
-	"remoteUser": "vscode"
+	"remoteUser": "vscode",
 
 	// the following commands are related to container lifecylce. More info: https://containers.dev/implementors/json_reference/#lifecycle-scripts
 
@@ -92,7 +93,7 @@
 	// "updateContentCommand": "",
 
 	// 10. Use 'postCreateCommand' to run commands after the container is created.
-	// "postCreateCommand": "pip3 install --user -r requirements.txt",
+	"postCreateCommand": "sudo apt -qq update && sudo apt install -qq ffmpeg -y"
 
 	// 11. Use 'postStartCommand' to run a command each time the container starts successfully.
 	// "postStartCommand": "",

desolver/backend/numpy_backend.py

@@ -6,6 +6,7 @@
 import scipy.sparse
 import scipy.sparse.linalg
 import autoray
+import contextlib
 
 
 def __solve_linear_system(A,b,overwrite_a=False,overwrite_b=False,check_finite=False,sparse=False):
@@ -16,3 +17,12 @@ def __solve_linear_system(A,b,overwrite_a=False,overwrite_b=False,check_finite=F
 
 
 autoray.register_function("numpy", "solve_linear_system", __solve_linear_system)
+
+
+@contextlib.contextmanager
+def __no_grad_ctx():
+    yield
+
+autoray.register_function("numpy", "no_grad", __no_grad_ctx)
+autoray.register_function("builtins", "no_grad", __no_grad_ctx)
+autoray.register_function("numpy", "clone", numpy.copy)

desolver/differential_system.py

@@ -517,7 +517,7 @@ def __init__(self, equ_rhs, y0, t=(0, 1), dense_output=False, dt=1.0, rtol=None,
         self.__rtol = rtol
         self.__atol = atol
         self.__consts = constants if constants is not None else dict()
-        self.__y = D.ar_numpy.copy(y0)[None]
+        self.__y = D.ar_numpy.clone(y0)[None]
         self.__t = D.ar_numpy.asarray(t[0], **self.__array_con_kwargs)[None]
         self.dim = D.ar_numpy.shape(self.__y[0])
         self.counter = 0
@@ -538,7 +538,6 @@ def __init__(self, equ_rhs, y0, t=(0, 1), dense_output=False, dt=1.0, rtol=None,
         if rhs_shape != y0_shape:
             raise ValueError(f"Expected rhs to return a shape that is compatible with y0, got shapes Shape[y0]={y0_shape} and Shape[dy]={rhs_shape}")
 
-        self.__move_to_device()
         self.__allocate_soln_space(self.__alloc_space_steps(self.tf))
         self.__fix_dt_dir(self.tf, self.t0)
         self.__events = []
@@ -669,7 +668,6 @@ def dt(self):
     @dt.setter
     def dt(self, new_dt):
         self.__dt = D.ar_numpy.asarray(new_dt, **self.__array_con_kwargs)
-        self.__move_to_device()
         self.__fix_dt_dir(self.tf, self.t0)
         return self.__dt
 
@@ -698,7 +696,6 @@ def t0(self, new_t0):
         if D.ar_numpy.abs(self.tf - new_t0) <= D.epsilon(self.__y[self.counter].dtype):
             raise ValueError("The start time of the integration cannot be greater than or equal to {}!".format(self.tf))
         self.__t0 = new_t0
-        self.__move_to_device()
         self.__fix_dt_dir(self.tf, self.t0)
 
     @property
@@ -726,7 +723,6 @@ def tf(self, new_tf):
         if D.ar_numpy.abs(self.t0 - new_tf) <= D.epsilon(self.__y[self.counter].dtype):
             raise ValueError("The end time of the integration cannot be equal to the start time: {}!".format(self.t0))
         self.__tf = new_tf
-        self.__move_to_device()
         self.__fix_dt_dir(self.tf, self.t0)
 
     def __fix_dt_dir(self, t1, t0):
@@ -735,34 +731,6 @@ def __fix_dt_dir(self, t1, t0):
         else:
             self.__dt = self.__dt
 
-    def __move_to_device(self):
-        if self.device is not None and self.__inferred_backend == 'torch':
-            self.__y[0] = self.__y[0].to(self.device)
-            self.__t[0] = self.__t[0].to(self.device)
-            self.__t0 = self.__t0.to(self.device)
-            self.__tf = self.__tf.to(self.device)
-            self.__dt = self.__dt.to(self.device)
-            self.__dt0 = self.__dt0.to(self.device)
-            try:
-                self.__atol = self.__atol.to(self.device)
-            except AttributeError:
-                pass
-            except:
-                raise
-            try:
-                self.__rtol = self.__rtol.to(self.device)
-            except AttributeError:
-                pass
-            except:
-                raise
-            try:
-                self.equ_rhs.rhs = self.equ_rhs.rhs.to(self.device)
-            except AttributeError:
-                pass
-            except:
-                raise
-        return
-
     def __alloc_space_steps(self, tf):
         """Returns the number of steps to allocate for a given final integration time
 
@@ -936,7 +904,6 @@ def reset(self):
         self.__sol = DenseOutput(None, None)
         self.dt = self.__dt0
         self.equ_rhs.nfev = 0
-        self.__move_to_device()
         self.__int_status = 0
         if self.__events:
             self.__events = []

desolver/integrators/components/runge_kutta_methods.py

@@ -42,7 +42,10 @@ def compute_step(rhs: typing.Callable, initial_time, initial_state, timestep, in
     """
 
     for stage in range(intermediate_stages_in.shape[-1]):
-        intermediate_dstate = timestep * D.ar_numpy.sum(intermediate_stages_in * rk_tableau[stage, 1:], axis=-1)
+        stage_coeffs = rk_tableau[stage, 1:]
+        nonzero_coeffs_mask = stage_coeffs != 0.0
+        nonzero_coeffs_mask = D.ar_numpy.where(~D.ar_numpy.any(nonzero_coeffs_mask), D.ar_numpy.ones_like(nonzero_coeffs_mask), nonzero_coeffs_mask)
+        intermediate_dstate = timestep * D.ar_numpy.sum(intermediate_stages_in[...,nonzero_coeffs_mask] * stage_coeffs[nonzero_coeffs_mask], axis=-1)
         intermediate_rhs = rhs(
             initial_time + timestep * rk_tableau[stage, 0], 
             initial_state + intermediate_dstate,

desolver/integrators/integrator_template.py

@@ -44,52 +44,53 @@ def dense_output(self):
         pass
 
     def update_timestep(self, ignore_custom_adaptation=False):
-        if self.adaptation_fn and not ignore_custom_adaptation:
-            return self.adaptation_fn(self)
-        initial_state = self.solver_dict['initial_state']
-        diff = self.solver_dict['diff']
-        timestep = self.solver_dict['timestep']
-        safety_factor = self.solver_dict['safety_factor']
-        atol = self.solver_dict['atol']
-        rtol = self.solver_dict['rtol']
-        dState = self.solver_dict['dState']
-        order = self.solver_dict['order']
-        if "system_scaling" in self.solver_dict:
-            self.solver_dict["system_scaling"] = 0.8 * self.solver_dict["system_scaling"] +  0.2 * D.ar_numpy.maximum(D.ar_numpy.abs(initial_state), D.ar_numpy.abs(dState / timestep))
-        else:
-            self.solver_dict["system_scaling"] = D.ar_numpy.maximum(D.ar_numpy.abs(initial_state), D.ar_numpy.abs(dState / timestep))
-        total_error_tolerance = (atol + rtol * self.solver_dict["system_scaling"])
-        with D.numpy.errstate(divide='ignore'):
-            epsilon_current = D.ar_numpy.reciprocal(D.ar_numpy.linalg.norm(diff / total_error_tolerance))
-        if "epsilon_last" in self.solver_dict:
-            epsilon_last = self.solver_dict["epsilon_last"]
-        else:
-            epsilon_last = None
-        if "epsilon_last_last" in self.solver_dict:
-            epsilon_last_last = self.solver_dict["epsilon_last_last"]
-        else:
-            epsilon_last_last = None
-        if epsilon_last is None:
-            corr = D.ar_numpy.where(epsilon_current > 0.0, epsilon_current ** (1.0 / order), 1.0)
-            self.solver_dict["epsilon_last"] = epsilon_current
-        elif epsilon_last is not None and epsilon_last_last is None:
-            corr = D.ar_numpy.where(epsilon_current > 0.0, epsilon_current ** (1.0 / order), 1.0)
-            corr = corr*D.ar_numpy.where(epsilon_last > 0.0, epsilon_last ** (1.0 / order), 1.0)
-            self.solver_dict["epsilon_last_last"], self.solver_dict["epsilon_last"] = epsilon_last, epsilon_current
-        elif epsilon_last is not None and epsilon_last_last is not None:
-            # Based on the triple product described in https://link.springer.com/article/10.1007/s42967-021-00159-w
-            # Eq. (6) with the coefficients from the second entry of Table 4
-            k1, k2, k3 = self.solver_dict.get("__adapt_k1", 0.55), self.solver_dict.get("__adapt_k2", -0.27), self.solver_dict.get("__adapt_k3", 0.05)
-            k1 = epsilon_current ** (k1 / order)
-            k2 = epsilon_last ** (k2 / order)
-            k3 = epsilon_last_last ** (k3 / order)
-            corr = D.ar_numpy.where(k1 > 0.0, k1, 1.0)
-            corr = corr*D.ar_numpy.where(k2 > 0.0, k2, 1.0)
-            corr = corr*D.ar_numpy.where(k3 > 0.0, k3, 1.0)
-            self.solver_dict["epsilon_last_last"], self.solver_dict["epsilon_last"] = epsilon_last, epsilon_current
-        corr = (1 + D.ar_numpy.arctan((safety_factor * corr - 1)))
-        timestep = corr * timestep
-        return timestep, bool(corr < 0.9**2)
+        with D.ar_numpy.no_grad(like=self.solver_dict['initial_state']):
+            if self.adaptation_fn and not ignore_custom_adaptation:
+                return self.adaptation_fn(self)
+            initial_state = self.solver_dict['initial_state']
+            diff = self.solver_dict['diff']
+            timestep = self.solver_dict['timestep']
+            safety_factor = self.solver_dict['safety_factor']
+            atol = self.solver_dict['atol']
+            rtol = self.solver_dict['rtol']
+            dState = self.solver_dict['dState']
+            order = self.solver_dict['order']
+            if "system_scaling" in self.solver_dict:
+                self.solver_dict["system_scaling"] = 0.8 * self.solver_dict["system_scaling"] +  0.2 * D.ar_numpy.maximum(D.ar_numpy.abs(initial_state), D.ar_numpy.abs(dState / timestep))
+            else:
+                self.solver_dict["system_scaling"] = D.ar_numpy.maximum(D.ar_numpy.abs(initial_state), D.ar_numpy.abs(dState / timestep))
+            total_error_tolerance = (atol + rtol * self.solver_dict["system_scaling"])
+            with D.numpy.errstate(divide='ignore'):
+                epsilon_current = D.ar_numpy.reciprocal(D.ar_numpy.linalg.norm(diff / total_error_tolerance))
+            if "epsilon_last" in self.solver_dict:
+                epsilon_last = self.solver_dict["epsilon_last"]
+            else:
+                epsilon_last = None
+            if "epsilon_last_last" in self.solver_dict:
+                epsilon_last_last = self.solver_dict["epsilon_last_last"]
+            else:
+                epsilon_last_last = None
+            if epsilon_last is None:
+                corr = D.ar_numpy.where(epsilon_current > 0.0, epsilon_current ** (1.0 / order), 1.0)
+                self.solver_dict["epsilon_last"] = epsilon_current
+            elif epsilon_last is not None and epsilon_last_last is None:
+                corr = D.ar_numpy.where(epsilon_current > 0.0, epsilon_current ** (1.0 / order), 1.0)
+                corr = corr*D.ar_numpy.where(epsilon_last > 0.0, epsilon_last ** (1.0 / order), 1.0)
+                self.solver_dict["epsilon_last_last"], self.solver_dict["epsilon_last"] = epsilon_last, epsilon_current
+            elif epsilon_last is not None and epsilon_last_last is not None:
+                # Based on the triple product described in https://link.springer.com/article/10.1007/s42967-021-00159-w
+                # Eq. (6) with the coefficients from the second entry of Table 4
+                k1, k2, k3 = self.solver_dict.get("__adapt_k1", 0.55), self.solver_dict.get("__adapt_k2", -0.27), self.solver_dict.get("__adapt_k3", 0.05)
+                k1 = epsilon_current ** (k1 / order)
+                k2 = epsilon_last ** (k2 / order)
+                k3 = epsilon_last_last ** (k3 / order)
+                corr = D.ar_numpy.where(k1 > 0.0, k1, 1.0)
+                corr = corr*D.ar_numpy.where(k2 > 0.0, k2, 1.0)
+                corr = corr*D.ar_numpy.where(k3 > 0.0, k3, 1.0)
+                self.solver_dict["epsilon_last_last"], self.solver_dict["epsilon_last"] = epsilon_last, epsilon_current
+            corr = (1 + D.ar_numpy.arctan((safety_factor * corr - 1)))
+            timestep = corr * timestep
+            return timestep, bool(corr < 0.9**2)
 
     def get_error_estimate(self):
         return 0.0

desolver/integrators/utilities.py

@@ -3,7 +3,6 @@
 from desolver.integrators.integrator_types import TableauIntegrator
 
 def implicit_aware_update_timestep(integrator: TableauIntegrator):
-    timestep = integrator.solver_dict['timestep']
     timestep_from_error, redo_step = integrator.update_timestep(ignore_custom_adaptation=True)
     if "niter0" in integrator.solver_dict.keys():
         # Adjust the timestep according to the computational cost of 

docs/examples/getting_started.py

@@ -1,5 +1,5 @@
 import desolver as de
-import desolver.backend as D
+import numpy as np
 
 @de.rhs_prettifier(
     equ_repr="[vx, -k*x/m]",
@@ -13,11 +13,11 @@
 """
 )
 def rhs(t, state, k, m, **kwargs):
-    return D.array([[0.0, 1.0], [-k/m,  0.0]])@state
+    return np.array([[0.0, 1.0], [-k/m,  0.0]])@state
 
-y_init = D.array([1., 0.])
+y_init = np.array([1., 0.])
 
-a = de.OdeSystem(rhs, y0=y_init, dense_output=True, t=(0, 2*D.pi), dt=0.01, rtol=1e-9, atol=1e-9, constants=dict(k=1.0, m=1.0))
+a = de.OdeSystem(rhs, y0=y_init, dense_output=True, t=(0, 2*np.pi), dt=0.01, rtol=1e-9, atol=1e-9, constants=dict(k=1.0, m=1.0))
 
 print(a)
 
@@ -29,4 +29,4 @@ def rhs(t, state, k, m, **kwargs):
 print("a[0].y  = {}".format(a[0].y))
 print("a[-1].y = {}".format(a[-1].y))
 
-print("Maximum difference from initial state after one oscillation cycle: {}".format(D.max(D.abs(a[0].y-a[-1].y))))
+print("Maximum difference from initial state after one oscillation cycle: {}".format(np.max(np.abs(a[0].y-a[-1].y))))